Coverage for Lib/asyncio/staggered.py: 95%

1"""Support for running coroutines in parallel with staggered start times.""" 

2 

3__all__ = 'staggered_race', 

4 

5import contextlib 

6 

7from . import events 

8from . import exceptions as exceptions_mod 

9from . import locks 

10from . import tasks 

11from . import futures 

12 

13 

14async def staggered_race(coro_fns, delay, *, loop=None): 

15 """Run coroutines with staggered start times and take the first to finish. 

16 

17 This method takes an iterable of coroutine functions. The first one is 

18 started immediately. From then on, whenever the immediately preceding one 

19 fails (raises an exception), or when *delay* seconds has passed, the next 

20 coroutine is started. This continues until one of the coroutines complete 

21 successfully, in which case all others are cancelled, or until all 

22 coroutines fail. 

23 

24 The coroutines provided should be well-behaved in the following way: 

25 

26 * They should only ``return`` if completed successfully. 

27 

28 * They should always raise an exception if they did not complete 

29 successfully. In particular, if they handle cancellation, they should 

30 probably reraise, like this:: 

31 

32 try: 

33 # do work 

34 except asyncio.CancelledError: 

35 # undo partially completed work 

36 raise 

37 

38 Args: 

39 coro_fns: an iterable of coroutine functions, i.e. callables that 

40 return a coroutine object when called. Use ``functools.partial`` or 

41 lambdas to pass arguments. 

42 

43 delay: amount of time, in seconds, between starting coroutines. If 

44 ``None``, the coroutines will run sequentially. 

45 

46 loop: the event loop to use. 

47 

48 Returns: 

49 tuple *(winner_result, winner_index, exceptions)* where 

50 

51 - *winner_result*: the result of the winning coroutine, or ``None`` 

52 if no coroutines won. 

53 

54 - *winner_index*: the index of the winning coroutine in 

55 ``coro_fns``, or ``None`` if no coroutines won. If the winning 

56 coroutine may return None on success, *winner_index* can be used 

57 to definitively determine whether any coroutine won. 

58 

59 - *exceptions*: list of exceptions returned by the coroutines. 

60 ``len(exceptions)`` is equal to the number of coroutines actually 

61 started, and the order is the same as in ``coro_fns``. The winning 

62 coroutine's entry is ``None``. 

63 

64 """ 

65 loop = loop or events.get_running_loop() 

66 parent_task = tasks.current_task(loop) 

67 enum_coro_fns = enumerate(coro_fns) 

68 winner_result = None 

69 winner_index = None 

70 unhandled_exceptions = [] 

71 exceptions = [] 

72 running_tasks = set() 

73 on_completed_fut = None 

74 

75 def task_done(task): 

76 running_tasks.discard(task) 

77 futures.future_discard_from_awaited_by(task, parent_task) 

78 if ( 

79 on_completed_fut is not None 

80 and not on_completed_fut.done() 

81 and not running_tasks 

82 ): 

83 on_completed_fut.set_result(None) 

84 

85 if task.cancelled(): 

86 return 

87 

88 exc = task.exception() 

89 if exc is None: 89 ↛ 91line 89 didn't jump to line 91 because the condition on line 89 was always true

90 return 

91 unhandled_exceptions.append(exc) 

92 

93 async def run_one_coro(previous_failed) -> None: 

94 # Wait for the previous task to finish, or for delay seconds 

95 if previous_failed is not None: 

96 with contextlib.suppress(exceptions_mod.TimeoutError): 

97 # Use asyncio.wait_for() instead of asyncio.wait() here, so 

98 # that if we get cancelled at this point, Event.wait() is also 

99 # cancelled, otherwise there will be a "Task destroyed but it is 

100 # pending" later. 

101 await tasks.wait_for(previous_failed.wait(), delay) 

102 # Get the next coroutine to run 

103 try: 

104 this_index, coro_fn = next(enum_coro_fns) 

105 except StopIteration: 

106 return 

107 # Start task that will run the next coroutine 

108 this_failed = locks.Event() 

109 next_task = loop.create_task( 

110 run_one_coro(this_failed), 

111 eager_start=False, 

112 ) 

113 futures.future_add_to_awaited_by(next_task, parent_task) 

114 running_tasks.add(next_task) 

115 next_task.add_done_callback(task_done) 

116 # Prepare place to put this coroutine's exceptions if not won 

117 exceptions.append(None) 

118 assert len(exceptions) == this_index + 1 

119 

120 try: 

121 result = await coro_fn() 

122 except (SystemExit, KeyboardInterrupt): 

123 raise 

124 except BaseException as e: 

125 exceptions[this_index] = e 

126 this_failed.set() # Kickstart the next coroutine 

127 else: 

128 # Store winner's results 

129 nonlocal winner_index, winner_result 

130 assert winner_index is None 

131 winner_index = this_index 

132 winner_result = result 

133 # Cancel all other tasks. We take care to not cancel the current 

134 # task as well. If we do so, then since there is no `await` after 

135 # here and CancelledError are usually thrown at one, we will 

136 # encounter a curious corner case where the current task will end 

137 # up as done() == True, cancelled() == False, exception() == 

138 # asyncio.CancelledError. This behavior is specified in 

139 # https://bugs.python.org/issue30048 

140 current_task = tasks.current_task(loop) 

141 for t in running_tasks: 

142 if t is not current_task: 

143 t.cancel() 

144 

145 propagate_cancellation_error = None 

146 try: 

147 first_task = loop.create_task(run_one_coro(None), eager_start=False) 

148 futures.future_add_to_awaited_by(first_task, parent_task) 

149 running_tasks.add(first_task) 

150 first_task.add_done_callback(task_done) 

151 # first_task has been appended to running_tasks before the event loop starts running it. 

152 propagate_cancellation_error = None 

153 # Make sure no tasks are left running if we leave this function 

154 while running_tasks: 

155 on_completed_fut = loop.create_future() 

156 try: 

157 await on_completed_fut 

158 except exceptions_mod.CancelledError as ex: 

159 propagate_cancellation_error = ex 

160 for task in running_tasks: 

161 task.cancel(*ex.args) 

162 on_completed_fut = None 

163 if __debug__ and unhandled_exceptions: 163 ↛ 166line 163 didn't jump to line 166 because the condition on line 163 was never true

164 # If run_one_coro raises an unhandled exception, it's probably a 

165 # programming error, and I want to see it. 

166 raise ExceptionGroup("staggered race failed", unhandled_exceptions) 

167 if propagate_cancellation_error is not None: 

168 raise propagate_cancellation_error 

169 return winner_result, winner_index, exceptions 

170 finally: 

171 del exceptions, propagate_cancellation_error, unhandled_exceptions, parent_task