This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new 936ebd0fec [CELEBORN-2501] Fix DataPushQueue dropping tasks when there
is no partition location map
936ebd0fec is described below
commit 936ebd0fec246bcc5039b0f67b5f54bf287e0962
Author: ForVic <[email protected]>
AuthorDate: Sun Oct 25 16:05:16 2026 +0800
[CELEBORN-2501] Fix DataPushQueue dropping tasks when there is no partition
location map
### What changes were proposed in this pull request?
Minor bug fix, ran into it while running a micro benchmark.
When `client.getPartitionLocation` returns `null`,
`DataPushQueue.takePushTasks` takes every queued task with:
```java
tasks.addAll(workingQueue);
workingQueue.clear();
```
The writer thread keeps adding to `workingQueue` while this runs. If it
adds a task after `addAll` has copied the queue and before `clear` runs,
`clear` removes that task even though it was never copied into `tasks`. The
task is never pushed, and its `PushTask` buffer never goes back to
`DataPusher`'s idle queue.
This PR replaces the two calls with `workingQueue.drainTo(tasks)`, which
removes only the tasks it returns. The branch behaved this way before
CELEBORN-1218 (#2220), which removed each task through `iterator.remove()`.
### Why are the changes needed?
Each lost task loses a pushed batch and one buffer from the idle pool. Once
every buffer is lost, the writer blocks forever in `DataPusher.addTask` waiting
for an idle buffer, and the push thread sleeps on an empty working queue. If
fewer are lost, `waitOnTermination` waits forever for the idle queue to fill up
again.
`ShuffleClientImpl.getPartitionLocation` never returns `null`: it returns
the registered map or throws. So production writers do not take this branch.
`DummyShuffleClient` returns `null` when no reduce partition map was
initialized, so tests and benchmarks built on it can hang.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
- New `DataPushQueueSuiteJ#testTaskAddedWhileTakingAllTasksIsNotLost`. A
package-private constructor lets the test pass in its own working queue. The
queue adds a second task right after `takePushTasks` reads the queued tasks,
which is exactly the gap between the old `addAll` and `clear`. The test checks
that `takePushTasks` returns the first task and that the second one is still
queued.
- Without the fix it failed 3 runs out of 3, because the second task was
cleared.
- With the fix it passes.
- Before switching to this test, a 200,000-task stress test through
`DataPusher` with no partition map hung the old code after 90,829 pushes, with
the writer parked in `addTask` and the push thread sleeping in `takePushTasks`.
- Existing `DataPushQueueSuiteJ`.
Closes #3895 from ForVic/ForVic/o5nq5.
Authored-by: ForVic <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/32
AI-Contributed/UT: 0/56
(cherry picked from commit d750eddeb3cfd90fca21d6d2e961b0eb53765c6f)
Signed-off-by: zhengtao <[email protected]>
---
.../celeborn/client/write/DataPushQueue.java | 32 +++++++++++--
.../celeborn/client/write/DataPushQueueSuiteJ.java | 56 ++++++++++++++++++++++
2 files changed, 84 insertions(+), 4 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
index aec6a69b5d..71e61d38f8 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
@@ -23,6 +23,7 @@ import java.util.concurrent.LinkedBlockingQueue;
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 class 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 class 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 class DataPushQueue {
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()) {
diff --git
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
index 841e04eee0..667d88b9bf 100644
---
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
+++
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
@@ -20,15 +20,19 @@ package org.apache.celeborn.client.write;
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;
@@ -178,6 +182,58 @@ public class DataPushQueueSuiteJ {
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);
