Allow OgaShutdown to be called multiple times

src/generators.cpp

@@ -51,21 +51,29 @@ OrtGlobals::OrtGlobals()
   env_->CreateAndRegisterAllocator(allocator_cpu.GetInfo(), *arena_config);
 }
 
+static std::unique_ptr<OrtGlobals> g_globals;
+static std::mutex g_globals_mutex;
+
 // Ensure Shutdown() has been called before process exit
 struct ValidateShutdown {
   ~ValidateShutdown() {
-    if (GetOrtGlobals()) {
+    if (g_globals) {
       std::cerr << "OGA Error: Shutdown must be called before process exit, please check the documentation for the proper API to call to ensure clean shutdown." << std::endl;
       std::abort();
     }
   }
-};
+} g_shutdown;  // This struct should stay immediately after the g_globals and g_globals_mutex
 
 std::unique_ptr<OrtGlobals>&
 GetOrtGlobals() {
-  static auto globals = std::make_unique<OrtGlobals>();
-  static auto validate = std::make_unique<ValidateShutdown>();  // Must be after the above line so the destructor runs before the above destructor
-  return globals;
+  // Initialize g_globals using the g_globals_mutex
+  if (!g_globals) {
+    std::lock_guard<std::mutex> lock(g_globals_mutex);
+    if (!g_globals)  // Now that we're in the mutex, double check
+      g_globals = std::make_unique<OrtGlobals>();
+  }
+
+  return g_globals;
 }
 
 // Used by Shutdown() to display the counts and types of any leaked objects
@@ -81,7 +89,7 @@ void Shutdown() {
     std::abort();
   }
 
-  GetOrtGlobals().reset();  // Delete now because on process exit is too late
+  g_globals.reset();
 }
 
 OrtEnv& GetOrtEnv() {
@@ -313,6 +321,8 @@ bool Generator::IsSessionTerminated() const {
 }
 
 void Generator::GenerateNextToken() {
+  if (search_->IsDone())
+    throw std::runtime_error("Search is already done, can't generate next token");
   ThrowErrorIfSessionTerminated(state_->session_terminated_);
   if (!computed_logits_)
     throw std::runtime_error("Must call ComputeLogits before GenerateNextToken");

test/c_api_tests.cpp

@@ -348,6 +348,10 @@ TEST(CAPITests, SetTerminate) {
 #endif
 }
 
+TEST(CAPITests, Shutdown) {
+  OgaShutdown(); // Shutdown in the middle, and see if the next test works properly
+}
+
 #if TEST_PHI2
 
 struct Phi2Test {