drift_rs/
websocket_account_subscriber.rs

1use std::{str::FromStr, sync::Arc};
2
3use drift_pubsub_client::PubsubClient;
4use futures_util::StreamExt;
5use log::warn;
6use solana_account_decoder_client_types::UiAccountEncoding;
7use solana_rpc_client::nonblocking::rpc_client::RpcClient;
8use solana_rpc_client_api::config::RpcAccountInfoConfig;
9use solana_sdk::{commitment_config::CommitmentConfig, pubkey::Pubkey};
10use tokio::sync::oneshot;
11
12use crate::{utils::get_http_url, SdkError, SdkResult, UnsubHandle};
13
14const LOG_TARGET: &str = "wsaccsub";
15
16#[derive(Clone, Debug)]
17pub struct AccountUpdate {
18    /// Address of the account
19    pub pubkey: Pubkey,
20    /// Owner of the account
21    pub owner: Pubkey,
22    pub lamports: u64,
23    /// Serialized account data (e.g. Anchor/Borsh)
24    pub data: Vec<u8>,
25    /// Slot retrieved
26    pub slot: u64,
27}
28
29#[derive(Clone)]
30pub struct WebsocketAccountSubscriber {
31    pubsub: Arc<PubsubClient>,
32    pub pubkey: Pubkey,
33    pub commitment: CommitmentConfig,
34}
35
36impl WebsocketAccountSubscriber {
37    pub fn new(pubsub: Arc<PubsubClient>, pubkey: Pubkey, commitment: CommitmentConfig) -> Self {
38        WebsocketAccountSubscriber {
39            pubsub,
40            pubkey,
41            commitment,
42        }
43    }
44
45    /// Start a Ws account subscription task
46    ///
47    /// * `subscription_name` - some user defined identifier for the subscription
48    /// * `sync` - true if subscription should fetch account data on start
49    /// * `on_update` - function to call on updates from the subscription
50    ///
51    /// Fetches the account to set the initial value, then uses event based updates
52    pub async fn subscribe<F>(
53        &self,
54        subscription_name: &'static str,
55        sync: bool,
56        on_update: F,
57    ) -> SdkResult<UnsubHandle>
58    where
59        F: 'static + Send + Fn(&AccountUpdate),
60    {
61        if sync {
62            // seed initial account state
63            log::debug!(target: LOG_TARGET, "seeding account: {subscription_name}-{:?}", self.pubkey);
64            let owner: Pubkey;
65            let rpc = RpcClient::new(get_http_url(self.pubsub.url().as_str())?);
66            match rpc
67                .get_account_with_commitment(&self.pubkey, self.commitment)
68                .await
69            {
70                Ok(response) => {
71                    if let Some(account) = response.value {
72                        owner = account.owner;
73                        on_update(&AccountUpdate {
74                            owner,
75                            lamports: account.lamports,
76                            pubkey: self.pubkey,
77                            data: account.data,
78                            slot: response.context.slot,
79                        });
80                    } else {
81                        warn!("seeding account failed: {response:?}");
82                        return Err(SdkError::InvalidAccount);
83                    }
84                }
85                Err(err) => {
86                    warn!("seeding account failed: {err:?}");
87                    return Err(err.into());
88                }
89            }
90            drop(rpc);
91        }
92
93        let (unsub_tx, mut unsub_rx) = oneshot::channel::<()>();
94        let account_config = RpcAccountInfoConfig {
95            commitment: Some(self.commitment),
96            encoding: Some(UiAccountEncoding::Base64Zstd),
97            ..RpcAccountInfoConfig::default()
98        };
99        let pubkey = self.pubkey;
100        let pubsub = Arc::clone(&self.pubsub);
101
102        tokio::spawn(async move {
103            loop {
104                log::debug!(target: LOG_TARGET, "spawn account subscriber: {subscription_name}-{:?}", pubkey);
105                let (mut account_updates, account_unsubscribe) = match pubsub
106                    .account_subscribe(&pubkey, Some(account_config.clone()))
107                    .await
108                {
109                    Ok(res) => res,
110                    Err(err) => {
111                        log::error!(target: LOG_TARGET, "account subscribe {pubkey} failed: {err:?}");
112                        continue;
113                    }
114                };
115                log::debug!(target: LOG_TARGET, "account subscribed: {subscription_name}-{pubkey:?}");
116                let mut latest_slot = 0;
117                let res = loop {
118                    tokio::select! {
119                        biased;
120                        message = account_updates.next() => {
121                            match message {
122                                Some(message) => {
123                                    let slot = message.context.slot;
124                                    if slot >= latest_slot {
125                                        latest_slot = slot;
126                                        let data = message.value.data.decode().expect("decoded");
127                                        let account_update = AccountUpdate {
128                                            owner: Pubkey::from_str(&message.value.owner).unwrap(),
129                                            lamports: message.value.lamports,
130                                            pubkey,
131                                            data,
132                                            slot,
133                                        };
134                                        on_update(&account_update);
135                                    }
136                                }
137                                None => {
138                                    log::error!(target: LOG_TARGET, "{subscription_name}: Ws ended unexpectedly: {pubkey:?}");
139                                    break Err(());
140                                }
141                            }
142                        }
143                        _ = &mut unsub_rx => {
144                            log::debug!(target: LOG_TARGET, "{subscription_name}: Unsubscribing from account stream: {pubkey:?}");
145                            account_unsubscribe().await;
146                            break Ok(());
147                        }
148                    }
149                };
150
151                if res.is_ok() {
152                    break;
153                }
154            }
155        });
156
157        Ok(unsub_tx)
158    }
159}
drift_rs/websocket_account_subscriber.rs

drift_rs/
websocket_account_subscriber.rs
