keep track storage bytes (#2154)

* refactor

* fix accountinfo creation + impl mem_stored

* add comment

* improve syntax

Co-Authored-By: thiolliere <gui.thiolliere@gmail.com>

* rename current_mem_stored -> storage_size

* more explaination + more readable code

* bump impl version of node + builds

* delete builds

substrate/node/runtime/src/lib.rs

@@ -60,7 +60,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	impl_name: create_runtime_str!("substrate-node"),
 	authoring_version: 10,
 	spec_version: 55,
-	impl_version: 55,
+	impl_version: 56,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/contract/src/account_db.rs

@@ -16,10 +16,9 @@
 
 //! Auxilliaries to help with managing partial changes to accounts state.
 
-use super::{CodeHash, CodeHashOf, Trait, AccountInfo, TrieId, AccountInfoOf, BalanceOf};
+use super::{CodeHash, CodeHashOf, Trait, TrieId, AccountInfoOf, BalanceOf, AccountInfo, TrieIdGenerator};
 use system;
 use rstd::cell::RefCell;
-use rstd::rc::Rc;
 use rstd::collections::btree_map::{BTreeMap, Entry};
 use rstd::prelude::*;
 use runtime_primitives::traits::Zero;
@@ -46,50 +45,13 @@ impl<T: Trait> Default for ChangeEntry<T> {
 
 pub type ChangeSet<T> = BTreeMap<<T as system::Trait>::AccountId, ChangeEntry<T>>;
 
-#[derive(Clone, Default)]
-pub struct AccountTrieIdMapping<A: Ord> {
-	to_account: BTreeMap<TrieId, A>,
-	to_key: BTreeMap<A, TrieId>,
-	// this lock is related to the way overlaydb stack
-	// if set it must be unset at the lower level
-	lock: bool,
-}
-
-impl<A: Clone + Ord> AccountTrieIdMapping<A> {
-
-	pub fn new() -> Self {
-		AccountTrieIdMapping {
-			to_account: BTreeMap::new(),
-			to_key: BTreeMap::new(),
-			lock: false,
-		}
-	}
-
-	pub fn lock(&mut self) {
-		self.lock = true;
-	}
-	pub fn unlock(&mut self) {
-		self.lock = false;
-	}
-	pub fn insert(&mut self, account: A, ks: TrieId) {
-		self.to_account.insert(ks.clone(), account.clone());
-		self.to_key.insert(account, ks);
-	}
-	pub fn get_trieid(&self, account: &A) -> Option<&TrieId> {
-		if self.lock { return None }
-		self.to_key.get(account)
-	}
-	pub fn get_account(&self, ks: &TrieId) -> Option<&A> {
-		if self.lock { return None }
-		self.to_account.get(ks)
-	}
-
-}
-
 pub trait AccountDb<T: Trait> {
-	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo>;
-	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId;
-	fn get_storage(&self, trie_id: &TrieId, location: &[u8]) -> Option<Vec<u8>>;
+	/// Account is used when overlayed otherwise trie_id must be provided.
+	/// This is for performance reason.
+	///
+	/// Trie id can be None iff account doesn't have an associated trie id in <AccountInfoOf<T>>.
+	/// Because DirectAccountDb bypass the lookup for this association.
+	fn get_storage(&self, account: &T::AccountId, trie_id: Option<&TrieId>, location: &[u8]) -> Option<Vec<u8>>;
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>>;
 	fn get_balance(&self, account: &T::AccountId) -> BalanceOf<T>;
 
@@ -98,18 +60,8 @@ pub trait AccountDb<T: Trait> {
 
 pub struct DirectAccountDb;
 impl<T: Trait> AccountDb<T> for DirectAccountDb {
-	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo> {
-		let res: Option<AccountInfo> = AccountInfoOf::<T>::get(account);
-		res
-	}
-	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId {
-		use super::TrieIdGenerator;
-		<Self as AccountDb<T>>::get_account_info(self, account)
-			.map(|s|s.trie_id)
-			.unwrap_or_else(||<T as Trait>::TrieIdGenerator::trie_id(account))
-	}
-	fn get_storage(&self, trie_id: &TrieId, location: &[u8]) -> Option<Vec<u8>> {
-		child::get_raw(trie_id, location)
+	fn get_storage(&self, _account: &T::AccountId, trie_id: Option<&TrieId>, location: &[u8]) -> Option<Vec<u8>> {
+		trie_id.and_then(|id| child::get_raw(id, location))
 	}
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>> {
 		<CodeHashOf<T>>::get(account)
@@ -120,7 +72,6 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 	fn commit(&mut self, s: ChangeSet<T>) {
 		let mut total_imbalance = SignedImbalance::zero();
 		for (address, changed) in s.into_iter() {
-			let trieid = <Self as AccountDb<T>>::get_or_create_trieid(&self, &address);
 			if let Some(balance) = changed.balance {
 				let (imbalance, outcome) = T::Currency::make_free_balance_be(&address, balance);
 				total_imbalance = total_imbalance.merge(imbalance);
@@ -131,18 +82,42 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 					continue;
 				}
 			}
-			if let Some(code) = changed.code {
-				if let Some(code) = code {
-					<CodeHashOf<T>>::insert(&address, code);
+			if changed.code.is_some() || !changed.storage.is_empty() {
+				let mut info = if !<AccountInfoOf<T>>::exists(&address) {
+					let info = AccountInfo {
+						trie_id: <T as Trait>::TrieIdGenerator::trie_id(&address),
+						storage_size: 0,
+					};
+					<AccountInfoOf<T>>::insert(&address, &info);
+					info
 				} else {
-					<CodeHashOf<T>>::remove(&address);
+					<AccountInfoOf<T>>::get(&address).unwrap()
+				};
+
+				if let Some(code) = changed.code {
+					if let Some(code) = code {
+						<CodeHashOf<T>>::insert(&address, code);
+					} else {
+						<CodeHashOf<T>>::remove(&address);
+					}
 				}
-			}
-			for (k, v) in changed.storage.into_iter() {
-				if let Some(value) = v {
-					child::put_raw(&trieid[..], &k, &value[..]);
-				} else {
-					child::kill(&trieid[..], &k);
+
+				let mut new_storage_size = info.storage_size;
+				for (k, v) in changed.storage.into_iter() {
+					if let Some(value) = child::get::<Vec<u8>>(&info.trie_id[..], &k) {
+						new_storage_size -= value.len() as u64;
+					}
+					if let Some(value) = v {
+						new_storage_size += value.len() as u64;
+						child::put_raw(&info.trie_id[..], &k, &value[..]);
+					} else {
+						child::kill(&info.trie_id[..], &k);
+					}
+				}
+
+				if new_storage_size != info.storage_size {
+					info.storage_size = new_storage_size;
+					<AccountInfoOf<T>>::insert(&address, info);
 				}
 			}
 		}
@@ -160,28 +135,18 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 }
 pub struct OverlayAccountDb<'a, T: Trait + 'a> {
 	local: RefCell<ChangeSet<T>>,
-	trie_account: Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>>,
-	trie_account_cache: bool,
 	underlying: &'a AccountDb<T>,
 }
 impl<'a, T: Trait> OverlayAccountDb<'a, T> {
 	pub fn new(
 		underlying: &'a AccountDb<T>,
-		trie_account: Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>>,
-		trie_account_cache: bool,
 	) -> OverlayAccountDb<'a, T> {
 		OverlayAccountDb {
 			local: RefCell::new(ChangeSet::new()),
-			trie_account,
-			trie_account_cache,
 			underlying,
 		}
 	}
 
-	pub fn reg_cache_new_rc(&self) -> Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>> {
-		self.trie_account.clone()
-	}
-
 	pub fn into_change_set(self) -> ChangeSet<T> {
 		self.local.into_inner()
 	}
@@ -216,39 +181,13 @@ impl<'a, T: Trait> OverlayAccountDb<'a, T> {
 }
 
 impl<'a, T: Trait> AccountDb<T> for OverlayAccountDb<'a, T> {
-	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo> {
-		let v = self.underlying.get_account_info(account);
-		if self.trie_account_cache {
-			v.as_ref().map(|v|self.trie_account.as_ref().borrow_mut().insert(account.clone(), v.trie_id.clone()));
-		}
-		v
-	}
-	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId {
-		if self.trie_account_cache {
-			let mut ka_mut = self.trie_account.as_ref().borrow_mut();
-			if let Some(v) = ka_mut.get_trieid(account) {
-				v.clone()
-			} else {
-				ka_mut.unlock();
-				let v = self.underlying.get_or_create_trieid(account);
-				ka_mut.insert(account.clone(), v.clone());
-				v
-			}
-		} else {
-			let res = self.trie_account.as_ref().borrow().get_trieid(account).map(|v|v.clone());
-			res.unwrap_or_else(|| {
-				self.trie_account.as_ref().borrow_mut().lock();
-				self.underlying.get_or_create_trieid(account)
-			})
-		}
-	}
-	fn get_storage(&self, ks: &TrieId, location: &[u8]) -> Option<Vec<u8>> {
-		self.trie_account.as_ref().borrow().get_account(ks).and_then(|account| self.local
+	fn get_storage(&self, account: &T::AccountId, trie_id: Option<&TrieId>, location: &[u8]) -> Option<Vec<u8>> {
+		self.local
 			.borrow()
-			.get(&account)
+			.get(account)
 			.and_then(|a| a.storage.get(location))
 			.cloned()
-			.unwrap_or_else(|| self.underlying.get_storage(ks, location)))
+			.unwrap_or_else(|| self.underlying.get_storage(account, trie_id, location))
 	}
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>> {
 		self.local

substrate/srml/contract/src/exec.rs

@@ -14,15 +14,13 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{CodeHash, Config, ContractAddressFor, Event, RawEvent, Trait, TrieId, BalanceOf};
-use crate::account_db::{AccountDb, DirectAccountDb, OverlayAccountDb, AccountTrieIdMapping};
+use super::{CodeHash, Config, ContractAddressFor, Event, RawEvent, Trait, TrieId, BalanceOf, AccountInfoOf};
+use crate::account_db::{AccountDb, DirectAccountDb, OverlayAccountDb};
 use crate::gas::{GasMeter, Token, approx_gas_for_balance};
 
 use rstd::prelude::*;
-use rstd::cell::RefCell;
-use rstd::rc::Rc;
 use runtime_primitives::traits::{CheckedAdd, CheckedSub, Zero};
-use srml_support::traits::{WithdrawReason, Currency};
+use srml_support::{StorageMap, traits::{WithdrawReason, Currency}};
 use timestamp;
 
 pub type AccountIdOf<T> = <T as system::Trait>::AccountId;
@@ -226,7 +224,7 @@ impl<T: Trait> Token<T> for ExecFeeToken {
 
 pub struct ExecutionContext<'a, T: Trait + 'a, V, L> {
 	pub self_account: T::AccountId,
-	pub self_trieid: TrieId,
+	pub self_trie_id: Option<TrieId>,
 	pub overlay: OverlayAccountDb<'a, T>,
 	pub depth: usize,
 	pub events: Vec<Event<T>>,
@@ -246,13 +244,11 @@ where
 	///
 	/// The specified `origin` address will be used as `sender` for
 	pub fn top_level(origin: T::AccountId, cfg: &'a Config<T>, vm: &'a V, loader: &'a L) -> Self {
-		let overlay = OverlayAccountDb::<T>::new(&DirectAccountDb, Rc::new(RefCell::new(AccountTrieIdMapping::new())), true);
-		let self_trieid = overlay.get_or_create_trieid(&origin);
 		ExecutionContext {
+			self_trie_id: <AccountInfoOf<T>>::get(&origin).map(|s| s.trie_id),
 			self_account: origin,
-			self_trieid,
+			overlay: OverlayAccountDb::<T>::new(&DirectAccountDb),
 			depth: 0,
-			overlay,
 			events: Vec::new(),
 			calls: Vec::new(),
 			config: &cfg,
@@ -262,11 +258,10 @@ where
 	}
 
 	fn nested(&self, overlay: OverlayAccountDb<'a, T>, dest: T::AccountId) -> Self {
-		let self_trieid = overlay.get_or_create_trieid(&dest);
 		ExecutionContext {
-			overlay,
+			self_trie_id: <AccountInfoOf<T>>::get(&dest).map(|s| s.trie_id),
 			self_account: dest,
-			self_trieid,
+			overlay,
 			depth: self.depth + 1,
 			events: Vec::new(),
 			calls: Vec::new(),
@@ -301,7 +296,7 @@ where
 
 		let (change_set, events, calls) = {
 			let mut nested = self.nested(
-				OverlayAccountDb::new(&self.overlay, self.overlay.reg_cache_new_rc(), false),
+				OverlayAccountDb::new(&self.overlay),
 				dest.clone()
 			);
 
@@ -376,7 +371,7 @@ where
 		}
 
 		let (change_set, events, calls) = {
-			let mut overlay = OverlayAccountDb::new(&self.overlay, self.overlay.reg_cache_new_rc(), false);
+			let mut overlay = OverlayAccountDb::new(&self.overlay);
 				
 			overlay.set_code(&dest, Some(code_hash.clone()));
 			let mut nested = self.nested(overlay, dest.clone());
@@ -561,7 +556,7 @@ where
 	type T = T;
 
 	fn get_storage(&self, key: &[u8]) -> Option<Vec<u8>> {
-		self.ctx.overlay.get_storage(&self.ctx.self_trieid, key)
+		self.ctx.overlay.get_storage(&self.ctx.self_account, self.ctx.self_trie_id.as_ref(), key)
 	}
 
 	fn set_storage(&mut self, key: &[u8], value: Option<Vec<u8>>) {

substrate/srml/contract/src/lib.rs

@@ -128,7 +128,7 @@ pub struct AccountInfo {
 	/// unique ID for the subtree encoded as a byte
 	pub trie_id: TrieId,
 	/// the size of stored value in octet
-	pub current_mem_stored: u64,
+	pub storage_size: u64,
 }
 
 /// Get a trie id (trie id must be unique and collision resistant depending upon its context)
@@ -450,9 +450,7 @@ decl_storage! {
 impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
 	fn on_free_balance_zero(who: &T::AccountId) {
 		<CodeHashOf<T>>::remove(who);
-		<DirectAccountDb as AccountDb<T>>::get_account_info(&DirectAccountDb, who).map(|subtrie| {
-			child::kill_storage(&subtrie.trie_id);
-		});
+		<AccountInfoOf<T>>::get(who).map(|info| child::kill_storage(&info.trie_id));
 	}
 }
 

substrate/srml/contract/src/tests.rs

@@ -248,7 +248,7 @@ fn account_removal_removes_storage() {
 				Balances::deposit_creating(&1, 110);
 				AccountInfoOf::<Test>::insert(1, &AccountInfo {
 					trie_id: unique_id1.to_vec(),
-					current_mem_stored: 0,
+					storage_size: 0,
 				});
 				child::put(&unique_id1[..], &b"foo".to_vec(), &b"1".to_vec());
 				assert_eq!(child::get(&unique_id1[..], &b"foo".to_vec()), Some(b"1".to_vec()));
@@ -257,7 +257,7 @@ fn account_removal_removes_storage() {
 				Balances::deposit_creating(&2, 110);
 				AccountInfoOf::<Test>::insert(2, &AccountInfo {
 					trie_id: unique_id2.to_vec(),
-					current_mem_stored: 0,
+					storage_size: 0,
 				});
 				child::put(&unique_id2[..], &b"hello".to_vec(), &b"3".to_vec());
 				child::put(&unique_id2[..], &b"world".to_vec(), &b"4".to_vec());
@@ -359,6 +359,8 @@ fn instantiate_and_call() {
 					event: MetaEvent::contract(RawEvent::Instantiated(ALICE, BOB))
 				}
 			]);
+
+			assert!(AccountInfoOf::<Test>::exists(BOB));
 		},
 	);
 }