Add `Store::call_hook` API #1144

emiltayl · 2024-07-29T16:28:10Z

My attempt at resolving #1083. I have tried to keep the public API similar to Wasmtime, and everything else as simple as possible.

This adds the function Store::call_hook that accepts a FnMut(&mut T, CallHook) -> Result<(), Error>, which is then called whenever a wasm function is called from the host, or a host function is called from wasm.

Let me know if you have any questions or you need me to explain or change something. Benchmarking on my own computer indicated a performance regresssion in the execute/br_table and execute/tiny_keccak benchmarks. If these regressions are unacceptable, a possible solution could be to add this behind a feature flag. This should be quite simple, so let me know if you would prefer that.

I did add a #[allow(clippy::type_complexity)] to please clippy to avoid having to create a type alias that is not used anywhere else, but it should be quite easy to fix that as well.

* Add the `Store::call_hook` function, which stores a call hook in `Store` * Create tests to verify call hook behavior

* Changed signature of call hooks to return a `wasmi::Error` instead of a `TrapCode`. * Use `assert_eq!` instead of `assert!` for tests, cosmetic change

* An `invoke_call_hook` function in Store to avoid problems with the borrow checker - we need a reference to the underlying data and the stored call hook * Invoke call hook on host -> wasm and wasm -> host calls

codecov · 2024-07-30T08:04:43Z

Codecov Report

Attention: Patch coverage is 91.22807% with 10 lines in your changes missing coverage. Please review.

Project coverage is 79.78%. Comparing base (7b6efbb) to head (574484a).

Files	Patch %	Lines
crates/wasmi/tests/e2e/v1/call_hook.rs	91.75%	8 Missing ⚠️
crates/wasmi/src/store.rs	80.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1144      +/-   ##
==========================================
+ Coverage   79.71%   79.78%   +0.07%     
==========================================
  Files         295      296       +1     
  Lines       25309    25422     +113     
==========================================
+ Hits        20176    20284     +108     
- Misses       5133     5138       +5

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Robbepop · 2024-07-30T08:05:34Z

@emiltayl Thanks a lot for this PR! Code looks pretty clean, too. :) Also great to have tests!

I will check it out locally and see how it performs and see if we can improve it. If we cannot fix the performance regressions we might have to put this feature behind a feature gate just like Wasmtime before merging.

Robbepop · 2024-07-30T09:22:00Z

I toyed a bit around and the fastest codegen I received was with this particular implementation:

/// Executes the callback set by [`Store::call_hook`] if any has been set.
///
/// # Note
///
/// - Returns the value returned by the call hook.
/// - Returns `Ok(())` if no call hook exists.
#[inline]
pub(crate) fn invoke_call_hook(&mut self, call_type: CallHook) -> Result<(), Error> {
    match self.call_hook.as_mut() {
        None => Ok(()),
        Some(call_hook) => Self::invoke_call_hook_impl(&mut self.data, call_type, call_hook),
    }
}

/// Invokes the [`Store::call_hook`] that is asserted to be available in this case.
#[cold]
fn invoke_call_hook_impl(
    data: &mut T,
    call_type: CallHook,
    call_hook: &mut CallHookWrapper<T>,
) -> Result<(), Error> {
    call_hook.0(data, call_type)
}

Robbepop

All in all LGTM.
Clippy warning needs to be fixed and some formatting itches resolves. :)
Also it would be great to base this on current main branch before we merge and run a perf comparison with my proposed changes to the invoke_call_hook implementation to see if it is beneficial on your system, too.

So far I couldn't see major performance regressions so I conclude that for now no crate feature is needed to guard this feature in Wasmi.

Robbepop · 2024-07-30T09:25:30Z

crates/wasmi/src/store.rs

+#[derive(Debug)]
+/// Argument to the callback set by [`Store::call_hook`] to indicate why the
+/// callback was invoked.
+pub enum CallHook {


Not important but it itches me a bit to see #[derive(Debug)] above doc comments. 🙃 Rest of the codebase has doc comments first.

Robbepop · 2024-07-30T09:26:26Z

crates/wasmi/src/store.rs

+    /// Executes the callback set by [`Store::call_hook`] if any has been set
+    /// and returns the value returned by the callback. If no callback has been
+    /// set, `Ok(())` is returned.
+    pub(crate) fn invoke_call_hook(&mut self, call_type: CallHook) -> Result<(), Error> {
+        match &mut self.call_hook {
+            None => Ok(()),
+            Some(call_hook) => call_hook.0(&mut self.data, call_type),
+        }
+    }


Please see my non-review comment for improvements (optimizations) here. Also could you please see how the suggested code performs on your machine? I only have an M2 macbook to run benchmarks on.

Robbepop · 2024-07-30T09:26:38Z

crates/wasmi/tests/e2e/v1/call_hook.rs

+#[derive(Default)]
+/// Number of times different callback events have fired.
+struct TimesCallbacksFired {


emiltayl · 2024-07-30T10:23:12Z

Thanks for the review!

I did not realize that clippy didn't check the tests when I executed it locally. I will fix the clippy errors and reposition the attributes appearing above doc comments. I will also try out your suggested changes and be back with benchmark results.

I do not have the time today, but should get this done before the weekend.

* Placed attributes after doc comments * Added hint to allow complex type for `generate_error_after_n_calls`

…tests

* Inline `Store::invoke_call_hook` to avoid extra function call * Add a `Store::invoke_call_hook_impl` function that is `#[cold]` to hint that the compiler should optimize for the scenario that there are no call hooks

emiltayl · 2024-07-31T18:00:35Z

I tried your suggested changes. My benchmark numbers were quite similar prior to the changes, I still had a slight regression in execute/tiny_keccak. I think I should disable SMT on my 3700x to get reliable benchmark numbers though, I get a regression on execute/tiny_keccak even when comparing main to main most of the time.

I'm not too happy with the fix I did in the tests to work around WasmTy not being implemented for ! (and could therefore not be "returned" from functions), but it does work and I do not think it is too bad.

Robbepop

All in all LGTM! Just one super minor nit.
Thanks for the tons of tests.

I'm not too happy with the fix I did in the tests to work around WasmTy not being implemented for ! (and could therefore not be "returned" from functions), but it does work and I do not think it is too bad.

Can you please point me to the code that you mean? I seem to fail to find it.

I'd be fine with merging your PR.

crates/wasmi/src/store.rs

Co-authored-by: Robin Freyler <[email protected]>

emiltayl · 2024-08-01T06:14:29Z

Great! Thank you for the help and reviews.

As for the changes I did that I am not that satisfied with myself (I still think the code is good to go, though), it is related to https://github.com/wasmi-labs/wasmi/actions/runs/10148134369/job/28090510038#step:8:84. should_not_run in the tests was implemented like this:

let should_not_run = Func::wrap(&mut store, |_: Caller<TimesCallbacksFired>| {
    panic!("Host function that should not run due to trap from call hook executed");
});

From what I understand (which may be flawed/incomplete), with the upcoming changes in Rust 2024, this would fail to compile. This is because the closure's return type would be ! instead of (), and WasmRet / WasmTy is not implemented for !, which is required for the return types of the closures passed to Func::wrap. I do not think it neccessarily makes sense to implement either WasmRet or WasmTy for ! though, I imagine it would only be useful for tests like these.

The new should_not_run can be seen at https://github.com/wasmi-labs/wasmi/pull/1144/files#diff-93c6a63dd2cba7296e30181521de10f97f315df932f3beb82c90d8879dac56deR175. This requires an additional assert! to check whether should_not_run has been executed or not, but I think it is a decent solution even though I think the old should_not_run is a tad more neat.

Robbepop · 2024-08-01T07:50:42Z

@emiltayl Couldn't the function return an error instead of panicking to provide the same functionality of not continuing the execution?

Anyway I am merging this now and any improvements or fixes can be done as follow-ups. :)
Thanks again for your work to implement this feature, @emiltayl ! 🚀

emiltayl added 4 commits July 12, 2024 22:35

Start implementing Store::call_hook (#1083)

aab855f

* Add the `Store::call_hook` function, which stores a call hook in `Store` * Create tests to verify call hook behavior

Let call hooks return any wasmi::Error

8f9851f

* Changed signature of call hooks to return a `wasmi::Error` instead of a `TrapCode`. * Use `assert_eq!` instead of `assert!` for tests, cosmetic change

Invoke call hooks on function calls

8652dd7

* An `invoke_call_hook` function in Store to avoid problems with the borrow checker - we need a reference to the underlying data and the stored call hook * Invoke call hook on host -> wasm and wasm -> host calls

Clarified documentation on error propagation for call hooks

3156199

Robbepop reviewed Jul 30, 2024

View reviewed changes

emiltayl and others added 4 commits July 31, 2024 18:16

Merge branch 'main' into call_hook

da53408

Fixed attribute placement and clippy errors

639e121

* Placed attributes after doc comments * Added hint to allow complex type for `generate_error_after_n_calls`

Fixed clippy warnings on implicit conversion of ! to () in call_hook …

7819bd3

…tests

Robbepop approved these changes Jul 31, 2024

View reviewed changes

crates/wasmi/src/store.rs Outdated Show resolved Hide resolved

Update crates/wasmi/src/store.rs

574484a

Co-authored-by: Robin Freyler <[email protected]>

Robbepop merged commit 0a6a083 into wasmi-labs:main Aug 1, 2024
18 checks passed

Robbepop mentioned this pull request Aug 1, 2024

Add Store::call_hook API #1083

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `Store::call_hook` API #1144

Add `Store::call_hook` API #1144

emiltayl commented Jul 29, 2024

codecov bot commented Jul 30, 2024 •

edited

Loading

Robbepop commented Jul 30, 2024

Robbepop commented Jul 30, 2024

Robbepop left a comment

Robbepop Jul 30, 2024

Robbepop Jul 30, 2024

Robbepop Jul 30, 2024

emiltayl commented Jul 30, 2024

emiltayl commented Jul 31, 2024 •

edited

Loading

Robbepop left a comment

emiltayl commented Aug 1, 2024

Robbepop commented Aug 1, 2024 •

edited

Loading

Add Store::call_hook API #1144

Add Store::call_hook API #1144

Conversation

emiltayl commented Jul 29, 2024

codecov bot commented Jul 30, 2024 • edited Loading

Codecov Report

Robbepop commented Jul 30, 2024

Robbepop commented Jul 30, 2024

Robbepop left a comment

Choose a reason for hiding this comment

Robbepop Jul 30, 2024

Choose a reason for hiding this comment

Robbepop Jul 30, 2024

Choose a reason for hiding this comment

Robbepop Jul 30, 2024

Choose a reason for hiding this comment

emiltayl commented Jul 30, 2024

emiltayl commented Jul 31, 2024 • edited Loading

Robbepop left a comment

Choose a reason for hiding this comment

emiltayl commented Aug 1, 2024

Robbepop commented Aug 1, 2024 • edited Loading

Add `Store::call_hook` API #1144

Add `Store::call_hook` API #1144

codecov bot commented Jul 30, 2024 •

edited

Loading

emiltayl commented Jul 31, 2024 •

edited

Loading

Robbepop commented Aug 1, 2024 •

edited

Loading