[Python-checkins] cpython (3.4): asyncio: Change as_completed() to use a Queue, to avoid O(N**2) behavior. Fixes

Mon Mar 17 07:31:36 CET 2014

http://hg.python.org/cpython/rev/b52113fb58a5
changeset: 89707:b52113fb58a5
branch: 3.4
user: Guido van Rossum <guido at python.org>
date: Wed Feb 12 17:58:19 2014 -0800
summary:
 asyncio: Change as_completed() to use a Queue, to avoid O(N**2) behavior. Fixes issue #20566.
files:
 Lib/asyncio/tasks.py | 53 ++++++++++------
 Lib/test/test_asyncio/test_tasks.py | 23 +++++++-
 2 files changed, 55 insertions(+), 21 deletions(-)

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -463,7 +463,11 @@
 
 # This is *not* a @coroutine! It is just an iterator (yielding Futures).
 def as_completed(fs, *, loop=None, timeout=None):
- """Return an iterator whose values, when waited for, are Futures.
+ """Return an iterator whose values are coroutines.
+
+ When waiting for the yielded coroutines you'll get the results (or
+ exceptions!) of the original Futures (or coroutines), in the order
+ in which and as soon as they complete.
 
 This differs from PEP 3148; the proper way to use this is:
 
@@ -471,8 +475,8 @@
 result = yield from f # The 'yield from' may raise.
 # Use result.
 
- Raises TimeoutError if the timeout occurs before all Futures are
- done.
+ If a timeout is specified, the 'yield from' will raise
+ TimeoutError when the timeout occurs before all Futures are done.
 
 Note: The futures 'f' are not necessarily members of fs.
 """
@@ -481,27 +485,36 @@
 loop = loop if loop is not None else events.get_event_loop()
 deadline = None if timeout is None else loop.time() + timeout
 todo = {async(f, loop=loop) for f in set(fs)}
- completed = collections.deque()
+ from .queues import Queue # Import here to avoid circular import problem.
+ done = Queue(loop=loop)
+ timeout_handle = None
+
+ def _on_timeout():
+ for f in todo:
+ f.remove_done_callback(_on_completion)
+ done.put_nowait(None) # Queue a dummy value for _wait_for_one().
+ todo.clear() # Can't do todo.remove(f) in the loop.
+
+ def _on_completion(f):
+ if not todo:
+ return # _on_timeout() was here first.
+ todo.remove(f)
+ done.put_nowait(f)
+ if not todo and timeout_handle is not None:
+ timeout_handle.cancel()
 
 @coroutine
 def _wait_for_one():
- while not completed:
- timeout = None
- if deadline is not None:
- timeout = deadline - loop.time()
- if timeout < 0:
- raise futures.TimeoutError()
- done, pending = yield from _wait(
- todo, timeout, FIRST_COMPLETED, loop)
- # Multiple callers might be waiting for the same events
- # and getting the same outcome. Dedupe by updating todo.
- for f in done:
- if f in todo:
- todo.remove(f)
- completed.append(f)
- f = completed.popleft()
- return f.result() # May raise.
+ f = yield from done.get()
+ if f is None:
+ # Dummy value from _on_timeout().
+ raise futures.TimeoutError
+ return f.result() # May raise f.exception().
 
+ for f in todo:
+ 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)):
 yield _wait_for_one()
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -779,7 +779,6 @@
 yield 0
 yield 0
 yield 0.1
- yield 0.02
 
 loop = test_utils.TestLoop(gen)
 self.addCleanup(loop.close)
@@ -791,6 +790,8 @@
 def foo():
 values = []
 for f in asyncio.as_completed([a, b], timeout=0.12, loop=loop):
+ if values:
+ loop.advance_time(0.02)
 try:
 v = yield from f
 values.append((1, v))
@@ -809,6 +810,26 @@
 loop.advance_time(10)
 loop.run_until_complete(asyncio.wait([a, b], loop=loop))
 
+ def test_as_completed_with_unused_timeout(self):
+
+ def gen():
+ yield
+ yield 0
+ yield 0.01
+
+ loop = test_utils.TestLoop(gen)
+ self.addCleanup(loop.close)
+
+ a = asyncio.sleep(0.01, 'a', loop=loop)
+
+ @asyncio.coroutine
+ def foo():
+ for f in asyncio.as_completed([a], timeout=1, loop=loop):
+ v = yield from f
+ self.assertEqual(v, 'a')
+
+ res = loop.run_until_complete(asyncio.Task(foo(), loop=loop))
+
 def test_as_completed_reverse_wait(self):
 
 def gen():
-- 
Repository URL: http://hg.python.org/cpython
