perf(runtime): improve callback registry throughput

src/Orleans.Core/Runtime/CallbackData.cs

@@ -128,7 +128,7 @@ private void OnCancellation(CancellationToken cancellationToken)
 
             RecordElapsedTime();
             SignalCancellation();
-            shared.Unregister(Message);
+            shared.Unregister(this);
             _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType());
             OrleansCallBackDataEvent.Instance.OnCanceled(Message);
             context.Complete(Response.FromException(new OperationCanceledException(cancellationToken)));
@@ -148,7 +148,7 @@ public void OnTimeout()
                 SignalCancellation();
             }
 
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
             _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType());
 
@@ -172,7 +172,7 @@ public void OnTargetSiloFail()
             }
 
             RecordElapsedTime();
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
 
             OrleansCallBackDataEvent.Instance.OnTargetSiloFail(this.Message);
@@ -191,7 +191,7 @@ public void OnHostShutdown()
             }
 
             RecordElapsedTime();
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
 
             var msg = this.Message;
@@ -200,10 +200,15 @@ public void OnHostShutdown()
         }
 
         public void DoCallback(Message response)
+        {
+            TryDoCallback(response);
+        }
+
+        internal bool TryDoCallback(Message response)
         {
             if (!TryComplete())
             {
-                return;
+                return false;
             }
 
             OrleansCallBackDataEvent.Instance.DoCallback(this.Message);
@@ -213,6 +218,7 @@ public void DoCallback(Message response)
 
             // do callback outside the CallbackData lock. Just not a good practice to hold a lock for this unrelated operation.
             ResponseCallback(response, this.context);
+            return true;
         }
 
         private bool TryComplete() => (Interlocked.Or(ref _state, StateCompleted) & StateCompleted) == 0;

src/Orleans.Core/Runtime/IRuntimeClient.cs

@@ -57,9 +57,14 @@ internal interface IRuntimeClient
         IGrainReferenceRuntime GrainReferenceRuntime { get; }
 
         void BreakOutstandingMessagesToSilo(SiloAddress deadSilo);
+    }
 
-        // For testing purposes only.
-        int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType);
+    /// <summary>
+    /// Exposes runtime request state to tests which synchronize request lifecycle transitions.
+    /// </summary>
+    internal interface IRuntimeClientTestAccessor
+    {
+        int GetRunningRequestCount(GrainInterfaceType grainInterfaceType);
     }
 
     /// <summary>

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -18,7 +18,7 @@
 
 namespace Orleans
 {
-    internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClusterConnectionStatusListener
+    internal partial class OutsideRuntimeClient : IRuntimeClient, IRuntimeClientTestAccessor, IDisposable, IClusterConnectionStatusListener
     {
         internal static bool TestOnlyThrowExceptionDuringInit { get; set; }
 
@@ -97,7 +97,7 @@ public OutsideRuntimeClient(
                     TimeSpan.FromSeconds(1)));
             this.callbackTimer = new PeriodicTimer(period, timeProvider);
             this.sharedCallbackData = new SharedCallbackData(
-                msg => this.UnregisterCallback(msg.Id),
+                this.UnregisterCallback,
                 this.loggerFactory.CreateLogger<CallbackData>(),
                 timeProvider,
                 this.clientMessagingOptions.ResponseTimeout,
@@ -312,6 +312,8 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp
                 callbacks.TryAdd(message.Id, callbackData);
                 callbackData.SubscribeForCancellation(cancellationToken);
 
+                // Shutdown sets _isStopping before sweeping callbacks. Recheck it after registration so that
+                // a callback published after the sweep passed it completes and removes itself here.
                 if (Volatile.Read(ref _isStopping) != 0)
                 {
                     callbackData.OnHostShutdown();
@@ -387,9 +389,9 @@ public void ReceiveResponse(Message response)
             }
         }
 
-        private void UnregisterCallback(CorrelationId id)
+        private void UnregisterCallback(CallbackData callback)
         {
-            callbacks.TryRemove(id, out _);
+            callbacks.TryRemove(KeyValuePair.Create(callback.Message.Id, callback));
         }
 
         private void ConstructorReset()
@@ -493,7 +495,7 @@ private void BreakOutstandingMessages()
             }
         }
 
-        public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType)
+        int IRuntimeClientTestAccessor.GetRunningRequestCount(GrainInterfaceType grainInterfaceType)
             => this.callbacks.Count(c => c.Value.Message.InterfaceType == grainInterfaceType);
 
         /// <inheritdoc />

src/Orleans.Core/Runtime/SharedCallbackData.cs

@@ -5,14 +5,14 @@ namespace Orleans.Runtime;
 
 internal sealed class SharedCallbackData
 {
-    public readonly Action<Message> Unregister;
+    public readonly Action<CallbackData> Unregister;
     public readonly ILogger Logger;
     public readonly TimeProvider TimeProvider;
     private TimeSpan _responseTimeout;
     public long ResponseTimeoutTimestampTicks;
 
     public SharedCallbackData(
-        Action<Message> unregister,
+        Action<CallbackData> unregister,
         ILogger logger,
         TimeProvider timeProvider,
         TimeSpan responseTimeout,

src/Orleans.Runtime/Core/CallbackRegistry.cs

@@ -0,0 +1,169 @@
+using System;
+using System.Buffers;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+namespace Orleans.Runtime;
+
+internal sealed class CallbackRegistry
+{
+    private const int StripeBits = 7;
+    private const ulong HashFactor = 11_400_714_819_323_198_485;
+
+    internal const int StripeCount = 1 << StripeBits;
+
+    private readonly Stripe[] _stripes = CreateStripes();
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal static int GetStripeIndex(CorrelationId correlationId)
+        => (int)(unchecked((ulong)correlationId.ToInt64() * HashFactor) >> (64 - StripeBits));
+
+    // MessageFactory assigns host-unique correlation ids, so the id is sufficient as the registry key.
+    public void Register(CallbackData callback)
+    {
+        var id = callback.Message.Id;
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            if (stripe.Callbacks.TryAdd(id, callback))
+            {
+                return;
+            }
+        }
+
+        throw new InvalidOperationException($"A callback with id '{id}' is already registered.");
+    }
+
+    public bool TryCompleteResponse(Message response)
+    {
+        var stripe = GetStripe(response.Id);
+        CallbackData callback;
+        lock (stripe.Lock)
+        {
+            if (!stripe.Callbacks.Remove(response.Id, out callback!))
+            {
+                return false;
+            }
+        }
+
+        _ = callback.TryDoCallback(response);
+        return true;
+    }
+
+    public bool TryGetResponseCallback(Message response, [NotNullWhen(true)] out CallbackData? callback)
+    {
+        var stripe = GetStripe(response.Id);
+        lock (stripe.Lock)
+        {
+            return stripe.Callbacks.TryGetValue(response.Id, out callback);
+        }
+    }
+
+    internal bool ContainsKey(CorrelationId id)
+    {
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            return stripe.Callbacks.ContainsKey(id);
+        }
+    }
+
+    public bool TryRemove(CallbackData callback)
+    {
+        var id = callback.Message.Id;
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            if (!stripe.Callbacks.TryGetValue(id, out var current) || !ReferenceEquals(current, callback))
+            {
+                return false;
+            }
+
+            return stripe.Callbacks.Remove(id);
+        }
+    }
+
+    // Test assertions synchronize the request lifecycle before calling this method. Locking one stripe
+    // at a time keeps this inspection path isolated from runtime operations and avoids global locking.
+    internal int GetRunningRequestCountForTest(GrainInterfaceType grainInterfaceType)
+    {
+        var result = 0;
+        foreach (var stripe in _stripes)
+        {
+            lock (stripe.Lock)
+            {
+                foreach (var callback in stripe.Callbacks.Values)
+                {
+                    if (callback.Message.InterfaceType == grainInterfaceType)
+                    {
+                        result++;
+                    }
+                }
+            }
+        }
+
+        return result;
+    }
+
+    public void ForEach<TState>(TState state, Action<CallbackData, TState> action)
+    {
+        foreach (var stripe in _stripes)
+        {
+            CallbackData[]? snapshot = null;
+            var snapshotCount = 0;
+            try
+            {
+                lock (stripe.Lock)
+                {
+                    if (stripe.Callbacks.Count == 0)
+                    {
+                        continue;
+                    }
+
+                    snapshot = ArrayPool<CallbackData>.Shared.Rent(stripe.Callbacks.Count);
+                    foreach (var callback in stripe.Callbacks.Values)
+                    {
+                        snapshot[snapshotCount++] = callback;
+                    }
+                }
+
+                for (var index = 0; index < snapshotCount; index++)
+                {
+                    action(snapshot[index], state);
+                }
+            }
+            finally
+            {
+                if (snapshot is not null)
+                {
+                    ArrayPool<CallbackData>.Shared.Return(snapshot, clearArray: true);
+                }
+            }
+        }
+    }
+
+    private static Stripe[] CreateStripes()
+    {
+        var result = new Stripe[StripeCount];
+        for (var index = 0; index < result.Length; index++)
+        {
+            result[index] = new Stripe();
+        }
+
+        return result;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private Stripe GetStripe(CorrelationId id) => _stripes[GetStripeIndex(id)];
+
+    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, CallbackData> Callbacks = new();
+    }
+}