Replace storage access by atomics

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::<Layer>(TRANSACTION_LEVEL_KEY)
-	}
-
-	fn set_transaction_level(level: &Layer) {
-		crate::storage::unhashed::put::<Layer>(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<StorageLayerGuard, ()> {
-		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 {