chore: improve subscription reconciler

magicblock-chainlink/src/remote_account_provider/chain_laser_actor.rs

@@ -2,7 +2,10 @@ use std::{
     collections::{HashMap, HashSet},
     fmt,
     pin::Pin,
-    sync::atomic::{AtomicU16, AtomicU64, Ordering},
+    sync::{
+        atomic::{AtomicU16, AtomicU64, Ordering},
+        Arc,
+    },
     time::Duration,
 };
 
@@ -23,6 +26,7 @@ use magicblock_metrics::metrics::{
     inc_per_program_account_updates_count,
     inc_program_subscription_account_updates_count,
 };
+use parking_lot::RwLock;
 use solana_account::Account;
 use solana_commitment_config::CommitmentLevel as SolanaCommitmentLevel;
 use solana_pubkey::Pubkey;
@@ -55,6 +59,8 @@ const SLOTS_BETWEEN_ACTIVATIONS: u64 =
     SUBSCRIPTION_ACTIVATION_INTERVAL_MILLIS / 400;
 const MAX_SLOTS_BACKFILL: u64 = 400;
 
+pub type SharedSubscriptions = Arc<RwLock<HashSet<Pubkey>>>;
+
 // -----------------
 // Slots
 // -----------------
@@ -116,8 +122,10 @@ impl fmt::Display for AccountUpdateSource {
 pub struct ChainLaserActor {
     /// Configuration used to create the laser client
     laser_client_config: LaserstreamConfig,
-    /// Requested subscriptions, some may not be active yet
-    subscriptions: HashSet<Pubkey>,
+    /// Requested subscriptions, some may not be active yet.
+    /// Shared with ChainLaserClientImpl for sync access to
+    /// subscription_count and subscriptions_union.
+    subscriptions: SharedSubscriptions,
     /// Pubkeys of currently active subscriptions
     active_subscription_pubkeys: HashSet<Pubkey>,
     /// Subscriptions that have been activated via the helius provider
@@ -155,6 +163,7 @@ impl ChainLaserActor {
         Self,
         mpsc::Sender<ChainPubsubActorMessage>,
         mpsc::Receiver<SubscriptionUpdate>,
+        SharedSubscriptions,
     ) {
         let channel_options = ChannelOptions {
             connect_timeout_secs: Some(5),
@@ -190,17 +199,21 @@ impl ChainLaserActor {
         Self,
         mpsc::Sender<ChainPubsubActorMessage>,
         mpsc::Receiver<SubscriptionUpdate>,
+        SharedSubscriptions,
     ) {
         let (subscription_updates_sender, subscription_updates_receiver) =
             mpsc::channel(SUBSCRIPTION_UPDATE_CHANNEL_SIZE);
         let (messages_sender, messages_receiver) =
             mpsc::channel(MESSAGE_CHANNEL_SIZE);
         let commitment = grpc_commitment_from_solana(commitment);
 
+        let subscriptions: SharedSubscriptions = Default::default();
+        let shared_subscriptions = Arc::clone(&subscriptions);
+
         let me = Self {
             laser_client_config,
             messages_receiver,
-            subscriptions: Default::default(),
+            subscriptions,
             active_subscriptions: Default::default(),
             active_subscription_pubkeys: Default::default(),
             program_subscriptions: Default::default(),
@@ -212,14 +225,19 @@ impl ChainLaserActor {
             rpc_client,
         };
 
-        (me, messages_sender, subscription_updates_receiver)
+        (
+            me,
+            messages_sender,
+            subscription_updates_receiver,
+            shared_subscriptions,
+        )
     }
 
     #[allow(dead_code)]
     #[instrument(skip(self), fields(client_id = %self.client_id))]
     fn shutdown(&mut self) {
         info!("Shutting down laser actor");
-        self.subscriptions.clear();
+        self.subscriptions.write().clear();
         self.active_subscriptions.clear();
         self.active_subscription_pubkeys.clear();
     }
@@ -256,7 +274,7 @@ impl ChainLaserActor {
                         None => {
                             debug!("Account subscription stream ended");
                             Self::signal_connection_issue(
-                                &mut self.subscriptions,
+                                &self.subscriptions,
                                 &mut self.active_subscriptions,
                                 &mut self.active_subscription_pubkeys,
                                 &mut self.program_subscriptions,
@@ -281,7 +299,7 @@ impl ChainLaserActor {
                         None => {
                             debug!("Program subscription stream ended");
                             Self::signal_connection_issue(
-                                &mut self.subscriptions,
+                                &self.subscriptions,
                                 &mut self.active_subscriptions,
                                 &mut self.active_subscription_pubkeys,
                                 &mut self.program_subscriptions,
@@ -338,7 +356,7 @@ impl ChainLaserActor {
             Shutdown { response } => {
                 info!(client_id = self.client_id, "Received Shutdown message");
                 Self::clear_subscriptions(
-                    &mut self.subscriptions,
+                    &self.subscriptions,
                     &mut self.active_subscriptions,
                     &mut self.active_subscription_pubkeys,
                     &mut self.program_subscriptions,
@@ -360,14 +378,28 @@ impl ChainLaserActor {
         pubkey: Pubkey,
         sub_response: oneshot::Sender<RemoteAccountProviderResult<()>>,
     ) {
-        if self.subscriptions.contains(&pubkey) {
-            debug!(pubkey = %pubkey, "Already subscribed to account");
+        let inserted = {
+            // Fast path: check with read lock first
+            let already_subscribed = {
+                let subs = self.subscriptions.read();
+                subs.contains(&pubkey)
+            };
+
+            if already_subscribed {
+                false
+            } else {
+                // Write lock only when we need to modify
+                let mut subs = self.subscriptions.write();
+                subs.insert(pubkey);
+                true
+            }
+        };
+        if !inserted {
+            trace!(pubkey = %pubkey, "Already subscribed to account");
             sub_response.send(Ok(())).unwrap_or_else(|_| {
                 warn!(pubkey = %pubkey, "Failed to send already subscribed response");
             });
         } else {
-            self.subscriptions.insert(pubkey);
-            // If this is the first sub for the clock sysvar we want to activate it immediately
             if self.active_subscriptions.is_empty() {
                 self.update_active_subscriptions();
             }
@@ -383,7 +415,8 @@ impl ChainLaserActor {
         pubkey: &Pubkey,
         unsub_response: oneshot::Sender<RemoteAccountProviderResult<()>>,
     ) {
-        match self.subscriptions.remove(pubkey) {
+        let removed = self.subscriptions.write().remove(pubkey);
+        match removed {
             true => {
                 trace!(pubkey = %pubkey, "Unsubscribed from account");
                 unsub_response.send(Ok(())).unwrap_or_else(|_| {
@@ -405,25 +438,28 @@ impl ChainLaserActor {
     }
 
     fn update_active_subscriptions(&mut self) {
-        // Check if the active subscriptions match what we already have
-        let new_pubkeys: HashSet<Pubkey> =
-            self.subscriptions.iter().copied().collect();
-        if new_pubkeys == self.active_subscription_pubkeys {
-            trace!(
-                count = self.subscriptions.len(),
-                "Active subscriptions already up to date"
-            );
-            return;
-        }
+        // Copy subscriptions and release the read lock immediately
+        let new_pubkeys: HashSet<Pubkey> = {
+            let subs = self.subscriptions.read();
+            // Check if the active subscriptions match what we already have
+            if subs.eq(&self.active_subscription_pubkeys) {
+                trace!(
+                    count = subs.len(),
+                    "Active subscriptions already up to date"
+                );
+                return;
+            }
+            subs.iter().copied().collect()
+        };
 
         inc_account_subscription_activations_count(&self.client_id);
 
         let mut new_subs: StreamMap<usize, LaserStream> = StreamMap::new();
 
         // Re-create streams for all subscriptions
-        let subs = self.subscriptions.iter().collect::<Vec<_>>();
+        let sub_refs = new_pubkeys.iter().collect::<Vec<_>>();
 
-        let chunks = subs
+        let chunks = sub_refs
             .chunks(PER_STREAM_SUBSCRIPTION_LIMIT)
             .map(|chunk| chunk.to_vec())
             .collect::<Vec<_>>();
@@ -435,7 +471,7 @@ impl ChainLaserActor {
 
         if tracing::enabled!(tracing::Level::TRACE) {
             trace!(
-                account_count = self.subscriptions.len(),
+                account_count = new_pubkeys.len(),
                 chain_slot,
                 from_slot,
                 stream_count = chunks.len(),
@@ -639,7 +675,7 @@ impl ChainLaserActor {
 
         error!(error = ?err, slots = ?self.slots, "Error in {} stream", source);
         Self::signal_connection_issue(
-            &mut self.subscriptions,
+            &self.subscriptions,
             &mut self.active_subscriptions,
             &mut self.active_subscription_pubkeys,
             &mut self.program_subscriptions,
@@ -710,12 +746,12 @@ impl ChainLaserActor {
     }
 
     fn clear_subscriptions(
-        subscriptions: &mut HashSet<Pubkey>,
+        subscriptions: &SharedSubscriptions,
         active_subscriptions: &mut StreamMap<usize, LaserStream>,
         active_subscription_pubkeys: &mut HashSet<Pubkey>,
         program_subscriptions: &mut Option<(HashSet<Pubkey>, LaserStream)>,
     ) {
-        subscriptions.clear();
+        subscriptions.write().clear();
         active_subscriptions.clear();
         active_subscription_pubkeys.clear();
         *program_subscriptions = None;
@@ -727,7 +763,7 @@ impl ChainLaserActor {
     /// we add this as a backup in case it is unable to do so
     #[instrument(skip(subscriptions, active_subscriptions, active_subscription_pubkeys, program_subscriptions, abort_sender), fields(client_id = %client_id))]
     async fn signal_connection_issue(
-        subscriptions: &mut HashSet<Pubkey>,
+        subscriptions: &SharedSubscriptions,
         active_subscriptions: &mut StreamMap<usize, LaserStream>,
         active_subscription_pubkeys: &mut HashSet<Pubkey>,
         program_subscriptions: &mut Option<(HashSet<Pubkey>, LaserStream)>,
@@ -837,8 +873,7 @@ impl ChainLaserActor {
             );
         }
 
-        if !self.subscriptions.contains(&pubkey) {
-            // Ignore updates for accounts we are not subscribed to
+        if !self.subscriptions.read().contains(&pubkey) {
             return;
         }
 

magicblock-chainlink/src/remote_account_provider/chain_laser_client.rs

@@ -11,7 +11,7 @@ use tokio::sync::{mpsc, oneshot};
 use tracing::*;
 
 use crate::remote_account_provider::{
-    chain_laser_actor::{ChainLaserActor, Slots},
+    chain_laser_actor::{ChainLaserActor, SharedSubscriptions, Slots},
     chain_rpc_client::ChainRpcClientImpl,
     pubsub_common::{ChainPubsubActorMessage, SubscriptionUpdate},
     ChainPubsubClient, ReconnectableClient, RemoteAccountProviderError,
@@ -48,6 +48,8 @@ pub struct ChainLaserClientImpl {
     updates: Arc<Mutex<Option<mpsc::Receiver<SubscriptionUpdate>>>>,
     /// Channel to send messages to the actor
     messages: mpsc::Sender<ChainPubsubActorMessage>,
+    /// Shared subscriptions with the actor for sync access
+    subscriptions: SharedSubscriptions,
     /// Client identifier
     client_id: String,
 }
@@ -62,18 +64,20 @@ impl ChainLaserClientImpl {
         slots: Slots,
         rpc_client: ChainRpcClientImpl,
     ) -> Self {
-        let (actor, messages, updates) = ChainLaserActor::new_from_url(
-            pubsub_url,
-            &client_id,
-            api_key,
-            commitment,
-            abort_sender,
-            slots,
-            rpc_client,
-        );
+        let (actor, messages, updates, subscriptions) =
+            ChainLaserActor::new_from_url(
+                pubsub_url,
+                &client_id,
+                api_key,
+                commitment,
+                abort_sender,
+                slots,
+                rpc_client,
+            );
         let client = Self {
             updates: Arc::new(Mutex::new(Some(updates))),
             messages,
+            subscriptions,
             client_id,
         };
         tokio::spawn(actor.run());
@@ -174,15 +178,8 @@ impl ChainPubsubClient for ChainLaserClientImpl {
             .expect("ChainLaserClientImpl::take_updates called more than once")
     }
 
-    async fn subscription_count(
-        &self,
-        _exclude: Option<&[Pubkey]>,
-    ) -> Option<(usize, usize)> {
-        None
-    }
-
-    fn subscriptions(&self) -> Option<Vec<Pubkey>> {
-        None
+    fn subscriptions_union(&self) -> HashSet<Pubkey> {
+        self.subscriptions.read().clone()
     }
 
     fn subs_immediately(&self) -> bool {

magicblock-chainlink/src/remote_account_provider/chain_pubsub_actor.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     sync::{
         atomic::{AtomicBool, AtomicU16, Ordering},
         Arc, Mutex,
@@ -169,26 +169,9 @@ impl ChainPubsubActor {
         }
     }
 
-    pub fn subscription_count(&self, filter: &[Pubkey]) -> usize {
+    pub fn subscriptions(&self) -> HashSet<Pubkey> {
         if !self.is_connected.load(Ordering::SeqCst) {
-            return 0;
-        }
-        let subs = self
-            .subscriptions
-            .lock()
-            .expect("subscriptions lock poisoned");
-        if filter.is_empty() {
-            subs.len()
-        } else {
-            subs.keys()
-                .filter(|pubkey| !filter.contains(pubkey))
-                .count()
-        }
-    }
-
-    pub fn subscriptions(&self) -> Vec<Pubkey> {
-        if !self.is_connected.load(Ordering::SeqCst) {
-            return vec![];
+            return HashSet::new();
         }
         let subs = self
             .subscriptions