From be1a504228db4185a4ad5defe1b57d4df2bc8b2f Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Wed, 19 Nov 2025 16:36:30 +1100
Subject: [PATCH] [orc-rt] Simplify Session shutdown. (#168664)

Moves all Session member variables dedicated to shutdown into a new
ShutdownInfo struct, and uses the presence / absence of this struct as
the flag to indicate that we've entered the "shutting down" state. This
simplifies the implementation of the shutdown process.
---
 orc-rt/include/orc-rt/Session.h | 15 ++++++------
 orc-rt/lib/executor/Session.cpp | 43 ++++++++++++++-------------------
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/orc-rt/include/orc-rt/Session.h b/orc-rt/include/orc-rt/Session.h
index 367cdb9a97b6..529aac6a2fad 100644
--- a/orc-rt/include/orc-rt/Session.h
+++ b/orc-rt/include/orc-rt/Session.h
@@ -75,21 +75,22 @@ public:
   }
 
 private:
-  void shutdownNext(Error Err,
-                    std::vector<std::unique_ptr<ResourceManager>> RemainingRMs);
+  struct ShutdownInfo {
+    bool Complete = false;
+    std::condition_variable CompleteCV;
+    std::vector<std::unique_ptr<ResourceManager>> ResourceMgrs;
+    std::vector<OnShutdownCompleteFn> OnCompletes;
+  };
 
+  void shutdownNext(Error Err);
   void shutdownComplete();
 
   std::unique_ptr<TaskDispatcher> Dispatcher;
   ErrorReporterFn ReportError;
 
-  enum class SessionState { Running, ShuttingDown, Shutdown };
-
   std::mutex M;
-  SessionState State = SessionState::Running;
-  std::condition_variable StateCV;
   std::vector<std::unique_ptr<ResourceManager>> ResourceMgrs;
-  std::vector<OnShutdownCompleteFn> ShutdownCallbacks;
+  std::unique_ptr<ShutdownInfo> SI;
 };
 
 inline orc_rt_SessionRef wrap(Session *S) noexcept {
diff --git a/orc-rt/lib/executor/Session.cpp b/orc-rt/lib/executor/Session.cpp
index fafa13b1cbb0..3ee9bad60c5b 100644
--- a/orc-rt/lib/executor/Session.cpp
+++ b/orc-rt/lib/executor/Session.cpp
@@ -17,66 +17,59 @@ namespace orc_rt {
 Session::~Session() { waitForShutdown(); }
 
 void Session::shutdown(OnShutdownCompleteFn OnShutdownComplete) {
-  std::vector<std::unique_ptr<ResourceManager>> ToShutdown;
-
   {
     std::scoped_lock<std::mutex> Lock(M);
-    ShutdownCallbacks.push_back(std::move(OnShutdownComplete));
-
-    // If somebody else has already called shutdown then there's nothing further
-    // for us to do here.
-    if (State >= SessionState::ShuttingDown)
+    if (SI) {
+      SI->OnCompletes.push_back(std::move(OnShutdownComplete));
       return;
+    }
 
-    State = SessionState::ShuttingDown;
-    std::swap(ResourceMgrs, ToShutdown);
+    SI = std::make_unique<ShutdownInfo>();
+    SI->OnCompletes.push_back(std::move(OnShutdownComplete));
+    std::swap(SI->ResourceMgrs, ResourceMgrs);
   }
 
-  shutdownNext(Error::success(), std::move(ToShutdown));
+  shutdownNext(Error::success());
 }
 
 void Session::waitForShutdown() {
   shutdown([]() {});
   std::unique_lock<std::mutex> Lock(M);
-  StateCV.wait(Lock, [&]() { return State == SessionState::Shutdown; });
+  SI->CompleteCV.wait(Lock, [&]() { return SI->Complete; });
 }
 
-void Session::shutdownNext(
-    Error Err, std::vector<std::unique_ptr<ResourceManager>> RemainingRMs) {
+void Session::shutdownNext(Error Err) {
   if (Err)
     reportError(std::move(Err));
 
-  if (RemainingRMs.empty())
+  if (SI->ResourceMgrs.empty())
     return shutdownComplete();
 
-  auto NextRM = std::move(RemainingRMs.back());
-  RemainingRMs.pop_back();
-  NextRM->shutdown(
-      [this, RemainingRMs = std::move(RemainingRMs)](Error Err) mutable {
-        shutdownNext(std::move(Err), std::move(RemainingRMs));
-      });
+  // Get the next ResourceManager to shut down.
+  auto NextRM = std::move(SI->ResourceMgrs.back());
+  SI->ResourceMgrs.pop_back();
+  NextRM->shutdown([this](Error Err) { shutdownNext(std::move(Err)); });
 }
 
 void Session::shutdownComplete() {
 
   std::unique_ptr<TaskDispatcher> TmpDispatcher;
-  std::vector<OnShutdownCompleteFn> TmpShutdownCallbacks;
   {
     std::lock_guard<std::mutex> Lock(M);
     TmpDispatcher = std::move(Dispatcher);
-    TmpShutdownCallbacks = std::move(ShutdownCallbacks);
   }
 
   TmpDispatcher->shutdown();
 
-  for (auto &OnShutdownComplete : TmpShutdownCallbacks)
+  for (auto &OnShutdownComplete : SI->OnCompletes)
     OnShutdownComplete();
 
   {
     std::lock_guard<std::mutex> Lock(M);
-    State = SessionState::Shutdown;
+    SI->Complete = true;
   }
-  StateCV.notify_all();
+
+  SI->CompleteCV.notify_all();
 }
 
 } // namespace orc_rt