From 4fc0e35a4d9d7df7402cc6ccecffcab4b1d9089c Mon Sep 17 00:00:00 2001 From: Chris Heaney Date: Mon, 16 Dec 2024 18:51:30 -0500 Subject: [PATCH] handle unsub for oracles --- src/driftpy/accounts/polling/drift_client.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/driftpy/accounts/polling/drift_client.py b/src/driftpy/accounts/polling/drift_client.py index b27275bb..9199a1b4 100644 --- a/src/driftpy/accounts/polling/drift_client.py +++ b/src/driftpy/accounts/polling/drift_client.py @@ -39,6 +39,7 @@ def __init__( self.program = program self.is_subscribed = False self.callbacks: dict[str, int] = {} + self.oracle_callbacks: dict[str, int] = {} self.perp_market_indexes = perp_market_indexes self.spot_market_indexes = spot_market_indexes @@ -148,7 +149,7 @@ async def add_oracle(self, oracle: Pubkey, oracle_source: OracleSource): callback_id = self.bulk_account_loader.add_account( oracle, self._get_oracle_callback(oracle_id, oracle_source) ) - self.callbacks[oracle_id] = callback_id + self.oracle_callbacks[oracle_id] = callback_id await self._wait_for_oracle(3, oracle_id) @@ -181,7 +182,14 @@ async def unsubscribe(self): self.bulk_account_loader.remove_account( Pubkey.from_string(pubkey_str), callback_id ) + + for oracle_id, callback_id in self.oracle_callbacks.items(): + self.bulk_account_loader.remove_account( + Pubkey.from_string(oracle_id.split("-")[0]), callback_id + ) + self.callbacks.clear() + self.oracle_callbacks.clear() def get_state_account_and_slot(self) -> Optional[DataAndSlot[StateAccount]]: return self.state