rune/runtime/
budget.rs

Help
1//! Budgeting module for Runestick.
2//!
3//! This module contains methods which allows for limiting the execution of the
4//! virtual machine to abide by the specified budget.
5//!
6//! By default the budget is disabled, but can be enabled by wrapping your
7//! function call in [with].
8
9#[cfg_attr(feature = "std", path = "budget/std.rs")]
10mod no_std;
11
12use core::future::Future;
13use core::pin::Pin;
14use core::task::{Context, Poll};
15
16use crate::alloc::callable::Callable;
17
18use pin_project::pin_project;
19
20/// Wrapper for something being [budgeted].
21///
22/// See [with].
23///
24/// [budgeted]: self
25#[pin_project]
26pub struct Budget<T> {
27    /// Instruction budget.
28    budget: usize,
29    /// The thing being budgeted.
30    #[pin]
31    value: T,
32}
33
34/// Wrap the given value with a budget.
35///
36/// Budgeting is only performed on a per-instruction basis in the virtual
37/// machine. What exactly constitutes an instruction might be a bit vague. But
38/// important to note is that without explicit co-operation from native
39/// functions the budget cannot be enforced. So care must be taken with the
40/// native functions that you provide to Rune to ensure that the limits you
41/// impose cannot be circumvented.
42///
43/// The following things can be wrapped:
44/// * A [`FnOnce`] closure, like `with(|| println!("Hello World")).call()`.
45/// * A [`Future`], like `with(async { /* async work */ }).await`;
46///
47/// It's also possible to wrap other wrappers which implement [`Callable`].
48///
49/// # Examples
50///
51/// ```no_run
52/// use rune::runtime::budget;
53/// use rune::Vm;
54///
55/// let mut vm: Vm = todo!();
56/// // The virtual machine and any tasks associated with it is only allowed to execute 100 budget.
57/// budget::with(100, || vm.call(&["main"], ())).call()?;
58/// # Ok::<(), rune::support::Error>(())
59/// ```
60///
61/// This budget can be conveniently combined with the memory [`limit`] module
62/// due to both wrappers implementing [`Callable`].
63///
64/// [`limit`]: crate::alloc::limit
65///
66/// ```
67/// use rune::runtime::budget;
68/// use rune::alloc::{limit, Vec};
69///
70/// #[derive(Debug, PartialEq)]
71/// struct Marker;
72///
73/// // Limit the given closure to run one instruction and allocate 1024 bytes.
74/// let f = budget::with(1, limit::with(1024, || {
75///     let mut budget = budget::acquire();
76///     assert!(budget.take());
77///     assert!(!budget.take());
78///     assert!(Vec::<u8>::try_with_capacity(1).is_ok());
79///     assert!(Vec::<u8>::try_with_capacity(1024).is_ok());
80///     assert!(Vec::<u8>::try_with_capacity(1025).is_err());
81///     Marker
82/// }));
83///
84/// assert_eq!(f.call(), Marker);
85/// ```
86pub fn with<T>(budget: usize, value: T) -> Budget<T> {
87    tracing::trace!(?budget);
88    Budget { budget, value }
89}
90
91/// Acquire the current budget.
92///
93/// Use [`BudgetGuard::take`] to take permites from the returned budget.
94#[inline(never)]
95pub fn acquire() -> BudgetGuard {
96    BudgetGuard(self::no_std::rune_budget_replace(usize::MAX))
97}
98
99/// A locally acquired budget.
100///
101/// This guard is acquired by calling [`take`] and can be used to take permits.
102///
103/// [`take`]: BudgetGuard::take
104#[repr(transparent)]
105pub struct BudgetGuard(usize);
106
107impl BudgetGuard {
108    /// Take a ticker from the budget.
109    pub fn take(&mut self) -> bool {
110        if self.0 == usize::MAX {
111            return true;
112        }
113
114        if self.0 == 0 {
115            return false;
116        }
117
118        self.0 -= 1;
119        true
120    }
121}
122
123impl Drop for BudgetGuard {
124    fn drop(&mut self) {
125        let _ = self::no_std::rune_budget_replace(self.0);
126    }
127}
128
129impl<T> Budget<T>
130where
131    T: Callable,
132{
133    /// Call the budgeted function.
134    pub fn call(self) -> T::Output {
135        Callable::call(self)
136    }
137}
138
139impl<T> Callable for Budget<T>
140where
141    T: Callable,
142{
143    type Output = T::Output;
144
145    #[inline]
146    fn call(self) -> Self::Output {
147        let _guard = BudgetGuard(self::no_std::rune_budget_replace(self.budget));
148        self.value.call()
149    }
150}
151
152impl<T> Future for Budget<T>
153where
154    T: Future,
155{
156    type Output = T::Output;
157
158    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
159        let this = self.project();
160
161        let _guard = BudgetGuard(self::no_std::rune_budget_replace(*this.budget));
162        let poll = this.value.poll(cx);
163        *this.budget = self::no_std::rune_budget_get();
164        poll
165    }
166}
rune/runtime/budget.rs

rune/runtime/
budget.rs
