Speed up timeout implementation in send_message_await_response_complex (#478)

This commit is contained in:
J. Nick Koston 2023-07-15 09:28:06 -10:00 committed by GitHub
parent 5dd12169f1
commit bbfa761aa0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 3 deletions

View File

@ -596,6 +596,12 @@ class APIConnection:
for msg_type in msg_types:
self._message_handlers.setdefault(msg_type, []).append(on_message)
def _handle_timeout(self, fut: asyncio.Future[None]) -> None:
"""Handle a timeout."""
if fut.done():
return
fut.set_exception(asyncio.TimeoutError())
async def send_message_await_response_complex(
self,
send_msg: message.Message,
@ -639,15 +645,15 @@ class APIConnection:
# We must not await without a finally or
# the message could fail to be removed if the
# the await is cancelled
timeout_handle = self._loop.call_later(timeout, self._handle_timeout, fut)
try:
async with async_timeout.timeout(timeout):
await fut
await fut
except asyncio.TimeoutError as err:
raise TimeoutAPIError(
f"Timeout waiting for response for {type(send_msg)} after {timeout}s"
) from err
finally:
timeout_handle.cancel()
for msg_type in msg_types:
with suppress(ValueError):
self._message_handlers[msg_type].remove(on_message)