Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 36 additions & 27 deletions src/solana/rpc/websocket_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,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:
Expand Down Expand Up @@ -154,12 +155,11 @@ async def account_unsubscribe(
"""Unsubscribe from account notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These docstrings (all nine) now document the request-ID form, but README.md:94 and the docs page mirroring it still pass first_resp[0].result — the server-assigned ID — and so do all the integration fixtures. Since the two forms are silently interchangeable and can collide, the public contract ends up documented one way and exercised another. Worth updating README/docs to whichever form is authoritative.

"""
req_id = self.increment_counter_and_get_id()
req = AccountUnsubscribe(subscription, req_id)
req = AccountUnsubscribe(self._pop_server_subscription(subscription), req_id)

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Local bookkeeping is now torn down before the request is sent.

_pop_server_subscription(...) is evaluated while building the request object, so the mapping and the self.subscriptions entry are gone before await self.send_request(req) runs. Previously the del ran only after a successful send.

If send_request raises (e.g. ConnectionClosed), a reconnect-and-resubscribe routine reading self.subscriptions sees no record of the subscription, and a retried slot_unsubscribe(req_id) now falls through to the pass-through path and sends the raw request ID as the subscription ID. Please resolve the ID, send, then clean up on success.

Same pattern at lines 192, 238, 294, 325, 345, 365, 385 and 405.

await self.send_request(req)
del self.subscriptions[subscription]

async def logs_subscribe(
self,
Expand All @@ -186,12 +186,11 @@ async def logs_unsubscribe(
"""Unsubscribe from transaction logging.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = LogsUnsubscribe(subscription, req_id)
req = LogsUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def block_subscribe(
self,
Expand Down Expand Up @@ -233,12 +232,11 @@ async def block_unsubscribe(
"""Unsubscribe from blocks.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = BlockUnsubscribe(subscription, req_id)
req = BlockUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def program_subscribe( # pylint: disable=too-many-arguments
self,
Expand Down Expand Up @@ -290,12 +288,11 @@ async def program_unsubscribe(
"""Unsubscribe from program account notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = ProgramUnsubscribe(subscription, req_id)
req = ProgramUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def signature_subscribe(
self,
Expand All @@ -322,12 +319,11 @@ async def signature_unsubscribe(
"""Unsubscribe from signature notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = SignatureUnsubscribe(subscription, req_id)
req = SignatureUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def slot_subscribe(self) -> int:
"""Subscribe to receive notification anytime a slot is processed by the validator."""
Expand All @@ -343,12 +339,11 @@ async def slot_unsubscribe(
"""Unsubscribe from slot notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = SlotUnsubscribe(subscription, req_id)
req = SlotUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def slots_updates_subscribe(self) -> int:
"""Subscribe to receive a notification from the validator on a variety of updates on every slot."""
Expand All @@ -364,12 +359,11 @@ async def slots_updates_unsubscribe(
"""Unsubscribe from slot update notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = SlotsUpdatesUnsubscribe(subscription, req_id)
req = SlotsUpdatesUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def root_subscribe(self) -> int:
"""Subscribe to receive notification anytime a new root is set by the validator."""
Expand All @@ -385,12 +379,11 @@ async def root_unsubscribe(
"""Unsubscribe from root notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = RootUnsubscribe(subscription, req_id)
req = RootUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

async def vote_subscribe(self) -> int:
"""Subscribe to receive notification anytime a new vote is observed in gossip."""
Expand All @@ -406,12 +399,11 @@ async def vote_unsubscribe(
"""Unsubscribe from vote notifications.

Args:
subscription: ID of subscription to cancel.
subscription: The request ID returned by the corresponding ``subscribe`` method.
"""
req_id = self.increment_counter_and_get_id()
req = VoteUnsubscribe(subscription, req_id)
req = VoteUnsubscribe(self._pop_server_subscription(subscription), req_id)
await self.send_request(req)
del self.subscriptions[subscription]

def _process_rpc_response(self, raw: str) -> List[Union[Notification, SubscriptionResult]]:
parsed = parse_websocket_message(raw)
Expand All @@ -422,8 +414,25 @@ 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
Comment thread
michaelhly marked this conversation as resolved.
return cast(List[Union[Notification, SubscriptionResult]], parsed)

def _pop_server_subscription(self, subscription: int) -> int:
"""Translate a subscribe request ID into its server-assigned subscription ID.

The RPC server assigns subscription IDs independently of the request IDs returned by the
``subscribe`` helpers, so an unconfirmed request ID is sent as-is.

Args:
subscription: The value returned by a ``subscribe`` helper, or a server-assigned ID.

Returns:
The server-assigned subscription ID, if known, else the input unchanged.
"""
server_subscription = self.request_ids_to_subscriptions.pop(subscription, subscription)

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Request IDs and server subscription IDs collide here, cancelling the wrong subscription.

Both spaces are plain int and the request-ID map is consulted first, so the "pass through unchanged if unconfirmed" rule can't distinguish an unconfirmed request ID from a server-assigned ID that happens to equal a live request ID. Reproduced on a fresh protocol object (small server IDs, as a local validator assigns):

slot_subscribe()  -> req 1, server sub 1;  slot_unsubscribe(1) consumes req id 2
root_subscribe()  -> req 3, server sub 2      # map {3: 2}
vote_subscribe()  -> req 4, server sub 3      # map {3: 2, 4: 3}
vote_unsubscribe(3)   # 3 = server id of the vote sub, the README form
wire: {"method":"voteUnsubscribe","id":5,"params":[2]}   # cancels the ROOT subscription
remaining: {3: VoteSubscribe...}                          # vote sub leaks

The PR description says callers that track server-assigned IDs "behave exactly as before"; that only holds while no live request ID equals the server ID being passed. README.md:94/:105 and every integration fixture use the server-ID form, so this is a real backward-compat break rather than a theoretical one.

Suggestion: disambiguate instead of guessing. Checking self.subscriptions first keeps the older contract authoritative:

if subscription in self.subscriptions:          # already a server-assigned ID
    self.request_ids_to_subscriptions = {k: v for k, v in self.request_ids_to_subscriptions.items() if v != subscription}
    self.subscriptions.pop(subscription, None)
    return subscription

Or, better long term, expose the mapping through a distinct return type/API so the two ID spaces are never interchangeable ints.

self.subscriptions.pop(server_subscription, None)

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: swapping del for pop(..., None) removes the only local signal that the unsubscribe target isn't a live subscription. Unsubscribing before confirmation (or twice) now silently sends an invalid subscription ID; the server's -32602 comes back as a SoldersSubscriptionError and is raised as SubscriptionError out of an unrelated later recv(), while the real subscription keeps streaming. Dropping the KeyError is the right call, but a returned bool or a logging.warning would let the caller learn the unsubscribe was a no-op at the call site.

return server_subscription


@asynccontextmanager
async def connect(uri: str = "ws://localhost:8900", **kwargs: Any) -> AsyncIterator[SolanaWsClientProtocol]:
Expand Down
48 changes: 48 additions & 0 deletions tests/unit/test_websocket_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,54 @@ 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 = SolanaWsClientProtocol.__new__(SolanaWsClientProtocol)
protocol.subscriptions = {}
protocol.sent_subscriptions = {}
protocol.failed_subscriptions = {}
protocol.request_ids_to_subscriptions = {}
protocol.request_counter = itertools.count()
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):
"""Unsubscribing an unconfirmed subscription should pass the given ID through instead of raising."""
protocol = SolanaWsClientProtocol.__new__(SolanaWsClientProtocol)
protocol.subscriptions = {}
protocol.sent_subscriptions = {}
protocol.failed_subscriptions = {}
protocol.request_ids_to_subscriptions = {}
protocol.request_counter = itertools.count()
sent_messages = []

async def fake_send_request(self, message):
sent_messages.append(message)

monkeypatch.setattr(SolanaWsClientProtocol, "send_request", fake_send_request)

await protocol.logs_unsubscribe(7)

assert '"params":[7]' in sent_messages[0].to_json()


async def test_connect_preserves_async_with_and_custom_connection(monkeypatch):
"""The connect helper should stay usable as an async context manager."""
captured = {}
Expand Down
Loading