diff --git a/Cargo.lock b/Cargo.lock index 7cda7836..0d1cc4e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12399,6 +12399,7 @@ name = "surfpool-types" version = "1.0.0-rc1" dependencies = [ "anchor-lang-idl", + "base64 0.22.1", "blake3", "chrono", "crossbeam-channel", diff --git a/crates/core/src/rpc/full.rs b/crates/core/src/rpc/full.rs index 03e04eb3..84e2f181 100644 --- a/crates/core/src/rpc/full.rs +++ b/crates/core/src/rpc/full.rs @@ -2626,7 +2626,7 @@ mod tests { setup .context .svm_locker - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); let res = setup @@ -3386,7 +3386,7 @@ mod tests { setup .context .svm_locker - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); } @@ -3431,7 +3431,7 @@ mod tests { setup .context .svm_locker - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); } diff --git a/crates/core/src/rpc/surfnet_cheatcodes.rs b/crates/core/src/rpc/surfnet_cheatcodes.rs index 16eca302..cd86bdfc 100644 --- a/crates/core/src/rpc/surfnet_cheatcodes.rs +++ b/crates/core/src/rpc/surfnet_cheatcodes.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use base64::{Engine as _, engine::general_purpose::STANDARD}; use jsonrpc_core::{BoxFuture, Error, Result, futures::future}; use jsonrpc_derive::rpc; @@ -14,9 +12,9 @@ use solana_system_interface::program as system_program; use solana_transaction::versioned::VersionedTransaction; use spl_associated_token_account_interface::address::get_associated_token_address_with_program_id; use surfpool_types::{ - AccountSnapshot, ClockCommand, ExportSnapshotConfig, GetStreamedAccountsResponse, - GetSurfnetInfoResponse, Idl, ResetAccountConfig, RpcProfileResultConfig, Scenario, - SimnetCommand, SimnetEvent, StreamAccountConfig, UiKeyedProfileResult, + ClockCommand, ExportSnapshotConfig, GetStreamedAccountsResponse, GetSurfnetInfoResponse, Idl, + ResetAccountConfig, RpcProfileResultConfig, Scenario, SimnetCommand, SimnetEvent, + SnapshotResult, StreamAccountConfig, UiKeyedProfileResult, types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature}, }; @@ -841,7 +839,7 @@ pub trait SurfnetCheatcodes { &self, meta: Self::Metadata, config: Option, - ) -> Result>>; + ) -> BoxFuture>>; /// A cheat code to simulate account streaming. /// When a transaction is processed, the accounts that are accessed are downloaded from the datasource and cached in the SVM. @@ -1127,7 +1125,7 @@ pub trait SurfnetCheatcodes { meta: Self::Metadata, scenario: Scenario, slot: Option, - ) -> Result>; + ) -> BoxFuture>>; } #[derive(Clone)] @@ -1794,13 +1792,30 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc { &self, meta: Self::Metadata, config: Option, - ) -> Result>> { + ) -> BoxFuture>> { let config = config.unwrap_or_default(); - let svm_locker = meta.get_svm_locker()?; - let snapshot = svm_locker.export_snapshot(config); - Ok(RpcResponse { - context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), - value: snapshot, + + Box::pin(async move { + let SurfnetRpcContext { + svm_locker, + remote_ctx, + } = meta.get_rpc_context(CommitmentConfig::confirmed())?; + + match &config.scope { + surfpool_types::ExportSnapshotScope::Network => todo!(), + surfpool_types::ExportSnapshotScope::PreTransaction(_) => todo!(), + surfpool_types::ExportSnapshotScope::Scenario(scenario) => { + svm_locker + .fetch_scenario_override_accounts(&remote_ctx, &scenario) + .await?; + } + }; + + let snapshot = svm_locker.export_snapshot(config); + Ok(RpcResponse { + context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), + value: snapshot, + }) }) } @@ -1809,18 +1824,29 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc { meta: Self::Metadata, scenario: Scenario, slot: Option, - ) -> Result> { - let svm_locker = meta.get_svm_locker()?; - svm_locker.register_scenario(scenario, slot)?; - Ok(RpcResponse { - context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), - value: (), + ) -> BoxFuture>> { + Box::pin(async move { + let SurfnetRpcContext { + svm_locker, + remote_ctx, + } = meta.get_rpc_context(CommitmentConfig::confirmed())?; + svm_locker + .fetch_scenario_override_accounts(&remote_ctx, &scenario) + .await?; + + svm_locker.register_scenario(scenario, slot)?; + Ok(RpcResponse { + context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), + value: (), + }) }) } } #[cfg(test)] mod tests { + use std::collections::BTreeMap; + use solana_account_decoder::{ UiAccountData, UiAccountEncoding, parse_account_data::ParsedAccount, }; @@ -1837,8 +1863,8 @@ mod tests { use spl_token_2022_interface::instruction::{initialize_mint2, mint_to, transfer_checked}; use spl_token_interface::state::Mint; use surfpool_types::{ - ExportSnapshotFilter, ExportSnapshotScope, RpcProfileDepth, UiAccountChange, - UiAccountProfileState, + AccountSnapshot, ExportSnapshotFilter, ExportSnapshotScope, RpcProfileDepth, + UiAccountChange, UiAccountProfileState, }; use super::*; @@ -2676,8 +2702,8 @@ mod tests { assert_eq!(expected_account.rent_epoch, account.rent_epoch); } - #[test] - fn test_export_snapshot() { + #[tokio::test] + async fn test_export_snapshot() { let client = TestSetup::new(SurfnetCheatcodesRpc); let pubkey1 = Pubkey::new_unique(); @@ -2705,15 +2731,18 @@ mod tests { let snapshot = client .rpc .export_snapshot(Some(client.context.clone()), None) + .await .expect("Failed to export snapshot") .value; + let snapshot = snapshot.as_accounts().unwrap(); + verify_snapshot_account(&snapshot, &pubkey1, &account1); verify_snapshot_account(&snapshot, &pubkey2, &account2); } - #[test] - fn test_export_snapshot_json_parsed() { + #[tokio::test] + async fn test_export_snapshot_json_parsed() { let client = TestSetup::new(SurfnetCheatcodesRpc); let pubkey1 = Pubkey::new_unique(); @@ -2762,10 +2791,14 @@ mod tests { scope: ExportSnapshotScope::Network, }), ) + .await .expect("Failed to export snapshot") .value; + let snapshot = snapshot.as_accounts().unwrap(); + verify_snapshot_account(&snapshot, &pubkey1, &account1); + let actual_account1 = snapshot .get(&pubkey1.to_string()) .expect("Account fixture not found"); @@ -2803,8 +2836,8 @@ mod tests { ); } - #[test] - fn test_export_snapshot_pre_transaction() { + #[tokio::test] + async fn test_export_snapshot_pre_transaction() { use std::collections::HashMap; use solana_signature::Signature; @@ -2890,9 +2923,12 @@ mod tests { scope: ExportSnapshotScope::PreTransaction(signature.to_string()), }), ) + .await .expect("Failed to export snapshot") .value; + let snapshot = snapshot.as_accounts().unwrap(); + // Verify that only account1 and account2 are in the snapshot assert!( snapshot.contains_key(&account1_pubkey.to_string()), @@ -2937,8 +2973,8 @@ mod tests { ); } - #[test] - fn test_export_snapshot_filtering() { + #[tokio::test] + async fn test_export_snapshot_filtering() { let system_account_pubkey = Pubkey::new_unique(); println!("System Account Pubkey: {}", system_account_pubkey); let excluded_system_account_pubkey = Pubkey::new_unique(); @@ -2979,8 +3015,12 @@ mod tests { let snapshot = client .rpc .export_snapshot(Some(client.context.clone()), None) + .await .expect("Failed to export snapshot") .value; + + let snapshot = snapshot.as_accounts().unwrap(); + assert!( !snapshot.contains_key(&program_account_pubkey.to_string()), "Program account should be excluded by default" @@ -3001,8 +3041,12 @@ mod tests { ..Default::default() }), ) + .await .expect("Failed to export snapshot") .value; + + let snapshot = snapshot.as_accounts().unwrap(); + assert!( !snapshot.contains_key(&program_account_pubkey.to_string()), "Program account should be excluded by default" @@ -3025,8 +3069,10 @@ mod tests { ..Default::default() }), ) + .await .expect("Failed to export snapshot") .value; + let snapshot = snapshot.as_accounts().unwrap(); assert!( snapshot.contains_key(&program_account_pubkey.to_string()), diff --git a/crates/core/src/runloops/mod.rs b/crates/core/src/runloops/mod.rs index d06cb8f6..2a5a2abe 100644 --- a/crates/core/src/runloops/mod.rs +++ b/crates/core/src/runloops/mod.rs @@ -262,7 +262,7 @@ pub async fn start_block_production_runloop( } SimnetCommand::UpdateInternalClock(_, clock) => { // Confirm the current block to materialize any scheduled overrides for this slot - if let Err(e) = svm_locker.confirm_current_block(&remote_client_with_commitment).await { + if let Err(e) = svm_locker.confirm_current_block().await { let _ = svm_locker.simnet_events_tx().send(SimnetEvent::error(format!( "Failed to confirm block after time travel: {}", e ))); @@ -281,7 +281,7 @@ pub async fn start_block_production_runloop( } SimnetCommand::UpdateInternalClockWithConfirmation(_, clock, response_tx) => { // Confirm the current block to materialize any scheduled overrides for this slot - if let Err(e) = svm_locker.confirm_current_block(&remote_client_with_commitment).await { + if let Err(e) = svm_locker.confirm_current_block().await { let _ = svm_locker.simnet_events_tx().send(SimnetEvent::error(format!( "Failed to confirm block after time travel: {}", e ))); @@ -348,9 +348,7 @@ pub async fn start_block_production_runloop( { if do_produce_block { - svm_locker - .confirm_current_block(&remote_client_with_commitment) - .await?; + svm_locker.confirm_current_block().await?; } } } diff --git a/crates/core/src/surfnet/locker.rs b/crates/core/src/surfnet/locker.rs index 0dcc7d95..183fd4b6 100644 --- a/crates/core/src/surfnet/locker.rs +++ b/crates/core/src/surfnet/locker.rs @@ -1,5 +1,6 @@ use std::{ collections::{BTreeMap, HashMap, HashSet}, + str::FromStr, sync::Arc, time::SystemTime, }; @@ -52,10 +53,10 @@ use solana_transaction_status::{ UiTransactionEncoding, }; use surfpool_types::{ - AccountSnapshot, ComputeUnitsEstimationResult, ExecutionCapture, ExportSnapshotConfig, Idl, - KeyedProfileResult, ProfileResult, RpcProfileResultConfig, RunbookExecutionStatusReport, - SimnetCommand, SimnetEvent, TransactionConfirmationStatus, TransactionStatusEvent, - UiKeyedProfileResult, UuidOrSignature, VersionedIdl, + ComputeUnitsEstimationResult, ExecutionCapture, ExportSnapshotConfig, Idl, KeyedProfileResult, + ProfileResult, RpcProfileResultConfig, RunbookExecutionStatusReport, Scenario, SimnetCommand, + SimnetEvent, SnapshotResult, TimeseriesSurfnetSnapshot, TransactionConfirmationStatus, + TransactionStatusEvent, UiKeyedProfileResult, UuidOrSignature, VersionedIdl, }; use tokio::sync::RwLock; use txtx_addon_kit::indexmap::IndexSet; @@ -2834,15 +2835,12 @@ impl SurfnetSvmLocker { } /// Confirms the current block on the underlying SVM, returning `Ok(())` or an error. - pub async fn confirm_current_block( - &self, - remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>, - ) -> SurfpoolResult<()> { + pub async fn confirm_current_block(&self) -> SurfpoolResult<()> { // Acquire write lock once and do both operations atomically // This prevents lock contention and potential deadlocks from mixing blocking and async locks let mut svm_writer = self.0.write().await; svm_writer.confirm_current_block()?; - svm_writer.materialize_overrides(remote_ctx).await + svm_writer.apply_overrides() } /// Subscribes for signature updates (confirmed/finalized) and returns a receiver of events. @@ -2908,11 +2906,101 @@ impl SurfnetSvmLocker { }); } - pub fn export_snapshot( + pub fn export_snapshot(&self, config: ExportSnapshotConfig) -> SnapshotResult { + self.with_svm_reader(|svm_reader| svm_reader.export_snapshot(config)) + } + + pub async fn fetch_scenario_override_accounts( + &self, + remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>, + scenario: &Scenario, + ) -> SurfpoolResult<()> { + for override_instance in scenario.overrides.iter() { + if !override_instance.enabled { + debug!("Skipping disabled override: {}", override_instance.id); + continue; + } + // Resolve account address + let account_pubkey = match &override_instance.account { + surfpool_types::AccountAddress::Pubkey(pubkey_str) => { + match Pubkey::from_str(pubkey_str) { + Ok(pubkey) => pubkey, + Err(e) => { + warn!( + "Failed to parse pubkey '{}' for override {}: {}", + pubkey_str, override_instance.id, e + ); + continue; + } + } + } + surfpool_types::AccountAddress::Pda { + program_id: _, + seeds: _, + } => unimplemented!(), + }; + + // Fetch fresh account data from remote if requested + if override_instance.fetch_before_use { + if let Some((client, _)) = remote_ctx { + debug!( + "Fetching fresh account data for {} from remote", + account_pubkey + ); + + match client + .get_account(&account_pubkey, CommitmentConfig::confirmed()) + .await + { + Ok(GetAccountResult::FoundAccount(_pubkey, remote_account, _)) => { + debug!( + "Fetched account {} from remote: {} lamports, {} bytes", + account_pubkey, + remote_account.lamports(), + remote_account.data().len() + ); + + // Set the fresh account data in the SVM + self.with_svm_writer(|svm_writer| { + if let Err(e) = + svm_writer.set_account(&account_pubkey, remote_account) + { + warn!( + "Failed to set account {} from remote: {}", + account_pubkey, e + ); + } + }); + } + Ok(GetAccountResult::None(_)) => { + debug!("Account {} not found on remote", account_pubkey); + } + Ok(_) => { + debug!("Account {} fetched (other variant)", account_pubkey); + } + Err(e) => { + warn!( + "Failed to fetch account {} from remote: {}", + account_pubkey, e + ); + } + } + } else { + debug!( + "fetch_before_use enabled but no remote client available for override {}", + override_instance.id + ); + } + } + } + Ok(()) + } + pub fn export_scenario_snapshot( &self, + scenario_id: String, config: ExportSnapshotConfig, - ) -> BTreeMap { - self.with_svm_reader(|svm_reader| svm_reader.export_snapshot(config)) + ) -> SurfpoolResult { + self.with_svm_writer(|svm_writer| svm_writer.snapshot_overrides(scenario_id, config)) } pub fn get_start_time(&self) -> SystemTime { diff --git a/crates/core/src/surfnet/svm.rs b/crates/core/src/surfnet/svm.rs index a07fa38e..5895c5f6 100644 --- a/crates/core/src/surfnet/svm.rs +++ b/crates/core/src/surfnet/svm.rs @@ -45,7 +45,7 @@ use solana_client::{ rpc_response::{RpcKeyedAccount, RpcLogsResponse, RpcPerfSample}, }; use solana_clock::{Clock, Slot}; -use solana_commitment_config::{CommitmentConfig, CommitmentLevel}; +use solana_commitment_config::CommitmentLevel; use solana_epoch_info::EpochInfo; use solana_feature_gate_interface::Feature; use solana_genesis_config::GenesisConfig; @@ -70,9 +70,9 @@ use spl_token_2022_interface::extension::{ }; use surfpool_types::{ AccountChange, AccountProfileState, AccountSnapshot, DEFAULT_PROFILING_MAP_CAPACITY, - DEFAULT_SLOT_TIME_MS, ExportSnapshotConfig, ExportSnapshotScope, FifoMap, Idl, - OverrideInstance, ProfileResult, RpcProfileDepth, RpcProfileResultConfig, - RunbookExecutionStatusReport, SimnetEvent, SvmFeature, SvmFeatureConfig, + DEFAULT_SLOT_TIME_MS, ExportSnapshotConfig, ExportSnapshotScope, FifoMap, Idl, ProfileResult, + RpcProfileDepth, RpcProfileResultConfig, RunbookExecutionStatusReport, SimnetEvent, + SnapshotResult, SvmFeature, SvmFeatureConfig, TimeseriesSurfnetSnapshot, TransactionConfirmationStatus, TransactionStatusEvent, UiAccountChange, UiAccountProfileState, UiProfileResult, VersionedIdl, types::{ @@ -270,7 +270,7 @@ pub struct SurfnetSvm { pub account_update_slots: HashMap, pub streamed_accounts: HashMap, pub recent_blockhashes: VecDeque<(SyntheticBlockhash, i64)>, - pub scheduled_overrides: HashMap>, + pub scheduled_overrides: HashMap>, /// Tracks accounts that have been explicitly closed by the user. /// These accounts will not be fetched from mainnet even if they don't exist in the local cache. pub closed_accounts: HashSet, @@ -1460,196 +1460,235 @@ impl SurfnetSvm { Ok(()) } - /// Materializes scheduled overrides for the current slot - /// - /// This function: - /// 1. Dequeues overrides scheduled for the current slot - /// 2. Resolves account addresses (Pubkey or PDA) - /// 3. Optionally fetches fresh account data from remote if `fetch_before_use` is enabled - /// 4. Applies the overrides to the account data - /// 5. Updates the SVM state - pub async fn materialize_overrides( + pub fn snapshot_overrides( &mut self, - remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>, - ) -> SurfpoolResult<()> { - let current_slot = self.latest_epoch_info.absolute_slot; - - // Remove and get overrides for this slot - let Some(overrides) = self.scheduled_overrides.remove(¤t_slot) else { - // No overrides for this slot - return Ok(()); + scenario_id: String, + config: ExportSnapshotConfig, + ) -> SurfpoolResult { + let mut result = BTreeMap::new(); + let encoding = if config.include_parsed_accounts.unwrap_or_default() { + UiAccountEncoding::JsonParsed + } else { + UiAccountEncoding::Base64 }; + debug!("\n\nCompiling overrides for scenario: {}", scenario_id); debug!( - "Materializing {} override(s) for slot {}", - overrides.len(), - current_slot + "Total scheduled overrides to process: {}", + self.scheduled_overrides + .values() + .map(|overrides| overrides.len()) + .sum::() ); + let overrides_for_scenario: HashMap<_, _> = self + .scheduled_overrides + .iter() + .map(|(slot, overrides)| { + ( + *slot, + overrides + .iter() + .filter_map(|(override_scenario_id, account_pubkey, modified_account)| { + if scenario_id.eq(override_scenario_id) { + Some((account_pubkey, modified_account)) + } else { + None + } + }) + .collect::>(), + ) + }) + .collect(); - for override_instance in overrides { - if !override_instance.enabled { - debug!("Skipping disabled override: {}", override_instance.id); - continue; + let first_slot = overrides_for_scenario.keys().min().cloned().unwrap_or(0); + + for (override_slot, overrides) in overrides_for_scenario.into_iter() { + let relative_slot = override_slot.saturating_sub(first_slot); + for (pubkey, account) in overrides { + // For token accounts, we need to provide the mint additional data + let additional_data = if account.owner == spl_token_interface::id() + || account.owner == spl_token_2022_interface::id() + { + if let Ok(token_account) = TokenAccount::unpack(&account.data) { + self.account_associated_data + .get(&token_account.mint()) + .cloned() + } else { + self.account_associated_data.get(pubkey).cloned() + } + } else { + self.account_associated_data.get(pubkey).cloned() + }; + + let ui_account = + self.encode_ui_account(pubkey, account, encoding, additional_data, None); + + let (base64, parsed_data) = match ui_account.data { + UiAccountData::Json(parsed_account) => { + (BASE64_STANDARD.encode(account.data()), Some(parsed_account)) + } + UiAccountData::Binary(base64, _) => (base64, None), + UiAccountData::LegacyBinary(_) => unreachable!(), + }; + + let account_snapshot = AccountSnapshot::new( + account.lamports, + account.owner.to_string(), + account.executable, + account.rent_epoch, + base64, + parsed_data, + ); + + result + .entry(relative_slot) + .or_insert_with(BTreeMap::new) + .insert(pubkey.to_string(), account_snapshot); } + } - // Resolve account address - let account_pubkey = match &override_instance.account { - surfpool_types::AccountAddress::Pubkey(pubkey_str) => { - match Pubkey::from_str(pubkey_str) { - Ok(pubkey) => pubkey, - Err(e) => { - warn!( - "Failed to parse pubkey '{}' for override {}: {}", - pubkey_str, override_instance.id, e - ); - continue; - } + Ok(result) + } + + pub fn compile_override_instance( + &self, + override_instance: &surfpool_types::OverrideInstance, + ) -> SurfpoolResult> { + // Resolve account address + let account_pubkey = match &override_instance.account { + surfpool_types::AccountAddress::Pubkey(pubkey_str) => { + match Pubkey::from_str(pubkey_str) { + Ok(pubkey) => pubkey, + Err(e) => { + warn!( + "Failed to parse pubkey '{}' for override {}: {}", + pubkey_str, override_instance.id, e + ); + return Ok(None); } } - surfpool_types::AccountAddress::Pda { - program_id: _, - seeds: _, - } => unimplemented!(), - }; + } + surfpool_types::AccountAddress::Pda { + program_id: _, + seeds: _, + } => unimplemented!(), + }; + debug!( + "Processing override {} for account {} (label: {:?})", + override_instance.id, account_pubkey, override_instance.label + ); + + // Apply the override values to the account data + if !override_instance.values.is_empty() { debug!( - "Processing override {} for account {} (label: {:?})", - override_instance.id, account_pubkey, override_instance.label + "Override {} applying {} field modification(s) to account {}", + override_instance.id, + override_instance.values.len(), + account_pubkey ); - // Fetch fresh account data from remote if requested - if override_instance.fetch_before_use { - if let Some((client, _)) = remote_ctx { - debug!( - "Fetching fresh account data for {} from remote", - account_pubkey - ); + // Get the account from the SVM + let Some(account) = self.inner.get_account(&account_pubkey) else { + warn!( + "Account {} not found in SVM for override {}, skipping modifications", + account_pubkey, override_instance.id + ); + return Ok(None); + }; - match client - .get_account(&account_pubkey, CommitmentConfig::confirmed()) - .await - { - Ok(GetAccountResult::FoundAccount(_pubkey, remote_account, _)) => { - debug!( - "Fetched account {} from remote: {} lamports, {} bytes", - account_pubkey, - remote_account.lamports(), - remote_account.data().len() - ); - - // Set the fresh account data in the SVM - if let Err(e) = self.inner.set_account(account_pubkey, remote_account) { - warn!( - "Failed to set account {} from remote: {}", - account_pubkey, e - ); - } - } - Ok(GetAccountResult::None(_)) => { - debug!("Account {} not found on remote", account_pubkey); - } - Ok(_) => { - debug!("Account {} fetched (other variant)", account_pubkey); - } - Err(e) => { - warn!( - "Failed to fetch account {} from remote: {}", - account_pubkey, e - ); - } - } - } else { - debug!( - "fetch_before_use enabled but no remote client available for override {}", - override_instance.id - ); - } - } + // Get the account owner (program ID) + let owner_program_id = account.owner(); - // Apply the override values to the account data - if !override_instance.values.is_empty() { - debug!( - "Override {} applying {} field modification(s) to account {}", - override_instance.id, - override_instance.values.len(), - account_pubkey + // Look up the IDL for the owner program + let Some(idl_versions) = self.registered_idls.get(owner_program_id) else { + warn!( + "No IDL registered for program {} (owner of account {}), skipping override {}", + owner_program_id, account_pubkey, override_instance.id ); + return Ok(None); + }; - // Get the account from the SVM - let Some(account) = self.inner.get_account(&account_pubkey) else { - warn!( - "Account {} not found in SVM for override {}, skipping modifications", - account_pubkey, override_instance.id - ); - continue; - }; + // Get the latest IDL version + let Some(versioned_idl) = idl_versions.peek() else { + warn!( + "IDL versions empty for program {}, skipping override {}", + owner_program_id, override_instance.id + ); + return Ok(None); + }; - // Get the account owner (program ID) - let owner_program_id = account.owner(); + let idl = &versioned_idl.1; - // Look up the IDL for the owner program - let Some(idl_versions) = self.registered_idls.get(owner_program_id) else { - warn!( - "No IDL registered for program {} (owner of account {}), skipping override {}", - owner_program_id, account_pubkey, override_instance.id - ); - continue; - }; + // Get account data + let account_data = account.data(); - // Get the latest IDL version - let Some(versioned_idl) = idl_versions.peek() else { + // Use get_forged_account_data to apply the overrides + let new_account_data = match self.get_forged_account_data( + &account_pubkey, + account_data, + idl, + &override_instance.values, + ) { + Ok(data) => data, + Err(e) => { warn!( - "IDL versions empty for program {}, skipping override {}", - owner_program_id, override_instance.id + "Failed to forge account data for {} (override {}): {}", + account_pubkey, override_instance.id, e ); - continue; - }; + return Ok(None); + } + }; - let idl = &versioned_idl.1; + // Create a new account with modified data + let modified_account = Account { + lamports: account.lamports(), + data: new_account_data, + owner: *account.owner(), + executable: account.executable(), + rent_epoch: account.rent_epoch(), + }; + Ok(Some((account_pubkey, modified_account))) + } else { + Ok(None) + } + } - // Get account data - let account_data = account.data(); + /// Materializes scheduled overrides for the current slot + /// + /// This function: + /// 1. Dequeues overrides scheduled for the current slot + /// 2. Resolves account addresses (Pubkey or PDA) + /// 3. Optionally fetches fresh account data from remote if `fetch_before_use` is enabled + /// 4. Applies the overrides to the account data + /// 5. Updates the SVM state + pub fn apply_overrides(&mut self) -> SurfpoolResult<()> { + let current_slot = self.latest_epoch_info.absolute_slot; - // Use get_forged_account_data to apply the overrides - let new_account_data = match self.get_forged_account_data( - &account_pubkey, - account_data, - idl, - &override_instance.values, - ) { - Ok(data) => data, - Err(e) => { - warn!( - "Failed to forge account data for {} (override {}): {}", - account_pubkey, override_instance.id, e - ); - continue; - } - }; + // Remove and get overrides for this slot + let Some(overrides) = self.scheduled_overrides.remove(¤t_slot) else { + // No overrides for this slot + return Ok(()); + }; - // Create a new account with modified data - let modified_account = Account { - lamports: account.lamports(), - data: new_account_data, - owner: *account.owner(), - executable: account.executable(), - rent_epoch: account.rent_epoch(), - }; + debug!( + "Materializing {} override(s) for slot {}", + overrides.len(), + current_slot + ); - // Update the account in the SVM - if let Err(e) = self.inner.set_account(account_pubkey, modified_account) { - warn!( - "Failed to set modified account {} in SVM: {}", - account_pubkey, e - ); - } else { - debug!( - "Successfully applied {} override(s) to account {} (override {})", - override_instance.values.len(), - account_pubkey, - override_instance.id - ); - } + for (scenario_id, account_pubkey, modified_account) in overrides { + // Update the account in the SVM + if let Err(e) = self.inner.set_account(account_pubkey, modified_account) { + warn!( + "Failed to set modified account {} in SVM: {}", + account_pubkey, e + ); + } else { + debug!( + "Successfully applied override to account {} (scenario {})", + account_pubkey, scenario_id + ); } } @@ -2484,11 +2523,7 @@ impl SurfnetSvm { /// /// # Returns /// A BTreeMap of pubkey -> AccountFixture that can be serialized to JSON. - pub fn export_snapshot( - &self, - config: ExportSnapshotConfig, - ) -> BTreeMap { - let mut fixtures = BTreeMap::new(); + pub fn export_snapshot(&self, config: ExportSnapshotConfig) -> SnapshotResult { let encoding = if config.include_parsed_accounts.unwrap_or_default() { UiAccountEncoding::JsonParsed } else { @@ -2506,14 +2541,14 @@ impl SurfnetSvm { } // Helper function to process an account and add it to fixtures - let mut process_account = |pubkey: &Pubkey, account: &Account| { + let process_account = |pubkey: &Pubkey, account: &Account| -> Option { let is_include_account = include_accounts.iter().any(|k| k.eq(&pubkey.to_string())); let is_exclude_account = exclude_accounts.iter().any(|k| k.eq(&pubkey.to_string())); let is_program_account = is_program_account(&account.owner); if is_exclude_account || ((is_program_account && !include_program_accounts) && !is_include_account) { - return; + return None; } // For token accounts, we need to provide the mint additional data @@ -2542,27 +2577,30 @@ impl SurfnetSvm { UiAccountData::LegacyBinary(_) => unreachable!(), }; - let account_snapshot = AccountSnapshot::new( + Some(AccountSnapshot::new( account.lamports, account.owner.to_string(), account.executable, account.rent_epoch, base64, parsed_data, - ); - - fixtures.insert(pubkey.to_string(), account_snapshot); + )) }; match &config.scope { ExportSnapshotScope::Network => { + let mut fixtures = BTreeMap::new(); // Export all network accounts (current behavior) for (pubkey, account_shared_data) in self.iter_accounts() { let account = Account::from(account_shared_data.clone()); - process_account(&pubkey, &account); + if let Some(account_snapshot) = process_account(&pubkey, &account) { + fixtures.insert(pubkey.to_string(), account_snapshot); + } } + SnapshotResult::Accounts(fixtures) } ExportSnapshotScope::PreTransaction(signature_str) => { + let mut fixtures = BTreeMap::new(); // Export accounts from a specific transaction's pre-execution state if let Ok(signature) = Signature::from_str(signature_str) { if let Some(profile) = self.executed_transaction_profiles.get(&signature) { @@ -2572,20 +2610,40 @@ impl SurfnetSvm { &profile.transaction_profile.pre_execution_capture { if let Some(account) = account_opt { - process_account(pubkey, account); + if let Some(account_snapshot) = process_account(pubkey, account) { + fixtures.insert(pubkey.to_string(), account_snapshot); + } } } // Also collect readonly account states (these don't change) for (pubkey, account) in &profile.readonly_account_states { - process_account(pubkey, account); + if let Some(account_snapshot) = process_account(pubkey, account) { + fixtures.insert(pubkey.to_string(), account_snapshot); + } + } + } + } + SnapshotResult::Accounts(fixtures) + } + ExportSnapshotScope::Scenario(scenario) => { + let mut fixtures = BTreeMap::new(); + for override_instance in scenario.overrides.iter() { + if let Some((account_pubkey, account)) = self + .compile_override_instance(&override_instance) + .unwrap_or(None) + { + if let Some(account_snapshot) = process_account(&account_pubkey, &account) { + fixtures + .entry(override_instance.scenario_relative_slot) + .or_insert_with(BTreeMap::new) + .insert(account_pubkey.to_string(), account_snapshot); } } } + SnapshotResult::TimeseriesSurfnet(fixtures) } } - - fixtures } /// Registers a scenario for execution by scheduling its overrides @@ -2618,10 +2676,14 @@ impl SurfnetSvm { absolute_slot, base_slot, scenario_relative_slot ); - self.scheduled_overrides - .entry(absolute_slot) - .or_insert_with(Vec::new) - .push(override_instance); + if let Some((account_pubkey, account)) = + self.compile_override_instance(&override_instance)? + { + self.scheduled_overrides + .entry(absolute_slot) + .or_insert_with(Vec::new) + .push((scenario.id.clone(), account_pubkey, account)); + } } Ok(()) diff --git a/crates/core/src/tests/integration.rs b/crates/core/src/tests/integration.rs index f40c198e..20bc6a2d 100644 --- a/crates/core/src/tests/integration.rs +++ b/crates/core/src/tests/integration.rs @@ -2898,7 +2898,7 @@ async fn test_profile_transaction_versioned_message() { .airdrop(&payer.pubkey(), 2 * lamports_to_send) .unwrap(); - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); // Create a transfer instruction let instruction = transfer(&payer.pubkey(), &recipient, lamports_to_send); @@ -2970,7 +2970,7 @@ async fn test_get_local_signatures_without_limit() { .unwrap(); svm_locker_for_context - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); @@ -3002,7 +3002,7 @@ async fn test_get_local_signatures_without_limit() { .unwrap(); // Confirm the block after creating the account svm_locker_for_context - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); @@ -3023,7 +3023,7 @@ async fn test_get_local_signatures_without_limit() { // Confirm the current block to create a block with the transaction signature svm_locker_for_context - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); @@ -3069,7 +3069,7 @@ async fn test_get_local_signatures_with_limit() { .unwrap(); svm_locker_for_context - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); @@ -3112,7 +3112,7 @@ async fn test_get_local_signatures_with_limit() { // Confirm the current block to create a new block with this transaction svm_locker_for_context - .confirm_current_block(&None) + .confirm_current_block() .await .unwrap(); } @@ -4014,13 +4014,13 @@ async fn test_reset_streamed_account() { svm_locker.airdrop(&p1.pubkey(), LAMPORTS_PER_SOL).unwrap(); // account is created in the SVM println!("Airdropped SOL to p1"); - let _ = svm_locker.confirm_current_block(&None).await; + let _ = svm_locker.confirm_current_block().await; // Account still exists assert!(!svm_locker.get_account_local(&p1.pubkey()).inner.is_none()); svm_locker.stream_account(p1.pubkey(), false).unwrap(); - let _ = svm_locker.confirm_current_block(&None).await; + let _ = svm_locker.confirm_current_block().await; // Account is cleaned up as soon as the block is processed assert!( svm_locker.get_account_local(&p1.pubkey()).inner.is_none(), @@ -4066,13 +4066,13 @@ async fn test_reset_streamed_account_cascade() { assert!(!svm_locker.get_account_local(&owner).inner.is_none()); assert!(!svm_locker.get_account_local(&owned).inner.is_none()); - let _ = svm_locker.confirm_current_block(&None).await; + let _ = svm_locker.confirm_current_block().await; // Accounts still exists assert!(!svm_locker.get_account_local(&owner).inner.is_none()); assert!(!svm_locker.get_account_local(&owned).inner.is_none()); svm_locker.stream_account(owner, true).unwrap(); - let _ = svm_locker.confirm_current_block(&None).await; + let _ = svm_locker.confirm_current_block().await; // Owner is deleted, owned account is deleted assert!(svm_locker.get_account_local(&owner).inner.is_none()); @@ -4364,7 +4364,7 @@ async fn test_ws_signature_subscribe(subscription_type: SignatureSubscriptionTyp match subscription_type { SignatureSubscriptionType::Commitment(CommitmentLevel::Confirmed) => { // confirm the block to trigger confirmed notification - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); } _ => {} } @@ -4507,7 +4507,7 @@ async fn test_ws_signature_subscribe_multiple_subscribers() { ); // confirm the block for confirmed subscription - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); assert!( notification_rx3 .recv_timeout(Duration::from_secs(5)) @@ -4692,7 +4692,7 @@ async fn test_ws_account_subscribe_multiple_changes() { // confirm the block to get fresh blockhash for next transaction if i < 2 { - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); } } } @@ -4915,7 +4915,7 @@ async fn test_ws_slot_subscribe_manual_advancement() { let initial_slot = svm_locker.get_latest_absolute_slot(); // manually advance slot by confirming a block - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); // should receive slot update notification let slot_update = slot_rx.recv_timeout(Duration::from_secs(5)); @@ -4946,7 +4946,7 @@ async fn test_ws_slot_subscribe_multiple_subscribers() { let slot_rx3 = svm_locker.subscribe_for_slot_updates(); // advance slot - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); // all subscribers should receive notification assert!( @@ -4974,7 +4974,7 @@ async fn test_ws_slot_subscribe_multiple_slot_changes() { // advance slot multiple times for i in 0..3 { - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); let slot_update = slot_rx.recv_timeout(Duration::from_secs(5)); assert!( @@ -5172,7 +5172,7 @@ async fn test_ws_logs_subscribe_confirmed_commitment() { .unwrap(); // confirm the block to trigger confirmed logs - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); // wait for confirmed logs notification let logs_notification = logs_rx.recv_timeout(Duration::from_secs(5)); @@ -5234,11 +5234,11 @@ async fn test_ws_logs_subscribe_finalized_commitment() { .unwrap(); // confirm and finalize the block - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); // advance enough slots to trigger finalization for _ in 0..FINALIZATION_SLOT_THRESHOLD { - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); } // wait for finalized logs notification @@ -5375,7 +5375,7 @@ async fn test_ws_logs_subscribe_multiple_subscribers() { ); // confirm block for confirmed subscriber - svm_locker.confirm_current_block(&None).await.unwrap(); + svm_locker.confirm_current_block().await.unwrap(); assert!( logs_rx3.recv_timeout(Duration::from_secs(5)).is_ok(), "Confirmed subscriber should receive logs" diff --git a/crates/types/Cargo.toml b/crates/types/Cargo.toml index 2991c4e6..b963eb97 100644 --- a/crates/types/Cargo.toml +++ b/crates/types/Cargo.toml @@ -14,6 +14,7 @@ path = "src/lib.rs" [dependencies] anchor-lang-idl = { workspace = true } +base64 = { workspace = true } blake3 = { workspace = true } chrono = { workspace = true } crossbeam-channel = { workspace = true } diff --git a/crates/types/src/types.rs b/crates/types/src/types.rs index a93f430a..cf34dd6e 100644 --- a/crates/types/src/types.rs +++ b/crates/types/src/types.rs @@ -25,7 +25,7 @@ use txtx_addon_kit::indexmap::IndexMap; use txtx_addon_network_svm_types::subgraph::SubgraphRequest; use uuid::Uuid; -use crate::SvmFeatureConfig; +use crate::{Scenario, SvmFeatureConfig}; pub const DEFAULT_RPC_URL: &str = "https://api.mainnet-beta.solana.com"; pub const DEFAULT_RPC_PORT: u16 = 8899; @@ -997,6 +997,19 @@ pub struct AccountSnapshot { pub parsed_data: Option, } +impl Into for AccountSnapshot { + fn into(self) -> Account { + use base64::{Engine, prelude::BASE64_STANDARD}; + Account { + lamports: self.lamports, + data: BASE64_STANDARD.decode(self.data).unwrap_or_default(), + owner: self.owner.parse().unwrap_or_default(), + executable: self.executable, + rent_epoch: self.rent_epoch, + } + } +} + impl AccountSnapshot { pub fn new( lamports: u64, @@ -1017,6 +1030,32 @@ impl AccountSnapshot { } } +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", untagged)] +pub enum SnapshotResult { + Accounts(AccountsSnapshot), + TimeseriesSurfnet(TimeseriesSurfnetSnapshot), +} + +impl SnapshotResult { + pub fn as_accounts(&self) -> Option<&AccountsSnapshot> { + match self { + SnapshotResult::Accounts(accounts) => Some(accounts), + _ => None, + } + } + + pub fn as_timeseries_surfnet(&self) -> Option<&TimeseriesSurfnetSnapshot> { + match self { + SnapshotResult::TimeseriesSurfnet(timeseries) => Some(timeseries), + _ => None, + } + } +} + +pub type AccountsSnapshot = BTreeMap; +pub type TimeseriesSurfnetSnapshot = BTreeMap; + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExportSnapshotConfig { @@ -1031,6 +1070,7 @@ pub enum ExportSnapshotScope { #[default] Network, PreTransaction(String), + Scenario(Scenario), } #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]