Skip to content

Commit 36c2c04

Browse files
authored
bpo-32355: Optimize asyncio.gather() (#4913)
1 parent a9d7e55 commit 36c2c04

File tree

3 files changed

+68
-47
lines changed

3 files changed

+68
-47
lines changed

Lib/asyncio/base_events.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -139,11 +139,12 @@ def _ipaddr_info(host, port, family, type, proto):
139139

140140

141141
def _run_until_complete_cb(fut):
142-
exc = fut._exception
143-
if isinstance(exc, BaseException) and not isinstance(exc, Exception):
144-
# Issue #22429: run_forever() already finished, no need to
145-
# stop it.
146-
return
142+
if not fut.cancelled():
143+
exc = fut.exception()
144+
if isinstance(exc, BaseException) and not isinstance(exc, Exception):
145+
# Issue #22429: run_forever() already finished, no need to
146+
# stop it.
147+
return
147148
fut._loop.stop()
148149

149150

Lib/asyncio/tasks.py

Lines changed: 61 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -575,8 +575,7 @@ def cancel(self):
575575

576576

577577
def gather(*coros_or_futures, loop=None, return_exceptions=False):
578-
"""Return a future aggregating results from the given coroutines
579-
or futures.
578+
"""Return a future aggregating results from the given coroutines/futures.
580579
581580
Coroutines will be wrapped in a future and scheduled in the event
582581
loop. They will not necessarily be scheduled in the same order as
@@ -605,56 +604,76 @@ def gather(*coros_or_futures, loop=None, return_exceptions=False):
605604
outer.set_result([])
606605
return outer
607606

608-
arg_to_fut = {}
609-
for arg in set(coros_or_futures):
610-
if not futures.isfuture(arg):
611-
fut = ensure_future(arg, loop=loop)
612-
if loop is None:
613-
loop = fut._loop
614-
# The caller cannot control this future, the "destroy pending task"
615-
# warning should not be emitted.
616-
fut._log_destroy_pending = False
617-
else:
618-
fut = arg
619-
if loop is None:
620-
loop = fut._loop
621-
elif fut._loop is not loop:
622-
raise ValueError("futures are tied to different event loops")
623-
arg_to_fut[arg] = fut
624-
625-
children = [arg_to_fut[arg] for arg in coros_or_futures]
626-
nchildren = len(children)
627-
outer = _GatheringFuture(children, loop=loop)
628-
nfinished = 0
629-
results = [None] * nchildren
630-
631-
def _done_callback(i, fut):
607+
def _done_callback(fut):
632608
nonlocal nfinished
609+
nfinished += 1
610+
633611
if outer.done():
634612
if not fut.cancelled():
635613
# Mark exception retrieved.
636614
fut.exception()
637615
return
638616

639-
if fut.cancelled():
640-
res = futures.CancelledError()
641-
if not return_exceptions:
642-
outer.set_exception(res)
643-
return
644-
elif fut._exception is not None:
645-
res = fut.exception() # Mark exception retrieved.
646-
if not return_exceptions:
647-
outer.set_exception(res)
617+
if not return_exceptions:
618+
if fut.cancelled():
619+
# Check if 'fut' is cancelled first, as
620+
# 'fut.exception()' will *raise* a CancelledError
621+
# instead of returning it.
622+
exc = futures.CancelledError()
623+
outer.set_exception(exc)
648624
return
649-
else:
650-
res = fut._result
651-
results[i] = res
652-
nfinished += 1
653-
if nfinished == nchildren:
625+
else:
626+
exc = fut.exception()
627+
if exc is not None:
628+
outer.set_exception(exc)
629+
return
630+
631+
if nfinished == nfuts:
632+
# All futures are done; create a list of results
633+
# and set it to the 'outer' future.
634+
results = []
635+
636+
for fut in children:
637+
if fut.cancelled():
638+
# Check if 'fut' is cancelled first, as
639+
# 'fut.exception()' will *raise* a CancelledError
640+
# instead of returning it.
641+
res = futures.CancelledError()
642+
else:
643+
res = fut.exception()
644+
if res is None:
645+
res = fut.result()
646+
results.append(res)
647+
654648
outer.set_result(results)
655649

656-
for i, fut in enumerate(children):
657-
fut.add_done_callback(functools.partial(_done_callback, i))
650+
arg_to_fut = {}
651+
children = []
652+
nfuts = 0
653+
nfinished = 0
654+
for arg in coros_or_futures:
655+
if arg not in arg_to_fut:
656+
fut = ensure_future(arg, loop=loop)
657+
if loop is None:
658+
loop = fut._loop
659+
if fut is not arg:
660+
# 'arg' was not a Future, therefore, 'fut' is a new
661+
# Future created specifically for 'arg'. Since the caller
662+
# can't control it, disable the "destroy pending task"
663+
# warning.
664+
fut._log_destroy_pending = False
665+
666+
nfuts += 1
667+
arg_to_fut[arg] = fut
668+
fut.add_done_callback(_done_callback)
669+
670+
else:
671+
# There's a duplicate Future object in coros_or_futures.
672+
fut = arg_to_fut[arg]
673+
674+
children.append(fut)
675+
676+
outer = _GatheringFuture(children, loop=loop)
658677
return outer
659678

660679

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Optimize asyncio.gather(); now up to 15% faster.

0 commit comments

Comments
 (0)