Enqaw drain pusher ws on reconnects (#2397)

* Draining pusher sockets on disconnects

* Optimize pusher ws; auto-recovery on failure

Author: beastoin

backend/routers/pusher.py

@@ -29,29 +29,6 @@ async def _websocket_util_trigger(
     websocket_active = True
     websocket_close_code = 1000
 
-    # heart beat
-    async def send_heartbeat():
-        print("pusher send_heartbeat", uid)
-        nonlocal websocket_active
-        nonlocal websocket_close_code
-        try:
-            while websocket_active:
-                if websocket.client_state == WebSocketState.CONNECTED:
-                    await websocket.send_json({"type": "ping"})
-                else:
-                    break
-                await asyncio.sleep(10)
-        except WebSocketDisconnect:
-            print("WebSocket disconnected")
-        except Exception as e:
-            print(f'Heartbeat error: {e}')
-            websocket_close_code = 1011
-        finally:
-            websocket_active = False
-
-    # start heart beat
-    heartbeat_task = asyncio.create_task(send_heartbeat())
-
     loop = asyncio.get_event_loop()
 
     # audio bytes
@@ -60,7 +37,7 @@ async def send_heartbeat():
     has_audio_apps_enabled = is_audio_bytes_app_enabled(uid)
 
     # task
-    async def receive_audio_bytes():
+    async def receive_tasks():
         nonlocal websocket_active
         nonlocal websocket_close_code
 
@@ -106,8 +83,8 @@ async def receive_audio_bytes():
             websocket_active = False
 
     try:
-        receive_task = asyncio.create_task(receive_audio_bytes())
-        await asyncio.gather(receive_task, heartbeat_task)
+        receive_task = asyncio.create_task(receive_tasks())
+        await asyncio.gather(receive_task)
 
     except Exception as e:
         print(f"Error during WebSocket operation: {e}")

backend/routers/transcribe.py

@@ -417,7 +417,6 @@ def create_pusher_task_handler():
         pusher_connected = False
 
         # Transcript
-        transcript_ws = None
         segment_buffers = []
         in_progress_conversation_id = None
 
@@ -431,11 +430,11 @@ async def transcript_consume():
             nonlocal websocket_active
             nonlocal segment_buffers
             nonlocal in_progress_conversation_id
-            nonlocal transcript_ws
+            nonlocal pusher_ws
             nonlocal pusher_connected
             while websocket_active or len(segment_buffers) > 0:
                 await asyncio.sleep(1)
-                if transcript_ws and len(segment_buffers) > 0:
+                if pusher_connected and pusher_ws and len(segment_buffers) > 0:
                     try:
                         # 102|data
                         data = bytearray()
@@ -444,17 +443,16 @@ async def transcript_consume():
                             bytes(json.dumps({"segments": segment_buffers, "memory_id": in_progress_conversation_id}),
                                   "utf-8"))
                         segment_buffers = []  # reset
-                        await transcript_ws.send(data)
+                        await pusher_ws.send(data)
                     except websockets.exceptions.ConnectionClosed as e:
                         print(f"Pusher transcripts Connection closed: {e}", uid)
-                        transcript_ws = None
                         pusher_connected = False
-                        await reconnect()
                     except Exception as e:
                         print(f"Pusher transcripts failed: {e}", uid)
+                if pusher_connected is False:
+                    await connect()
 
         # Audio bytes
-        audio_bytes_ws = None
         audio_buffers = bytearray()
         audio_bytes_enabled = bool(get_audio_bytes_webhook_seconds(uid)) or is_audio_bytes_app_enabled(uid)
 
@@ -465,52 +463,56 @@ def audio_bytes_send(audio_bytes):
         async def audio_bytes_consume():
             nonlocal websocket_active
             nonlocal audio_buffers
-            nonlocal audio_bytes_ws
+            nonlocal pusher_ws
             nonlocal pusher_connected
             while websocket_active or len(audio_buffers) > 0:
                 await asyncio.sleep(1)
-                if audio_bytes_ws and len(audio_buffers) > 0:
+                if pusher_connected and pusher_ws and len(audio_buffers) > 0:
                     try:
                         # 101|data
                         data = bytearray()
                         data.extend(struct.pack("I", 101))
                         data.extend(audio_buffers.copy())
                         audio_buffers = bytearray()  # reset
-                        await audio_bytes_ws.send(data)
+                        await pusher_ws.send(data)
                     except websockets.exceptions.ConnectionClosed as e:
                         print(f"Pusher audio_bytes Connection closed: {e}", uid)
-                        audio_bytes_ws = None
                         pusher_connected = False
-                        await reconnect()
                     except Exception as e:
                         print(f"Pusher audio_bytes failed: {e}", uid)
+                if pusher_connected is False:
+                    await connect()
 
-        async def reconnect():
+        async def connect():
             nonlocal pusher_connected
             nonlocal pusher_connect_lock
+            nonlocal pusher_ws
             async with pusher_connect_lock:
                 if pusher_connected:
                     return
-                await connect()
+                # drain
+                if pusher_ws:
+                    try:
+                        await pusher_ws.close()
+                        pusher_ws = None
+                    except Exception as e:
+                        print(f"Pusher draining failed: {e}", uid)
+                # connect
+                await _connect()
 
-        async def connect():
+        async def _connect():
             nonlocal pusher_ws
-            nonlocal transcript_ws
-            nonlocal audio_bytes_ws
-            nonlocal audio_bytes_enabled
             nonlocal pusher_connected
 
             try:
                 pusher_ws = await connect_to_trigger_pusher(uid, sample_rate)
                 pusher_connected = True
-                transcript_ws = pusher_ws
-                if audio_bytes_enabled:
-                    audio_bytes_ws = pusher_ws
             except Exception as e:
                 print(f"Exception in connect: {e}")
 
         async def close(code: int = 1000):
-            await pusher_ws.close(code)
+            if pusher_ws:
+                await pusher_ws.close(code)
 
         return (connect, close,
                 transcript_send, transcript_consume,

backend/utils/pusher.py

@@ -24,7 +24,8 @@ async def _connect_to_trigger_pusher(uid: str, sample_rate: int = 8000):
     try:
         print("Connecting to Pusher transcripts trigger WebSocket...", uid)
         ws_host = PusherAPI.replace("http", "ws")
-        socket = await websockets.connect(f"{ws_host}/v1/trigger/listen?uid={uid}&sample_rate={sample_rate}")
+        socket = await websockets.connect(f"{ws_host}/v1/trigger/listen?uid={uid}&sample_rate={sample_rate}",
+                                          ping_interval=15,)
         print("Connected to Pusher transcripts trigger WebSocket.", uid)
         return socket
     except Exception as e: