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,15 @@ 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:
+                            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,18 @@ 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_after_context_manager_shutdown(self):
+        with self.executor_type(max_workers=1) as executor:
+            res = executor.map(str, range(4), buffersize=2)
+        assert next(res) == "0"
+        assert next(res) == "1"
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "cannot schedule new futures after shutdown",
+        ):
+            next(res)
+
     def test_map_buffersize_on_empty_iterable(self):
         res = self.executor.map(str, [], buffersize=2)
         self.assertIsNone(next(res, None))