Implement nested storage transactions

frame/support/src/storage/mod.rs

@@ -29,6 +29,35 @@ pub mod child;
 pub mod generator;
 pub mod migration;
 
+/// Describes whether a storage transaction should be committed or rolled back.
+pub enum TransactionOutcome {
+	/// Transaction should be committed.
+	Commit,
+	/// Transaction should be rolled back.
+	Rollback,
+}
+
+/// Execute the supplied function in a new storage transaction.
+///
+/// All changes to storage performed by the supplied function are discarded if the returned
+/// outcome is `TransactionOutcome::Rollback`.
+///
+/// Transactions can be nested to any depth. Commits happen to the parent transaction.
+pub fn with_transaction<R>(f: impl FnOnce() -> (R, TransactionOutcome)) -> R {
+	use sp_io::storage::{
+		start_transaction, commit_transaction, rollback_transaction,
+	};
+	use TransactionOutcome::*;
+
+	start_transaction();
+	let (result, outcome) = f();
+	match outcome {
+		Commit => commit_transaction(),
+		Rollback => rollback_transaction(),
+	}
+	result
+}
+
 /// A trait for working with macro-generated storage values under the substrate storage API.
 ///
 /// Details on implementation can be found at

frame/support/test/tests/storage_transaction.rs

@@ -0,0 +1,159 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2019-2020 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use codec::{Encode, Decode, EncodeLike};
+use frame_support::{
+	StorageMap, StorageValue, storage::{with_transaction, TransactionOutcome::*},
+};
+use sp_io::TestExternalities;
+
+pub trait Trait {
+	type Origin;
+	type BlockNumber: Encode + Decode + EncodeLike + Default + Clone;
+}
+
+frame_support::decl_module! {
+	pub struct Module<T: Trait> for enum Call where origin: T::Origin {}
+}
+
+frame_support::decl_storage!{
+	trait Store for Module<T: Trait> as StorageTransactions {
+		pub Value: u32;
+		pub Map: map hasher(twox_64_concat) String => u32;
+	}
+}
+
+
+#[test]
+fn storage_transaction_basic_commit() {
+	TestExternalities::default().execute_with(|| {
+
+		assert_eq!(Value::get(), 0);
+		assert!(!Map::contains_key("val0"));
+
+		with_transaction(|| {
+			Value::set(99);
+			Map::insert("val0", 99);
+			assert_eq!(Value::get(), 99);
+			assert_eq!(Map::get("val0"), 99);
+			((), Commit)
+		});
+
+		assert_eq!(Value::get(), 99);
+		assert_eq!(Map::get("val0"), 99);
+	});
+}
+
+#[test]
+fn storage_transaction_basic_rollback() {
+	TestExternalities::default().execute_with(|| {
+
+		assert_eq!(Value::get(), 0);
+		assert_eq!(Map::get("val0"), 0);
+
+		with_transaction(|| {
+			Value::set(99);
+			Map::insert("val0", 99);
+			assert_eq!(Value::get(), 99);
+			assert_eq!(Map::get("val0"), 99);
+			((), Rollback)
+		});
+
+		assert_eq!(Value::get(), 0);
+		assert_eq!(Map::get("val0"), 0);
+	});
+}
+
+#[test]
+fn storage_transaction_rollback_then_commit() {
+	TestExternalities::default().execute_with(|| {
+		Value::set(1);
+		Map::insert("val1", 1);
+
+		with_transaction(|| {
+			Value::set(2);
+			Map::insert("val1", 2);
+			Map::insert("val2", 2);
+
+			with_transaction(|| {
+				Value::set(3);
+				Map::insert("val1", 3);
+				Map::insert("val2", 3);
+				Map::insert("val3", 3);
+
+				assert_eq!(Value::get(), 3);
+				assert_eq!(Map::get("val1"), 3);
+				assert_eq!(Map::get("val2"), 3);
+				assert_eq!(Map::get("val3"), 3);
+
+				((), Rollback)
+			});
+
+			assert_eq!(Value::get(), 2);
+			assert_eq!(Map::get("val1"), 2);
+			assert_eq!(Map::get("val2"), 2);
+			assert_eq!(Map::get("val3"), 0);
+
+			((), Commit)
+		});
+
+		assert_eq!(Value::get(), 2);
+		assert_eq!(Map::get("val1"), 2);
+		assert_eq!(Map::get("val2"), 2);
+		assert_eq!(Map::get("val3"), 0);
+	});
+}
+
+#[test]
+fn storage_transaction_commit_then_rollback() {
+	TestExternalities::default().execute_with(|| {
+		Value::set(1);
+		Map::insert("val1", 1);
+
+		with_transaction(|| {
+			Value::set(2);
+			Map::insert("val1", 2);
+			Map::insert("val2", 2);
+
+			with_transaction(|| {
+				Value::set(3);
+				Map::insert("val1", 3);
+				Map::insert("val2", 3);
+				Map::insert("val3", 3);
+
+				assert_eq!(Value::get(), 3);
+				assert_eq!(Map::get("val1"), 3);
+				assert_eq!(Map::get("val2"), 3);
+				assert_eq!(Map::get("val3"), 3);
+
+				((), Commit)
+			});
+
+			assert_eq!(Value::get(), 3);
+			assert_eq!(Map::get("val1"), 3);
+			assert_eq!(Map::get("val2"), 3);
+			assert_eq!(Map::get("val3"), 3);
+
+			((), Rollback)
+		});
+
+		assert_eq!(Value::get(), 1);
+		assert_eq!(Map::get("val1"), 1);
+		assert_eq!(Map::get("val2"), 0);
+		assert_eq!(Map::get("val3"), 0);
+	});
+}

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

@@ -257,6 +257,7 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 				&self,
 				map_call: F,
 			) -> std::result::Result<R, E> where Self: Sized {
+				self.changes.borrow_mut().start_transaction();
 				*self.commit_on_success.borrow_mut() = false;
 				let res = map_call(self);
 				*self.commit_on_success.borrow_mut() = true;
@@ -366,6 +367,9 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 				&self,
 				call_api_at: F,
 			) -> std::result::Result<#crate_::NativeOrEncoded<R>, E> {
+				if *self.commit_on_success.borrow() {
+					self.changes.borrow_mut().start_transaction();
+				}
 				let res = call_api_at(
 					&self.call,
 					self,
@@ -381,11 +385,16 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 			}
 
 			fn commit_on_ok<R, E>(&self, res: &std::result::Result<R, E>) {
+				let proof = "\
+					We only close a transaction when we opened one ourself.
+					Other parts of the runtime that make use of transactions (state-machine)
+					also balance their transactions. The runtime cannot close client initiated
+					transactions. qed";
 				if *self.commit_on_success.borrow() {
 					if res.is_err() {
-						self.changes.borrow_mut().discard_prospective();
+						self.changes.borrow_mut().rollback_transaction().expect(proof);
 					} else {
-						self.changes.borrow_mut().commit_prospective();
+						self.changes.borrow_mut().commit_transaction().expect(proof);
 					}
 				}
 			}

primitives/externalities/src/lib.rs

@@ -195,6 +195,29 @@ pub trait Externalities: ExtensionStore {
 	/// The returned hash is defined by the `Block` and is SCALE encoded.
 	fn storage_changes_root(&mut self, parent: &[u8]) -> Result<Option<Vec<u8>>, ()>;
 
+	/// Start a new nested transaction.
+	///
+	/// This allows to either commit or roll back all changes made after this call to the
+	/// top changes or the default child changes. For every transaction there cam be a
+	/// matching call to either `storage_rollback_transaction` or `storage_commit_transaction`.
+	/// Any transactions that are still open after returning from runtime are committed
+	/// automatically.
+	///
+	/// Changes made without any open transaction are committed immediately.
+	fn storage_start_transaction(&mut self);
+
+	/// Rollback the last transaction started by `storage_start_transaction`.
+	///
+	/// Any changes made during that storage transaction are discarded. Returns an error when
+	/// no transaction is open that can be closed.
+	fn storage_rollback_transaction(&mut self) -> Result<(), ()>;
+
+	/// Commit the last transaction started by `storage_start_transaction`.
+	///
+	/// Any changes made during that storage transaction are committed. Returns an error when
+	/// no transaction is open that can be closed.
+	fn storage_commit_transaction(&mut self) -> Result<(), ()>;
+
 	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 	/// Benchmarking related functionality and shouldn't be used anywhere else!
 	/// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!