jc211
/
tractor
forked from goodboy/tractor
1
0
Fork 0
Code Pull Requests Activity

Add support for cancelling remote tasks via a msg

remote_task_cancelling
Tyler Goodlet 2018-12-10 23:12:46 -05:00
parent c0cdb3945a
commit 4dccb44c67
1 changed files with 24 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
tractor/_actor.py
View File

@ -103,7 +103,7 @@ async def _invoke(
# TODO: we should really support a proper # TODO: we should really support a proper
# `StopAsyncIteration` system here for returning a final # `StopAsyncIteration` system here for returning a final
# value if desired # value if desired
await chan.send({'stop': None, 'cid': cid}) await chan.send({'stop': True, 'cid': cid})
else: else:
if treat_as_gen: if treat_as_gen:
await chan.send({'functype': 'asyncgen', 'cid': cid}) await chan.send({'functype': 'asyncgen', 'cid': cid})
@ -117,7 +117,7 @@ async def _invoke(
if not cs.cancelled_caught: if not cs.cancelled_caught:
# task was not cancelled so we can instruct the # task was not cancelled so we can instruct the
# far end async gen to tear down # far end async gen to tear down
await chan.send({'stop': None, 'cid': cid}) await chan.send({'stop': True, 'cid': cid})
else: else:
await chan.send({'functype': 'asyncfunction', 'cid': cid}) await chan.send({'functype': 'asyncfunction', 'cid': cid})
with trio.open_cancel_scope() as cs: with trio.open_cancel_scope() as cs:
@ -141,7 +141,7 @@ async def _invoke(
tasks = actor._rpc_tasks.get(chan, None) tasks = actor._rpc_tasks.get(chan, None)
if tasks: if tasks:
try: try:
tasks.remove((cs, func)) scope, func = tasks.pop(cid)
except ValueError: except ValueError:
# If we're cancelled before the task returns then the # If we're cancelled before the task returns then the
# cancel scope will not have been inserted yet # cancel scope will not have been inserted yet
@ -197,7 +197,7 @@ class Actor:
self._no_more_rpc_tasks.set() self._no_more_rpc_tasks.set()
self._rpc_tasks: Dict[ self._rpc_tasks: Dict[
Channel, Channel,
List[Tuple[trio._core._run.CancelScope, typing.Callable]] Dict[str, Tuple[trio._core._run.CancelScope, typing.Callable]]
] = {} ] = {}
# map {uids -> {callids -> waiter queues}} # map {uids -> {callids -> waiter queues}}
self._actors2calls: Dict[Tuple[str, str], Dict[str, trio.Queue]] = {} self._actors2calls: Dict[Tuple[str, str], Dict[str, trio.Queue]] = {}
@ -344,9 +344,10 @@ class Actor:
if msg is None: # terminate sentinel if msg is None: # terminate sentinel
log.debug( log.debug(
f"Cancelling all tasks for {chan} from {chan.uid}") f"Cancelling all tasks for {chan} from {chan.uid}")
for scope, func in self._rpc_tasks.pop(chan, ()): for cid, (scope, func) in self._rpc_tasks.pop(
chan, {}
).items():
scope.cancel() scope.cancel()
log.debug( log.debug(
f"Msg loop signalled to terminate for" f"Msg loop signalled to terminate for"
f" {chan} from {chan.uid}") f" {chan} from {chan.uid}")
@ -354,6 +355,16 @@ class Actor:
log.debug(f"Received msg {msg} from {chan.uid}") log.debug(f"Received msg {msg} from {chan.uid}")
cid = msg.get('cid') cid = msg.get('cid')
if cid: if cid:
cancel = msg.get('cancel')
if cancel:
# right now this is only implicitly used by
# async generator IPC
scope, func = self._rpc_tasks[chan][cid]
log.debug(
f"Received cancel request for task {cid}"
f" from {chan.uid}")
scope.cancel()
else:
# deliver response to local caller/waiter # deliver response to local caller/waiter
await self._push_result(chan.uid, cid, msg) await self._push_result(chan.uid, cid, msg)
log.debug( log.debug(
@ -403,7 +414,7 @@ class Actor:
log.info(f"RPC func is {func}") log.info(f"RPC func is {func}")
# store cancel scope such that the rpc task can be # store cancel scope such that the rpc task can be
# cancelled gracefully if requested # cancelled gracefully if requested
self._rpc_tasks.setdefault(chan, []).append((cs, func)) self._rpc_tasks.setdefault(chan, {})[cid] = (cs, func)
log.debug( log.debug(
f"Waiting on next msg for {chan} from {chan.uid}") f"Waiting on next msg for {chan} from {chan.uid}")
else: else:
@ -611,9 +622,9 @@ class Actor:
""" """
tasks = self._rpc_tasks tasks = self._rpc_tasks
log.info(f"Cancelling all {len(tasks)} rpc tasks:\n{tasks}") log.info(f"Cancelling all {len(tasks)} rpc tasks:\n{tasks}")
for chan, scopes in tasks.items(): for chan, cids2scopes in tasks.items():
log.debug(f"Cancelling all tasks for {chan.uid}") log.debug(f"Cancelling all tasks for {chan.uid}")
for scope, func in scopes: for cid, (scope, func) in cids2scopes.items():
log.debug(f"Cancelling task for {func}") log.debug(f"Cancelling task for {func}")
scope.cancel() scope.cancel()
if tasks: if tasks: