referrerpolicy=no-referrer-when-downgrade

pallet_contracts/
transient_storage.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! This module contains routines for accessing and altering a contract transient storage.
19
20use crate::{
21	exec::{AccountIdOf, Key},
22	storage::WriteOutcome,
23	Config, Error,
24};
25use alloc::{collections::BTreeMap, vec::Vec};
26use codec::Encode;
27use core::{marker::PhantomData, mem};
28use frame_support::DefaultNoBound;
29use sp_runtime::{DispatchError, DispatchResult, Saturating};
30
31/// Meter entry tracks transaction allocations.
32#[derive(Default, Debug)]
33pub struct MeterEntry {
34	/// Allocations made in the current transaction.
35	pub amount: u32,
36	/// Allocations limit in the current transaction.
37	pub limit: u32,
38}
39
40impl MeterEntry {
41	/// Create a new entry.
42	fn new(limit: u32) -> Self {
43		Self { limit, amount: Default::default() }
44	}
45
46	/// Check if the allocated amount exceeds the limit.
47	fn exceeds_limit(&self, amount: u32) -> bool {
48		self.amount.saturating_add(amount) > self.limit
49	}
50
51	/// Absorb the allocation amount of the nested entry into the current entry.
52	fn absorb(&mut self, rhs: Self) {
53		self.amount.saturating_accrue(rhs.amount)
54	}
55}
56
57// The storage meter enforces a limit for each transaction,
58// which is calculated as free_storage * (1 - 1/16) for each subsequent frame.
59#[derive(DefaultNoBound)]
60pub struct StorageMeter<T: Config> {
61	nested_meters: Vec<MeterEntry>,
62	root_meter: MeterEntry,
63	_phantom: PhantomData<T>,
64}
65
66impl<T: Config> StorageMeter<T> {
67	const STORAGE_FRACTION_DENOMINATOR: u32 = 16;
68	/// Create a new storage allocation meter.
69	fn new(memory_limit: u32) -> Self {
70		Self { root_meter: MeterEntry::new(memory_limit), ..Default::default() }
71	}
72
73	/// Charge the allocated amount of transaction storage from the meter.
74	fn charge(&mut self, amount: u32) -> DispatchResult {
75		let meter = self.current_mut();
76		if meter.exceeds_limit(amount) {
77			return Err(Error::<T>::OutOfTransientStorage.into());
78		}
79		meter.amount.saturating_accrue(amount);
80		Ok(())
81	}
82
83	/// Revert a transaction meter.
84	fn revert(&mut self) {
85		self.nested_meters.pop().expect(
86			"A call to revert a meter must be preceded by a corresponding call to start a meter;
87			the code within this crate makes sure that this is always the case; qed",
88		);
89	}
90
91	/// Start a transaction meter.
92	fn start(&mut self) {
93		let meter = self.current();
94		let mut transaction_limit = meter.limit.saturating_sub(meter.amount);
95		if !self.nested_meters.is_empty() {
96			// Allow use of (1 - 1/STORAGE_FRACTION_DENOMINATOR) of free storage for subsequent
97			// calls.
98			transaction_limit.saturating_reduce(
99				transaction_limit.saturating_div(Self::STORAGE_FRACTION_DENOMINATOR),
100			);
101		}
102
103		self.nested_meters.push(MeterEntry::new(transaction_limit));
104	}
105
106	/// Commit a transaction meter.
107	fn commit(&mut self) {
108		let transaction_meter = self.nested_meters.pop().expect(
109			"A call to commit a meter must be preceded by a corresponding call to start a meter;
110			the code within this crate makes sure that this is always the case; qed",
111		);
112		self.current_mut().absorb(transaction_meter)
113	}
114
115	/// The total allocated amount of memory.
116	#[cfg(test)]
117	fn total_amount(&self) -> u32 {
118		self.nested_meters
119			.iter()
120			.fold(self.root_meter.amount, |acc, e| acc.saturating_add(e.amount))
121	}
122
123	/// A mutable reference to the current meter entry.
124	pub fn current_mut(&mut self) -> &mut MeterEntry {
125		self.nested_meters.last_mut().unwrap_or(&mut self.root_meter)
126	}
127
128	/// A reference to the current meter entry.
129	pub fn current(&self) -> &MeterEntry {
130		self.nested_meters.last().unwrap_or(&self.root_meter)
131	}
132}
133
134/// An entry representing a journal change.
135struct JournalEntry {
136	key: Vec<u8>,
137	prev_value: Option<Vec<u8>>,
138}
139
140impl JournalEntry {
141	/// Create a new change.
142	fn new(key: Vec<u8>, prev_value: Option<Vec<u8>>) -> Self {
143		Self { key, prev_value }
144	}
145
146	/// Revert the change.
147	fn revert(self, storage: &mut Storage) {
148		storage.write(&self.key, self.prev_value);
149	}
150}
151
152/// A journal containing transient storage modifications.
153struct Journal(Vec<JournalEntry>);
154
155impl Journal {
156	/// Create a new journal.
157	fn new() -> Self {
158		Self(Default::default())
159	}
160
161	/// Add a change to the journal.
162	fn push(&mut self, entry: JournalEntry) {
163		self.0.push(entry);
164	}
165
166	/// Length of the journal.
167	fn len(&self) -> usize {
168		self.0.len()
169	}
170
171	/// Roll back all journal changes until the chackpoint
172	fn rollback(&mut self, storage: &mut Storage, checkpoint: usize) {
173		self.0.drain(checkpoint..).rev().for_each(|entry| entry.revert(storage));
174	}
175}
176
177/// Storage for maintaining the current transaction state.
178#[derive(Default)]
179struct Storage(BTreeMap<Vec<u8>, Vec<u8>>);
180
181impl Storage {
182	/// Read the storage entry.
183	fn read(&self, key: &Vec<u8>) -> Option<Vec<u8>> {
184		self.0.get(key).cloned()
185	}
186
187	/// Write the storage entry.
188	fn write(&mut self, key: &Vec<u8>, value: Option<Vec<u8>>) -> Option<Vec<u8>> {
189		if let Some(value) = value {
190			// Insert storage entry.
191			self.0.insert(key.clone(), value)
192		} else {
193			// Remove storage entry.
194			self.0.remove(key)
195		}
196	}
197}
198
199/// Transient storage behaves almost identically to regular storage but is discarded after each
200/// transaction. It consists of a `BTreeMap` for the current state, a journal of all changes, and a
201/// list of checkpoints. On entry to the `start_transaction` function, a marker (checkpoint) is
202/// added to the list. New values are written to the current state, and the previous value is
203/// recorded in the journal (`write`). When the `commit_transaction` function is called, the marker
204/// to the journal index (checkpoint) of when that call was entered is discarded.
205/// On `rollback_transaction`, all entries are reverted up to the last checkpoint.
206pub struct TransientStorage<T: Config> {
207	// The storage and journal size is limited by the storage meter.
208	storage: Storage,
209	journal: Journal,
210	// The size of the StorageMeter is limited by the stack depth.
211	meter: StorageMeter<T>,
212	// The size of the checkpoints is limited by the stack depth.
213	checkpoints: Vec<usize>,
214}
215
216impl<T: Config> TransientStorage<T> {
217	/// Create new transient storage with the supplied memory limit.
218	pub fn new(memory_limit: u32) -> Self {
219		TransientStorage {
220			storage: Default::default(),
221			journal: Journal::new(),
222			checkpoints: Default::default(),
223			meter: StorageMeter::new(memory_limit),
224		}
225	}
226
227	/// Read the storage value. If the entry does not exist, `None` is returned.
228	pub fn read(&self, account: &AccountIdOf<T>, key: &Key<T>) -> Option<Vec<u8>> {
229		self.storage.read(&Self::storage_key(&account.encode(), &key.hash()))
230	}
231
232	/// Write a value to storage.
233	///
234	/// If the `value` is `None`, then the entry is removed. If `take` is true,
235	/// a [`WriteOutcome::Taken`] is returned instead of a [`WriteOutcome::Overwritten`].
236	/// If the entry did not exist, [`WriteOutcome::New`] is returned.
237	pub fn write(
238		&mut self,
239		account: &AccountIdOf<T>,
240		key: &Key<T>,
241		value: Option<Vec<u8>>,
242		take: bool,
243	) -> Result<WriteOutcome, DispatchError> {
244		let key = Self::storage_key(&account.encode(), &key.hash());
245		let prev_value = self.storage.read(&key);
246		// Skip if the same value is being set.
247		if prev_value != value {
248			// Calculate the allocation size.
249			if let Some(value) = &value {
250				// Charge the key, value and journal entry.
251				// If a new value is written, a new journal entry is created. The previous value is
252				// moved to the journal along with its key, and the new value is written to
253				// storage.
254				let key_len = key.capacity();
255				let mut amount = value
256					.capacity()
257					.saturating_add(key_len)
258					.saturating_add(mem::size_of::<JournalEntry>());
259				if prev_value.is_none() {
260					// Charge a new storage entry.
261					// If there was no previous value, a new entry is added to storage (BTreeMap)
262					// containing a Vec for the key and a Vec for the value. The value was already
263					// included in the amount.
264					amount.saturating_accrue(key_len.saturating_add(mem::size_of::<Vec<u8>>()));
265				}
266				self.meter.charge(amount as _)?;
267			}
268			self.storage.write(&key, value);
269			// Update the journal.
270			self.journal.push(JournalEntry::new(key, prev_value.clone()));
271		}
272
273		Ok(match (take, prev_value) {
274			(_, None) => WriteOutcome::New,
275			(false, Some(prev_value)) => WriteOutcome::Overwritten(prev_value.len() as _),
276			(true, Some(prev_value)) => WriteOutcome::Taken(prev_value),
277		})
278	}
279
280	/// Start a new nested transaction.
281	///
282	/// This allows to either commit or roll back all changes that are made after this call.
283	/// For every transaction there must be a matching call to either `rollback_transaction`
284	/// or `commit_transaction`.
285	pub fn start_transaction(&mut self) {
286		self.meter.start();
287		self.checkpoints.push(self.journal.len());
288	}
289
290	/// Rollback the last transaction started by `start_transaction`.
291	///
292	/// Any changes made during that transaction are discarded.
293	///
294	/// # Panics
295	///
296	/// Will panic if there is no open transaction.
297	pub fn rollback_transaction(&mut self) {
298		let checkpoint = self
299			.checkpoints
300			.pop()
301			.expect(
302				"A call to rollback_transaction must be preceded by a corresponding call to start_transaction;
303				the code within this crate makes sure that this is always the case; qed"
304			);
305		self.meter.revert();
306		self.journal.rollback(&mut self.storage, checkpoint);
307	}
308
309	/// Commit the last transaction started by `start_transaction`.
310	///
311	/// Any changes made during that transaction are committed.
312	///
313	/// # Panics
314	///
315	/// Will panic if there is no open transaction.
316	pub fn commit_transaction(&mut self) {
317		self.checkpoints
318			.pop()
319			.expect(
320				"A call to commit_transaction must be preceded by a corresponding call to start_transaction;
321				the code within this crate makes sure that this is always the case; qed"
322			);
323		self.meter.commit();
324	}
325
326	/// The storage allocation meter used for transaction metering.
327	#[cfg(any(test, feature = "runtime-benchmarks"))]
328	pub fn meter(&mut self) -> &mut StorageMeter<T> {
329		return &mut self.meter
330	}
331
332	fn storage_key(account: &[u8], key: &[u8]) -> Vec<u8> {
333		let mut storage_key = Vec::with_capacity(account.len() + key.len());
334		storage_key.extend_from_slice(&account);
335		storage_key.extend_from_slice(&key);
336		storage_key
337	}
338}
339
340#[cfg(test)]
341mod tests {
342	use super::*;
343	use crate::{
344		tests::{Test, ALICE, BOB, CHARLIE},
345		Error,
346	};
347	use core::u32::MAX;
348
349	// Calculate the allocation size for the given entry.
350	fn allocation_size(
351		account: &AccountIdOf<Test>,
352		key: &Key<Test>,
353		value: Option<Vec<u8>>,
354	) -> u32 {
355		let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(MAX);
356		storage
357			.write(account, key, value, false)
358			.expect("Could not write to transient storage.");
359		storage.meter().current().amount
360	}
361
362	#[test]
363	fn read_write_works() {
364		let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(2048);
365		assert_eq!(
366			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
367			Ok(WriteOutcome::New)
368		);
369		assert_eq!(
370			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![2]), true),
371			Ok(WriteOutcome::New)
372		);
373		assert_eq!(
374			storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![3]), false),
375			Ok(WriteOutcome::New)
376		);
377		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
378		assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![2]));
379		assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![3]));
380		// Overwrite values.
381		assert_eq!(
382			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![4, 5]), false),
383			Ok(WriteOutcome::Overwritten(1))
384		);
385		assert_eq!(
386			storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![6, 7]), true),
387			Ok(WriteOutcome::Taken(vec![3]))
388		);
389		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
390		assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![4, 5]));
391		assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![6, 7]));
392
393		// Check for an empty value.
394		assert_eq!(
395			storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![]), true),
396			Ok(WriteOutcome::Taken(vec![6, 7]))
397		);
398		assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![]));
399
400		assert_eq!(
401			storage.write(&BOB, &Key::Fix([3; 32]), None, true),
402			Ok(WriteOutcome::Taken(vec![]))
403		);
404		assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), None);
405	}
406
407	#[test]
408	fn read_write_with_var_sized_keys_works() {
409		let mut storage = TransientStorage::<Test>::new(2048);
410		assert_eq!(
411			storage.write(
412				&ALICE,
413				&Key::<Test>::try_from_var([1; 64].to_vec()).unwrap(),
414				Some(vec![1]),
415				false
416			),
417			Ok(WriteOutcome::New)
418		);
419		assert_eq!(
420			storage.write(
421				&BOB,
422				&Key::<Test>::try_from_var([2; 64].to_vec()).unwrap(),
423				Some(vec![2, 3]),
424				false
425			),
426			Ok(WriteOutcome::New)
427		);
428		assert_eq!(
429			storage.read(&ALICE, &Key::<Test>::try_from_var([1; 64].to_vec()).unwrap()),
430			Some(vec![1])
431		);
432		assert_eq!(
433			storage.read(&BOB, &Key::<Test>::try_from_var([2; 64].to_vec()).unwrap()),
434			Some(vec![2, 3])
435		);
436		// Overwrite values.
437		assert_eq!(
438			storage.write(
439				&ALICE,
440				&Key::<Test>::try_from_var([1; 64].to_vec()).unwrap(),
441				Some(vec![4, 5]),
442				false
443			),
444			Ok(WriteOutcome::Overwritten(1))
445		);
446		assert_eq!(
447			storage.read(&ALICE, &Key::<Test>::try_from_var([1; 64].to_vec()).unwrap()),
448			Some(vec![4, 5])
449		);
450	}
451
452	#[test]
453	fn rollback_transaction_works() {
454		let mut storage = TransientStorage::<Test>::new(1024);
455
456		storage.start_transaction();
457		assert_eq!(
458			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
459			Ok(WriteOutcome::New)
460		);
461		storage.rollback_transaction();
462		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None)
463	}
464
465	#[test]
466	fn commit_transaction_works() {
467		let mut storage = TransientStorage::<Test>::new(1024);
468
469		storage.start_transaction();
470		assert_eq!(
471			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
472			Ok(WriteOutcome::New)
473		);
474		storage.commit_transaction();
475		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]))
476	}
477
478	#[test]
479	fn overwrite_and_commmit_transaction_works() {
480		let mut storage = TransientStorage::<Test>::new(1024);
481		storage.start_transaction();
482		assert_eq!(
483			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
484			Ok(WriteOutcome::New)
485		);
486		assert_eq!(
487			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1, 2]), false),
488			Ok(WriteOutcome::Overwritten(1))
489		);
490		storage.commit_transaction();
491		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1, 2]))
492	}
493
494	#[test]
495	fn rollback_in_nested_transaction_works() {
496		let mut storage = TransientStorage::<Test>::new(1024);
497		storage.start_transaction();
498		assert_eq!(
499			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
500			Ok(WriteOutcome::New)
501		);
502		storage.start_transaction();
503		assert_eq!(
504			storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1]), false),
505			Ok(WriteOutcome::New)
506		);
507		storage.rollback_transaction();
508		storage.commit_transaction();
509		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
510		assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None)
511	}
512
513	#[test]
514	fn commit_in_nested_transaction_works() {
515		let mut storage = TransientStorage::<Test>::new(1024);
516		storage.start_transaction();
517		assert_eq!(
518			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
519			Ok(WriteOutcome::New)
520		);
521		storage.start_transaction();
522		assert_eq!(
523			storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
524			Ok(WriteOutcome::New)
525		);
526		storage.start_transaction();
527		assert_eq!(
528			storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
529			Ok(WriteOutcome::New)
530		);
531		storage.commit_transaction();
532		storage.commit_transaction();
533		storage.commit_transaction();
534		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
535		assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), Some(vec![2]));
536		assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), Some(vec![3]));
537	}
538
539	#[test]
540	fn rollback_all_transactions_works() {
541		let mut storage = TransientStorage::<Test>::new(1024);
542		storage.start_transaction();
543		assert_eq!(
544			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
545			Ok(WriteOutcome::New)
546		);
547		storage.start_transaction();
548		assert_eq!(
549			storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
550			Ok(WriteOutcome::New)
551		);
552		storage.start_transaction();
553		assert_eq!(
554			storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
555			Ok(WriteOutcome::New)
556		);
557		storage.commit_transaction();
558		storage.commit_transaction();
559		storage.rollback_transaction();
560		assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None);
561		assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None);
562		assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), None);
563	}
564
565	#[test]
566	fn metering_transactions_works() {
567		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
568		let mut storage = TransientStorage::<Test>::new(size * 2);
569		storage.start_transaction();
570		assert_eq!(
571			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
572			Ok(WriteOutcome::New)
573		);
574		let limit = storage.meter().current().limit;
575		storage.commit_transaction();
576
577		storage.start_transaction();
578		assert_eq!(storage.meter().current().limit, limit - size);
579		assert_eq!(storage.meter().current().limit - storage.meter().current().amount, size);
580		assert_eq!(
581			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
582			Ok(WriteOutcome::New)
583		);
584		assert_eq!(storage.meter().current().amount, size);
585		storage.commit_transaction();
586		assert_eq!(storage.meter().total_amount(), size * 2);
587	}
588
589	#[test]
590	fn metering_nested_transactions_works() {
591		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
592		let mut storage = TransientStorage::<Test>::new(size * 3);
593
594		storage.start_transaction();
595		let limit = storage.meter().current().limit;
596		assert_eq!(
597			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
598			Ok(WriteOutcome::New)
599		);
600		storage.start_transaction();
601		assert_eq!(storage.meter().total_amount(), size);
602		assert!(storage.meter().current().limit < limit - size);
603		assert_eq!(
604			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
605			Ok(WriteOutcome::New)
606		);
607		storage.commit_transaction();
608		assert_eq!(storage.meter().current().limit, limit);
609		assert_eq!(storage.meter().total_amount(), storage.meter().current().amount);
610		storage.commit_transaction();
611	}
612
613	#[test]
614	fn metering_transaction_fails() {
615		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
616		let mut storage = TransientStorage::<Test>::new(size - 1);
617		storage.start_transaction();
618		assert_eq!(
619			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
620			Err(Error::<Test>::OutOfTransientStorage.into())
621		);
622		assert_eq!(storage.meter.current().amount, 0);
623		storage.commit_transaction();
624		assert_eq!(storage.meter.total_amount(), 0);
625	}
626
627	#[test]
628	fn metering_nested_transactions_fails() {
629		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
630		let mut storage = TransientStorage::<Test>::new(size * 2);
631
632		storage.start_transaction();
633		assert_eq!(
634			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
635			Ok(WriteOutcome::New)
636		);
637		storage.start_transaction();
638		assert_eq!(
639			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
640			Err(Error::<Test>::OutOfTransientStorage.into())
641		);
642		storage.commit_transaction();
643		storage.commit_transaction();
644		assert_eq!(storage.meter.total_amount(), size);
645	}
646
647	#[test]
648	fn metering_nested_transaction_with_rollback_works() {
649		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
650		let mut storage = TransientStorage::<Test>::new(size * 2);
651
652		storage.start_transaction();
653		let limit = storage.meter.current().limit;
654		storage.start_transaction();
655		assert_eq!(
656			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
657			Ok(WriteOutcome::New)
658		);
659		storage.rollback_transaction();
660
661		assert_eq!(storage.meter.total_amount(), 0);
662		assert_eq!(storage.meter.current().limit, limit);
663		assert_eq!(
664			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
665			Ok(WriteOutcome::New)
666		);
667		let amount = storage.meter().current().amount;
668		assert_eq!(storage.meter().total_amount(), amount);
669		storage.commit_transaction();
670	}
671
672	#[test]
673	fn metering_with_rollback_works() {
674		let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
675		let mut storage = TransientStorage::<Test>::new(size * 5);
676
677		storage.start_transaction();
678		assert_eq!(
679			storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
680			Ok(WriteOutcome::New)
681		);
682		let amount = storage.meter.total_amount();
683		storage.start_transaction();
684		assert_eq!(
685			storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
686			Ok(WriteOutcome::New)
687		);
688		storage.start_transaction();
689		assert_eq!(
690			storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
691			Ok(WriteOutcome::New)
692		);
693		storage.commit_transaction();
694		storage.rollback_transaction();
695		assert_eq!(storage.meter.total_amount(), amount);
696		storage.commit_transaction();
697	}
698}