perf(runtime): cache local message targets with directory entries

src/Orleans.Core.Abstractions/Runtime/GrainReference.cs

@@ -267,6 +267,18 @@ public class GrainReference : IAddressable, IEquatable<GrainReference>, ISpanFor
         [NonSerialized]
         private readonly IdSpan _key;
 
+        [NonSerialized]
+        private object? _messageTargetCache;
+
+        internal object? MessageTargetCache
+        {
+            get => Volatile.Read(ref _messageTargetCache);
+            set => Volatile.Write(ref _messageTargetCache, value);
+        }
+
+        internal void ClearMessageTargetCache(object expected)
+            => Interlocked.CompareExchange(ref _messageTargetCache, null, expected);
+
         /// <summary>
         /// Gets the grain reference functionality which is shared by all grain references of a given type.
         /// </summary>

src/Orleans.Core/Caching/ConcurrentLruCache.cs

@@ -64,6 +64,12 @@ public ConcurrentLruCache(int capacity) : this(capacity, comparer: null)
     {
     }
 
+    protected virtual void UpdateItem(LruItem item, V value) => item.Value = value;
+
+    protected virtual void OnItemRemoved(LruItem item)
+    {
+    }
+
     /// <summary>
     /// Initializes a new instance of the ConcurrentLruCore class with the specified capacity and expire-after-access time to live.
     /// </summary>
@@ -161,15 +167,25 @@ public V Get(K key)
     ///<inheritdoc/>
     public bool TryGet(K key, [MaybeNullWhen(false)] out V value)
     {
-        if (_dictionary.TryGetValue(key, out var item))
+        if (TryGetItem(key, out var item))
         {
             value = item.Value;
-            Touch(item);
-            _telemetryPolicy.IncrementHit();
             return true;
         }
 
         value = default;
+        return false;
+    }
+
+    protected bool TryGetItem(K key, [NotNullWhen(true)] out LruItem? item)
+    {
+        if (_dictionary.TryGetValue(key, out item))
+        {
+            TouchItem(item);
+            _telemetryPolicy.IncrementHit();
+            return true;
+        }
+
         _telemetryPolicy.IncrementMiss();
         return false;
     }
@@ -326,6 +342,7 @@ private void OnRemove(LruItem item, ItemRemovedReason reason)
         // from the queue.
         item.WasAccessed = false;
         item.WasRemoved = true;
+        OnItemRemoved(item);
 
         if (reason == ItemRemovedReason.Evicted)
         {
@@ -351,7 +368,7 @@ public bool TryUpdate(K key, V value)
                 {
                     var oldValue = existing.Value;
 
-                    existing.Value = value;
+                    UpdateItem(existing, value);
                     UpdateTimestamp(existing);
 
                     _telemetryPolicy.IncrementUpdated();
@@ -853,7 +870,7 @@ private static ItemDestination RouteCold(LruItem item)
     /// <param name="value">The value.</param>
     // NOTE: Internal for testing
     [DebuggerDisplay("[{Key}] = {Value}")]
-    internal sealed class LruItem(K key, V value, long timestamp = 0)
+    internal class LruItem(K key, V value, long timestamp = 0)
     {
         private V _data = value;
 
@@ -1057,11 +1074,13 @@ private async Task RunExpirationLoop()
         }
     }
 
-    private LruItem CreateItem(K key, V value) =>
-        new(key, value, _expiresAfterAccess ? _timeProvider.GetTimestamp() : 0);
+    protected virtual LruItem CreateItem(K key, V value) =>
+        new(key, value, GetCurrentTimestamp());
+
+    protected long GetCurrentTimestamp() => _expiresAfterAccess ? _timeProvider.GetTimestamp() : 0;
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private void Touch(LruItem item)
+    protected void TouchItem(LruItem item)
     {
         if (_expiresAfterAccess)
         {

src/Orleans.Core/Runtime/GrainReferenceRuntime.cs

@@ -107,7 +107,13 @@ public object Cast(IAddressable grain, Type grainInterface)
             }
 
             var interfaceType = this.interfaceTypeResolver.GetGrainInterfaceType(grainInterface);
-            return this.referenceActivator.CreateReference(grainId, interfaceType);
+            var result = this.referenceActivator.CreateReference(grainId, interfaceType);
+            if (grain is GrainReference source)
+            {
+                result.MessageTargetCache = source.MessageTargetCache;
+            }
+
+            return result;
         }
 
         /// <summary>

src/Orleans.Runtime/Core/InsideRuntimeClient.cs

@@ -216,7 +216,7 @@ public void SendRequest(
             }
 
             this.messagingTrace.OnSendRequest(message);
-            this.MessageCenter.AddressAndSendMessage(message);
+            this.MessageCenter.AddressAndSendMessage(message, target);
         }
 
         public void SendResponse(Message request, Response response)

src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs

@@ -218,6 +218,20 @@ public void UpdateCache(GrainId grainId, SiloAddress siloAddress)
         }
         public void InvalidateCache(GrainId grainId) => cache.Remove(grainId);
         public void InvalidateCache(GrainAddress address) => cache.Remove(address);
+
+        internal bool TryGetCacheEntry(GrainId grainId, SiloAddress siloAddress, [NotNullWhen(true)] out GrainDirectoryCacheEntry? entry)
+        {
+            if (cache is IGrainDirectoryCacheEntrySource entrySource
+                && entrySource.TryGetEntry(grainId, out entry)
+                && entry.Address.SiloAddress?.Equals(siloAddress) == true)
+            {
+                return true;
+            }
+
+            entry = null;
+            return false;
+        }
+
         public bool TryLookupInCache(GrainId grainId, [NotNullWhen(true)] out GrainAddress? address)
         {
             var grainType = grainId.Type;

src/Orleans.Runtime/GrainDirectory/DhtGrainLocator.cs

@@ -81,6 +81,17 @@ public static DhtGrainLocator FromLocalGrainDirectory(LocalGrainDirectory localG
         public void InvalidateCache(GrainAddress address) => _localGrainDirectory.InvalidateCacheEntry(address);
         public bool TryLookupInCache(GrainId grainId, [NotNullWhen(true)] out GrainAddress? address) => _localGrainDirectory.TryLocalLookup(grainId, out address);
 
+        internal bool TryGetCacheEntry(GrainId grainId, SiloAddress siloAddress, [NotNullWhen(true)] out GrainDirectoryCacheEntry? entry)
+        {
+            if (_localGrainDirectory is LocalGrainDirectory directory)
+            {
+                return directory.TryGetCacheEntry(grainId, siloAddress, out entry);
+            }
+
+            entry = null;
+            return false;
+        }
+
         private class BatchedDeregistrationWorker
         {
             private const int OperationBatchSizeLimit = 2_000;

src/Orleans.Runtime/GrainDirectory/GrainDirectoryCacheEntry.cs

@@ -0,0 +1,180 @@
+using System.Threading;
+using Orleans.Caching;
+
+namespace Orleans.Runtime.GrainDirectory;
+
+internal sealed class GrainDirectoryCacheEntry
+    : ConcurrentLruCache<GrainId, (GrainAddress ActivationAddress, int Version)>.LruItem,
+      IDisposable
+{
+    private static readonly object Invalidated = new();
+    private static readonly object Updating = new();
+    private readonly LruGrainDirectoryCache? _owner;
+    private readonly WeakReference<GrainDirectoryCacheEntry> _referenceHandle;
+    private object? _messageTarget;
+
+    public GrainDirectoryCacheEntry(GrainAddress address, int version)
+        : this(owner: null, address.GrainId, (address, version), timestamp: 0)
+    {
+    }
+
+    public GrainDirectoryCacheEntry(
+        LruGrainDirectoryCache? owner,
+        GrainId grainId,
+        (GrainAddress ActivationAddress, int Version) value,
+        long timestamp)
+        : base(grainId, value, timestamp)
+    {
+        _owner = owner;
+        _referenceHandle = new(this);
+    }
+
+    public GrainAddress Address => Value.ActivationAddress;
+
+    public int Version => Value.Version;
+
+    public WeakReference<GrainDirectoryCacheEntry> ReferenceHandle => _referenceHandle;
+
+    public bool TryTouch()
+    {
+        if (!IsValid)
+        {
+            return false;
+        }
+
+        _owner?.Touch(this);
+        return IsValid;
+    }
+
+    public bool IsValid
+    {
+        get
+        {
+            var target = Volatile.Read(ref _messageTarget);
+            return !ReferenceEquals(target, Invalidated) && !ReferenceEquals(target, Updating);
+        }
+    }
+
+    public bool TryGetMessageTarget(out object? messageTarget)
+    {
+        var target = Volatile.Read(ref _messageTarget);
+        if (ReferenceEquals(target, Invalidated) || ReferenceEquals(target, Updating))
+        {
+            messageTarget = null;
+            return false;
+        }
+
+        messageTarget = target;
+        return messageTarget is not null;
+    }
+
+    public bool TrySetMessageTarget(object messageTarget, GrainAddress expectedAddress)
+    {
+        ArgumentNullException.ThrowIfNull(messageTarget);
+        ArgumentNullException.ThrowIfNull(expectedAddress);
+        if (!Address.Matches(expectedAddress) || !TrySetMessageTargetCore(messageTarget))
+        {
+            return false;
+        }
+
+        if (Address.Matches(expectedAddress))
+        {
+            return true;
+        }
+
+        ClearMessageTarget(messageTarget);
+        return false;
+    }
+
+    public bool TrySetMessageTarget(object messageTarget, SiloAddress expectedSilo)
+    {
+        ArgumentNullException.ThrowIfNull(messageTarget);
+        ArgumentNullException.ThrowIfNull(expectedSilo);
+        if (Address.SiloAddress?.Equals(expectedSilo) != true || !TrySetMessageTargetCore(messageTarget))
+        {
+            return false;
+        }
+
+        if (Address.SiloAddress?.Equals(expectedSilo) == true)
+        {
+            return true;
+        }
+
+        ClearMessageTarget(messageTarget);
+        return false;
+    }
+
+    private bool TrySetMessageTargetCore(object messageTarget)
+    {
+        var current = Volatile.Read(ref _messageTarget);
+        if (ReferenceEquals(current, Invalidated) || ReferenceEquals(current, Updating))
+        {
+            return false;
+        }
+
+        return ReferenceEquals(current, messageTarget)
+            || current is null && Interlocked.CompareExchange(ref _messageTarget, messageTarget, null) is null;
+    }
+
+    public void ClearMessageTarget(object messageTarget)
+    {
+        ArgumentNullException.ThrowIfNull(messageTarget);
+        Interlocked.CompareExchange(ref _messageTarget, null, messageTarget);
+    }
+
+    public void Invalidate() => Interlocked.Exchange(ref _messageTarget, Invalidated);
+
+    public void Dispose() => Invalidate();
+
+    internal void Update((GrainAddress ActivationAddress, int Version) value)
+    {
+        var updateStarted = TryBeginUpdate();
+        try
+        {
+            Value = value;
+        }
+        finally
+        {
+            if (updateStarted)
+            {
+                EndUpdate();
+            }
+        }
+    }
+
+    public void ClearMessageTarget()
+    {
+        while (true)
+        {
+            var current = Volatile.Read(ref _messageTarget);
+            if (current is null || ReferenceEquals(current, Invalidated))
+            {
+                return;
+            }
+
+            if (ReferenceEquals(Interlocked.CompareExchange(ref _messageTarget, null, current), current))
+            {
+                return;
+            }
+        }
+    }
+
+    internal bool TryBeginUpdate()
+    {
+        while (true)
+        {
+            var current = Volatile.Read(ref _messageTarget);
+            if (ReferenceEquals(current, Invalidated))
+            {
+                return false;
+            }
+
+            if (ReferenceEquals(Interlocked.CompareExchange(ref _messageTarget, Updating, current), current))
+            {
+                return true;
+            }
+        }
+    }
+
+    internal void EndUpdate() => Interlocked.CompareExchange(ref _messageTarget, null, Updating);
+}

src/Orleans.Runtime/GrainDirectory/GrainLocator.cs

@@ -54,6 +54,26 @@ public GrainLocator(GrainLocatorResolver grainLocatorResolver, DirectoryInstrume
 
         public bool TryLookupInCache(GrainId grainId, [NotNullWhen(true)] out GrainAddress? address) => GetGrainLocator(grainId.Type).TryLookupInCache(grainId, out address);
 
+        internal bool TryGetCacheEntry(
+            GrainId grainId,
+            SiloAddress siloAddress,
+            [NotNullWhen(true)] out GrainDirectoryCacheEntry? entry)
+        {
+            var grainLocator = GetGrainLocator(grainId.Type);
+            return grainLocator switch
+            {
+                CachedGrainLocator cached => cached.TryGetCacheEntry(grainId, siloAddress, out entry),
+                DhtGrainLocator dht => dht.TryGetCacheEntry(grainId, siloAddress, out entry),
+                _ => ReturnFalse(out entry),
+            };
+
+            static bool ReturnFalse(out GrainDirectoryCacheEntry? result)
+            {
+                result = null;
+                return false;
+            }
+        }
+
         public void InvalidateCache(GrainId grainId) => GetGrainLocator(grainId.Type).InvalidateCache(grainId);
 
         public void InvalidateCache(GrainAddress address) => GetGrainLocator(address.GrainId.Type).InvalidateCache(address);

src/Orleans.Runtime/GrainDirectory/IGrainDirectoryCache.cs

@@ -65,4 +65,9 @@ public static bool LookUp(this IGrainDirectoryCache cache, GrainId key, [NotNull
             return cache.LookUp(key, out result, out _);
         }
     }
+
+    internal interface IGrainDirectoryCacheEntrySource
+    {
+        bool TryGetEntry(GrainId key, [NotNullWhen(true)] out GrainDirectoryCacheEntry? entry);
+    }
 }