fix: use mutiple connections for ws subscriptions (#365)

This fix rewrites the way account subscriptions work in the cloning
pipeline, instead of using a single connection per shard, we now use
pool of connections, which are used to distribute subscriptions, giving
us more subscription capacity, before threshold imposed by RPC provider
is reached.

closes gh-362
<!-- greptile_comment -->

## Greptile Summary

This PR implements a WebSocket connection pool for account subscriptions
to increase capacity before hitting RPC provider limits. The changes
replace single-connection architecture with a distributed approach.

- Added `PubsubPool` in `remote_account_updates_shard.rs` to manage
multiple WebSocket connections (8 per RPC upstream)
- Modified test setup in `remote_account_updates.rs` to include 5-second
delay for connection establishment
- Added connection tracking and subscription distribution logic across
pool
- Uses unsafe `transmute` for stream lifetime extension which should be
reviewed carefully
- Implements proper cleanup of subscriptions and connections during
shutdown



<!-- /greptile_comment -->

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: bmuddha

magicblock-account-updates/src/remote_account_updates_shard.rs

@@ -6,12 +6,12 @@ use std::{
     sync::{Arc, RwLock},
 };
 
-use futures_util::StreamExt;
+use futures_util::{stream::FuturesUnordered, Stream, StreamExt};
 use log::*;
 use magicblock_metrics::metrics;
-use solana_account_decoder::{UiAccountEncoding, UiDataSliceConfig};
+use solana_account_decoder::{UiAccount, UiAccountEncoding, UiDataSliceConfig};
 use solana_pubsub_client::nonblocking::pubsub_client::PubsubClient;
-use solana_rpc_client_api::config::RpcAccountInfoConfig;
+use solana_rpc_client_api::{config::RpcAccountInfoConfig, response::Response};
 use solana_sdk::{
     clock::{Clock, Slot},
     commitment_config::{CommitmentConfig, CommitmentLevel},
@@ -23,6 +23,13 @@ use tokio::sync::mpsc::Receiver;
 use tokio_stream::StreamMap;
 use tokio_util::sync::CancellationToken;
 
+type BoxFn = Box<
+    dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send,
+>;
+
+type SubscriptionStream =
+    Pin<Box<dyn Stream<Item = Response<UiAccount>> + Send + 'static>>;
+
 #[derive(Debug, Error)]
 pub enum RemoteAccountUpdatesShardError {
     #[error(transparent)]
@@ -66,11 +73,9 @@ impl RemoteAccountUpdatesShard {
     ) -> Result<(), RemoteAccountUpdatesShardError> {
         // Create a pubsub client
         info!("Shard {}: Starting", self.shard_id);
-        let pubsub_client = PubsubClient::new(&self.url)
-            .await
-            .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
+        let ws_url = self.url.as_str();
         // For every account, we only want the updates, not the actual content of the accounts
-        let rpc_account_info_config = Some(RpcAccountInfoConfig {
+        let config = RpcAccountInfoConfig {
             commitment: self
                 .commitment
                 .map(|commitment| CommitmentConfig { commitment }),
@@ -80,21 +85,13 @@ impl RemoteAccountUpdatesShard {
                 length: 0,
             }),
             min_context_slot: None,
-        });
+        };
+        let mut pool = PubsubPool::new(ws_url, config).await?;
         // Subscribe to the clock from the RPC (to figure out the latest slot)
-        let (mut clock_stream, clock_unsubscribe) = pubsub_client
-            .account_subscribe(&clock::ID, rpc_account_info_config.clone())
-            .await
-            .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
+        let mut clock_stream = pool.subscribe(clock::ID).await?;
         let mut clock_slot = 0;
         // We'll store useful maps for each of the account subscriptions
         let mut account_streams = StreamMap::new();
-        // rust compiler is not yet smart enough to figure out the exact type
-        type BoxFn = Box<
-            dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>>
-                + Send,
-        >;
-        let mut account_unsubscribes: HashMap<Pubkey, BoxFn> = HashMap::new();
         const LOG_CLOCK_FREQ: u64 = 100;
         let mut log_clock_count = 0;
 
@@ -118,19 +115,17 @@ impl RemoteAccountUpdatesShard {
                     } else {
                         warn!("Shard {}: Received empty clock data", self.shard_id);
                     }
+                    self.try_to_override_last_known_update_slot(clock::ID, clock_slot);
                 }
                 // When we receive a message to start monitoring an account
                 Some((pubkey, unsub)) = self.monitoring_request_receiver.recv() => {
                     if unsub {
-                        let Some(request) = account_unsubscribes.remove(&pubkey) else {
-                            continue;
-                        };
                         account_streams.remove(&pubkey);
                         metrics::set_subscriptions_count(account_streams.len(), &self.shard_id);
-                        request().await;
+                        pool.unsubscribe(&pubkey).await;
                         continue;
                     }
-                    if account_unsubscribes.contains_key(&pubkey) {
+                    if pool.subscribed(&pubkey) {
                         continue;
                     }
                     debug!(
@@ -139,12 +134,10 @@ impl RemoteAccountUpdatesShard {
                         pubkey,
                         clock_slot
                     );
-                    let (stream, unsubscribe) = pubsub_client
-                        .account_subscribe(&pubkey, rpc_account_info_config.clone())
-                        .await
-                        .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
+                    let stream = pool
+                        .subscribe(pubkey)
+                        .await?;
                     account_streams.insert(pubkey, stream);
-                    account_unsubscribes.insert(pubkey, unsubscribe);
                     metrics::set_subscriptions_count(account_streams.len(), &self.shard_id);
                     self.try_to_override_first_subscribed_slot(pubkey, clock_slot);
                 }
@@ -164,17 +157,9 @@ impl RemoteAccountUpdatesShard {
             }
         }
         // Cleanup all subscriptions and wait for proper shutdown
-        for (pubkey, account_unsubscribes) in account_unsubscribes.into_iter() {
-            info!(
-                "Shard {}: Account monitoring killed: {:?}",
-                self.shard_id, pubkey
-            );
-            account_unsubscribes().await;
-        }
-        clock_unsubscribe().await;
         drop(account_streams);
         drop(clock_stream);
-        pubsub_client.shutdown().await?;
+        pool.shutdown().await;
         info!("Shard {}: Stopped", self.shard_id);
         // Done
         Ok(())
@@ -236,3 +221,98 @@ impl RemoteAccountUpdatesShard {
         }
     }
 }
+
+struct PubsubPool {
+    clients: Vec<PubSubConnection>,
+    unsubscribes: HashMap<Pubkey, (usize, BoxFn)>,
+    config: RpcAccountInfoConfig,
+}
+
+impl PubsubPool {
+    async fn new(
+        url: &str,
+        config: RpcAccountInfoConfig,
+    ) -> Result<Self, RemoteAccountUpdatesShardError> {
+        // 8 is pretty much arbitrary, but a sane value for the number
+        // of connections per RPC upstream, we don't overcomplicate things
+        // here, as the whole cloning pipeline will be rewritten quite soon
+        const CONNECTIONS_PER_POOL: usize = 8;
+        let mut clients = Vec::with_capacity(CONNECTIONS_PER_POOL);
+        let mut connections: FuturesUnordered<_> = (0..CONNECTIONS_PER_POOL)
+            .map(|_| PubSubConnection::new(url))
+            .collect();
+        while let Some(c) = connections.next().await {
+            clients.push(c?);
+        }
+        Ok(Self {
+            clients,
+            unsubscribes: HashMap::new(),
+            config,
+        })
+    }
+
+    async fn subscribe(
+        &mut self,
+        pubkey: Pubkey,
+    ) -> Result<SubscriptionStream, RemoteAccountUpdatesShardError> {
+        let (index, client) = self
+            .clients
+            .iter_mut()
+            .enumerate()
+            .min_by(|a, b| a.1.subs.cmp(&b.1.subs))
+            .expect("clients vec is always greater than 0");
+        let (stream, unsubscribe) = client
+            .inner
+            .account_subscribe(&pubkey, Some(self.config.clone()))
+            .await
+            .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
+        client.subs += 1;
+        // SAFETY:
+        // we never drop the PubsubPool before the returned subscription stream
+        // so the lifetime of the stream can be safely extended to 'static
+        #[allow(clippy::missing_transmute_annotations)]
+        let stream = unsafe { std::mem::transmute(stream) };
+        self.unsubscribes.insert(pubkey, (index, unsubscribe));
+        Ok(stream)
+    }
+
+    async fn unsubscribe(&mut self, pubkey: &Pubkey) {
+        let Some((index, callback)) = self.unsubscribes.remove(pubkey) else {
+            return;
+        };
+        callback().await;
+        let Some(client) = self.clients.get_mut(index) else {
+            return;
+        };
+        client.subs = client.subs.saturating_sub(1);
+    }
+
+    fn subscribed(&mut self, pubkey: &Pubkey) -> bool {
+        self.unsubscribes.contains_key(pubkey)
+    }
+
+    async fn shutdown(&mut self) {
+        // Cleanup all subscriptions and wait for proper shutdown
+        for (pubkey, (_, callback)) in self.unsubscribes.drain() {
+            info!("Account monitoring killed: {:?}", pubkey);
+            callback().await;
+        }
+        for client in self.clients.drain(..) {
+            let _ = client.inner.shutdown().await;
+        }
+    }
+}
+
+struct PubSubConnection {
+    inner: PubsubClient,
+    subs: usize,
+}
+
+impl PubSubConnection {
+    async fn new(url: &str) -> Result<Self, RemoteAccountUpdatesShardError> {
+        let inner = PubsubClient::new(url)
+            .await
+            .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
+        Ok(Self { inner, subs: 0 })
+    }
+}

magicblock-account-updates/tests/remote_account_updates.rs

@@ -14,7 +14,7 @@ use test_tools::skip_if_devnet_down;
 use tokio::time::sleep;
 use tokio_util::sync::CancellationToken;
 
-fn setup() -> (
+async fn setup() -> (
     RemoteAccountUpdatesClient,
     CancellationToken,
     tokio::task::JoinHandle<()>,
@@ -23,7 +23,7 @@ fn setup() -> (
     let mut worker = RemoteAccountUpdatesWorker::new(
         vec![RpcProviderConfig::devnet().ws_url().into(); 2],
         Some(solana_sdk::commitment_config::CommitmentLevel::Confirmed),
-        Duration::from_secs(1), // We constantly refresh stuff to make it struggle
+        Duration::from_secs(50 * 60),
     );
     let client = RemoteAccountUpdatesClient::new(&worker);
     // Run the worker in a separate task
@@ -36,6 +36,8 @@ fn setup() -> (
                 .await
         })
     };
+    // wait a bit for websocket connections to establish
+    sleep(Duration::from_millis(5_000)).await;
     // Ready to run
     (client, cancellation_token, worker_handle)
 }
@@ -44,11 +46,9 @@ fn setup() -> (
 async fn test_devnet_monitoring_clock_sysvar_changes_over_time() {
     skip_if_devnet_down!();
     // Create account updates worker and client
-    let (client, cancellation_token, worker_handle) = setup();
+    let (client, cancellation_token, worker_handle) = setup().await;
     // The clock will change every slots, perfect for testing updates
     let sysvar_clock = clock::ID;
-    // Before starting the monitoring, we should know nothing about the clock
-    assert!(client.get_last_known_update_slot(&sysvar_clock).is_none());
     // Start the monitoring
     assert!(client
         .ensure_account_monitoring(&sysvar_clock)
@@ -75,15 +75,14 @@ async fn test_devnet_monitoring_clock_sysvar_changes_over_time() {
 async fn test_devnet_monitoring_multiple_accounts_at_the_same_time() {
     skip_if_devnet_down!();
     // Create account updates worker and client
-    let (client, cancellation_token, worker_handle) = setup();
+    let (client, cancellation_token, worker_handle) = setup().await;
     // Devnet accounts to be monitored for this test
     let sysvar_rent = rent::ID;
     let sysvar_sh = slot_hashes::ID;
     let sysvar_clock = clock::ID;
     // We shouldnt known anything about the accounts until we subscribe
     assert!(client.get_last_known_update_slot(&sysvar_rent).is_none());
     assert!(client.get_last_known_update_slot(&sysvar_sh).is_none());
-    assert!(client.get_last_known_update_slot(&sysvar_clock).is_none());
     // Start monitoring the accounts now
     assert!(client.ensure_account_monitoring(&sysvar_rent).await.is_ok());
     assert!(client.ensure_account_monitoring(&sysvar_sh).await.is_ok());
@@ -106,15 +105,14 @@ async fn test_devnet_monitoring_multiple_accounts_at_the_same_time() {
 async fn test_devnet_monitoring_some_accounts_only() {
     skip_if_devnet_down!();
     // Create account updates worker and client
-    let (client, cancellation_token, worker_handle) = setup();
+    let (client, cancellation_token, worker_handle) = setup().await;
     // Devnet accounts for this test
     let sysvar_rent = rent::ID;
     let sysvar_sh = slot_hashes::ID;
     let sysvar_clock = clock::ID;
     // We shouldnt known anything about the accounts until we subscribe
     assert!(client.get_last_known_update_slot(&sysvar_rent).is_none());
     assert!(client.get_last_known_update_slot(&sysvar_sh).is_none());
-    assert!(client.get_last_known_update_slot(&sysvar_clock).is_none());
     // Start monitoring only some of the accounts
     assert!(client.ensure_account_monitoring(&sysvar_rent).await.is_ok());
     assert!(client.ensure_account_monitoring(&sysvar_sh).await.is_ok());
@@ -123,7 +121,7 @@ async fn test_devnet_monitoring_some_accounts_only() {
     // Check that we detected the accounts changes only on the accounts we monitored
     assert!(client.get_last_known_update_slot(&sysvar_rent).is_none()); // Rent doesn't change
     assert!(client.get_last_known_update_slot(&sysvar_sh).is_some());
-    assert!(client.get_last_known_update_slot(&sysvar_clock).is_none());
+    assert!(client.get_last_known_update_slot(&sysvar_clock).is_some());
     // Cleanup everything correctly
     cancellation_token.cancel();
     assert!(worker_handle.await.is_ok());
@@ -133,7 +131,7 @@ async fn test_devnet_monitoring_some_accounts_only() {
 async fn test_devnet_monitoring_invalid_and_immutable_and_program_account() {
     skip_if_devnet_down!();
     // Create account updates worker and client
-    let (client, cancellation_token, worker_handle) = setup();
+    let (client, cancellation_token, worker_handle) = setup().await;
     // Devnet accounts for this test (none of them should change)
     let new_account = Keypair::new().pubkey();
     let system_program = system_program::ID;