From d61a238597dd52e2a6e39b4cc315987b0ce4c764 Mon Sep 17 00:00:00 2001 From: Darwin D Wu Date: Fri, 8 May 2026 13:58:10 -0700 Subject: [PATCH 1/3] fix(connect): guard stale websocket send failures --- .../inngest/connect/_internal/models.py | 18 ++++ .../inngest/connect/_internal/ws_utils.py | 18 +++- .../connect/_internal/ws_utils_test.py | 98 +++++++++++++++++++ 3 files changed, 130 insertions(+), 4 deletions(-) create mode 100644 pkg/inngest/inngest/connect/_internal/ws_utils_test.py diff --git a/pkg/inngest/inngest/connect/_internal/models.py b/pkg/inngest/inngest/connect/_internal/models.py index 8328fcb3..dd6ec097 100644 --- a/pkg/inngest/inngest/connect/_internal/models.py +++ b/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): diff --git a/pkg/inngest/inngest/connect/_internal/ws_utils.py b/pkg/inngest/inngest/connect/_internal/ws_utils.py index f1eb072b..f59ebac8 100644 --- a/pkg/inngest/inngest/connect/_internal/ws_utils.py +++ b/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 diff --git a/pkg/inngest/inngest/connect/_internal/ws_utils_test.py b/pkg/inngest/inngest/connect/_internal/ws_utils_test.py new file mode 100644 index 00000000..94d77275 --- /dev/null +++ b/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 From 3108435ca52cbbf782ab4b0ecd4b1018aba661cc Mon Sep 17 00:00:00 2001 From: Darwin D Wu Date: Fri, 8 May 2026 13:58:23 -0700 Subject: [PATCH 2/3] fix(connect): harden request ack and flush handling --- .../inngest/connect/_internal/buffer.py | 12 ++ .../connect/_internal/execution_handler.py | 37 ++++-- .../_internal/execution_handler_test.py | 123 ++++++++++++++++++ 3 files changed, 158 insertions(+), 14 deletions(-) create mode 100644 pkg/inngest/inngest/connect/_internal/execution_handler_test.py diff --git a/pkg/inngest/inngest/connect/_internal/buffer.py b/pkg/inngest/inngest/connect/_internal/buffer.py index 4f7c7e35..f3768d67 100644 --- a/pkg/inngest/inngest/connect/_internal/buffer.py +++ b/pkg/inngest/inngest/connect/_internal/buffer.py @@ -131,3 +131,15 @@ def get_older_than(self, seconds: float) -> list[tuple[str, bytes]]: result.append((item.id, item.data)) return result + + 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 diff --git a/pkg/inngest/inngest/connect/_internal/execution_handler.py b/pkg/inngest/inngest/connect/_internal/execution_handler.py index 52e211b1..8fd74043 100644 --- a/pkg/inngest/inngest/connect/_internal/execution_handler.py +++ b/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. diff --git a/pkg/inngest/inngest/connect/_internal/execution_handler_test.py b/pkg/inngest/inngest/connect/_internal/execution_handler_test.py new file mode 100644 index 00000000..c0b98293 --- /dev/null +++ b/pkg/inngest/inngest/connect/_internal/execution_handler_test.py @@ -0,0 +1,123 @@ +import logging +import typing +import unittest +from unittest import mock + +import httpx +import websockets + +from inngest._internal import comm_lib, net + +from . import connect_pb2 +from . import ws_utils as ws_utils_module +from .execution_handler import ExecutionHandler +from .models import ConnectionState, State +from .value_watcher import ValueWatcher + + +class _FakeCommHandler: + called = False + + async def post( + self, + req: comm_lib.CommRequest, + ) -> comm_lib.CommResponse: + self.called = True + raise AssertionError("post should not be called") + + +class _FakeWS: + async def send(self, message: bytes) -> None: + pass + + +class _TestExecutionHandler(ExecutionHandler): + flush_results: list[Exception | None] + + async def _flush_message(self, msg: bytes) -> Exception | None: + return self.flush_results.pop(0) + + +def _state() -> State: + return State( + 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, _FakeWS())), + ) + + +def _handler() -> _TestExecutionHandler: + handler = _TestExecutionHandler( + 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=_state(), + ) + handler.flush_results = [] + return handler + + +class TestExecutionHandler(unittest.IsolatedAsyncioTestCase): + async def test_ack_failure_abandons_request_without_error_reply( + self, + ) -> None: + handler = _handler() + comm_handler = _FakeCommHandler() + req_data = connect_pb2.GatewayExecutorRequestData( + account_id="account", + app_id="app", + env_id="env", + function_slug="fn", + request_id="req", + ) + send_calls: list[bytes] = [] + + async def fail_ack( + logger: object, + state: State, + message: bytes, + ) -> Exception | None: + send_calls.append(message) + return OSError("connection reset by peer") + + with mock.patch.object( + ws_utils_module, + "safe_send", + fail_ack, + ): + await handler._execute_request( + req_data, + typing.cast(comm_lib.CommHandler, comm_handler), + ) + + assert len(send_calls) == 1 + assert comm_handler.called is False + assert handler._buffer.get(req_data.request_id) is None + + async def test_failed_flush_keeps_message_for_retry(self) -> None: + handler = _handler() + handler.flush_results = [Exception("temporary failure")] + 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: + handler = _handler() + handler.flush_results = [None] + handler._buffer.add("req", b"reply") + + await handler._flush_ready_messages(0) + + assert handler._buffer.get("req") is None From a4b384af15ea0b5a20e21fe24e56b746345f5374 Mon Sep 17 00:00:00 2001 From: Aaron Harper Date: Sat, 9 May 2026 14:09:11 -0400 Subject: [PATCH 3/3] Tweak execution_handler_test.py --- .../inngest/connect/_internal/buffer.py | 7 + .../_internal/execution_handler_test.py | 167 ++++++++++-------- 2 files changed, 98 insertions(+), 76 deletions(-) diff --git a/pkg/inngest/inngest/connect/_internal/buffer.py b/pkg/inngest/inngest/connect/_internal/buffer.py index f3768d67..1eafa666 100644 --- a/pkg/inngest/inngest/connect/_internal/buffer.py +++ b/pkg/inngest/inngest/connect/_internal/buffer.py @@ -132,6 +132,13 @@ def get_older_than(self, seconds: float) -> list[tuple[str, bytes]]: 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. diff --git a/pkg/inngest/inngest/connect/_internal/execution_handler_test.py b/pkg/inngest/inngest/connect/_internal/execution_handler_test.py index c0b98293..f9270bee 100644 --- a/pkg/inngest/inngest/connect/_internal/execution_handler_test.py +++ b/pkg/inngest/inngest/connect/_internal/execution_handler_test.py @@ -1,7 +1,6 @@ import logging import typing import unittest -from unittest import mock import httpx import websockets @@ -9,70 +8,77 @@ from inngest._internal import comm_lib, net from . import connect_pb2 -from . import ws_utils as ws_utils_module from .execution_handler import ExecutionHandler from .models import ConnectionState, State from .value_watcher import ValueWatcher -class _FakeCommHandler: - called = False - - async def post( - self, - req: comm_lib.CommRequest, - ) -> comm_lib.CommResponse: - self.called = True - raise AssertionError("post should not be called") - - class _FakeWS: - async def send(self, message: bytes) -> None: - pass - + def __init__(self) -> None: + self.sent: list[bytes] = [] -class _TestExecutionHandler(ExecutionHandler): - flush_results: list[Exception | None] - - async def _flush_message(self, msg: bytes) -> Exception | None: - return self.flush_results.pop(0) - - -def _state() -> State: - return State( - 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, _FakeWS())), - ) + 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)), + ) -def _handler() -> _TestExecutionHandler: - handler = _TestExecutionHandler( - 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=_state(), - ) - handler.flush_results = [] - return handler +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: - handler = _handler() - comm_handler = _FakeCommHandler() + """ + 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", @@ -80,33 +86,34 @@ async def test_ack_failure_abandons_request_without_error_reply( function_slug="fn", request_id="req", ) - send_calls: list[bytes] = [] - - async def fail_ack( - logger: object, - state: State, - message: bytes, - ) -> Exception | None: - send_calls.append(message) - return OSError("connection reset by peer") - - with mock.patch.object( - ws_utils_module, - "safe_send", - fail_ack, - ): - await handler._execute_request( - req_data, - typing.cast(comm_lib.CommHandler, comm_handler), - ) - - assert len(send_calls) == 1 + + 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 - assert handler._buffer.get(req_data.request_id) is None + + # Nothing buffered + assert handler._buffer.length() == 0 async def test_failed_flush_keeps_message_for_retry(self) -> None: - handler = _handler() - handler.flush_results = [Exception("temporary failure")] + """ + 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) @@ -114,8 +121,16 @@ async def test_failed_flush_keeps_message_for_retry(self) -> None: assert handler._buffer.get("req") == b"reply" async def test_successful_flush_deletes_message(self) -> None: - handler = _handler() - handler.flush_results = [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)