Executor.map: avoid temporarily exceeding the buffersize while collecting the next result

Author: ebonnal

Lib/concurrent/futures/_base.py

@@ -628,17 +628,18 @@ def result_iterator():
                 # reverse to keep finishing order
                 fs.reverse()
                 while fs:
+                    # Careful not to keep a reference to the popped future
+                    if timeout is None:
+                        result = _result_or_cancel(fs.pop())
+                    else:
+                        result = _result_or_cancel(fs.pop(), end_time - time.monotonic())
                     if (
                         buffersize
                         and (executor := executor_weakref())
                         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
             finally:
                 for future in fs:
                     future.cancel()

Lib/test/test_concurrent_futures/executor.py

@@ -1,9 +1,12 @@
 import itertools
+import operator
 import threading
 import time
 import weakref
 from concurrent import futures
 from operator import add
+from functools import partial
+from contextlib import suppress
 from test import support
 from test.support import Py_GIL_DISABLED
 
@@ -143,6 +146,29 @@ def test_map_buffersize_when_buffer_is_full(self):
             msg="should have fetched only `buffersize` elements from `ints`.",
         )
 
+    def test_map_buffersize_when_error(self):
+        ints = [1, 2, 3, 0, 4, 5, 6]
+        index_of_zero = ints.index(0)
+        ints_iter = iter(ints)
+        buffersize = 2
+        reciprocal = partial(operator.truediv, 1)
+        results = []
+        with suppress(ZeroDivisionError):
+            for result in self.executor.map(
+                reciprocal, ints_iter, buffersize=buffersize
+            ):
+                results.append(result)
+        self.assertEqual(
+            len(results),
+            index_of_zero,
+            msg="should have mapped until reaching the zero.",
+        )
+        self.assertEqual(
+            len(results) + buffersize + len(list(ints_iter)),
+            len(ints),
+            msg="ints should be either processed, or buffered, or not fetched.",
+        )
+
     def test_shutdown_race_issue12456(self):
         # Issue #12456: race condition at shutdown where trying to post a
         # sentinel in the call queue blocks (the queue is full while processes