Add incremental snapshotting support to wizer

[Gentle]

Rationale: If you keep the instrumentation information inside the output wasm file, then you can at any time use wizer to snapshot it again, allowing you to effectively save and resume later or on another machine

• adds a keep_instrumentation option that preserves __wizer_* exports in the output module
• adds a parse_instrumented method that re-parses a previously snapshotted module so it can be snapshotted again.
• parser::parse has been adjusted to either reject or require existing __wizer_* exports depending on new bool arg

[github-actions]

Subscribe to Label Action

cc @fitzgen

Details

This issue or pull request has been labeled: "wizer"

Thus the following users have been cc'd because of the following labels:

• fitzgen: wizer

To subscribe or unsubscribe from this label, edit the .github/subscribe-to-label.json configuration file.

Learn more.

[bjorn3]

parse_instrumented skips validating the wasm module, the logic behind this is that if the wasm file was successfully instrumented before, then it must be valid.

Is that safe to do when the wasm module comes from an untrusted source?

[Gentle]

well no, at that point you are trusting that whoever added the __wizer_* exports did not lie and did validate the wasm before adding them

it's possible that my microbenchmarks were bad, but with a 20MB wasm file, validate, parse, then snapshot took roughly 6 times as long as just parse and snapshot. It doesn't break anything to add self.wasm_validate() in parse_instrumented, I just think it's redundant but I'd be happy to implement it differently if there are concerns

[tschneidereit]

well no, at that point you are trusting that whoever added the __wizer_* exports did not lie and did validate the wasm before adding them

Validation is there in part for security reasons: the following steps can then assume they're handling valid wasm content. If the validation step is skipped, that's not a safe assumption anymore, so any code depending on it might do wrong things—potentially in ways that an attacker can exploit.

If you think about it, if we trusted the author of a module with these exports to create a valid module, why wouldn't we do that for all content in general?

with a 20MB wasm file, validate, parse, then snapshot took roughly 6 times as long as just parse and snapshot.

6x overhead for validation seems like a lot. Would you be able to share the module in question, by any chance?

[Gentle]

I'm sorry for the confusion, my test setup was indeed flawed, in a clean reproduction I can see that validate takes reasonable time even with large code sections. I added it back to parse_instrumented

(there was an issue in the allocator of my runtime, so I actually found a bug in the way I use wasmtime in my code, but that was entirely unrelated to wizer)

[alexcrichton]

Could you detail a bit more your use case here? I think it should work today to wizen a module/component twice, e.g. running one export on one machine and another export on another machine. The only downside I can think of to doing this is that the instrumentation takes a small amount of time to generate, but I would expect that to be negligible compared to compilation as a whole.

So, for more info, could you describe if performance is a primary concern here? Or if wizening a module twice is or isn't appropriate? Or if you're doing something else with the instrumented wizer artifact?

[Gentle]

I use this for basically durable execution and stateful modules

I have an RPC system and there are complex tasks that need processing. I always keep instrumentation info in the modules. When starting a task, the fresh module is instantiated, told to load the input task it should process, then it runs the task, possibly sending out RPC requests. The guest runs eagerly either until it finishes the task or until it only has pending requests that await responses.

In the case of nothing left to do but not finished yet, I take the current state in memory and the last snapshot I made and use these to snapshot the long-running instance, then store that file and shut down the worker. When the RPC responses arrive, I wake the worker up, resolve the response, snapshot the module again and save it until the next response arrives. This can loop until the task is finished even if some outgoing requests may take days or require human intervention

[Gentle]

from previous attempts, the instrumentation info has to be in the wasm file when it is instantiated or it will fail to snapshot, I tried naively just running the regular wasm file, then running the file through instrument; snapshot to learn that this won't work

so at a minimum I need --keep-instrumentation and parser::parse needs to not reject files that were already instrumented to make this feature work, but skipping instrument when possible would be optimal

Edit: to clarify, my use case is

• instantiate
• at some later point load wizer only if needed
• snapshot if not finished

in that order. This is because I am actually running the wasm as V8 WebAssembly.Instance and I use wizer compiled to wasm, so effectively I can dynamically load wizer only if required, but otherwise run the wasm file regularly. But for that to work the file has to always include instrumentation exports

[alexcrichton]

I think what I'm confused about is it sounds like your use case is satisfied today with the API of wasmtime-wizer, so I'm not sure why this PR is needed. For example:

• When you load a module, you call Wizer::instrument to get a context + wasm
• You execute the wasm provided with RPCs and such
• Eventually it's decided the wasm should be serialized, so Wizer::snapshot is used
• The result of Wizer::snapshot is saved
• Eventually the process is repeated when the module is reloaded.

Am I missing something though? Would this flow not work for your use case?

[Gentle]

for that approach I would need to somehow save ModuleContext to a file that can be loaded again on demand

what's currently possible:

• wizer always instruments and we need to keep ctx alive
• if the running wasm wants to suspend, ctx and the instance are given to Wizer::snapshot

what I want to do/am doing with the branch:

• instantiate the module
• if the module wants to suspend:
  • instantiate wizer.wasm
  • snapshot the module using wizer.wasm
  • throw away wizer instance, freeing memory

and I am using wizer compiled to wasm so that this can be run on cloudflare and browsers