diff --git a/CHANGES.md b/CHANGES.md index 7f013581..6e5d614a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -17,7 +17,8 @@ - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set frequently. See the docstring for caveats. - https://github.com/joblib/threadpoolctl/issues/257 + https://github.com/joblib/threadpoolctl/pull/257 + https://github.com/joblib/threadpoolctl/pull/285 - Added official support for Python 3.15, and declared support for Python 3.14 in package metadata. diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 080cc244..5a49ecf3 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -7,6 +7,7 @@ import shutil import subprocess import sys +import time from threading import Thread, current_thread import pytest @@ -21,6 +22,7 @@ ThreadpoolController, _determine_thread_limit_scope, get_cached_controller, + reset_cached_controller, threadpool_info, threadpool_limits, ) @@ -1345,6 +1347,48 @@ def test_get_cached_controller(): assert process.returncode == 17 +def test_get_cached_controller_concurrent_reset_never_returns_none(): + """Concurrent reset must not make ``get_cached_controller()`` return None. + + One thread reads the cache while another invalidates it. A stale + ``ThreadpoolController`` is acceptable; ``None`` is not, because callers + do ``get_cached_controller().limit(...)``. + """ + deadline = time.monotonic() + 0.2 + failures = [] + observed = 0 + + def reader(): + nonlocal observed + while True: + controller = get_cached_controller() + observed += 1 + if controller is None or not isinstance(controller, ThreadpoolController): + failures.append(controller) + return + if time.monotonic() >= deadline: + return + + def resetter(): + while time.monotonic() < deadline: + reset_cached_controller() + + reader_thread = Thread(target=reader) + reset_thread = Thread(target=resetter) + try: + reader_thread.start() + reset_thread.start() + reader_thread.join(timeout=2) + reset_thread.join(timeout=2) + assert not reader_thread.is_alive() + assert not reset_thread.is_alive() + finally: + reset_cached_controller() + + assert observed > 0 + assert failures == [] + + @pytest.mark.skipif( not sys.platform.startswith("linux"), reason="ctypes.util is only avoided on Linux (#225) in Python 3.14", diff --git a/threadpoolctl.py b/threadpoolctl.py index 47bc8e15..ea20990c 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -2026,9 +2026,13 @@ def get_cached_controller() -> ThreadpoolController: `reset_cached_controller()` module level function. """ global _CACHED_CONTROLLER, _CACHED_SYS_MODULES_LEN - # Fast path, with no lock: - if _CACHED_SYS_MODULES_LEN == len(sys.modules) and _CACHED_CONTROLLER is not None: - return _CACHED_CONTROLLER + # Fast path with no lock: snapshot before the `None` check as + # `reset_cached_controller()` and `FlexiBLASController.switch_backend()` + # can concurrently set the global to `None`. A stale controller instance is + # acceptable but returning `None` would break the caller. + controller = _CACHED_CONTROLLER + if _CACHED_SYS_MODULES_LEN == len(sys.modules) and controller is not None: + return controller with _CACHE_LOCK: # Do this once, so there's no race condition with imports in other # threads: @@ -2036,7 +2040,7 @@ def get_cached_controller() -> ThreadpoolController: if _CACHED_CONTROLLER is None or num_modules != _CACHED_SYS_MODULES_LEN: _CACHED_CONTROLLER = ThreadpoolController() _CACHED_SYS_MODULES_LEN = num_modules - return _CACHED_CONTROLLER + return _CACHED_CONTROLLER def reset_cached_controller() -> None: