gh-124309: Modernize the staggered_race implementation to support eager task factories

[ZeroIntensity]

• Issue: Happy Eyeballs crashes with Eager Task Factory #124309

[graingert]

You're using private methods of TaskGroup and starting tasks on the loop rather than the TaskGroup

[ZeroIntensity]

I think I'm just going to refactor this to not use TaskGroup. It's causing more problems than solutions.

[graingert]

I think it's worth persevering with TaskGroup, you just need to write it without using add_done_callback or private attributes

[ZeroIntensity]

I'll try it, but I'm worried that it isn't possible when considering an eager task factory. The previous implementation used a variation of a task group (a list containing tasks, since it predated asyncio.TaskGroup) but broke due to the reliance of tasks not starting before the event loop is called.

While we're here, staggered_race is undocumented -- might that be something worth addressing in this PR?

[graingert]

A demo of what I mean wrt TaskGroup:

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

__all__ = 'staggered_race',

from . import locks
from . import tasks
from . import taskgroups


async def staggered_race(coro_fns, delay, *, loop=None):
    """Run coroutines with staggered start times and take the first to finish.

    This method takes an iterable of coroutine functions. The first one is
    started immediately. From then on, whenever the immediately preceding one
    fails (raises an exception), or when *delay* seconds has passed, the next
    coroutine is started. This continues until one of the coroutines complete
    successfully, in which case all others are cancelled, or until all
    coroutines fail.

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

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

    * They should always raise an exception if they did not complete
      successfully. In particular, if they handle cancellation, they should
      probably reraise, like this::

        try:
            # do work
        except asyncio.CancelledError:
            # undo partially completed work
            raise

    Args:
        coro_fns: an iterable of coroutine functions, i.e. callables that
            return a coroutine object when called. Use ``functools.partial`` or
            lambdas to pass arguments.

        delay: amount of time, in seconds, between starting coroutines. If
            ``None``, the coroutines will run sequentially.

        loop: the event loop to use.

    Returns:
        tuple *(winner_result, winner_index, exceptions)* where

        - *winner_result*: the result of the winning coroutine, or ``None``
          if no coroutines won.

        - *winner_index*: the index of the winning coroutine in
          ``coro_fns``, or ``None`` if no coroutines won. If the winning
          coroutine may return None on success, *winner_index* can be used
          to definitively determine whether any coroutine won.

        - *exceptions*: list of exceptions returned by the coroutines.
          ``len(exceptions)`` is equal to the number of coroutines actually
          started, and the order is the same as in ``coro_fns``. The winning
          coroutine's entry is ``None``.

    """
    # TODO: when we have aiter() and anext(), allow async iterables in coro_fns.
    winner_result = None
    winner_index = None
    exceptions = []

    class _Done(Exception):
        pass

    async def run_one_coro(this_index, coro_fn, this_failed):
        try:
            result = await coro_fn()
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as e:
            exceptions[this_index] = e
            this_failed.set()  # Kickstart the next coroutine
        else:
            # Store winner's results
            nonlocal winner_index, winner_result
            # There could be more than one winner
            winner_index = this_index
            winner_result = result
            raise _Done

    try:
        async with taskgroups.TaskGroup() as tg:
            for this_index, coro_fn in enumerate(coro_fns):
                this_failed = locks.Event()
                exceptions.append(None)
                tg.create_task(run_one_coro(this_index, coro_fn, this_failed))
                try:
                    await tasks.wait_for(this_failed.wait(), delay)
                except TimeoutError:
                    pass
    except* _Done:
        pass

    return winner_result, winner_index, exceptions

[ZeroIntensity]

Thank you for the review, Thomas!

[willingc]

Thanks @ZeroIntensity.

[miss-islington-app]

Thanks @ZeroIntensity for the PR, and @kumaraditya303 for merging it 🌮🎉.. I'm working now to backport this PR to: 3.12, 3.13.
🐍🍒⛏🤖

[miss-islington-app]

Sorry, @ZeroIntensity and @kumaraditya303, I could not cleanly backport this to 3.12 due to a conflict.
Please backport using cherry_picker on command line.

cherry_picker de929f353c413459834a2a37b2d9b0240673d874 3.12

[bedevere-app]

GH-124573 is a backport of this pull request to the 3.13 branch.

[bedevere-app]

GH-124574 is a backport of this pull request to the 3.12 branch.