Introduce ReplayProofSizeProvider, RecordingProofProvider & transactional extensions

prdoc/pr_9930.prdoc

@@ -0,0 +1,19 @@
+title: Introduce `ReplayProofSizeProvider`, `RecordingProofProvider` & transactional
+  extensions
+doc:
+- audience: Node Dev
+  description: |-
+    The `ProofSizeExt` extension is used to serve the proof size to the runtime. It uses the proof recorder to request the current proof size. The `RecordingProofProvider` extension can record the calls to the proof size function. Later the `ReplayProofSizeProvider` can be used to replay these recorded proof sizes. So, the proof recorder is not required anymore.
+
+    Extensions are now also hooked into the transactional system. This means they are called when a new transaction is created and informed when a transaction is committed or reverted.
+crates:
+- name: sp-api-proc-macro
+  bump: major
+- name: sp-api
+  bump: major
+- name: sp-externalities
+  bump: major
+- name: sp-state-machine
+  bump: major
+- name: sp-trie
+  bump: major

substrate/frame/support/test/tests/construct_runtime_ui/both_use_and_excluded_parts.stderr

@@ -20,3 +20,9 @@ error[E0412]: cannot find type `Runtime` in this scope
    |
 37 | impl pallet::Config for Runtime {}
    |                         ^^^^^^^ not found in this scope
+   |
+help: there is an enum variant `sp_api::__private::TransactionType::Runtime`; try using the variant's enum
+   |
+37 - impl pallet::Config for Runtime {}
+37 + impl pallet::Config for sp_api::__private::TransactionType {}
+   |

substrate/frame/support/test/tests/construct_runtime_ui/exclude_undefined_part.stderr

@@ -20,3 +20,9 @@ error[E0412]: cannot find type `Runtime` in this scope
    |
 42 | impl pallet::Config for Runtime {}
    |                         ^^^^^^^ not found in this scope
+   |
+help: there is an enum variant `sp_api::__private::TransactionType::Runtime`; try using the variant's enum
+   |
+42 - impl pallet::Config for Runtime {}
+42 + impl pallet::Config for sp_api::__private::TransactionType {}
+   |

substrate/frame/support/test/tests/construct_runtime_ui/use_undefined_part.stderr

@@ -20,3 +20,9 @@ error[E0412]: cannot find type `Runtime` in this scope
    |
 42 | impl pallet::Config for Runtime {}
    |                         ^^^^^^^ not found in this scope
+   |
+help: there is an enum variant `sp_api::__private::TransactionType::Runtime`; try using the variant's enum
+   |
+42 - impl pallet::Config for Runtime {}
+42 + impl pallet::Config for sp_api::__private::TransactionType {}
+   |

substrate/primitives/api/proc-macro/src/impl_runtime_apis.rs

@@ -412,6 +412,11 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 							&mut std::cell::RefCell::borrow_mut(&self.changes)
 						);
 
+						#crate_::Extensions::commit_transaction(
+							&mut std::cell::RefCell::borrow_mut(&self.extensions),
+							#crate_::TransactionType::Host,
+						);
+
 						// Will panic on an `Err` below, however we should call commit
 						// on the recorder and the changes together.
 						std::result::Result::and(res, std::result::Result::map_err(res2, drop))
@@ -426,6 +431,11 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 							&mut std::cell::RefCell::borrow_mut(&self.changes)
 						);
 
+						#crate_::Extensions::rollback_transaction(
+							&mut std::cell::RefCell::borrow_mut(&self.extensions),
+							#crate_::TransactionType::Host,
+						);
+
 						// Will panic on an `Err` below, however we should call commit
 						// on the recorder and the changes together.
 						std::result::Result::and(res, std::result::Result::map_err(res2, drop))
@@ -441,6 +451,11 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 					if let Some(recorder) = &self.recorder {
 						#crate_::ProofRecorder::<Block>::start_transaction(&recorder);
 					}
+
+					#crate_::Extensions::start_transaction(
+						&mut std::cell::RefCell::borrow_mut(&self.extensions),
+						#crate_::TransactionType::Host,
+					);
 				}
 			}
 		}

substrate/primitives/api/src/lib.rs

@@ -81,7 +81,7 @@ pub mod __private {
 	mod std_imports {
 		pub use hash_db::Hasher;
 		pub use sp_core::traits::CallContext;
-		pub use sp_externalities::{Extension, Extensions};
+		pub use sp_externalities::{Extension, Extensions, TransactionType};
 		pub use sp_runtime::StateVersion;
 		pub use sp_state_machine::{
 			Backend as StateBackend, InMemoryBackend, OverlayedChanges, StorageProof, TrieBackend,

substrate/primitives/api/test/Cargo.toml

@@ -25,6 +25,7 @@ sc-block-builder = { workspace = true, default-features = true }
 scale-info = { features = ["derive"], workspace = true }
 sp-api = { workspace = true, default-features = true }
 sp-consensus = { workspace = true, default-features = true }
+sp-externalities = { workspace = true, default-features = true }
 sp-metadata-ir = { workspace = true, default-features = true }
 sp-runtime = { workspace = true, default-features = true }
 sp-state-machine = { workspace = true, default-features = true }

substrate/primitives/api/test/tests/runtime_calls.rs

@@ -15,10 +15,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::panic::UnwindSafe;
+use std::{
+	panic::UnwindSafe,
+	sync::{
+		atomic::{AtomicUsize, Ordering},
+		Arc,
+	},
+};
 
 use sc_block_builder::BlockBuilderBuilder;
 use sp_api::{ApiExt, Core, ProvideRuntimeApi};
+use sp_externalities::{decl_extension, TransactionType};
 use sp_runtime::{
 	traits::{HashingFor, Header as HeaderT},
 	TransactionOutcome,
@@ -182,14 +189,44 @@ fn disable_logging_works() {
 // Ensure that the type is not unwind safe!
 static_assertions::assert_not_impl_any!(<TestClient as ProvideRuntimeApi<_>>::Api: UnwindSafe);
 
+#[derive(Default)]
+struct TransactionTesterInner {
+	started: AtomicUsize,
+	committed: AtomicUsize,
+	rolled_back: AtomicUsize,
+}
+
+decl_extension! {
+	struct TransactionTester(Arc<TransactionTesterInner>);
+
+	impl TransactionTester {
+		fn start_transaction(&mut self, ty: TransactionType) {
+			assert_eq!(ty, TransactionType::Host);
+			self.0.started.fetch_add(1, Ordering::Relaxed);
+		}
+
+		fn commit_transaction(&mut self, ty: TransactionType) {
+			assert_eq!(ty, TransactionType::Host);
+			self.0.committed.fetch_add(1, Ordering::Relaxed);
+		}
+
+		fn rollback_transaction(&mut self, ty: TransactionType) {
+			assert_eq!(ty, TransactionType::Host);
+			self.0.rolled_back.fetch_add(1, Ordering::Relaxed);
+		}
+	}
+}
+
 #[test]
 fn ensure_transactional_works() {
 	const KEY: &[u8] = b"test";
 
 	let client = TestClientBuilder::new().build();
 	let best_hash = client.chain_info().best_hash;
+	let transaction_tester = Arc::new(TransactionTesterInner::default());
 
-	let runtime_api = client.runtime_api();
+	let mut runtime_api = client.runtime_api();
+	runtime_api.register_extension(TransactionTester(transaction_tester.clone()));
 	runtime_api.execute_in_transaction(|api| {
 		api.write_key_value(best_hash, KEY.to_vec(), vec![1, 2, 3], false).unwrap();
 
@@ -207,7 +244,8 @@ fn ensure_transactional_works() {
 		.unwrap();
 	assert_eq!(changes.main_storage_changes[0].1, Some(vec![1, 2, 3, 4]));
 
-	let runtime_api = client.runtime_api();
+	let mut runtime_api = client.runtime_api();
+	runtime_api.register_extension(TransactionTester(transaction_tester.clone()));
 	runtime_api.execute_in_transaction(|api| {
 		assert!(api.write_key_value(best_hash, KEY.to_vec(), vec![1, 2, 3], true).is_err());
 
@@ -218,4 +256,21 @@ fn ensure_transactional_works() {
 		.into_storage_changes(&client.state_at(best_hash).unwrap(), best_hash)
 		.unwrap();
 	assert_eq!(changes.main_storage_changes[0].1, Some(vec![1, 2, 3]));
+
+	let mut runtime_api = client.runtime_api();
+	runtime_api.register_extension(TransactionTester(transaction_tester.clone()));
+	runtime_api.execute_in_transaction(|api| {
+		assert!(api.write_key_value(best_hash, KEY.to_vec(), vec![1, 2], true).is_err());
+
+		TransactionOutcome::Rollback(())
+	});
+
+	let changes = runtime_api
+		.into_storage_changes(&client.state_at(best_hash).unwrap(), best_hash)
+		.unwrap();
+	assert!(changes.main_storage_changes.is_empty());
+
+	assert_eq!(transaction_tester.started.load(Ordering::Relaxed), 4);
+	assert_eq!(transaction_tester.committed.load(Ordering::Relaxed), 3);
+	assert_eq!(transaction_tester.rolled_back.load(Ordering::Relaxed), 1);
 }

substrate/primitives/externalities/src/extensions.rs

@@ -32,6 +32,27 @@ use core::{
 	ops::DerefMut,
 };
 
+/// Informs [`Extension`] about what type of transaction is started, committed or rolled back.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum TransactionType {
+	/// A transaction started by the host.
+	Host,
+	/// A transaction started by the runtime.
+	Runtime,
+}
+
+impl TransactionType {
+	/// Is `self` set to [`Self::Host`].
+	pub fn is_host(self) -> bool {
+		matches!(self, Self::Host)
+	}
+
+	/// Is `self` set to [`Self::Runtime`].
+	pub fn is_runtime(self) -> bool {
+		matches!(self, Self::Runtime)
+	}
+}
+
 /// Marker trait for types that should be registered as [`Externalities`](crate::Externalities)
 /// extension.
 ///
@@ -40,11 +61,26 @@ use core::{
 pub trait Extension: Send + 'static {
 	/// Return the extension as `&mut dyn Any`.
 	///
-	/// This is a trick to make the trait type castable into an `Any`.
+	/// This is a trick to make the trait type castable into an [`Any`].
 	fn as_mut_any(&mut self) -> &mut dyn Any;
 
 	/// Get the [`TypeId`] of this `Extension`.
 	fn type_id(&self) -> TypeId;
+
+	/// Start a transaction of type `ty`.
+	fn start_transaction(&mut self, ty: TransactionType) {
+		let _ty = ty;
+	}
+
+	/// Commit a transaction of type `ty`.
+	fn commit_transaction(&mut self, ty: TransactionType) {
+		let _ty = ty;
+	}
+
+	/// Rollback a transaction of type `ty`.
+	fn rollback_transaction(&mut self, ty: TransactionType) {
+		let _ty = ty;
+	}
 }
 
 impl Extension for Box<dyn Extension> {
@@ -55,6 +91,18 @@ impl Extension for Box<dyn Extension> {
 	fn type_id(&self) -> TypeId {
 		(**self).type_id()
 	}
+
+	fn start_transaction(&mut self, ty: TransactionType) {
+		(**self).start_transaction(ty);
+	}
+
+	fn commit_transaction(&mut self, ty: TransactionType) {
+		(**self).commit_transaction(ty);
+	}
+
+	fn rollback_transaction(&mut self, ty: TransactionType) {
+		(**self).rollback_transaction(ty);
+	}
 }
 
 /// Macro for declaring an extension that usable with [`Extensions`].
@@ -70,11 +118,37 @@ impl Extension for Box<dyn Extension> {
 ///     struct TestExt(String);
 /// }
 /// ```
+///
+/// The [`Extension`] trait provides hooks that are called when starting, committing or rolling back
+/// a transaction. These can be implemented with the macro as well:
+/// ```
+/// # use sp_externalities::{decl_extension, TransactionType};
+/// decl_extension! {
+///     /// Some test extension
+///     struct TestExtWithCallback(String);
+///
+///     impl TestExtWithCallback {
+///         fn start_transaction(&mut self, ty: TransactionType) {
+///             // do something cool
+///         }
+///
+///         // The other methods `commit_transaction` and `rollback_transaction` can also
+///         // be implemented in the same way.
+///     }
+/// }
+/// ```
 #[macro_export]
 macro_rules! decl_extension {
 	(
 		$( #[ $attr:meta ] )*
 		$vis:vis struct $ext_name:ident ($inner:ty);
+		$(
+			impl $ext_name_impl:ident {
+				$(
+					$impls:tt
+				)*
+			}
+		)*
 	) => {
 		$( #[ $attr ] )*
 		$vis struct $ext_name (pub $inner);
@@ -87,6 +161,12 @@ macro_rules! decl_extension {
 			fn type_id(&self) -> core::any::TypeId {
 				core::any::Any::type_id(self)
 			}
+
+			$(
+				$(
+					$impls
+				)*
+			)*
 		}
 
 		impl core::ops::Deref for $ext_name {
@@ -220,6 +300,21 @@ impl Extensions {
 	pub fn merge(&mut self, other: Self) {
 		self.extensions.extend(other.extensions);
 	}
+
+	/// Start a transaction of type `ty`.
+	pub fn start_transaction(&mut self, ty: TransactionType) {
+		self.extensions.values_mut().for_each(|e| e.start_transaction(ty));
+	}
+
+	/// Commit a transaction of type `ty`.
+	pub fn commit_transaction(&mut self, ty: TransactionType) {
+		self.extensions.values_mut().for_each(|e| e.commit_transaction(ty));
+	}
+
+	/// Rollback a transaction of type `ty`.
+	pub fn rollback_transaction(&mut self, ty: TransactionType) {
+		self.extensions.values_mut().for_each(|e| e.rollback_transaction(ty));
+	}
 }
 
 impl Extend<Extensions> for Extensions {

substrate/primitives/externalities/src/lib.rs

@@ -32,7 +32,7 @@ use core::any::{Any, TypeId};
 
 use sp_storage::{ChildInfo, StateVersion, TrackedStorageKey};
 
-pub use extensions::{Extension, ExtensionStore, Extensions};
+pub use extensions::{Extension, ExtensionStore, Extensions, TransactionType};
 pub use scope_limited::{set_and_run_with_externalities, with_externalities};
 
 mod extensions;