small fix to websocketchannel and relay

This commit is contained in:
Timothy Pogue 2022-11-17 12:03:11 -07:00
parent 875e9ab447
commit ce43fa5f43
1 changed files with 17 additions and 14 deletions

View File

@ -77,21 +77,24 @@ class WebSocketChannel:
# until self.drain_timeout for the relay to drain the outgoing queue
# We can't use asyncio.wait_for here because the queue may have been created with a
# different eventloop
start = time.time()
while self.queue.full():
await asyncio.sleep(1)
if (time.time() - start) > self.drain_timeout:
if not self.is_closed():
start = time.time()
while self.queue.full():
await asyncio.sleep(1)
if (time.time() - start) > self.drain_timeout:
return False
# If for some reason the queue is still full, just return False
try:
self.queue.put_nowait(data)
except asyncio.QueueFull:
return False
# If for some reason the queue is still full, just return False
try:
self.queue.put_nowait(data)
except asyncio.QueueFull:
# If we got here everything is ok
return True
else:
return False
# If we got here everything is ok
return True
async def recv(self):
"""
Receive data on the wrapped websocket
@ -109,14 +112,14 @@ class WebSocketChannel:
Close the WebSocketChannel
"""
self._closed.set()
self._relay_task.cancel()
try:
await self.raw_websocket.close()
except Exception:
pass
self._closed.set()
self._relay_task.cancel()
def is_closed(self) -> bool:
"""
Closed flag