solana_accounts_db/
account_locks.rs

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

solana_accounts_db/
account_locks.rs
