Support the cancel_futures parameter to executor.shutdown()

Author: kristjanvalur

src/qasync/__init__.py

@@ -216,10 +216,17 @@ def submit(self, callback, *args, **kwargs):
     def map(self, func, *iterables, timeout=None):
         raise NotImplementedError("use as_completed on the event loop")
 
-    def shutdown(self, wait=True):
+    def shutdown(self, wait=True, *, cancel_futures=False):
         with self.__shutdown_lock:
             self.__been_shutdown = True
             self._logger.debug("Shutting down")
+            if cancel_futures:
+                # pop all the futures and cancel them
+                while not self.__queue.empty():
+                    item = self.__queue.get_nowait()
+                    if item is not None:
+                        future, _, _, _ = item
+                        future.cancel()
             for i in range(len(self.__workers)):
                 # Signal workers to stop
                 self.__queue.put(None)

tests/test_qthreadexec.py

@@ -4,7 +4,9 @@
 # BSD License
 import logging
 import threading
+import time
 import weakref
+from concurrent.futures import CancelledError
 
 import pytest
 
@@ -118,3 +120,28 @@ def test_context(executor):
         # but will fail when we submit
         with pytest.raises(RuntimeError):
             executor.submit(lambda: 42)
+
+
+@pytest.mark.parametrize("cancel", [True, False])
+def test_shutdown_cancel_futures(executor, cancel):
+    """Test that shutdown with cancel_futures=True cancels all remaining futures in the queue."""
+
+    def task():
+        time.sleep(0.01)
+
+    # Submit ten tasks to the executor
+    futures = [executor.submit(task) for _ in range(10)]
+    # shut it down
+    executor.shutdown(cancel_futures=cancel)
+
+    cancels = 0
+    for future in futures:
+        try:
+            future.result(timeout=0.01)
+        except CancelledError:
+            cancels += 1
+
+    if cancel:
+        assert cancels > 0
+    else:
+        assert cancels == 0