BitFaster.Caching Orleans IGrainDirectoryCache implementation

ConcurrentLruGrainDirectoryCache.cs

using BitFaster.Caching;
using BitFaster.Caching.Lru;
using Microsoft.Extensions.Options;
using Orleans.Configuration;
using Orleans.Runtime;
using Orleans.Runtime.GrainDirectory;
using System.Diagnostics.CodeAnalysis;

namespace Orleans.BitFasterCaching.GrainDirectoryCache;

internal class ConcurrentLruGrainDirectoryCache(IOptions<GrainDirectoryOptions> options) : IGrainDirectoryCache
{
    private readonly ICache<GrainId, (GrainAddress ActivationAddress, int Version)> _cache = new ConcurrentLruBuilder<GrainId, (GrainAddress ActivationAddress, int Version)>()
        .WithMetrics()
        .WithCapacity(options.Value.CacheSize)
        .WithExpireAfterAccess(options.Value.MaximumCacheTTL)
        .Build();

    public void AddOrUpdate(GrainAddress activationAddress, int version)
    {
        // ignore the version number
        _cache.AddOrUpdate(activationAddress.GrainId, (activationAddress, version));
    }

    public bool Remove(GrainId key) => _cache.TryRemove(key);

    public bool Remove(GrainAddress grainAddress)
    {
        if (_cache.TryGet(grainAddress.GrainId, out var item))
        {
            if (MatchesForRemoval(grainAddress, item.ActivationAddress))
            {
                return _cache.TryRemove(new KeyValuePair<GrainId, (GrainAddress ActivationAddress, int Version)>(grainAddress.GrainId, item));
            }
        }

        return false;
    }

    public void Clear() => _cache.Clear();

    public bool LookUp(GrainId key, [NotNullWhen(true)] out GrainAddress? result, out int version)
    {
        if (_cache.TryGet(key, out var entry))
        {
            version = entry.Version;
            result = entry.ActivationAddress;
            return true;
        }

        version = default;
        result = default;
        return false;
    }

    public IEnumerable<(GrainAddress ActivationAddress, int Version)> KeyValues
    {
        get
        {
            foreach (var entry in _cache)
            {
                yield return (entry.Value.ActivationAddress, entry.Value.Version);
            }
        }
    }

    private static bool MatchesForRemoval(GrainAddress address, GrainAddress other)
    {
        return other is not null && address.GrainId.Equals(other.GrainId) && (address.SiloAddress?.Equals(other.SiloAddress) ?? other.SiloAddress is null);
    }
}

Orleans.BitFaster.Caching.csproj

<Project Sdk="Microsoft.NET.Sdk">

  <PropertyGroup>
    <TargetFramework>net8.0</TargetFramework>
    <ImplicitUsings>enable</ImplicitUsings>
    <Nullable>enable</Nullable>
  </PropertyGroup>

  <ItemGroup>
    <PackageReference Include="BitFaster.Caching" Version="2.4.1" />
    <PackageReference Include="Microsoft.Orleans.Runtime" Version="8.0.0" />
  </ItemGroup>

</Project>

SiloBuilderGrainDirectoryCacheExtensions.cs

using Microsoft.Extensions.DependencyInjection;
using Orleans.Configuration;
using Orleans.Hosting;
using Orleans.Runtime.GrainDirectory;

namespace Orleans.BitFasterCaching.GrainDirectoryCache;

public static class SiloBuilderGrainDirectoryCacheExtensions
{
    public static ISiloBuilder UseConcurrentLruGrainDirectoryCache(this ISiloBuilder builder)
    {
        var services = builder.Services;
        services.Configure<GrainDirectoryOptions>(options =>
        {
            options.CachingStrategy = GrainDirectoryOptions.CachingStrategyType.Custom;
        });

        services.AddSingleton<IGrainDirectoryCache, ConcurrentLruGrainDirectoryCache>();
        return builder;
    }
}