gh-108518: Do not cancel remaining calls in Executor.map() on error (GH-109497)

If a call raises an exception, the remaining calls are no longer cancelled
and iteration can be continued.  Use the close() method of the returned
iterator to cancel them.

Co-authored-by: Meng Xiangzhuo <aumo@foxmail.com>
Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>

Author: 3 people

Doc/library/concurrent.futures.rst

@@ -61,11 +61,18 @@ Executor Objects
       The returned iterator raises a :exc:`TimeoutError`
       if :meth:`~iterator.__next__` is called and the result isn't available
       after *timeout* seconds from the original call to :meth:`Executor.map`.
-      *timeout* can be an int or a float.  If *timeout* is not specified or
+      *timeout* can be an int or a float.
+      It cancels all future calls of *fn* and closes the iterator.
+      If *timeout* is not specified or
       ``None``, there is no limit to the wait time.
 
       If a *fn* call raises an exception, then that exception will be
       raised when its value is retrieved from the iterator.
+      It does not cancel future calls of *fn*.
+
+      The returned iterator has method :meth:`!close` which cancels all
+      future calls of *fn* and discards the results of already finished calls
+      if they are available.
 
       When using :class:`ProcessPoolExecutor`, this method chops *iterables*
       into a number of chunks which it submits to the pool as separate
@@ -82,6 +89,10 @@ Executor Objects
       .. versionchanged:: 3.14
          Added the *buffersize* parameter.
 
+      .. versionchanged:: next
+         The returned iterator is no longer automatically closed if a *fn*
+         call raises an exception.
+
    .. method:: shutdown(wait=True, *, cancel_futures=False)
 
       Signal the executor that it should free any resources that it is using

Doc/whatsnew/3.16.rst

@@ -276,6 +276,15 @@ ctypes
   (Contributed by Peter Bierma in :gh:`153903`.)
 
 
+concurrent.futures
+------------------
+
+* The iterator returned by :meth:`concurrent.futures.Executor.map` is no longer
+  automatically closed if a function call raises an exception.
+  Use method :meth:`!close` to explicitly close the iterator.
+  (Contributed by xzmeng and Serhiy Storchaka in :gh:`108518`.)
+
+
 encodings
 ---------
 

Lib/concurrent/futures/_base.py

@@ -309,7 +309,11 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
 def _result_or_cancel(fut, timeout=None):
     try:
         try:
-            return fut.result(timeout)
+            return (fut.result(timeout), None)
+        except TimeoutError:
+            raise
+        except BaseException as exc:
+            return (None, exc)
         finally:
             fut.cancel()
     finally:
@@ -592,6 +596,7 @@ def _get_snapshot(self):
 
     __class_getitem__ = classmethod(types.GenericAlias)
 
+
 class Executor(object):
     """This is an abstract base class for concrete asynchronous executors."""
 
@@ -638,7 +643,10 @@ def map(self, fn, *iterables, timeout=None, chunksize=1, buffersize=None):
             raise TypeError("buffersize must be an integer or None")
         if buffersize is not None and buffersize < 1:
             raise ValueError("buffersize must be None or > 0")
+        return _MapResultIterator(self._map(fn, *iterables, timeout=timeout,
+                                            buffersize=buffersize))
 
+    def _map(self, fn, *iterables, timeout=None, buffersize=None):
         if timeout is not None:
             end_time = timeout + time.monotonic()
 
@@ -701,6 +709,24 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
 
+class _MapResultIterator:
+    """The iterator returned by map()."""
+    def __init__(self, gen):
+        self.gen = gen
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value, exc = next(self.gen)
+        if exc is not None:
+            raise exc
+        return value
+
+    def close(self):
+        self.gen.close()
+
+
 class BrokenExecutor(RuntimeError):
     """
     Raised when an executor has become non-functional after a severe failure.

Lib/concurrent/futures/process.py

@@ -200,7 +200,14 @@ def _process_chunk(fn, chunk):
     This function is run in a separate process.
 
     """
-    return [fn(*args) for args in chunk]
+    results = []
+    for args in chunk:
+        try:
+            result = (fn(*args), None)
+        except BaseException as exc:
+            result = (None, exc)
+        results.append(result)
+    return results
 
 
 def _sendback_result(result_queue, work_id, result=None, exception=None,
@@ -963,7 +970,7 @@ def map(self, fn, *iterables, timeout=None, chunksize=1, buffersize=None):
                               itertools.batched(zip(*iterables), chunksize),
                               timeout=timeout,
                               buffersize=buffersize)
-        return _chain_from_iterable_of_lists(results)
+        return _base._MapResultIterator(_chain_from_iterable_of_lists(results))
 
     def shutdown(self, wait=True, *, cancel_futures=False):
         with self._shutdown_lock:

Lib/test/test_concurrent_futures/executor.py

@@ -71,21 +71,30 @@ def test_map(self):
 
     @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
     def test_map_exception(self):
-        i = self.executor.map(divmod, [1, 1, 1, 1], [2, 3, 0, 5])
-        self.assertEqual(i.__next__(), (0, 1))
-        self.assertEqual(i.__next__(), (0, 1))
-        with self.assertRaises(ZeroDivisionError):
-            i.__next__()
+        i = self.executor.map(divmod, [5, 5, 5, 5], [2, 3, 0, 5])
+        self.assertEqual(next(i), (2, 1))
+        self.assertEqual(next(i), (1, 2))
+        self.assertRaises(ZeroDivisionError, next, i)
+        self.assertEqual(next(i), (1, 0))
+        self.assertRaises(StopIteration, next, i)
+        self.assertRaises(StopIteration, next, i)
+
+        i = self.executor.map(divmod, [5, 5, 5, 5], [2, 0, 3, 5], chunksize=3)
+        self.assertEqual(next(i), (2, 1))
+        self.assertRaises(ZeroDivisionError, next, i)
+        self.assertEqual(next(i), (1, 2))
+        self.assertEqual(next(i), (1, 0))
+        self.assertRaises(StopIteration, next, i)
+        self.assertRaises(StopIteration, next, i)
 
     @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
     @support.requires_resource('walltime')
     def test_map_timeout(self):
         results = []
+        i = self.executor.map(time.sleep, [0, 0, 6], timeout=5)
         try:
-            for i in self.executor.map(time.sleep,
-                                       [0, 0, 6],
-                                       timeout=5):
-                results.append(i)
+            for result in i:
+                results.append(result)
         except futures.TimeoutError:
             pass
         else:
@@ -95,6 +104,24 @@ def test_map_timeout(self):
         # take longer than the specified timeout.
         self.assertIn(results, ([None, None], [None], []))
 
+        # The remaining calls are cancelled, so the iterator is exhausted.
+        self.assertRaises(StopIteration, next, i)
+        self.assertRaises(StopIteration, next, i)
+
+    @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+    def test_map_close(self):
+        i = self.executor.map(divmod, [5, 5, 5, 5], [2, 0, 3, 5])
+        self.assertEqual(next(i), (2, 1))
+        i.close()
+        self.assertRaises(StopIteration, next, i)
+        self.assertRaises(StopIteration, next, i)
+
+        i = self.executor.map(divmod, [5, 5, 5, 5], [2, 0, 3, 5], chunksize=3)
+        self.assertEqual(next(i), (2, 1))
+        i.close()
+        self.assertRaises(StopIteration, next, i)
+        self.assertRaises(StopIteration, next, i)
+
     def test_map_buffersize_type_validation(self):
         for buffersize in ("foo", 2.0):
             with self.subTest(buffersize=buffersize):

Misc/NEWS.d/next/Library/2024-02-04-13-56-48.gh-issue-108518.6NCPk_.rst

@@ -0,0 +1,3 @@
+The iterator returned by :meth:`concurrent.futures.Executor.map` is no longer
+automatically closed if a function call raises an exception.
+Use method :meth:`!close` to explicitly close the iterator.