Fix shutdown from non-main thread.

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);
 

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();

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__)

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();