diff --git a/README.md b/README.md index 89ac9f7a..2af8c1dc 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,11 @@ async def main(): asyncio.run(main()) ``` +`*_unsubscribe()` takes the server-assigned subscription ID (`first_resp[0].result`, as above). +As a convenience it also accepts the request ID returned by the matching `*_subscribe()` helper, +which is translated once the subscription confirmation has been received. Server-assigned IDs are +resolved first, so prefer that form when you have it. + ## 🔨 Development ### Setup diff --git a/src/solana/rpc/websocket_api.py b/src/solana/rpc/websocket_api.py index ff51a0b3..750d5f7b 100644 --- a/src/solana/rpc/websocket_api.py +++ b/src/solana/rpc/websocket_api.py @@ -3,8 +3,9 @@ from __future__ import annotations import itertools +import logging +from collections.abc import Callable, Sequence from contextlib import asynccontextmanager -from collections.abc import Sequence from typing import Any, AsyncIterator, Dict, List, Optional, Union, cast from solders.account_decoder import UiDataSliceConfig @@ -43,7 +44,7 @@ VoteUnsubscribe, batch_to_json, ) -from solders.rpc.responses import Notification +from solders.rpc.responses import Notification, SignatureNotification from solders.rpc.responses import SubscriptionError as SoldersSubscriptionError from solders.rpc.responses import SubscriptionResult, parse_websocket_message from solders.signature import Signature @@ -62,6 +63,8 @@ _TX_ENCODING_TO_SOLDERS, ) +logger = logging.getLogger(__name__) + class SubscriptionError(Exception): """Raise when subscribing to an RPC feed fails.""" @@ -89,6 +92,7 @@ def __init__(self, *args, **kwargs): self.subscriptions: Dict[int, Body] = {} self.sent_subscriptions: Dict[int, Body] = {} self.failed_subscriptions = {} + self.request_ids_to_subscriptions: Dict[int, int] = {} self.request_counter = itertools.count() def increment_counter_and_get_id(self) -> int: @@ -104,13 +108,17 @@ async def send_request(self, message: Union[Body, List[Body]]) -> None: message: The request(s) to send. """ if isinstance(message, list): + reqs: List[Body] = message to_send = batch_to_json(message) - for req in message: - self.sent_subscriptions[req.id] = req else: + reqs = [message] to_send = message.to_json() - self.sent_subscriptions[message.id] = message + for req in reqs: + self.sent_subscriptions[req.id] = req await self.send(to_send) + for req in reqs: + if hasattr(req, "subscription_id"): + self._forget_subscription(req.subscription_id) async def recv( # type: ignore self, @@ -154,12 +162,10 @@ async def account_unsubscribe( """Unsubscribe from account notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = AccountUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(AccountUnsubscribe, subscription) async def logs_subscribe( self, @@ -186,12 +192,10 @@ async def logs_unsubscribe( """Unsubscribe from transaction logging. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = LogsUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(LogsUnsubscribe, subscription) async def block_subscribe( self, @@ -233,12 +237,10 @@ async def block_unsubscribe( """Unsubscribe from blocks. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = BlockUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(BlockUnsubscribe, subscription) async def program_subscribe( # pylint: disable=too-many-arguments self, @@ -290,12 +292,10 @@ async def program_unsubscribe( """Unsubscribe from program account notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = ProgramUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(ProgramUnsubscribe, subscription) async def signature_subscribe( self, @@ -322,12 +322,10 @@ async def signature_unsubscribe( """Unsubscribe from signature notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = SignatureUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(SignatureUnsubscribe, subscription) async def slot_subscribe(self) -> int: """Subscribe to receive notification anytime a slot is processed by the validator.""" @@ -343,12 +341,10 @@ async def slot_unsubscribe( """Unsubscribe from slot notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = SlotUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(SlotUnsubscribe, subscription) async def slots_updates_subscribe(self) -> int: """Subscribe to receive a notification from the validator on a variety of updates on every slot.""" @@ -364,12 +360,10 @@ async def slots_updates_unsubscribe( """Unsubscribe from slot update notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = SlotsUpdatesUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(SlotsUpdatesUnsubscribe, subscription) async def root_subscribe(self) -> int: """Subscribe to receive notification anytime a new root is set by the validator.""" @@ -385,12 +379,10 @@ async def root_unsubscribe( """Unsubscribe from root notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = RootUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(RootUnsubscribe, subscription) async def vote_subscribe(self) -> int: """Subscribe to receive notification anytime a new vote is observed in gossip.""" @@ -406,12 +398,10 @@ async def vote_unsubscribe( """Unsubscribe from vote notifications. Args: - subscription: ID of subscription to cancel. + subscription: Request ID from the matching ``subscribe`` helper, or the + server-assigned ID from the confirmation (``recv()[0].result``). """ - req_id = self.increment_counter_and_get_id() - req = VoteUnsubscribe(subscription, req_id) - await self.send_request(req) - del self.subscriptions[subscription] + await self._unsubscribe(VoteUnsubscribe, subscription) def _process_rpc_response(self, raw: str) -> List[Union[Notification, SubscriptionResult]]: parsed = parse_websocket_message(raw) @@ -422,8 +412,47 @@ def _process_rpc_response(self, raw: str) -> List[Union[Notification, Subscripti raise SubscriptionError(item, subscription) if isinstance(item, SubscriptionResult): self.subscriptions[item.result] = self.sent_subscriptions[item.id] + self.request_ids_to_subscriptions[item.id] = item.result + elif isinstance(item, SignatureNotification): + # Signature subscriptions expire server-side after the notification. + self._forget_subscription(item.subscription) return cast(List[Union[Notification, SubscriptionResult]], parsed) + async def _unsubscribe(self, constructor: Callable[[int, int], Body], subscription: int) -> None: + """Send an unsubscribe request after resolving the server-assigned subscription ID.""" + server_subscription = self._resolve_server_subscription(subscription) + req_id = self.increment_counter_and_get_id() + req = constructor(server_subscription, req_id) + await self.send_request(req) + self._forget_subscription(server_subscription) + + def _resolve_server_subscription(self, subscription: int) -> int: + """Translate a subscribe request ID into its server-assigned subscription ID. + + Server-assigned IDs already recorded in ``self.subscriptions`` win, so the README + form (``recv()[0].result``) is never rewritten through the request-ID map. An + unconfirmed request ID is sent as-is. + """ + if subscription in self.subscriptions: + return subscription + mapped = self.request_ids_to_subscriptions.get(subscription) + if mapped is not None: + return mapped + logger.warning( + "Unsubscribe target %s is not a known server subscription ID or confirmed request ID", + subscription, + ) + return subscription + + def _forget_subscription(self, server_subscription: int) -> None: + """Drop local bookkeeping for a server-assigned subscription ID.""" + self.subscriptions.pop(server_subscription, None) + self.request_ids_to_subscriptions = { + req_id: sub_id + for req_id, sub_id in self.request_ids_to_subscriptions.items() + if sub_id != server_subscription + } + @asynccontextmanager async def connect(uri: str = "ws://localhost:8900", **kwargs: Any) -> AsyncIterator[SolanaWsClientProtocol]: diff --git a/tests/unit/test_websocket_api.py b/tests/unit/test_websocket_api.py index 4d9532fc..89ec843f 100644 --- a/tests/unit/test_websocket_api.py +++ b/tests/unit/test_websocket_api.py @@ -3,18 +3,31 @@ from __future__ import annotations import itertools +import logging +import pytest from solders.account_decoder import UiAccountEncoding, UiDataSliceConfig from solders.commitment_config import CommitmentLevel from solders.pubkey import Pubkey from solders.rpc.config import RpcAccountInfoConfig, RpcProgramAccountsConfig from solders.rpc.filter import Memcmp +from solders.rpc.requests import LogsUnsubscribe from solana.rpc.commitment import Processed from solana.rpc.models import DataSliceOpts, MemcmpOpts from solana.rpc.websocket_api import SolanaWsClientProtocol, connect +def _ws_protocol() -> SolanaWsClientProtocol: + protocol = SolanaWsClientProtocol.__new__(SolanaWsClientProtocol) + protocol.subscriptions = {} + protocol.sent_subscriptions = {} + protocol.failed_subscriptions = {} + protocol.request_ids_to_subscriptions = {} + protocol.request_counter = itertools.count() + return protocol + + async def test_account_subscribe_returns_request_id(monkeypatch): """Subscription helpers should return the request id used for the request.""" protocol = SolanaWsClientProtocol.__new__(SolanaWsClientProtocol) @@ -82,6 +95,153 @@ async def fake_send_request(self, message): assert sent_messages[0].config == expected_config +async def test_unsubscribe_sends_server_assigned_subscription_id(monkeypatch): + """Unsubscribe helpers should translate a subscribe() request ID into the server-assigned subscription ID.""" + protocol = _ws_protocol() + sent_messages = [] + + async def fake_send_request(self, message): + sent_messages.append(message) + + monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request) + + request_id = await protocol.slot_subscribe() + protocol.sent_subscriptions[request_id] = sent_messages[0] + server_subscription = 9710270 + confirmation = f'{{"jsonrpc":"2.0","result":{server_subscription},"id":{request_id}}}' + protocol._process_rpc_response(confirmation) + + await protocol.slot_unsubscribe(request_id) + + assert f'"params":[{server_subscription}]' in sent_messages[1].to_json() + assert server_subscription not in protocol.subscriptions + assert request_id not in protocol.request_ids_to_subscriptions + + +async def test_unsubscribe_before_confirmation_sends_given_id(monkeypatch, caplog): + """Unsubscribing an unconfirmed subscription should pass the given ID through instead of raising.""" + protocol = _ws_protocol() + sent_messages = [] + + async def fake_send_request(self, message): + sent_messages.append(message) + + monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request) + + with caplog.at_level(logging.WARNING, logger="solana.rpc.websocket_api"): + await protocol.logs_unsubscribe(7) + + assert '"params":[7]' in sent_messages[0].to_json() + assert "7" in caplog.text + + +async def test_unsubscribe_server_id_wins_when_it_collides_with_request_id(monkeypatch): + """A server-assigned ID must not be rewritten through the request-ID map.""" + protocol = _ws_protocol() + sent_messages = [] + + async def fake_send_request(self, message): + sent_messages.append(message) + + monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request) + + slot_req = await protocol.slot_subscribe() + protocol.sent_subscriptions[slot_req] = sent_messages[-1] + protocol._process_rpc_response('{"jsonrpc":"2.0","result":1,"id":1}') + await protocol.slot_unsubscribe(1) + + root_req = await protocol.root_subscribe() + protocol.sent_subscriptions[root_req] = sent_messages[-1] + protocol._process_rpc_response('{"jsonrpc":"2.0","result":2,"id":3}') + + vote_req = await protocol.vote_subscribe() + protocol.sent_subscriptions[vote_req] = sent_messages[-1] + protocol._process_rpc_response('{"jsonrpc":"2.0","result":3,"id":4}') + + await protocol.vote_unsubscribe(3) + + assert '"method":"voteUnsubscribe"' in sent_messages[-1].to_json() + assert '"params":[3]' in sent_messages[-1].to_json() + assert 3 not in protocol.subscriptions + assert 2 in protocol.subscriptions + assert protocol.request_ids_to_subscriptions == {3: 2} + + +async def test_unsubscribe_by_server_id_drops_request_id_mapping(monkeypatch): + """Unsubscribing with the README server-ID form should still drop the reverse map.""" + protocol = _ws_protocol() + sent_messages = [] + + async def fake_send_request(self, message): + sent_messages.append(message) + + monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request) + + request_id = await protocol.slot_subscribe() + protocol.sent_subscriptions[request_id] = sent_messages[0] + protocol._process_rpc_response('{"jsonrpc":"2.0","result":99,"id":1}') + + await protocol.slot_unsubscribe(99) + + assert '"params":[99]' in sent_messages[-1].to_json() + assert protocol.subscriptions == {} + assert protocol.request_ids_to_subscriptions == {} + + +async def test_unsubscribe_keeps_bookkeeping_if_send_fails(monkeypatch): + """Local state must stay intact when the unsubscribe request never leaves the client.""" + protocol = _ws_protocol() + sent_messages = [] + + async def fake_send_request(self, message): + sent_messages.append(message) + if len(sent_messages) > 1: + raise ConnectionError("closed") + + monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request) + + request_id = await protocol.slot_subscribe() + protocol.sent_subscriptions[request_id] = sent_messages[0] + protocol._process_rpc_response('{"jsonrpc":"2.0","result":99,"id":1}') + + with pytest.raises(ConnectionError): + await protocol.slot_unsubscribe(request_id) + + assert 99 in protocol.subscriptions + assert protocol.request_ids_to_subscriptions == {1: 99} + + +async def test_raw_unsubscribe_send_request_forgets_mapping(): + """Raw batched unsubscribes should drop the request-ID map, not only the helper path.""" + protocol = _ws_protocol() + protocol.subscriptions[9] = object() + protocol.request_ids_to_subscriptions[1] = 9 + sent = [] + + async def fake_send(data): + sent.append(data) + + protocol.send = fake_send # type: ignore[method-assign] + await protocol.send_request(LogsUnsubscribe(9, 2)) + + assert sent + assert 9 not in protocol.subscriptions + assert protocol.request_ids_to_subscriptions == {} + + +async def test_signature_notification_forgets_expired_subscription(): + """Signature subscriptions expire server-side after the notification.""" + protocol = _ws_protocol() + protocol.subscriptions[42] = object() + protocol.request_ids_to_subscriptions[1] = 42 + protocol._process_rpc_response( + '{"jsonrpc":"2.0","method":"signatureNotification",' + '"params":{"result":{"context":{"slot":1},"value":{"err":null}},"subscription":42}}' + ) + assert 42 not in protocol.subscriptions + assert protocol.request_ids_to_subscriptions == {} + + async def test_connect_preserves_async_with_and_custom_connection(monkeypatch): """The connect helper should stay usable as an async context manager.""" captured = {}