concurrent.futures.Executor.map: yield buffered results and raise after an executor shutdown

Author: ebonnal

Lib/concurrent/futures/_base.py

@@ -305,10 +305,12 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
     return DoneAndNotDoneFutures(done, fs - done)
 
 
-def _result_or_cancel(fut, timeout=None):
+def _result_or_cancel(fut, end_time=None):
     try:
         try:
-            return fut.result(timeout)
+            if end_time is None:
+                return fut.result()
+            return fut.result(end_time - time.monotonic())
         finally:
             fut.cancel()
     finally:
@@ -633,20 +635,20 @@ def map(self, fn, *iterables, timeout=None, chunksize=1, buffersize=None):
         if buffersize is not None and buffersize < 1:
             raise ValueError("buffersize must be None or > 0")
 
-        if timeout is not None:
+        if timeout is None:
+            end_time = None
+        else:
             end_time = timeout + time.monotonic()
 
         zipped_iterables = zip(*iterables)
         if buffersize:
             fs = collections.deque(
                 self.submit(fn, *args) for args in islice(zipped_iterables, buffersize)
             )
+            executor = self
         else:
             fs = [self.submit(fn, *args) for args in zipped_iterables]
-
-        # Use a weak reference to ensure that the executor can be garbage
-        # collected independently of the result_iterator closure.
-        executor_weakref = weakref.ref(self)
+            executor = None
 
         # Yield must be hidden in closure so that the futures are submitted
         # before the first iterator value is required.
@@ -655,17 +657,17 @@ def result_iterator():
                 # reverse to keep finishing order
                 fs.reverse()
                 while fs:
-                    if (
-                        buffersize
-                        and (executor := executor_weakref())
-                        and (args := next(zipped_iterables, None))
-                    ):
-                        fs.appendleft(executor.submit(fn, *args))
+                    if buffersize and (args := next(zipped_iterables, None)):
+                        try:
+                            fs.appendleft(executor.submit(fn, *args))
+                        except RuntimeError:
+                            # 'cannot schedule new futures after shutdown' error
+                            # yield from buffer:
+                            while fs:
+                                yield _result_or_cancel(fs.pop(), end_time)
+                            raise
                     # Careful not to keep a reference to the popped future
-                    if timeout is None:
-                        yield _result_or_cancel(fs.pop())
-                    else:
-                        yield _result_or_cancel(fs.pop(), end_time - time.monotonic())
+                    yield _result_or_cancel(fs.pop(), end_time)
             finally:
                 for future in fs:
                     future.cancel()

Lib/test/test_concurrent_futures/executor.py

@@ -148,6 +148,48 @@ def test_map_buffersize_on_multiple_infinite_iterables(self):
         self.assertEqual(next(res, None), 2)
         self.assertEqual(next(res, None), 4)
 
+    @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+    def test_map_buffersize_shutdown_yield_buffer_and_raise(self):
+        ints = range(4)
+        fn = str
+
+        def test(buffersize, nexts_before_shutdown, wait):
+            expected_results = map(fn, ints)
+
+            executor = self.executor_type(max_workers=1)
+            results = executor.map(fn, ints, buffersize=buffersize)
+            # iterate before shutdown
+            for _ in range(nexts_before_shutdown):
+                self.assertEqual(next(results), next(expected_results))
+            time.sleep(0.1)
+            executor.shutdown(wait=wait)
+            # iterate over buffered results after shutdown
+            for _ in range(min(buffersize, len(ints) - nexts_before_shutdown)):
+                self.assertEqual(next(results), next(expected_results))
+
+            if buffersize + nexts_before_shutdown >= len(ints):
+                # case where all tasks were submitted before shutdown
+                with self.assertRaises(StopIteration):
+                    next(results)
+            else:
+                # buffer exhausted, the error is raised next
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "cannot schedule new futures after shutdown",
+                ):
+                    next(results)
+                with self.assertRaises(StopIteration):
+                    next(results)
+
+        for buffersize in [3]:
+            for nexts_before_shutdown in range(buffersize):
+                for wait in (True, False):
+                    with self.subTest(
+                        buffersize=buffersize,
+                        nexts_before_shutdown=nexts_before_shutdown,
+                    ):
+                        test(buffersize, nexts_before_shutdown, wait)
+
     def test_map_buffersize_on_empty_iterable(self):
         res = self.executor.map(str, [], buffersize=2)
         self.assertIsNone(next(res, None))