bpo-30782: Allow limiting the number of concurrent tasks in asyncio.as_completed

Lib/asyncio/tasks.py

@@ -9,6 +9,7 @@
 import concurrent.futures
 import functools
 import inspect
+import itertools
 import warnings
 import weakref
 
@@ -405,7 +406,7 @@ def _on_completion(f):
 
 
 # This is *not* a @coroutine!  It is just an iterator (yielding Futures).
-def as_completed(fs, *, loop=None, timeout=None):
+def as_completed(fs, *, loop=None, timeout=None, limit=None):
     """Return an iterator whose values are coroutines.
 
     When waiting for the yielded coroutines you'll get the results (or
@@ -421,12 +422,25 @@ def as_completed(fs, *, loop=None, timeout=None):
     If a timeout is specified, the 'yield from' will raise
     TimeoutError when the timeout occurs before all Futures are done.
 
+    If a limit is specified, and the supplied Futures (or coroutines)
+    have not already been scheduled for execution, they will be scheduled
+    in a limited way, ensuring that the number that are executed
+    concurrently will not go above the limit.
+
     Note: The futures 'f' are not necessarily members of fs.
     """
-    if futures.isfuture(fs) or coroutines.iscoroutine(fs):
+    if futures.isfuture(fs):
         raise TypeError("expect a list of futures, not %s" % type(fs).__name__)
     loop = loop if loop is not None else events.get_event_loop()
-    todo = {ensure_future(f, loop=loop) for f in set(fs)}
+    if limit is not None:
+        fiter = iter(fs)
+        initial_futures = set(itertools.islice(fiter, 0, limit))
+    else:
+        fiter = None
+        initial_futures = set(fs)
+    todo = {ensure_future(f, loop=loop) for f in initial_futures}
+    received = len(todo)
+    yielded = 0
     from .queues import Queue  # Import here to avoid circular import problem.
     done = Queue(loop=loop)
     timeout_handle = None
@@ -438,10 +452,19 @@ def _on_timeout():
         todo.clear()  # Can't do todo.remove(f) in the loop.
 
     def _on_completion(f):
+        nonlocal received, fiter
         if not todo:
             return  # _on_timeout() was here first.
         todo.remove(f)
         done.put_nowait(f)
+        if fiter is not None:
+            try:
+                newf = ensure_future(next(fiter), loop=loop)
+                received += 1
+                todo.add(newf)
+                newf.add_done_callback(_on_completion)
+            except StopIteration:
+                fiter = None  # We've finished all the work we were given
         if not todo and timeout_handle is not None:
             timeout_handle.cancel()
 
@@ -457,8 +480,9 @@ def _wait_for_one():
         f.add_done_callback(_on_completion)
     if todo and timeout is not None:
         timeout_handle = loop.call_later(timeout, _on_timeout)
-    for _ in range(len(todo)):
+    while received > yielded:
         yield _wait_for_one()
+        yielded += 1
 
 
 @coroutine

Lib/test/test_asyncio/test_tasks.py

@@ -1075,6 +1075,82 @@ def foo():
         res = loop.run_until_complete(self.new_task(loop, foo()))
         self.assertAlmostEqual(0.15, loop.time())
 
+    def test_as_completed_over_iterable(self):
+
+        def gen():
+            for i in range(3):
+                yield 0.1
+
+        loop = self.new_test_loop(gen)
+
+        fs_it = (asyncio.sleep(i/10.0, i, loop=loop) for i in range(0, 3))
+
+        @asyncio.coroutine
+        def foo():
+            values = []
+            for r in asyncio.as_completed(fs_it, loop=loop):
+                v = yield from r
+                values.append(v)
+            return values
+
+        res = list(sorted(loop.run_until_complete(self.new_task(loop, foo()))))
+        self.assertEqual([0, 1, 2], res)
+
+    def test_as_completed_with_limit(self):
+        """
+        Use as_completed to get a number of results
+        concurrently, but supply the limit parameter
+        to ensure the number of concurrent
+        executions is kept to a given value.
+        """
+        num_futures = 10
+        limit = 3
+
+        def gen():
+            yield 0
+            for i in range(num_futures):
+                yield 1
+
+        loop = self.new_test_loop(gen)
+
+        concurrent = 0
+        max_concurrent = 0
+
+        @asyncio.coroutine
+        def sleeper(dt):
+            nonlocal concurrent, max_concurrent
+            concurrent += 1
+            if concurrent > max_concurrent:
+                max_concurrent = concurrent
+            yield from asyncio.sleep(dt * 0.1, loop=loop)
+            concurrent -= 1
+            return str(dt)
+
+        times = list(range(1, num_futures + 1))
+        # We supply the coroutines in the opposite
+        # order from the order in which we expect
+        # the results to come, to confirm things
+        # happen concurrently.
+        fs = (sleeper(t) for t in reversed(times))
+
+        @asyncio.coroutine
+        def foo():
+            values = []
+            for f in asyncio.as_completed(fs, loop=loop, limit=limit):
+                values.append((yield from f))
+            return values
+
+        res = loop.run_until_complete(self.new_task(loop, foo()))
+
+        # The loop ran to the end, and we got all the
+        # values out.  They are in a weird order because
+        # we supplied them slow-first, so we had to wait
+        # for the slowest ones before we hit the faster
+        # ones.
+        self.assertListEqual(sorted(res), sorted([str(t) for t in times]))
+        # We never started more than limit concurrent executions
+        self.assertEqual(max_concurrent, limit)
+
     def test_as_completed_with_timeout(self):
 
         def gen():
@@ -1676,10 +1752,6 @@ def test_as_completed_invalid_args(self):
         # as_completed() expects a list of futures, not a future instance
         self.assertRaises(TypeError, self.loop.run_until_complete,
             asyncio.as_completed(fut, loop=self.loop))
-        coro = coroutine_function()
-        self.assertRaises(TypeError, self.loop.run_until_complete,
-            asyncio.as_completed(coro, loop=self.loop))
-        coro.close()
 
     def test_wait_invalid_args(self):
         fut = self.new_future(self.loop)