fix(core): refactor account update to prevent race condition

crates/core/src/rpc/accounts_data.rs

@@ -383,8 +383,6 @@ impl AccountsData for SurfpoolAccountsDataRpc {
             if let Some(m) = crate::telemetry::metrics() {
                 m.record_rpc_request("getAccountInfo", rpc_start.elapsed().as_millis() as u64);
             }
-            svm_locker.write_account_update(account_update.clone());
-
             let ui_account = if let Some(((pubkey, account), token_data)) =
                 account_update.map_account_with_token_data()
             {
@@ -453,8 +451,6 @@ impl AccountsData for SurfpoolAccountsDataRpc {
                 );
             }
 
-            svm_locker.write_multiple_account_updates(&account_updates);
-
             // Convert account updates to UI accounts, order is already preserved by get_multiple_accounts
             let mut ui_accounts = vec![];
             for account_update in account_updates.into_iter() {
@@ -544,8 +540,6 @@ impl AccountsData for SurfpoolAccountsDataRpc {
                 .await?
                 .inner;
 
-            svm_locker.write_account_update(token_account_result.clone());
-
             let token_account = token_account_result.map_account()?;
 
             let (mint_pubkey, _amount) = if is_supported_token_program(&token_account.owner) {
@@ -571,8 +565,6 @@ impl AccountsData for SurfpoolAccountsDataRpc {
                 .get_account(&remote_ctx, &mint_pubkey, None)
                 .await?;
 
-            svm_locker.write_account_update(mint_account_result.clone());
-
             let mint_account = mint_account_result.map_account()?;
 
             let token_decimals = if is_supported_token_program(&mint_account.owner) {
@@ -635,8 +627,6 @@ impl AccountsData for SurfpoolAccountsDataRpc {
                 .get_account(&remote_ctx, &mint_pubkey, None)
                 .await?;
 
-            svm_locker.write_account_update(mint_account_result.clone());
-
             let mint_account = mint_account_result.map_account()?;
 
             if !is_supported_token_program(&mint_account.owner) {
@@ -715,7 +705,9 @@ mod tests {
 
     use super::*;
     use crate::{
-        surfnet::{GetAccountResult, remote::SurfnetRemoteClient},
+        surfnet::{
+            AccountSource, GetAccountResult, remote::SurfnetRemoteClient, svm::AccountUpdatePolicy,
+        },
         tests::helpers::TestSetup,
         types::SyntheticBlockhash,
     };
@@ -754,7 +746,11 @@ mod tests {
         setup
             .context
             .svm_locker
-            .write_account_update(GetAccountResult::FoundAccount(mint_pk, mint_account, true));
+            .apply_account_update(
+                GetAccountResult::FoundAccount(mint_pk, mint_account, AccountSource::Generated),
+                AccountUpdatePolicy::Authoritative,
+            )
+            .unwrap();
 
         let token_account_pk = Pubkey::new_unique();
 
@@ -786,11 +782,15 @@ mod tests {
         setup
             .context
             .svm_locker
-            .write_account_update(GetAccountResult::FoundAccount(
-                token_account_pk,
-                token_account,
-                true,
-            ));
+            .apply_account_update(
+                GetAccountResult::FoundAccount(
+                    token_account_pk,
+                    token_account,
+                    AccountSource::Generated,
+                ),
+                AccountUpdatePolicy::Authoritative,
+            )
+            .unwrap();
 
         let res = setup
             .rpc
@@ -1567,11 +1567,19 @@ mod tests {
         setup
             .context
             .svm_locker
-            .write_account_update(GetAccountResult::FoundAccount(pk1, account1, true));
+            .apply_account_update(
+                GetAccountResult::FoundAccount(pk1, account1, AccountSource::Generated),
+                AccountUpdatePolicy::Authoritative,
+            )
+            .unwrap();
         setup
             .context
             .svm_locker
-            .write_account_update(GetAccountResult::FoundAccount(pk3, account3, true));
+            .apply_account_update(
+                GetAccountResult::FoundAccount(pk3, account3, AccountSource::Generated),
+                AccountUpdatePolicy::Authoritative,
+            )
+            .unwrap();
 
         // Request accounts in order: [pk1, pk2, pk3]
         // pk1 and pk3 are local, pk2 is missing (will try remote fetch and fail)

crates/core/src/rpc/full.rs

@@ -48,7 +48,7 @@ use crate::{
     error::{SurfpoolError, SurfpoolResult},
     rpc::utils::{adjust_default_transaction_config, get_default_transaction_config},
     surfnet::{
-        FINALIZATION_SLOT_THRESHOLD, GetAccountResult, GetTransactionResult,
+        CoupledAccount, FINALIZATION_SLOT_THRESHOLD, GetAccountResult, GetTransactionResult,
         locker::SvmAccessContext, svm::MAX_RECENT_BLOCKHASHES_STANDARD,
     },
     types::{SurfnetTransactionStatus, surfpool_tx_metadata_to_litesvm_tx_metadata},
@@ -1901,9 +1901,10 @@ impl Full for SurfpoolFullRpc {
                         }
                     }
                     // According to SIMD 0186, program data is tracked as well as program accounts
-                    GetAccountResult::FoundProgramAccount(
+                    GetAccountResult::FoundCoupledAccount(
                         (pubkey, account),
-                        (pd_pubkey, pd_account),
+                        CoupledAccount::ProgramData(pd_pubkey, pd_account),
+                        _,
                     ) => {
                         if seen_accounts.insert(*pubkey) {
                             loaded_accounts_data_size += account.data.len() as u64;
@@ -1914,9 +1915,10 @@ impl Full for SurfpoolFullRpc {
                             }
                         }
                     }
-                    GetAccountResult::FoundTokenAccount(
+                    GetAccountResult::FoundCoupledAccount(
                         (pubkey, account),
-                        (td_pubkey, td_account),
+                        CoupledAccount::Mint(td_pubkey, td_account),
+                        _,
                     ) => {
                         if seen_accounts.insert(*pubkey) {
                             loaded_accounts_data_size += account.data.len() as u64;
@@ -1937,8 +1939,6 @@ impl Full for SurfpoolFullRpc {
                 track_accounts_data_size(res);
             }
 
-            svm_locker.write_multiple_account_updates(&account_updates);
-
             // Convert TransactionLoadedAddresses to LoadedAddresses before it gets consumed
             let loaded_addresses_data = loaded_addresses.as_ref().map(|la| la.loaded_addresses());
 
@@ -1950,7 +1950,6 @@ impl Full for SurfpoolFullRpc {
                 for res in alt_updates.iter() {
                     track_accounts_data_size(res);
                 }
-                svm_locker.write_multiple_account_updates(&alt_updates);
             }
 
             let replacement_blockhash = if config.replace_recent_blockhash {

crates/core/src/rpc/jito.rs

@@ -1047,8 +1047,7 @@ async fn snapshot_accounts(
                 }
             }
             crate::surfnet::GetAccountResult::FoundAccount(_, account, _)
-            | crate::surfnet::GetAccountResult::FoundProgramAccount((_, account), _)
-            | crate::surfnet::GetAccountResult::FoundTokenAccount((_, account), _) => {
+            | crate::surfnet::GetAccountResult::FoundCoupledAccount((_, account), _, _) => {
                 account.clone()
             }
         };

crates/core/src/rpc/minimal.rs

@@ -626,13 +626,10 @@ impl Minimal for SurfpoolMinimalRpc {
 
             let balance = match &account_update {
                 GetAccountResult::FoundAccount(_, account, _)
-                | GetAccountResult::FoundProgramAccount((_, account), _)
-                | GetAccountResult::FoundTokenAccount((_, account), _) => account.lamports,
+                | GetAccountResult::FoundCoupledAccount((_, account), _, _) => account.lamports,
                 GetAccountResult::None(_) => 0,
             };
 
-            svm_locker.write_account_update(account_update);
-
             #[cfg(feature = "prometheus")]
             if let Some(m) = crate::telemetry::metrics() {
                 m.record_rpc_request("getBalance", rpc_start.elapsed().as_millis() as u64);

crates/core/src/rpc/surfnet_cheatcodes.rs

@@ -30,7 +30,9 @@ use crate::{
         State,
         utils::{decode_and_deserialize, verify_pubkey, verify_pubkeys},
     },
-    surfnet::{GetAccountResult, locker::SvmAccessContext},
+    surfnet::{
+        AccountSource, GetAccountResult, locker::SvmAccessContext, svm::AccountUpdatePolicy,
+    },
     types::{
         TimeTravelConfig, TokenAccount, build_confidential_token_account_data,
         mint_has_transfer_fee_config,
@@ -1422,7 +1424,7 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
         Box::pin(async move {
             let (account_to_set, latest_absolute_slot) = if let Some(account) = account_update_opt {
                 (
-                    GetAccountResult::FoundAccount(pubkey, account, true),
+                    GetAccountResult::FoundAccount(pubkey, account, AccountSource::Generated),
                     svm_locker.get_latest_absolute_slot(),
                 )
             } else {
@@ -1445,15 +1447,15 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
                                     rent_epoch: 0,
                                     data: vec![],
                                 },
-                                true, // indicate that the account should be updated in the SVM, since it's new
+                                AccountSource::Generated,
                             )
                 }))).await?;
 
                 update.apply_ext(&mut account_result_to_update)?;
                 (account_result_to_update, slot)
             };
 
-            svm_locker.write_account_update(account_to_set);
+            svm_locker.apply_account_update(account_to_set, AccountUpdatePolicy::Authoritative)?;
 
             Ok(RpcResponse {
                 context: RpcResponseContext::new(latest_absolute_slot),
@@ -1642,8 +1644,6 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
             let mint_has_transfer_fee = confidential.is_some()
                 && !get_mint_result.is_none()
                 && mint_has_transfer_fee_config(get_mint_result.expected_data());
-            svm_locker.write_account_update(get_mint_result);
-
             let minimum_rent = svm_locker.with_svm_reader(|svm_reader| {
                 svm_reader.inner.minimum_balance_for_rent_exemption(
                     TokenAccount::get_packed_len_for_token_program_id(&token_program_id),
@@ -1678,7 +1678,7 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
                                 rent_epoch: 0,
                                 data,
                             },
-                            true, // indicate that the account should be updated in the SVM, since it's new
+                            AccountSource::Generated,
                         )
                     })),
                 )
@@ -1722,7 +1722,7 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
                 account.data = final_account_bytes.clone();
                 Ok(())
             })?;
-            svm_locker.write_account_update(token_account);
+            svm_locker.apply_account_update(token_account, AccountUpdatePolicy::Authoritative)?;
 
             Ok(RpcResponse {
                 context: RpcResponseContext::new(slot),

crates/core/src/runloops/mod.rs

@@ -45,7 +45,7 @@ use crate::{
         surfnet_cheatcodes::SurfnetCheatcodes, ws::Rpc,
     },
     surfnet::{
-        GetAccountResult, GeyserEvent, PluginCommand, locker::SurfnetSvmLocker,
+        AccountSource, GetAccountResult, GeyserEvent, PluginCommand, locker::SurfnetSvmLocker,
         remote::SurfnetRemoteClient,
     },
 };
@@ -577,11 +577,6 @@ pub async fn start_block_production_runloop(
                                 .await
                             {
                                 Ok(account_updates) => {
-                                    // The locker holds one write guard while applying the complete
-                                    // batch, so Ready cannot expose a partially installed clone set.
-                                    svm_locker
-                                        .write_multiple_account_updates(&account_updates.inner);
-
                                     // A cloned account the datasource does not have is not a
                                     // failure: hydration did complete, and some workflows clone
                                     // addresses that do not exist yet. Warn, though, because the
@@ -1171,7 +1166,7 @@ mod absent_after_hydration_tests {
         let results = vec![
             GetAccountResult::None(missing),
             GetAccountResult::None(offline),
-            GetAccountResult::FoundAccount(found, Account::default(), true),
+            GetAccountResult::FoundAccount(found, Account::default(), AccountSource::Generated),
         ];
 
         assert_eq!(