From fd98d0ca8bfb6137aa33e1ab6b41ff0f45c88712 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 28 Aug 2026 04:57:46 +0000 Subject: [PATCH] fix by claude --- cpp/velox/memory/GlutenDirectBufferedInput.h | 82 ++++++++++++++++++-- 1 file changed, 76 insertions(+), 6 deletions(-) diff --git a/cpp/velox/memory/GlutenDirectBufferedInput.h b/cpp/velox/memory/GlutenDirectBufferedInput.h index d5b588770ec..86c34d6dba5 100644 --- a/cpp/velox/memory/GlutenDirectBufferedInput.h +++ b/cpp/velox/memory/GlutenDirectBufferedInput.h @@ -17,11 +17,59 @@ #pragma once +#include + #include "velox/dwio/common/DirectBufferedInput.h" +#include "velox/dwio/common/ExecutorBarrier.h" namespace gluten { -class GlutenDirectBufferedInput : public facebook::velox::dwio::common::DirectBufferedInput { +namespace detail { + +// Owns the ExecutorBarrier that wraps the IO executor passed to +// DirectBufferedInput. DirectBufferedInput::readRegions() enqueues an +// AsyncLoadHolder closure per planned load, and that closure keeps a +// shared_ptr on the reader's MemoryPool. Cancelling a load only flips its +// state, it does not dequeue or destroy the closure, so the pool reference +// survives until the executor happens to drain that entry. Routing all the +// enqueues through a barrier lets the destructor wait for them explicitly. +// +// This must be listed as a base before DirectBufferedInput so that it is +// constructed first (the barrier pointer is handed to the base constructor) +// and destructed last (the barrier has to outlive any closure referring to +// it). +class ExecutorBarrierHolder { + public: + explicit ExecutorBarrierHolder(folly::Executor* executor) + : rawExecutor_(executor), barrier_(makeBarrier(executor)) {} + + protected: + // The unwrapped executor, to be handed to clones instead of this object's + // barrier. + folly::Executor* rawExecutor() const { + return rawExecutor_; + } + + facebook::velox::dwio::common::ExecutorBarrier* barrier() const { + return barrier_.get(); + } + + private: + static std::unique_ptr makeBarrier(folly::Executor* executor) { + if (executor == nullptr) { + return nullptr; + } + return std::make_unique(folly::getKeepAliveToken(executor)); + } + + folly::Executor* const rawExecutor_; + const std::unique_ptr barrier_; +}; + +} // namespace detail + +class GlutenDirectBufferedInput : private detail::ExecutorBarrierHolder, + public facebook::velox::dwio::common::DirectBufferedInput { public: GlutenDirectBufferedInput( std::shared_ptr readFile, @@ -34,7 +82,8 @@ class GlutenDirectBufferedInput : public facebook::velox::dwio::common::DirectBu folly::Executor* executor, const facebook::velox::io::ReaderOptions& readerOptions, folly::F14FastMap fileReadOps = {}) - : DirectBufferedInput( + : ExecutorBarrierHolder(executor), + DirectBufferedInput( std::move(readFile), metricsLog, std::move(fileNum), @@ -42,13 +91,16 @@ class GlutenDirectBufferedInput : public facebook::velox::dwio::common::DirectBu std::move(groupId), std::move(ioStatistics), std::move(ioStats), - executor, + barrier(), readerOptions, std::move(fileReadOps)) {} ~GlutenDirectBufferedInput() override { requests_.clear(); // Cancel all the planned loads as soon as possible to avoid unnecessary IO. + // Only kPlanned loads may be cancelled: cancel() overwrites the state + // unconditionally, so cancelling a kLoading load would hide it from the + // wait below while its IO is still in flight. for (auto& load : coalescedLoads_) { if (load->state() == facebook::velox::cache::CoalescedLoad::State::kPlanned) { load->cancel(); @@ -66,11 +118,28 @@ class GlutenDirectBufferedInput : public facebook::velox::dwio::common::DirectBu } } coalescedLoads_.clear(); + // The cancelled loads above are still referenced by the AsyncLoadHolder + // closures queued on the executor, and those closures hold a shared_ptr on + // the memory pool. Wait until the executor has run and destroyed them so + // that the pool reference is released before this destructor returns, + // instead of on an IO thread after the task and its memory manager are + // gone. + if (barrier() != nullptr) { + try { + barrier()->waitAll(); + } catch (const std::exception& e) { + // waitAll() rethrows an exception raised by any of the loads. It must + // not escape the destructor: the loads were cancelled anyway. + LOG(WARNING) << "Async load failed while destructing GlutenDirectBufferedInput: " << e.what(); + } + } } std::unique_ptr clone() const override { + // Pass the unwrapped executor: the clone has its own lifetime and must not + // enqueue onto this object's barrier. return std::unique_ptr(new GlutenDirectBufferedInput( - input_, fileNum_, tracker_, groupId_, ioStatistics_, ioStats_, executor_, options_)); + input_, fileNum_, tracker_, groupId_, ioStatistics_, ioStats_, rawExecutor(), options_)); } private: @@ -84,14 +153,15 @@ class GlutenDirectBufferedInput : public facebook::velox::dwio::common::DirectBu std::shared_ptr ioStats, folly::Executor* executor, const facebook::velox::io::ReaderOptions& readerOptions) - : DirectBufferedInput( + : ExecutorBarrierHolder(executor), + DirectBufferedInput( std::move(input), std::move(fileNum), std::move(tracker), std::move(groupId), std::move(ioStatistics), std::move(ioStats), - executor, + barrier(), readerOptions) {} };