Storage metering

polkadot/node/core/pvf/common/src/executor_intf.rs

@@ -291,6 +291,10 @@ impl sp_externalities::Externalities for ValidationExternalities {
 		panic!("storage_commit_transaction: unsupported feature for parachain validation")
 	}
 
+    fn transaction_storage_size(&mut self) -> u64 {
+		panic!("storage_commit_transaction: unsupported feature for parachain validation")
+    }
+
 	fn wipe(&mut self) {
 		panic!("wipe: unsupported feature for parachain validation")
 	}

substrate/frame/support/src/storage/transactional.rs

@@ -26,7 +26,9 @@
 //!
 //! [`with_transaction`] provides a way to run a given closure in a transactional context.
 
-use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction};
+use sp_io::storage::{
+	commit_transaction, rollback_transaction, start_transaction, transaction_storage_size,
+};
 use sp_runtime::{DispatchError, TransactionOutcome, TransactionalError};
 
 /// The type that is being used to store the current number of active layers.
@@ -38,6 +40,10 @@ pub const TRANSACTION_LEVEL_KEY: &[u8] = b":transaction_level:";
 /// The maximum number of nested layers.
 pub const TRANSACTIONAL_LIMIT: Layer = 255;
 
+/// If the transaction storage use exceeds MAX_TRANSACTION_STORAGE_SIZE,
+/// the transaction would be rejected/rolled back.
+pub const MAX_TRANSACTION_STORAGE_SIZE: Option<u64> = Some(64_000_000);
+
 /// Returns the current number of nested transactional layers.
 fn get_transaction_level() -> Layer {
 	crate::storage::unhashed::get_or_default::<Layer>(TRANSACTION_LEVEL_KEY)
@@ -117,7 +123,20 @@ where
 
 	match f() {
 		TransactionOutcome::Commit(res) => {
-			commit_transaction();
+			// Rollback if the storage usage exceeds threshold.
+			let usage = transaction_storage_size();
+			let commit =
+				MAX_TRANSACTION_STORAGE_SIZE.as_ref().map_or(true, |limit| usage <= *limit);
+			if commit {
+				commit_transaction();
+			} else {
+				log::warn!(
+					"with_transaction(): transaction storage exceeds limit: {}, {:?}",
+					usage,
+					MAX_TRANSACTION_STORAGE_SIZE,
+				);
+				rollback_transaction();
+			}
 			res
 		},
 		TransactionOutcome::Rollback(res) => {

substrate/primitives/externalities/src/lib.rs

@@ -239,6 +239,9 @@ pub trait Externalities: ExtensionStore {
 	/// no transaction is open that can be closed.
 	fn storage_commit_transaction(&mut self) -> Result<(), ()>;
 
+	/// Get the total storage used by the current transaction.
+	fn transaction_storage_size(&mut self) -> u64;
+
 	/// Index specified transaction slice and store it.
 	fn storage_index_transaction(&mut self, _index: u32, _hash: &[u8], _size: u32) {
 		unimplemented!("storage_index_transaction");

substrate/primitives/io/src/lib.rs

@@ -367,6 +367,11 @@ pub trait Storage {
 		self.storage_commit_transaction()
 			.expect("No open transaction that can be committed.");
 	}
+
+	/// Get the total storage used by the current transaction.
+	fn transaction_storage_size(&mut self) -> u64 {
+		self.transaction_storage_size()
+	}
 }
 
 /// Interface for accessing the child storage for default child trie,

substrate/primitives/state-machine/src/basic.rs

@@ -303,6 +303,10 @@ impl Externalities for BasicExternalities {
 		self.overlay.commit_transaction().map_err(drop)
 	}
 
+	fn transaction_storage_size(&mut self) -> u64 {
+		self.overlay.transaction_storage_size()
+	}
+
 	fn wipe(&mut self) {}
 
 	fn commit(&mut self) {}

substrate/primitives/state-machine/src/ext.rs

@@ -592,6 +592,10 @@ where
 		self.overlay.commit_transaction().map_err(|_| ())
 	}
 
+	fn transaction_storage_size(&mut self) -> u64 {
+		self.overlay.transaction_storage_size()
+	}
+
 	fn wipe(&mut self) {
 		for _ in 0..self.overlay.transaction_depth() {
 			self.overlay.rollback_transaction().expect(BENCHMARKING_FN);

substrate/primitives/state-machine/src/overlayed_changes/mod.rs

@@ -478,6 +478,15 @@ impl<H: Hasher> OverlayedChanges<H> {
 		Ok(())
 	}
 
+	/// Get the total storage used by the current transaction.
+	pub fn transaction_storage_size(&mut self) -> u64 {
+		let mut total = 0;
+		for (_, overlayed_value) in self.changes() {
+			total += overlayed_value.value().map_or(0, |v| v.len() as u64)
+		}
+		total
+	}
+
 	/// Call this before transfering control to the runtime.
 	///
 	/// This protects all existing transactions from being removed by the runtime.

substrate/primitives/state-machine/src/read_only.rs

@@ -195,6 +195,10 @@ where
 		unimplemented!("Transactions are not supported by ReadOnlyExternalities");
 	}
 
+	fn transaction_storage_size(&mut self) -> u64 {
+		unimplemented!("Transactions are not supported by ReadOnlyExternalities");
+	}
+
 	fn wipe(&mut self) {}
 
 	fn commit(&mut self) {}