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

@@ -8,7 +8,7 @@
 import subprocess
 import sys
 import time
-from threading import Thread, current_thread
+from threading import Lock, Thread, current_thread
 
 import pytest
 
@@ -792,11 +792,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

@@ -557,24 +557,26 @@ def set_num_threads(self, num_threads):
         # A shared process-wide thread pool, so don't want different Python
         # threads fighting over limits:
         this_thread = current_thread()
-
-        match self._SHARED_STATE.controlling_thread:
-            case None:
-                # 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:
-                        self._SHARED_STATE.controlling_thread = this_thread
-                # Controlling thread should now be set, to either this thread
-                # or another, so try again:
-                return self.set_num_threads(num_threads)
-            case controlling_thread if controlling_thread is this_thread:
-                # This thread controls OpenBLAS:
-                return self._real_set_num_threads(num_threads)
-            case _:
-                # Some other thread controlls OpenBLAS:
-                return
+        controlling_thread = self._SHARED_STATE.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
+                    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:
+            return self.set_num_threads(num_threads)
+        elif controlling_thread is this_thread:
+            # This thread controls OpenBLAS:
+            return self._real_set_num_threads(num_threads)
+        else:
+            # Some other thread controlls OpenBLAS: do nothing.
+            return
 
     def _real_set_num_threads(self, num_threads):
         # The OpenBLAS limiting API is process-wide, and we want current thread