Skip to content

Commit

Permalink
Do renaming, implement lowlevel Outcome sending
Browse files Browse the repository at this point in the history
As was listed in the many todos, this changes the `.start_soon()` impl
to instead (manually) `.send()` into the user defined
`@task_scope_manager` an `Outcome` from the spawned task. In this case
the task manager wraps that in a user defined (and renamed)
`TaskOutcome` and delivers that + a containing `trio.CancelScope` to the
`.start_soon()` caller. Here the user defined `TaskOutcome` defines
a `.wait_for_result()` method that can be used to await the task's exit
and handle it's underlying returned value or raised error; the
implementation could be different and subject to the user's own whims.

Note that by default, if this was added to `trio`'s core, the
`@task_scope_manager` would simply be implemented as either a `None`
yielding single-yield-generator but more likely just entirely ignored
by the runtime (as in no manual task outcome collecting, generator
calling and sending is done at all) by default if the user does not provide
the `task_scope_manager` to the nursery at open time.
  • Loading branch information
goodboy committed May 19, 2023
1 parent 3f1bf47 commit bc2bd2d
Showing 1 changed file with 91 additions and 61 deletions.
152 changes: 91 additions & 61 deletions tractor/trionics/_supervisor.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,17 @@
Nursery,
)

class MaybeOutcome(Struct):

_ready: Event = trio.Event()
_outcome: Outcome | None = None
_result: Any | None = None
class TaskOutcome(Struct):
'''
The outcome of a scheduled ``trio`` task which includes an interface
for synchronizing to the completion of the task's runtime and access
to the eventual boxed result/value or raised exception.
'''
_exited: Event = trio.Event() # as per `trio.Runner.task_exited()`
_outcome: Outcome | None = None # as per `outcome.Outcome`
_result: Any | None = None # the eventual maybe-returned-value

@property
def result(self) -> Any:
Expand All @@ -55,27 +61,35 @@ def result(self) -> Any:
raise RuntimeError(
# f'Task {task.name} is not complete.\n'
f'Outcome is not complete.\n'
'wait on `await MaybeOutcome.unwrap()` first!'
'wait on `await TaskOutcome.wait_for_result()` first!'
)
return self._result

def _set_outcome(
self,
outcome: Outcome,
):
'''
Set the ``Outcome`` for this task.
This method should only ever be called by the task's supervising
nursery implemenation.
'''
self._outcome = outcome
self._result = outcome.unwrap()
self._ready.set()

# TODO: maybe a better name like,
# - .wait_and_unwrap()
# - .wait_unwrap()
# - .aunwrap() ?
async def unwrap(self) -> Any:
if self._ready.is_set():
self._exited.set()

async def wait_for_result(self) -> Any:
'''
Unwind the underlying task's ``Outcome`` by async waiting for
the task to first complete and then unwrap it's result-value.
'''
if self._exited.is_set():
return self._result

await self._ready.wait()
await self._exited.wait()

out = self._outcome
if out is None:
Expand All @@ -84,13 +98,6 @@ async def unwrap(self) -> Any:
return self.result


class TaskHandle(Struct):
task: Task
cs: CancelScope
exited: Event | None = None
_outcome: Outcome | None = None


class ScopePerTaskNursery(Struct):
_n: Nursery
_scopes: dict[
Expand Down Expand Up @@ -122,73 +129,97 @@ async def start_soon(
cs = CancelScope()
new_task: Task | None = None
to_return: tuple[Any] | None = None
maybe_outcome = MaybeOutcome()

sm = self.scope_manager
if sm is None:
mngr = nullcontext([cs])
else:
mngr = sm(
nursery=n,
scope=cs,
maybe_outcome=maybe_outcome,
)
# NOTE: what do we enforce as a signature for the
# `@task_scope_manager` here?
mngr = sm(nursery=n, scope=cs)

async def _start_wrapped_in_scope(
task_status: TaskStatus[
tuple[CancelScope, Task]
] = trio.TASK_STATUS_IGNORED,

) -> None:
nonlocal maybe_outcome
nonlocal to_return

# TODO: this was working before?!
# nonlocal to_return

with cs:

task = trio.lowlevel.current_task()
self._scopes[cs] = task

# TODO: instead we should probably just use
# `Outcome.send(mngr)` here no and inside a custom
# decorator `@trio.cancel_scope_manager` enforce
# that it's a single yield generator?
with mngr as to_return:

# TODO: relay through whatever the
# started task passes back via `.started()` ?
# seems like that won't work with also returning
# a "task handle"?
task_status.started()
# execute up to the first yield
try:
to_return: tuple[Any] = next(mngr)
except StopIteration:
raise RuntimeError("task manager didn't yield") from None

# TODO: how do we support `.start()` style?
# - relay through whatever the
# started task passes back via `.started()` ?
# seems like that won't work with also returning
# a "task handle"?
# - we were previously binding-out this `to_return` to
# the parent's lexical scope, why isn't that working
# now?
task_status.started(to_return)

# invoke underlying func now that cs is entered.
outcome = await acapture(async_fn, *args)

# execute from the 1st yield to return and expect
# generator-mngr `@task_scope_manager` thinger to
# terminate!
try:
mngr.send(outcome)

# NOTE: this will instead send the underlying
# `.value`? Not sure if that's better or not?
# I would presume it's better to have a handle to
# the `Outcome` entirely? This method sends *into*
# the mngr this `Outcome.value`; seems like kinda
# weird semantics for our purposes?
# outcome.send(mngr)

except StopIteration:
return
else:
raise RuntimeError(f"{mngr} didn't stop!")

to_return = await n.start(_start_wrapped_in_scope)
assert to_return is not None

# invoke underlying func now that cs is entered.
outcome = await acapture(async_fn, *args)
# TODO: use the fancy type-check-time type signature stuff from
# mypy i guess..to like, relay the type of whatever the
# generator yielded through? betcha that'll be un-grokable XD
return to_return

# TODO: instead, mngr.send(outcome) so that we don't
# tie this `.start_soon()` impl to the
# `MaybeOutcome` type? Use `Outcome.send(mngr)`
# right?
maybe_outcome._set_outcome(outcome)

await n.start(_start_wrapped_in_scope)
assert to_return is not None

# TODO: better way to concat the values delivered by the user
# provided `.scope_manager` and the outcome?
return tuple([maybe_outcome] + to_return)
# TODO: you could wrap your output task handle in this?
# class TaskHandle(Struct):
# task: Task
# cs: CancelScope
# outcome: TaskOutcome


# TODO: maybe just make this a generator with a single yield that also
# delivers a value (of some std type) from the yield expression?
# @trio.cancel_scope_manager
@cm
# @trio.task_scope_manager
def add_task_handle_and_crash_handling(
nursery: Nursery,
scope: CancelScope,
maybe_outcome: MaybeOutcome,
# maybe_outcome: TaskOutcome,

) -> Generator[None, list[Any]]:

cs: CancelScope = CancelScope()
task_outcome = TaskOutcome()

# if you need it you can ask trio for the task obj
task: Task = trio.lowlevel.current_task()
Expand All @@ -197,12 +228,11 @@ def add_task_handle_and_crash_handling(
try:
# yields back when task is terminated, cancelled, returns?
with cs:
# the yielded values here are what are returned to the
# nursery's `.start_soon()` caller

# TODO: actually make this work so that `MaybeOutcome` isn't
# tied to the impl of `.start_soon()` on our custom nursery!
task_outcome: Outcome = yield [cs]
# the yielded value(s) here are what are returned to the
# nursery's `.start_soon()` caller B)
lowlevel_outcome: Outcome = yield (task_outcome, cs)
task_outcome._set_outcome(lowlevel_outcome)

except Exception as err:
# Adds "crash handling" from `pdbp` by entering
Expand Down Expand Up @@ -247,7 +277,7 @@ async def main():

val: str = 'yoyoyo'
val_outcome, cs = await sn.start_soon(sleep_then_return_val, val)
res = await val_outcome.unwrap()
res = await val_outcome.wait_for_result()
assert res == val
print(f'GOT EXPECTED TASK VALUE: {res}')

Expand Down

0 comments on commit bc2bd2d

Please sign in to comment.