feat: implement a new type FifoMap and associated tests (#318)

Co-authored-by: MicaiahReid <[email protected]>

Author: BretasArthur1

crates/cli/src/cli/mod.rs

@@ -197,15 +197,19 @@ pub struct StartSimnet {
     /// Start surfpool without a remote RPC client to simulate an offline environment (eg. surfpool start --offline)
     #[clap(long = "offline", action=ArgAction::SetTrue, default_value = "false")]
     pub offline: bool,
-    /// Disable instruction profiling (default: false)
+    /// Disable instruction profiling (eg. surfpool start --disable-instruction-profiling)
     #[clap(long = "disable-instruction-profiling", action=ArgAction::SetTrue)]
     pub disable_instruction_profiling: bool,
-    /// The log level to use for simnet logs. Options are "trace", "debug", "info", "warn", "error".
+    /// The log level to use for simnet logs. Options are "trace", "debug", "info", "warn", "error". (eg. surfpool start --log-level debug)
     #[arg(long = "log-level", short = 'l', default_value = "info")]
     pub log_level: String,
-    /// The directory to put simnet logs.
+    /// The directory to put simnet logs. (eg. surfpool start --log-path ./logs)
     #[arg(long = "log-path", default_value = DEFAULT_LOG_DIR.as_str())]
     pub log_dir: String,
+    /// The maximum number of transaction profiles to hold in memory.
+    /// Changing this will affect the memory usage of surfpool. (eg. surfpool start --max-profiles 2000)
+    #[arg(long = "max-profiles", short = 'c', default_value = "200")]
+    pub max_profiles: usize,
 }
 
 #[derive(clap::ValueEnum, PartialEq, Clone, Debug)]
@@ -325,6 +329,7 @@ impl StartSimnet {
             expiry: None,
             offline_mode: self.offline,
             instruction_profiling_enabled: !self.disable_instruction_profiling,
+            max_profiles: self.max_profiles,
         }
     }
 

crates/core/src/error.rs

@@ -424,16 +424,6 @@ impl SurfpoolError {
         Self(error)
     }
 
-    pub fn transaction_not_found_in_svm<S>(signature: S) -> Self
-    where
-        S: Display,
-    {
-        let mut error = Error::internal_error();
-        error.message =
-            format!("Transaction with signature '{signature}' was not found in the SVM");
-        Self(error)
-    }
-
     pub fn tag_not_found(tag: &str) -> Self {
         let mut error = Error::internal_error();
         error.message = format!("Profile result associated with tag '{tag}' not found in the SVM");

crates/core/src/surfnet/locker.rs

@@ -744,6 +744,10 @@ impl SurfnetSvmLocker {
         self.with_svm_reader(|svm_reader| svm_reader.instruction_profiling_enabled)
     }
 
+    pub fn get_profiling_map_capacity(&self) -> usize {
+        self.with_svm_reader(|svm_reader| svm_reader.max_profiles)
+    }
+
     pub async fn process_transaction(
         &self,
         remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
@@ -2327,22 +2331,10 @@ impl SurfnetSvmLocker {
         config: &RpcProfileResultConfig,
     ) -> SurfpoolResult<Option<UiKeyedProfileResult>> {
         let result = match &signature_or_uuid {
-            UuidOrSignature::Signature(signature) => {
-                let profile = self.with_svm_reader(|svm| {
-                    svm.executed_transaction_profiles.get(signature).cloned()
-                });
-                let transaction_exists =
-                    self.with_svm_reader(|svm| svm.transactions.contains_key(signature));
-                if profile.is_none() && transaction_exists {
-                    return Err(SurfpoolError::transaction_not_found_in_svm(signature));
-                } else {
-                    profile
-                }
-            }
+            UuidOrSignature::Signature(signature) => self
+                .with_svm_reader(|svm| svm.executed_transaction_profiles.get(signature).cloned()),
             UuidOrSignature::Uuid(uuid) => {
-                let profile = self
-                    .with_svm_reader(|svm| svm.simulated_transaction_profiles.get(uuid).cloned());
-                profile
+                self.with_svm_reader(|svm| svm.simulated_transaction_profiles.get(uuid).cloned())
             }
         };
         Ok(result.map(|profile| self.encode_ui_keyed_profile_result(profile, config)))

crates/core/src/surfnet/svm.rs

@@ -1,4 +1,5 @@
 use std::{
+    cmp::max,
     collections::{BinaryHeap, HashMap, HashSet, VecDeque},
     str::FromStr,
 };
@@ -46,9 +47,10 @@ use spl_token_2022::extension::{
     scaled_ui_amount::ScaledUiAmountConfig,
 };
 use surfpool_types::{
-    AccountChange, AccountProfileState, DEFAULT_SLOT_TIME_MS, Idl, ProfileResult, RpcProfileDepth,
-    RpcProfileResultConfig, SimnetEvent, TransactionConfirmationStatus, TransactionStatusEvent,
-    UiAccountChange, UiAccountProfileState, UiProfileResult, VersionedIdl,
+    AccountChange, AccountProfileState, DEFAULT_PROFILING_MAP_CAPACITY, DEFAULT_SLOT_TIME_MS,
+    FifoMap, Idl, ProfileResult, RpcProfileDepth, RpcProfileResultConfig, SimnetEvent,
+    TransactionConfirmationStatus, TransactionStatusEvent, UiAccountChange, UiAccountProfileState,
+    UiProfileResult, VersionedIdl,
     types::{
         ComputeUnitsEstimationResult, KeyedProfileResult, UiKeyedProfileResult, UuidOrSignature,
     },
@@ -112,7 +114,7 @@ pub struct SurfnetSvm {
     pub slot_subscriptions: Vec<Sender<SlotInfo>>,
     pub profile_tag_map: HashMap<String, Vec<UuidOrSignature>>,
     pub simulated_transaction_profiles: HashMap<Uuid, KeyedProfileResult>,
-    pub executed_transaction_profiles: HashMap<Signature, KeyedProfileResult>,
+    pub executed_transaction_profiles: FifoMap<Signature, KeyedProfileResult>,
     pub logs_subscriptions: Vec<LogsSubscriptionData>,
     pub updated_at: u64,
     pub slot_time: u64,
@@ -136,6 +138,7 @@ pub struct SurfnetSvm {
     pub registered_idls: HashMap<Pubkey, BinaryHeap<VersionedIdl>>,
     pub feature_set: FeatureSet,
     pub instruction_profiling_enabled: bool,
+    pub max_profiles: usize,
 }
 
 pub const FEATURE: Feature = Feature {
@@ -186,7 +189,7 @@ impl SurfnetSvm {
             slot_subscriptions: Vec::new(),
             profile_tag_map: HashMap::new(),
             simulated_transaction_profiles: HashMap::new(),
-            executed_transaction_profiles: HashMap::new(),
+            executed_transaction_profiles: FifoMap::default(),
             logs_subscriptions: Vec::new(),
             updated_at: Utc::now().timestamp_millis() as u64,
             slot_time: DEFAULT_SLOT_TIME_MS,
@@ -207,6 +210,7 @@ impl SurfnetSvm {
             registered_idls: HashMap::new(),
             feature_set,
             instruction_profiling_enabled: true,
+            max_profiles: DEFAULT_PROFILING_MAP_CAPACITY,
         };
 
         // Generate the initial synthetic blockhash
@@ -240,6 +244,7 @@ impl SurfnetSvm {
         self.updated_at = Utc::now().timestamp_millis() as u64;
         self.slot_time = slot_time;
         self.instruction_profiling_enabled = do_profile_instructions;
+        self.set_profiling_map_capacity(self.max_profiles);
 
         if let Some(remote_client) = remote_ctx {
             let _ = self
@@ -265,6 +270,12 @@ impl SurfnetSvm {
         self.instruction_profiling_enabled = do_profile_instructions;
     }
 
+    pub fn set_profiling_map_capacity(&mut self, capacity: usize) {
+        let clamped_capacity = max(1, capacity);
+        self.max_profiles = clamped_capacity;
+        self.executed_transaction_profiles = FifoMap::new(clamped_capacity);
+    }
+
     /// Airdrops a specified amount of lamports to a single public key.
     ///
     /// # Arguments
@@ -2345,4 +2356,20 @@ mod tests {
             assert_eq!(ui_account, expected_account);
         }
     }
+
+    #[test]
+    fn test_profiling_map_capacity_default() {
+        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::new();
+        assert_eq!(
+            svm.executed_transaction_profiles.capacity(),
+            DEFAULT_PROFILING_MAP_CAPACITY
+        );
+    }
+
+    #[test]
+    fn test_profiling_map_capacity_set() {
+        let (mut svm, _events_rx, _geyser_rx) = SurfnetSvm::new();
+        svm.set_profiling_map_capacity(10);
+        assert_eq!(svm.executed_transaction_profiles.capacity(), 10);
+    }
 }