Break out "ServerSession" type

Author: blast-hardcheese

src/replit_river/client_transport.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import Awaitable, Callable
-from typing import Generic
+from typing import Generic, Mapping
 
 import websockets
 from pydantic import ValidationError
@@ -36,6 +36,7 @@
     IgnoreMessageException,
     InvalidMessageException,
 )
+from replit_river.session import Session
 from replit_river.transport import Transport
 from replit_river.transport_options import (
     HandshakeMetadataType,
@@ -47,6 +48,8 @@
 
 
 class ClientTransport(Transport, Generic[HandshakeMetadataType]):
+    _sessions: dict[str, ClientSession]
+
     def __init__(
         self,
         uri_and_metadata_factory: Callable[[], Awaitable[UriAndMetadata]],
@@ -59,6 +62,7 @@ def __init__(
             transport_options=transport_options,
             is_server=False,
         )
+        self._sessions = {}
         self._uri_and_metadata_factory = uri_and_metadata_factory
         self._client_id = client_id
         self._server_id = server_id
@@ -70,7 +74,7 @@ def __init__(
 
     async def close(self) -> None:
         self._rate_limiter.close()
-        await self._close_all_sessions()
+        await self._close_all_sessions(self._get_all_sessions)
 
     async def get_or_create_session(self) -> ClientSession:
         async with self._create_session_lock:
@@ -207,13 +211,13 @@ async def _create_new_session(
             handlers={},
         )
 
-        self._set_session(new_session)
+        self._sessions[new_session._to_id] = new_session
         await new_session.start_serve_responses()
         return new_session
 
     async def _retry_connection(self) -> ClientSession:
         if not self._transport_options.transparent_reconnect:
-            await self._close_all_sessions()
+            await self._close_all_sessions(self._get_all_sessions)
         return await self.get_or_create_session()
 
     async def _send_handshake_request(
@@ -352,3 +356,11 @@ async def _establish_handshake(
                 + f"{handshake_response.status.reason}",
             )
         return handshake_request, handshake_response
+
+    def _get_all_sessions(self) -> Mapping[str, Session]:
+        return self._sessions
+
+    async def _delete_session(self, session: Session) -> None:
+        async with self._session_lock:
+            if session._to_id in self._sessions:
+                del self._sessions[session._to_id]

src/replit_river/server.py

@@ -8,8 +8,8 @@
 
 from replit_river.messages import WebsocketClosedException
 from replit_river.seq_manager import SessionStateMismatchException
+from replit_river.server_session import ServerSession
 from replit_river.server_transport import ServerTransport
-from replit_river.session import Session
 from replit_river.transport import TransportOptions
 
 from .rpc import (
@@ -41,7 +41,7 @@ def add_rpc_handlers(
 
     async def _handshake_to_get_session(
         self, websocket: WebSocketServerProtocol
-    ) -> Session | None:
+    ) -> ServerSession | None:
         """This is a wrapper to make sentry happy, sentry doesn't recognize the
         exception handling outside of a task or asyncio.wait_for. So we need to catch
         the errors specifically here.

src/replit_river/server_session.py

@@ -0,0 +1,20 @@
+import logging
+
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from replit_river.session import Session
+
+from .rpc import (
+    TransportMessageTracingSetter,
+)
+
+logger = logging.getLogger(__name__)
+
+trace_propagator = TraceContextTextMapPropagator()
+trace_setter = TransportMessageTracingSetter()
+
+
+class ServerSession(Session):
+    """A transport object that handles the websocket connection with a client."""
+
+    pass

src/replit_river/server_transport.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any
+from typing import Any, Mapping
 
 import nanoid  # type: ignore  # type: ignore
 from pydantic import ValidationError
@@ -27,6 +27,7 @@
     InvalidMessageException,
     SessionStateMismatchException,
 )
+from replit_river.server_session import ServerSession
 from replit_river.session import Session
 from replit_river.transport import Transport
 from replit_river.transport_options import TransportOptions
@@ -35,6 +36,8 @@
 
 
 class ServerTransport(Transport):
+    _sessions: dict[str, ServerSession]
+
     def __init__(
         self,
         transport_id: str,
@@ -45,11 +48,12 @@ def __init__(
             transport_options=transport_options,
             is_server=True,
         )
+        self._sessions = {}
 
     async def handshake_to_get_session(
         self,
         websocket: WebSocketServerProtocol,
-    ) -> Session:
+    ) -> ServerSession:
         async for message in websocket:
             try:
                 msg = parse_transport_msg(message, self._transport_options)
@@ -88,23 +92,23 @@ async def handshake_to_get_session(
         raise WebsocketClosedException("No handshake message received")
 
     async def close(self) -> None:
-        await self._close_all_sessions()
+        await self._close_all_sessions(self._get_all_sessions)
 
     async def _get_or_create_session(
         self,
         transport_id: str,
         to_id: str,
         session_id: str,
         websocket: WebSocketCommonProtocol,
-    ) -> Session:
+    ) -> ServerSession:
         async with self._session_lock:
             session_to_close: Session | None = None
-            new_session: Session | None = None
+            new_session: ServerSession | None = None
             if to_id not in self._sessions:
                 logger.info(
                     'Creating new session with "%s" using ws: %s', to_id, websocket.id
                 )
-                new_session = Session(
+                new_session = ServerSession(
                     transport_id,
                     to_id,
                     session_id,
@@ -125,7 +129,7 @@ async def _get_or_create_session(
                         old_session.session_id,
                     )
                     session_to_close = old_session
-                    new_session = Session(
+                    new_session = ServerSession(
                         transport_id,
                         to_id,
                         session_id,
@@ -152,7 +156,7 @@ async def _get_or_create_session(
             if session_to_close:
                 logger.info("Closing stale session %s", session_to_close.session_id)
                 await session_to_close.close()
-            self._set_session(new_session)
+            self._sessions[new_session._to_id] = new_session
         return new_session
 
     async def _send_handshake_response(
@@ -293,3 +297,11 @@ async def _establish_handshake(
             )
 
             return handshake_request, handshake_response
+
+    def _get_all_sessions(self) -> Mapping[str, Session]:
+        return self._sessions
+
+    async def _delete_session(self, session: Session) -> None:
+        async with self._session_lock:
+            if session._to_id in self._sessions:
+                del self._sessions[session._to_id]

src/replit_river/session.py

@@ -43,7 +43,7 @@
 
 
 class Session:
-    """A transport object that handles the websocket connection with a client."""
+    """Common functionality shared between client_session and server_session"""
 
     def __init__(
         self,
@@ -253,9 +253,6 @@ async def replace_with_new_websocket(
                 await old_wrapper.close()
             self._ws_wrapper = WebsocketWrapper(new_ws)
         await self._send_buffered_messages(new_ws)
-        # Server will call serve itself.
-        if not self._is_server:
-            await self.start_serve_responses()
 
     async def _get_current_time(self) -> float:
         return asyncio.get_event_loop().time()

src/replit_river/transport.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from typing import Callable, Mapping
 
 import nanoid  # type: ignore
 
@@ -22,12 +23,14 @@ def __init__(
         self._transport_id = transport_id
         self._transport_options = transport_options
         self._is_server = is_server
-        self._sessions: dict[str, Session] = {}
         self._handlers: dict[tuple[str, str], tuple[str, GenericRpcHandler]] = {}
         self._session_lock = asyncio.Lock()
 
-    async def _close_all_sessions(self) -> None:
-        sessions = self._sessions.values()
+    async def _close_all_sessions(
+        self,
+        get_all_sessions: Callable[[], Mapping[str, Session]],
+    ) -> None:
+        sessions = get_all_sessions().values()
         logger.info(
             f"start closing sessions {self._transport_id}, number sessions : "
             f"{len(sessions)}"
@@ -41,13 +44,5 @@ async def _close_all_sessions(self) -> None:
 
         logger.info(f"Transport closed {self._transport_id}")
 
-    async def _delete_session(self, session: Session) -> None:
-        async with self._session_lock:
-            if session._to_id in self._sessions:
-                del self._sessions[session._to_id]
-
-    def _set_session(self, session: Session) -> None:
-        self._sessions[session._to_id] = session
-
     def generate_nanoid(self) -> str:
         return str(nanoid.generate())