Support embedder-provided CoW memory initialization images

[benjamincburns]

Feature

Make it possible to initialize an instance's linear memory from an embedder-provided, file-backed CoW image, independently of the compiled module.

Benefit

I'm running largish (hundreds of MiB) pre-initialized modules using Wasmtime's copy-on-write memory initialization. My setup also uses lifecycle hooks, similar to AWS Lambda SnapStart. I call an init hook, let it do application-specific setup, then capture the resulting state.

This saves initialization work, but increases the size of the pre-initialized module, increasing the time that cold requests spend faulting those pages into memory.

I'd like to compile the module once, run the init hook while observing which code and linear-memory pages become resident, and capture the resulting linear memory separately. New instances would use that captured memory as their CoW initialization image, with the compiled code unchanged.

That would let me prefetch the observed working set without recompiling and changing the layout I'm measuring. It also lets me pack the selected pages together on disk while preserving their guest addresses, for added prefetch efficiency.

I can't find a public API for supplying that replacement memory image, which is why I'm requesting this feature.

Implementation

I have a Linux-specific proof-of-concept that installs an immutable file-backed image into an existing memory's CoW slot. It supports page-aligned extents mapping file offsets to linear-memory offsets, including images whose physical page order differs from their logical order.

Obviously my proof-of-concept approach is less efficient than one that supplies the image on instantiation, that way we'd avoid redundant memory initialization work. All the same, here's a snippet showing use of the proof-of-concept. CowMemoryImage and Memory::restore_cow_image are my additions.

// snapshot_file contains the captured linear memory.
// Its contents must remain unchanged while any image or instance uses it.
let image = unsafe {
    CowMemoryImage::from_file(snapshot_file, 0, captured_length_bytes)?
};

// `memory` belongs to a newly instantiated, idle `store`.
// Grow it to the captured size before installing the image.
let current_pages = memory.size(&store);
if captured_pages > current_pages {
    memory.grow(&mut store, captured_pages - current_pages)?;
}

if let Err(error) = unsafe {
    memory.restore_cow_image(&mut store, &image)
} {
    // The prototype does not promise an unchanged memory on failure.
    drop(store);
    return Err(error);
}

// Restore captured globals and other required instance state
// before calling into the guest.

I'm open to whatever alternative API you guys might prefer.

I'm only asking for the linear-memory mechanism here. I haven't tested it yet, but I think I can accomplish a similar packing/prefetching scheme for code pages without modification to Wasmtime via Module::deserialize_raw and CustomCodeMemory. If that's not the case, I'll likely open a new ticket for a strategy to support a similar prefetch/packing scheme for code pages.

Alternatives

I couldn't think of any strong alternatives that don't remove some benefit offered by the above.

Baking the captured memory into data segments and recompiling could work, provided I can map the linear-memory offsets into the final artifact. However, packing selected pages together would then require custom loading rather than using Wasmtime's normal file-loading path. Further, unless recompilation preserves the compiled code's layout, the code-page offsets collected during init would need a separate translation before I could use them for prefetching and packing.

Copying the captured bytes into each new instance loses the sharing and demand-paging benefits of file-backed CoW initialization.

I guess theoretically I could do something via MemoryCreator, but I'd have to bring my own CoW machinery, and it wouldn't directly support the pooling allocator we use.

[benjamincburns]

Ah I see in #12517 (comment) that the approach of supplying memory on init is likely not going to be acceptable due to the need to preserve wasm memory init semantics.

And reading the comments further, I think the strategy in this comment could work for us, except we'd lose the further optimization of packing pre-fetched pages into a contiguous region of the file.

I'll have to go searching to see if there's some mechanism we can use to control the layout of backing files on disk, and measure the effect of prefetching a contiguous region of the file again.

I'll do some experiments and close this if the strategy works.

[alexcrichton]

Reading over this I was about to basically say the same thing as the comment you already linked, and yeah if that works I think that would work best. Basically modeling situations as wasm primitives is generally far easier to support and reason about where possible.

If this problem is reduced to a storage issue of where the pages are source from though that's where I think we could consider adding something. So long as wasm semantics are preserved we've got a lot of leeway of how exactly they're implemented, for example if you wanted to provide a custom CoW image we'd internally verify the image matches the wasm module and could then tweak things internally. Basically storage options are flexible and optimizations as well so long as it's always framed within the context of preserving preexisting wasm semantics.

[benjamincburns]

@alexcrichton I think for the moment I won't have the time to run those experiments, but on further reflection I don't think it will perform as well as the scheme I'm already running. In particular, I'm not sure whether I'd be able to combine this with the learned prefetch optimisation as described above, and even if I do figure that out, I worry that the extra step on activation would offset the performance improvement that I'd see from prefetching.

If you guys are open to allowing captured CoW images on init, by allowing for embedders to define where the pages are sourced from, I think that would be much easier to consume and implement vs a custom memory importing scheme, and I suspect that it will perform better (at least in my case), so long as the validation steps you mentioned weren't particularly onerous.

We're doing some structural checking when loading the memory image (dimensions, alignment, bounds, coverage), if that's what you had in mind? We also rely on a cryptographic signature and fs-verity in prod to ensure the authenticity and integrity of our serialized guest template bundle at time of registration and time of use, but I'd imagine that's likely well beyond what you were thinking.

All of that is to say, if you were thinking that wasmtime would own the serialization format for these images, I'd just ask that it be extensible and/or support embedder-defined extent mappings so that we can pack prefetched pages in to a contiguous region of the file.

[alexcrichton]

How exactly this would get added to Wasmtime I'm not entirely sure, what I was vaguely envisioning is that given a component or module during deserialization/compilation you'd provide some sort of "ok but use the data segments from this fd" or something like that, probably specified as some sort of set of fds and offsets. Given that Wasmtime would verify the contents of those fds against the actual data segments expected and then go from there. Or... something like that, it's all pretty vague and nebulous and I'm not sure if it'd work out.

Overall though this feature is definitely something that we're going to want to have in one form or another at some point. Historically the use case I've heard of for this is an interpreted language (e.g. JS) which has an initial interpreter/memory snapshot. Wizening would then load an application, and then the final heap snapshot is in theory a smaller diff from the original snapshot than the entire snapshot. Beyond historical discussions of "we probably want layers" there hasn't been much discussion of how exactly to implement such a feature and/or model it. For example Wasmtime supports separately compiled code and memory insofar as components can import modules, but beyond that the support is pretty limited. Wasmtime won't do any sort of layering/sharing data sections across components.

In the end though I personally still think this is something we're going to want to model as wasms overall. That'll require engine changes to fit the performance profile regardless, but without changing the wasm toolchain/pipeline I think there's limited options available to Wasmtime in terms of optimizations.