Out-of-band fuel metering #4466
Conversation
| panic!("unsupported platform") | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
The platform-specific code here and elsewhere in this PR relies on quite some C&P.
I wondered if it would be a good idea to reshuffle the low-level code into a separate sys module. It will then contain the platform-specific parts for traphandlers, out-of-band fuel and whatnot.
Besides, traphandlers and it's CallThreadState has grown out IMO from being only about traphandlers.
| // | ||
| // Writing to COOKIE is safe since no data race is possible because `platform_init` is | ||
| // executed only once. | ||
| COOKIE = psm::stack_pointer() as usize; |
There was a problem hiding this comment.
Alternatively, we could use a magic constant.
github-actions
bot
added
wasmtime:api
Subscribe to Label Actioncc @peterhuene
This issue or pull request has been labeled: "wasmtime:api", "wasmtime:config"
Thus the following users have been cc'd because of the following labels:
To subscribe or unsubscribe from this label, edit the |
Label Messager: wasmtime:configIt looks like you are changing Wasmtime's configuration options. Make sure to
To modify this label's message, edit the To add new label messages or remove existing label messages, edit the |
There was a problem hiding this comment.
Before giving too much more of a detailed review I think that there's still a few major pieces which should come together:
- Can you run the
callbenchmarks (or at least a subset of them) to test the performance impact of this? Even with this feature disabled I fear that the new branches and code added may have a nontrivial impact. - I believe libcalls made right now don't have any affordance for storing/reloading fuel
- I don't think
TypedFuncis using trampolines to enter wasm to store/reload fuel. - I don't think
Func::wrapis using trampolines to exit wasm to store/reload fuel.
The signalling bits can probably work out one way or another so I'm not personally too worried about the specifics. Otherwise though I'm not sure how best to do this but I think it would be good to confirm with others the trajectory of this implementation before going too much further.
| /// Saves the current fuel counter saved in a pinned register into `VMRuntimeLimits` in memory. | ||
| fn spill_fuel( | ||
| &self, | ||
| offsets: &wasmtime_environ::VMOffsets<u8>, | ||
| builder: &mut FunctionBuilder, | ||
| vmctx: ir::Value, | ||
| ) { | ||
| let isa = &*self.isa; | ||
| let pointer_type = isa.pointer_type(); | ||
| let fuel_consumed = builder.ins().get_pinned_reg(ir::types::I64); | ||
|
|
||
| let vmctx_runtime_limits_offset = i32::try_from(offsets.vmctx_runtime_limits()).unwrap(); | ||
| let vmruntime_limits_fuel_consumed_offset = | ||
| i32::try_from(offsets.vmruntime_limits_fuel_consumed()).unwrap(); | ||
| let runtime_limits_ptr = builder.ins().load( | ||
| pointer_type, | ||
| ir::MemFlags::trusted(), | ||
| vmctx, | ||
| vmctx_runtime_limits_offset, | ||
| ); | ||
| builder.ins().store( | ||
| ir::MemFlags::trusted(), | ||
| fuel_consumed, | ||
| runtime_limits_ptr, | ||
| vmruntime_limits_fuel_consumed_offset, | ||
| ); | ||
| } |
There was a problem hiding this comment.
Should this also check fuel during spilling? Otherwise I think nondeterminism might creep in where host functions could execute when the wasm module has no fuel left?
| @@ -507,14 +524,23 @@ impl<'module_environment> FuncEnvironment<'module_environment> { | |||
| let fuel = builder | |||
| .ins() | |||
| .load(ir::types::I64, ir::MemFlags::trusted(), addr, offset); | |||
| builder.def_var(self.fuel_var, fuel); | |||
|
|
|||
| if self.tunables.outband_fuel { | |||
There was a problem hiding this comment.
Is this function ever executed with outband_fuel? This seems like a confusing implementation of this function since the purpose of the pinned register is that it's always valid and doesn't need reloading.
| } | ||
|
|
||
| /// Stores the fuel consumption value from `self.fuel_var` into | ||
| /// `VMRuntimeLimits`. | ||
| fn fuel_save_from_var(&mut self, builder: &mut FunctionBuilder<'_>) { | ||
| let (addr, offset) = self.fuel_addr_offset(builder); | ||
| let fuel_consumed = builder.use_var(self.fuel_var); | ||
| let fuel_consumed = if self.tunables.outband_fuel { | ||
| builder.ins().get_pinned_reg(ir::types::I64) |
There was a problem hiding this comment.
Similar to fuel_load_into_var I think this function may never be executed if outband_fuel is true?
| @@ -546,7 +573,11 @@ impl<'module_environment> FuncEnvironment<'module_environment> { | |||
| // Compare to see if our fuel is positive, and if so we ran out of gas. | |||
| // Otherwise we can continue on like usual. | |||
| let zero = builder.ins().iconst(ir::types::I64, 0); | |||
| let fuel = builder.use_var(self.fuel_var); | |||
| let fuel = if self.tunables.outband_fuel { | |||
There was a problem hiding this comment.
Similar to above I think that fuel_check is never executed with outband_fuel?
| mod linux; | ||
| use linux as sys; | ||
| } else { | ||
| // TODO: I need to make it compile on other platforms. |
There was a problem hiding this comment.
Mind leaving a compile_error! for now?
|
|
||
| /// This function is required to be called before a call to wasm code that supports out-of-band | ||
| /// fuel metering. | ||
| pub fn init_outband_fuel(is_wasm_pc: fn(usize) -> bool) { |
There was a problem hiding this comment.
Can this be folded into the preexisting trap initialization?
There was a problem hiding this comment.
Oh I see that this is a slightly different query, but could the query be passed through a function parameter? Something like must_be_trap: bool but probably with an enum of some kind to be more descriptive.
|
Thank you for taking a look! Yes, @alexcrichton, this PR is very raw right, way much less that I prefer right now. Sorry about that. I was eager to create it because of two reasons:
Specifically, to give the context for this question #4431 (comment). I was scared that it made this PR impossible or very hard to implement. Right now, I am rebasing this PR onto Nick's and trying to figure out what to do with the trampolines. The best I came up with so far is under the spoiler below. I hope my programming license won't be revoked. struct VMRuntimeLimits {
// ...
/// A flag that indicates if the out-of-band fuel metering is enabled.
pub outband_fuel: UnsafeCell<bool>,
/// host-to-wasm return address.
pub orig_entry_trampoline_return_address: UnsafeCell<usize>,
/// wasm-to-host return address.
pub orig_exit_trampoline_return_address: UnsafeCell<usize>,
}
// ...
asm_func!(
"host_to_wasm_trampoline",
"
.cfi_startproc simple
.cfi_def_cfa_offset 0
// Load the pointer to `VMRuntimeLimits` in `r10`.
mov r10, 8[rsi]
// Check to see if this is a core `VMContext` (MAGIC == 'core').
cmp DWORD PTR [rdi], 0x65726f63
// Store the last Wasm SP into the `last_wasm_entry_sp` in the limits, if this
// was core Wasm, otherwise store an invalid sentinal value.
mov r11, -1
cmove r11, rsp
mov 40[r10], r11
// Check if the out-of-band fuel metering is enabled. If it is not enabled, skip to the
// tail call directly.
cmp QWORD PTR 48[r10], 0x00
je host_to_wasm_trampoline__perform_tail_call
// OK, here we know that the out-of-band fuel metering is enabled.
//
// Dump the `fuel_consumed` into the fuel register.
mov r15, 8[r10]
// Next, we have to do some shenanigans. Below we call into the wasm with a tail call, to
// not mess with the stack layout. But at the same time, after returning from wasm, we
// have to our fuel register r15 back into `fuel_consumed` of `VMRuntimeLimits`. To do that
// we stash the original return address (which is currently located at `[rsp]`) that points
// back to the caller from the host side and change it with the continuation defined
// below. The continuation will dump the r15 back into the memory and unstash the original
// return address and jump at it.
mov r11, [rsp]
mov 56[r10], r11
lea r11, [rip+host_to_wasm_trampoline__cont]
mov [rsp], r11
host_to_wasm_trampoline__perform_tail_call:
// Tail call to the callee function pointer in the vmctx.
jmp 16[rsi]
host_to_wasm_trampoline__cont:
// Here we're in a inconvenient situation, where the callee's frame is popped and all the
// arguments are gone. We have to improvise our way out of this situation.
// Open a temporary stack frame.
//
// The frame layout:
//
// rsp + 0: The original return address. Written by `wasmtime_outband_fuel_on_entry`.
// rsp + 8: r15
// rsp + 16: rax
// rsp + 24: padding
// rsp + 32: xmm0
sub rsp, 48
mov 8[rsp], r15
mov 16[rsp], rax
movdqu XMMWORD PTR 32[rsp], xmm0
// Finally, move the current rsp into the first argument.
mov rdi, rsp
call wasmtime_outband_fuel_on_entry
// Load the original return address into r10
mov r10, 0[rsp]
// Restore the initial values of rax and xmm0.
mov rax, 16[rsp]
movdqu xmm0, XMMWORD PTR 32[rsp]
// Release the allocated stack area.
add rsp, 48
// Jump at the original return address.
jmp r10
.cfi_endproc
",
);
/// Dumps fuel into the `VMRuntimeLimits->fuel_consumed` and unstashes the original return address
/// for the trampoline.
///
/// Only used if out-of-band fuel is used.
#[no_mangle]
pub unsafe extern "C" fn wasmtime_outband_fuel_on_entry(sp: *mut usize) {
use crate::traphandlers::tls;
tls::with(|info| {
let info = match info {
Some(info) => info,
None => {
// Entering trampoline must set the TLS.
std::hint::unreachable_unchecked();
}
};
#[repr(C)]
struct StackArgs {
return_address: usize,
fuel: i64,
}
let args = sp as *mut StackArgs;
(*args).return_address = *(*info.runtime_limits())
.orig_entry_trampoline_return_address
.get();
*(*info.runtime_limits()).fuel_consumed.get() = (*args).fuel;
if (*args).fuel > 0 {
// Fuel ran out and we bail through unwinding.
crate::traphandlers::raise_lib_trap(wasmtime_environ::TrapCode::Interrupt);
}
});
}This seems to be working and I hope I am not missing anything. I am in the process of replicating the same construction for the wasm-to-host and libcall trampolines. You can see there it should address the concerns you expressed wrt not checking the fuel. To be clear, the trampoline generation in I ack the other notes. I hope the next push will feature not as raw code as this one. Also will run the benchmarks on it. |
|
Oh no worries! I mostly want to make sure I'm clear on expectations and such, it's totally fine to push up work-in-progress PRs. Also wow that is some gnarly assembly. I... can see how it might work though! I suspect one day we're going to need to severly improve our trampoline story, even with "just maintain frame pointers" it was already complicated enough... In any case though not a blocker for this work, just something for us to ponder. |
|
One thing I just thought of I want to note before I forget about, right now "is this a wasm pc" is pretty crude where it just checks if the instruction pointer is in the code section of a loaded module (this is preexisting in Wasmtime) but I think that we'll want to enhance that because if an interrupt happens during an entry trampoline we can't guarantee that the fuel register is configured yet, so we'll want to make sure that the "is this a wasm pc" check may exclude trampolines. |
|
Oh, yes, definitely. On top of that, the query should be a bit more smarter. It should only return true if I've been battling with the trampolines for quite some time now, and it does not seem good. The problem is I am not sure how to make the system unwinder work through this trampoline. So at first, I was getting the SIGABRT from the unwinder since it couldn't find the CFA. I fixed that by replacing I fixed that by slapping We need to annotate where we get the original return address to fix that. The return address is stored behind TLS. There is no way, we can come up with a DWARF expression that fetches it from the |
|
Ah, regarding the fuel register being initialized, I think it should already work. Here we ask for the |
|
I'm pretty sure our dwarf unwind information is not "async correct" where every single ip in a function guarantees a correct backtrace. With the fast-stack-walking branch and fp-based unwinds that may get better though. |
Oh, sure, I remember that. It should be fine though since I was not planning on capturing the backtraces in the fuel interrupts, and I still keep in mind the landing pads idea. The problem I am dealing with now is not about async traps. I should've been clearer on this. That SIGABRT would happen during capturing the stack trace iff the trampoline on the stack, and it went through the return address swizzling of the continuation. I came up with a standalone minimized version of the problem which can be found under the spoiler: // extern crate backtrace;
use std::arch;
use std::mem;
arch::global_asm!(
"
.p2align 4
.global trampoline
.type trampoline, @function
trampoline:
.cfi_startproc simple
.cfi_def_cfa rsp, 0
// move RDI, the original function, into r11
mov r11, rdi
// the return address pushed by the `call` instruction calling this trampoline
// is located at *rsp. Move the value to rdi.
mov rdi, [rsp]
// Then, replace the return value with the continuation at the label `cont`.
lea r10, cont[rip]
mov [rsp], r10
//
// Important!!
//
// Mark the register 16, aka the return address, as undefined. That stops the unwinder here.
// If that was not here, the unwinder would fall into an infinite loop.
//
// An unfortunate side-effect is the stack traces are also chopped here.
.cfi_undefined 16
// Finally, perform the jump.
jmp r11
cont:
jmp rax
.cfi_endproc
.size trampoline, .-trampoline
"
);
extern "C" {
fn trampoline();
}
type Func = extern "C" fn(*const (), usize) -> *const ();
#[inline]
unsafe fn prepare_trampoline(dest: Func) -> (*const (), Func) {
let dest = dest as *const Func as *const ();
let func = mem::transmute_copy(&(trampoline as usize));
(dest, func)
}
fn main() {
unsafe {
let (dest, f) = prepare_trampoline(host_func);
f(dest, 48);
}
}
extern "C" fn host_func(rdi: *const (), rsi: usize) -> *const () {
println!("{}", rsi);
let mut count = 0;
backtrace::trace(|frame| {
let ip = frame.ip();
println!("{:X}", ip as usize);
if count == 100 {
// likely we stuck in a loop. Break into the debugger.
unsafe {
arch::asm!("ud2");
}
}
count += 1;
true
});
// we don't have such luxury in the real thing.
rdi
}If you remove the |
|
Oh right yeah with the trampoline you gisted above I was wondering how backtrace info would work out and it seems like it doesn't. I don't know how to solve that myself, the trampoline there is quite gnarly. |
If wasmtime-flavoured ABIs enshrined the caller vmctx and guaranteed it so that it is still usable after the callee returns, it would've made things simpler: no need for reaching to tls for the original return address. The dwarf expression would require a few dereferences but otherwise seems to be doable. If you (or anyone from BA) think that's acceptable, I can try to research this implementation strategy. Otherwise, I think I ran out of ideas how to make this work. |
|
Yeah unfortunately I don't have a great answer myself on that. I feel like with this and the stacks PR our trampoline story is pretty lacking, but I don't know how best to rectify it. Short of that it may be that a pinned register won't work at all until we figure out a better story for trampolines. Ideally all the pinned register/fuel stuff would be part of cranelift-compiled trampolines and we wouldn't have to tinker with |
This is a prototype of the solution for #4109 1. This is very rough and is not intended to be landed as is. Rather, this PR here is to validate that this approach is sensible and to source some preliminary feedback.
Introduction
The regular fuel metering holds the fuel in the
vmctx->rumtime_limits->fuel_consumed. At the beginning of each function, the value is loaded into a local variable. Roughly every basic block the value is increased with the cost of that basic block. The value is checked for overflow at function entries and loop headers. If fuel overflowed, then a certain libcall handles it. Before leaving a function (normally, through a call or before a trap), the fuel is dumped into the VMRuntimeLimits.WIth the out-of-band fuel metering, the fuel is now promoted to a dedicated register tapping to the
pinned_regcranelift feature. The value is still increased every basic block, but the value does not leave the pinned register within wasm. Only at the wasm-host boundaries, i.e. trampolines & libcalls (not implemented as of this PR, coming later), is the fuel value loaded in or flushed from the pinned register into theVMRuntimeLimits.Also, no checks are performed in the wasm. The checks are meant to be performed either when crossing the wasm-host boundary or asynchronously. Specifically, on Linux, the check is performed by sending a signal each, e.g., 1ms. The signal handler checks if the signal came from the wasmtime (on a best effort basis) and if the program counter points at some the JIT code. If it does, then that means the pinned register holds the currently consumed fuel value. If the fuel value is overflown, we bail out unwinding the wasm stack.
This kind of mechanism showed a great improvement in performance on our tests while still being deterministic as long as the in-wasm state is irrelevant in the case of the OOG.
Now, the prototype here right now targets x86_64 Linux. There is a plan to support aarch64 and macOS. Windows should also be possible to implement. The prototype does not support async. It would be great to support it, but additional work is required.
Implementation Notes and Rationale
Mutex
Right now, before entering we save the tid of the calling thread. This is because theoretically the store can be called from different threads. I also wanted to prepare for the async: potentially the future can be polled on any thread, with each fiber switch we can find ourselves on a new thread.
The problem is with Linux, it turns out that sending signals is a bit of a hassle. The signal's sender cannot know if the destination thread is dead or alive. Moreover, the tid can theoretically be reused and thus a signal could be sent to the wrong thread.
At first, I thought it might be a problem performance-wise, but now I don't think so. The reason is: that the mutex does not get too contended. The mutex is taken on wasm entry & exit and also during the out-of-band fuel check. The latter also uses
try_lock. The interesting case is when the wasm tries to exit to host but the mutex is held: in that case, the exit will be delayed until the out-of-band check request is finished.rt_tgsigqueueinfoI resorted to using a raw syscall
rt_tgsigqueueinfoon Linux to send the signal.I thought about using
pthread_sigqueue(in constrast to justpthread_kill) because it allows to send asival. This is helpful to tell if the signal is coming from wasmtime or not. However, turns out that at least glibc does a bunch of syscalls that we probably don't want to have inside of the out-of-band fuel check request. So I decided to go straight forrt_tgsigqueueinfo. It takes thesiginfo_tbut it seems like the kernel does not use that and passes it as is, so I used this opportunity pass dummy values.Another potential problem that I am not sure needs to be tackled: the
pidis cached during the creation of the out-of-band check handle. This is not entirely correct since theoretically, it can change, but I figured it does not warrant worrying.Future Work
If this gets a green light, then several things will need to be done in the future:
As I mentioned above it should work on other platforms, namely aarch64, macOS, and possibly Windows.
Then, make it work with async. That is actually a bit tricky. The main problem revolves around handling the yields. Specifically with Linux, if a signal handler interrupted the wasm code and figured it was OOG, it should yield the execution. Not sure how good is the idea to switch fibers from inside a signal handler. With macOS/Windows it's not any better: the check thread should manipulate the target thread so that it's possible to switch the fiber.
In case, that works, we can think of applying the same technique we use here for a high-performance substitution for the epoch interrupts.
Big thanks to @alexcrichton & @cfallin for their great support !
Footnotes
I decided to change the name from slacked metering. First, I wanted to avoid mentioning async, so that's why it's not async fuel metering, since that async is orthogonal to the async currently used in wasmtime (as in
Func::call_async). At the same time I don't know if slacked conveys the meaning (English is not my mother tongue). So I figured that that "out-of-band metering" is a better name. I contract it tooutbandin code, I assume it's fine since I saw people using it elsewhere. Please let me know if you have a better name! ↩