From ee3a0a34f7a013f858e21e01a304ae66a2f7a42f Mon Sep 17 00:00:00 2001
From: Pratik Joseph Dabre <pdabre12@gmail.com>
Date: Tue, 15 Sep 2026 15:58:25 -0700
Subject: [PATCH] fix(native): Drain httpSrvCpuExecutor before driverExecutor
 on shutdown

---
 .../presto_cpp/main/PrestoServer.cpp          | 32 ++++++++--------
 .../main/tests/ShutdownOrderTest.cpp          | 37 +++++++++++++++++++
 2 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
index 25372543b465a..6c1cd0866f4f7 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -830,11 +830,22 @@ void PrestoServer::stopAnnouncer() {
 }
 
 void PrestoServer::joinExecutors() {
-  // Join exchange HTTP CPU executor first. Exchange CPU threads run
-  // PrestoExchangeSource::handleDataResponse which dispatches callbacks to
-  // driverExecutor_ (MonitoredExecutor) via InMemoryExchangeClient. We must
-  // drain these threads before destroying driverExecutor_ to avoid
-  // use-after-free.
+  // Drain httpSrvCpuExecutor_ first: /v1/expressions tasks run here and pass
+  // driverExecutor_.get() into QueryCtx; they must finish before the driver
+  // pool is stopped.
+  if (httpSrvCpuExecutor_ != nullptr) {
+    PRESTO_SHUTDOWN_LOG(INFO)
+        << "Joining HTTP Server CPU Executor '"
+        << httpSrvCpuExecutor_->getName()
+        << "': threads: " << httpSrvCpuExecutor_->numActiveThreads() << "/"
+        << httpSrvCpuExecutor_->numThreads()
+        << ", task queue: " << httpSrvCpuExecutor_->getTaskQueueSize();
+    httpSrvCpuExecutor_->join();
+  }
+
+  // Join exchange HTTP CPU executor before the driver executor.  Exchange CPU
+  // threads dispatch callbacks to driverExecutor_ via InMemoryExchangeClient;
+  // drain them before destroying driverExecutor_ to avoid use-after-free.
   PRESTO_SHUTDOWN_LOG(INFO)
       << "Joining Exchange Http CPU executor '"
       << exchangeHttpCpuExecutor_->getName()
@@ -873,15 +884,6 @@ void PrestoServer::joinExecutors() {
     connectorIoExecutor_->join();
   }
 
-  if (httpSrvCpuExecutor_ != nullptr) {
-    PRESTO_SHUTDOWN_LOG(INFO)
-        << "Joining HTTP Server CPU Executor '"
-        << httpSrvCpuExecutor_->getName()
-        << "': threads: " << httpSrvCpuExecutor_->numActiveThreads() << "/"
-        << httpSrvCpuExecutor_->numThreads()
-        << ", task queue: " << httpSrvCpuExecutor_->getTaskQueueSize();
-    httpSrvCpuExecutor_->join();
-  }
   if (httpSrvIoExecutor_ != nullptr) {
     PRESTO_SHUTDOWN_LOG(INFO)
         << "Joining HTTP Server IO Executor '" << httpSrvIoExecutor_->getName()
@@ -1994,7 +1996,7 @@ void PrestoServer::registerSidecarEndpoints() {
                   .thenValue([](auto&& result) {
                     // Serialize on the CPU executor so the I/O thread only
                     // transmits pre-built bytes.
-                    return util::dumpJson(json(result));
+                    return util::dumpJson(json(std::move(result)));
                   })
                   .via(
                       folly::getKeepAliveToken(
diff --git a/presto-native-execution/presto_cpp/main/tests/ShutdownOrderTest.cpp b/presto-native-execution/presto_cpp/main/tests/ShutdownOrderTest.cpp
index b8e82046b9892..87cbb6764ee60 100644
--- a/presto-native-execution/presto_cpp/main/tests/ShutdownOrderTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/ShutdownOrderTest.cpp
@@ -173,4 +173,41 @@ TEST(ShutdownOrderTest, promiseChainDispatchesSafelyDuringShutdown) {
   // call would dispatch to a freed executor.
 }
 
+// /v1/expressions tasks run on httpSrvCpuExecutor_ and pass
+// driverExecutor_.get() into QueryCtx.  httpSrvCpuExecutor_ must be drained
+// before driverExecutor_ is reset; httpSrvCpu->join() serialises the task's
+// driverRawPtr->add() against driverExecutor.reset() — swapping those two calls
+// would be use-after-free.
+TEST(ShutdownOrderTest, httpSrvTasksDrainBeforeDriverExecutorJoined) {
+  auto driverExecutor = std::make_unique<NoKeepAliveExecutor>(
+      std::make_unique<folly::CPUThreadPoolExecutor>(
+          2, std::make_shared<folly::NamedThreadFactory>("TestDriver")));
+  auto httpSrvCpu = std::make_unique<folly::CPUThreadPoolExecutor>(
+      2, std::make_shared<folly::NamedThreadFactory>("TestHttpSrvCPU"));
+
+  auto* driverRawPtr = driverExecutor.get();
+  auto sync = std::make_shared<SyncState>();
+
+  // Simulate a /v1/expressions task: runs on httpSrvCpu, dispatches to
+  // driverExecutor via raw pointer (use-after-free if driverExecutor is gone).
+  httpSrvCpu->add([sync, driverRawPtr]() {
+    sync->exchangeReadyFlag = true;
+    sync->exchangeReadyEvent.notifyAll();
+    sync->proceedEvent.await([sync]() { return sync->proceedFlag.load(); });
+    driverRawPtr->add([]() {});
+    sync->completed = true;
+  });
+
+  sync->exchangeReadyEvent.await(
+      [sync]() { return sync->exchangeReadyFlag.load(); });
+  sync->proceedFlag = true;
+  sync->proceedEvent.notifyAll();
+
+  // Correct order: drain httpSrvCpu before destroying driverExecutor.
+  httpSrvCpu->join();
+  driverExecutor.reset();
+  httpSrvCpu.reset();
+
+  EXPECT_TRUE(sync->completed);
+}
 } // namespace facebook::presto