Caller Function Cache

[martindevans]

Added caching to the Store for Function and Memory objects. The Caller uses this to retrieve Function/Memory objects, this saves allocating lots of objects in caller contexts.

This is an alternative to #233 and #224. I prefer this approach to #233 since it gives access to the entire Function API, instead of a tiny subset exposed through the Caller. I prefer it to #224 since the Memory object remains a class instead of becoming a struct, which is more idiomatic.

[martindevans]

@kpreisser I've updated this PR with some extra work, basically adding the same caching for Memory as well as Function. That means it supersedes #224 as well (and we can conveniently ignore the MacOS CI weirdness).

I'd appreciate your re-review :)

[kpreisser]

I added a few notes, but generally this looks good to me, thanks!

However, we will probably need to wait for @peterhuene, to verify e.g. that the usage of only the index field of the ExternXyz is correct/sufficient to match the object from the dictionary (I would think so, as the dictionary will only store objects that belong to the current Store), or if it would technically be required to compare the whole ExternXyz (store + index).

[martindevans]

I've addressed those three comments.

However, we will probably need to wait for @peterhuene, to verify e.g. that the usage of only the index field of the ExternXyz is correct/sufficient to match the object from the dictionary (I would think so, as the dictionary will only store objects that belong to the current Store), or if it would technically be required to compare the whole ExternXyz (store + index).

I would like to add a sanity check that extern.store == this_store, but I'm not sure exactly what I should compare with. The docs simply say it's an "internal identifier" and as far as I can see there's no way to retrieve that identifier for a given wasmtime_store*.

[peterhuene]

So I don't think you'll be able to assert on store ids (there is no way to get the id from the wasmtime_store_t*, but could we assert that the item (memory/function) being return from the cache has the same store object reference as this?

Edit: actually, I don't think that buys us much at all. Wasmtime has a bunch of places it checks for cross-store referencing and will panic, so I'm fairly confident the places where we're assuming these objects come from the same store are safe.

[peterhuene]

LGTM.

I'll merge after @kpreisser completes their review.

[martindevans]

I don't think checking the Store field of the Memory/Function will help much, since those objects are all constructed inside the Store (referencing this) in the first place.

If it's not possible to check then I think we're probably good enough as we are. Any objects constructed in the cache will use the store context internally and will presumably detect errors for us. For example Memory does this:

Native.wasmtime_memory_type(store.Context.handle, this.memory);

Which I assume will internally sanity check if the store context and the memory are mismatched?

[kpreisser]

👍

Additionally, to be on the safe side regarding the uint64_t store_id (ulong store) value of externs that Wasmtime uses, we could use both fields from the Extern... as key (ValueTuple) for the dictionary, so that we don't need to assume that Wasmtime always uses the same ulong store value for Extern... objects for the current store, but instead we simply use the same values to identify the extern as Wasmtime does. E.g.

private readonly ConcurrentDictionary<(ulong store, nuint index), Function> _funcCache = new();
internal Function GetCachedExtern(ExternFunc @extern)
{
    var key = (@extern.store, @extern.index);
    if (!_funcCache.TryGetValue(key, out var func))
    {
        func = new Function(this, @extern);
        func = _funcCache.GetOrAdd(key, func);
    }

    return func;
}

This should work as ValueTuple<...> provides an implementation of Equals and GetHashCode that combines the members of the tuple.

What do you think?
Thanks!

[martindevans]

I've changed the key to a tuple of store/index as @kpreisser suggested.

[kpreisser]

Would you mind adding tests that e.g. verify that two consecutive calls of Caller.GetFunction, Caller.GetMemory with the same name actually return the same Function/Memory instance (so the cache works)?

Apart from that and the comment, this looks good to me, thanks!

A possible optimization idea for the future (especially when a cache is also added for other externs like Table), would be to check whether the multiple ConcurrentDictionarys can be combined into a single one, where the key additionally consists of ExternKind, because AFAIK creating a ConcurrentDictionary involes a number of allocations, which could then be saved.

[martindevans]

Would you mind adding tests that...

I've added tests for:

• Getting memory multiple times always returns the same object.
• Getting a non-existant memory returns null, even after getting something else from the cache.
• Getting a function multiple time always returns the same object.
• Getting two different functions does not return the same object.

A possible optimization idea ... key of ExternKind

I did look into this but it's made a bit difficult by the way Extern is built at the moment - the ExternUnion doesn't know which variant it represents so equality/GetHashCode can't be implemented on it.

I'll probably come back to look at this again soon. It looks like it can be fixed by merging the fields of ExternUnion into Extern so that all the necessary information is in one place. This would create lots of small changes across the rest of the codebase, so I didn't want to do that in this PR.

[kpreisser]

Thanks!

I did look into this but it's made a bit difficult by the way Extern is built at the moment - the ExternUnion doesn't know which variant it represents so equality/GetHashCode can't be implemented on it.

Actually, what I meant was to use something like

ConcurrentDictionary<(ExternKind kind, ulong store, nuint index), IExternal> _externCache

i.e. add ExternKind to the key, which I think should work since currently all externs have the same fields (ulong store and nuint index), so that only a single ConcurrentDictionary would be needed. The functions could then look like this:

internal Function GetCachedExtern(ExternFunc @extern)
{
    var key = (ExternKind.Func, @extern.store, @extern.index);

    if (!_externCache.TryGetValue(key, out var func))
    {
        func = new Function(this, @extern);
        func = _externCache.GetOrAdd(key, func);
    }

    return (Function)func;
}

internal Memory GetCachedExtern(ExternMemory @extern)
{
    var key = (ExternKind.Memory, @extern.store, @extern.index);

    if (!_externCache.TryGetValue(key, out var mem))
    {
        mem = new Memory(this, @extern);
        mem = _externCache.GetOrAdd(key, mem);
    }

    return (Memory)mem;
}

However, a downside is the additional cast e.g. to Function or Memory, since the dictionary's value type parameter is IExternal.
What do you think?

[martindevans]

Actually, what I meant was...

That's much simpler than what I had in mind, good idea. I've added that.

[peterhuene]

Looks good, thanks again for contributing this!