fix(connect): harden websocket reconnect and reply delivery

pkg/inngest/inngest/connect/_internal/buffer.py

@@ -131,3 +131,22 @@ def get_older_than(self, seconds: float) -> list[tuple[str, bytes]]:
                 result.append((item.id, item.data))
 
         return result
+
+    def length(self) -> int:
+        """
+        Get number of items in buffer.
+        """
+
+        return len(self._items)
+
+    def touch(self, item_id: str) -> bool:
+        """
+        Refresh an item's timestamp so it is retried after the next TTL window.
+        """
+
+        item = self._items.get(item_id)
+        if item is None:
+            return False
+
+        item.timestamp = time.time()
+        return True

pkg/inngest/inngest/connect/_internal/execution_handler.py

@@ -264,9 +264,13 @@ async def _execute_request(
                 comm_res = await asyncio.wrap_future(future)
             else:
                 self._logger.error(
-                    "Execution failed", extra={"error": str(err)}
+                    "Failed to acknowledge executor request",
+                    extra={
+                        "error": str(err),
+                        "request_id": req_data.request_id,
+                    },
                 )
-                comm_res = comm_lib.CommResponse.from_error(self._logger, err)
+                return
 
             body = comm_res.body_bytes()
             if isinstance(body, Exception):
@@ -484,20 +488,25 @@ async def _unacked_msg_flush_poller(self) -> None:
         flush_ttl = await self._state.extend_lease_interval.wait_for_not_none()
 
         while self.closed_event.is_set() is False:
-            for request_id, reply_msg in self._buffer.get_older_than(flush_ttl):
-                try:
-                    err = await self._flush_message(reply_msg)
-                    if err is not None:
-                        self._logger.error(
-                            "Failed to flush message", extra={"error": str(err)}
-                        )
-                finally:
-                    # We only attempt to flush once, so we can delete the
-                    # message.
-                    self._buffer.delete(request_id)
-
+            await self._flush_ready_messages(flush_ttl)
             await asyncio.sleep(1)
 
+    async def _flush_ready_messages(self, flush_ttl: int | float) -> None:
+        for request_id, reply_msg in self._buffer.get_older_than(flush_ttl):
+            err = await self._flush_message(reply_msg)
+            if err is None:
+                self._buffer.delete(request_id)
+                continue
+
+            self._logger.error(
+                "Failed to flush message",
+                extra={
+                    "error": str(err),
+                    "request_id": request_id,
+                },
+            )
+            self._buffer.touch(request_id)
+
     async def _flush_message(self, msg: bytes) -> types.MaybeError[None]:
         """
         Flush a single message via HTTP.

pkg/inngest/inngest/connect/_internal/execution_handler_test.py

@@ -0,0 +1,138 @@
+import logging
+import typing
+import unittest
+
+import httpx
+import websockets
+
+from inngest._internal import comm_lib, net
+
+from . import connect_pb2
+from .execution_handler import ExecutionHandler
+from .models import ConnectionState, State
+from .value_watcher import ValueWatcher
+
+
+class _FakeWS:
+    def __init__(self) -> None:
+        self.sent: list[bytes] = []
+
+    async def send(self, message: bytes) -> None:
+        self.sent.append(message)
+
+
+class _FakeState(State):
+    def __init__(self, ws: _FakeWS) -> None:
+        super().__init__(
+            conn_id=ValueWatcher(None),
+            conn_init=ValueWatcher(None),
+            conn_state=ValueWatcher(ConnectionState.ACTIVE),
+            exclude_gateways=ValueWatcher([]),
+            extend_lease_interval=ValueWatcher(1),
+            fatal_error=ValueWatcher(None),
+            init_handshake_complete=ValueWatcher(True),
+            pending_request_count=ValueWatcher(0),
+            ws=ValueWatcher(typing.cast(websockets.ClientConnection, ws)),
+        )
+
+
+class _FakeExecutionHandler(ExecutionHandler):
+    def __init__(self, ws: _FakeWS | None = None) -> None:
+        super().__init__(
+            api_origin="http://127.0.0.1",
+            comm_handlers={},
+            http_client=typing.cast(net.ThreadAwareAsyncHTTPClient, object()),
+            http_client_sync=typing.cast(httpx.Client, object()),
+            logger=logging.getLogger(__name__),
+            signing_key=None,
+            signing_key_fallback=None,
+            state=_FakeState(ws or _FakeWS()),
+        )
+
+
+class TestExecutionHandler(unittest.IsolatedAsyncioTestCase):
+    async def test_ack_failure_abandons_request_without_error_reply(
+        self,
+    ) -> None:
+        """
+        If there's an error when sending the execution request ack, then we
+        don't process the execution request and we don't buffer.
+        """
+
+        class WS(_FakeWS):
+            async def send(self, message: bytes) -> None:
+                await super().send(message)
+                raise OSError("connection reset by peer")
+
+        ws = WS()
+        handler = _FakeExecutionHandler(ws)
+
+        class CommHandler:
+            called = False
+
+            async def post(
+                self,
+                req: comm_lib.CommRequest,
+            ) -> comm_lib.CommResponse:
+                print("yo")
+                self.called = True
+                raise Exception("unreachable")
+
+        comm_handler = CommHandler()
+        req_data = connect_pb2.GatewayExecutorRequestData(
+            account_id="account",
+            app_id="app",
+            env_id="env",
+            function_slug="fn",
+            request_id="req",
+        )
+
+        await handler._execute_request(
+            req_data,
+            typing.cast(comm_lib.CommHandler, comm_handler),
+        )
+
+        # We attempted to send the execution request ack
+        assert len(ws.sent) == 1
+        msg = connect_pb2.ConnectMessage()
+        msg.ParseFromString(ws.sent[0])
+        assert msg.kind == connect_pb2.GatewayMessageType.WORKER_REQUEST_ACK
+
+        # CommHandler was not called since the ack failed to send
+        assert comm_handler.called is False
+
+        # Nothing buffered
+        assert handler._buffer.length() == 0
+
+    async def test_failed_flush_keeps_message_for_retry(self) -> None:
+        """
+        If a message flush fails then the buffer retains the message.
+        """
+
+        class Handler(_FakeExecutionHandler):
+            async def _flush_message(self, msg: bytes) -> Exception | None:
+                return Exception("temporary failure")
+
+        handler = Handler()
+        handler._buffer.add("req", b"reply")
+
+        await handler._flush_ready_messages(0)
+
+        assert handler._buffer.get("req") == b"reply"
+
+    async def test_successful_flush_deletes_message(self) -> None:
+        """
+        If a message flush succeeds then the buffer deletes the message.
+        """
+
+        class Handler(_FakeExecutionHandler):
+            async def _flush_message(self, msg: bytes) -> Exception | None:
+                return None
+
+        handler = Handler()
+
+        handler._buffer.add("req", b"reply")
+
+        await handler._flush_ready_messages(0)
+
+        assert handler._buffer.get("req") is None

pkg/inngest/inngest/connect/_internal/models.py

@@ -48,10 +48,28 @@ def close_ws(self) -> None:
         Close the WebSocket connection
         """
 
+        self.close_ws_if_current(None)
+
+    def close_ws_if_current(
+        self,
+        ws: websockets.ClientConnection | None,
+    ) -> bool:
+        """
+        Clear the WebSocket if it is still current.
+
+        Args:
+            ws: If provided, only clear the WebSocket when it still matches this
+                connection. If omitted, always clear it.
+        """
+
+        if ws is not None and self.ws.value is not ws:
+            return False
+
         if self.allow_reconnect():
             self.conn_state.value = ConnectionState.RECONNECTING
         self.conn_init.value = None
         self.ws.value = None
+        return True
 
 
 class ConnectionState(enum.Enum):

pkg/inngest/inngest/connect/_internal/ws_utils.py

@@ -6,6 +6,17 @@
 from . import models
 
 
+def _is_connection_fatal_send_error(err: Exception) -> bool:
+    return isinstance(
+        err,
+        (
+            websockets.exceptions.ConnectionClosed,
+            OSError,
+            EOFError,
+        ),
+    )
+
+
 async def safe_send(
     logger: types.Logger,
     state: models.State,
@@ -17,17 +28,16 @@ async def safe_send(
     connection to trigger a reconnect.
     """
 
+    ws: websockets.ClientConnection | None = None
     try:
         ws = state.ws.value
         if ws is None:
             return Exception("No WebSocket connection")
         await ws.send(message)
-    except websockets.exceptions.ConnectionClosed as e:
-        logger.error(f"Error sending message: {e!s}", extra={"error": str(e)})
-        state.close_ws()
-        return e
     except Exception as e:
         logger.error(f"Error sending message: {e!s}", extra={"error": str(e)})
+        if ws is not None and _is_connection_fatal_send_error(e):
+            state.close_ws_if_current(ws)
         return e
 
     return None

pkg/inngest/inngest/connect/_internal/ws_utils_test.py

@@ -0,0 +1,98 @@
+import logging
+import typing
+import unittest
+
+import websockets
+
+from . import ws_utils
+from .models import ConnectionState, State
+from .value_watcher import ValueWatcher
+
+
+class _FakeWS:
+    def __init__(
+        self,
+        err: Exception | None = None,
+        on_send: typing.Callable[[], None] | None = None,
+    ) -> None:
+        self._err = err
+        self._on_send = on_send
+        self.messages: list[bytes] = []
+
+    async def send(self, message: bytes) -> None:
+        self.messages.append(message)
+        if self._on_send is not None:
+            self._on_send()
+        if self._err is not None:
+            raise self._err
+
+
+def _state(
+    ws: websockets.ClientConnection | None,
+    *,
+    conn_state: ConnectionState = ConnectionState.ACTIVE,
+) -> State:
+    return State(
+        conn_id=ValueWatcher(None),
+        conn_init=ValueWatcher(None),
+        conn_state=ValueWatcher(conn_state),
+        exclude_gateways=ValueWatcher([]),
+        extend_lease_interval=ValueWatcher(None),
+        fatal_error=ValueWatcher(None),
+        init_handshake_complete=ValueWatcher(True),
+        pending_request_count=ValueWatcher(0),
+        ws=ValueWatcher(ws),
+    )
+
+
+class TestSafeSend(unittest.IsolatedAsyncioTestCase):
+    async def test_closes_current_ws_for_connection_fatal_error(self) -> None:
+        ws = typing.cast(
+            websockets.ClientConnection,
+            _FakeWS(OSError("connection reset by peer")),
+        )
+        state = _state(ws)
+
+        err = await ws_utils.safe_send(
+            logging.getLogger(__name__), state, b"msg"
+        )
+
+        assert isinstance(err, OSError)
+        assert state.ws.value is None
+        assert state.conn_state.value == ConnectionState.RECONNECTING
+
+    async def test_does_not_close_replacement_ws(self) -> None:
+        replacement_ws = typing.cast(websockets.ClientConnection, _FakeWS())
+        state = _state(None)
+
+        stale_ws = typing.cast(
+            websockets.ClientConnection,
+            _FakeWS(
+                OSError("connection reset by peer"),
+                on_send=lambda: setattr(state.ws, "value", replacement_ws),
+            ),
+        )
+        state.ws.value = stale_ws
+
+        err = await ws_utils.safe_send(
+            logging.getLogger(__name__), state, b"msg"
+        )
+
+        assert isinstance(err, OSError)
+        assert state.ws.value is replacement_ws
+        assert state.conn_state.value == ConnectionState.ACTIVE
+
+    async def test_non_fatal_error_does_not_close_ws(self) -> None:
+        ws = typing.cast(
+            websockets.ClientConnection,
+            _FakeWS(ValueError("bad message")),
+        )
+        state = _state(ws)
+
+        err = await ws_utils.safe_send(
+            logging.getLogger(__name__), state, b"msg"
+        )
+
+        assert isinstance(err, ValueError)
+        assert state.ws.value is ws
+        assert state.conn_state.value == ConnectionState.ACTIVE