perf(runtime): reduce callback tracking contention

src/Orleans.Core/Messaging/StripedCallbackDictionary.cs

@@ -0,0 +1,191 @@
+#nullable enable
+using System;
+using System.Buffers;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+namespace Orleans.Runtime;
+
+/// <summary>
+/// A striped dictionary that distributes entries across multiple internal dictionaries
+/// to reduce lock contention by hashing correlation ids across stripes.
+/// </summary>
+/// <typeparam name="TValue">The type of values stored in the dictionary.</typeparam>
+internal sealed class StripedCallbackDictionary<TValue>
+    where TValue : notnull
+{
+    private const int StripeBits = 7;
+    // Fibonacci hashing spreads sequential and strided ids using one multiply and shift.
+    private const ulong HashFactor = 11_400_714_819_323_198_485;
+
+    /// <summary>
+    /// The number of stripes.
+    /// </summary>
+    public const int StripeCount = 1 << StripeBits;
+
+    private readonly Stripe[] _stripes;
+
+    public StripedCallbackDictionary()
+    {
+        _stripes = new Stripe[StripeCount];
+        for (int i = 0; i < StripeCount; i++)
+        {
+            _stripes[i] = new Stripe();
+        }
+    }
+
+    /// <summary>
+    /// Computes the stripe index for a correlation id.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static int GetStripeIndex(CorrelationId correlationId)
+        => (int)(unchecked((ulong)correlationId.ToInt64() * HashFactor) >> (64 - StripeBits));
+
+    /// <summary>
+    /// Gets the stripe for the given callback id.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private Stripe GetStripe(CorrelationId correlationId)
+    {
+        return _stripes[GetStripeIndex(correlationId)];
+    }
+
+    /// <summary>
+    /// Attempts to add the specified key and value to the dictionary.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public bool TryAdd(CorrelationId id, TValue value)
+    {
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            return stripe.Dictionary.TryAdd(id, value);
+        }
+    }
+
+    /// <summary>
+    /// Attempts to get the value associated with the specified key.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public bool TryGetValue(CorrelationId id, [NotNullWhen(true)] out TValue? value)
+    {
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            return stripe.Dictionary.TryGetValue(id, out value);
+        }
+    }
+
+    /// <summary>
+    /// Attempts to remove the value with the specified key.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public bool TryRemove(CorrelationId id, [NotNullWhen(true)] out TValue? value)
+    {
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            return stripe.Dictionary.Remove(id, out value);
+        }
+    }
+
+    /// <summary>
+    /// Gets the approximate total count of items across all stripes.
+    /// </summary>
+    public int Count
+    {
+        get
+        {
+            int count = 0;
+            foreach (var stripe in _stripes)
+            {
+                lock (stripe.Lock)
+                {
+                    count += stripe.Dictionary.Count;
+                }
+            }
+            return count;
+        }
+    }
+
+    /// <summary>
+    /// Counts items matching a predicate across all stripes.
+    /// </summary>
+    public int CountWhere(Func<TValue, bool> predicate)
+        => CountWhere(predicate, static (value, predicate) => predicate(value));
+
+    /// <summary>
+    /// Counts items matching a predicate across all stripes.
+    /// </summary>
+    public int CountWhere<TState>(TState state, Func<TValue, TState, bool> predicate)
+    {
+        int count = 0;
+        foreach (var stripe in _stripes)
+        {
+            lock (stripe.Lock)
+            {
+                foreach (var value in stripe.Dictionary.Values)
+                {
+                    if (predicate(value, state))
+                    {
+                        count++;
+                    }
+                }
+            }
+        }
+        return count;
+    }
+
+    /// <summary>
+    /// Visits a snapshot of the values in each stripe.
+    /// </summary>
+    public void ForEach<TState>(TState state, Action<TValue, TState> action)
+    {
+        foreach (var stripe in _stripes)
+        {
+            TValue[]? snapshot = null;
+            var snapshotCount = 0;
+            try
+            {
+                lock (stripe.Lock)
+                {
+                    if (stripe.Dictionary.Count == 0)
+                    {
+                        continue;
+                    }
+
+                    snapshot = ArrayPool<TValue>.Shared.Rent(stripe.Dictionary.Count);
+                    foreach (var value in stripe.Dictionary.Values)
+                    {
+                        snapshot[snapshotCount++] = value;
+                    }
+                }
+
+                for (var i = 0; i < snapshotCount; i++)
+                {
+                    action(snapshot[i], state);
+                }
+            }
+            finally
+            {
+                if (snapshot is not null)
+                {
+                    ArrayPool<TValue>.Shared.Return(
+                        snapshot,
+                        clearArray: RuntimeHelpers.IsReferenceOrContainsReferences<TValue>());
+                }
+            }
+        }
+    }
+
+    private sealed class Stripe
+    {
+#if NET9_0_OR_GREATER
+        public readonly System.Threading.Lock Lock = new();
+#else
+        public readonly object Lock = new();
+#endif
+        public readonly Dictionary<CorrelationId, TValue> Dictionary = new();
+    }
+}

src/Orleans.Runtime/Core/InsideRuntimeClient.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
@@ -31,7 +30,8 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa
         private readonly ILogger invokeExceptionLogger;
         private readonly ILoggerFactory loggerFactory;
         private readonly SiloMessagingOptions messagingOptions;
-        private readonly ConcurrentDictionary<(GrainId, CorrelationId), CallbackData> callbacks;
+        // MessageFactory assigns unique correlation ids to every request created by this runtime client.
+        private readonly StripedCallbackDictionary<CallbackData> callbacks;
         private readonly InterfaceToImplementationMappingCache interfaceToImplementationMapping;
         private readonly SharedCallbackData sharedCallbackData;
         private readonly SharedCallbackData systemSharedCallbackData;
@@ -74,7 +74,7 @@ public InsideRuntimeClient(
             this._applicationRequestInstruments = new(orleansInstruments);
             this.ServiceProvider = serviceProvider;
             this.MySilo = siloDetails.SiloAddress;
-            this.callbacks = new ConcurrentDictionary<(GrainId, CorrelationId), CallbackData>();
+            this.callbacks = new StripedCallbackDictionary<CallbackData>();
             this.messageFactory = messageFactory;
             this.ConcreteGrainFactory = new GrainFactory(this, referenceActivator, interfaceIdResolver, interfaceToTypeResolver);
             this.logger = loggerFactory.CreateLogger<InsideRuntimeClient>();
@@ -88,15 +88,15 @@ public InsideRuntimeClient(
 
             var callbackDataLogger = loggerFactory.CreateLogger<CallbackData>();
             this.sharedCallbackData = new SharedCallbackData(
-                msg => this.UnregisterCallback(msg.SendingGrain, msg.Id),
+                msg => this.UnregisterCallback(msg.Id),
                 callbackDataLogger,
                 this.messagingOptions.ResponseTimeout,
                 this.messagingOptions.CancelRequestOnTimeout,
                 this.messagingOptions.WaitForCancellationAcknowledgement,
                 cancellationManager: null!);
 
             this.systemSharedCallbackData = new SharedCallbackData(
-                msg => this.UnregisterCallback(msg.SendingGrain, msg.Id),
+                msg => this.UnregisterCallback(msg.Id),
                 callbackDataLogger,
                 this.messagingOptions.SystemResponseTimeout,
                 cancelOnTimeout: false,
@@ -195,7 +195,7 @@ public void SendRequest(
                     return;
                 }
 
-                callbacks.TryAdd((message.SendingGrain, message.Id), callbackData);
+                callbacks.TryAdd(message.Id, callbackData);
                 callbackData.SubscribeForCancellation(cancellationToken);
             }
             else
@@ -236,9 +236,9 @@ public void SendResponse(Message request, Response response)
         /// <summary>
         /// UnRegister a callback.
         /// </summary>
-        private void UnregisterCallback(GrainId grainId, CorrelationId correlationId)
+        private void UnregisterCallback(CorrelationId correlationId)
         {
-            callbacks.TryRemove((grainId, correlationId), out _);
+            callbacks.TryRemove(correlationId, out _);
         }
 
         public void SniffIncomingMessage(Message message)
@@ -468,7 +468,7 @@ public void ReceiveResponse(Message message)
 
         private void ProcessResponseCallback(Message message)
         {
-            if (callbacks.TryRemove((message.TargetGrain, message.Id), out var callbackData))
+            if (callbacks.TryRemove(message.Id, out var callbackData))
             {
                 // IMPORTANT: we do not schedule the response callback via the scheduler, since the only thing it does
                 // is to resolve/break the resolver. The continuations/waits that are based on this resolution will be scheduled as work items.
@@ -483,7 +483,7 @@ private void ProcessResponseCallback(Message message)
         private void ProcessStatusResponse(Message message)
         {
             var status = (StatusResponse)message.BodyObject!;
-            callbacks.TryGetValue((message.TargetGrain, message.Id), out var callback);
+            callbacks.TryGetValue(message.Id, out var callback);
             var request = callback?.Message;
             if (request is not null)
             {
@@ -566,17 +566,17 @@ private async Task OnRuntimeInitializeStop(CancellationToken tc)
 
         private void BreakOutstandingMessages()
         {
-            foreach (var (_, callback) in callbacks)
+            callbacks.ForEach(this, static (callback, self) =>
             {
                 try
                 {
                     callback.OnHostShutdown();
                 }
                 catch (Exception exception)
                 {
-                    LogWarningWhileProcessingCallbackExpiry(this.logger, exception);
+                    LogWarningWhileProcessingCallbackExpiry(self.logger, exception);
                 }
-            }
+            });
         }
 
         private Task OnRuntimeInitializeStart(CancellationToken tc)
@@ -600,13 +600,13 @@ override public string ToString()
 
         public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo)
         {
-            foreach (var callback in callbacks)
+            callbacks.ForEach(deadSilo, static (callback, deadSilo) =>
             {
-                if (deadSilo.Equals(callback.Value.Message.TargetSilo))
+                if (deadSilo.Equals(callback.Message.TargetSilo))
                 {
-                    callback.Value.OnTargetSiloFail();
+                    callback.OnTargetSiloFail();
                 }
-            }
+            });
         }
 
         public void Participate(ISiloLifecycle lifecycle)
@@ -616,7 +616,9 @@ public void Participate(ISiloLifecycle lifecycle)
         }
 
         public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType)
-            => this.callbacks.Count(c => c.Value.Message.InterfaceType == grainInterfaceType);
+            => this.callbacks.CountWhere(
+                grainInterfaceType,
+                static (callback, grainInterfaceType) => callback.Message.InterfaceType == grainInterfaceType);
 
         private async Task MonitorCallbackExpiry()
         {
@@ -625,18 +627,18 @@ private async Task MonitorCallbackExpiry()
                 try
                 {
                     var currentStopwatchTicks = ValueStopwatch.GetTimestamp();
-                    foreach (var (_, callback) in callbacks)
+                    callbacks.ForEach(currentStopwatchTicks, static (callback, currentStopwatchTicks) =>
                     {
                         if (callback.IsCompleted)
                         {
-                            continue;
+                            return;
                         }
 
                         if (callback.IsExpired(currentStopwatchTicks))
                         {
                             callback.OnTimeout();
                         }
-                    }
+                    });
                 }
                 catch (Exception ex)
                 {