concurrent.futures.Executor.map: raise after 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:
@@ -606,20 +608,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.
@@ -628,17 +630,10 @@ def result_iterator():
                 # reverse to keep finishing order
                 fs.reverse()
                 while fs:
-                    if (
-                        buffersize
-                        and (executor := executor_weakref())
-                        and (args := next(zipped_iterables, None))
-                    ):
+                    if buffersize and (args := next(zipped_iterables, None)):
                         fs.appendleft(executor.submit(fn, *args))
                     # 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

@@ -139,6 +139,24 @@ 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_raise_on_shutdown(self):
+        ints = range(4)
+        fn = str
+        buffersize = 2
+        expected_results = map(fn, ints)
+        for nexts in [0, 1]:
+            with self.subTest(nexts=nexts):
+                with self.executor_type(max_workers=1) as executor:
+                    results = executor.map(fn, ints, buffersize=buffersize)
+                    for _ in range(nexts):
+                        self.assertEqual(next(results), next(expected_results))
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "cannot schedule new futures after shutdown",
+                ):
+                    next(results)
+
     def test_map_buffersize_on_empty_iterable(self):
         res = self.executor.map(str, [], buffersize=2)
         self.assertIsNone(next(res, None))