Better Caller

[martindevans]

For the last couple of days I've been contemplating how to improve the Caller system, what I've come up with entails a large refactor of how things are done so I thought I'd lay out my thoughts and get some feedback. @peterhuene and @kpreisser I'd really value your input :)

Context

At the moment a callback from WASM into C# can optionally take a Caller parameter which gives the call some context (e.g. access to Memory etc). For example:

Linker.DefineFunction("env", "check_string", (Caller caller, int address, int length) =>
{
    caller.GetMemory("mem").ReadString(address, length).Should().Be("Hello World");
});

The Problem

At the moment this mechanism always allocates a Caller for every call and will also often be used to fetch a Memory or Function which is in turn allocated. That's not great for performance (particularly in Unity, which is where I'm using wasmtime).

The Solution?

The obvious solution to this is to make Caller a ref struct. This also better models the intended lifetime of the Caller as well - you're not meant to hold it beyond that one method call.

However there is a problem with this - Caller implements IStore and in turns passes itself into Memory and Function constructors as the store reference. A ref struct cannot implement an interface so this doesn't work (it also can't be passed into a generic method, so we can't refactor to something like Invoke<T> where T : IStore either). So if we want to get a Memory or Function we'd be back needing two allocations, one for the store and one for the object itself.

To solve both problems we could introduce new memory/function types which are ref structs (e.g. RefStructMemory) which are returned by new methods (e.g. Caller.GetRefMemory). Since these would be ref structs they could contain the StoreContext directly. Again this also models the intended lifetime better.

Obviously this is a lot of duplicated code, but I think we could reduce that by putting all of the actual work inside the RefStructXXX types and the current class types would become wrappers which internally create the relevant RefStructXXX, call into it and immediately discard it.

Thoughts?

I realise this is a huge amount of churn, but given that it improves the handling of lifetimes as well as performance I think it's probably worth it.

[martindevans]

I just tried sketching some of this out in code and it looks like it won't work. A ref struct can't be used as a type parameter, so the nice type safe way we define all of our functions kills this idea I think.

[kpreisser]

A ref struct can't be used as a type parameter

I think I stumbled over exactly that same issue too when trying to avoid the Caller allocation as part of #186/#163 😄
(At least, the Caller is now only allocated when the callback actually uses it as parameter, or when it has a Function parameter.)

A possible way I can think of to solve this, might be to define additional delegate types similar to the built-in Action<...> and Func<...> that have a fixed Caller parameter, like this:

public delegate void CallerAction(Caller caller);
public delegate void CallerAction<in T>(Caller caller, T arg);
public delegate void CallerAction<in T1, in T2>(Caller caller, T1 arg1, T2 arg2);
// (...)
public delegate TResult CallerFunc<out TResult>(Caller caller);
public delegate TResult CallerFunc<in T, out TResult>(Caller caller, T arg);
public delegate TResult CallerFunc<in T1, in T2, out TResult>(Caller caller, T1 arg1, T2 arg2);
// (...)

Then, instead of

public static Function FromCallback<T, TResult>(IStore store, Func<Caller, T?, TResult> callback)

we would have

public static Function FromCallback<T, TResult>(IStore store, CallerFunc<T?, TResult> callback)

for the delegates that take a Caller, which I think should work if Caller is a ref struct.

[martindevans]

That's an interesting idea, thanks! I'll have a fiddle with it over the weekend to see if it looks like it'll work.

[kpreisser]

I could imagine a possible alternative to the RefStructMemory approach (but keeping the ref struct Caller), based on the changes in #200/#201 (but from the discussion it's not yet clear whether this is the way to go), which wouldn't require changes e.g. in Memory itself.
With these changes, the Caller would be able to get the Store it originates from.

As far as I understand the Extern types like ExternMemory (wasmtime_memory_t), they identify the object using a store_id (that identifies the store), and an index that is relative to the store and ExternKind. While a store exists, only new objects can be added to it, but existing ones cannot be removed (they would live as long as the Store lives); at least for the wasmtime_... externs.

So the Store could maintain a Dictionary<nuint, Memory> memories, Dictionary<nuint, Function> functions etc., or a single one for all types like Dictionary<(ExternKind kind, nuint index), IExternal> externs, using the index as key, where wrappers are stored once they are created. Then e.g. Caller.GetMemory, Instance.GetMemory etc. could do Store.Memories.TryGet(item.of.memory.index, out var existingMemory) to try to get the Memory with this index if it was already created previously, or create a new instance and add it to the dictionary otherwise.

That way, there would still be wrapper objects (Memory etc.) being allocated like it is now, but they would only be allocated once while the Store exists. So calls like caller.GetMemory("mem") would not cause an allocation every time it is called (only the first time the specific Memory is accessed). In this sense, the lifetime of the .NET wrappers would match more closely the lifetime of the native wasmtime_... objects.
(Though, there will be a few additional allocations when the Store is created, for the dictionary/-ies.)

However, I don't know Wasmtime good enough to determine whether this approach would be OK or whether e.g. it might rely too much on internals.
Thanks!

[martindevans]

I did consider a caching scheme like that. However since the lifetime of things returned from the caller seems to be deliberately tied to the lifetime of the caller I assumed it wouldn't be ok.

[martindevans]

I experimented with this over the weekend and have it mostly working. Currently failing just 7 tests due to 3 NotImplementedExceptions that I haven't resolved yet.

Two of these are the Caller.GetFunction and Caller.GetMemory implementations and I don't anticipate any major problems there (just a lot of refactoring).

However the final one is this method which invokes a method using reflection, it seems to be used as a fallback when the generic methods can't be used (for example too many arguments or too many return values). As far as I can see it's not possible to pass a ref struct to a method invoked in this way (the arguments are boxed, but ref structs cannot be boxed).

The only solution I can see is simply to fail in the case that the fallback is used and a Caller is required. Is that a sufficently weird edge case that we're ok with not supporting it?

[peterhuene]

Closed with #214.