feat(cloning): make cloning concurrency configurable

magicblock-account-cloner/src/remote_account_cloner_worker.rs

@@ -17,7 +17,9 @@ use magicblock_account_fetcher::AccountFetcher;
 use magicblock_account_updates::{AccountUpdates, AccountUpdatesResult};
 use magicblock_accounts_api::InternalAccountProvider;
 use magicblock_committor_service::ChangesetCommittor;
-use magicblock_config::{AccountsCloneConfig, PrepareLookupTables};
+use magicblock_config::{
+    AccountsCloneConfig, PrepareLookupTables, ReplayConfig,
+};
 use magicblock_metrics::metrics;
 use magicblock_mutator::idl::{get_pubkey_anchor_idl, get_pubkey_shank_idl};
 use solana_sdk::{
@@ -110,6 +112,7 @@ pub struct RemoteAccountClonerWorker<IAP, AFE, AUP, ADU, CC> {
     validator_identity: Pubkey,
     monitored_accounts: RefCell<LruCache<Pubkey, ()>>,
     clone_config: AccountsCloneConfig,
+    replay_config: ReplayConfig,
 }
 
 // SAFETY:
@@ -149,6 +152,7 @@ where
         validator_authority: Pubkey,
         max_monitored_accounts: usize,
         clone_config: AccountsCloneConfig,
+        replay_config: ReplayConfig,
     ) -> Self {
         let (clone_request_sender, clone_request_receiver) = flume::unbounded();
         let fetch_retries = 50;
@@ -173,6 +177,7 @@ where
             validator_identity: validator_authority,
             monitored_accounts: LruCache::new(max_monitored_accounts).into(),
             clone_config,
+            replay_config,
         }
     }
 
@@ -294,34 +299,36 @@ where
         // retry resulting in overall slower hydration.
         // If the optimal rate here is desired we might make this configurable in the
         // future.
-        // TODO(GabrielePicco): Make the concurrency configurable
         let result = stream
             .map(Ok::<_, AccountClonerError>)
-            .try_for_each_concurrent(10, |(pubkey, owner)| async move {
-                trace!("Hydrating '{}'", pubkey);
-                let res = self
-                    .do_clone_and_update_cache(
-                        &pubkey,
-                        ValidatorStage::Hydrating {
-                            validator_identity: self.validator_identity,
-                            account_owner: owner,
-                        },
-                    )
-                    .await;
-                match res {
-                    Ok(output) => {
-                        trace!("Cloned '{}': {:?}", pubkey, output);
-                        Ok(())
-                    }
-                    Err(err) => {
-                        error!("Failed to clone {} ('{:?}')", pubkey, err);
-                        // NOTE: the account fetch already has retries built in, so
-                        // we don't to retry here
+            .try_for_each_concurrent(
+                self.replay_config.hydration_concurrency,
+                |(pubkey, owner)| async move {
+                    trace!("Hydrating '{}'", pubkey);
+                    let res = self
+                        .do_clone_and_update_cache(
+                            &pubkey,
+                            ValidatorStage::Hydrating {
+                                validator_identity: self.validator_identity,
+                                account_owner: owner,
+                            },
+                        )
+                        .await;
+                    match res {
+                        Ok(output) => {
+                            trace!("Cloned '{}': {:?}", pubkey, output);
+                            Ok(())
+                        }
+                        Err(err) => {
+                            error!("Failed to clone {} ('{:?}')", pubkey, err);
+                            // NOTE: the account fetch already has retries built in, so
+                            // we don't to retry here
 
-                        Err(err)
+                            Err(err)
+                        }
                     }
-                }
-            })
+                },
+            )
             .await;
         info!("On-startup account ensurance is complete: {count}");
         result

magicblock-account-cloner/tests/remote_account_cloner.rs

@@ -11,7 +11,7 @@ use magicblock_account_fetcher::AccountFetcherStub;
 use magicblock_account_updates::AccountUpdatesStub;
 use magicblock_accounts_api::InternalAccountProviderStub;
 use magicblock_committor_service::stubs::ChangesetCommittorStub;
-use magicblock_config::AccountsCloneConfig;
+use magicblock_config::{AccountsCloneConfig, ReplayConfig};
 use magicblock_mutator::idl::{get_pubkey_anchor_idl, get_pubkey_shank_idl};
 use solana_sdk::{
     bpf_loader_upgradeable::get_program_data_address,
@@ -51,6 +51,7 @@ fn setup_custom(
         Pubkey::new_unique(),
         1024,
         AccountsCloneConfig::default(),
+        ReplayConfig::default(),
     );
     let cloner_client = RemoteAccountClonerClient::new(&cloner_worker);
     // Run the worker in a separate task

magicblock-accounts/tests/ensure_accounts.rs

@@ -18,7 +18,7 @@ use magicblock_accounts::{
 };
 use magicblock_accounts_api::InternalAccountProviderStub;
 use magicblock_committor_service::stubs::ChangesetCommittorStub;
-use magicblock_config::AccountsCloneConfig;
+use magicblock_config::{AccountsCloneConfig, ReplayConfig};
 use solana_sdk::pubkey::Pubkey;
 use stubs::{
     account_committer_stub::AccountCommitterStub,
@@ -62,6 +62,7 @@ fn setup_with_lifecycle(
         Pubkey::new_unique(),
         1024,
         AccountsCloneConfig::default(),
+        ReplayConfig::default(),
     );
     let remote_account_cloner_client =
         RemoteAccountClonerClient::new(&remote_account_cloner_worker);

magicblock-api/src/magic_validator.rs

@@ -355,6 +355,7 @@ impl MagicValidator {
             identity_keypair.pubkey(),
             config.validator_config.accounts.max_monitored_accounts,
             config.validator_config.accounts.clone.clone(),
+            config.validator_config.ledger.replay.clone(),
         );
 
         let accounts_manager = Self::init_accounts_manager(

magicblock-config/src/accounts.rs

@@ -65,14 +65,12 @@ impl AccountsConfig {
         if self.db == default.db && other.db != default.db {
             self.db = other.db;
         }
-        if self.clone == default.clone && other.clone != default.clone {
-            self.clone = other.clone;
-        }
         if self.max_monitored_accounts == default.max_monitored_accounts
             && other.max_monitored_accounts != default.max_monitored_accounts
         {
             self.max_monitored_accounts = other.max_monitored_accounts;
         }
+        self.clone.merge(other.clone);
     }
 }
 
@@ -247,6 +245,18 @@ pub struct AccountsCloneConfig {
     pub prepare_lookup_tables: PrepareLookupTables,
 }
 
+impl AccountsCloneConfig {
+    pub fn merge(&mut self, other: AccountsCloneConfig) {
+        let default = Self::default();
+
+        if self.prepare_lookup_tables == default.prepare_lookup_tables
+            && other.prepare_lookup_tables != default.prepare_lookup_tables
+        {
+            self.prepare_lookup_tables = other.prepare_lookup_tables;
+        }
+    }
+}
+
 #[derive(
     Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize, Args,
 )]

magicblock-config/src/ledger.rs

@@ -32,6 +32,9 @@ pub struct LedgerConfig {
     #[arg(help = "The size under which it's desired to keep ledger in bytes.")]
     #[serde(default = "default_ledger_size")]
     pub size: u64,
+    #[serde(default)]
+    #[command(flatten)]
+    pub replay: ReplayConfig,
 }
 
 impl LedgerConfig {
@@ -47,6 +50,7 @@ impl LedgerConfig {
         {
             self.size = other.size;
         }
+        self.replay.merge(other.replay);
     }
 }
 
@@ -56,6 +60,38 @@ impl Default for LedgerConfig {
             reset: bool_true(),
             path: Default::default(),
             size: DEFAULT_LEDGER_SIZE_BYTES,
+            replay: ReplayConfig::default(),
+        }
+    }
+}
+
+#[clap_prefix("replay")]
+#[clap_from_serde]
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize, Args)]
+#[serde(deny_unknown_fields)]
+pub struct ReplayConfig {
+    /// The number of threads to use for cloning accounts.
+    #[derive_env_var]
+    #[serde(default = "default_cloning_concurrency")]
+    pub hydration_concurrency: usize,
+}
+
+impl ReplayConfig {
+    pub fn merge(&mut self, other: Self) {
+        let default = Self::default();
+
+        if self.hydration_concurrency == default.hydration_concurrency
+            && other.hydration_concurrency != default.hydration_concurrency
+        {
+            self.hydration_concurrency = other.hydration_concurrency;
+        }
+    }
+}
+
+impl Default for ReplayConfig {
+    fn default() -> Self {
+        Self {
+            hydration_concurrency: default_cloning_concurrency(),
         }
     }
 }
@@ -64,16 +100,23 @@ const fn default_ledger_size() -> u64 {
     DEFAULT_LEDGER_SIZE_BYTES
 }
 
+const fn default_cloning_concurrency() -> usize {
+    10
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_merge_with_default() {
+    fn test_ledger_merge_with_default() {
         let mut config = LedgerConfig {
             reset: false,
             path: Some("ledger.example.com".to_string()),
             size: 1000000000,
+            replay: ReplayConfig {
+                hydration_concurrency: 20,
+            },
         };
         let original_config = config.clone();
         let other = LedgerConfig::default();
@@ -84,12 +127,15 @@ mod tests {
     }
 
     #[test]
-    fn test_merge_default_with_non_default() {
+    fn test_ledger_merge_default_with_non_default() {
         let mut config = LedgerConfig::default();
         let other = LedgerConfig {
             reset: false,
             path: Some("ledger.example.com".to_string()),
             size: 1000000000,
+            replay: ReplayConfig {
+                hydration_concurrency: 20,
+            },
         };
 
         config.merge(other.clone());
@@ -98,17 +144,63 @@ mod tests {
     }
 
     #[test]
-    fn test_merge_non_default() {
+    fn test_ledger_merge_non_default() {
         let mut config = LedgerConfig {
             reset: false,
             path: Some("ledger.example.com".to_string()),
             size: 1000000000,
+            replay: ReplayConfig {
+                hydration_concurrency: 20,
+            },
         };
         let original_config = config.clone();
         let other = LedgerConfig {
             reset: true,
             path: Some("ledger2.example.com".to_string()),
             size: 10000,
+            replay: ReplayConfig {
+                hydration_concurrency: 150,
+            },
+        };
+
+        config.merge(other);
+
+        assert_eq!(config, original_config);
+    }
+
+    #[test]
+    fn test_replay_merge_with_default() {
+        let mut config = ReplayConfig {
+            hydration_concurrency: 20,
+        };
+        let original_config = config.clone();
+        let other = ReplayConfig::default();
+
+        config.merge(other);
+
+        assert_eq!(config, original_config);
+    }
+
+    #[test]
+    fn test_replay_merge_default_with_non_default() {
+        let mut config = ReplayConfig::default();
+        let other = ReplayConfig {
+            hydration_concurrency: 20,
+        };
+
+        config.merge(other.clone());
+
+        assert_eq!(config, other);
+    }
+
+    #[test]
+    fn test_replay_merge_non_default() {
+        let mut config = ReplayConfig {
+            hydration_concurrency: 20,
+        };
+        let original_config = config.clone();
+        let other = ReplayConfig {
+            hydration_concurrency: 150,
         };
 
         config.merge(other);