diff --git a/.gitignore b/.gitignore index 2790fa6..63fd416 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ tests/* !tests/test_positions.py !tests/test_rh_lighter.py !tests/test_mempool.py +!tests/test_websocket_sends.py !tests/_mock_api.py !tests/test_timestamps.py !tests/test_spec_fixes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 475be9c..4bc0d2b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,15 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm `ws_endpoint`, `None` means the default endpoint, `wss://api.0xarchive.io/ws`. For `plans`, `None` means every plan. +### Fixed + +- `OxArchiveWs.disconnect()` now sends the subscribe and unsubscribe calls + made just before it (waiting up to 2 seconds), so an unsubscribe followed + by `disconnect()` reaches the server. A background send that finds the + connection already closed is dropped quietly instead of logging + `Task exception was never retrieved`; any other send failure goes to the + `on_error` handler. + ## [1.12.0] - 2026-10-05 Builds on 1.11.0. Two venues, Hyperliquid and Lighter; Lighter now has two diff --git a/oxarchive/websocket.py b/oxarchive/websocket.py index e3ce59c..f8fc96e 100644 --- a/oxarchive/websocket.py +++ b/oxarchive/websocket.py @@ -59,7 +59,7 @@ import warnings from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Callable, Literal, Mapping, Optional, Set, Union +from typing import Any, Callable, Coroutine, Literal, Mapping, Optional, Set, Union from urllib.parse import urlencode try: @@ -828,6 +828,10 @@ def __init__(self, options: WsOptions): self._running = False self._ping_task: Optional[asyncio.Task] = None self._receive_task: Optional[asyncio.Task] = None + # Subscribe and unsubscribe frames sent in the background by the + # non-async methods; held here because asyncio keeps only weak + # references to tasks, and flushed by disconnect(). + self._pending_sends: Set[asyncio.Task[None]] = set() # Event handlers self._on_message: Optional[MessageHandler] = None @@ -920,7 +924,14 @@ def _connection_url(self) -> str: return f"{self.options.ws_url}{separator}{query}" async def disconnect(self) -> None: - """Disconnect from the WebSocket server.""" + """Disconnect from the WebSocket server. + + Subscribe and unsubscribe calls made just before are sent first (for + up to 2 seconds), so an unsubscribe followed by ``disconnect()`` reaches + the server. + """ + if self._pending_sends: + await asyncio.wait(set(self._pending_sends), timeout=2) self._running = False self._set_state("disconnected") @@ -969,7 +980,7 @@ def subscribe( self._remember_subscription(channel, coin, interval_ms) if self.is_connected: - asyncio.create_task(self._send_subscribe(channel, coin, interval_ms)) + self._send_in_background(self._send_subscribe(channel, coin, interval_ms)) async def subscribe_async( self, @@ -1012,7 +1023,7 @@ def unsubscribe(self, channel: WsChannel, coin: Optional[str] = None) -> None: self._subscription_options.pop(key, None) if self.is_connected: - asyncio.create_task(self._send_unsubscribe(channel, coin)) + self._send_in_background(self._send_unsubscribe(channel, coin)) async def unsubscribe_async(self, channel: WsChannel, coin: Optional[str] = None) -> None: """Unsubscribe from a channel (async version).""" @@ -1959,6 +1970,28 @@ def _remember_subscription( else: self._subscription_options.pop(key, None) + def _send_in_background(self, send: Coroutine[Any, Any, None]) -> None: + """Run a subscribe or unsubscribe send without blocking the caller.""" + task = asyncio.create_task(send) + self._pending_sends.add(task) + task.add_done_callback(self._send_finished) + + def _send_finished(self, task: asyncio.Task[None]) -> None: + """Retrieve a background send's outcome. A send that lost the + connection is dropped: the server forgets a closed connection's + subscriptions, and a reconnect resubscribes from the remembered set. + Any other failure goes to the error handler.""" + self._pending_sends.discard(task) + if task.cancelled(): + return + exc = task.exception() + if not isinstance(exc, Exception) or isinstance(exc, ConnectionClosed): + return + if self._on_error: + self._on_error(exc) + else: + logger.warning("WebSocket send failed: %s", exc) + async def _send_subscribe( self, channel: WsChannel, diff --git a/tests/test_websocket_sends.py b/tests/test_websocket_sends.py new file mode 100644 index 0000000..48fae76 --- /dev/null +++ b/tests/test_websocket_sends.py @@ -0,0 +1,92 @@ +"""Background subscribe and unsubscribe sends: flushed by disconnect(), and +quiet when the connection is already gone.""" + +import asyncio +import json + +import pytest + +pytest.importorskip("websockets") + +from websockets.exceptions import ConnectionClosedError # noqa: E402 +from websockets.protocol import State # noqa: E402 + +from oxarchive.websocket import OxArchiveWs, WsOptions # noqa: E402 + + +class FakeConnection: + def __init__(self, fail_with=None): + self.state = State.OPEN + self.sent = [] + self.closed = False + self.fail_with = fail_with + + async def send(self, text): + await asyncio.sleep(0.01) + if self.fail_with is not None: + raise self.fail_with + if self.closed: + raise ConnectionClosedError(None, None) + self.sent.append(json.loads(text)) + + async def close(self, code=1000, reason=""): + self.closed = True + self.state = State.CLOSED + + +def client(conn): + ws = OxArchiveWs(WsOptions(api_key="0xa_" + "0" * 64)) + ws._ws = conn + return ws + + +def unretrieved(loop_errors): + return [e for e in loop_errors if "exception was never retrieved" in str(e.get("message", ""))] + + +@pytest.mark.asyncio +async def test_unsubscribe_then_disconnect_sends_the_unsubscribe_first(): + loop = asyncio.get_running_loop() + errors = [] + loop.set_exception_handler(lambda _loop, ctx: errors.append(ctx)) + conn = FakeConnection() + ws = client(conn) + ws.subscribe_mempool("BTC") + ws.unsubscribe_mempool("BTC") + await ws.disconnect() + assert [m["op"] for m in conn.sent] == ["subscribe", "unsubscribe"] + assert conn.closed + assert not ws._pending_sends + await asyncio.sleep(0.05) + assert not unretrieved(errors) + + +@pytest.mark.asyncio +async def test_a_send_that_finds_the_connection_closed_is_dropped_quietly(): + loop = asyncio.get_running_loop() + errors = [] + loop.set_exception_handler(lambda _loop, ctx: errors.append(ctx)) + ws = client(FakeConnection(fail_with=ConnectionClosedError(None, None))) + seen = [] + ws.on_error(seen.append) + ws.unsubscribe("trades", "BTC") + await asyncio.sleep(0.05) + import gc + + gc.collect() + await asyncio.sleep(0) + assert seen == [] + assert not ws._pending_sends + assert not unretrieved(errors) + + +@pytest.mark.asyncio +async def test_any_other_send_failure_reaches_the_error_handler(): + boom = RuntimeError("boom") + ws = client(FakeConnection(fail_with=boom)) + seen = [] + ws.on_error(seen.append) + ws.subscribe("trades", "BTC") + await asyncio.sleep(0.05) + assert seen == [boom] + assert not ws._pending_sends