diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs index c9814e28a7..106d3e3f45 100644 --- a/substrate/frame/support/src/storage/mod.rs +++ b/substrate/frame/support/src/storage/mod.rs @@ -48,46 +48,48 @@ pub mod unhashed; pub mod weak_bounded_vec; mod transaction_level_tracker { - use core::sync::atomic::{AtomicU32, Ordering}; - type Layer = u32; - static NUM_LEVELS: AtomicU32 = AtomicU32::new(0); + const TRANSACTION_LEVEL_KEY: &'static [u8] = b":transaction_level:"; const TRANSACTIONAL_LIMIT: Layer = 255; pub fn get_transaction_level() -> Layer { - NUM_LEVELS.load(Ordering::SeqCst) + 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); } /// 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 { - 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(|_| ())?; + 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)); Ok(StorageLayerGuard) } fn dec_transaction_level() { - 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(); + 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)); + } } pub fn is_transactional() -> bool {