feat(core): add support for loading account snapshots at startup (solana-foundation#464)

Author: MicaiahReid

crates/cli/src/cli/mod.rs

@@ -1,4 +1,4 @@
-use std::{env, fs::File, path::PathBuf, process, str::FromStr};
+use std::{collections::BTreeMap, env, fs::File, path::PathBuf, process, str::FromStr};
 
 use chrono::Local;
 use clap::{ArgAction, CommandFactory, Parser, Subcommand};
@@ -13,7 +13,7 @@ use solana_pubkey::Pubkey;
 use solana_signer::{EncodableKey, Signer};
 use surfpool_mcp::McpOptions;
 use surfpool_types::{
-    BlockProductionMode, CHANGE_TO_DEFAULT_STUDIO_PORT_ONCE_SUPERVISOR_MERGED,
+    AccountSnapshot, BlockProductionMode, CHANGE_TO_DEFAULT_STUDIO_PORT_ONCE_SUPERVISOR_MERGED,
     DEFAULT_NETWORK_HOST, DEFAULT_RPC_PORT, DEFAULT_SLOT_TIME_MS, DEFAULT_WS_PORT, RpcConfig,
     SimnetConfig, SimnetEvent, StudioConfig, SubgraphConfig, SurfpoolConfig, SvmFeature,
     SvmFeatureConfig,
@@ -245,6 +245,13 @@ pub struct StartSimnet {
     /// A set of inputs to use for the runbook (eg. surfpool start --runbook-input myInputs.json)
     #[arg(long = "runbook-input", short = 'i')]
     pub runbook_input: Vec<String>,
+    /// Path to JSON snapshot file(s) to preload accounts from. Can be specified multiple times.
+    /// (eg. surfpool start --snapshot ./snapshot1.json --snapshot ./snapshot2.json)
+    /// The snapshot format matches the output of surfnet_exportSnapshot RPC method.
+    /// Account values can be null to fetch the account from the remote RPC instead.
+    /// When multiple files are provided, later files override earlier ones for duplicate keys.
+    #[arg(long = "snapshot")]
+    pub snapshot: Vec<String>,
 }
 
 fn parse_svm_feature(s: &str) -> Result<SvmFeature, String> {
@@ -372,7 +379,11 @@ impl StartSimnet {
         config
     }
 
-    pub fn simnet_config(&self, airdrop_addresses: Vec<Pubkey>) -> SimnetConfig {
+    pub fn simnet_config(
+        &self,
+        airdrop_addresses: Vec<Pubkey>,
+        snapshot: BTreeMap<String, Option<AccountSnapshot>>,
+    ) -> SimnetConfig {
         let remote_rpc_url = if !self.offline {
             Some(self.datasource_rpc_url())
         } else {
@@ -396,6 +407,7 @@ impl StartSimnet {
             },
             feature_config: self.feature_config(),
             skip_signature_verification: false,
+            snapshot,
         }
     }
 
@@ -415,15 +427,19 @@ impl StartSimnet {
         SubgraphConfig {}
     }
 
-    pub fn surfpool_config(&self, airdrop_addresses: Vec<Pubkey>) -> SurfpoolConfig {
+    pub fn surfpool_config(
+        &self,
+        airdrop_addresses: Vec<Pubkey>,
+        snapshot: BTreeMap<String, Option<AccountSnapshot>>,
+    ) -> SurfpoolConfig {
         let plugin_config_path = self
             .plugin_config_path
             .iter()
             .map(PathBuf::from)
             .collect::<Vec<_>>();
 
         SurfpoolConfig {
-            simnets: vec![self.simnet_config(airdrop_addresses)],
+            simnets: vec![self.simnet_config(airdrop_addresses, snapshot)],
             rpc: self.rpc_config(),
             subgraph: self.subgraph_config(),
             studio: self.studio_config(),

crates/cli/src/cli/simnet/mod.rs

@@ -82,8 +82,39 @@ pub async fn handle_start_local_surfnet_command(
         Some(keypair)
     };
 
+    // Parse and merge snapshot files (multiple files supported, later files override earlier ones)
+    // The actual loading happens in the runloop after the locker is created
+    let snapshot = {
+        let mut merged_snapshot: std::collections::BTreeMap<
+            String,
+            Option<surfpool_types::AccountSnapshot>,
+        > = std::collections::BTreeMap::new();
+
+        for snapshot_path in &cmd.snapshot {
+            let file_location = FileLocation::from_path(std::path::PathBuf::from(snapshot_path));
+            let content = file_location
+                .read_content_as_utf8()
+                .map_err(|e| format!("Failed to read snapshot file '{}': {}", snapshot_path, e))?;
+            let snapshot_data: std::collections::BTreeMap<
+                String,
+                Option<surfpool_types::AccountSnapshot>,
+            > = serde_json::from_str(&content)
+                .map_err(|e| format!("Failed to parse snapshot JSON '{}': {}", snapshot_path, e))?;
+            let _ = simnet_events_tx.send(SimnetEvent::info(format!(
+                "Loaded {} accounts from snapshot file: {}",
+                snapshot_data.len(),
+                snapshot_path
+            )));
+
+            // Merge into the combined snapshot (later files override earlier ones)
+            merged_snapshot.extend(snapshot_data);
+        }
+
+        merged_snapshot
+    };
+
     // Build config
-    let config = cmd.surfpool_config(airdrop_addresses);
+    let config = cmd.surfpool_config(airdrop_addresses, snapshot);
 
     let studio_binding_address = config.studio.get_studio_base_url();
 
@@ -155,6 +186,8 @@ pub async fn handle_start_local_surfnet_command(
         })
         .map_err(|e| format!("{}", e))?;
 
+    // Collect events that occur before Ready so we can re-send them to the TUI
+    let mut early_events = Vec::new();
     loop {
         match simnet_events_rx.recv() {
             Ok(SimnetEvent::Aborted(error)) => {
@@ -163,10 +196,16 @@ pub async fn handle_start_local_surfnet_command(
             }
             Ok(SimnetEvent::Shutdown) => return Ok(()),
             Ok(SimnetEvent::Ready) => break,
-            _other => continue,
+            Ok(other) => early_events.push(other),
+            Err(_) => continue,
         }
     }
 
+    // Re-send early events (like snapshot loading messages) so the TUI receives them
+    for event in early_events {
+        let _ = simnet_events_tx.send(event);
+    }
+
     for event in airdrop_events {
         let _ = simnet_events_tx.send(event);
     }

crates/core/src/runloops/mod.rs

@@ -109,6 +109,36 @@ pub async fn start_local_surfnet_runloop(
         .await?;
 
     svm_locker.airdrop_pubkeys(simnet.airdrop_token_amount, &simnet.airdrop_addresses);
+
+    // Load snapshot accounts if provided
+    if !simnet.snapshot.is_empty() {
+        match svm_locker
+            .load_snapshot(
+                &simnet.snapshot,
+                remote_rpc_client.as_ref(),
+                CommitmentConfig::confirmed(),
+            )
+            .await
+        {
+            Ok(loaded_count) => {
+                let _ = svm_locker.with_svm_reader(|svm| {
+                    svm.simnet_events_tx.send(SimnetEvent::info(format!(
+                        "Preloaded {} accounts from snapshot(s) into SVM",
+                        loaded_count
+                    )))
+                });
+            }
+            Err(e) => {
+                let _ = svm_locker.with_svm_reader(|svm| {
+                    svm.simnet_events_tx.send(SimnetEvent::warn(format!(
+                        "Error loading snapshot accounts: {}",
+                        e
+                    )))
+                });
+            }
+        }
+    }
+
     let simnet_events_tx_cc = svm_locker.simnet_events_tx();
 
     let (plugin_manager_commands_rx, _rpc_handle, _ws_handle) = start_rpc_servers_runloop(
@@ -726,6 +756,40 @@ fn start_geyser_runloop(
                             }
                         }
                     }
+                    Ok(GeyserEvent::StartupAccountUpdate(account_update)) => {
+                        let GeyserAccountUpdate {
+                            pubkey,
+                            account,
+                            slot,
+                            sanitized_transaction,
+                            write_version,
+                        } = account_update;
+
+                        let account_replica = ReplicaAccountInfoV3 {
+                            pubkey: pubkey.as_ref(),
+                            lamports: account.lamports,
+                            owner: account.owner.as_ref(),
+                            executable: account.executable,
+                            rent_epoch: account.rent_epoch,
+                            data: account.data.as_ref(),
+                            write_version,
+                            txn: sanitized_transaction.as_ref(),
+                        };
+
+                        // Send startup account updates with is_startup=true
+                        for plugin in surfpool_plugin_manager.iter() {
+                            if let Err(e) = plugin.update_account(ReplicaAccountInfoVersions::V0_0_3(&account_replica), slot, true) {
+                                let _ = simnet_events_tx.send(SimnetEvent::error(format!("Failed to send startup account update to Geyser plugin: {:?}", e)));
+                            }
+                        }
+
+                        #[cfg(feature = "geyser_plugin")]
+                        for plugin in plugin_manager.plugins.iter() {
+                            if let Err(e) = plugin.update_account(ReplicaAccountInfoVersions::V0_0_3(&account_replica), slot, true) {
+                                let _ = simnet_events_tx.send(SimnetEvent::error(format!("Failed to send startup account update to Geyser plugin: {:?}", e)));
+                            }
+                        }
+                    }
                 }
             }
         };