1use 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#[derive(Default, Debug)]
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)]
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
134struct JournalEntry {
136 key: Vec<u8>,
137 prev_value: Option<Vec<u8>>,
138}
139
140impl JournalEntry {
141 fn new(key: Vec<u8>, prev_value: Option<Vec<u8>>) -> Self {
143 Self { key, prev_value }
144 }
145
146 fn revert(self, storage: &mut Storage) {
148 storage.write(&self.key, self.prev_value);
149 }
150}
151
152struct Journal(Vec<JournalEntry>);
154
155impl Journal {
156 fn new() -> Self {
158 Self(Default::default())
159 }
160
161 fn push(&mut self, entry: JournalEntry) {
163 self.0.push(entry);
164 }
165
166 fn len(&self) -> usize {
168 self.0.len()
169 }
170
171 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#[derive(Default)]
179struct Storage(BTreeMap<Vec<u8>, Vec<u8>>);
180
181impl Storage {
182 fn read(&self, key: &Vec<u8>) -> Option<Vec<u8>> {
184 self.0.get(key).cloned()
185 }
186
187 fn write(&mut self, key: &Vec<u8>, value: Option<Vec<u8>>) -> Option<Vec<u8>> {
189 if let Some(value) = value {
190 self.0.insert(key.clone(), value)
192 } else {
193 self.0.remove(key)
195 }
196 }
197}
198
199pub struct TransientStorage<T: Config> {
207 storage: Storage,
209 journal: Journal,
210 meter: StorageMeter<T>,
212 checkpoints: Vec<usize>,
214}
215
216impl<T: Config> TransientStorage<T> {
217 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 pub fn read(&self, account: &AccountIdOf<T>, key: &Key) -> Option<Vec<u8>> {
229 self.storage.read(&Self::storage_key(&account.encode(), &key.hash()))
230 }
231
232 pub fn write(
238 &mut self,
239 account: &AccountIdOf<T>,
240 key: &Key,
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 if prev_value != value {
248 if let Some(value) = &value {
250 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 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 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 pub fn start_transaction(&mut self) {
286 self.meter.start();
287 self.checkpoints.push(self.journal.len());
288 }
289
290 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 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 #[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::{test_utils::*, tests::Test, Error};
344 use core::u32::MAX;
345
346 fn allocation_size(account: &AccountIdOf<Test>, key: &Key, value: Option<Vec<u8>>) -> u32 {
348 let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(MAX);
349 storage
350 .write(account, key, value, false)
351 .expect("Could not write to transient storage.");
352 storage.meter().current().amount
353 }
354
355 #[test]
356 fn read_write_works() {
357 let mut storage: TransientStorage<Test> = TransientStorage::<Test>::new(2048);
358 assert_eq!(
359 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
360 Ok(WriteOutcome::New)
361 );
362 assert_eq!(
363 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![2]), true),
364 Ok(WriteOutcome::New)
365 );
366 assert_eq!(
367 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![3]), false),
368 Ok(WriteOutcome::New)
369 );
370 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
371 assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![2]));
372 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![3]));
373 assert_eq!(
375 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![4, 5]), false),
376 Ok(WriteOutcome::Overwritten(1))
377 );
378 assert_eq!(
379 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![6, 7]), true),
380 Ok(WriteOutcome::Taken(vec![3]))
381 );
382 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
383 assert_eq!(storage.read(&ALICE, &Key::Fix([2; 32])), Some(vec![4, 5]));
384 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![6, 7]));
385
386 assert_eq!(
388 storage.write(&BOB, &Key::Fix([3; 32]), Some(vec![]), true),
389 Ok(WriteOutcome::Taken(vec![6, 7]))
390 );
391 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), Some(vec![]));
392
393 assert_eq!(
394 storage.write(&BOB, &Key::Fix([3; 32]), None, true),
395 Ok(WriteOutcome::Taken(vec![]))
396 );
397 assert_eq!(storage.read(&BOB, &Key::Fix([3; 32])), None);
398 }
399
400 #[test]
401 fn read_write_with_var_sized_keys_works() {
402 let mut storage = TransientStorage::<Test>::new(2048);
403 assert_eq!(
404 storage.write(
405 &ALICE,
406 &Key::try_from_var([1; 64].to_vec()).unwrap(),
407 Some(vec![1]),
408 false
409 ),
410 Ok(WriteOutcome::New)
411 );
412 assert_eq!(
413 storage.write(
414 &BOB,
415 &Key::try_from_var([2; 64].to_vec()).unwrap(),
416 Some(vec![2, 3]),
417 false
418 ),
419 Ok(WriteOutcome::New)
420 );
421 assert_eq!(
422 storage.read(&ALICE, &Key::try_from_var([1; 64].to_vec()).unwrap()),
423 Some(vec![1])
424 );
425 assert_eq!(
426 storage.read(&BOB, &Key::try_from_var([2; 64].to_vec()).unwrap()),
427 Some(vec![2, 3])
428 );
429 assert_eq!(
431 storage.write(
432 &ALICE,
433 &Key::try_from_var([1; 64].to_vec()).unwrap(),
434 Some(vec![4, 5]),
435 false
436 ),
437 Ok(WriteOutcome::Overwritten(1))
438 );
439 assert_eq!(
440 storage.read(&ALICE, &Key::try_from_var([1; 64].to_vec()).unwrap()),
441 Some(vec![4, 5])
442 );
443 }
444
445 #[test]
446 fn rollback_transaction_works() {
447 let mut storage = TransientStorage::<Test>::new(1024);
448
449 storage.start_transaction();
450 assert_eq!(
451 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
452 Ok(WriteOutcome::New)
453 );
454 storage.rollback_transaction();
455 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None)
456 }
457
458 #[test]
459 fn commit_transaction_works() {
460 let mut storage = TransientStorage::<Test>::new(1024);
461
462 storage.start_transaction();
463 assert_eq!(
464 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
465 Ok(WriteOutcome::New)
466 );
467 storage.commit_transaction();
468 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]))
469 }
470
471 #[test]
472 fn overwrite_and_commmit_transaction_works() {
473 let mut storage = TransientStorage::<Test>::new(1024);
474 storage.start_transaction();
475 assert_eq!(
476 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
477 Ok(WriteOutcome::New)
478 );
479 assert_eq!(
480 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1, 2]), false),
481 Ok(WriteOutcome::Overwritten(1))
482 );
483 storage.commit_transaction();
484 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1, 2]))
485 }
486
487 #[test]
488 fn rollback_in_nested_transaction_works() {
489 let mut storage = TransientStorage::<Test>::new(1024);
490 storage.start_transaction();
491 assert_eq!(
492 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
493 Ok(WriteOutcome::New)
494 );
495 storage.start_transaction();
496 assert_eq!(
497 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1]), false),
498 Ok(WriteOutcome::New)
499 );
500 storage.rollback_transaction();
501 storage.commit_transaction();
502 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
503 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None)
504 }
505
506 #[test]
507 fn commit_in_nested_transaction_works() {
508 let mut storage = TransientStorage::<Test>::new(1024);
509 storage.start_transaction();
510 assert_eq!(
511 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
512 Ok(WriteOutcome::New)
513 );
514 storage.start_transaction();
515 assert_eq!(
516 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
517 Ok(WriteOutcome::New)
518 );
519 storage.start_transaction();
520 assert_eq!(
521 storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
522 Ok(WriteOutcome::New)
523 );
524 storage.commit_transaction();
525 storage.commit_transaction();
526 storage.commit_transaction();
527 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), Some(vec![1]));
528 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), Some(vec![2]));
529 assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), Some(vec![3]));
530 }
531
532 #[test]
533 fn rollback_all_transactions_works() {
534 let mut storage = TransientStorage::<Test>::new(1024);
535 storage.start_transaction();
536 assert_eq!(
537 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1]), false),
538 Ok(WriteOutcome::New)
539 );
540 storage.start_transaction();
541 assert_eq!(
542 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![2]), false),
543 Ok(WriteOutcome::New)
544 );
545 storage.start_transaction();
546 assert_eq!(
547 storage.write(&CHARLIE, &Key::Fix([1; 32]), Some(vec![3]), false),
548 Ok(WriteOutcome::New)
549 );
550 storage.commit_transaction();
551 storage.commit_transaction();
552 storage.rollback_transaction();
553 assert_eq!(storage.read(&ALICE, &Key::Fix([1; 32])), None);
554 assert_eq!(storage.read(&BOB, &Key::Fix([1; 32])), None);
555 assert_eq!(storage.read(&CHARLIE, &Key::Fix([1; 32])), None);
556 }
557
558 #[test]
559 fn metering_transactions_works() {
560 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
561 let mut storage = TransientStorage::<Test>::new(size * 2);
562 storage.start_transaction();
563 assert_eq!(
564 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
565 Ok(WriteOutcome::New)
566 );
567 let limit = storage.meter().current().limit;
568 storage.commit_transaction();
569
570 storage.start_transaction();
571 assert_eq!(storage.meter().current().limit, limit - size);
572 assert_eq!(storage.meter().current().limit - storage.meter().current().amount, size);
573 assert_eq!(
574 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
575 Ok(WriteOutcome::New)
576 );
577 assert_eq!(storage.meter().current().amount, size);
578 storage.commit_transaction();
579 assert_eq!(storage.meter().total_amount(), size * 2);
580 }
581
582 #[test]
583 fn metering_nested_transactions_works() {
584 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
585 let mut storage = TransientStorage::<Test>::new(size * 3);
586
587 storage.start_transaction();
588 let limit = storage.meter().current().limit;
589 assert_eq!(
590 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
591 Ok(WriteOutcome::New)
592 );
593 storage.start_transaction();
594 assert_eq!(storage.meter().total_amount(), size);
595 assert!(storage.meter().current().limit < limit - size);
596 assert_eq!(
597 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
598 Ok(WriteOutcome::New)
599 );
600 storage.commit_transaction();
601 assert_eq!(storage.meter().current().limit, limit);
602 assert_eq!(storage.meter().total_amount(), storage.meter().current().amount);
603 storage.commit_transaction();
604 }
605
606 #[test]
607 fn metering_transaction_fails() {
608 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
609 let mut storage = TransientStorage::<Test>::new(size - 1);
610 storage.start_transaction();
611 assert_eq!(
612 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
613 Err(Error::<Test>::OutOfTransientStorage.into())
614 );
615 assert_eq!(storage.meter.current().amount, 0);
616 storage.commit_transaction();
617 assert_eq!(storage.meter.total_amount(), 0);
618 }
619
620 #[test]
621 fn metering_nested_transactions_fails() {
622 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
623 let mut storage = TransientStorage::<Test>::new(size * 2);
624
625 storage.start_transaction();
626 assert_eq!(
627 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
628 Ok(WriteOutcome::New)
629 );
630 storage.start_transaction();
631 assert_eq!(
632 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
633 Err(Error::<Test>::OutOfTransientStorage.into())
634 );
635 storage.commit_transaction();
636 storage.commit_transaction();
637 assert_eq!(storage.meter.total_amount(), size);
638 }
639
640 #[test]
641 fn metering_nested_transaction_with_rollback_works() {
642 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
643 let mut storage = TransientStorage::<Test>::new(size * 2);
644
645 storage.start_transaction();
646 let limit = storage.meter.current().limit;
647 storage.start_transaction();
648 assert_eq!(
649 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
650 Ok(WriteOutcome::New)
651 );
652 storage.rollback_transaction();
653
654 assert_eq!(storage.meter.total_amount(), 0);
655 assert_eq!(storage.meter.current().limit, limit);
656 assert_eq!(
657 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
658 Ok(WriteOutcome::New)
659 );
660 let amount = storage.meter().current().amount;
661 assert_eq!(storage.meter().total_amount(), amount);
662 storage.commit_transaction();
663 }
664
665 #[test]
666 fn metering_with_rollback_works() {
667 let size = allocation_size(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]));
668 let mut storage = TransientStorage::<Test>::new(size * 5);
669
670 storage.start_transaction();
671 assert_eq!(
672 storage.write(&ALICE, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
673 Ok(WriteOutcome::New)
674 );
675 let amount = storage.meter.total_amount();
676 storage.start_transaction();
677 assert_eq!(
678 storage.write(&ALICE, &Key::Fix([2; 32]), Some(vec![1u8; 4096]), false),
679 Ok(WriteOutcome::New)
680 );
681 storage.start_transaction();
682 assert_eq!(
683 storage.write(&BOB, &Key::Fix([1; 32]), Some(vec![1u8; 4096]), false),
684 Ok(WriteOutcome::New)
685 );
686 storage.commit_transaction();
687 storage.rollback_transaction();
688 assert_eq!(storage.meter.total_amount(), amount);
689 storage.commit_transaction();
690 }
691}