Fix async task being destroyed before cancelled (#2787)

This commit is contained in:
Nuno Campos 2023-04-12 20:38:38 +01:00 committed by GitHub
parent 819d72614a
commit 1bea9ea4be
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -47,14 +47,19 @@ class AsyncIteratorCallbackHandler(AsyncCallbackHandler):
while not self.queue.empty() or not self.done.is_set():
# Wait for the next token in the queue,
# but stop waiting if the done event is set
done, _ = await asyncio.wait(
done, other = await asyncio.wait(
[
# NOTE: If you add other tasks here, update the code below,
# which assumes each set has exactly one task each
asyncio.ensure_future(self.queue.get()),
asyncio.ensure_future(self.done.wait()),
],
return_when=asyncio.FIRST_COMPLETED,
)
# Cancel the other task
other.pop().cancel()
# Extract the value of the first completed task
token_or_done = cast(Union[str, Literal[True]], done.pop().result())