Bmuddha/fix/limit num monitoring accounts (#356)

This PR introduces a quick fix to existing cloning pipeline, to remove
undelegated account from validator (both from database and from
websocket monitoring) once certain preconfigured threshold is reached.
The eviction of those accounts happen using LRU strategy. Every
transaction affects this cache by touching the account and renewing its
"stay", thus only stale/unused accounts are removed. Delegated accounts
are not affected by this change.
 
<!-- greptile_comment -->

## Greptile Summary

Implemented account monitoring limits using an LRU cache to prevent
unbounded memory growth in the validator.

- Added `monitored_accounts` LRU cache in `RemoteAccountClonerWorker`
with configurable capacity (default 2048)
- Implemented account eviction logic in `track_undelegated_account` that
removes accounts from AccountsDB and websocket subscriptions
- Added metrics tracking for monitored accounts count, subscriptions per
shard, and evicted accounts
- Modified account update system to support unsubscribing from account
monitoring when accounts are evicted
- Added `remove_account` capability to AccountsDB and providers to
properly cleanup evicted accounts



<!-- /greptile_comment -->

---------

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

Author: 3 people

Cargo.lock

@@ -22,5 +22,6 @@ solana-sdk = { workspace = true }
 tokio = { workspace = true }
 tokio-util = { workspace = true }
 thiserror = { workspace = true }
+lru = "0.14"
 
 [dev-dependencies]

magicblock-account-cloner/Cargo.toml

@@ -1,4 +1,5 @@
 use std::{
+    cell::RefCell,
     collections::{hash_map::Entry, HashMap, HashSet},
     sync::{Arc, RwLock},
     time::Duration,
@@ -14,9 +15,10 @@ use futures_util::{
     stream::{self, StreamExt, TryStreamExt},
 };
 use log::*;
+use lru::LruCache;
 use magicblock_account_dumper::AccountDumper;
 use magicblock_account_fetcher::AccountFetcher;
-use magicblock_account_updates::AccountUpdates;
+use magicblock_account_updates::{AccountUpdates, AccountUpdatesResult};
 use magicblock_accounts_api::InternalAccountProvider;
 use magicblock_metrics::metrics;
 use magicblock_mutator::idl::{get_pubkey_anchor_idl, get_pubkey_shank_idl};
@@ -110,6 +112,22 @@ pub struct RemoteAccountClonerWorker<IAP, AFE, AUP, ADU> {
     clone_listeners: Arc<RwLock<HashMap<Pubkey, AccountClonerListeners>>>,
     last_clone_output: CloneOutputMap,
     validator_identity: Pubkey,
+    monitored_accounts: RefCell<LruCache<Pubkey, ()>>,
+}
+
+// SAFETY:
+// we never keep references to monitored_accounts around,
+// especially across await points, so this type is Send
+unsafe impl<IAP, AFE, AUP, ADU> Send
+    for RemoteAccountClonerWorker<IAP, AFE, AUP, ADU>
+{
+}
+// SAFETY:
+// we never produce references to RefCell in monitored_accounts
+// especially not across await points, so this type is Sync
+unsafe impl<IAP, AFE, AUP, ADU> Sync
+    for RemoteAccountClonerWorker<IAP, AFE, AUP, ADU>
+{
 }
 
 impl<IAP, AFE, AUP, ADU> RemoteAccountClonerWorker<IAP, AFE, AUP, ADU>
@@ -131,10 +149,14 @@ where
         validator_charges_fees: ValidatorCollectionMode,
         permissions: AccountClonerPermissions,
         validator_authority: Pubkey,
+        max_monitored_accounts: usize,
     ) -> Self {
         let (clone_request_sender, clone_request_receiver) =
             unbounded_channel();
         let fetch_retries = 50;
+        let max_monitored_accounts = max_monitored_accounts
+            .try_into()
+            .expect("max number of monitored accounts cannot be 0");
         Self {
             internal_account_provider,
             account_fetcher,
@@ -151,6 +173,7 @@ where
             clone_listeners: Default::default(),
             last_clone_output: Default::default(),
             validator_identity: validator_authority,
+            monitored_accounts: LruCache::new(max_monitored_accounts).into(),
         }
     }
 
@@ -169,16 +192,16 @@ where
     }
 
     pub async fn start_clone_request_processing(
-        &mut self,
+        mut self,
         cancellation_token: CancellationToken,
     ) {
+        let mut requests = vec![];
         loop {
-            let mut requests = vec![];
             tokio::select! {
                 _ = self.clone_request_receiver.recv_many(&mut requests, 100) => {
                     join_all(
                         requests
-                            .into_iter()
+                            .drain(..)
                             .map(|request| self.process_clone_request(request))
                     ).await;
                 }
@@ -323,6 +346,7 @@ where
             .account_updates
             .get_last_known_update_slot(pubkey)
             .unwrap_or(u64::MIN);
+        self.monitored_accounts.borrow_mut().promote(pubkey);
         // Check for the happy/fast path, we may already have cloned this account before
         match self.get_last_clone_output_from_pubkey(pubkey) {
             // If we already cloned this account, check what the output of the clone was
@@ -405,6 +429,39 @@ where
         Ok(updated_clone_output)
     }
 
+    /// Put the account's key into cache of monitored accounts, which has a limited capacity.
+    /// Once the cache capacity exceeds the preconfigured limit, it will trigger an eviction,
+    /// followed by account's removal from AccountsDB and termination of its ws subscription
+    async fn track_not_delegated_account(
+        &self,
+        pubkey: Pubkey,
+    ) -> AccountUpdatesResult<()> {
+        let evicted = self
+            .monitored_accounts
+            .borrow_mut()
+            .push(pubkey, ())
+            .filter(|(pk, _)| *pk != pubkey);
+        if let Some((evicted, _)) = evicted {
+            self.last_clone_output
+                .write()
+                .expect("last accounts clone output map is poisoned")
+                .remove(&evicted);
+            self.internal_account_provider.remove_account(&evicted);
+            self.clone_listeners
+                .write()
+                .expect("clone listeners map is poisoned")
+                .remove(&evicted);
+            self.account_updates
+                .stop_account_monitoring(&evicted)
+                .await?;
+            metrics::inc_evicted_accounts_count();
+        }
+        metrics::adjust_monitored_accounts_count(
+            self.monitored_accounts.borrow().len(),
+        );
+        Ok(())
+    }
+
     async fn do_clone(
         &self,
         pubkey: &Pubkey,
@@ -428,8 +485,7 @@ where
             //  - we may not want to track lamport changes, especially for payers
             self.account_updates
                 .ensure_account_monitoring(pubkey)
-                .await
-                .map_err(AccountClonerError::AccountUpdatesError)?;
+                .await?;
 
             // Fetch the account, repeat and retry until we have a satisfactory response
             let mut fetch_count = 0;
@@ -491,6 +547,8 @@ where
                     });
                 }
 
+                // Fee payer accounts are non-delegated ones, so we keep track of them as well
+                self.track_not_delegated_account(*pubkey).await?;
                 match self.validator_charges_fees {
                     ValidatorCollectionMode::NoFees => self
                         .do_clone_feepayer_account_for_non_charging_validator(
@@ -602,6 +660,9 @@ where
                             at_slot: account_chain_snapshot.at_slot,
                         });
                     }
+                    // Keep track of non-delegated accounts, removing any stale ones,
+                    // which were evicted from monitored accounts cache
+                    self.track_not_delegated_account(*pubkey).await?;
                     self.do_clone_undelegated_account(pubkey, account)?
                 }
             }
@@ -612,6 +673,13 @@ where
                 delegation_record,
                 ..
             } => {
+                // Just in case if the account was promoted from not delegated to delegated state, we
+                // remove it from list of monitored accounts, to avoid removal on eviction
+                self.monitored_accounts.borrow_mut().pop(pubkey);
+                metrics::adjust_monitored_accounts_count(
+                    self.monitored_accounts.borrow().len(),
+                );
+
                 if !self.permissions.allow_cloning_delegated_accounts {
                     return Ok(AccountClonerOutput::Unclonable {
                         pubkey: *pubkey,

magicblock-account-cloner/src/remote_account_cloner_worker.rs

@@ -37,7 +37,7 @@ fn setup_custom(
     // Default configuration
     let payer_init_lamports = Some(1_000 * LAMPORTS_PER_SOL);
     // Create account cloner worker and client
-    let mut cloner_worker = RemoteAccountClonerWorker::new(
+    let cloner_worker = RemoteAccountClonerWorker::new(
         internal_account_provider,
         account_fetcher,
         account_updates,
@@ -48,17 +48,17 @@ fn setup_custom(
         ValidatorCollectionMode::NoFees,
         permissions,
         Pubkey::new_unique(),
+        1024,
     );
     let cloner_client = RemoteAccountClonerClient::new(&cloner_worker);
     // Run the worker in a separate task
     let cancellation_token = CancellationToken::new();
     let cloner_worker_handle = {
         let cloner_cancellation_token = cancellation_token.clone();
-        tokio::spawn(async move {
+        tokio::spawn(
             cloner_worker
-                .start_clone_request_processing(cloner_cancellation_token)
-                .await
-        })
+                .start_clone_request_processing(cloner_cancellation_token),
+        )
     };
     // Ready to run
     (cloner_client, cancellation_token, cloner_worker_handle)

magicblock-account-cloner/tests/remote_account_cloner.rs

@@ -8,6 +8,7 @@ license.workspace = true
 edition.workspace = true
 
 [dependencies]
+magicblock-metrics = { workspace = true }
 conjunto-transwise = { workspace = true }
 futures-util = { workspace = true }
 log = { workspace = true }

magicblock-account-updates/Cargo.toml

@@ -1,10 +1,11 @@
 use solana_sdk::{clock::Slot, pubkey::Pubkey};
 use thiserror::Error;
+use tokio::sync::mpsc::error::SendError;
 
 #[derive(Debug, Clone, Error)]
 pub enum AccountUpdatesError {
     #[error(transparent)]
-    SendError(#[from] tokio::sync::mpsc::error::SendError<Pubkey>),
+    SendError(#[from] SendError<(Pubkey, bool)>),
 }
 
 pub type AccountUpdatesResult<T> = Result<T, AccountUpdatesError>;
@@ -15,6 +16,13 @@ pub trait AccountUpdates {
         &self,
         pubkey: &Pubkey,
     ) -> AccountUpdatesResult<()>;
+    #[allow(async_fn_in_trait)]
+    async fn stop_account_monitoring(
+        &self,
+        _pubkey: &Pubkey,
+    ) -> AccountUpdatesResult<()> {
+        Ok(())
+    }
     fn get_first_subscribed_slot(&self, pubkey: &Pubkey) -> Option<Slot>;
     fn get_last_known_update_slot(&self, pubkey: &Pubkey) -> Option<Slot>;
 }

magicblock-account-updates/src/account_updates.rs

@@ -9,7 +9,7 @@ use tokio::sync::mpsc::Sender;
 use crate::{AccountUpdates, AccountUpdatesError, RemoteAccountUpdatesWorker};
 
 pub struct RemoteAccountUpdatesClient {
-    monitoring_request_sender: Sender<Pubkey>,
+    monitoring_request_sender: Sender<(Pubkey, bool)>,
     first_subscribed_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
     last_known_update_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
 }
@@ -30,17 +30,29 @@ impl AccountUpdates for RemoteAccountUpdatesClient {
         pubkey: &Pubkey,
     ) -> Result<(), AccountUpdatesError> {
         self.monitoring_request_sender
-            .send(*pubkey)
+            .send((*pubkey, false))
             .await
-            .map_err(AccountUpdatesError::SendError)
+            .map_err(Into::into)
     }
+
+    async fn stop_account_monitoring(
+        &self,
+        pubkey: &Pubkey,
+    ) -> Result<(), AccountUpdatesError> {
+        self.monitoring_request_sender
+            .send((*pubkey, true))
+            .await
+            .map_err(Into::into)
+    }
+
     fn get_first_subscribed_slot(&self, pubkey: &Pubkey) -> Option<Slot> {
         self.first_subscribed_slots
             .read()
             .expect("RwLock of RemoteAccountUpdatesClient.first_subscribed_slots poisoned")
             .get(pubkey)
             .cloned()
     }
+
     fn get_last_known_update_slot(&self, pubkey: &Pubkey) -> Option<Slot> {
         self.last_known_update_slots
             .read()