Avoid deadlocking client if streams don't clean up after themselves

Author: blast-hardcheese

src/replit_river/transport_options.py

@@ -27,6 +27,7 @@ class TransportOptions(BaseModel):
     connection_retry_options: ConnectionRetryOptions = ConnectionRetryOptions()
     buffer_size: int = 1_000
     transparent_reconnect: bool = True
+    shutdown_all_streams_timeout_ms: float = 10_000
 
     def websocket_disconnect_grace_ms(self) -> float:
         return self.heartbeat_ms * self.heartbeats_until_dead
@@ -39,11 +40,16 @@ def create_from_env(cls) -> "TransportOptions":
         )
         heartbeat_ms = float(os.getenv("HEARTBEAT_MS", 2_000))
         heartbeats_to_dead = int(os.getenv("HEARTBEATS_UNTIL_DEAD", 2))
+        shutdown_all_streams_timeout_ms = float(
+            os.getenv("SHUTDOWN_STREAMS_TIMEOUT_MS", 10_000)
+        )
+
         return TransportOptions(
             handshake_timeout_ms=handshake_timeout_ms,
             session_disconnect_grace_ms=session_disconnect_grace_ms,
             heartbeat_ms=heartbeat_ms,
             heartbeats_until_dead=heartbeats_to_dead,
+            shutdown_all_streams_timeout_ms=shutdown_all_streams_timeout_ms,
         )
 
 

src/replit_river/v2/session.py

@@ -262,8 +262,8 @@ def close_session(reason: Exception | None) -> None:
                 # during the cleanup procedure.
 
                 self._terminating_task = asyncio.create_task(
-                        self.close(reason, current_state=current_state),
-                    )
+                    self.close(reason, current_state=current_state),
+                )
 
         def transition_connecting() -> None:
             if self._state in TerminalStates:
@@ -396,7 +396,9 @@ async def _enqueue_message(
         # Wake up buffered_message_sender
         self._process_messages.set()
 
-    async def close(self, reason: Exception | None = None, current_state: SessionState | None = None ) -> None:
+    async def close(
+        self, reason: Exception | None = None, current_state: SessionState | None = None
+    ) -> None:
         """Close the session and all associated streams."""
         logger.info(
             f"{self.session_id} closing session to {self._server_id}, ws: {self._ws}"
@@ -435,9 +437,28 @@ async def close(self, reason: Exception | None = None, current_state: SessionSta
                 )
             stream_meta["release_backpressured_waiter"]()
         # Before we GC the streams, let's wait for all tasks to be closed gracefully.
-        await asyncio.gather(
-            *[stream_meta["output"].join() for stream_meta in self._streams.values()]
-        )
+        try:
+            async with asyncio.timeout(
+                self._transport_options.shutdown_all_streams_timeout_ms
+            ):
+                # Block for backpressure and emission errors from the ws
+                await asyncio.gather(
+                    *[
+                        stream_meta["output"].join()
+                        for stream_meta in self._streams.values()
+                    ]
+                )
+        except asyncio.TimeoutError:
+            spans: list[Span] = [
+                stream_meta["span"]
+                for stream_meta in self._streams.values()
+                if not stream_meta["output"].closed()
+            ]
+            span_ids = [span.get_span_context().span_id for span in spans]
+            logger.exception(
+                "Timeout waiting for output streams to finallize",
+                extra={"span_ids": span_ids},
+            )
         self._streams.clear()
 
         if self._ws: