fix(runtime): preserve callback ownership

Author: ReubenBond

src/Orleans.Core/Messaging/CorrelationId.cs

@@ -16,7 +16,7 @@ namespace Orleans.Runtime
 
         public static CorrelationId GetNext() => new(System.Threading.Interlocked.Increment(ref lastUsed));
 
-        public override int GetHashCode() => HashCode.Combine(id);
+        public override int GetHashCode() => id.GetHashCode();
 
         public override bool Equals(object? obj) => obj is CorrelationId correlationId && Equals(correlationId);
 

src/Orleans.Core/Messaging/MessageFactory.cs

@@ -49,8 +49,7 @@ public Message CreateMessage(object? body, InvokeMethodOptions options)
         private CorrelationId GetNextCorrelationId()
         {
             var id = _seed ^ Interlocked.Increment(ref _nextId);
-            var stripeIndex = StripedCallbackDictionary<object>.GetCurrentThreadStripeIndex();
-            return StripedCallbackDictionary<object>.CreateCorrelationId(unchecked((long)id), stripeIndex);
+            return new CorrelationId(unchecked((long)id));
         }
 
         public Message CreateResponseMessage(Message request)

src/Orleans.Core/Messaging/StripedCallbackDictionary.cs

@@ -10,31 +10,20 @@ namespace Orleans.Runtime;
 
 /// <summary>
 /// A striped dictionary that distributes entries across multiple internal dictionaries
-/// to reduce lock contention. The stripe is determined by bits embedded in the CorrelationId.
+/// 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> : IEnumerable<KeyValuePair<CorrelationId, TValue>>
+internal sealed class StripedCallbackDictionary<TValue> : IEnumerable<TValue>
     where TValue : notnull
 {
-    /// <summary>
-    /// The number of bits used to identify the stripe (stored in the upper bits of the CorrelationId).
-    /// </summary>
-    public const int StripeBits = 7;
-
-    /// <summary>
-    /// The number of stripes (must be a power of 2).
-    /// </summary>
-    public const int StripeCount = 1 << StripeBits; // 128 stripes
-
-    /// <summary>
-    /// Mask to extract the stripe index from the upper bits.
-    /// </summary>
-    private const long StripeMask = (long)(StripeCount - 1) << (64 - StripeBits);
+    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 shift amount to move the stripe bits to the lowest position.
+    /// The number of stripes.
     /// </summary>
-    private const int StripeShift = 64 - StripeBits;
+    public const int StripeCount = 1 << StripeBits;
 
     private readonly Stripe[] _stripes;
 
@@ -48,40 +37,14 @@ public StripedCallbackDictionary()
     }
 
     /// <summary>
-    /// Encodes a stripe index into the upper bits of a base value to create a CorrelationId.
-    /// </summary>
-    /// <param name="baseValue">The base value (e.g., from an incrementing counter XORed with a seed).</param>
-    /// <param name="stripeIndex">The stripe index (typically derived from thread id).</param>
-    /// <returns>A CorrelationId with the stripe encoded in the upper bits.</returns>
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static CorrelationId CreateCorrelationId(long baseValue, int stripeIndex)
-    {
-        // Clear the upper StripeBits of the base value and set the stripe index there
-        long maskedBase = baseValue & ~StripeMask;
-        long stripeValue = (long)(stripeIndex & (StripeCount - 1)) << StripeShift;
-        return new CorrelationId(maskedBase | stripeValue);
-    }
-
-    /// <summary>
-    /// Extracts the stripe index from a CorrelationId.
+    /// Computes the stripe index for a correlation id.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public static int GetStripeIndex(CorrelationId correlationId)
-    {
-        return (int)((correlationId.ToInt64() & StripeMask) >>> StripeShift);
-    }
+        => (int)(unchecked((ulong)correlationId.ToInt64() * HashFactor) >> (64 - StripeBits));
 
     /// <summary>
-    /// Gets the stripe index for the current thread. Use this when creating new CorrelationIds.
-    /// </summary>
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static int GetCurrentThreadStripeIndex()
-    {
-        return Environment.CurrentManagedThreadId & (StripeCount - 1);
-    }
-
-    /// <summary>
-    /// Gets the stripe for the given correlation id.
+    /// Gets the stripe for the given callback id.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private Stripe GetStripe(CorrelationId correlationId)
@@ -93,38 +56,38 @@ private Stripe GetStripe(CorrelationId correlationId)
     /// Attempts to add the specified key and value to the dictionary.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public bool TryAdd(CorrelationId key, TValue value)
+    public bool TryAdd(GrainId owner, CorrelationId id, TValue value)
     {
-        var stripe = GetStripe(key);
+        var stripe = GetStripe(id);
         lock (stripe.Lock)
         {
-            return stripe.Dictionary.TryAdd(key, value);
+            return stripe.Dictionary.TryAdd(new(owner, id), value);
         }
     }
 
     /// <summary>
     /// Attempts to get the value associated with the specified key.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public bool TryGetValue(CorrelationId key, [NotNullWhen(true)] out TValue? value)
+    public bool TryGetValue(GrainId owner, CorrelationId id, [NotNullWhen(true)] out TValue? value)
     {
-        var stripe = GetStripe(key);
+        var stripe = GetStripe(id);
         lock (stripe.Lock)
         {
-            return stripe.Dictionary.TryGetValue(key, out value);
+            return stripe.Dictionary.TryGetValue(new(owner, id), out value);
         }
     }
 
     /// <summary>
     /// Attempts to remove the value with the specified key.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public bool TryRemove(CorrelationId key, [NotNullWhen(true)] out TValue? value)
+    public bool TryRemove(GrainId owner, CorrelationId id, [NotNullWhen(true)] out TValue? value)
     {
-        var stripe = GetStripe(key);
+        var stripe = GetStripe(id);
         lock (stripe.Lock)
         {
-            return stripe.Dictionary.Remove(key, out value);
+            return stripe.Dictionary.Remove(new(owner, id), out value);
         }
     }
 
@@ -150,16 +113,16 @@ public int Count
     /// <summary>
     /// Counts items matching a predicate across all stripes.
     /// </summary>
-    public int CountWhere(Func<KeyValuePair<CorrelationId, TValue>, bool> predicate)
+    public int CountWhere(Func<TValue, bool> predicate)
     {
         int count = 0;
         foreach (var stripe in _stripes)
         {
             lock (stripe.Lock)
             {
-                foreach (var kvp in stripe.Dictionary)
+                foreach (var value in stripe.Dictionary.Values)
                 {
-                    if (predicate(kvp))
+                    if (predicate(value))
                     {
                         count++;
                     }
@@ -175,21 +138,23 @@ public int CountWhere(Func<KeyValuePair<CorrelationId, TValue>, bool> predicate)
     /// </summary>
     public Enumerator GetEnumerator() => new(this);
 
-    IEnumerator<KeyValuePair<CorrelationId, TValue>> IEnumerable<KeyValuePair<CorrelationId, TValue>>.GetEnumerator() => GetEnumerator();
+    IEnumerator<TValue> IEnumerable<TValue>.GetEnumerator() => GetEnumerator();
 
     IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
     private sealed class Stripe
     {
         public readonly object Lock = new();
-        public readonly Dictionary<CorrelationId, TValue> Dictionary = new();
+        public readonly Dictionary<CallbackKey, TValue> Dictionary = new();
     }
 
-    public sealed class Enumerator : IEnumerator<KeyValuePair<CorrelationId, TValue>>
+    private readonly record struct CallbackKey(GrainId Owner, CorrelationId Id);
+
+    public sealed class Enumerator : IEnumerator<TValue>
     {
         private readonly StripedCallbackDictionary<TValue> _dictionary;
         private int _stripeIndex;
-        private KeyValuePair<CorrelationId, TValue>[]? _currentSnapshot;
+        private TValue[]? _currentSnapshot;
         private int _snapshotCount;
         private int _snapshotIndex;
 
@@ -202,7 +167,7 @@ internal Enumerator(StripedCallbackDictionary<TValue> dictionary)
             _snapshotIndex = -1;
         }
 
-        public KeyValuePair<CorrelationId, TValue> Current => _currentSnapshot![_snapshotIndex];
+        public TValue Current => _currentSnapshot![_snapshotIndex];
 
         object IEnumerator.Current => Current;
 
@@ -236,11 +201,11 @@ public bool MoveNext()
                 {
                     if (stripe.Dictionary.Count > 0)
                     {
-                        _currentSnapshot = ArrayPool<KeyValuePair<CorrelationId, TValue>>.Shared.Rent(stripe.Dictionary.Count);
+                        _currentSnapshot = ArrayPool<TValue>.Shared.Rent(stripe.Dictionary.Count);
                         _snapshotCount = 0;
-                        foreach (var pair in stripe.Dictionary)
+                        foreach (var value in stripe.Dictionary.Values)
                         {
-                            _currentSnapshot[_snapshotCount++] = pair;
+                            _currentSnapshot[_snapshotCount++] = value;
                         }
                         _snapshotIndex = -1;
                     }
@@ -265,9 +230,9 @@ private void ReturnSnapshot()
         {
             if (_currentSnapshot is { } snapshot)
             {
-                ArrayPool<KeyValuePair<CorrelationId, TValue>>.Shared.Return(
+                ArrayPool<TValue>.Shared.Return(
                     snapshot,
-                    clearArray: RuntimeHelpers.IsReferenceOrContainsReferences<KeyValuePair<CorrelationId, TValue>>());
+                    clearArray: RuntimeHelpers.IsReferenceOrContainsReferences<TValue>());
                 _currentSnapshot = null;
                 _snapshotCount = 0;
             }

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -92,7 +92,7 @@ public OutsideRuntimeClient(
                     TimeSpan.FromSeconds(1)));
             this.callbackTimer = new PeriodicTimer(period, timeProvider);
             this.sharedCallbackData = new SharedCallbackData(
-                msg => this.UnregisterCallback(msg.Id),
+                msg => this.UnregisterCallback(msg.SendingGrain, msg.Id),
                 this.loggerFactory.CreateLogger<CallbackData>(),
                 this.clientMessagingOptions.ResponseTimeout,
                 this.clientMessagingOptions.CancelRequestOnTimeout,
@@ -296,7 +296,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp
                     return;
                 }
 
-                callbacks.TryAdd(message.Id, callbackData);
+                callbacks.TryAdd(message.SendingGrain, message.Id, callbackData);
                 callbackData.SubscribeForCancellation(cancellationToken);
 
                 if (Volatile.Read(ref _isStopping) != 0)
@@ -327,7 +327,7 @@ public void ReceiveResponse(Message response)
             if (response.Result is Message.ResponseTypes.Status)
             {
                 var status = (StatusResponse)response.BodyObject!;
-                callbacks.TryGetValue(response.Id, out var callback);
+                callbacks.TryGetValue(response.TargetGrain, response.Id, out var callback);
                 var request = callback?.Message;
                 if (request is not null)
                 {
@@ -360,7 +360,7 @@ public void ReceiveResponse(Message response)
             }
 
             CallbackData? callbackData;
-            var found = callbacks.TryRemove(response.Id, out callbackData);
+            var found = callbacks.TryRemove(response.TargetGrain, response.Id, out callbackData);
             if (found)
             {
                 // We need to import the RequestContext here as well.
@@ -374,9 +374,9 @@ public void ReceiveResponse(Message response)
             }
         }
 
-        private void UnregisterCallback(CorrelationId id)
+        private void UnregisterCallback(GrainId owner, CorrelationId id)
         {
-            callbacks.TryRemove(id, out _);
+            callbacks.TryRemove(owner, id, out _);
         }
 
         private void ConstructorReset()
@@ -448,16 +448,16 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo)
         {
             foreach (var callback in callbacks)
             {
-                if (deadSilo.Equals(callback.Value.Message.TargetSilo))
+                if (deadSilo.Equals(callback.Message.TargetSilo))
                 {
-                    callback.Value.OnTargetSiloFail();
+                    callback.OnTargetSiloFail();
                 }
             }
         }
 
         private void BreakOutstandingMessages()
         {
-            foreach (var (_, callback) in callbacks)
+            foreach (var callback in callbacks)
             {
                 try
                 {
@@ -471,7 +471,7 @@ private void BreakOutstandingMessages()
         }
 
         public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType)
-            => this.callbacks.CountWhere(c => c.Value.Message.InterfaceType == grainInterfaceType);
+            => this.callbacks.CountWhere(c => c.Message.InterfaceType == grainInterfaceType);
 
         /// <inheritdoc />
         public void NotifyClusterConnectionLost()
@@ -515,7 +515,7 @@ private async Task MonitorCallbackExpiry()
                 try
                 {
                     var currentStopwatchTicks = ValueStopwatch.GetTimestamp();
-                    foreach (var (_, callback) in callbacks)
+                    foreach (var callback in callbacks)
                     {
                         if (callback.IsCompleted)
                         {