Moving setup_heartbeat and check_to_close_session over to v1 session.py

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -1,7 +1,7 @@
 import asyncio
 import enum
 import logging
-from typing import Any, Awaitable, Callable, Coroutine, Protocol
+from typing import Any, Callable, Coroutine, Protocol
 
 from opentelemetry.trace import Span
 from websockets import WebSocketCommonProtocol
@@ -12,7 +12,7 @@
     WebsocketClosedException,
     send_transport_message,
 )
-from replit_river.rpc import ACK_BIT, TransportMessage
+from replit_river.rpc import TransportMessage
 
 logger = logging.getLogger(__name__)
 
@@ -52,86 +52,6 @@ class SessionState(enum.Enum):
 TerminalStates = set([SessionState.CLOSING, SessionState.CLOSED])
 
 
-async def setup_heartbeat(
-    session_id: str,
-    heartbeat_ms: float,
-    heartbeats_until_dead: int,
-    get_state: Callable[[], SessionState],
-    get_closing_grace_period: Callable[[], float | None],
-    close_websocket: Callable[[], Awaitable[None]],
-    send_message: SendMessage,
-    increment_and_get_heartbeat_misses: Callable[[], int],
-) -> None:
-    while True:
-        await asyncio.sleep(heartbeat_ms / 1000)
-        state = get_state()
-        if state == SessionState.CONNECTING:
-            logger.debug("Websocket is not connected, not sending heartbeat")
-            continue
-        if state in TerminalStates:
-            logger.debug(
-                "Session is closed, no need to send heartbeat, state : "
-                "%r close_session_after_this: %r",
-                {state},
-                {get_closing_grace_period()},
-            )
-            # session is closing / closed, no need to send heartbeat anymore
-            return
-        try:
-            await send_message(
-                stream_id="heartbeat",
-                # TODO: make this a message class
-                # https://github.com/replit/river/blob/741b1ea6d7600937ad53564e9cf8cd27a92ec36a/transport/message.ts#L42
-                payload={
-                    "ack": 0,
-                },
-                control_flags=ACK_BIT,
-                procedure_name=None,
-                service_name=None,
-                span=None,
-            )
-
-            if increment_and_get_heartbeat_misses() > heartbeats_until_dead:
-                if get_closing_grace_period() is not None:
-                    # already in grace period, no need to set again
-                    continue
-                logger.info(
-                    "%r closing websocket because of heartbeat misses",
-                    session_id,
-                )
-                await close_websocket()
-                continue
-        except FailedSendingMessageException:
-            # this is expected during websocket closed period
-            continue
-
-
-async def check_to_close_session(
-    transport_id: str,
-    close_session_check_interval_ms: float,
-    get_state: Callable[[], SessionState],
-    get_current_time: Callable[[], Awaitable[float]],
-    get_close_session_after_time_secs: Callable[[], float | None],
-    do_close: Callable[[], Awaitable[None]],
-) -> None:
-    while True:
-        await asyncio.sleep(close_session_check_interval_ms / 1000)
-        if get_state() in TerminalStates:
-            # already closing
-            return
-        # calculate the value now before comparing it so that there are no
-        # await points between the check and the comparison to avoid a TOCTOU
-        # race.
-        current_time = await get_current_time()
-        close_session_after_time_secs = get_close_session_after_time_secs()
-        if not close_session_after_time_secs:
-            continue
-        if current_time > close_session_after_time_secs:
-            logger.info("Grace period ended for %s, closing session", transport_id)
-            await do_close()
-            return
-
-
 async def buffered_message_sender(
     connection_condition: asyncio.Condition,
     message_enqueued: asyncio.Semaphore,

src/replit_river/session.py

@@ -9,9 +9,9 @@
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 from replit_river.common_session import (
+    SendMessage,
     SessionState,
-    check_to_close_session,
-    setup_heartbeat,
+    TerminalStates,
 )
 from replit_river.message_buffer import MessageBuffer, MessageBufferClosedError
 from replit_river.messages import (
@@ -27,6 +27,7 @@
 from replit_river.websocket_wrapper import WebsocketWrapper, WsState
 
 from .rpc import (
+    ACK_BIT,
     TransportMessage,
     TransportMessageTracingSetter,
 )
@@ -337,3 +338,83 @@ async def close(self) -> None:
                 self._streams.clear()
 
             self._state = SessionState.CLOSED
+
+
+async def check_to_close_session(
+    transport_id: str,
+    close_session_check_interval_ms: float,
+    get_state: Callable[[], SessionState],
+    get_current_time: Callable[[], Awaitable[float]],
+    get_close_session_after_time_secs: Callable[[], float | None],
+    do_close: Callable[[], Awaitable[None]],
+) -> None:
+    while True:
+        await asyncio.sleep(close_session_check_interval_ms / 1000)
+        if get_state() in TerminalStates:
+            # already closing
+            return
+        # calculate the value now before comparing it so that there are no
+        # await points between the check and the comparison to avoid a TOCTOU
+        # race.
+        current_time = await get_current_time()
+        close_session_after_time_secs = get_close_session_after_time_secs()
+        if not close_session_after_time_secs:
+            continue
+        if current_time > close_session_after_time_secs:
+            logger.info("Grace period ended for %s, closing session", transport_id)
+            await do_close()
+            return
+
+
+async def setup_heartbeat(
+    session_id: str,
+    heartbeat_ms: float,
+    heartbeats_until_dead: int,
+    get_state: Callable[[], SessionState],
+    get_closing_grace_period: Callable[[], float | None],
+    close_websocket: Callable[[], Awaitable[None]],
+    send_message: SendMessage,
+    increment_and_get_heartbeat_misses: Callable[[], int],
+) -> None:
+    while True:
+        await asyncio.sleep(heartbeat_ms / 1000)
+        state = get_state()
+        if state == SessionState.CONNECTING:
+            logger.debug("Websocket is not connected, not sending heartbeat")
+            continue
+        if state in TerminalStates:
+            logger.debug(
+                "Session is closed, no need to send heartbeat, state : "
+                "%r close_session_after_this: %r",
+                {state},
+                {get_closing_grace_period()},
+            )
+            # session is closing / closed, no need to send heartbeat anymore
+            return
+        try:
+            await send_message(
+                stream_id="heartbeat",
+                # TODO: make this a message class
+                # https://github.com/replit/river/blob/741b1ea6d7600937ad53564e9cf8cd27a92ec36a/transport/message.ts#L42
+                payload={
+                    "ack": 0,
+                },
+                control_flags=ACK_BIT,
+                procedure_name=None,
+                service_name=None,
+                span=None,
+            )
+
+            if increment_and_get_heartbeat_misses() > heartbeats_until_dead:
+                if get_closing_grace_period() is not None:
+                    # already in grace period, no need to set again
+                    continue
+                logger.info(
+                    "%r closing websocket because of heartbeat misses",
+                    session_id,
+                )
+                await close_websocket()
+                continue
+        except FailedSendingMessageException:
+            # this is expected during websocket closed period
+            continue