Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 17 additions & 15 deletions presto-native-execution/presto_cpp/main/PrestoServer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reordering creates the mirror-image hazard

httpSrvCpuExecutor_ is not only the sidecar endpoints' executor. It is also TaskManager's and TaskResource's continuation executor: TaskManager.cpp lines 1104, 1122, 1143, 1157, 1179, 1186, 1213, 1228, 1251, 1294, 1314, 1327 all do .via(httpSrvCpuExecutor_). In particular prestoTask->task->stateChangeFuture(maxWaitMicros).via(httpSrvCpuExecutor_) (TaskManager.cpp:1143) is fulfilled by a driver thread on task state change.

So the dependency runs both ways:

  • sidecar tasks on httpSrvCpu → driverExecutor_ (the direction this PR fixes), and
  • driver threads → continuations scheduled onto httpSrvCpu (the direction this PR breaks).

Joining httpSrvCpu here, before driverCpuExecutor_->join() below, means a driver thread can touch an already-joined pool. What folly does in that case:

  • stopAndJoinAllThreads sets maxThreads_ = 0 and activeThreads_ = 0 (folly/executors/ThreadPoolExecutor.cpp:283-284). A later add() enqueues fine, then ensureActiveThreads() sees active >= total and starts no thread (ThreadPoolExecutor.cpp:531-536) — the task is silently never run.
  • .via(Executor*) acquires a keepalive, and DefaultKeepAliveExecutor::keepAliveAcquire asserts DCHECK(keepAliveCount > 0) (folly/DefaultKeepAliveExecutor.h:141-145). The count is 0 after join, so this fails in debug builds; in release it increments from 0 and the matching release re-posts keepAliveReleaseBaton_.

The only thing making this survivable today is that shutdownServer() destroys httpServer_ and taskManager_ before calling joinExecutors() (lines 947-948, 956). That is not a guarantee: TaskManager::shutdown() waits only on kRunning tasks, and the teardown path explicitly tolerates Velox tasks with outstanding references, logging "Velox task has pending reference on destruction" and continuing (TaskManager.cpp:1540-1555). Driver-side work can therefore still be live when httpSrvCpu is already joined.

Ordering alone cannot satisfy both directions. Suggest removing the raw-pointer dependency instead: hand QueryCtx::create an executor held by a keepalive, or run sidecar expression evaluation on an executor that is not the driver pool. getOptimizedExpressions (lines 224-267) is fully synchronous and only needs an executor for QueryCtx, so it does not need the driver pool specifically.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pdabre12 : I'm generally in favor of the suggestion
"Run sidecar expression evaluation on an executor that is not the driver pool. getOptimizedExpressions (lines 224-267) is fully synchronous and only needs an executor for QueryCtx, so it does not need the driver pool specifically."

This change might be changing current problems to something else. If you can implement this suggestion then that works best.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shutdown-hang risk moves to the front of joinExecutors()

ThreadPoolExecutor::join() calls joinKeepAliveOnce()joinKeepAlive(), which blocks on keepAliveReleaseBaton_.wait() (folly/DefaultKeepAliveExecutor.h:51-55) until every KeepAlive token on this executor is released. Pending long-poll chains (getTaskInfo/getResults with maxWait) hold exactly those tokens.

Previously this wait ran last, after the driver, connector and exchange CPU pools were already drained. It is now the first thing joinExecutors() does, with every other pool still up and no timeout on the wait. Any token whose release depends on a stage that is joined later turns shutdown into a hang rather than a crash. Worth a comment recording the counter-constraint, so the next person does not reorder this blind — and note the existing exchange-before-driver ordering has no such note either.

}

// 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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

safe, but not applied to the sibling handler

The move is safe: result is the future's json::array_t&& and is not read afterwards.

The /v1/velox/plan handler below still does json(response) without the move (around line 2049), and response is equally dead after that point. Either apply the same change there or drop this hunk, so the two sidecar handlers stay consistent.

Minor: the PR description's "Better performance" claim rests entirely on this one line, not on the shutdown reordering.

})
.via(
folly::getKeepAliveToken(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Comment on lines +206 to +207

@pdabre12 pdabre12 Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thats how the above test cases do it as well in ShutdownOrderTest so I think its okay to test it like this for now.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

test does not pin the production ordering

Agreeing with Copilot's comment: this hardcodes httpSrvCpu->join(); driverExecutor.reset(); in the test body, so reverting the PrestoServer.cpp change leaves the test green.

On the "the other tests in this file do it the same way" reply — that is accurate, but weak justification here. Pinning the shutdown order is this file's only purpose, and with two opposing constraints now in play (see the comment on PrestoServer.cpp) nothing stops the next reorder from reintroducing the bug in either direction. Extracting the ordering into a helper that joinExecutors() calls, and driving that helper from the test, would cost little and would actually protect the fix.

driverExecutor.reset();
httpSrvCpu.reset();

EXPECT_TRUE(sync->completed);
}
} // namespace facebook::presto
Loading