Imrpove timeout behavior and cleanliness

Pipfile

@@ -10,6 +10,7 @@ toml = "*"
 aioredis = {extras = ["hiredis"], version = "*"}
 httpx = "*"
 aioscheduler = "*"
+async-timeout = "*"
 
 [dev-packages]
 python-language-server = {extras = ["all"], version = "*"}

Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "c70908392b9cfefc6dd5a39017e71d95379cdac7c457ba15a1139871590a36df"
+            "sha256": "a8a7084ed6a8ffcda9f47292ff89f4279aa9123e09c4a3e7114786048ff7d68e"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -56,7 +56,7 @@
                 "sha256:2163e1640ddb52b7a8c80d0a67a08587e5d245cc9c553a74a847056bc2976b15",
                 "sha256:8ca1e4fcf50d07413d66d1a5e416e42cfdf5851c981d679a09851a6853383b3c"
             ],
-            "markers": "python_version >= '3.6'",
+            "index": "pypi",
             "version": "==4.0.2"
         },
         "blinker": {

app.py

@@ -11,6 +11,7 @@ import toml
 import aioredis
 import httpx
 from aioscheduler import TimedScheduler
+from async_timeout import timeout
 
 ENV_PREFIX = 'EVENTHOOK_'
 
@@ -193,16 +194,21 @@ async def client_websocket(client_id):
         await asyncio.gather(producer, consumer, heartbeat)
 
     except asyncio.CancelledError as e:
-        await websocket.close(1002)
+        ws_status = 1002
         raise e
     except asyncio.TimeoutError as e:
         app.logger.error(f'Timing out websocket for {client_id}')
-        await websocket.close(1005)
+        ws_status = 1005
     except Exception as e:
-        await websocket.close(1001)
+        ws_status = 1001
         app.logger.error(f'Websocket handle failed! - {type(e).__name__}: {e}')
         app.logger.error(''.join(format_exception(None, e, e.__traceback__)))
     finally:
+        try:
+            async with timeout(3):
+                await websocket.close(ws_status)
+        except asyncio.TimeoutError:
+            pass
         clients = await app.redis.decr(client_tracker_key.format(client_id))
         app.logger.info(f'End: ws for {client_id}. {clients} clients alive')
         if clients == 0:
@@ -253,7 +259,8 @@ async def ws_heartbeat(pong_event):
     while True:
         await asyncio.sleep(30)
         await websocket.send('PING')
-        response = await asyncio.wait_for(pong_event.wait(), timeout=3)
+        async with timeout(3):
+            await pong_event.wait()
         pong_event.clear()
 
 if __name__ == "__main__":