forked from goodboy/tractor
Change to `gather_contexts()`, use event for graceful exit
The api we've made here is actually closer to `asyncio.gather()` but with opening async context managers instead of funcs. Use another event to allow for graceful teardown of children on non-cancellation exits and add a doc string.graceful_gather
parent
ebf080b8a2
commit
d0f5c7a5e2
|
@ -3,7 +3,7 @@ import itertools
|
||||||
import trio
|
import trio
|
||||||
import tractor
|
import tractor
|
||||||
from tractor import open_actor_cluster
|
from tractor import open_actor_cluster
|
||||||
from tractor.trionics import async_enter_all
|
from tractor.trionics import gather_contexts
|
||||||
|
|
||||||
from conftest import tractor_test
|
from conftest import tractor_test
|
||||||
|
|
||||||
|
@ -25,10 +25,10 @@ async def worker(ctx: tractor.Context) -> None:
|
||||||
async def test_streaming_to_actor_cluster() -> None:
|
async def test_streaming_to_actor_cluster() -> None:
|
||||||
async with (
|
async with (
|
||||||
open_actor_cluster(modules=[__name__]) as portals,
|
open_actor_cluster(modules=[__name__]) as portals,
|
||||||
async_enter_all(
|
gather_contexts(
|
||||||
mngrs=[p.open_context(worker) for p in portals.values()],
|
mngrs=[p.open_context(worker) for p in portals.values()],
|
||||||
) as contexts,
|
) as contexts,
|
||||||
async_enter_all(
|
gather_contexts(
|
||||||
mngrs=[ctx[0].open_stream() for ctx in contexts],
|
mngrs=[ctx[0].open_stream() for ctx in contexts],
|
||||||
) as streams,
|
) as streams,
|
||||||
):
|
):
|
||||||
|
|
|
@ -2,12 +2,12 @@
|
||||||
Sugary patterns for trio + tractor designs.
|
Sugary patterns for trio + tractor designs.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
from ._mngrs import async_enter_all
|
from ._mngrs import gather_contexts
|
||||||
from ._broadcast import broadcast_receiver, BroadcastReceiver, Lagged
|
from ._broadcast import broadcast_receiver, BroadcastReceiver, Lagged
|
||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
'async_enter_all',
|
'gather_contexts',
|
||||||
'broadcast_receiver',
|
'broadcast_receiver',
|
||||||
'BroadcastReceiver',
|
'BroadcastReceiver',
|
||||||
'Lagged',
|
'Lagged',
|
||||||
|
|
|
@ -14,11 +14,15 @@ T = TypeVar("T")
|
||||||
|
|
||||||
|
|
||||||
async def _enter_and_wait(
|
async def _enter_and_wait(
|
||||||
|
|
||||||
mngr: AsyncContextManager[T],
|
mngr: AsyncContextManager[T],
|
||||||
unwrapped: dict[int, T],
|
unwrapped: dict[int, T],
|
||||||
all_entered: trio.Event,
|
all_entered: trio.Event,
|
||||||
|
parent_exit: trio.Event,
|
||||||
|
|
||||||
) -> None:
|
) -> None:
|
||||||
'''Open the async context manager deliver it's value
|
'''
|
||||||
|
Open the async context manager deliver it's value
|
||||||
to this task's spawner and sleep until cancelled.
|
to this task's spawner and sleep until cancelled.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
|
@ -28,16 +32,31 @@ async def _enter_and_wait(
|
||||||
if all(unwrapped.values()):
|
if all(unwrapped.values()):
|
||||||
all_entered.set()
|
all_entered.set()
|
||||||
|
|
||||||
await trio.sleep_forever()
|
await parent_exit.wait()
|
||||||
|
|
||||||
|
|
||||||
@acm
|
@acm
|
||||||
async def async_enter_all(
|
async def gather_contexts(
|
||||||
|
|
||||||
mngrs: Sequence[AsyncContextManager[T]],
|
mngrs: Sequence[AsyncContextManager[T]],
|
||||||
|
|
||||||
) -> AsyncGenerator[tuple[T, ...], None]:
|
) -> AsyncGenerator[tuple[T, ...], None]:
|
||||||
|
'''
|
||||||
|
Concurrently enter a sequence of async context managers, each in
|
||||||
|
a separate ``trio`` task and deliver the unwrapped values in the
|
||||||
|
same order once all managers have entered. On exit all contexts are
|
||||||
|
subsequently and concurrently exited.
|
||||||
|
|
||||||
|
This function is somewhat similar to common usage of
|
||||||
|
``contextlib.AsyncExitStack.enter_async_context()`` (in a loop) in
|
||||||
|
combo with ``asyncio.gather()`` except the managers are concurrently
|
||||||
|
entered and exited cancellation just works.
|
||||||
|
|
||||||
|
'''
|
||||||
unwrapped = {}.fromkeys(id(mngr) for mngr in mngrs)
|
unwrapped = {}.fromkeys(id(mngr) for mngr in mngrs)
|
||||||
|
|
||||||
all_entered = trio.Event()
|
all_entered = trio.Event()
|
||||||
|
parent_exit = trio.Event()
|
||||||
|
|
||||||
async with trio.open_nursery() as n:
|
async with trio.open_nursery() as n:
|
||||||
for mngr in mngrs:
|
for mngr in mngrs:
|
||||||
|
@ -46,6 +65,7 @@ async def async_enter_all(
|
||||||
mngr,
|
mngr,
|
||||||
unwrapped,
|
unwrapped,
|
||||||
all_entered,
|
all_entered,
|
||||||
|
parent_exit,
|
||||||
)
|
)
|
||||||
|
|
||||||
# deliver control once all managers have started up
|
# deliver control once all managers have started up
|
||||||
|
@ -53,4 +73,6 @@ async def async_enter_all(
|
||||||
|
|
||||||
yield tuple(unwrapped.values())
|
yield tuple(unwrapped.values())
|
||||||
|
|
||||||
n.cancel_scope.cancel()
|
# we don't need a try/finally since cancellation will be triggered
|
||||||
|
# by the surrounding nursery on error.
|
||||||
|
parent_exit.set()
|
||||||
|
|
Loading…
Reference in New Issue