forked from goodboy/tractor
Close all portal created async gens on shutdown
parent
db85e13657
commit
eb6e82f577
|
@ -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,13 +272,18 @@ async def open_portal(
|
|||
|
||||
nursery.start_soon(actor._process_messages, channel)
|
||||
portal = Portal(channel)
|
||||
yield portal
|
||||
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():
|
||||
await portal.close()
|
||||
# cancel remote channel-msg loop
|
||||
if channel.connected():
|
||||
await portal.close()
|
||||
|
||||
# cancel background msg loop task
|
||||
nursery.cancel_scope.cancel()
|
||||
if was_connected:
|
||||
await channel.aclose()
|
||||
# cancel background msg loop task
|
||||
nursery.cancel_scope.cancel()
|
||||
if was_connected:
|
||||
await channel.aclose()
|
||||
|
|
Loading…
Reference in New Issue