From b70dd3bbafa4bee1121d6c0296b8b310bcb8fb8e Mon Sep 17 00:00:00 2001
From: Delio Vicini <vicini@google.com>
Date: Mon, 9 Sep 2024 12:58:13 +0200
Subject: [PATCH] Guard against GIL acquisition post Python shutdown.

---
 src/python/main.cpp | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/python/main.cpp b/src/python/main.cpp
index b68567fd0..f5c6456b3 100644
--- a/src/python/main.cpp
+++ b/src/python/main.cpp
@@ -1,3 +1,5 @@
+#include <shared_mutex>
+
 #include <mitsuba/core/bitmap.h>
 #include <mitsuba/core/jit.h>
 #include <mitsuba/core/logger.h>
@@ -6,6 +8,11 @@
 #include <mitsuba/core/profiler.h>
 #include <mitsuba/python/python.h>
 
+
+// Flag & mutex to explicitly track if Python is still available.
+static bool ready_flag = true;
+static std::shared_mutex ready_mutex;
+
 // core
 MI_PY_DECLARE(atomic);
 MI_PY_DECLARE(filesystem);
@@ -80,6 +87,13 @@ NB_MODULE(mitsuba_ext, m) {
             Py_INCREF(o);
         },
         [](PyObject *o) noexcept {
+            std::shared_lock ready_guard(ready_mutex);
+            /* If the Python interpreter has already been shut down, we can no longer use
+             * its reference counting mechanism. This leaks memory on interpreter shutdown.
+             * However, this should only affect static, thread local objects, for which
+             * enforcing an orderly shutdown is difficult. */
+            if (!ready_flag)
+                return;
             nb::gil_scoped_acquire guard;
             Py_DECREF(o);
         }
@@ -178,6 +192,14 @@ NB_MODULE(mitsuba_ext, m) {
             Thread::static_shutdown();
             Thread::static_initialization();
         }
+
+        /* After this point, we can no longer guarantee that the Python
+         * interpreter is available on all threads. */
+        {
+            nb::gil_scoped_release g;
+            std::unique_lock ready_guard(ready_mutex);
+            ready_flag = false;
+        }
     }));
 
     /* Callback function cleanup static data strucutres, this should be called