1use crate::{
21 Config, Error,
22 exec::{AccountIdOf, Key},
23 storage::WriteOutcome,
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#[derive(Default, Debug, Clone)]
33pub struct MeterEntry {
34 pub amount: u32,
36 pub limit: u32,
38}
39
40impl MeterEntry {
41 fn new(limit: u32) -> Self {
43 Self { limit, amount: Default::default() }
44 }
45
46 fn exceeds_limit(&self, amount: u32) -> bool {
48 self.amount.saturating_add(amount) > self.limit
49 }
50
51 fn absorb(&mut self, rhs: Self) {
53 self.amount.saturating_accrue(rhs.amount)
54 }
55}
56
57#[derive(DefaultNoBound, Clone)]
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 fn new(memory_limit: u32) -> Self {
70 Self { root_meter: MeterEntry::new(memory_limit), ..Default::default() }
71 }
72
73 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 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 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 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 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 #[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 pub fn current_mut(&mut self) -> &mut MeterEntry {
125 self.nested_meters.last_mut().unwrap_or(&mut self.root_meter)
126 }
127
128 pub fn current(&self) -> &MeterEntry {
130 self.nested_meters.last().unwrap_or(&self.root_meter)
131 }
132}
133
134#[derive(Clone)]
136struct JournalEntry {
137 key: Vec<u8>,
138 prev_value: Option<Vec<u8>>,
139}
140
141impl JournalEntry {
142 fn new(key: Vec<u8>, prev_value: Option<Vec<u8>>) -> Self {
144 Self { key, prev_value }
145 }
146
147 fn revert(self, storage: &mut Storage) {
149 storage.write(&self.key, self.prev_value);
150 }
151}
152
153#[derive(Clone)]
155struct Journal(Vec<JournalEntry>);
156
157impl Journal {
158 fn new() -> Self {
160 Self(Default::default())
161 }
162
163 fn push(&mut self, entry: JournalEntry) {
165 self.0.push(entry);
166 }
167
168 fn len(&self) -> usize {
170 self.0.len()
171 }
172
173 fn rollback(&mut self, storage: &mut Storage, checkpoint: usize) {
175 self.0.drain(checkpoint..).rev().for_each(|entry| entry.revert(storage));
176 }
177}
178
179#[derive(Default, Clone)]
181struct Storage(BTreeMap<Vec<u8>, Vec<u8>>);
182
183impl Storage {
184 fn read(&self, key: &Vec<u8>) -> Option<Vec<u8>> {
186 self.0.get(key).cloned()
187 }
188
189 fn write(&mut self, key: &Vec<u8>, value: Option<Vec<u8>>) -> Option<Vec<u8>> {
191 if let Some(value) = value {
192 self.0.insert(key.clone(), value)
194 } else {
195 self.0.remove(key)
197 }
198 }
199}
200
201#[derive(Clone)]
209pub struct TransientStorage<T: Config> {
210 storage: Storage,
212 journal: Journal,
213 meter: StorageMeter<T>,
215 checkpoints: Vec<usize>,
217}
218
219impl<T: Config> TransientStorage<T> {
220 pub fn new(memory_limit: u32) -> Self {
222 TransientStorage {
223 storage: Default::default(),
224 journal: Journal::new(),
225 checkpoints: Default::default(),
226 meter: StorageMeter::new(memory_limit),
227 }
228 }
229
230 pub fn read(&self, account: &AccountIdOf<T>, key: &Key) -> Option<Vec<u8>> {
232 self.storage.read(&Self::storage_key(&account.encode(), &key.hash()))
233 }
234
235 pub fn write(
241 &mut self,
242 account: &AccountIdOf<T>,
243 key: &Key,
244 value: Option<Vec<u8>>,
245 take: bool,
246 ) -> Result<WriteOutcome, DispatchError> {
247 let key = Self::storage_key(&account.encode(), &key.hash());
248 let prev_value = self.storage.read(&key);
249 if prev_value != value {
251 if let Some(value) = &value {
253 let key_len = key.capacity();
258 let mut amount = value
259 .capacity()
260 .saturating_add(key_len)
261 .saturating_add(mem::size_of::<JournalEntry>());
262 if prev_value.is_none() {
263 amount.saturating_accrue(key_len.saturating_add(mem::size_of::<Vec<u8>>()));
268 }
269 self.meter.charge(amount as _)?;
270 }
271 self.storage.write(&key, value);
272 self.journal.push(JournalEntry::new(key, prev_value.clone()));
274 }
275
276 Ok(match (take, prev_value) {
277 (_, None) => WriteOutcome::New,
278 (false, Some(prev_value)) => WriteOutcome::Overwritten(prev_value.len() as _),
279 (true, Some(prev_value)) => WriteOutcome::Taken(prev_value),
280 })
281 }
282
283 pub fn start_transaction(&mut self) {
289 self.meter.start();
290 self.checkpoints.push(self.journal.len());
291 }
292
293 pub fn rollback_transaction(&mut self) {
301 let checkpoint = self
302 .checkpoints
303 .pop()
304 .expect(
305 "A call to rollback_transaction must be preceded by a corresponding call to start_transaction;
306 the code within this crate makes sure that this is always the case; qed"
307 );
308 self.meter.revert();
309 self.journal.rollback(&mut self.storage, checkpoint);
310 }
311
312 pub fn commit_transaction(&mut self) {
320 self.checkpoints
321 .pop()
322 .expect(
323 "A call to commit_transaction must be preceded by a corresponding call to start_transaction;
324 the code within this crate makes sure that this is always the case; qed"
325 );
326 self.meter.commit();
327 }
328
329 #[cfg(any(test, feature = "runtime-benchmarks"))]
331 pub fn meter(&mut self) -> &mut StorageMeter<T> {
332 return &mut self.meter;
333 }
334
335 fn storage_key(account: &[u8], key: &[u8]) -> Vec<u8> {
336 let mut storage_key = Vec::with_capacity(account.len() + key.len());
337 storage_key.extend_from_slice(&account);
338 storage_key.extend_from_slice(&key);
339 storage_key
340 }
341}
342
343#[cfg(test)]
344mod tests {
345 use super::*;
346 use crate::{Error, test_utils::*, tests::Test};
347 use core::u32::MAX;
348
349 fn allocation_size(account: &AccountIdOf<Test>, key: &Key, value: Option<Vec<u8>>) -> u32 {
351 let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(MAX);
352 storage
353 .write(account, key, value, false)
354 .expect("Could not write to transient storage.");
355 storage.meter().current().amount
356 }
357
358 #[test]
359 fn read_write_works() {
360 let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(2048);
361 assert_eq!(
362 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
363 Ok(WriteOutcome::New)
364 );
365 assert_eq!(
366 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![2]), true),
367 Ok(WriteOutcome::New)
368 );
369 assert_eq!(
370 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![3]), false),
371 Ok(WriteOutcome::New)
372 );
373 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
374 assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![2]));
375 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![3]));
376 assert_eq!(
378 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![4, 5]), false),
379 Ok(WriteOutcome::Overwritten(1))
380 );
381 assert_eq!(
382 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![6, 7]), true),
383 Ok(WriteOutcome::Taken(vec![3]))
384 );
385 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
386 assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![4, 5]));
387 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![6, 7]));
388
389 assert_eq!(
391 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![]), true),
392 Ok(WriteOutcome::Taken(vec![6, 7]))
393 );
394 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![]));
395
396 assert_eq!(
397 storage.write(&BOB, &Key::Fix([3; 32]), None, true),
398 Ok(WriteOutcome::Taken(vec![]))
399 );
400 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), None);
401 }
402
403 #[test]
404 fn read_write_with_var_sized_keys_works() {
405 let mut storage = TransientStorage::<Test>::new(2048);
406 assert_eq!(
407 storage.write(
408 &ALICE,
409 &Key::try_from_var([1; 64].to_vec()).unwrap(),
410 Some(vec![1]),
411 false
412 ),
413 Ok(WriteOutcome::New)
414 );
415 assert_eq!(
416 storage.write(
417 &BOB,
418 &Key::try_from_var([2; 64].to_vec()).unwrap(),
419 Some(vec![2, 3]),
420 false
421 ),
422 Ok(WriteOutcome::New)
423 );
424 assert_eq!(
425 storage.read(&ALICE, &Key::try_from_var([1; 64].to_vec()).unwrap()),
426 Some(vec![1])
427 );
428 assert_eq!(
429 storage.read(&BOB, &Key::try_from_var([2; 64].to_vec()).unwrap()),
430 Some(vec![2, 3])
431 );
432 assert_eq!(
434 storage.write(
435 &ALICE,
436 &Key::try_from_var([1; 64].to_vec()).unwrap(),
437 Some(vec![4, 5]),
438 false
439 ),
440 Ok(WriteOutcome::Overwritten(1))
441 );
442 assert_eq!(
443 storage.read(&ALICE, &Key::try_from_var([1; 64].to_vec()).unwrap()),
444 Some(vec![4, 5])
445 );
446 }
447
448 #[test]
449 fn rollback_transaction_works() {
450 let mut storage = TransientStorage::<Test>::new(1024);
451
452 storage.start_transaction();
453 assert_eq!(
454 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
455 Ok(WriteOutcome::New)
456 );
457 storage.rollback_transaction();
458 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None)
459 }
460
461 #[test]
462 fn commit_transaction_works() {
463 let mut storage = TransientStorage::<Test>::new(1024);
464
465 storage.start_transaction();
466 assert_eq!(
467 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
468 Ok(WriteOutcome::New)
469 );
470 storage.commit_transaction();
471 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]))
472 }
473
474 #[test]
475 fn overwrite_and_commmit_transaction_works() {
476 let mut storage = TransientStorage::<Test>::new(1024);
477 storage.start_transaction();
478 assert_eq!(
479 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
480 Ok(WriteOutcome::New)
481 );
482 assert_eq!(
483 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1, 2]), false),
484 Ok(WriteOutcome::Overwritten(1))
485 );
486 storage.commit_transaction();
487 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1, 2]))
488 }
489
490 #[test]
491 fn rollback_in_nested_transaction_works() {
492 let mut storage = TransientStorage::<Test>::new(1024);
493 storage.start_transaction();
494 assert_eq!(
495 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
496 Ok(WriteOutcome::New)
497 );
498 storage.start_transaction();
499 assert_eq!(
500 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1]), false),
501 Ok(WriteOutcome::New)
502 );
503 storage.rollback_transaction();
504 storage.commit_transaction();
505 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
506 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None)
507 }
508
509 #[test]
510 fn commit_in_nested_transaction_works() {
511 let mut storage = TransientStorage::<Test>::new(1024);
512 storage.start_transaction();
513 assert_eq!(
514 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
515 Ok(WriteOutcome::New)
516 );
517 storage.start_transaction();
518 assert_eq!(
519 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
520 Ok(WriteOutcome::New)
521 );
522 storage.start_transaction();
523 assert_eq!(
524 storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
525 Ok(WriteOutcome::New)
526 );
527 storage.commit_transaction();
528 storage.commit_transaction();
529 storage.commit_transaction();
530 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
531 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), Some(vec![2]));
532 assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), Some(vec![3]));
533 }
534
535 #[test]
536 fn rollback_all_transactions_works() {
537 let mut storage = TransientStorage::<Test>::new(1024);
538 storage.start_transaction();
539 assert_eq!(
540 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
541 Ok(WriteOutcome::New)
542 );
543 storage.start_transaction();
544 assert_eq!(
545 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
546 Ok(WriteOutcome::New)
547 );
548 storage.start_transaction();
549 assert_eq!(
550 storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
551 Ok(WriteOutcome::New)
552 );
553 storage.commit_transaction();
554 storage.commit_transaction();
555 storage.rollback_transaction();
556 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None);
557 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None);
558 assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), None);
559 }
560
561 #[test]
562 fn metering_transactions_works() {
563 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
564 let mut storage = TransientStorage::<Test>::new(size * 2);
565 storage.start_transaction();
566 assert_eq!(
567 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
568 Ok(WriteOutcome::New)
569 );
570 let limit = storage.meter().current().limit;
571 storage.commit_transaction();
572
573 storage.start_transaction();
574 assert_eq!(storage.meter().current().limit, limit - size);
575 assert_eq!(storage.meter().current().limit - storage.meter().current().amount, size);
576 assert_eq!(
577 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
578 Ok(WriteOutcome::New)
579 );
580 assert_eq!(storage.meter().current().amount, size);
581 storage.commit_transaction();
582 assert_eq!(storage.meter().total_amount(), size * 2);
583 }
584
585 #[test]
586 fn metering_nested_transactions_works() {
587 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
588 let mut storage = TransientStorage::<Test>::new(size * 3);
589
590 storage.start_transaction();
591 let limit = storage.meter().current().limit;
592 assert_eq!(
593 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
594 Ok(WriteOutcome::New)
595 );
596 storage.start_transaction();
597 assert_eq!(storage.meter().total_amount(), size);
598 assert!(storage.meter().current().limit < limit - size);
599 assert_eq!(
600 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
601 Ok(WriteOutcome::New)
602 );
603 storage.commit_transaction();
604 assert_eq!(storage.meter().current().limit, limit);
605 assert_eq!(storage.meter().total_amount(), storage.meter().current().amount);
606 storage.commit_transaction();
607 }
608
609 #[test]
610 fn metering_transaction_fails() {
611 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
612 let mut storage = TransientStorage::<Test>::new(size - 1);
613 storage.start_transaction();
614 assert_eq!(
615 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
616 Err(Error::<Test>::OutOfTransientStorage.into())
617 );
618 assert_eq!(storage.meter.current().amount, 0);
619 storage.commit_transaction();
620 assert_eq!(storage.meter.total_amount(), 0);
621 }
622
623 #[test]
624 fn metering_nested_transactions_fails() {
625 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
626 let mut storage = TransientStorage::<Test>::new(size * 2);
627
628 storage.start_transaction();
629 assert_eq!(
630 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
631 Ok(WriteOutcome::New)
632 );
633 storage.start_transaction();
634 assert_eq!(
635 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
636 Err(Error::<Test>::OutOfTransientStorage.into())
637 );
638 storage.commit_transaction();
639 storage.commit_transaction();
640 assert_eq!(storage.meter.total_amount(), size);
641 }
642
643 #[test]
644 fn metering_nested_transaction_with_rollback_works() {
645 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
646 let mut storage = TransientStorage::<Test>::new(size * 2);
647
648 storage.start_transaction();
649 let limit = storage.meter.current().limit;
650 storage.start_transaction();
651 assert_eq!(
652 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
653 Ok(WriteOutcome::New)
654 );
655 storage.rollback_transaction();
656
657 assert_eq!(storage.meter.total_amount(), 0);
658 assert_eq!(storage.meter.current().limit, limit);
659 assert_eq!(
660 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
661 Ok(WriteOutcome::New)
662 );
663 let amount = storage.meter().current().amount;
664 assert_eq!(storage.meter().total_amount(), amount);
665 storage.commit_transaction();
666 }
667
668 #[test]
669 fn metering_with_rollback_works() {
670 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
671 let mut storage = TransientStorage::<Test>::new(size * 5);
672
673 storage.start_transaction();
674 assert_eq!(
675 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
676 Ok(WriteOutcome::New)
677 );
678 let amount = storage.meter.total_amount();
679 storage.start_transaction();
680 assert_eq!(
681 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
682 Ok(WriteOutcome::New)
683 );
684 storage.start_transaction();
685 assert_eq!(
686 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
687 Ok(WriteOutcome::New)
688 );
689 storage.commit_transaction();
690 storage.rollback_transaction();
691 assert_eq!(storage.meter.total_amount(), amount);
692 storage.commit_transaction();
693 }
694}