feat: account hydration on background task (#407)

- **deps: add flume properly**
- **feat: put account hydration on a background task**


<!-- greptile_comment -->

## Greptile Summary

Refactors account hydration to run as a background task, improving
validator startup performance by making the process asynchronous.

- Switches from `tokio::mpsc` to `flume` channels for better
multi-producer/consumer performance in `magicblock-account-cloner`
- Reduces concurrent account hydration from 30 to 10 tasks in
`remote_account_cloner_worker.rs` for better resource management
- Improves async task handling in test suites with proper await patterns
and timing adjustments
- Wraps `RemoteAccountClonerWorker` in `Arc` for safe concurrent access
in `magic_validator.rs`
- Adds 5-second sleep in integration tests to accommodate background
hydration timing



<!-- /greptile_comment -->

Author: thlorenz

Cargo.lock

@@ -65,9 +65,10 @@ isocountry = "0.3.2"
 crossbeam-channel = "0.5.11"
 enum-iterator = "1.5.0"
 env_logger = "0.11.2"
-magic-domain-program = { git = "https://github.com/magicblock-labs/magic-domain-program.git", rev = "ea04d46", default-features = false}
+magic-domain-program = { git = "https://github.com/magicblock-labs/magic-domain-program.git", rev = "ea04d46", default-features = false }
 magicblock-delegation-program = { git = "https://github.com/magicblock-labs/delegation-program.git", rev = "4af7f1c" }
 fd-lock = "4.0.2"
+flume = "0.11"
 fs_extra = "1.3.0"
 futures-util = "0.3.30"
 geyser-grpc-proto = { path = "./geyser-grpc-proto" }
@@ -150,7 +151,7 @@ solana-rpc = "2.2"
 solana-rpc-client = { version = "2.2" }
 solana-rpc-client-api = { version = "2.2" }
 solana-sdk = { version = "2.2" }
-solana-svm = { version = "2.2", features = [ "dev-context-only-utils" ] }
+solana-svm = { version = "2.2", features = ["dev-context-only-utils"] }
 solana-svm-transaction = { version = "2.2" }
 solana-storage-proto = { path = "storage-proto" }
 solana-system-program = { version = "2.2" }

Cargo.toml

@@ -9,6 +9,7 @@ edition.workspace = true
 
 [dependencies]
 conjunto-transwise = { workspace = true }
+flume = { workspace = true }
 futures-util = { workspace = true }
 log = { workspace = true }
 magicblock-account-fetcher = { workspace = true }

magicblock-account-cloner/Cargo.toml

@@ -16,7 +16,7 @@ use tokio::sync::oneshot::Sender;
 #[derive(Debug, Clone, Error)]
 pub enum AccountClonerError {
     #[error(transparent)]
-    SendError(#[from] tokio::sync::mpsc::error::SendError<Pubkey>),
+    SendError(#[from] flume::SendError<Pubkey>),
 
     #[error(transparent)]
     RecvError(#[from] tokio::sync::oneshot::error::RecvError),

magicblock-account-cloner/src/account_cloner.rs

@@ -12,15 +12,15 @@ use magicblock_account_fetcher::AccountFetcher;
 use magicblock_account_updates::AccountUpdates;
 use magicblock_accounts_api::InternalAccountProvider;
 use solana_sdk::pubkey::Pubkey;
-use tokio::sync::{mpsc::UnboundedSender, oneshot::channel};
+use tokio::sync::oneshot::channel;
 
 use crate::{
     AccountCloner, AccountClonerError, AccountClonerListeners,
     AccountClonerOutput, AccountClonerResult, RemoteAccountClonerWorker,
 };
 
 pub struct RemoteAccountClonerClient {
-    clone_request_sender: UnboundedSender<Pubkey>,
+    clone_request_sender: flume::Sender<Pubkey>,
     clone_listeners: Arc<RwLock<HashMap<Pubkey, AccountClonerListeners>>>,
 }
 

magicblock-account-cloner/src/remote_account_cloner_client.rs

@@ -3,17 +3,13 @@ use std::{
     collections::{hash_map::Entry, HashMap, HashSet},
     sync::{Arc, RwLock},
     time::Duration,
-    vec,
 };
 
 use conjunto_transwise::{
     AccountChainSnapshot, AccountChainSnapshotShared, AccountChainState,
     DelegationRecord,
 };
-use futures_util::{
-    future::join_all,
-    stream::{self, StreamExt, TryStreamExt},
-};
+use futures_util::stream::{self, StreamExt, TryStreamExt};
 use log::*;
 use lru::LruCache;
 use magicblock_account_dumper::AccountDumper;
@@ -29,10 +25,7 @@ use solana_sdk::{
     pubkey::Pubkey,
     signature::Signature,
 };
-use tokio::{
-    sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
-    time::sleep,
-};
+use tokio::time::sleep;
 use tokio_util::sync::CancellationToken;
 
 use crate::{
@@ -107,8 +100,8 @@ pub struct RemoteAccountClonerWorker<IAP, AFE, AUP, ADU> {
     validator_charges_fees: ValidatorCollectionMode,
     permissions: AccountClonerPermissions,
     fetch_retries: u64,
-    clone_request_receiver: UnboundedReceiver<Pubkey>,
-    clone_request_sender: UnboundedSender<Pubkey>,
+    clone_request_sender: flume::Sender<Pubkey>,
+    clone_request_receiver: flume::Receiver<Pubkey>,
     clone_listeners: Arc<RwLock<HashMap<Pubkey, AccountClonerListeners>>>,
     last_clone_output: CloneOutputMap,
     validator_identity: Pubkey,
@@ -151,8 +144,7 @@ where
         validator_authority: Pubkey,
         max_monitored_accounts: usize,
     ) -> Self {
-        let (clone_request_sender, clone_request_receiver) =
-            unbounded_channel();
+        let (clone_request_sender, clone_request_receiver) = flume::unbounded();
         let fetch_retries = 50;
         let max_monitored_accounts = max_monitored_accounts
             .try_into()
@@ -177,7 +169,7 @@ where
         }
     }
 
-    pub fn get_clone_request_sender(&self) -> UnboundedSender<Pubkey> {
+    pub fn get_clone_request_sender(&self) -> flume::Sender<Pubkey> {
         self.clone_request_sender.clone()
     }
 
@@ -192,18 +184,18 @@ where
     }
 
     pub async fn start_clone_request_processing(
-        mut self,
+        &self,
         cancellation_token: CancellationToken,
     ) {
-        let mut requests = vec![];
         loop {
             tokio::select! {
-                _ = self.clone_request_receiver.recv_many(&mut requests, 100) => {
-                    join_all(
-                        requests
-                            .drain(..)
-                            .map(|request| self.process_clone_request(request))
-                    ).await;
+                res = self.clone_request_receiver.recv_async() => {
+                    match res {
+                        Ok(req) => self.process_clone_request(req).await,
+                        Err(err) => {
+                            error!("Failed to receive clone request: {:?}", err);
+                        }
+                    }
                 }
                 _ = cancellation_token.cancelled() => {
                     return;
@@ -299,7 +291,7 @@ where
         // TODO(GabrielePicco): Make the concurrency configurable
         let result = stream
             .map(Ok::<_, AccountClonerError>)
-            .try_for_each_concurrent(30, |(pubkey, owner)| async move {
+            .try_for_each_concurrent(10, |(pubkey, owner)| async move {
                 trace!("Hydrating '{}'", pubkey);
                 let res = self
                     .do_clone_and_update_cache(

magicblock-account-cloner/src/remote_account_cloner_worker.rs

@@ -55,10 +55,11 @@ fn setup_custom(
     let cancellation_token = CancellationToken::new();
     let cloner_worker_handle = {
         let cloner_cancellation_token = cancellation_token.clone();
-        tokio::spawn(
+        tokio::spawn(async move {
             cloner_worker
-                .start_clone_request_processing(cloner_cancellation_token),
-        )
+                .start_clone_request_processing(cloner_cancellation_token)
+                .await
+        })
     };
     // Ready to run
     (cloner_client, cancellation_token, cloner_worker_handle)

magicblock-account-cloner/tests/remote_account_cloner.rs

@@ -62,10 +62,11 @@ fn setup_with_lifecycle(
         RemoteAccountClonerClient::new(&remote_account_cloner_worker);
     let remote_account_cloner_worker_handle = {
         let cloner_cancellation_token = cancellation_token.clone();
-        tokio::spawn(
+        tokio::spawn(async move {
             remote_account_cloner_worker
-                .start_clone_request_processing(cloner_cancellation_token),
-        )
+                .start_clone_request_processing(cloner_cancellation_token)
+                .await
+        })
     };
 
     let external_account_manager = ExternalAccountsManager {

magicblock-accounts/tests/ensure_accounts.rs

@@ -138,11 +138,13 @@ pub struct MagicValidator {
     remote_account_updates_worker: Option<RemoteAccountUpdatesWorker>,
     remote_account_updates_handle: Option<tokio::task::JoinHandle<()>>,
     remote_account_cloner_worker: Option<
-        RemoteAccountClonerWorker<
-            BankAccountProvider,
-            RemoteAccountFetcherClient,
-            RemoteAccountUpdatesClient,
-            AccountDumperBank,
+        Arc<
+            RemoteAccountClonerWorker<
+                BankAccountProvider,
+                RemoteAccountFetcherClient,
+                RemoteAccountUpdatesClient,
+                AccountDumperBank,
+            >,
         >,
     >,
     remote_account_cloner_handle: Option<tokio::task::JoinHandle<()>>,
@@ -354,7 +356,9 @@ impl MagicValidator {
             remote_account_fetcher_handle: None,
             remote_account_updates_worker: Some(remote_account_updates_worker),
             remote_account_updates_handle: None,
-            remote_account_cloner_worker: Some(remote_account_cloner_worker),
+            remote_account_cloner_worker: Some(Arc::new(
+                remote_account_cloner_worker,
+            )),
             remote_account_cloner_handle: None,
             pubsub_handle: Default::default(),
             pubsub_close_handle: Default::default(),
@@ -720,8 +724,20 @@ impl MagicValidator {
             self.remote_account_cloner_worker.take()
         {
             if !self.config.ledger.reset {
-                remote_account_cloner_worker.hydrate().await?;
-                info!("Validator hydration complete (bank hydrate, replay, account clone)");
+                let remote_account_cloner_worker =
+                    remote_account_cloner_worker.clone();
+                tokio::spawn(async move {
+                    let _ = remote_account_cloner_worker
+                        .hydrate()
+                        .await
+                        .inspect_err(|err| {
+                            error!(
+                                "Failed to hydrate validator accounts: {:?}",
+                                err
+                            );
+                        });
+                    info!("Validator hydration complete (bank hydrate, replay, account clone)");
+                });
             }
 
             let cancellation_token = self.token.clone();

magicblock-api/src/magic_validator.rs

@@ -14,12 +14,12 @@ bs58 = { workspace = true }
 expiring-hashmap = { workspace = true }
 geyser-grpc-proto = { workspace = true }
 hostname = { workspace = true }
-flume = "0.11"
+flume = { workspace = true }
 log = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 magicblock-transaction-status = { workspace = true }
-scc = "2.3" 
+scc = "2.3"
 solana-geyser-plugin-interface = { workspace = true }
 solana-sdk = { workspace = true }
 spl-token-2022 = { workspace = true, features = ["no-entrypoint"] }