This is an automated email from the ASF dual-hosted git repository. SteNicholas pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/celeborn.git
commit 93de9b3042f419fa914144cd640e18e8995a3b8b Author: afterincomparableyum <[email protected]> AuthorDate: Mon Jun 8 14:03:26 2026 +0800 [CELEBORN-2332] Fix self join deadlock in C++ WorkerPartitionReader fetch callbacks ### What changes were proposed in this pull request? push-merged data, which exercises this fetch path more aggressively and reliably triggered an EDEADLK abort. The onSuccess_/onFailure_ callbacks are invoked on the TransportClient's IO thread and capture a weak_ptr that is lifted to a shared_ptr inside the callback body. When that local shared_ptr happens to hold the last reference, dropping it inline runs WorkerPartitionReader on the IO executor's own thread, which transitively destroys the embedded TransportClient and its IOThreadPoolExecutor. The executor then attempts to pthread_join the thread that is currently executing the callback an [...] Hand the final reference off to the global CPU executor so destruction of the reader (and the IO executor underneath it) always happens on a different thread than the one running the callback. ### Why are the changes needed? When running the bytedance bolt celeborn e2e tests for push-merged data I am working on, I run into this error. ### Does this PR resolve a correctness bug? Yes ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? I ran the Celeborn bolt e2e tests with this change and the tests passed with push merged data support. Closes #3693 from afterincomparableyum/celeborn-2332. Authored-by: afterincomparableyum <[email protected]> Signed-off-by: Nicholas Jiang <[email protected]> --- .../client/reader/WorkerPartitionReader.cpp | 35 ++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/cpp/celeborn/client/reader/WorkerPartitionReader.cpp b/cpp/celeborn/client/reader/WorkerPartitionReader.cpp index 7fb2c5de7..a0bb577bb 100644 --- a/cpp/celeborn/client/reader/WorkerPartitionReader.cpp +++ b/cpp/celeborn/client/reader/WorkerPartitionReader.cpp @@ -17,8 +17,29 @@ #include "celeborn/client/reader/WorkerPartitionReader.h" +#include <folly/executors/CPUThreadPoolExecutor.h> + namespace celeborn { namespace client { +namespace { +// Off-IO-thread executor for releasing the last shared_ptr to a +// WorkerPartitionReader. The fetch callbacks run on the +// embedded TransportClient's IOThreadPoolExecutor, and dropping the final +// reference inline would have that executor pthread_join its own thread. +// A directly constructed CPUThreadPoolExecutor avoids the folly singleton +// vault and so does not require folly::init() to have been called. +// +// TODO: this posts a task on every fetch callback, even when the reader is +// still owned elsewhere and no off thread destruction is needed. Optimize to +// hand off only when this drop would actually destroy the reader, but do it +// race free (drain in-flight fetches before the owner releases, so the +// reader is always destroyed on the consumer thread). A use_count() based +// "last ref" guard is NOT a valid optimization here -- see the call sites. +folly::CPUThreadPoolExecutor& destructionExecutor() { + static folly::CPUThreadPoolExecutor instance{1}; + return instance; +} +} // namespace std::shared_ptr<WorkerPartitionReader> WorkerPartitionReader::create( const std::shared_ptr<const conf::CelebornConf>& conf, const std::string& shuffleKey, @@ -122,6 +143,15 @@ void WorkerPartitionReader::initAndCheck() { shared_this->chunkQueue_.enqueue(std::move(chunk)); VLOG(1) << "WorkerPartitionReader::onSuccess: " << streamChunkSlice.toString(); + // Always hand the final reference off, unconditionally. A + // use_count()-based "only offload if last ref" check is NOT safe here. + // This callback always runs on the IO thread, and use_count() is a stale + // snapshot. The owner (CelebornInputStream::currReader_) can drop its + // reference on another thread right after the check, leaving this + // callback to destroy the reader inline on the IO thread and re-trigger + // the EDEADLK self-join error. + destructionExecutor().add( + [s = std::move(shared_this)]() mutable { s.reset(); }); }; onFailure_ = [weak_this = weak_from_this()]( @@ -138,6 +168,11 @@ void WorkerPartitionReader::initAndCheck() { auto exp = shared_this->exception_.wlock(); *exp = std::move(exception); } + // See onSuccess_ above. The off thread handoff is unconditional on + // purpose. A use_count() "last ref" guard would be racy and could + // re-trigger the EDEADLK self-join error. + destructionExecutor().add( + [s = std::move(shared_this)]() mutable { s.reset(); }); }; }
