Close all portal created async gens on shutdown

remote_task_cancelling
Tyler Goodlet 2018-12-15 02:20:55 -05:00
parent db85e13657
commit eb6e82f577
1 changed files with 18 additions and 10 deletions

View File

@ -65,6 +65,7 @@ class Portal:
self._expect_result: Optional[
Tuple[str, Any, str, Dict[str, Any]]
] = None
self._agens: Set(AsyncGenerator) = set()
async def aclose(self) -> None:
log.debug(f"Closing {self}")
@ -142,14 +143,16 @@ class Portal:
except GeneratorExit:
# for now this msg cancels an ongoing remote task
await self.channel.send({'cancel': True, 'cid': cid})
log.debug(
log.warn(
f"Cancelling async gen call {cid} to "
f"{self.channel.uid}")
raise
# TODO: use AsyncExitStack to aclose() all agens
# on teardown
return yield_from_q()
agen = yield_from_q()
self._agens.add(agen)
return agen
elif resptype == 'return':
msg = await q.get()
@ -269,7 +272,12 @@ async def open_portal(
nursery.start_soon(actor._process_messages, channel)
portal = Portal(channel)
try:
yield portal
finally:
# tear down all async generators
for agen in portal._agens:
await agen.aclose()
# cancel remote channel-msg loop
if channel.connected():