solana_accounts_db/
account_locks.rs

1#[cfg(feature = "dev-context-only-utils")]
2use qualifier_attr::qualifiers;
3use {
4    ahash::{AHashMap, AHashSet},
5    solana_pubkey::Pubkey,
6    solana_sdk::{
7        message::AccountKeys,
8        transaction::{TransactionError, MAX_TX_ACCOUNT_LOCKS},
9    },
10    std::{cell::RefCell, collections::hash_map},
11};
12
13#[derive(Debug, Default)]
14pub struct AccountLocks {
15    write_locks: AHashSet<Pubkey>,
16    readonly_locks: AHashMap<Pubkey, u64>,
17}
18
19impl AccountLocks {
20    /// Lock the account keys in `keys` for a transaction.
21    /// The bool in the tuple indicates if the account is writable.
22    /// Returns an error if any of the accounts are already locked in a way
23    /// that conflicts with the requested lock.
24    pub fn try_lock_accounts<'a>(
25        &mut self,
26        keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone,
27    ) -> Result<(), TransactionError> {
28        for (key, writable) in keys.clone() {
29            if writable {
30                if !self.can_write_lock(key) {
31                    return Err(TransactionError::AccountInUse);
32                }
33            } else if !self.can_read_lock(key) {
34                return Err(TransactionError::AccountInUse);
35            }
36        }
37
38        for (key, writable) in keys {
39            if writable {
40                self.lock_write(key);
41            } else {
42                self.lock_readonly(key);
43            }
44        }
45
46        Ok(())
47    }
48
49    /// Unlock the account keys in `keys` after a transaction.
50    /// The bool in the tuple indicates if the account is writable.
51    /// In debug-mode this function will panic if an attempt is made to unlock
52    /// an account that wasn't locked in the way requested.
53    pub fn unlock_accounts<'a>(&mut self, keys: impl Iterator<Item = (&'a Pubkey, bool)>) {
54        for (k, writable) in keys {
55            if writable {
56                self.unlock_write(k);
57            } else {
58                self.unlock_readonly(k);
59            }
60        }
61    }
62
63    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
64    fn is_locked_readonly(&self, key: &Pubkey) -> bool {
65        self.readonly_locks.get(key).is_some_and(|count| *count > 0)
66    }
67
68    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
69    fn is_locked_write(&self, key: &Pubkey) -> bool {
70        self.write_locks.contains(key)
71    }
72
73    fn can_read_lock(&self, key: &Pubkey) -> bool {
74        // If the key is not write-locked, it can be read-locked
75        !self.is_locked_write(key)
76    }
77
78    fn can_write_lock(&self, key: &Pubkey) -> bool {
79        // If the key is not read-locked or write-locked, it can be write-locked
80        !self.is_locked_readonly(key) && !self.is_locked_write(key)
81    }
82
83    fn lock_readonly(&mut self, key: &Pubkey) {
84        *self.readonly_locks.entry(*key).or_default() += 1;
85    }
86
87    fn lock_write(&mut self, key: &Pubkey) {
88        self.write_locks.insert(*key);
89    }
90
91    fn unlock_readonly(&mut self, key: &Pubkey) {
92        if let hash_map::Entry::Occupied(mut occupied_entry) = self.readonly_locks.entry(*key) {
93            let count = occupied_entry.get_mut();
94            *count -= 1;
95            if *count == 0 {
96                occupied_entry.remove_entry();
97            }
98        } else {
99            debug_assert!(
100                false,
101                "Attempted to remove a read-lock for a key that wasn't read-locked"
102            );
103        }
104    }
105
106    fn unlock_write(&mut self, key: &Pubkey) {
107        let removed = self.write_locks.remove(key);
108        debug_assert!(
109            removed,
110            "Attempted to remove a write-lock for a key that wasn't write-locked"
111        );
112    }
113}
114
115/// Validate account locks before locking.
116pub fn validate_account_locks(
117    account_keys: AccountKeys,
118    tx_account_lock_limit: usize,
119) -> Result<(), TransactionError> {
120    if account_keys.len() > tx_account_lock_limit {
121        Err(TransactionError::TooManyAccountLocks)
122    } else if has_duplicates(account_keys) {
123        Err(TransactionError::AccountLoadedTwice)
124    } else {
125        Ok(())
126    }
127}
128
129thread_local! {
130    static HAS_DUPLICATES_SET: RefCell<AHashSet<Pubkey>> = RefCell::new(AHashSet::with_capacity(MAX_TX_ACCOUNT_LOCKS));
131}
132
133/// Check for duplicate account keys.
134fn has_duplicates(account_keys: AccountKeys) -> bool {
135    // Benchmarking has shown that for sets of 32 or more keys, it is faster to
136    // use a HashSet to check for duplicates.
137    // For smaller sets a brute-force O(n^2) check seems to be faster.
138    const USE_ACCOUNT_LOCK_SET_SIZE: usize = 32;
139    if account_keys.len() >= USE_ACCOUNT_LOCK_SET_SIZE {
140        HAS_DUPLICATES_SET.with_borrow_mut(|set| {
141            let has_duplicates = account_keys.iter().any(|key| !set.insert(*key));
142            set.clear();
143            has_duplicates
144        })
145    } else {
146        for (idx, key) in account_keys.iter().enumerate() {
147            for jdx in idx + 1..account_keys.len() {
148                if key == &account_keys[jdx] {
149                    return true;
150                }
151            }
152        }
153        false
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use {super::*, solana_sdk::message::v0::LoadedAddresses};
160
161    #[test]
162    fn test_account_locks() {
163        let mut account_locks = AccountLocks::default();
164
165        let key1 = Pubkey::new_unique();
166        let key2 = Pubkey::new_unique();
167
168        // Add write and read-lock.
169        let result = account_locks.try_lock_accounts([(&key1, true), (&key2, false)].into_iter());
170        assert!(result.is_ok());
171
172        // Try to add duplicate write-lock.
173        let result = account_locks.try_lock_accounts([(&key1, true)].into_iter());
174        assert_eq!(result, Err(TransactionError::AccountInUse));
175
176        // Try to add write lock on read-locked account.
177        let result = account_locks.try_lock_accounts([(&key2, true)].into_iter());
178        assert_eq!(result, Err(TransactionError::AccountInUse));
179
180        // Try to add read lock on write-locked account.
181        let result = account_locks.try_lock_accounts([(&key1, false)].into_iter());
182        assert_eq!(result, Err(TransactionError::AccountInUse));
183
184        // Add read lock on read-locked account.
185        let result = account_locks.try_lock_accounts([(&key2, false)].into_iter());
186        assert!(result.is_ok());
187
188        // Unlock write and read locks.
189        account_locks.unlock_accounts([(&key1, true), (&key2, false)].into_iter());
190
191        // No more remaining write-locks. Read-lock remains.
192        assert!(!account_locks.is_locked_write(&key1));
193        assert!(account_locks.is_locked_readonly(&key2));
194
195        // Unlock read lock.
196        account_locks.unlock_accounts([(&key2, false)].into_iter());
197        assert!(!account_locks.is_locked_readonly(&key2));
198    }
199
200    #[test]
201    fn test_validate_account_locks_valid_no_dynamic() {
202        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
203        let account_keys = AccountKeys::new(static_keys, None);
204        assert!(validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS).is_ok());
205    }
206
207    #[test]
208    fn test_validate_account_locks_too_many_no_dynamic() {
209        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
210        let account_keys = AccountKeys::new(static_keys, None);
211        assert_eq!(
212            validate_account_locks(account_keys, 1),
213            Err(TransactionError::TooManyAccountLocks)
214        );
215    }
216
217    #[test]
218    fn test_validate_account_locks_duplicate_no_dynamic() {
219        let duplicate_key = Pubkey::new_unique();
220        let static_keys = &[duplicate_key, Pubkey::new_unique(), duplicate_key];
221        let account_keys = AccountKeys::new(static_keys, None);
222        assert_eq!(
223            validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS),
224            Err(TransactionError::AccountLoadedTwice)
225        );
226    }
227
228    #[test]
229    fn test_validate_account_locks_valid_dynamic() {
230        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
231        let dynamic_keys = LoadedAddresses {
232            writable: vec![Pubkey::new_unique()],
233            readonly: vec![Pubkey::new_unique()],
234        };
235        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
236        assert!(validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS).is_ok());
237    }
238
239    #[test]
240    fn test_validate_account_locks_too_many_dynamic() {
241        let static_keys = &[Pubkey::new_unique()];
242        let dynamic_keys = LoadedAddresses {
243            writable: vec![Pubkey::new_unique()],
244            readonly: vec![Pubkey::new_unique()],
245        };
246        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
247        assert_eq!(
248            validate_account_locks(account_keys, 2),
249            Err(TransactionError::TooManyAccountLocks)
250        );
251    }
252
253    #[test]
254    fn test_validate_account_locks_duplicate_dynamic() {
255        let duplicate_key = Pubkey::new_unique();
256        let static_keys = &[duplicate_key];
257        let dynamic_keys = LoadedAddresses {
258            writable: vec![Pubkey::new_unique()],
259            readonly: vec![duplicate_key],
260        };
261        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
262        assert_eq!(
263            validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS),
264            Err(TransactionError::AccountLoadedTwice)
265        );
266    }
267
268    #[test]
269    fn test_has_duplicates_small() {
270        let mut keys = (0..16).map(|_| Pubkey::new_unique()).collect::<Vec<_>>();
271        let account_keys = AccountKeys::new(&keys, None);
272        assert!(!has_duplicates(account_keys));
273
274        keys[14] = keys[3]; // Duplicate key
275        let account_keys = AccountKeys::new(&keys, None);
276        assert!(has_duplicates(account_keys));
277    }
278
279    #[test]
280    fn test_has_duplicates_large() {
281        let mut keys = (0..64).map(|_| Pubkey::new_unique()).collect::<Vec<_>>();
282        let account_keys = AccountKeys::new(&keys, None);
283        assert!(!has_duplicates(account_keys));
284
285        keys[47] = keys[3]; // Duplicate key
286        let account_keys = AccountKeys::new(&keys, None);
287        assert!(has_duplicates(account_keys));
288    }
289}
solana_accounts_db/account_locks.rs

solana_accounts_db/
account_locks.rs
