Break out heartbeat lifecycle

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -0,0 +1,87 @@
+import asyncio
+import enum
+import logging
+from typing import Any, Awaitable, Callable, Protocol
+
+from opentelemetry.trace import Span
+
+from replit_river.messages import FailedSendingMessageException
+from replit_river.rpc import ACK_BIT
+
+
+logger = logging.getLogger(__name__)
+
+class SendMessage(Protocol):
+    async def __call__(
+        self,
+        *,
+        stream_id: str,
+        payload: dict[Any, Any] | str,
+        control_flags: int,
+        service_name: str | None,
+        procedure_name: str | None,
+        span: Span | None,
+    ) -> None: ...
+
+
+class SessionState(enum.Enum):
+    """The state a session can be in.
+
+    Can only transition from ACTIVE to CLOSING to CLOSED.
+    """
+
+    ACTIVE = 0
+    CLOSING = 1
+    CLOSED = 2
+
+
+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:
+    logger.debug("Start heartbeat")
+    while True:
+        await asyncio.sleep(heartbeat_ms / 1000)
+        state = get_state()
+        if state != SessionState.ACTIVE:
+            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

src/replit_river/session.py

@@ -1,7 +1,6 @@
 import asyncio
-import enum
 import logging
-from typing import Any, Awaitable, Callable, Coroutine, Protocol
+from typing import Any, Callable, Coroutine
 
 import nanoid  # type: ignore
 import websockets
@@ -10,6 +9,7 @@
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from websockets.exceptions import ConnectionClosed
 
+from replit_river.common_session import SessionState, setup_heartbeat
 from replit_river.message_buffer import MessageBuffer, MessageBufferClosedError
 from replit_river.messages import (
     FailedSendingMessageException,
@@ -42,30 +42,6 @@
 trace_setter = TransportMessageTracingSetter()
 
 
-class SendMessage(Protocol):
-    async def __call__(
-        self,
-        *,
-        stream_id: str,
-        payload: dict[Any, Any] | str,
-        control_flags: int,
-        service_name: str | None,
-        procedure_name: str | None,
-        span: Span | None,
-    ) -> None: ...
-
-
-class SessionState(enum.Enum):
-    """The state a session can be in.
-
-    Can only transition from ACTIVE to CLOSING to CLOSED.
-    """
-
-    ACTIVE = 0
-    CLOSING = 1
-    CLOSED = 2
-
-
 class Session:
     """A transport object that handles the websocket connection with a client."""
 
@@ -131,7 +107,7 @@ def increment_and_get_heartbeat_misses() -> int:
             return self._heartbeat_misses
 
         self._task_manager.create_task(
-            self._heartbeat(
+            setup_heartbeat(
                 self.session_id,
                 self._transport_options.heartbeat_ms,
                 self._transport_options.heartbeats_until_dead,
@@ -309,58 +285,6 @@ async def _check_to_close_session(self) -> None:
                 await self.close()
                 return
 
-    async def _heartbeat(
-        self,
-        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:
-        logger.debug("Start heartbeat")
-        while True:
-            await asyncio.sleep(heartbeat_ms / 1000)
-            state = get_state()
-            if state != SessionState.ACTIVE:
-                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 _send_buffered_messages(
         self, websocket: websockets.WebSocketCommonProtocol
     ) -> None: