diff --git a/README.md b/README.md index 30bfdab..578ee55 100644 --- a/README.md +++ b/README.md @@ -292,6 +292,8 @@ async for row in client.run_eventql_query( *Note that each row returned by the generator matches the projection specified in your query.* +While the query runs, the server sends a heartbeat every second whenever there is no row to send. If neither a row nor a heartbeat arrives for 30 seconds, the generator closes the connection and raises a `HeartbeatTimeoutError`, which is a subclass of `ServerError`. For details, see [Detecting a Stalled Connection](#detecting-a-stalled-connection). + #### Aborting a Query If you need to abort a query use `break` or `return` within the `async for` loop. However, this only works if there is currently an iteration going on. @@ -456,6 +458,26 @@ async for event in client.observe_events( *Note that `from_latest_event` and `lower_bound` can not be provided at the same time.* +#### Detecting a Stalled Connection + +While observing, the server sends a heartbeat every second whenever there is no new event. If neither an event nor a heartbeat arrives for 30 seconds, for example because a proxy keeps the connection open but no longer passes anything on, the generator closes the connection and raises a `HeartbeatTimeoutError`. Since `HeartbeatTimeoutError` is a subclass of `ServerError`, you can handle it specifically or together with other server errors: + +```python +from eventsourcingdb import HeartbeatTimeoutError, ObserveEventsOptions + +try: + async for event in client.observe_events( + subject = '/books/42', + options = ObserveEventsOptions( + recursive = False + ), + ): + pass +except HeartbeatTimeoutError: + # Neither an event nor a heartbeat arrived for 30 seconds. + pass +``` + #### Aborting Observing If you need to abort observing use `break` or `return` within the `async for` loop. However, this only works if there is currently an iteration going on. diff --git a/eventsourcingdb/__init__.py b/eventsourcingdb/__init__.py index 99ca576..da15e94 100644 --- a/eventsourcingdb/__init__.py +++ b/eventsourcingdb/__init__.py @@ -4,6 +4,7 @@ from .errors import ( ClientError, CustomError, + HeartbeatTimeoutError, InternalError, ServerError, ValidationError, @@ -39,6 +40,7 @@ "Event", "EventCandidate", "EventType", + "HeartbeatTimeoutError", "IfEventIsMissingDuringObserve", "IfEventIsMissingDuringRead", "InternalError", diff --git a/eventsourcingdb/client.py b/eventsourcingdb/client.py index d5d9ce4..3807785 100644 --- a/eventsourcingdb/client.py +++ b/eventsourcingdb/client.py @@ -16,6 +16,7 @@ from .parse_raw_message import parse_raw_message from .read_event_types import EventType, is_event_type from .read_events import ReadEventsOptions +from .read_lines_with_heartbeat_timeout import read_lines_with_heartbeat_timeout from .read_subjects import is_subject from .write_events import Precondition @@ -185,7 +186,7 @@ async def run_eventql_query(self, query: str) -> AsyncGenerator[Any]: async with response: self._validate_response(response) - async for raw_message in response.body: + async for raw_message in read_lines_with_heartbeat_timeout(response.body): message = parse_raw_message(raw_message) if is_heartbeat(message): @@ -222,7 +223,7 @@ async def observe_events( async with response: self._validate_response(response) - async for raw_message in response.body: + async for raw_message in read_lines_with_heartbeat_timeout(response.body): message = parse_raw_message(raw_message) if is_heartbeat(message): diff --git a/eventsourcingdb/errors/__init__.py b/eventsourcingdb/errors/__init__.py index d2b3228..cc3648f 100644 --- a/eventsourcingdb/errors/__init__.py +++ b/eventsourcingdb/errors/__init__.py @@ -1,5 +1,6 @@ from .client_error import ClientError from .custom_error import CustomError +from .heartbeat_timeout_error import HeartbeatTimeoutError from .internal_error import InternalError from .server_error import ServerError from .validation_error import ValidationError @@ -7,6 +8,7 @@ __all__ = [ "ClientError", "CustomError", + "HeartbeatTimeoutError", "InternalError", "ServerError", "ValidationError", diff --git a/eventsourcingdb/errors/heartbeat_timeout_error.py b/eventsourcingdb/errors/heartbeat_timeout_error.py new file mode 100644 index 0000000..7418e02 --- /dev/null +++ b/eventsourcingdb/errors/heartbeat_timeout_error.py @@ -0,0 +1,6 @@ +from .server_error import ServerError + + +class HeartbeatTimeoutError(ServerError): + def __init__(self) -> None: + super().__init__("No event and no heartbeat arrived for 30 seconds") diff --git a/eventsourcingdb/read_lines_with_heartbeat_timeout.py b/eventsourcingdb/read_lines_with_heartbeat_timeout.py new file mode 100644 index 0000000..b61110c --- /dev/null +++ b/eventsourcingdb/read_lines_with_heartbeat_timeout.py @@ -0,0 +1,27 @@ +import asyncio +from collections.abc import AsyncGenerator + +from aiohttp import StreamReader + +from .errors import HeartbeatTimeoutError + +HEARTBEAT_TIMEOUT_SECONDS = 30 + + +async def read_lines_with_heartbeat_timeout(stream: StreamReader) -> AsyncGenerator[bytes, None]: + while True: + timeout = asyncio.timeout(HEARTBEAT_TIMEOUT_SECONDS) + try: + async with timeout: + line = await stream.readline() + except TimeoutError as error: + # Only this timeout means that the server went silent, any other + # timeout (for example the one of the HTTP session) is passed on. + if not timeout.expired(): + raise + raise HeartbeatTimeoutError() from error + + if line == b'': + return + + yield line diff --git a/tests/shared/stream_server.py b/tests/shared/stream_server.py new file mode 100644 index 0000000..add3ef0 --- /dev/null +++ b/tests/shared/stream_server.py @@ -0,0 +1,75 @@ +import asyncio +from collections.abc import Awaitable, Callable +from types import TracebackType +from typing import Self + +from aiohttp import web + +from eventsourcingdb import Client + +WriteStream = Callable[[web.StreamResponse], Awaitable[None]] + + +class StreamServer: + """ + A local HTTP server that stands in for EventSourcingDB. It answers + observe-events and run-eventql-query with whatever the given function + writes to the stream, so that tests can control every single line. + """ + + def __init__(self, write_stream: WriteStream) -> None: + self.__write_stream = write_stream + self.__runner: web.AppRunner | None = None + self.__client: Client | None = None + self.client_disconnected = asyncio.Event() + + async def __aenter__(self) -> Self: + application = web.Application() + application.router.add_post('/api/v1/observe-events', self.__handle) + application.router.add_post('/api/v1/run-eventql-query', self.__handle) + + # Cancelling the handler when the client disconnects is what lets + # tests notice that the client closed the connection. + self.__runner = web.AppRunner(application, handler_cancellation=True) + await self.__runner.setup() + await web.TCPSite(self.__runner, '127.0.0.1', 0).start() + host, port = self.__runner.addresses[0][:2] + + self.__client = Client(base_url=f'http://{host}:{port}', api_token='secret') + await self.__client.__aenter__() + + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None = None, + exc_val: BaseException | None = None, + exc_tb: TracebackType | None = None, + ) -> None: + if self.__client is not None: + await self.__client.__aexit__(exc_type, exc_val, exc_tb) + if self.__runner is not None: + await self.__runner.cleanup() + + def get_client(self) -> Client: + if self.__client is None: + raise RuntimeError('Stream server must be started before getting a client.') + + return self.__client + + async def __handle(self, request: web.Request) -> web.StreamResponse: + response = web.StreamResponse( + headers={ + 'Server': 'EventSourcingDB/test', + 'Content-Type': 'application/x-ndjson', + }, + ) + await response.prepare(request) + + try: + await self.__write_stream(response) + except asyncio.CancelledError: + self.client_disconnected.set() + raise + + return response diff --git a/tests/test_heartbeat_timeout.py b/tests/test_heartbeat_timeout.py new file mode 100644 index 0000000..b2bf583 --- /dev/null +++ b/tests/test_heartbeat_timeout.py @@ -0,0 +1,208 @@ +import asyncio +import json +import time +from collections.abc import AsyncGenerator, Callable +from typing import Any + +import pytest +from aiohttp import web + +import eventsourcingdb.read_lines_with_heartbeat_timeout as heartbeat_timeout_module +from eventsourcingdb import Client, HeartbeatTimeoutError, ObserveEventsOptions +from eventsourcingdb.read_lines_with_heartbeat_timeout import ( + read_lines_with_heartbeat_timeout, +) + +from .shared.stream_server import StreamServer + +SHORT_HEARTBEAT_TIMEOUT_SECONDS = 0.5 +GUARD_TIMEOUT_SECONDS = 5 + +HEARTBEAT_LINE = b'{"type":"heartbeat","payload":{}}\n' +EVENT_LINE = json.dumps({ + 'type': 'event', + 'payload': { + 'specversion': '1.0', + 'id': '0', + 'time': '2026-10-03T12:00:00.000000000Z', + 'source': 'https://www.eventsourcingdb.io', + 'subject': '/books/42', + 'type': 'io.eventsourcingdb.library.book-acquired', + 'datacontenttype': 'application/json', + 'data': {'title': 'Dune'}, + 'predecessorhash': '0' * 64, + 'hash': '1' * 64, + 'traceparent': None, + 'tracestate': None, + 'signature': None, + }, +}).encode('utf-8') + b'\n' +ROW_LINE = b'{"type":"row","payload":{"title":"Dune"}}\n' + +StartStream = Callable[[Client], AsyncGenerator[Any]] + + +def observe_events(client: Client) -> AsyncGenerator[Any]: + return client.observe_events('/', ObserveEventsOptions(recursive=True)) + + +def run_eventql_query(client: Client) -> AsyncGenerator[Any]: + return client.run_eventql_query('FROM e IN events PROJECT INTO e') + + +STREAMS = [ + pytest.param(observe_events, id='observe_events'), + pytest.param(run_eventql_query, id='run_eventql_query'), +] +STREAMS_WITH_ITEMS = [ + pytest.param(observe_events, EVENT_LINE, id='observe_events'), + pytest.param(run_eventql_query, ROW_LINE, id='run_eventql_query'), +] + + +async def stall() -> None: + await asyncio.Event().wait() + + +@pytest.fixture +def short_heartbeat_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + heartbeat_timeout_module, + 'HEARTBEAT_TIMEOUT_SECONDS', + SHORT_HEARTBEAT_TIMEOUT_SECONDS, + ) + + +@pytest.mark.usefixtures('short_heartbeat_timeout') +class TestHeartbeatTimeout: + @staticmethod + @pytest.mark.asyncio + @pytest.mark.parametrize('start_stream', STREAMS) + async def test_ends_stream_with_error_if_neither_event_nor_heartbeat_arrives( + start_stream: StartStream, + ) -> None: + async def write_stream(response: web.StreamResponse) -> None: + await response.write(HEARTBEAT_LINE) + await stall() + + async with StreamServer(write_stream) as server: + started_at = time.monotonic() + + with pytest.raises( + HeartbeatTimeoutError, + match='No event and no heartbeat arrived for 30 seconds', + ): + async with asyncio.timeout(GUARD_TIMEOUT_SECONDS): + async for _ in start_stream(server.get_client()): + pass + + elapsed_seconds = time.monotonic() - started_at + assert elapsed_seconds >= SHORT_HEARTBEAT_TIMEOUT_SECONDS + assert elapsed_seconds < SHORT_HEARTBEAT_TIMEOUT_SECONDS + 1 + + async with asyncio.timeout(GUARD_TIMEOUT_SECONDS): + await server.client_disconnected.wait() + + @staticmethod + @pytest.mark.asyncio + @pytest.mark.parametrize(('start_stream', 'item_line'), STREAMS_WITH_ITEMS) + async def test_keeps_stream_open_while_heartbeats_arrive( + start_stream: StartStream, item_line: bytes + ) -> None: + heartbeat_interval_seconds = 0.1 + heartbeat_count = 15 + + async def write_stream(response: web.StreamResponse) -> None: + for _ in range(heartbeat_count): + await response.write(HEARTBEAT_LINE) + await asyncio.sleep(heartbeat_interval_seconds) + await response.write(item_line) + await stall() + + async with StreamServer(write_stream) as server: + started_at = time.monotonic() + items = [] + + async with asyncio.timeout(GUARD_TIMEOUT_SECONDS): + async for item in start_stream(server.get_client()): + items.append(item) + break + + elapsed_seconds = time.monotonic() - started_at + assert len(items) == 1 + assert elapsed_seconds > SHORT_HEARTBEAT_TIMEOUT_SECONDS + + @staticmethod + @pytest.mark.asyncio + @pytest.mark.parametrize(('start_stream', 'item_line'), STREAMS_WITH_ITEMS) + async def test_delivers_items_that_arrive_within_the_timeout( + start_stream: StartStream, item_line: bytes + ) -> None: + item_interval_seconds = 0.2 + item_count = 4 + + async def write_stream(response: web.StreamResponse) -> None: + for _ in range(item_count): + await response.write(item_line) + await asyncio.sleep(item_interval_seconds) + + async with StreamServer(write_stream) as server: + items = [] + + async with asyncio.timeout(GUARD_TIMEOUT_SECONDS): + async for item in start_stream(server.get_client()): + items.append(item) + + assert len(items) == item_count + + @staticmethod + @pytest.mark.asyncio + @pytest.mark.parametrize('start_stream', STREAMS) + async def test_cancelling_ends_stream_without_heartbeat_timeout( + start_stream: StartStream, + ) -> None: + async def write_stream(response: web.StreamResponse) -> None: + await response.write(HEARTBEAT_LINE) + await stall() + + async with StreamServer(write_stream) as server: + async def consume_stream() -> None: + async for _ in start_stream(server.get_client()): + pass + + task = asyncio.create_task(consume_stream()) + await asyncio.sleep(SHORT_HEARTBEAT_TIMEOUT_SECONDS / 2) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert task.cancelled(), "Task should be marked as cancelled" + + @staticmethod + @pytest.mark.asyncio + @pytest.mark.parametrize('start_stream', STREAMS) + async def test_does_not_report_timeout_of_caller_as_heartbeat_timeout( + start_stream: StartStream, + ) -> None: + async def write_stream(response: web.StreamResponse) -> None: + await response.write(HEARTBEAT_LINE) + await stall() + + async with StreamServer(write_stream) as server: + with pytest.raises(TimeoutError): + async with asyncio.timeout(SHORT_HEARTBEAT_TIMEOUT_SECONDS / 2): + async for _ in start_stream(server.get_client()): + pass + + @staticmethod + @pytest.mark.asyncio + async def test_passes_on_timeouts_that_are_not_heartbeat_timeouts() -> None: + class TimingOutStream: + @staticmethod + async def readline() -> bytes: + raise TimeoutError('Session timed out') + + with pytest.raises(TimeoutError, match='Session timed out'): + async for _ in read_lines_with_heartbeat_timeout(TimingOutStream()): # type: ignore + pass