This is an automated email from the ASF dual-hosted git repository.

FelixYBW pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git


The following commit(s) were added to refs/heads/main by this push:
     new 810810ce94 [GLUTEN-12911][VL] Fix memory pool holding in async thread 
(#12919)
810810ce94 is described below

commit 810810ce94b75f881205ab6d3e113fb409a26aaf
Author: BInwei Yang <[email protected]>
AuthorDate: Fri Sep 11 11:18:45 2026 -0700

    [GLUTEN-12911][VL] Fix memory pool holding in async thread (#12919)
    
    Without the barrier, it's possible we mark the load as canceled, but it's 
never scheduled before the timeout, so it holds the memory pool pointer all the 
time.
    
    The barrier will wait until the load is scheduled and exits.
---
 cpp/velox/memory/GlutenDirectBufferedInput.h | 81 +++++++++++++++++++++++++---
 1 file changed, 75 insertions(+), 6 deletions(-)

diff --git a/cpp/velox/memory/GlutenDirectBufferedInput.h 
b/cpp/velox/memory/GlutenDirectBufferedInput.h
index d5b588770e..7c33c8a879 100644
--- a/cpp/velox/memory/GlutenDirectBufferedInput.h
+++ b/cpp/velox/memory/GlutenDirectBufferedInput.h
@@ -17,11 +17,58 @@
 
 #pragma once
 
+#include <glog/logging.h>
+
 #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<facebook::velox::dwio::common::ExecutorBarrier> 
makeBarrier(folly::Executor* executor) {
+    if (executor == nullptr) {
+      return nullptr;
+    }
+    return 
std::make_unique<facebook::velox::dwio::common::ExecutorBarrier>(folly::getKeepAliveToken(executor));
+  }
+
+  folly::Executor* const rawExecutor_;
+  const std::unique_ptr<facebook::velox::dwio::common::ExecutorBarrier> 
barrier_;
+};
+
+} // namespace detail
+
+class GlutenDirectBufferedInput : private detail::ExecutorBarrierHolder,
+                                  public 
facebook::velox::dwio::common::DirectBufferedInput {
  public:
   GlutenDirectBufferedInput(
       std::shared_ptr<facebook::velox::ReadFile> readFile,
@@ -34,7 +81,8 @@ class GlutenDirectBufferedInput : public 
facebook::velox::dwio::common::DirectBu
       folly::Executor* executor,
       const facebook::velox::io::ReaderOptions& readerOptions,
       folly::F14FastMap<std::string, std::string> fileReadOps = {})
-      : DirectBufferedInput(
+      : ExecutorBarrierHolder(executor),
+        DirectBufferedInput(
             std::move(readFile),
             metricsLog,
             std::move(fileNum),
@@ -42,13 +90,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 +117,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<facebook::velox::dwio::common::BufferedInput> 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<facebook::velox::dwio::common::BufferedInput>(new 
GlutenDirectBufferedInput(
-        input_, fileNum_, tracker_, groupId_, ioStatistics_, ioStats_, 
executor_, options_));
+        input_, fileNum_, tracker_, groupId_, ioStatistics_, ioStats_, 
rawExecutor(), options_));
   }
 
  private:
@@ -84,14 +152,15 @@ class GlutenDirectBufferedInput : public 
facebook::velox::dwio::common::DirectBu
       std::shared_ptr<facebook::velox::IoStats> 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) {}
 };
 


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to