[websocket] use background tasks instead of cleanup function (#100)

Why
===
* The cleanup function added a burden on the caller to make sure the
tasks were awaited.
* Instead, follow the advice of
https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
and store them in the background properly

What changed
===
* Revert cleanup commit
* Add background websocket task storage

Test plan
===
* Ran a pytest, and didn't see "Pending task destroyed" exceptions

Author: ryantm

replit_river/client.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from typing import Any, Generic, Optional, Union
@@ -39,11 +38,10 @@ def __init__(
             transport_options=transport_options,
         )
 
-    async def close(self) -> asyncio.Task | None:
+    async def close(self) -> None:
         logger.info(f"river client {self._client_id} start closing")
-        cleanup_task = await self._transport.close()
+        await self._transport.close()
         logger.info(f"river client {self._client_id} closed")
-        return cleanup_task
 
     async def ensure_connected(self) -> None:
         await self._transport.get_or_create_session()

replit_river/client_transport.py

@@ -68,9 +68,9 @@ def __init__(
         # We want to make sure there's only one session creation at a time
         self._create_session_lock = asyncio.Lock()
 
-    async def close(self) -> asyncio.Task:
+    async def close(self) -> None:
         self._rate_limiter.close()
-        return await self._close_all_sessions()
+        await self._close_all_sessions()
 
     async def get_or_create_session(self) -> ClientSession:
         async with self._create_session_lock:

replit_river/session.py

@@ -434,18 +434,15 @@ async def _send_responses_from_output_stream(
 
     async def close_websocket(
         self, ws_wrapper: WebsocketWrapper, should_retry: bool
-    ) -> asyncio.Task | None:
+    ) -> None:
         """Mark the websocket as closed, close the websocket, and retry if needed."""
-        cleanup_websocket_task: asyncio.Task | None = None
         async with self._ws_lock:
             # Already closed.
             if not await ws_wrapper.is_open():
-                logger.info("websocket wrapper already closed")
-                return cleanup_websocket_task
-            cleanup_websocket_task = await ws_wrapper.close()
+                return
+            await ws_wrapper.close()
         if should_retry and self._retry_connection_callback:
             self._task_manager.create_task(self._retry_connection_callback())
-        return cleanup_websocket_task
 
     async def _open_stream_and_call_handler(
         self,
@@ -526,9 +523,8 @@ async def _remove_acked_messages_in_buffer(self) -> None:
     async def start_serve_responses(self) -> None:
         self._task_manager.create_task(self.serve())
 
-    async def close(self) -> asyncio.Task | None:
+    async def close(self) -> None:
         """Close the session and all associated streams."""
-        cleanup_websocket_task: asyncio.Task | None = None
         logger.info(
             f"{self._transport_id} closing session "
             f"to {self._to_id}, ws: {self._ws_wrapper.id}, "
@@ -537,14 +533,12 @@ async def close(self) -> asyncio.Task | None:
         async with self._state_lock:
             if self._state != SessionState.ACTIVE:
                 # already closing
-                return cleanup_websocket_task
+                return
             self._state = SessionState.CLOSING
             self._reset_session_close_countdown()
             await self._task_manager.cancel_all_tasks()
 
-            cleanup_websocket_task = await self.close_websocket(
-                self._ws_wrapper, should_retry=False
-            )
+            await self.close_websocket(self._ws_wrapper, should_retry=False)
 
             await self._buffer.close()
 
@@ -559,4 +553,3 @@ async def close(self) -> asyncio.Task | None:
                 self._streams.clear()
 
             self._state = SessionState.CLOSED
-        return cleanup_websocket_task

replit_river/transport.py

@@ -27,8 +27,7 @@ def __init__(
         self._handlers: Dict[Tuple[str, str], Tuple[str, GenericRpcHandler]] = {}
         self._session_lock = asyncio.Lock()
 
-    async def _close_all_sessions(self) -> asyncio.Task:
-        cleanup_tasks: list[asyncio.Task] = []
+    async def _close_all_sessions(self) -> None:
         sessions = self._sessions.values()
         logger.info(
             f"start closing sessions {self._transport_id}, number sessions : "
@@ -39,18 +38,10 @@ async def _close_all_sessions(self) -> asyncio.Task:
         # closing sessions requires access to the session lock, so we need to close
         # them one by one to be safe
         for session in sessions_to_close:
-            cleanup_task = await session.close()
-            if cleanup_task:
-                cleanup_tasks.append(cleanup_task)
+            await session.close()
 
         logger.info(f"Transport closed {self._transport_id}")
 
-        async def cleanup() -> None:
-            for cleanup_task in cleanup_tasks:
-                await cleanup_task
-
-        return asyncio.create_task(cleanup())
-
     async def _delete_session(self, session: Session) -> None:
         async with self._session_lock:
             if session._to_id in self._sessions:

replit_river/websocket_wrapper.py

@@ -5,6 +5,7 @@
 from websockets import WebSocketCommonProtocol
 
 logger = logging.getLogger(__name__)
+_background_tasks: set[asyncio.Task] = set()
 
 
 class WsState(enum.Enum):
@@ -24,14 +25,11 @@ async def is_open(self) -> bool:
         async with self.ws_lock:
             return self.ws_state == WsState.OPEN
 
-    async def close(self) -> asyncio.Task | None:
+    async def close(self) -> None:
         async with self.ws_lock:
             if self.ws_state == WsState.OPEN:
                 self.ws_state = WsState.CLOSING
                 task = asyncio.create_task(self.ws.close())
-                task.add_done_callback(
-                    lambda _: logger.debug("old websocket %s closed.", self.ws.id)
-                )
+                _background_tasks.add(task)
+                task.add_done_callback(_background_tasks.discard)
                 self.ws_state = WsState.CLOSED
-                return task
-            return None