Repository navigation
Conversation
ForVic
changed the title
ForVic
marked this pull request as ready for review
October 22, 2026 22:33
Contributor
Author
|
cc @zaynt4606 or @SteNicholas if you get a chance to review sometime |
zaynt4606
approved these changes
Oct 23, 2026
zaynt4606
left a comment
zaynt4606
left a comment
Contributor
There was a problem hiding this comment.
LGTM, with one non-blocking test nit below.
| writer.submit( | ||
| () -> { | ||
| byte[] b = new byte[4]; | ||
| for (int i = 0; i < numTasks; i++) { |
Contributor
There was a problem hiding this comment.
Nit (non-blocking): Could we replace the 200,000-task stress loop with a deterministic regression test that explicitly controls the producer/consumer interleaving, for example with latches? The current test relies on scheduling to expose the race. Ideally, the test should force an enqueue between the old snapshot and clear operations and verify that the task is preserved, failing reliably with the old implementation and passing with drainTo.
Contributor
Author
There was a problem hiding this comment.
good callout, improved the test
…tition location map takePushTasks copied the working queue with addAll and then called clear. A task the writer added between the two calls was cleared without being returned, so it was never pushed and its buffer never went back to the idle queue. Use drainTo, which removes only the tasks it returns.
zaynt4606
pushed a commit
that referenced
this pull request
Oct 25, 2026
…tition 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 <victor.lakers0@gmail.com> Signed-off-by: zhengtao <shuaizhentao.szt@alibaba-inc.com> AI-Contributed/Feature: 0/32 AI-Contributed/UT: 0/56 (cherry picked from commit d750edd) Signed-off-by: zhengtao <shuaizhentao.szt@alibaba-inc.com>
Contributor
Contributor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
Minor bug fix, ran into it while running a micro benchmark.
When
client.getPartitionLocationreturnsnull,DataPushQueue.takePushTaskstakes every queued task with:The writer thread keeps adding to
workingQueuewhile this runs. If it adds a task afteraddAllhas copied the queue and beforeclearruns,clearremoves that task even though it was never copied intotasks. The task is never pushed, and itsPushTaskbuffer never goes back toDataPusher'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 throughiterator.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.addTaskwaiting for an idle buffer, and the push thread sleeps on an empty working queue. If fewer are lost,waitOnTerminationwaits forever for the idle queue to fill up again.ShuffleClientImpl.getPartitionLocationnever returnsnull: it returns the registered map or throws. So production writers do not take this branch.DummyShuffleClientreturnsnullwhen no reduce partition map was initialized, so tests and benchmarks built on it can hang.Does this PR resolve a correctness bug?
Does this PR introduce any user-facing change?
How was this patch tested?
DataPushQueueSuiteJ#testTaskAddedWhileTakingAllTasksIsNotLost. A package-private constructor lets the test pass in its own working queue. The queue adds a second task right aftertakePushTasksreads the queued tasks, which is exactly the gap between the oldaddAllandclear. The test checks thattakePushTasksreturns the first task and that the second one is still queued.DataPusherwith no partition map hung the old code after 90,829 pushes, with the writer parked inaddTaskand the push thread sleeping intakePushTasks.DataPushQueueSuiteJ.