Skip to main content

carbon_rpc_program_subscribe_datasource/
lib.rs

1use {
2    async_trait::async_trait,
3    carbon_core::{
4        datasource::{AccountUpdate, Datasource, DatasourceId, Update, UpdateType},
5        error::CarbonResult,
6        metrics::{Counter, Histogram, MetricsRegistry},
7    },
8    futures::StreamExt,
9    solana_account::Account,
10    solana_client::{
11        nonblocking::pubsub_client::PubsubClient, rpc_config::RpcProgramAccountsConfig,
12    },
13    solana_pubkey::Pubkey,
14    std::{str::FromStr, sync::LazyLock, time::Duration},
15    tokio::sync::mpsc::Sender,
16    tokio_util::sync::CancellationToken,
17};
18
19const MAX_RECONNECTION_ATTEMPTS: u32 = 10;
20const RECONNECTION_DELAY_MS: u64 = 3000;
21
22static ACCOUNT_PROCESS_TIME_NANOS: LazyLock<Histogram> = LazyLock::new(|| {
23    Histogram::new(
24        "program_subscribe_account_process_time_nanoseconds",
25        "Time to process account in nanoseconds",
26        vec![
27            1_000.0,
28            10_000.0,
29            100_000.0,
30            1_000_000.0,
31            10_000_000.0,
32            100_000_000.0,
33            1_000_000_000.0,
34        ],
35    )
36});
37static ACCOUNTS_PROCESSED: Counter = Counter::new(
38    "program_subscribe_accounts_processed_total",
39    "Accounts processed by program subscribe",
40);
41
42fn register_program_subscribe_metrics() {
43    let registry = MetricsRegistry::global();
44    registry.register_counter(&ACCOUNTS_PROCESSED);
45    registry.register_histogram(&ACCOUNT_PROCESS_TIME_NANOS);
46}
47
48#[derive(Debug, Clone)]
49pub struct Filters {
50    pub pubkey: Pubkey,
51    pub program_subscribe_config: Option<RpcProgramAccountsConfig>,
52}
53
54impl Filters {
55    pub const fn new(
56        pubkey: Pubkey,
57        program_subscribe_config: Option<RpcProgramAccountsConfig>,
58    ) -> Self {
59        Filters {
60            pubkey,
61            program_subscribe_config,
62        }
63    }
64}
65
66pub struct RpcProgramSubscribe {
67    pub rpc_ws_url: String,
68    pub filters: Filters,
69}
70
71impl RpcProgramSubscribe {
72    pub const fn new(rpc_ws_url: String, filters: Filters) -> Self {
73        Self {
74            rpc_ws_url,
75            filters,
76        }
77    }
78}
79
80#[async_trait]
81impl Datasource for RpcProgramSubscribe {
82    async fn consume(
83        &self,
84        id: DatasourceId,
85        sender: Sender<(Update, DatasourceId)>,
86        cancellation_token: CancellationToken,
87    ) -> CarbonResult<()> {
88        register_program_subscribe_metrics();
89
90        let mut reconnection_attempts = 0;
91
92        loop {
93            if cancellation_token.is_cancelled() {
94                log::info!("Cancellation requested, stopping reconnection attempts");
95                break;
96            }
97
98            let client = match PubsubClient::new(&self.rpc_ws_url).await {
99                Ok(client) => client,
100                Err(err) => {
101                    log::error!("Failed to create RPC subscribe client: {err}");
102                    reconnection_attempts += 1;
103                    if reconnection_attempts >= MAX_RECONNECTION_ATTEMPTS {
104                        return Err(carbon_core::error::Error::Custom(format!(
105                            "Failed to create RPC subscribe client after {MAX_RECONNECTION_ATTEMPTS} attempts: {err}"
106                        )));
107                    }
108                    tokio::time::sleep(Duration::from_millis(RECONNECTION_DELAY_MS)).await;
109                    continue;
110                }
111            };
112
113            let filters = self.filters.clone();
114            let sender_clone = sender.clone();
115            let id_for_loop = id.clone();
116
117            let (mut program_stream, _program_unsub) = match client
118                .program_subscribe(&filters.pubkey, filters.program_subscribe_config)
119                .await
120            {
121                Ok(subscription) => subscription,
122                Err(err) => {
123                    log::error!("Failed to subscribe to program updates: {err:?}");
124                    reconnection_attempts += 1;
125                    if reconnection_attempts > MAX_RECONNECTION_ATTEMPTS {
126                        return Err(carbon_core::error::Error::Custom(format!(
127                            "Failed to subscribe after {MAX_RECONNECTION_ATTEMPTS} attempts: {err}"
128                        )));
129                    }
130                    tokio::time::sleep(Duration::from_millis(RECONNECTION_DELAY_MS)).await;
131                    continue;
132                }
133            };
134
135            reconnection_attempts = 0;
136
137            loop {
138                tokio::select! {
139                    _ = cancellation_token.cancelled() => {
140                        log::info!("Cancellation requested, stopping subscription...");
141                        return Ok(());
142                    }
143                    event_result = program_stream.next() => {
144                        match event_result {
145                            Some(acc_event) => {
146                                let start_time = std::time::Instant::now();
147                                let decoded_account: Account = match acc_event.value.account.to_account() {
148                                    Some(account_data) => account_data,
149                                    None => {
150                                        log::error!("Error decoding account event");
151                                        continue;
152                                    }
153                                };
154
155                                let Ok(account_pubkey) = Pubkey::from_str(&acc_event.value.pubkey) else {
156                                    log::error!("Error parsing account pubkey. Value: {}", &acc_event.value.pubkey);
157                                    continue;
158                                };
159
160                                let update = AccountUpdate {
161                                    pubkey: account_pubkey,
162                                    account: decoded_account,
163                                    slot: acc_event.context.slot,
164                                    transaction_signature: None,
165                                }
166                                .into_update();
167
168                                ACCOUNT_PROCESS_TIME_NANOS.record(start_time.elapsed().as_nanos() as f64);
169                                ACCOUNTS_PROCESSED.inc();
170
171                                if let Err(err) = sender_clone.try_send((update, id_for_loop.clone())) {
172                                    log::error!("Error sending account update: {err:?}");
173                                    break;
174                                }
175                            }
176                            None => {
177                                log::warn!("Program accounts stream has been closed, attempting to reconnect...");
178                                break;
179                            }
180                        }
181                    }
182                }
183            }
184
185            tokio::time::sleep(Duration::from_millis(RECONNECTION_DELAY_MS)).await;
186        }
187
188        Ok(())
189    }
190
191    fn update_types(&self) -> Vec<UpdateType> {
192        vec![UpdateType::AccountUpdate, UpdateType::AccountDeletion]
193    }
194}