diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 3c5f27e226c7f..c7a1b7a35066a 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -48,48 +48,46 @@ pub mod unhashed; pub mod weak_bounded_vec; mod transaction_level_tracker { + use core::sync::atomic::{AtomicU32, Ordering}; + type Layer = u32; - const TRANSACTION_LEVEL_KEY: &'static [u8] = b":transaction_level:"; + static NUM_LEVELS: AtomicU32 = AtomicU32::new(0); const TRANSACTIONAL_LIMIT: Layer = 255; pub fn get_transaction_level() -> Layer { - crate::storage::unhashed::get_or_default::(TRANSACTION_LEVEL_KEY) - } - - fn set_transaction_level(level: &Layer) { - crate::storage::unhashed::put::(TRANSACTION_LEVEL_KEY, level); - } - - fn kill_transaction_level() { - crate::storage::unhashed::kill(TRANSACTION_LEVEL_KEY); + NUM_LEVELS.load(Ordering::SeqCst) } /// Increments the transaction level. Returns an error if levels go past the limit. /// /// Returns a guard that when dropped decrements the transaction level automatically. pub fn inc_transaction_level() -> Result { - let existing_levels = get_transaction_level(); - if existing_levels >= TRANSACTIONAL_LIMIT { - return Err(()) - } - // Cannot overflow because of check above. - set_transaction_level(&(existing_levels + 1)); + NUM_LEVELS + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |existing_levels| { + if existing_levels >= TRANSACTIONAL_LIMIT { + return None + } + // Cannot overflow because of check above. + Some(existing_levels + 1) + }) + .map_err(|_| ())?; Ok(StorageLayerGuard) } fn dec_transaction_level() { - let existing_levels = get_transaction_level(); - if existing_levels == 0 { - log::warn!( - "We are underflowing with calculating transactional levels. Not great, but let's not panic...", - ); - } else if existing_levels == 1 { - // Don't leave any trace of this storage item. - kill_transaction_level(); - } else { - // Cannot underflow because of checks above. - set_transaction_level(&(existing_levels - 1)); - } + NUM_LEVELS + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |existing_levels| { + if existing_levels == 0 { + log::warn!( + "We are underflowing with calculating transactional levels. Not great, but let's not panic...", + ); + None + } else { + // Cannot underflow because of checks above. + Some(existing_levels - 1) + } + }) + .ok(); } pub fn is_transactional() -> bool {