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<T>) -> 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<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 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::{
344 tests::{Test, ALICE, BOB, CHARLIE},
345 Error,
346 };
347 use core::u32::MAX;
348
349 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 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 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 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}