FIX Dead threads lose control

CHANGES.md

@@ -13,6 +13,7 @@
   thread controlls OpenBLAS in these cases, and decorator and context manager
   limit APIs will automatically call that API when they are done.
   https://github.com/joblib/threadpoolctl/pull/275
+  https://github.com/joblib/threadpoolctl/pull/289
 
 - Add a new API to get a cached `ThreadpoolController` instance
   (`get_cached_controller()`), as an optimization for cases where limits are set

tests/test_threadpoolctl.py

@@ -7,7 +7,7 @@
 import shutil
 import subprocess
 import sys
-from threading import Thread, current_thread
+from threading import Lock, Thread, current_thread
 
 import pytest
 
@@ -801,11 +801,41 @@ def clear_and_set():
     run_in_thread(openblas_lib.set_num_threads, 1)
     assert openblas_lib.get_num_threads() == 1
 
+    # Dead threads lose control:
+    openblas_lib.set_num_threads(orig_num_threads)
+    assert openblas_lib.get_num_threads() == orig_num_threads
+
+    # Release control:
+    openblas_lib.clear_controlling_thread()
+
     # clear_controlling_thread(force=True) allows any thread to clear:
+    was_set, finish = Lock(), Lock()
+    was_set.acquire()
+    finish.acquire()
+
+    def set_and_exit_after_lock_released():
+        openblas_lib.set_num_threads(1)
+        was_set.release()
+        finish.acquire()
+
+    t = Thread(target=set_and_exit_after_lock_released)
+    t.start()
+    was_set.acquire()
+    assert openblas_lib.get_num_threads() == 1
+
+    # This thread can't set since it's not the controller:
+    openblas_lib.set_num_threads(orig_num_threads)
+    assert openblas_lib.get_num_threads() == 1
+
+    # But when forced clear happens, we can set:
     openblas_lib.clear_controlling_thread(force=True)
     openblas_lib.set_num_threads(orig_num_threads)
     assert openblas_lib.get_num_threads() == orig_num_threads
 
+    # Clean up
+    finish.release()
+    t.join()
+
 
 def _openblas_controlling_thread():
     openblas_controller = ThreadpoolController().select(internal_api="openblas")

threadpoolctl.py

@@ -555,12 +555,17 @@ def set_num_threads(self, num_threads):
         this_thread = current_thread()
 
         match self._SHARED_STATE.controlling_thread:
-            case None:
+            case controlling_thread if (
+                controlling_thread is None or not controlling_thread.is_alive()
+            ):
                 # No thread controls OpenBLAS, so try to take control:
                 with self._SHARED_STATE.lock:
                     # Another thread might have changed it, so ensure that's
                     # not the case:
-                    if self._SHARED_STATE.controlling_thread is None:
+                    if (
+                        self._SHARED_STATE.controlling_thread is None
+                        or not self._SHARED_STATE.controlling_thread.is_alive()
+                    ):
                         self._SHARED_STATE.controlling_thread = this_thread
                 # Controlling thread should now be set, to either this thread
                 # or another, so try again: