Add release callbacks to QueryCtx (facebookincubator#15858)

Summary:

Add ability to provide callbacks to be executed before QueryCtx is destructed

Reviewed By: xiaoxmeng

Differential Revision: D89797291

Author: shrinidhijoshi

velox/core/QueryCtx.cpp

@@ -21,6 +21,13 @@
 
 namespace facebook::velox::core {
 
+QueryCtx::~QueryCtx() {
+  for (auto& cb : releaseCallbacks_) {
+    cb();
+  }
+  VELOX_CHECK(!underArbitration_);
+}
+
 // static
 std::shared_ptr<QueryCtx> QueryCtx::create(
     folly::Executor* executor,

velox/core/QueryCtx.h

@@ -18,6 +18,8 @@
 
 #include <folly/Executor.h>
 #include <folly/executors/CPUThreadPoolExecutor.h>
+#include <deque>
+#include <functional>
 #include "velox/common/caching/AsyncDataCache.h"
 #include "velox/common/memory/Memory.h"
 #include "velox/core/QueryConfig.h"
@@ -32,8 +34,11 @@ namespace facebook::velox::core {
 
 class QueryCtx : public std::enable_shared_from_this<QueryCtx> {
  public:
-  ~QueryCtx() {
-    VELOX_CHECK(!underArbitration_);
+  using ReleaseCallback = std::function<void()>;
+
+  ~QueryCtx();
+  void addReleaseCallback(ReleaseCallback callback) {
+    releaseCallbacks_.push_back(std::move(callback));
   }
 
   /// QueryCtx is used in different places. When used with `Task::start()`, it's
@@ -234,6 +239,8 @@ class QueryCtx : public std::enable_shared_from_this<QueryCtx> {
   std::atomic_bool underArbitration_{false};
   std::vector<ContinuePromise> arbitrationPromises_;
   std::shared_ptr<filesystems::TokenProvider> fsTokenProvider_;
+  // Callbacks invoked before destruction to clean up external resources.
+  std::deque<ReleaseCallback> releaseCallbacks_;
 };
 
 // Represents the state of one thread of query execution.

velox/core/tests/QueryCtxTest.cpp

@@ -51,4 +51,36 @@ TEST_F(QueryCtxTest, withSysRootPool) {
   ASSERT_FALSE(
       queryPool->reclaimer()->reclaimableBytes(*queryPool, reclaimableBytes));
 }
+
+TEST_F(QueryCtxTest, releaseCallbacks) {
+  int callbackCount = 0;
+  std::string capturedQueryId;
+
+  {
+    auto queryCtx = QueryCtx::create(
+        nullptr,
+        QueryConfig{{}},
+        std::unordered_map<std::string, std::shared_ptr<config::ConfigBase>>{},
+        nullptr,
+        nullptr,
+        nullptr,
+        "test_query_id");
+
+    // Add multiple callbacks.
+    queryCtx->addReleaseCallback([&callbackCount]() { ++callbackCount; });
+
+    queryCtx->addReleaseCallback(
+        [&callbackCount, &capturedQueryId, id = queryCtx->queryId()]() {
+          ++callbackCount;
+          capturedQueryId = id;
+        });
+
+    // Callbacks should not be invoked yet.
+    ASSERT_EQ(callbackCount, 0);
+  }
+
+  // After QueryCtx destruction, all callbacks should have been invoked.
+  ASSERT_EQ(callbackCount, 2);
+  ASSERT_EQ(capturedQueryId, "test_query_id");
+}
 } // namespace facebook::velox::core::test