diff --git a/Cargo.lock b/Cargo.lock index 0605e85ef5130..1b478fcc18a78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2251,7 +2251,6 @@ version = "4.0.0-dev" dependencies = [ "assert_matches", "bitflags", - "environmental", "frame-metadata", "frame-support-procedural", "frame-system", diff --git a/frame/support/Cargo.toml b/frame/support/Cargo.toml index 620fc226bbf94..3161e4f4db69b 100644 --- a/frame/support/Cargo.toml +++ b/frame/support/Cargo.toml @@ -35,7 +35,6 @@ impl-trait-for-tuples = "0.2.1" smallvec = "1.8.0" log = { version = "0.4.14", default-features = false } sp-core-hashing-proc-macro = { version = "5.0.0", path = "../../primitives/core/hashing/proc-macro" } -environmental = "1.1.3" [dev-dependencies] assert_matches = "1.3.0" diff --git a/frame/support/procedural/src/transactional.rs b/frame/support/procedural/src/transactional.rs index 96cad89828d1b..b3984cac031f5 100644 --- a/frame/support/procedural/src/transactional.rs +++ b/frame/support/procedural/src/transactional.rs @@ -24,13 +24,12 @@ pub fn transactional(_attr: TokenStream, input: TokenStream) -> Result Result<(), ()> { - if *existing_levels >= TRANSACTIONAL_LIMIT || *existing_levels == Layer::MAX { - return Err(()) + pub fn get_transaction_level() -> Layer { + crate::storage::unhashed::get_or_default::(TRANSACTION_LEVEL_KEY) } - // Cannot overflow because of check above. - *existing_levels = *existing_levels + 1; - Ok(()) -} -fn dec_transaction_level(existing_levels: &mut Layer) { - if *existing_levels == 0 { - crate::defensive!( - "We are underflowing with calculating transactional levels. Not great, but let's not panic..." - ); - } else { - *existing_levels = *existing_levels - 1; + 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 { + let existing_levels = get_transaction_level(); + if existing_levels >= TRANSACTIONAL_LIMIT || existing_levels == Layer::MAX { + return Err(()) + } + // Cannot overflow because of check above. + set_transaction_level(&(existing_levels + 1)); + 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)); + } + } + + pub fn is_transactional() -> bool { + get_transaction_level() > 0 + } + + pub struct StorageLayerGuard; + + impl Drop for StorageLayerGuard { + fn drop(&mut self) { + dec_transaction_level() + } } } /// Check if the current call is within a transactional layer. pub fn is_transactional() -> bool { - current_level::with(|_| {}).is_some() + transaction_level_tracker::is_transactional() } /// Execute the supplied function in a new storage transaction. @@ -85,40 +118,27 @@ pub fn is_transactional() -> bool { /// Transactions can be nested up to 10 times; more than that will result in an error. /// /// Commits happen to the parent transaction. -pub fn with_transaction( - f: impl FnOnce() -> TransactionOutcome>, - existing_levels: &mut Layer, -) -> Result +pub fn with_transaction(f: impl FnOnce() -> TransactionOutcome>) -> Result where E: From, { use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction}; use TransactionOutcome::*; - let execute_with_transaction_levels = || { - current_level::with(|existing_levels| { - let _ = inc_transaction_level(existing_levels) - .map_err(|()| DispatchError::TransactionalLimit.into())?; - start_transaction(); - let res = match f() { - Commit(res) => { - commit_transaction(); - res - }, - Rollback(res) => { - rollback_transaction(); - res - }, - }; - dec_transaction_level(existing_levels); + + let _guard = transaction_level_tracker::inc_transaction_level() + .map_err(|()| DispatchError::TransactionalLimit.into())?; + + start_transaction(); + + match f() { + Commit(res) => { + commit_transaction(); res - }) - .unwrap() - }; - if is_transactional() { - execute_with_transaction_levels() - } else { - let mut init = Zero::zero(); - current_level::using(&mut init, || execute_with_transaction_levels()) + }, + Rollback(res) => { + rollback_transaction(); + res + }, } } @@ -1574,28 +1594,20 @@ mod test { }) } - fn get_transaction_level_or_default() -> Layer { - if is_transactional() { - current_level::with(|x| *x).unwrap() - } else { - 0 - } - } - #[test] fn transaction_limit_should_work() { TestExternalities::default().execute_with(|| { - assert_eq!(get_transaction_level_or_default(), 0); + assert_eq!(transaction_level_tracker::get_transaction_level(), 0); assert_ok!(with_transaction(|| -> TransactionOutcome { - assert_eq!(get_transaction_level_or_default(), 1); + assert_eq!(transaction_level_tracker::get_transaction_level(), 1); TransactionOutcome::Commit(Ok(())) })); assert_ok!(with_transaction(|| -> TransactionOutcome { - assert_eq!(get_transaction_level_or_default(), 1); + assert_eq!(transaction_level_tracker::get_transaction_level(), 1); let res = with_transaction(|| -> TransactionOutcome { - assert_eq!(get_transaction_level_or_default(), 2); + assert_eq!(transaction_level_tracker::get_transaction_level(), 2); TransactionOutcome::Commit(Ok(())) }); TransactionOutcome::Commit(res) @@ -1607,7 +1619,7 @@ mod test { sp_runtime::DispatchError::TransactionalLimit ); - assert_eq!(get_transaction_level_or_default(), 0); + assert_eq!(transaction_level_tracker::get_transaction_level(), 0); }); }