Support concurrent plugin execution through instances

[clarkmcc]

See Discord thread.

It seems like there should be a Plugin type that is thread-safe, and a PluginInstance type that is not thread-safe. The Plugin contains the runtime and the compiled module, and you'd have something like func(p *Plugin) NewInstance() *PluginInstance which would instantiate a wazero module that would then be used in a single-threaded context.

The idea here is I can compile my plugin one time, and call functions from the module as many times as I want, concurrently, without having to re-compile.

[mhmd-azeez]

I like the idea, but I think it'd better if:

1. We keep the plugin struct as is (with the same NewPlugin function we have now; i.e. keep API back-compat)
2. As @bhelx has mentioned, maybe we can create a new PluginModule or Module or PrecompiledPlugin struct that acts as a intermediate state where you can intialize Plugins from quickly. And NewPlugin can use PrecompiledPlugin under the covers
3. Do we have any benchmarks showing the performance gains of having this intermediate step?
4. We will need to sure all of the tests still pass, maybe we need some more tests to ensure the thread safety of the intermediate struct

[evacchi]

Another concern for me is that (while this might make sense in a general sense) this changes the general model of interaction with the Extism API, so this change might need to be replicated on all other SDKs for consistency 🤔

[clarkmcc]

@mhmd-azeez

We keep the plugin struct as is (with the same NewPlugin function we have now; i.e. keep API back-compat)

Some of this design is mutually exclusive with these changes. Under this design, you can't get the benefits of a shared plugin with multiple plugin instances while also preserving the API. What I would propose for preserving the API is moving basically all of these changes to an extismcore package, and having the existing API today just call into the extismcore package. It preserves the API, while still giving power users the ability to work with plugins and instances.

As @bhelx has mentioned, maybe we can create a new PluginModule or Module or PrecompiledPlugin struct that acts as a intermediate state where you can intialize Plugins from quickly. And NewPlugin can use PrecompiledPlugin under the covers

You're suggesting that the instance be called Plugin and the "factory" for lack of a better word would me the module?

Do we have any benchmarks showing the performance gains of having this intermediate step?

No, mainly because in my opinion the performance gains are obvious in the sense that we're not talking about a small loop optimization or something that may have negligible impact, we're talking about eliminating the need for multiple compilations when you need multiple concurrent instances. For example, if compiling my module takes 80MB of memory (which it actually does, this is a real example) and I need to run 10 instances concurrently, then prior to the changes in this PR, that would be 800MB of allocated memory just for compilation. The vast majority of that goes away when I can compile once and instantiate 10 times instead. We could definitely write some benchmarks to quantify exactly what the benefits are and the scenarios in which you would see those benefits if you feel that would be helpful, but if the question is "does this approach improve performance", I would argue -- definitely yes.

We will need to sure all of the tests still pass, maybe we need some more tests to ensure the thread safety of the intermediate struct

Agreed. I haven't even audited my changes to make sure they're conceptually thread-safe so once we're comfortable with an API, we'll want to do a thorough pass through this in addition to tests to make sure everything checks out.

@evacchi

Another concern for me is that (while this might make sense in a general sense) this changes the general model of interaction with the Extism API, so this change might need to be replicated on all other SDKs for consistency

Hm... yeah I see that all the SDKs are identical in their verbiage and usability. I'm sort of surprised if this use case -- concurrent execution and isolated execution instances -- hasn't been an issue across any of the other SDKs.

[zshipko]

For example, if compiling my module takes 80MB of memory (which it actually does, this is a real example) and I need to run 10 instances concurrently, then prior to the changes in this PR, that would be 800MB of allocated memory just for compilation. The vast majority of that goes away when I can compile once and instantiate 10 times instead.

Yeah, this change could help memory usage a lot! It seems worth porting this to the Rust/C SDKs too, but I think the naming is important if we're going to consider that.

I agree with @mhmd-azeez that something like CompiledPlugin which can be used to create Plugin seems best. That way Plugin stays the same and we could add a function like NewPluginFromCompiled or something that NewPlugin uses.

[clarkmcc]

@zshipko understood on the naming, I'll get that adjusted.

What's the direction we want to take on preserving the API?

[zshipko]

What's the direction we want to take on preserving the API?

It would be nice if we could mostly preserve the current API and just add to it, but that might not be possible.

I will experiment a little with the Rust SDK and come up with some more concrete suggestions.

[clarkmcc]

@zshipko curious if you've had a chance to give this any more thought.

Are there areas of the Go SDK we can keep moving forward on in the meantime?

[zshipko]

I just pushed up a rough implementation of this for Rust/C yesterday: extism/extism#784 - I'm still working on some tests and benchmarks.

Based on that it seems like we could expand the API with a CompiledPlugin type and NewCompiledPlugin/NewPluginFromCompiled functions. Then the existing API should be possible to implemented by calling NewPluginFromCompiled(NewCompiledPlugin(...)) in NewPlugin.

I think the CompiledPlugin should contain all the information needed to build a plugin - so this might include a manifest, host functions and anything else needed to create the plugin.

I'm also open to suggestions if there is something Go-specific that might make more sense.

[clarkmcc]

Based on that it seems like we could expand the API with a CompiledPlugin type and NewCompiledPlugin/NewPluginFromCompiled functions. Then the existing API should be possible to implemented by calling NewPluginFromCompiled(NewCompiledPlugin(...)) in NewPlugin.

@zshipko Hm... I'm not sure how I feel about this. If I create NewPluginFromCompiled or NewCompiledPlugin they'll both look identical by data structure, but one of them will be the owner of the Runtime, while the other just hold a reference to it. What happens when you want to Close a plugin? Do we close the runtime or not? It's not safe to close the runtime while there are still outstanding references to it by plugins created with NewPluginFromCompiled.

Let me illustrate

p1 := NewCompiledPlugin(...)
p2 := NewPluginFromCompiled(p1)

p2.Close(ctx) // should only close p2
p1.Close(ctx) // inadvertantly closes p2 as well

I'm in the process of renaming the types in my PR to more closely align with what you're looking for, but I personally feel that this API does not clearly denote the parent/child relationship that needs to exist between compiled modules and instantiated modules.

[clarkmcc]

So I've been testing the changes in my branch at a few of our production customers and the performance improvements are fantastic -- single compilation + thread safety seems to be working well.

I am noticing that there's some very inconsistent and elusive bugs happening specifically around the plugin.ReadBytes function. My running theory at this point is that some of the extism functions (like length) are stateful, but the state is contained in the extism module which is shared across all InstantiatedPlugins. I suspect that what's happening is the concurrent instances of my module are clobbering this shared state. In fact, it's possible that all memory allocations used to communicate between host/guest, across all plugin instances are happening against the same WASM linear memory in the extism module.

Edit: moving the extism module into the instance, and compiling a separate one for each plugin instance seems to have resolved the issue. I suspect there are ways to avoid that redundancy, but it would require a fundamental redesign of the approach extism takes here.

[clarkmcc]

@zshipko Alright, here's the latest update. After doing quite a bit of production testing with some key customers using these changes, we're preparing to release to our entire customer base behind a feature flag for more wide-spread testing.

As for an explanation behind my most recent changes:

• I've moved the env and extism modules into the CompiledPlugin, this means we only need to compile these once. The env module is stateless so not only do we compile it, but we also instantiate it into the runtime at the time the CompiledPlugin is created.
• Previously, the extism module also had to be instantiated during compilation because it was referenced by the env module, but since the extism module is stateful, it couldn't be safely shared to multiple InstantiatedPlugins without having what were essentially segmentation faults within the WASM linear memory. I've removed the env modules dependency on the extism module by (1) instantiating the pre-compiled extism module anonymously. Since it is anonymous, functions can't be called from within the WASM module, but this is okay because as far as I am aware, the only callers are actually Go host functions in the env module. So to continue, (2) now all the env functions access the extism module by context instead of directly. The extism module is injected into the context on every function call through the Extism API.

Loading

flowchart TD
    subgraph CompiledModule
        direction TB
        CM_MainWASM["Compiled Main WASM"]
        CM_ExtismWASM["Compiled Extism WASM"]
        CM_EnvModule["Instantiated Stateless Env Module"]
    end

    subgraph InstantiatedModule
        direction TB
        IM_MainWASM["Instantiated Main WASM"]
        IM_AnonExtismWASM["Instantiated Extism WASM (Anonymous)"]
        IM_EnvModuleRef["Reference to Env Module in CompiledModule"]
    end

    CM_MainWASM --> IM_MainWASM
    CM_ExtismWASM --> IM_AnonExtismWASM
    CM_EnvModule --> IM_EnvModuleRef
    IM_EnvModuleRef -- "Calls functions in Extism by passing context reference" --> IM_AnonExtismWASM
    IM_MainWASM -- "Calls Env Module Functions" --> IM_EnvModuleRef

[zshipko]

Thanks, this looks great! I am still just concerned about the naming since it should be consistent across all the SDKs, and right now the other SDKs don't have a concept of Compiled vs Instantiated plugins. It seems like renaming InstantiatedPlugin to Plugin could allow us to maintain the backwards compatible API by storing a CompiledPlugin somewhere only when it's owned by a single Plugin.

It also looks like your base commit is missing some changes from main - I could probably help rebase this if it's complicated.

[clarkmcc]

storing a CompiledPlugin somewhere only when it's owned by a single Plugin

@zshipko Could you possibly provide a small pseudo-example of what you have in mind? I'm not sure that I follow exactly how this would work/look.

[zshipko]

It looks like that isn't even needed - but I have a commit here that shows how we would like to preserve most of the existing Plugin API: fa6208c

I also noticed that host functions are moved into PluginConfig instead of being passed as arguments - this kind of breaking change might be possible, but in an ideal world we would avoid making any backward-incompatible changes.

[clarkmcc]

Got it, this makes sense.

I guess my only concern is this leaks memory. Instances of plugins can't close compiled plugins (for good reason if there are multiple instances). I think we can sort of hack around this by passing a closer that the instance that closes the compiled plugin when the instance is closed. We can make that private so that it can't be misused.

[zshipko]

Great point - some way to include that closer is definitely necessary!

[clarkmcc]

@zshipko okay, I've made those changes. The other thing that feels a little clunky is this. Any ideas?

Edit: for context, the reason we need the module config per-instance, not per-compiled plugin is because (for example) the module config is used to specify the stdout/stdin. When running multiple instances you most likely want those outputs to be separate io.Writers.

[zshipko]

Hm, I can't really see anywhere else to put it. I think it's okay to leave that how it is since it's explained in the comment.

[zshipko]

@clarkmcc - If you're done making changes, I can do one more review and work on rebasing this PR

[clarkmcc]

@zshipko I am done. I was partially through a rebase but one or two unit tests started failing and I haven't had a change to dig in to figure out why. If you have a spare second to attempt a rebase, that would be very much appreciated.

[zshipko]

just rebased and fixed the failing test - I will do a final review and merge this after discussing with the rest of the team early next week

thanks so much for your work on this!

[clarkmcc]

Awesome, thanks for pulling everything together and for the reviews. Let me know if there's anything else I can do to help on this!