rune/runtime/
budget.rs

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

rune/runtime/
budget.rs
