Only pop contexts on decorated entrypoints

stricter_context_starting
Tyler Goodlet 2021-12-06 13:47:49 -05:00
parent 58805a0430
commit 1f8e1cccbb
1 changed files with 10 additions and 6 deletions

View File

@ -180,6 +180,14 @@ async def _invoke(
raise raise
finally:
# XXX: only pop the context tracking if
# a ``@tractor.context`` entrypoint was called
assert chan.uid
ctx = actor._contexts.pop((chan.uid, cid))
if ctx:
log.runtime(f'Context entrypoint for {func} was terminated:\n{ctx}')
assert cs assert cs
if cs.cancelled_caught: if cs.cancelled_caught:
@ -255,15 +263,11 @@ async def _invoke(
task_status.started(err) task_status.started(err)
finally: finally:
assert chan.uid
ctx = actor._contexts.pop((chan.uid, cid))
if ctx:
log.cancel(f'{ctx} was terminated')
# RPC task bookeeping # RPC task bookeeping
try: try:
scope, func, is_complete = actor._rpc_tasks.pop((chan, cid)) scope, func, is_complete = actor._rpc_tasks.pop((chan, cid))
is_complete.set() is_complete.set()
except KeyError: except KeyError:
if is_rpc: if is_rpc:
# If we're cancelled before the task returns then the # If we're cancelled before the task returns then the
@ -696,7 +700,7 @@ class Actor:
if not ctx._stream_opened: if not ctx._stream_opened:
lines.insert( lines.insert(
1, 1,
f'\n*** No stream open on {self.uid[0]} side! ***\n' f'\n*** No stream open on `{self.uid[0]}` side! ***\n'
) )
text = '\n'.join(lines) text = '\n'.join(lines)