mysten_common/rpc_format/
meter.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4/// A trait for tracking the resource limits consumed while serializing a value.
5pub trait Meter {
6    type Nested<'a>: Meter
7    where
8        Self: 'a;
9
10    /// Produce a new meter for nested contexts, consuming one unit of nesting budget.
11    fn nest(&mut self) -> Result<Self::Nested<'_>, MeterError>;
12
13    /// Charge `amount` units against the remaining size budget.
14    fn charge(&mut self, amount: usize) -> Result<(), MeterError>;
15}
16
17#[derive(Clone, Copy, Default)]
18pub struct Unmetered;
19
20/// A meter backed by mutable local budget state.
21///
22/// "Local" means nested visitors share a single remaining size budget reference while each meter
23/// handle keeps its own remaining depth budget.
24pub struct LocalMeter<'a> {
25    size_budget: &'a mut usize,
26    depth_budget: usize,
27}
28
29#[derive(thiserror::Error, Debug, Clone, Copy)]
30pub enum MeterError {
31    #[error("Deserialized value too large")]
32    TooBig,
33
34    #[error("Exceeded maximum depth")]
35    TooDeep,
36}
37
38impl<'a> LocalMeter<'a> {
39    pub fn new(size_budget: &'a mut usize, depth_budget: usize) -> Self {
40        Self {
41            size_budget,
42            depth_budget,
43        }
44    }
45}
46
47impl Meter for Unmetered {
48    type Nested<'a>
49        = Self
50    where
51        Self: 'a;
52
53    fn nest(&mut self) -> Result<Self::Nested<'_>, MeterError> {
54        Ok(*self)
55    }
56
57    fn charge(&mut self, _amount: usize) -> Result<(), MeterError> {
58        Ok(())
59    }
60}
61
62impl Meter for LocalMeter<'_> {
63    type Nested<'a>
64        = LocalMeter<'a>
65    where
66        Self: 'a;
67
68    fn nest(&mut self) -> Result<Self::Nested<'_>, MeterError> {
69        if self.depth_budget == 0 {
70            Err(MeterError::TooDeep)
71        } else {
72            Ok(LocalMeter::new(self.size_budget, self.depth_budget - 1))
73        }
74    }
75
76    fn charge(&mut self, amount: usize) -> Result<(), MeterError> {
77        if *self.size_budget < amount {
78            Err(MeterError::TooBig)
79        } else {
80            *self.size_budget -= amount;
81            Ok(())
82        }
83    }
84}