From 79cb013f1f67d5565700e528cb5d8ef079fe5e8a Mon Sep 17 00:00:00 2001
From: Nikita Shulga <nikita.shulga@gmail.com>
Date: Sat, 19 Dec 2020 12:45:19 -0800
Subject: [PATCH] fix: allow users to avoid thread termination in
 scoped_released (#2657)

* Avoid thread termination in scoped_released

Do not call `PyEval_RestoreThread()` from `~gil_scoped_release()` if python runtime is finalizing, as it will result in thread termination in Python runtime newer than 3.6, as documented in https://docs.python.org/3/c-api/init.html#c.PyEval_RestoreThread
Similarly do not call `PyThreadState_DeleteCurrent` from `~gil_scoped_acquire()` if runtime is finalizing.

Discovered while debugging PyTorch crash using Python-3.9 described in  https://github.com/pytorch/pytorch/issues/47776

* Simplify _Py_IsFinalizing() availability check

* Fix typo

* Add version agnostic `detail::finalization_guard()`

* Move `finalization_guard` to detail/common.h

And rename it to `is_finalizing`

* Move `is_finalizing()` back to pybind11.h

* Simplify `is_finalizing()` check

One should follow documentation rather than make any assumptions

* feat: disarm

* docs: fix comment

Co-authored-by: Henry Schreiner <henryschreineriii@gmail.com>
---
 include/pybind11/pybind11.h | 38 +++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/include/pybind11/pybind11.h b/include/pybind11/pybind11.h
index 96f3ff805..104f32206 100644
--- a/include/pybind11/pybind11.h
+++ b/include/pybind11/pybind11.h
@@ -2113,12 +2113,22 @@ public:
                     pybind11_fail("scoped_acquire::dec_ref(): internal error!");
             #endif
             PyThreadState_Clear(tstate);
-            PyThreadState_DeleteCurrent();
+            if (active)
+                PyThreadState_DeleteCurrent();
             PYBIND11_TLS_DELETE_VALUE(detail::get_internals().tstate);
             release = false;
         }
     }
 
+    /// This method will disable the PyThreadState_DeleteCurrent call and the
+    /// GIL won't be acquired. This method should be used if the interpreter
+    /// could be shutting down when this is called, as thread deletion is not
+    /// allowed during shutdown. Check _Py_IsFinalizing() on Python 3.7+, and
+    /// protect subsequent code.
+    PYBIND11_NOINLINE void disarm() {
+        active = false;
+    }
+
     PYBIND11_NOINLINE ~gil_scoped_acquire() {
         dec_ref();
         if (release)
@@ -2127,6 +2137,7 @@ public:
 private:
     PyThreadState *tstate = nullptr;
     bool release = true;
+    bool active = true;
 };
 
 class gil_scoped_release {
@@ -2142,10 +2153,22 @@ public:
             PYBIND11_TLS_DELETE_VALUE(key);
         }
     }
+
+    /// This method will disable the PyThreadState_DeleteCurrent call and the
+    /// GIL won't be acquired. This method should be used if the interpreter
+    /// could be shutting down when this is called, as thread deletion is not
+    /// allowed during shutdown. Check _Py_IsFinalizing() on Python 3.7+, and
+    /// protect subsequent code.
+    PYBIND11_NOINLINE void disarm() {
+        active = false;
+    }
+
     ~gil_scoped_release() {
         if (!tstate)
             return;
-        PyEval_RestoreThread(tstate);
+        // `PyEval_RestoreThread()` should not be called if runtime is finalizing
+        if (active)
+            PyEval_RestoreThread(tstate);
         if (disassoc) {
             auto key = detail::get_internals().tstate;
             PYBIND11_TLS_REPLACE_VALUE(key, tstate);
@@ -2154,6 +2177,7 @@ public:
 private:
     PyThreadState *tstate;
     bool disassoc;
+    bool active = true;
 };
 #elif defined(PYPY_VERSION)
 class gil_scoped_acquire {
@@ -2161,6 +2185,7 @@ class gil_scoped_acquire {
 public:
     gil_scoped_acquire() { state = PyGILState_Ensure(); }
     ~gil_scoped_acquire() { PyGILState_Release(state); }
+    void disarm() {}
 };
 
 class gil_scoped_release {
@@ -2168,10 +2193,15 @@ class gil_scoped_release {
 public:
     gil_scoped_release() { state = PyEval_SaveThread(); }
     ~gil_scoped_release() { PyEval_RestoreThread(state); }
+    void disarm() {}
 };
 #else
-class gil_scoped_acquire { };
-class gil_scoped_release { };
+class gil_scoped_acquire {
+    void disarm() {}
+};
+class gil_scoped_release {
+    void disarm() {}
+};
 #endif
 
 error_already_set::~error_already_set() {