Fix Batch bulk task retry race (#50881)

* Fix Batch bulk task retry race

* Updated versions and changelog

Author: vcolin7

sdk/batch/azure-compute-batch/CHANGELOG.md

@@ -16,6 +16,10 @@
 - Replaced the public `BatchFileProperties(HttpHeaders)` constructor with a static factory
   `BatchFileProperties.fromHeaders(HttpHeaders)`.
 
+### Bugs Fixed
+
+- Fixed a race condition that could cause bulk task creation to return before retrying tasks with server errors.
+
 ### Other Changes
 
 This is the first generally available release of the `azure-compute-batch` library. It includes all features and breaking changes introduced in `1.0.0-beta.6`.

sdk/batch/azure-compute-batch/src/main/java/com/azure/compute/batch/implementation/task/TaskManager.java

@@ -220,8 +220,8 @@ public static void createTasks(TaskSubmitter taskSubmitter, String jobId, List<B
         List<ThreadPairInfo> semaphoreThreads = new ArrayList<>();
         Exception innerException = null;
 
-        // Continue looping while there are still tasks left to submit OR while any active threads are still processing tasks.
-        while (!pendingList.isEmpty() || taskSemaphore.availablePermits() < threadNumber) {
+        // Check workers first so tasks requeued before their permits are released cannot be missed.
+        while (taskSemaphore.availablePermits() < threadNumber || !pendingList.isEmpty()) {
             try {
                 if (taskSemaphore.tryAcquire(30, TimeUnit.SECONDS)) {
                     if (!pendingList.isEmpty()) {

sdk/batch/azure-compute-batch/src/test/java/com/azure/compute/batch/TaskManagerTests.java

@@ -4,19 +4,35 @@
 package com.azure.compute.batch;
 
 import com.azure.compute.batch.implementation.task.TaskManager;
+import com.azure.compute.batch.implementation.task.TaskSubmitter;
 import com.azure.compute.batch.models.BatchTaskBulkCreateOptions;
 import com.azure.compute.batch.models.BatchTaskCreateParameters;
 import com.azure.compute.batch.models.CreateTasksErrorException;
-import org.junit.jupiter.api.*;
+import org.junit.jupiter.api.RepeatedTest;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.parallel.Execution;
 import org.junit.jupiter.api.parallel.ExecutionMode;
 import org.junit.jupiter.api.parallel.Isolated;
-
-import java.util.*;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @Execution(ExecutionMode.SAME_THREAD)
 @Isolated
@@ -28,7 +44,7 @@ private static List<BatchTaskCreateParameters> makeTasks(String prefix, int coun
             .collect(Collectors.toList());
     }
 
-    private static void runCreateTasks(FakeTaskSubmitter submitter, List<BatchTaskCreateParameters> tasks) {
+    private static void runCreateTasks(TaskSubmitter submitter, List<BatchTaskCreateParameters> tasks) {
         TaskManager.createTasks(submitter, "job-1", tasks, new BatchTaskBulkCreateOptions().setMaxConcurrency(4));
     }
 
@@ -44,11 +60,11 @@ public void serverErrorsAreRetriedUntilSuccessNoCap() {
 
         assertDoesNotThrow(() -> runCreateTasks(fake, tasks));
 
-        assertTrue(fake.getSubmissionCount("A") > 1, "A should have been retried at least once");
+        assertEquals(11, fake.getSubmissionCount("A"), "A should succeed after ten retries");
         assertEquals(1, fake.getSubmissionCount("B"), "B should submit once");
     }
 
-    @Test
+    @RepeatedTest(1000)
     public void successAfterOneServerErrorDidRetry() {
         Map<String, Integer> serverErrors = new HashMap<>();
         serverErrors.put("X", 1); // one failure then success
@@ -57,7 +73,45 @@ public void successAfterOneServerErrorDidRetry() {
         List<BatchTaskCreateParameters> tasks = Collections.singletonList(new BatchTaskCreateParameters("X", "echo X"));
 
         assertDoesNotThrow(() -> runCreateTasks(fake, tasks));
-        assertTrue(fake.getSubmissionCount("X") >= 2, "X should have been retried at least once");
+        assertEquals(2, fake.getSubmissionCount("X"), "X should succeed after one retry");
+    }
+
+    @ParameterizedTest
+    @ValueSource(ints = { 1, 4 })
+    public void retriesAcrossMultipleChunks(int maxConcurrency) {
+        List<BatchTaskCreateParameters> tasks = makeTasks("T", 201);
+        Map<String, Integer> serverErrors = new HashMap<>();
+        tasks.forEach(task -> serverErrors.put(task.getId(), 1));
+        FakeTaskSubmitter fake = new FakeTaskSubmitter(Collections.emptySet(), serverErrors, 0);
+
+        assertDoesNotThrow(() -> TaskManager.createTasks(fake, "job-1", tasks,
+            new BatchTaskBulkCreateOptions().setMaxConcurrency(maxConcurrency)));
+
+        tasks.forEach(task -> assertEquals(2, fake.getSubmissionCount(task.getId()),
+            task.getId() + " should succeed after one retry"));
+    }
+
+    @RepeatedTest(100)
+    public void submissionExceptionIsPropagated() {
+        RuntimeException expected = new IllegalStateException("Submission failed");
+        TaskSubmitter submitter = (jobId, taskCollection) -> {
+            throw expected;
+        };
+
+        assertSame(expected, assertThrows(IllegalStateException.class,
+            () -> runCreateTasks(submitter, Collections.singletonList(new BatchTaskCreateParameters("X", "echo X")))));
+    }
+
+    @Test
+    public void checkedSubmissionExceptionIsPropagated() {
+        IOException expected = new IOException("Submission failed");
+        TaskSubmitter submitter = (jobId, taskCollection) -> {
+            throw expected;
+        };
+
+        RuntimeException actual = assertThrows(RuntimeException.class,
+            () -> runCreateTasks(submitter, Collections.singletonList(new BatchTaskCreateParameters("X", "echo X"))));
+        assertSame(expected, actual.getCause());
     }
 
     @Test