Lock free chainlink

magicblock-chainlink/Cargo.toml

@@ -9,10 +9,10 @@ bincode = { workspace = true }
 env_logger = { workspace = true }
 futures-util = { workspace = true }
 log = { workspace = true }
-lru = { workspace = true }
 magicblock-core = { workspace = true }
 magicblock-magic-program-api = { workspace = true }
 magicblock-delegation-program = { workspace = true }
+parking_lot = { workspace = true }
 serde_json = { workspace = true }
 solana-account = { workspace = true }
 solana-account-decoder = { workspace = true }
@@ -27,6 +27,7 @@ solana-sdk = { workspace = true }
 solana-sdk-ids = { workspace = true }
 solana-system-interface = { workspace = true }
 solana-transaction-error = { workspace = true }
+scc = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 tokio-stream = { workspace = true }

magicblock-chainlink/src/accounts_bank.rs

@@ -1,33 +1,33 @@
 #[cfg(any(test, feature = "dev-context"))]
 pub mod mock {
-    use std::{collections::HashMap, fmt, sync::Mutex};
+    use std::fmt;
 
     use log::*;
     use magicblock_core::traits::AccountsBank;
+    use scc::HashMap;
     use solana_account::{AccountSharedData, WritableAccount};
     use solana_pubkey::Pubkey;
 
     use crate::blacklisted_accounts;
 
     #[derive(Default)]
     pub struct AccountsBankStub {
-        pub accounts: Mutex<HashMap<Pubkey, AccountSharedData>>,
+        pub accounts: HashMap<Pubkey, AccountSharedData>,
     }
 
     impl AccountsBankStub {
         pub fn insert(&self, pubkey: Pubkey, account: AccountSharedData) {
             trace!("Inserting account: {pubkey}");
-            self.accounts.lock().unwrap().insert(pubkey, account);
+            self.accounts.upsert(pubkey, account);
         }
 
         pub fn get(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
-            self.accounts.lock().unwrap().get(pubkey).cloned()
+            self.accounts.get(pubkey).map(|acc| acc.get().clone())
         }
 
         pub fn set_owner(&self, pubkey: &Pubkey, owner: Pubkey) -> &Self {
             trace!("Setting owner for account: {pubkey} to {owner}");
-            let mut accounts = self.accounts.lock().unwrap();
-            if let Some(account) = accounts.get_mut(pubkey) {
+            if let Some(mut account) = self.accounts.get(pubkey) {
                 account.set_owner(owner);
             } else {
                 panic!("Account not found in bank: {pubkey}");
@@ -37,8 +37,7 @@ pub mod mock {
 
         fn set_delegated(&self, pubkey: &Pubkey, delegated: bool) -> &Self {
             trace!("Setting delegated for account: {pubkey} to {delegated}");
-            let mut accounts = self.accounts.lock().unwrap();
-            if let Some(account) = accounts.get_mut(pubkey) {
+            if let Some(mut account) = self.accounts.get(pubkey) {
                 account.set_delegated(delegated);
             } else {
                 panic!("Account not found in bank: {pubkey}");
@@ -70,51 +69,42 @@ pub mod mock {
             output.push_str("AccountsBank {\n");
             let blacklisted_accounts =
                 blacklisted_accounts(&Pubkey::default(), &Pubkey::default());
-            for pubkey in self.accounts.lock().unwrap().keys() {
-                if !include_blacklisted && blacklisted_accounts.contains(pubkey)
-                {
-                    continue;
+            self.accounts.scan(|pk, _| {
+                if !include_blacklisted && blacklisted_accounts.contains(pk) {
+                    return;
                 }
-                output.push_str(&format!("{pubkey},\n"));
-            }
+                output.push_str(&format!("{pk},\n"));
+            });
             output.push_str("} ");
-            output.push_str(&format!(
-                "{} total",
-                self.accounts.lock().unwrap().len()
-            ));
+            output.push_str(&format!("{} total", self.accounts.len()));
             output
         }
     }
 
     impl AccountsBank for AccountsBankStub {
         fn get_account(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
-            self.accounts.lock().unwrap().get(pubkey).cloned()
+            self.get(pubkey)
         }
         fn remove_account(&self, pubkey: &Pubkey) {
-            self.accounts.lock().unwrap().remove(pubkey);
+            self.accounts.remove(pubkey);
         }
         fn remove_where(
             &self,
             predicate: impl Fn(&Pubkey, &AccountSharedData) -> bool,
         ) -> usize {
-            let mut accounts = self.accounts.lock().unwrap();
-            let initial_len = accounts.len();
-            accounts.retain(|k, v| !predicate(k, v));
-            initial_len - accounts.len()
+            let initial_len = self.accounts.len();
+            self.accounts.retain(|k, v| !predicate(k, v));
+            initial_len - self.accounts.len()
         }
     }
 
     impl fmt::Display for AccountsBankStub {
         fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
             write!(f, "AccountsBankStub {{")?;
-            for (pubkey, acc) in self.accounts.lock().unwrap().iter() {
-                write!(f, "\n  - {pubkey}{acc:?}")?;
-            }
-            write!(
-                f,
-                "}}\nTotal {} accounts",
-                self.accounts.lock().unwrap().len()
-            )
+            self.accounts.scan(|pubkey, acc| {
+                let _ = write!(f, "\n  - {pubkey}{acc:?}");
+            });
+            write!(f, "}}\nTotal {} accounts", self.accounts.len())
         }
     }
 }

magicblock-chainlink/src/chainlink/fetch_cloner.rs

@@ -1,9 +1,9 @@
 use std::{
-    collections::{HashMap, HashSet},
+    collections::HashSet,
     fmt,
     sync::{
         atomic::{AtomicU64, Ordering},
-        Arc, Mutex,
+        Arc,
     },
 };
 
@@ -12,6 +12,7 @@ use dlp::{
 };
 use log::*;
 use magicblock_core::traits::AccountsBank;
+use scc::{hash_map::Entry, HashMap};
 use solana_account::{AccountSharedData, ReadableAccount};
 use solana_pubkey::Pubkey;
 use tokio::{
@@ -49,7 +50,7 @@ where
     remote_account_provider: Arc<RemoteAccountProvider<T, U>>,
     /// Tracks pending account fetch requests to avoid duplicate fetches in parallel
     /// Once an account is fetched and cloned into the bank, it's removed from here
-    pending_requests: Arc<Mutex<HashMap<Pubkey, RemoteAccountRequests>>>,
+    pending_requests: Arc<HashMap<Pubkey, RemoteAccountRequests>>,
     /// Counter to track the number of fetch operations for testing deduplication
     fetch_count: Arc<AtomicU64>,
 
@@ -146,7 +147,7 @@ where
             accounts_bank: accounts_bank.clone(),
             cloner: cloner.clone(),
             validator_pubkey,
-            pending_requests: Arc::new(Mutex::new(HashMap::new())),
+            pending_requests: Arc::new(HashMap::new()),
             fetch_count: Arc::new(AtomicU64::new(0)),
             blacklisted_accounts,
         });
@@ -955,11 +956,6 @@ where
 
         // Check pending requests and bank synchronously
         {
-            let mut pending = self
-                .pending_requests
-                .lock()
-                .expect("pending_requests lock poisoned");
-
             for &pubkey in pubkeys {
                 // Check synchronously if account is in bank
                 if self.accounts_bank.get_account(&pubkey).is_some() {
@@ -969,21 +965,22 @@ where
                 }
 
                 // Check if account fetch is already pending
-                if let Some(requests) = pending.get_mut(&pubkey) {
-                    let (sender, receiver) = oneshot::channel();
-                    requests.push(sender);
-                    await_pending.push((pubkey, receiver));
-                    continue;
-                }
+                match self.pending_requests.entry(pubkey) {
+                    Entry::Occupied(mut requests) => {
+                        let (sender, receiver) = oneshot::channel();
+                        requests.get_mut().push(sender);
+                        await_pending.push((pubkey, receiver));
+                        continue;
+                    }
+                    Entry::Vacant(e) => {
+                        // or reserve an entry if the fetch request is new
+                        e.insert_entry(vec![]);
+                    }
+                };
 
                 // Account needs to be fetched - add to fetch list
                 fetch_new.push(pubkey);
             }
-
-            // Create pending entries for accounts we need to fetch
-            for &pubkey in &fetch_new {
-                pending.insert(pubkey, vec![]);
-            }
         }
 
         // If we have accounts to fetch, delegate to the existing implementation
@@ -1005,12 +1002,10 @@ where
         // Clear pending requests for fetched accounts - pending requesters can get
         // the accounts from the bank now since fetch_and_clone_accounts succeeded
         {
-            let mut pending = self
-                .pending_requests
-                .lock()
-                .expect("pending_requests lock poisoned");
             for &pubkey in &fetch_new {
-                if let Some(requests) = pending.remove(&pubkey) {
+                if let Some((_, requests)) =
+                    self.pending_requests.remove(&pubkey)
+                {
                     // We signal completion but don't send the actual account data since:
                     // 1. The account is now in the bank if it was successfully cloned
                     // 2. If there was an error, the result will contain the error info