added test scenario - concurrent threads trying to acquire jobs

Author: natrajmolala

build.gradle

@@ -6,7 +6,7 @@ plugins {
 }
 
 group = 'uk.gov.hmcts.cp'
-version = System.getProperty('API_SPEC_VERSION') ?: '1.0.7'
+version = System.getProperty('API_SPEC_VERSION') ?: '1.0.0'
 description = 'Spring Job Scheduler'
 
 def githubActor = project.findProperty("github.actor") ?: System.getenv("GITHUB_ACTOR")

integration-tests/src/test/java/uk/gov/hmcts/cp/taskmanager/integration/JobConcurrencyIntegrationTest.java

@@ -0,0 +1,83 @@
+package uk.gov.hmcts.cp.taskmanager.integration;
+
+import static java.util.UUID.randomUUID;
+import static java.util.concurrent.Executors.newFixedThreadPool;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import uk.gov.hmcts.cp.taskmanager.integration.config.IntegrationTestConfiguration;
+import uk.gov.hmcts.cp.taskmanager.persistence.entity.Job;
+import uk.gov.hmcts.cp.taskmanager.persistence.service.JobService;
+
+import java.time.ZonedDateTime;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import jakarta.json.Json;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.transaction.support.TransactionTemplate;
+
+@IntegrationTestConfiguration
+@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD)
+class JobConcurrencyIntegrationTest extends PostgresIntegrationTestBase {
+
+    @Autowired
+    private JobService jobService;
+
+    @Autowired
+    private TransactionTemplate transactionTemplate;
+
+    @Test
+    void testReplicateRaceConditionWithPostgres() throws InterruptedException {
+        // 1. Setup: Insert a single job into the real Postgres container
+        final UUID jobId = randomUUID();
+        transactionTemplate.executeWithoutResult(status -> {
+            final Job job = new Job();
+            job.setJobId(jobId);
+            job.setAssignedTaskName("test-task");
+            job.setJobData(Json.createObjectBuilder()
+                    .add("testKey", "testValue")
+                    .add("testNumber", 42)
+                    .build());
+            job.setAssignedTaskStartTime(ZonedDateTime.now().minusMinutes(1));
+            job.setPriority(1);
+            job.setRetryAttemptsRemaining(3);
+            jobService.insertJob(job);
+        });
+
+        final int threadCount = 2;
+        final ExecutorService executor = newFixedThreadPool(threadCount);
+        final CyclicBarrier barrier = new CyclicBarrier(threadCount);
+        final AtomicInteger successfulClaims = new AtomicInteger(0);
+        final CountDownLatch latch = new CountDownLatch(threadCount);
+
+        // 2. Run two threads simultaneously
+        for (int i = 0; i < threadCount; i++) {
+            executor.submit(() -> {
+                try {
+                    barrier.await(); // Sync threads to hit the DB at once
+
+                    final var jobs = jobService.assignJobsToWorkerBatch(jobId, 10);
+                    if (!jobs.isEmpty()) {
+                        successfulClaims.incrementAndGet();
+                    }
+                } catch (Exception e) {
+                    System.err.println("Thread failed: " + e.getMessage());
+                } finally {
+                    latch.countDown();
+                }
+            });
+        }
+
+        latch.await(10, TimeUnit.SECONDS);
+
+        // 3. Verification
+        assertEquals(1, successfulClaims.get(),
+                "RACE CONDITION DETECTED: Multiple threads claimed the same job!");
+    }
+}

task-manager-service/src/main/java/uk/gov/hmcts/cp/taskmanager/domain/executor/JobExecutor.java

@@ -156,7 +156,7 @@ public void checkAndAssignJobs() {
         try {
             logger.debug("Checking for unassigned jobs (batch size: {})...", batchSize);
             final UUID workerId = UUID.randomUUID();
-            final List<Job> assignedJobs = jobService.assignJobsToWorkerBatch(workerId, ZonedDateTime.now(), ZonedDateTime.now(), batchSize);
+            final List<Job> assignedJobs = jobService.assignJobsToWorkerBatch(workerId, batchSize);
 
             if (assignedJobs.isEmpty()) {
                 logger.debug("No unassigned jobs found");

task-manager-service/src/main/java/uk/gov/hmcts/cp/taskmanager/persistence/service/JobService.java

@@ -65,7 +65,7 @@ public class JobService {
      * @param jsonObjectConverter the JSON object converter, must not be null
      */
     @Autowired
-    public JobService(JobsRepository jobsRepository, JsonObjectConverter jsonObjectConverter) {
+    public JobService(final JobsRepository jobsRepository, final JsonObjectConverter jsonObjectConverter) {
         this.jobsRepository = jobsRepository;
         this.jsonObjectConverter = jsonObjectConverter;
     }
@@ -84,14 +84,14 @@ public JobService(JobsRepository jobsRepository, JsonObjectConverter jsonObjectC
      *         ordered by priority and start time
      */
     @Transactional
-    public List<Job> getUnassignedJobs(int batchSize) {
+    public List<Job> getUnassignedJobs(final int batchSize) {
         Pageable pageable = PageRequest.of(0, batchSize);
         return jobsRepository.findUnassignedJobsWithLimit(ZonedDateTime.now(), pageable);
     }
 
     @Transactional
-    public List<Job> assignJobsToWorkerBatch(UUID workerId,ZonedDateTime lockTime, ZonedDateTime currentTime,int batchSize) {
-        return jobsRepository.assignJobsToWorkerBatch(workerId, lockTime, currentTime, batchSize);
+    public List<Job> assignJobsToWorkerBatch(final UUID workerId, final int batchSize) {
+        return jobsRepository.assignJobsToWorkerBatch(workerId, ZonedDateTime.now(), ZonedDateTime.now(), batchSize);
     }
 
     /**
@@ -104,7 +104,7 @@ public List<Job> assignJobsToWorkerBatch(UUID workerId,ZonedDateTime lockTime, Z
      * @throws RuntimeException if the job is not found
      */
     @Transactional
-    public void decrementRetryAttempts(UUID jobId) {
+    public void decrementRetryAttempts(final UUID jobId) {
         Job job = jobsRepository.findByJobId(jobId)
                 .orElseThrow(() -> new RuntimeException("Job not found with id: " + jobId));
 
@@ -122,7 +122,7 @@ public void decrementRetryAttempts(UUID jobId) {
      * @param job the job to insert, must not be null
      */
     @Transactional
-    public void insertJob(Job job) {
+    public void insertJob(final Job job) {
         jobsRepository.save(job);
     }
 

task-manager-service/src/test/java/uk/gov/hmcts/cp/taskmanager/domain/executor/JobExecutorTest.java

@@ -82,7 +82,7 @@ void testInit() {
 
     @Test
     void testCheckAndAssignJobsWithNoJobs() {
-        when(jobService.assignJobsToWorkerBatch(any(UUID.class), any(ZonedDateTime.class), any(ZonedDateTime.class), eq(50))).thenReturn(Collections.emptyList());
+        when(jobService.assignJobsToWorkerBatch(any(UUID.class), eq(50))).thenReturn(Collections.emptyList());
 
         jobExecutor.checkAndAssignJobs();
 
@@ -103,20 +103,20 @@ void testCheckAndAssignJobsWithMultipleJobs() {
         );
         List<Job> jobs = Arrays.asList(testJob, job2);
 
-        when(jobService.assignJobsToWorkerBatch(any(UUID.class), any(ZonedDateTime.class), any(ZonedDateTime.class), eq(50)))
+        when(jobService.assignJobsToWorkerBatch(any(UUID.class), eq(50)))
                 .thenReturn(jobs);
 
         jobExecutor.checkAndAssignJobs();
 
-        verify(jobService).assignJobsToWorkerBatch(any(UUID.class), any(ZonedDateTime.class), any(ZonedDateTime.class), eq(50));
+        verify(jobService).assignJobsToWorkerBatch(any(UUID.class), eq(50));
         verify(executor, times(2)).execute(any());
     }
 
     @Test
     void testCheckAndAssignJobsWithExecutionFailure() {
         List<Job> jobs = Collections.singletonList(testJob);
 
-        when(jobService.assignJobsToWorkerBatch(any(UUID.class), any(ZonedDateTime.class), any(ZonedDateTime.class), eq(50)))
+        when(jobService.assignJobsToWorkerBatch(any(UUID.class), eq(50)))
                 .thenReturn(jobs);
         doThrow(new RuntimeException("Executor failed")).when(executor).execute(any(Runnable.class));
 
@@ -129,7 +129,7 @@ void testCheckAndAssignJobsWithExecutionFailure() {
     void testCheckAndAssignJobsWithDecrementFailure() {
         List<Job> jobs = Collections.singletonList(testJob);
 
-        when(jobService.assignJobsToWorkerBatch(any(UUID.class), any(ZonedDateTime.class), any(ZonedDateTime.class), eq(50)))
+        when(jobService.assignJobsToWorkerBatch(any(UUID.class), eq(50)))
                 .thenReturn(jobs);
         doThrow(new RuntimeException("Executor failed")).when(executor).execute(any(Runnable.class));
         doThrow(new RuntimeException("Decrement failed")).when(jobService).decrementRetryAttempts(any(UUID.class));
@@ -142,7 +142,7 @@ void testCheckAndAssignJobsWithDecrementFailure() {
 
     @Test
     void testCheckAndAssignJobsWithServiceException() {
-        when(jobService.assignJobsToWorkerBatch(any(), any(), any(), eq(50)))
+        when(jobService.assignJobsToWorkerBatch(any(), eq(50)))
                 .thenThrow(new RuntimeException("Service error"));
 
         // Should not throw exception, should handle gracefully