[CELEBORN-2501] Fix DataPushQueue dropping tasks when there is no partition location map

client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java

@@ -23,6 +23,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -62,6 +63,29 @@ public DataPushQueue(
       int attemptId,
       int numMappers,
       int numPartitions) {
+    this(
+        conf,
+        dataPusher,
+        client,
+        shuffleId,
+        mapId,
+        attemptId,
+        numMappers,
+        numPartitions,
+        new LinkedBlockingQueue<>(conf.clientPushQueueCapacity()));
+  }
+
+  @VisibleForTesting
+  DataPushQueue(
+      CelebornConf conf,
+      DataPusher dataPusher,
+      ShuffleClient client,
+      int shuffleId,
+      int mapId,
+      int attemptId,
+      int numMappers,
+      int numPartitions,
+      LinkedBlockingQueue<PushTask> workingQueue) {
     this.shuffleId = shuffleId;
     this.numMappers = numMappers;
     this.numPartitions = numPartitions;
@@ -71,8 +95,7 @@ public DataPushQueue(
     this.pushState = client.getPushState(mapKey);
     this.takeTaskWaitIntervalMs = conf.clientPushTakeTaskWaitIntervalMs();
     this.takeTaskMaxWaitAttempts = conf.clientPushTakeTaskMaxWaitAttempts();
-    final int capacity = conf.clientPushQueueCapacity();
-    workingQueue = new LinkedBlockingQueue<>(capacity);
+    this.workingQueue = workingQueue;
   }
 
   /*
@@ -91,8 +114,9 @@ public ArrayList<PushTask> takePushTasks() throws IOException, InterruptedExcept
       Map<Integer, PartitionLocation> partitionLocationMap =
           client.getPartitionLocation(shuffleId, numMappers, numPartitions);
       if (partitionLocationMap == null) {
-        tasks.addAll(workingQueue);
-        workingQueue.clear();
+        // drainTo removes exactly the tasks it returns, so a task added concurrently by the
+        // writer thread is not dropped as it would be by addAll followed by clear.
+        workingQueue.drainTo(tasks);
       } else {
         Iterator<PushTask> iterator = workingQueue.iterator();
         while (iterator.hasNext()) {

client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java

@@ -20,15 +20,19 @@
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.LongAdder;
 
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.mockito.Mockito;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -185,6 +189,58 @@ protected void pushData(PushTask task) throws IOException {
     client.shutdown();
   }
 
+  @Test
+  public void testTaskAddedWhileTakingAllTasksIsNotLost() throws Exception {
+    CelebornConf conf = new CelebornConf();
+    final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+    // No reduce partition map is initialized, so getPartitionLocation returns null and
+    // takePushTasks takes every queued task.
+    DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+    DataPusher dataPusher = Mockito.mock(DataPusher.class);
+    Mockito.when(dataPusher.stillRunning()).thenReturn(true);
+
+    PushTask queuedTask = new PushTask(4);
+    PushTask concurrentTask = new PushTask(4);
+    // Adds concurrentTask right after takePushTasks reads the queued tasks, as the writer thread
+    // can. toArray is how ArrayList.addAll reads the queue, drainTo is how it is read with the fix.
+    LinkedBlockingQueue<PushTask> workingQueue =
+        new LinkedBlockingQueue<PushTask>(conf.clientPushQueueCapacity()) {
+          private boolean added = false;
+
+          private void addConcurrentTaskOnce() {
+            if (!added) {
+              added = true;
+              super.offer(concurrentTask);
+            }
+          }
+
+          @Override
+          public Object[] toArray() {
+            Object[] tasks = super.toArray();
+            addConcurrentTaskOnce();
+            return tasks;
+          }
+
+          @Override
+          public int drainTo(Collection<? super PushTask> c) {
+            int n = super.drainTo(c);
+            addConcurrentTaskOnce();
+            return n;
+          }
+        };
+    DataPushQueue dataPushQueue =
+        new DataPushQueue(conf, dataPusher, client, 0, 0, 0, 1, 1, workingQueue);
+
+    try {
+      Assert.assertTrue(dataPushQueue.addPushTask(queuedTask));
+      Assert.assertEquals(Collections.singletonList(queuedTask), dataPushQueue.takePushTasks());
+      // The task added while taking must stay queued for the next takePushTasks call.
+      Assert.assertEquals(Collections.singletonList(concurrentTask), new ArrayList<>(workingQueue));
+    } finally {
+      client.shutdown();
+    }
+  }
+
   public static byte[] intToBytes(int value) {
     byte[] src = new byte[4];
     src[0] = (byte) (value & 0xFF);