diff --git a/crates/core/src/surfnet/locker.rs b/crates/core/src/surfnet/locker.rs index cd2da724..21bb64c4 100644 --- a/crates/core/src/surfnet/locker.rs +++ b/crates/core/src/surfnet/locker.rs @@ -69,13 +69,15 @@ use super::{ use crate::{ error::{AirdropError, SurfpoolError, SurfpoolResult}, helpers::time_travel::calculate_time_travel_clock, - rpc::full::{ - ComparisonFilter, RpcGetTransactionsForAddressConfig, RpcTransactionForAddressEntry, - RpcTransactionForAddressFullInfo, RpcTransactionForAddressSignatureInfo, - RpcTransactionsForAddressResult, SortOrder, TransactionsForAddressDetails, - TransactionsForAddressStatusFilter, TransactionsForAddressTokenFilter, + rpc::{ + full::{ + ComparisonFilter, RpcGetTransactionsForAddressConfig, RpcTransactionForAddressEntry, + RpcTransactionForAddressFullInfo, RpcTransactionForAddressSignatureInfo, + RpcTransactionsForAddressResult, SortOrder, TransactionsForAddressDetails, + TransactionsForAddressStatusFilter, TransactionsForAddressTokenFilter, + }, + utils::{convert_transaction_metadata_from_canonical, verify_pubkey}, }, - rpc::utils::{convert_transaction_metadata_from_canonical, verify_pubkey}, storage::StorageResult, surfnet::FINALIZATION_SLOT_THRESHOLD, types::{ @@ -225,6 +227,26 @@ impl SurfnetSvmLocker { }) } + /// Executes a write operation and captures the SVM context from the same lock scope. + fn with_contextualized_svm_writer(&self, writer: F) -> SvmAccessContext + where + F: FnOnce(&mut SurfnetSvm) -> T + Send + Sync, + T: Send + 'static, + { + let write_lock = self.0.clone(); + tokio::task::block_in_place(move || { + let mut write_guard = write_lock.blocking_write(); + let res = writer(&mut write_guard); + + SvmAccessContext::new( + write_guard.get_latest_absolute_slot(), + write_guard.latest_epoch_info(), + write_guard.latest_blockhash(), + res, + ) + }) + } + /// Executes a write operation on the underlying `SurfnetSvm` by acquiring a blocking write lock. /// Accepts a closure that receives a mutable reference to `SurfnetSvm` and returns a value. /// @@ -292,8 +314,19 @@ impl SurfnetSvmLocker { /// Retrieves a local account from the SVM cache, returning a contextualized result. pub fn get_account_local(&self, pubkey: &Pubkey) -> SvmAccessContext { - self.with_contextualized_svm_reader(|svm_reader| { - return svm_reader.inner.get_account_result(pubkey).unwrap(); + let result = self.with_contextualized_svm_reader(|svm_reader| { + svm_reader.inner.get_account_result(pubkey).unwrap() + }); + + if !result.inner.requires_update() { + return result; + } + + let pubkey = *pubkey; + self.with_contextualized_svm_writer(move |svm_writer| { + let update = svm_writer.inner.get_account_result(&pubkey).unwrap(); + svm_writer.write_account_update(update); + svm_writer.inner.get_account_result(&pubkey).unwrap() }) } @@ -314,13 +347,12 @@ impl SurfnetSvmLocker { if !is_offline { let offline_owners = self.get_offline_account_owners(); let remote_account = client.get_account(pubkey, commitment_config).await?; - Ok( - result.with_new_value(Self::filter_downloaded_account_result( - pubkey, - remote_account, - &offline_owners, - )), - ) + let remote_account = + Self::filter_downloaded_account_result(pubkey, remote_account, &offline_owners); + + Ok(self.with_contextualized_svm_writer(move |svm_writer| { + svm_writer.hydrate_account_update(remote_account) + })) } else { Ok(result) } @@ -356,15 +388,27 @@ impl SurfnetSvmLocker { &self, pubkeys: &[Pubkey], ) -> SvmAccessContext> { - self.with_contextualized_svm_reader(|svm_reader| { - let mut accounts = vec![]; + let results = self.with_contextualized_svm_reader(|svm_reader| { + pubkeys + .iter() + .map(|pubkey| svm_reader.inner.get_account_result(pubkey).unwrap()) + .collect::>() + }); - for pubkey in pubkeys { - let result = svm_reader.inner.get_account_result(pubkey).unwrap(); - if result.is_none() {}; - accounts.push(result); + if results.inner.iter().all(|result| !result.requires_update()) { + return results; + } + + let pubkeys = pubkeys.to_vec(); + self.with_contextualized_svm_writer(move |svm_writer| { + for pubkey in &pubkeys { + let update = svm_writer.inner.get_account_result(pubkey).unwrap(); + svm_writer.write_account_update(update); } - accounts + pubkeys + .iter() + .map(|pubkey| svm_writer.inner.get_account_result(pubkey).unwrap()) + .collect() }) } @@ -416,9 +460,8 @@ impl SurfnetSvmLocker { .get_multiple_accounts(&missing_accounts, commitment_config) .await?; - // Build map of pubkey -> remote result for O(1) lookup let offline_owners = self.get_offline_account_owners(); - let remote_map: HashMap = missing_accounts + let remote_results = missing_accounts .iter() .copied() .zip(remote_results.into_iter()) @@ -432,33 +475,30 @@ impl SurfnetSvmLocker { ), ) }) - .collect(); - - // Replace None entries with remote results while preserving order - // We iterate through original pubkeys array to ensure order is explicit - let combined_results: Vec = pubkeys - .iter() - .zip(local_results.into_iter()) - .map(|(pubkey, local_result)| { - match local_result { - GetAccountResult::None(_) => remote_map - .get(pubkey) - .cloned() - .unwrap_or(GetAccountResult::None(*pubkey)), - found => { - debug!("Keeping local account: {}", pubkey); - found - } // Keep found accounts (no clone, just move) - } - }) - .collect(); + .collect::>(); + + // Remote data is hydration, not an authoritative update. Recheck and insert each + // fetched account under the write lock so local writes that happened during the + // network request always win. Companion mint/program-data accounts follow the same rule. + let requested_pubkeys = pubkeys.to_vec(); + Ok(self.with_contextualized_svm_writer(move |svm_writer| { + let mut hydrated_results = HashMap::new(); + for (pubkey, update) in remote_results { + hydrated_results.insert(pubkey, svm_writer.hydrate_account_update(update)); + } - Ok(SvmAccessContext::new( - slot, - latest_epoch_info, - latest_blockhash, - combined_results, - )) + requested_pubkeys + .iter() + .map(|pubkey| { + let local_result = svm_writer.inner.get_account_result(pubkey).unwrap(); + if local_result.is_none() { + hydrated_results.remove(pubkey).unwrap_or(local_result) + } else { + local_result + } + }) + .collect() + })) } /// Retrieves multiple accounts, using local or remote context and applying factory defaults if provided. diff --git a/crates/core/src/surfnet/svm.rs b/crates/core/src/surfnet/svm.rs index d10bd46f..45897904 100644 --- a/crates/core/src/surfnet/svm.rs +++ b/crates/core/src/surfnet/svm.rs @@ -2251,47 +2251,11 @@ impl SurfnetSvm { /// # Arguments /// * `account_update` - The account update result to process. pub fn write_account_update(&mut self, account_update: GetAccountResult) { - let init_programdata_account = |program_account: &Account| { - if !program_account.executable { - return None; - } - if !program_account - .owner - .eq(&solana_sdk_ids::bpf_loader_upgradeable::id()) - { - return None; - } - let Ok(UpgradeableLoaderState::Program { - programdata_address, - }) = bincode::deserialize::(&program_account.data) - else { - return None; - }; - - let programdata_state = UpgradeableLoaderState::ProgramData { - upgrade_authority_address: Some(system_program::id()), - slot: self.get_latest_absolute_slot(), - }; - let mut data = bincode::serialize(&programdata_state).unwrap(); - - data.extend_from_slice(crate::surfnet::noop_program::NOOP_PROGRAM_ELF); - let lamports = self.inner.minimum_balance_for_rent_exemption(data.len()); - Some(( - programdata_address, - Account { - lamports, - data, - owner: solana_sdk_ids::bpf_loader_upgradeable::id(), - executable: false, - rent_epoch: 0, - }, - )) - }; match account_update { GetAccountResult::FoundAccount(pubkey, account, do_update_account) => { if do_update_account { if let Some((programdata_address, programdata_account)) = - init_programdata_account(&account) + self.default_programdata_account(&account) { match self.get_account(&programdata_address) { Ok(None) => { @@ -2320,7 +2284,7 @@ impl SurfnetSvm { } GetAccountResult::FoundProgramAccount((pubkey, account), (_, None)) => { if let Some((programdata_address, programdata_account)) = - init_programdata_account(&account) + self.default_programdata_account(&account) { match self.get_account(&programdata_address) { Ok(None) => { @@ -2377,6 +2341,182 @@ impl SurfnetSvm { } } + /// Hydrates an account fetched from a remote RPC without replacing local fork state. + /// + /// The local lookup and conditional insert happen while the caller holds the SVM write lock. + /// This closes the check/fetch/write race: a remote response can initialize an account that is + /// still absent, but it cannot replace an account written locally while the request was in + /// flight. Companion mint and program-data accounts are merged independently. + pub fn hydrate_account_update(&mut self, account_update: GetAccountResult) -> GetAccountResult { + let requested_pubkey = match &account_update { + GetAccountResult::None(pubkey) | GetAccountResult::FoundAccount(pubkey, _, _) => { + *pubkey + } + GetAccountResult::FoundProgramAccount((pubkey, _), _) + | GetAccountResult::FoundTokenAccount((pubkey, _), _) => *pubkey, + }; + let remote_read_result = match &account_update { + GetAccountResult::None(pubkey) => GetAccountResult::None(*pubkey), + GetAccountResult::FoundAccount(pubkey, account, _) => { + GetAccountResult::FoundAccount(*pubkey, account.clone(), false) + } + GetAccountResult::FoundProgramAccount((pubkey, account), _) + | GetAccountResult::FoundTokenAccount((pubkey, account), _) => { + GetAccountResult::FoundAccount(*pubkey, account.clone(), false) + } + }; + + // The requested account may have been written while the remote request was in flight. + // Re-read it under the write lock and materialize that complete local result (including + // any DB-backed companion accounts) instead of applying any part of the remote result. + match self.inner.get_account_result(&requested_pubkey) { + Ok(local_update) if !local_update.is_none() => { + self.write_account_update(local_update); + return match self.inner.get_account_result(&requested_pubkey) { + Ok(local_result) if !local_result.is_none() => local_result, + _ => remote_read_result, + }; + } + Ok(_) => {} + Err(e) => { + let _ = self + .simnet_events_tx + .send(SimnetEvent::error(e.to_string())); + return remote_read_result; + } + } + + match account_update { + GetAccountResult::FoundAccount(pubkey, account, do_update_account) => { + if do_update_account { + if let Some((programdata_address, programdata_account)) = + self.default_programdata_account(&account) + { + self.hydrate_account_if_missing(programdata_address, programdata_account); + } + self.hydrate_account_if_missing(pubkey, account); + } + } + GetAccountResult::FoundProgramAccount((pubkey, account), (_, None)) => { + if let Some((programdata_address, programdata_account)) = + self.default_programdata_account(&account) + { + self.hydrate_account_if_missing(programdata_address, programdata_account); + } + self.hydrate_account_if_missing(pubkey, account); + } + GetAccountResult::FoundTokenAccount((pubkey, account), (_, None)) => { + self.hydrate_account_if_missing(pubkey, account); + } + GetAccountResult::FoundProgramAccount( + (pubkey, account), + (coupled_pubkey, Some(coupled_account)), + ) + | GetAccountResult::FoundTokenAccount( + (pubkey, account), + (coupled_pubkey, Some(coupled_account)), + ) => { + self.hydrate_account_if_missing(coupled_pubkey, coupled_account); + self.hydrate_account_if_missing(pubkey, account); + } + GetAccountResult::None(_) => {} + } + + match self.inner.get_account_result(&requested_pubkey) { + Ok(local_result) if !local_result.is_none() => local_result, + _ => remote_read_result, + } + } + + fn hydrate_account_if_missing(&mut self, pubkey: Pubkey, remote_account: Account) { + if self.inner.get_account_no_db(&pubkey).is_some() { + debug!("Keeping local account {} during remote hydration", pubkey); + return; + } + + match self.inner.get_account(&pubkey) { + Ok(Some(local_account)) => { + // The account survived in the backing store after LiteSVM garbage collection. + // Materialize that local value rather than replacing it with remote state. + if let Err(e) = self.set_account(&pubkey, local_account) { + let _ = self + .simnet_events_tx + .send(SimnetEvent::error(e.to_string())); + } + } + Ok(None) => { + if self.is_remote_hydration_blocked(&pubkey, &remote_account) { + return; + } + if let Err(e) = self.set_account(&pubkey, remote_account) { + let _ = self + .simnet_events_tx + .send(SimnetEvent::error(e.to_string())); + } + } + Err(e) => { + let _ = self + .simnet_events_tx + .send(SimnetEvent::error(e.to_string())); + } + } + } + + fn is_remote_hydration_blocked(&self, pubkey: &Pubkey, account: &Account) -> bool { + if self + .offline_accounts + .contains_key(&pubkey.to_string()) + .unwrap_or(false) + { + return true; + } + + self.offline_accounts + .into_iter() + .map(|mut entries| { + entries.any(|(offline_pubkey, config)| { + config.include_owned_accounts + && Pubkey::from_str(&offline_pubkey) + .map(|owner| owner == account.owner) + .unwrap_or(false) + }) + }) + .unwrap_or(false) + } + + fn default_programdata_account(&self, program_account: &Account) -> Option<(Pubkey, Account)> { + if !program_account.executable + || program_account.owner != solana_sdk_ids::bpf_loader_upgradeable::id() + { + return None; + } + let Ok(UpgradeableLoaderState::Program { + programdata_address, + }) = bincode::deserialize::(&program_account.data) + else { + return None; + }; + + let programdata_state = UpgradeableLoaderState::ProgramData { + upgrade_authority_address: Some(system_program::id()), + slot: self.get_latest_absolute_slot(), + }; + let mut data = bincode::serialize(&programdata_state).unwrap(); + + data.extend_from_slice(crate::surfnet::noop_program::NOOP_PROGRAM_ELF); + let lamports = self.inner.minimum_balance_for_rent_exemption(data.len()); + Some(( + programdata_address, + Account { + lamports, + data, + owner: solana_sdk_ids::bpf_loader_upgradeable::id(), + executable: false, + rent_epoch: 0, + }, + )) + } + pub fn confirm_current_block(&mut self) -> SurfpoolResult<()> { let slot = self.get_latest_absolute_slot(); // `slotsUpdatesSubscribe` clients expect millisecond-precision Unix @@ -4414,6 +4554,97 @@ mod tests { ) } + #[test_case(TestType::sqlite(); "with on-disk sqlite db")] + #[test_case(TestType::in_memory(); "with in-memory sqlite db")] + #[test_case(TestType::no_db(); "with no db")] + #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] + fn test_remote_hydration_preserves_local_account_state(test_type: TestType) { + let (mut svm, _events_rx, _geyser_rx) = test_type.initialize_svm(); + let pubkey = Pubkey::new_unique(); + let owner = Pubkey::new_unique(); + let local_account = Account { + lamports: 2, + data: b"local-state".to_vec(), + owner, + executable: false, + rent_epoch: 0, + }; + let remote_account = Account { + lamports: 1, + data: b"remote-state".to_vec(), + owner, + executable: false, + rent_epoch: 0, + }; + + svm.set_account(&pubkey, local_account.clone()).unwrap(); + svm.hydrate_account_update(GetAccountResult::FoundAccount(pubkey, remote_account, true)); + + assert_eq!(svm.get_account(&pubkey).unwrap(), Some(local_account)); + } + + #[test_case(TestType::sqlite(); "with on-disk sqlite db")] + #[test_case(TestType::in_memory(); "with in-memory sqlite db")] + #[test_case(TestType::no_db(); "with no db")] + #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] + fn test_remote_hydration_merges_companion_accounts_independently(test_type: TestType) { + let (mut svm, _events_rx, _geyser_rx) = test_type.initialize_svm(); + + let token_pubkey = Pubkey::new_unique(); + let mint_pubkey = Pubkey::new_unique(); + let local_mint = Account { + lamports: 2, + data: b"local-mint".to_vec(), + owner: spl_token_interface::id(), + executable: false, + rent_epoch: 0, + }; + let remote_mint = Account { + lamports: 1, + data: b"remote-mint".to_vec(), + ..local_mint.clone() + }; + let remote_token = Account { + lamports: 1, + data: b"remote-token".to_vec(), + owner: spl_token_interface::id(), + executable: false, + rent_epoch: 0, + }; + + svm.set_account(&mint_pubkey, local_mint.clone()).unwrap(); + svm.hydrate_account_update(GetAccountResult::FoundTokenAccount( + (token_pubkey, remote_token.clone()), + (mint_pubkey, Some(remote_mint)), + )); + + assert_eq!(svm.get_account(&token_pubkey).unwrap(), Some(remote_token)); + assert_eq!(svm.get_account(&mint_pubkey).unwrap(), Some(local_mint)); + + let (program_pubkey, remote_program, programdata_pubkey, remote_programdata) = + create_program_accounts(); + let local_programdata = Account { + lamports: remote_programdata.lamports + 1, + ..remote_programdata.clone() + }; + + svm.set_account(&programdata_pubkey, local_programdata.clone()) + .unwrap(); + svm.hydrate_account_update(GetAccountResult::FoundProgramAccount( + (program_pubkey, remote_program.clone()), + (programdata_pubkey, Some(remote_programdata)), + )); + + assert_eq!( + svm.get_account(&program_pubkey).unwrap(), + Some(remote_program) + ); + assert_eq!( + svm.get_account(&programdata_pubkey).unwrap(), + Some(local_programdata) + ); + } + #[test_case(TestType::sqlite(); "with on-disk sqlite db")] #[test_case(TestType::in_memory(); "with in-memory sqlite db")] #[test_case(TestType::no_db(); "with no db")]