Protect shutdown with lock. Allow shutdown more than once.

Author: kristjanvalur

src/qasync/__init__.py

@@ -22,6 +22,7 @@
 import time
 from concurrent.futures import Future
 from queue import Queue
+from threading import Lock
 from typing import TYPE_CHECKING, Literal, Tuple, cast, get_args
 
 logger = logging.getLogger(__name__)
@@ -172,45 +173,42 @@ def __init__(self, max_workers=10, stack_size=None):
         self.__workers = [
             _QThreadWorker(self.__queue, i + 1, stack_size) for i in range(max_workers)
         ]
+        self.__shutdown_lock = Lock()
         self.__been_shutdown = False
 
         for w in self.__workers:
             w.start()
 
     def submit(self, callback, *args, **kwargs):
-        if self.__been_shutdown:
-            raise RuntimeError("QThreadExecutor has been shutdown")
+        with self.__shutdown_lock:
+            if self.__been_shutdown:
+                raise RuntimeError("QThreadExecutor has been shutdown")
 
-        future = Future()
-        self._logger.debug(
-            "Submitting callback %s with args %s and kwargs %s to thread worker queue",
-            callback,
-            args,
-            kwargs,
-        )
-        self.__queue.put((future, callback, args, kwargs))
-        return future
+            future = Future()
+            self._logger.debug(
+                "Submitting callback %s with args %s and kwargs %s to thread worker queue",
+                callback,
+                args,
+                kwargs,
+            )
+            self.__queue.put((future, callback, args, kwargs))
+            return future
 
     def map(self, func, *iterables, timeout=None):
         raise NotImplementedError("use as_completed on the event loop")
 
     def shutdown(self, wait=True):
-        if self.__been_shutdown:
-            raise RuntimeError("QThreadExecutor has been shutdown")
-
-        self.__been_shutdown = True
-
-        self._logger.debug("Shutting down")
-        for i in range(len(self.__workers)):
-            # Signal workers to stop
-            self.__queue.put(None)
-        if wait:
-            for w in self.__workers:
-                w.wait()
+        with self.__shutdown_lock:
+            self.__been_shutdown = True
+            self._logger.debug("Shutting down")
+            for i in range(len(self.__workers)):
+                # Signal workers to stop
+                self.__queue.put(None)
+            if wait:
+                for w in self.__workers:
+                    w.wait()
 
     def __enter__(self, *args):
-        if self.__been_shutdown:
-            raise RuntimeError("QThreadExecutor has been shutdown")
         return self
 
     def __exit__(self, *args):

tests/test_qthreadexec.py

@@ -44,15 +44,16 @@ def shutdown_executor():
     return exe
 
 
-def test_shutdown_after_shutdown(shutdown_executor):
-    with pytest.raises(RuntimeError):
-        shutdown_executor.shutdown()
+@pytest.mark.parametrize("wait", [True, False])
+def test_shutdown_after_shutdown(shutdown_executor, wait):
+    # it is safe to shutdown twice
+    shutdown_executor.shutdown(wait=wait)
 
 
 def test_ctx_after_shutdown(shutdown_executor):
-    with pytest.raises(RuntimeError):
-        with shutdown_executor:
-            pass
+    # it is safe to enter and exit the context after shutdown
+    with shutdown_executor:
+        pass
 
 
 def test_submit_after_shutdown(shutdown_executor):
@@ -104,3 +105,16 @@ def test_no_stale_reference_as_result(executor, disable_executor_logging):
     assert collected is True, (
         "Stale reference to executor result not collected within timeout."
     )
+
+
+def test_context(executor):
+    """Test that the context manager will shutdown executor"""
+    with executor:
+        f = executor.submit(lambda: 42)
+        assert f.result() == 42
+
+    # it can be entered again
+    with executor:
+        # but will fail when we submit
+        with pytest.raises(RuntimeError):
+            executor.submit(lambda: 42)