fix(runtime): preserve exact callback ownership

src/Orleans.Core/Messaging/StripedCallbackDictionary.cs

@@ -0,0 +1,212 @@
+#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>
+    /// Attempts to remove the value with the specified key if it is the expected instance or value.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public bool TryRemove(CorrelationId id, TValue expected)
+    {
+        var stripe = GetStripe(id);
+        lock (stripe.Lock)
+        {
+            if (!stripe.Dictionary.TryGetValue(id, out var value)
+                || (typeof(TValue).IsValueType
+                    ? !EqualityComparer<TValue>.Default.Equals(value, expected)
+                    : !ReferenceEquals(value, expected)))
+            {
+                return false;
+            }
+
+            return stripe.Dictionary.Remove(id);
+        }
+    }
+
+    /// <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.Core/Runtime/CallbackData.cs

@@ -6,13 +6,17 @@
 
 namespace Orleans.Runtime
 {
+    internal interface ICallbackDataTarget
+    {
+        void Unregister(CallbackData callback);
+    }
+
     internal sealed partial class CallbackData
     {
         private const int StateNone = 0;
         private const int StateCompleted = 1;
         private const int StateCancellationRegistrationPending = 2;
         private const int StateCancellationRegistrationPublished = 4;
-
         private readonly SharedCallbackData shared;
         private readonly IResponseCompletionSource context;
         private readonly ApplicationRequestInstruments _applicationRequestInstruments;
@@ -129,7 +133,7 @@ private void OnCancellation(CancellationToken cancellationToken)
 
             stopwatch.Stop();
             SignalCancellation();
-            shared.Unregister(Message);
+            shared.Unregister(this);
             _applicationRequestInstruments.OnAppRequestsEnd((long)stopwatch.Elapsed.TotalMilliseconds);
             _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType());
             OrleansCallBackDataEvent.Instance.OnCanceled(Message);
@@ -150,7 +154,7 @@ public void OnTimeout()
                 SignalCancellation();
             }
 
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
             _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds);
             _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType());
@@ -175,7 +179,7 @@ public void OnTargetSiloFail()
             }
 
             this.stopwatch.Stop();
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
             _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds);
 
@@ -195,7 +199,7 @@ public void OnHostShutdown()
             }
 
             this.stopwatch.Stop();
-            this.shared.Unregister(this.Message);
+            this.shared.Unregister(this);
             DisposeCancellationRegistration();
             _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds);
 

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -18,7 +18,7 @@
 
 namespace Orleans
 {
-    internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClusterConnectionStatusListener
+    internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClusterConnectionStatusListener, ICallbackDataTarget
     {
         internal static bool TestOnlyThrowExceptionDuringInit { get; set; }
 
@@ -93,7 +93,7 @@ public OutsideRuntimeClient(
                     TimeSpan.FromSeconds(1)));
             this.callbackTimer = new PeriodicTimer(period, timeProvider);
             this.sharedCallbackData = new SharedCallbackData(
-                msg => this.UnregisterCallback(msg.Id),
+                this,
                 this.loggerFactory.CreateLogger<CallbackData>(),
                 this.clientMessagingOptions.ResponseTimeout,
                 this.clientMessagingOptions.CancelRequestOnTimeout,
@@ -297,7 +297,11 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp
                     return;
                 }
 
-                callbacks.TryAdd(message.Id, callbackData);
+                if (!callbacks.TryAdd(message.Id, callbackData))
+                {
+                    throw new InvalidOperationException($"A callback with correlation id {message.Id} is already registered.");
+                }
+
                 callbackData.SubscribeForCancellation(cancellationToken);
 
                 if (Volatile.Read(ref _isStopping) != 0)
@@ -375,10 +379,9 @@ public void ReceiveResponse(Message response)
             }
         }
 
-        private void UnregisterCallback(CorrelationId id)
-        {
-            callbacks.TryRemove(id, out _);
-        }
+        void ICallbackDataTarget.Unregister(CallbackData callback) =>
+            ((ICollection<KeyValuePair<CorrelationId, CallbackData>>)callbacks)
+                .Remove(KeyValuePair.Create(callback.Message.Id, callback));
 
         private void ConstructorReset()
         {

src/Orleans.Core/Runtime/SharedCallbackData.cs

@@ -6,20 +6,20 @@ namespace Orleans.Runtime;
 
 internal sealed class SharedCallbackData
 {
-    public readonly Action<Message> Unregister;
+    private readonly ICallbackDataTarget _target;
     public readonly ILogger Logger;
     private TimeSpan _responseTimeout;
     public long ResponseTimeoutStopwatchTicks;
 
     public SharedCallbackData(
-        Action<Message> unregister,
+        ICallbackDataTarget target,
         ILogger logger,
         TimeSpan responseTimeout,
         bool cancelOnTimeout,
         bool waitForCancellationAcknowledgement,
         IGrainCallCancellationManager? cancellationManager)
     {
-        Unregister = unregister;
+        _target = target;
         Logger = logger;
         ResponseTimeout = responseTimeout;
         CancelRequestOnTimeout = cancelOnTimeout;
@@ -42,4 +42,6 @@ public TimeSpan ResponseTimeout
     public bool CancelRequestOnTimeout { get; }
 
     public bool WaitForCancellationAcknowledgement { get; }
+
+    public void Unregister(CallbackData callback) => _target.Unregister(callback);
 }