Shield alloc & free from cancellation (#163)
A handler's RPC code may be cancelled due to a request timeout or a client closing the connection. Before this PR: - If `.cancel()` happens while waiting for `hivemind.utils.enter_asynchronously()`, the lock will never be released. - If `.cancel()` happens while doing that before freeing memory, the memory will never be freed. This PR fixes it by deferring the cancellation with [asyncio.shield()](https://docs.python.org/3/library/asyncio-task.html#asyncio.shield). Now, the cancellation will happen only when all locks are released and alloc/free has completed.pull/165/head
parent
d6992fca63
commit
9997ada3bb
@ -0,0 +1,21 @@
|
|||||||
|
import asyncio
|
||||||
|
|
||||||
|
|
||||||
|
async def shield_and_wait(task):
|
||||||
|
"""
|
||||||
|
Works like asyncio.shield(), but waits for the task to finish before raising CancelledError to the caller.
|
||||||
|
"""
|
||||||
|
|
||||||
|
if not isinstance(task, asyncio.Task):
|
||||||
|
task = asyncio.create_task(task)
|
||||||
|
|
||||||
|
cancel_exc = None
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
result = await asyncio.shield(task)
|
||||||
|
break
|
||||||
|
except asyncio.CancelledError as e:
|
||||||
|
cancel_exc = e
|
||||||
|
if cancel_exc is not None:
|
||||||
|
raise cancel_exc
|
||||||
|
return result
|
Loading…
Reference in New Issue