From 40f3119af75aed783d8a3ebac61a9c4590a91473 Mon Sep 17 00:00:00 2001
From: Delio Vicini <vicini@google.com>
Date: Mon, 9 Sep 2024 16:14:04 +0200
Subject: [PATCH] Fix shutdown from non-main thread.

---
 include/mitsuba/core/thread.h |  3 +++
 src/core/logger.cpp           |  5 ++---
 src/core/thread.cpp           |  9 +++++++--
 src/python/main.cpp           | 15 +++++++--------
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/include/mitsuba/core/thread.h b/include/mitsuba/core/thread.h
index 37c50fff4..b93789307 100644
--- a/include/mitsuba/core/thread.h
+++ b/include/mitsuba/core/thread.h
@@ -132,6 +132,9 @@ class MI_EXPORT_LIB Thread : public Object {
     /// Return the current thread
     static Thread *thread();
 
+    /// Return whether the current thread data structure has been initialized.
+    static bool has_initialized_thread();
+
     /// Sleep for a certain amount of time (in milliseconds)
     static void sleep(uint32_t ms);
 
diff --git a/src/core/logger.cpp b/src/core/logger.cpp
index 79ab8c9d9..fbf1c4f21 100644
--- a/src/core/logger.cpp
+++ b/src/core/logger.cpp
@@ -122,9 +122,8 @@ void Logger::static_initialization() {
 #endif
 }
 
-void Logger::static_shutdown() {
-    Thread::thread()->set_logger(nullptr);
-}
+// Removal of logger from main thread is handled in thread.cpp.
+void Logger::static_shutdown() { }
 
 size_t Logger::appender_count() const {
     return d->appenders.size();
diff --git a/src/core/thread.cpp b/src/core/thread.cpp
index 63a43a3e8..fce115b85 100644
--- a/src/core/thread.cpp
+++ b/src/core/thread.cpp
@@ -186,6 +186,10 @@ Thread* Thread::thread() {
     return self_val;
 }
 
+bool Thread::has_initialized_thread() {
+    return self != nullptr;
+}
+
 bool Thread::is_running() const {
     return d->running;
 }
@@ -558,8 +562,9 @@ void Thread::static_shutdown() {
         }
         thread_registry.clear();
     }
-
-    thread()->d->running = false;
+    main_thread->d->logger    = nullptr;
+    main_thread->d->fresolver = nullptr;
+    main_thread->d->running   = false;
     self = nullptr;
     main_thread = nullptr;
     #if defined(__linux__) || defined(__APPLE__)
diff --git a/src/python/main.cpp b/src/python/main.cpp
index f5c6456b3..171c6b45f 100644
--- a/src/python/main.cpp
+++ b/src/python/main.cpp
@@ -182,13 +182,12 @@ NB_MODULE(mitsuba_ext, m) {
         Class::static_remove_functors();
         StructConverter::static_shutdown();
 
-        /* When the main thread's lifetime was shared with Python, it would be
-         * detected as a nanbonid leak as the last reference is held by C++.
-         * Calling `Thread::static_shutdown()` deletes this last reference,
-         * however the main thread is still needed to clean up other parts of
-         * the system. The solution is therefore to temporarily clean it up
-         * and then re-build it such that Python can no longer track it. */
-        if(Thread::thread()->self_py()) {
+        /* Potentially re-initialize the threading system:
+         * 1) Deleting and re-initializing threading prevents a Nanobind leak
+         * if the lifetime of the main thread was shared with Python.
+         * 2) Additionally, this can ensure correct shutdown if the shutdown
+         * happens on another thread than the initialization. */
+        if (!Thread::has_initialized_thread() || Thread::thread()->self_py()) {
             Thread::static_shutdown();
             Thread::static_initialization();
         }
@@ -202,7 +201,7 @@ NB_MODULE(mitsuba_ext, m) {
         }
     }));
 
-    /* Callback function cleanup static data strucutres, this should be called
+    /* Callback function cleanup static data structures, this should be called
      * when the module is being deallocated */
     nanobind_module_def_mitsuba_ext.m_free = [](void *) {
         Profiler::static_shutdown();