1use std::future::Future;
24use std::sync::Arc;
25use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
26
27use wasmtime::{AsContextMut, Caller, Func, FuncType, Trap, Val};
28
29use super::StoreData;
30
31#[derive(Clone, Copy, Debug)]
35pub struct Rate {
36 fuel: u64,
37 per: u64,
38}
39
40impl Rate {
41 const fn per_unit(fuel: u64) -> Rate {
42 Rate { fuel, per: 1 }
43 }
44
45 const fn per_units(fuel: u64, per: u64) -> Rate {
46 Rate { fuel, per }
47 }
48
49 pub fn cost(self, n: u64) -> u64 {
51 n.saturating_mul(self.fuel).div_ceil(self.per)
52 }
53}
54
55pub const CALL: u64 = 16;
59pub const COPY: Rate = Rate::per_units(1, 8);
61pub const SCAN: Rate = Rate::per_unit(1);
64pub const PARSE: Rate = Rate::per_unit(6);
67pub const ELEM: Rate = Rate::per_unit(10);
69pub const REGEX: Rate = Rate::per_unit(8);
73pub const REGEX_COMPILE: u64 = 2_000;
76pub const HASH: Rate = Rate::per_unit(4);
78pub const IO: Rate = Rate::per_units(1, 4);
80pub const SYSCALL: Rate = Rate::per_unit(400);
82pub const TZ: u64 = 200;
84pub const GATE: u64 = 400;
86
87pub fn sort_cost(n: u64) -> u64 {
90 let log = u64::from(n.max(2).ilog2());
91 ELEM.cost(n.saturating_mul(log))
92}
93
94pub fn bigint_product_cost(a: u64, b: u64) -> u64 {
99 let (small, large) = (a.min(b).max(1), a.max(b).max(1));
100 let schoolbook = small.saturating_mul(large.min(64));
101 let beyond = large
102 .saturating_sub(64)
103 .saturating_mul(u64::from(large.ilog2()) + 1);
104 ELEM.cost(schoolbook.saturating_add(beyond))
105}
106
107pub fn charge_call(ctx: impl AsContextMut<Data = StoreData>) -> wasmtime::Result<()> {
109 charge_host_fuel(ctx, CALL)
110}
111
112pub fn charge(
114 ctx: impl AsContextMut<Data = StoreData>,
115 rate: Rate,
116 n: u64,
117) -> wasmtime::Result<()> {
118 charge_host_fuel(ctx, rate.cost(n))
119}
120
121pub(crate) fn host_func(
125 mut store: impl AsContextMut<Data = StoreData>,
126 ty: FuncType,
127 body: impl Fn(&mut Caller<'_, StoreData>, &[Val], &mut [Val]) -> wasmtime::Result<()>
128 + Send
129 + Sync
130 + 'static,
131) -> Func {
132 let abi = ty.clone();
133 Func::new(&mut store, ty, move |mut caller, params, results| {
134 super::host::check_host_abi(&abi, params, results)?;
135 charge_call(&mut caller)?;
136 body(&mut caller, params, results)
137 })
138}
139
140pub(crate) fn host_func_async<F>(
143 mut store: impl AsContextMut<Data = StoreData>,
144 ty: FuncType,
145 body: F,
146) -> Func
147where
148 F: for<'a> Fn(
149 Caller<'a, StoreData>,
150 &'a [Val],
151 &'a mut [Val],
152 ) -> Box<dyn Future<Output = wasmtime::Result<()>> + Send + 'a>
153 + Send
154 + Sync
155 + 'static,
156{
157 let abi = ty.clone();
158 Func::new_async(&mut store, ty, move |mut caller, params, results| {
159 if let Err(error) = super::host::check_host_abi(&abi, params, results) {
160 return Box::new(async move { Err(error) });
161 }
162 if let Err(error) = charge_call(&mut caller) {
163 return Box::new(async move { Err(error) });
164 }
165 body(caller, params, results)
166 })
167}
168
169pub fn settle(
174 ctx: impl AsContextMut<Data = StoreData>,
175 rate: Rate,
176 n: u64,
177) -> wasmtime::Result<()> {
178 settle_host_fuel(ctx, rate.cost(n))
179}
180
181pub(crate) fn settle_result<T, C: AsContextMut<Data = StoreData>>(
185 caller: &mut C,
186 body: impl FnOnce(&mut C) -> wasmtime::Result<T>,
187) -> wasmtime::Result<T> {
188 let previous = caller.as_context().data().settling_host_result;
189 caller.as_context_mut().data_mut().settling_host_result = true;
190 let result = body(caller);
191 caller.as_context_mut().data_mut().settling_host_result = previous;
192 result
193}
194
195pub fn settle_host_fuel(
197 mut ctx: impl AsContextMut<Data = StoreData>,
198 units: u64,
199) -> wasmtime::Result<()> {
200 let mut ctx = ctx.as_context_mut();
201 let pending = ctx.data().host_fuel_pending;
202 let available = ctx.get_fuel()?.saturating_sub(pending);
203 if units <= available {
204 return charge_host_fuel(ctx, units);
205 }
206 ctx.set_fuel(0)?;
207 let data = ctx.data_mut();
208 data.host_fuel_pending = 0;
209 data.host_fuel = data.host_fuel.saturating_add(available);
210 Ok(())
211}
212
213pub const HOST_FUEL_BATCH: u64 = 10_000;
219pub const HOST_FUEL_PENDING_CAP: u64 = 1_000 * HOST_FUEL_BATCH;
224
225pub fn charge_host_fuel(
236 mut ctx: impl AsContextMut<Data = StoreData>,
237 units: u64,
238) -> wasmtime::Result<()> {
239 if units == 0 {
240 return Ok(());
241 }
242 let mut ctx = ctx.as_context_mut();
243 let pending = ctx.data().host_fuel_pending;
244 let engine_fuel = ctx.get_fuel()?;
245 let available = engine_fuel.saturating_sub(pending);
246 if units > available {
247 ctx.set_fuel(0)?;
248 let data = ctx.data_mut();
249 data.host_fuel_pending = 0;
250 data.host_fuel = data.host_fuel.saturating_add(available);
251 if data.settling_host_result {
252 return Ok(());
253 }
254 return Err(Trap::OutOfFuel.into());
255 }
256 let batch = pending + units;
258 let data = ctx.data_mut();
259 data.host_fuel = data.host_fuel.saturating_add(units);
260 let wasm_burned = data
265 .host_fuel_applied_at
266 .map_or(u64::MAX, |mark| mark.saturating_sub(engine_fuel));
267 let apply = batch >= HOST_FUEL_PENDING_CAP
268 || (batch >= HOST_FUEL_BATCH && wasm_burned > HOST_FUEL_BATCH);
269 if !apply {
270 data.host_fuel_pending = batch;
271 return Ok(());
272 }
273 data.host_fuel_pending = 0;
274 data.host_fuel_applied_at = Some(engine_fuel - batch);
275 ctx.set_fuel(engine_fuel - batch)
276}
277
278#[derive(Debug, thiserror::Error)]
279#[error("host operation exhausted its fuel budget")]
280pub struct Exhausted;
281
282#[derive(Debug)]
286pub struct WorkBudget {
287 limit: u64,
288 spent: AtomicU64,
289 after_effect: AtomicBool,
290 cancel: Option<Arc<AtomicBool>>,
291 exhausted: AtomicBool,
292}
293impl WorkBudget {
294 pub fn new(limit: u64) -> Self {
295 Self {
296 limit,
297 spent: AtomicU64::new(0),
298 after_effect: AtomicBool::new(false),
299 cancel: None,
300 exhausted: AtomicBool::new(false),
301 }
302 }
303 pub fn with_cancellation(mut self, cancel: Arc<AtomicBool>) -> Self {
306 self.cancel = Some(cancel);
307 self
308 }
309
310 pub fn charge(&self, rate: Rate, count: u64) -> Result<(), Exhausted> {
311 self.charge_fuel(rate.cost(count))
312 }
313 pub fn charge_fuel(&self, amount: u64) -> Result<(), Exhausted> {
314 let settled = self.after_effect.load(Ordering::Acquire);
315 let mut refused = false;
316 let _ = self
317 .spent
318 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |spent| {
319 let remaining = self.limit.saturating_sub(spent);
320 refused = !settled && amount > remaining;
321 let admitted = if settled {
322 amount
323 } else {
324 amount.min(remaining)
325 };
326 Some(spent.saturating_add(admitted))
327 });
328 if refused { Err(Exhausted) } else { Ok(()) }
329 }
330 pub fn record_fuel(&self, amount: u64) {
333 let _ = self
334 .spent
335 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |spent| {
336 Some(spent.saturating_add(amount))
337 });
338 if !self.after_effect.load(Ordering::Acquire) && self.spent() > self.limit {
339 self.exhausted.store(true, Ordering::Release);
340 if let Some(cancel) = &self.cancel {
341 cancel.store(true, Ordering::Release);
342 }
343 }
344 }
345
346 pub fn before_effect(&self) {
351 self.after_effect.store(true, Ordering::Release);
352 }
353 pub fn spent(&self) -> u64 {
354 self.spent.load(Ordering::Acquire)
355 }
356
357 pub fn is_exhausted(&self) -> bool {
359 self.exhausted.load(Ordering::Acquire)
360 }
361
362 pub fn settle(&self, ctx: impl AsContextMut<Data = StoreData>) -> wasmtime::Result<()> {
363 settle_host_fuel(ctx, self.spent())
364 }
365}
366
367pub fn remaining_fuel(mut ctx: impl AsContextMut<Data = StoreData>) -> wasmtime::Result<u64> {
369 let ctx = ctx.as_context_mut();
370 Ok(ctx.get_fuel()?.saturating_sub(ctx.data().host_fuel_pending))
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376 use crate::runtime::{RuntimeConfig, Vfs};
377 use wasmtime::{Func, FuncType, Val, ValType};
378
379 #[test]
380 fn refuses_before_effect_and_settles_after_publication_starts() {
381 let budget = WorkBudget::new(HASH.cost(4));
382 budget.charge(HASH, 3).unwrap();
383 assert!(budget.charge(HASH, 2).is_err());
384 assert_eq!(budget.spent(), HASH.cost(4));
385 budget.before_effect();
386 budget.charge(HASH, 2).unwrap();
387 assert_eq!(budget.spent(), HASH.cost(6));
388 }
389
390 #[test]
391 fn reported_work_cancels_before_effect_and_keeps_its_full_cost() {
392 let cancel = Arc::new(AtomicBool::new(false));
393 let budget = WorkBudget::new(10).with_cancellation(cancel.clone());
394 budget.record_fuel(10);
395 assert!(!cancel.load(Ordering::Acquire));
396 budget.record_fuel(3);
397 assert!(cancel.load(Ordering::Acquire));
398 assert!(budget.is_exhausted());
399 assert_eq!(budget.spent(), 13);
400 assert!(budget.charge_fuel(1).is_err());
401 assert_eq!(budget.spent(), 13, "admission must not erase reported work");
402 budget.before_effect();
403 budget.record_fuel(4);
404 budget.charge_fuel(5).unwrap();
405 assert_eq!(budget.spent(), 22);
406 assert!(
407 cancel.load(Ordering::Acquire),
408 "never clear an existing cancellation"
409 );
410 }
411
412 #[test]
413 fn effect_completion_does_not_request_fuel_cancellation() {
414 let cancel = Arc::new(AtomicBool::new(false));
415 let budget = WorkBudget::new(10).with_cancellation(cancel.clone());
416 budget.charge_fuel(10).unwrap();
417 assert!(budget.charge_fuel(1).is_err());
418 assert!(
419 !cancel.load(Ordering::Acquire),
420 "admission returns its own error"
421 );
422 budget.before_effect();
423 budget.record_fuel(20);
424 budget.charge_fuel(30).unwrap();
425 assert_eq!(budget.spent(), 60);
426 assert!(!cancel.load(Ordering::Acquire));
427 assert!(!budget.is_exhausted());
428 }
429
430 #[test]
431 fn worker_charges_share_one_limit_and_saturate_after_effects() {
432 let budget = WorkBudget::new(100);
433 std::thread::scope(|scope| {
434 let workers: Vec<_> = (0..4)
435 .map(|_| scope.spawn(|| (0..100).filter(|_| budget.charge_fuel(1).is_ok()).count()))
436 .collect();
437 let admitted: usize = workers
438 .into_iter()
439 .map(|worker| worker.join().unwrap())
440 .sum();
441 assert_eq!(admitted, 100);
442 });
443 assert_eq!(budget.spent(), 100);
444 budget.before_effect();
445 budget.charge_fuel(u64::MAX).unwrap();
446 budget.record_fuel(1);
447 assert_eq!(budget.spent(), u64::MAX);
448 let budget = WorkBudget::new(u64::MAX);
449 budget.charge_fuel(u64::MAX).unwrap();
450 assert!(budget.charge_fuel(1).is_err());
451 }
452
453 #[test]
454 fn worker_budget_accounts_for_pending_charges_and_settles_once_after_effect() {
455 let config = RuntimeConfig::default();
456 let engine = config.engine().unwrap();
457 let mut store = config
458 .store(&engine, StoreData::with_vfs("test-run".into(), Vfs::none()))
459 .unwrap();
460 store.set_fuel(100).unwrap();
461 charge_host_fuel(&mut store, 20).unwrap();
462 assert_eq!(remaining_fuel(&mut store).unwrap(), 80);
463 let budget = WorkBudget::new(remaining_fuel(&mut store).unwrap());
464 budget.charge_fuel(80).unwrap();
465 budget.before_effect();
466 budget.record_fuel(10);
467 budget.settle(&mut store).unwrap();
468 assert_eq!(store.get_fuel().unwrap(), 0);
469 assert_eq!(budget.spent(), 90);
470 assert_eq!(store.data().host_fuel, 100);
471 }
472
473 #[test]
474 fn result_settlement_preserves_values_and_restores_refusing_charges() {
475 let config = RuntimeConfig::default();
476 let engine = config.engine().unwrap();
477 let mut store = config
478 .store(&engine, StoreData::with_vfs("test-run".into(), Vfs::none()))
479 .unwrap();
480 let complete = Func::new(
481 &mut store,
482 FuncType::new(&engine, [], [ValType::I32]),
483 |mut caller, _, results| {
484 results[0] = settle_result(&mut caller, |caller| {
485 charge(caller, PARSE, 100)?;
486 Ok(Val::I32(42))
487 })?;
488 assert!(!caller.data().settling_host_result);
489 Ok(())
490 },
491 );
492 store.set_fuel(1).unwrap();
493 let mut results = [Val::I32(0)];
494 complete.call(&mut store, &[], &mut results).unwrap();
495 assert_eq!(results[0].i32(), Some(42));
496 assert_eq!(store.get_fuel().unwrap(), 0);
497 assert_eq!(store.data().host_fuel, 1);
498 assert!(charge_call(&mut store).unwrap_err().is::<Trap>());
499
500 let fail = Func::new(
501 &mut store,
502 FuncType::new(&engine, [], []),
503 |mut caller, _, _| {
504 let result = settle_result(&mut caller, |caller| {
505 charge(caller, COPY, 100)?;
506 Err::<(), _>(crate::runtime::host::range_error("copy entry limit"))
507 });
508 assert!(!caller.data().settling_host_result);
509 result
510 },
511 );
512 let error = fail.call(&mut store, &[], &mut []).unwrap_err();
513 assert!(error.to_string().contains("copy entry limit"), "{error}");
514 assert!(!store.data().settling_host_result);
515 assert!(charge_call(&mut store).unwrap_err().is::<Trap>());
516 }
517
518 #[test]
519 fn settling_never_refuses_and_leaves_a_short_budget_at_zero() {
520 let config = RuntimeConfig::default();
521 let engine = config.engine().unwrap();
522 let mut store = config
523 .store(&engine, StoreData::with_vfs("test-run".into(), Vfs::none()))
524 .unwrap();
525 let settle = Func::new(
526 &mut store,
527 FuncType::new(&engine, [ValType::I64], []),
528 |mut caller, params, _| settle_host_fuel(&mut caller, params[0].unwrap_i64() as u64),
529 );
530 let charge = Func::new(
531 &mut store,
532 FuncType::new(&engine, [ValType::I64], []),
533 |mut caller, params, _| charge_host_fuel(&mut caller, params[0].unwrap_i64() as u64),
534 );
535
536 store.set_fuel(1_000).unwrap();
538 settle.call(&mut store, &[Val::I64(300)], &mut []).unwrap();
539 assert_eq!(store.get_fuel().unwrap(), 1_000);
540 assert_eq!(store.data().host_fuel_pending, 300);
541 assert_eq!(store.data().host_fuel, 300);
542
543 settle
546 .call(&mut store, &[Val::I64(5_000)], &mut [])
547 .unwrap();
548 assert_eq!(store.get_fuel().unwrap(), 0);
549 assert_eq!(store.data().host_fuel_pending, 0);
550 assert_eq!(store.data().host_fuel, 1_000);
551 let error = charge
552 .call(&mut store, &[Val::I64(1)], &mut [])
553 .unwrap_err();
554 assert_eq!(error.downcast_ref::<Trap>(), Some(&Trap::OutOfFuel));
555 }
556
557 #[test]
558 fn the_batch_covers_the_default_yield_interval() {
559 assert!(RuntimeConfig::default().async_yield_fuel <= Some(HOST_FUEL_BATCH));
560 }
561
562 #[test]
563 fn rates_round_up_and_never_make_work_free() {
564 assert_eq!(COPY.cost(0), 0);
565 assert_eq!(COPY.cost(1), 1);
566 assert_eq!(COPY.cost(8), 1);
567 assert_eq!(COPY.cost(9), 2);
568 assert_eq!(PARSE.cost(3), 18);
569 assert_eq!(SCAN.cost(u64::MAX), u64::MAX);
570 assert_eq!(sort_cost(0), 0);
571 assert_eq!(sort_cost(1), ELEM.cost(1));
572 assert_eq!(sort_cost(1024), ELEM.cost(10 * 1024));
573 }
574
575 #[test]
576 fn host_charges_are_batched_counted_and_refused_before_the_work() {
577 use crate::runtime::limits::ExecutionUsage;
578 let config = RuntimeConfig::default();
579 let engine = config.engine().unwrap();
580 let mut store = config
581 .store(&engine, StoreData::with_vfs("test-run".into(), Vfs::none()))
582 .unwrap();
583 let charge = Func::new(
584 &mut store,
585 FuncType::new(&engine, [ValType::I64], []),
586 |mut caller, params, _| charge_host_fuel(&mut caller, params[0].unwrap_i64() as u64),
587 );
588 let charge_units = |store: &mut wasmtime::Store<StoreData>, units: u64| {
589 charge.call(&mut *store, &[Val::I64(units as i64)], &mut [])
590 };
591 let burn = |store: &mut wasmtime::Store<StoreData>, units: u64| {
593 let left = store.get_fuel().unwrap() - units;
594 store.set_fuel(left).unwrap();
595 };
596
597 store.set_fuel(100_000).unwrap();
600 charge_units(&mut store, 9_999).unwrap();
601 assert_eq!(store.get_fuel().unwrap(), 100_000);
602 assert_eq!(store.data().host_fuel, 9_999);
603 assert_eq!(store.data().host_fuel_pending, 9_999);
604 let usage = ExecutionUsage::capture(&store, 100_000).unwrap();
605 assert_eq!(
606 (usage.fuel, usage.wasm_fuel, usage.host_fuel),
607 (9_999, 0, 9_999)
608 );
609 charge_units(&mut store, 1).unwrap();
611 assert_eq!(store.get_fuel().unwrap(), 90_000);
612 assert_eq!(store.data().host_fuel_pending, 0);
613 charge_units(&mut store, 20_000).unwrap();
616 assert_eq!(store.get_fuel().unwrap(), 90_000);
617 assert_eq!(store.data().host_fuel_pending, 20_000);
618 burn(&mut store, 10_000);
619 charge_units(&mut store, 1).unwrap();
620 assert_eq!(store.get_fuel().unwrap(), 80_000);
621 burn(&mut store, 1);
622 charge_units(&mut store, 1).unwrap();
623 assert_eq!(store.get_fuel().unwrap(), 59_997);
624 assert_eq!(store.data().host_fuel_pending, 0);
625 let charged_so_far = store.data().host_fuel;
626
627 store.set_fuel(300).unwrap();
631 charge_units(&mut store, 100).unwrap();
632 let error = charge_units(&mut store, 201).unwrap_err();
633 assert_eq!(error.downcast_ref::<Trap>(), Some(&Trap::OutOfFuel));
634 assert_eq!(store.get_fuel().unwrap(), 0);
635 assert_eq!(store.data().host_fuel, charged_so_far + 100 + 200);
636 assert_eq!(store.data().host_fuel_pending, 0);
637
638 store.set_fuel(1_000).unwrap();
642 store.data_mut().host_fuel = 0;
643 charge_units(&mut store, 600).unwrap();
644 burn(&mut store, 700);
645 let error = charge_units(&mut store, 1).unwrap_err();
646 assert_eq!(error.downcast_ref::<Trap>(), Some(&Trap::OutOfFuel));
647 assert_eq!(store.data().host_fuel, 600);
648 let usage = ExecutionUsage::capture(&store, 1_000).unwrap();
649 assert_eq!(
650 (usage.fuel, usage.wasm_fuel, usage.host_fuel),
651 (1_000, 400, 600)
652 );
653
654 store.set_fuel(50).unwrap();
656 charge_units(&mut store, 0).unwrap();
657 assert_eq!(store.get_fuel().unwrap(), 50);
658 }
659}