diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs
new file mode 100644
index 00000000000..84b669382a5
--- /dev/null
+++ b/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;
+
+///
+/// A striped dictionary that distributes entries across multiple internal dictionaries
+/// to reduce lock contention by hashing correlation ids across stripes.
+///
+/// The type of values stored in the dictionary.
+internal sealed class StripedCallbackDictionary
+ 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;
+
+ ///
+ /// The number of stripes.
+ ///
+ 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();
+ }
+ }
+
+ ///
+ /// Computes the stripe index for a correlation id.
+ ///
+ [MethodImpl(MethodImplOptions.AggressiveInlining)]
+ public static int GetStripeIndex(CorrelationId correlationId)
+ => (int)(unchecked((ulong)correlationId.ToInt64() * HashFactor) >> (64 - StripeBits));
+
+ ///
+ /// Gets the stripe for the given callback id.
+ ///
+ [MethodImpl(MethodImplOptions.AggressiveInlining)]
+ private Stripe GetStripe(CorrelationId correlationId)
+ {
+ return _stripes[GetStripeIndex(correlationId)];
+ }
+
+ ///
+ /// Attempts to add the specified key and value to the dictionary.
+ ///
+ [MethodImpl(MethodImplOptions.AggressiveInlining)]
+ public bool TryAdd(CorrelationId id, TValue value)
+ {
+ var stripe = GetStripe(id);
+ lock (stripe.Lock)
+ {
+ return stripe.Dictionary.TryAdd(id, value);
+ }
+ }
+
+ ///
+ /// Attempts to get the value associated with the specified key.
+ ///
+ [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);
+ }
+ }
+
+ ///
+ /// Attempts to remove the value with the specified key.
+ ///
+ [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);
+ }
+ }
+
+ ///
+ /// Attempts to remove the value with the specified key if it is the expected instance or value.
+ ///
+ [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.Default.Equals(value, expected)
+ : !ReferenceEquals(value, expected)))
+ {
+ return false;
+ }
+
+ return stripe.Dictionary.Remove(id);
+ }
+ }
+
+ ///
+ /// Gets the approximate total count of items across all stripes.
+ ///
+ public int Count
+ {
+ get
+ {
+ int count = 0;
+ foreach (var stripe in _stripes)
+ {
+ lock (stripe.Lock)
+ {
+ count += stripe.Dictionary.Count;
+ }
+ }
+ return count;
+ }
+ }
+
+ ///
+ /// Counts items matching a predicate across all stripes.
+ ///
+ public int CountWhere(Func predicate)
+ => CountWhere(predicate, static (value, predicate) => predicate(value));
+
+ ///
+ /// Counts items matching a predicate across all stripes.
+ ///
+ public int CountWhere(TState state, Func 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;
+ }
+
+ ///
+ /// Visits a snapshot of the values in each stripe.
+ ///
+ public void ForEach(TState state, Action action)
+ {
+ foreach (var stripe in _stripes)
+ {
+ TValue[]? snapshot = null;
+ var snapshotCount = 0;
+ try
+ {
+ lock (stripe.Lock)
+ {
+ if (stripe.Dictionary.Count == 0)
+ {
+ continue;
+ }
+
+ snapshot = ArrayPool.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.Shared.Return(
+ snapshot,
+ clearArray: RuntimeHelpers.IsReferenceOrContainsReferences());
+ }
+ }
+ }
+ }
+
+ 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 Dictionary = new();
+ }
+}
diff --git a/src/Orleans.Core/Runtime/CallbackData.cs b/src/Orleans.Core/Runtime/CallbackData.cs
index bfbfb7a28f4..3383c0aba6a 100644
--- a/src/Orleans.Core/Runtime/CallbackData.cs
+++ b/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);
diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs
index bb86127f633..545233b4f3d 100644
--- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs
+++ b/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(),
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>)callbacks)
+ .Remove(KeyValuePair.Create(callback.Message.Id, callback));
private void ConstructorReset()
{
diff --git a/src/Orleans.Core/Runtime/SharedCallbackData.cs b/src/Orleans.Core/Runtime/SharedCallbackData.cs
index aef64491fbb..f83d57a8ad2 100644
--- a/src/Orleans.Core/Runtime/SharedCallbackData.cs
+++ b/src/Orleans.Core/Runtime/SharedCallbackData.cs
@@ -6,20 +6,20 @@ namespace Orleans.Runtime;
internal sealed class SharedCallbackData
{
- public readonly Action Unregister;
+ private readonly ICallbackDataTarget _target;
public readonly ILogger Logger;
private TimeSpan _responseTimeout;
public long ResponseTimeoutStopwatchTicks;
public SharedCallbackData(
- Action 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);
}
\ No newline at end of file
diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs
index d605549c37e..478e6ca9255 100644
--- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs
+++ b/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;
@@ -25,13 +24,14 @@ namespace Orleans.Runtime
///
/// Internal class for system grains to get access to runtime object
///
- internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecycleParticipant
+ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecycleParticipant, ICallbackDataTarget
{
private readonly ILogger logger;
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 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();
this.messageFactory = messageFactory;
this.ConcreteGrainFactory = new GrainFactory(this, referenceActivator, interfaceIdResolver, interfaceToTypeResolver);
this.logger = loggerFactory.CreateLogger();
@@ -88,7 +88,7 @@ public InsideRuntimeClient(
var callbackDataLogger = loggerFactory.CreateLogger();
this.sharedCallbackData = new SharedCallbackData(
- msg => this.UnregisterCallback(msg.SendingGrain, msg.Id),
+ this,
callbackDataLogger,
this.messagingOptions.ResponseTimeout,
this.messagingOptions.CancelRequestOnTimeout,
@@ -96,7 +96,7 @@ public InsideRuntimeClient(
cancellationManager: null!);
this.systemSharedCallbackData = new SharedCallbackData(
- msg => this.UnregisterCallback(msg.SendingGrain, msg.Id),
+ this,
callbackDataLogger,
this.messagingOptions.SystemResponseTimeout,
cancelOnTimeout: false,
@@ -195,7 +195,11 @@ public void SendRequest(
return;
}
- callbacks.TryAdd((message.SendingGrain, 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);
}
else
@@ -233,13 +237,8 @@ public void SendResponse(Message request, Response response)
this.MessageCenter.SendResponse(request, response);
}
- ///
- /// UnRegister a callback.
- ///
- private void UnregisterCallback(GrainId grainId, CorrelationId correlationId)
- {
- callbacks.TryRemove((grainId, correlationId), out _);
- }
+ void ICallbackDataTarget.Unregister(CallbackData callback) =>
+ callbacks.TryRemove(callback.Message.Id, callback);
public void SniffIncomingMessage(Message message)
{
@@ -468,7 +467,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 +482,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,7 +565,7 @@ private async Task OnRuntimeInitializeStop(CancellationToken tc)
private void BreakOutstandingMessages()
{
- foreach (var (_, callback) in callbacks)
+ callbacks.ForEach(this, static (callback, self) =>
{
try
{
@@ -574,9 +573,9 @@ private void BreakOutstandingMessages()
}
catch (Exception exception)
{
- LogWarningWhileProcessingCallbackExpiry(this.logger, exception);
+ LogWarningWhileProcessingCallbackExpiry(self.logger, exception);
}
- }
+ });
}
private Task OnRuntimeInitializeStart(CancellationToken tc)
@@ -600,13 +599,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 +615,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 +626,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)
{
diff --git a/test/Orleans.Runtime.Tests/CallbackDataTests.cs b/test/Orleans.Runtime.Tests/CallbackDataTests.cs
index dedcc46b756..8d5da45c884 100644
--- a/test/Orleans.Runtime.Tests/CallbackDataTests.cs
+++ b/test/Orleans.Runtime.Tests/CallbackDataTests.cs
@@ -10,11 +10,12 @@
namespace Tester;
+[TestSuite("BVT")]
+[TestProvider("None")]
+[TestCategory("BVT")]
public class CallbackDataTests
{
- [TestSuite("BVT")]
- [TestProvider("None")]
- [Fact, TestCategory("BVT")]
+ [Fact]
public void AlreadyCanceledTokenCompletesCallback()
{
using var serviceProvider = CreateServiceProvider();
@@ -35,9 +36,7 @@ public void AlreadyCanceledTokenCompletesCallback()
Assert.Equal(cancellation.Token, exception.CancellationToken);
}
- [TestSuite("BVT")]
- [TestProvider("None")]
- [Fact, TestCategory("BVT")]
+ [Fact]
public void CancellationSubscriptionAfterCompletionDoesNotRetainCallback()
{
using var serviceProvider = CreateServiceProvider();
@@ -55,6 +54,128 @@ public void CancellationSubscriptionAfterCompletionDoesNotRetainCallback()
GC.KeepAlive(cancellation);
}
+ [Fact]
+ public void TimeoutAndResponseRaceCompletesExactlyOnce()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var completion = new TestResponseCompletionSource();
+ var callback = registry.Register(new CorrelationId(1), completion);
+ var response = CreateResponse(callback.Message);
+
+ Parallel.Invoke(
+ callback.OnTimeout,
+ () => registry.TryCompleteResponse(response));
+
+ Assert.Equal(1, completion.CompletionCount);
+ Assert.Equal(0, registry.Count);
+ }
+
+ [Fact]
+ public void CancellationAndResponseRaceCompletesExactlyOnce()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ using var cancellation = new CancellationTokenSource();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var completion = new TestResponseCompletionSource();
+ var callback = registry.Register(new CorrelationId(2), completion);
+ callback.SubscribeForCancellation(cancellation.Token);
+ var response = CreateResponse(callback.Message);
+
+ Parallel.Invoke(
+ cancellation.Cancel,
+ () => registry.TryCompleteResponse(response));
+
+ Assert.Equal(1, completion.CompletionCount);
+ Assert.Equal(0, registry.Count);
+ }
+
+ [Fact]
+ public void StaleCancellationDoesNotRemoveReplacement()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ using var cancellation = new CancellationTokenSource();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var id = new CorrelationId(3);
+ var staleCompletion = new TestResponseCompletionSource();
+ var stale = registry.Register(id, staleCompletion);
+ stale.SubscribeForCancellation(cancellation.Token);
+ Assert.True(registry.TryTake(id, out var removed));
+ Assert.Same(stale, removed);
+ var replacementCompletion = new TestResponseCompletionSource();
+ var replacement = registry.Register(id, replacementCompletion);
+
+ cancellation.Cancel();
+
+ Assert.Same(replacement, registry.Take(id));
+ replacement.DoCallback(CreateResponse(replacement.Message));
+ Assert.IsType(staleCompletion.Response.Exception);
+ Assert.Same(Response.Completed, replacementCompletion.Response);
+ Assert.Equal(0, registry.Count);
+ }
+
+ [Fact]
+ public void StaleTimeoutDoesNotRemoveReplacement()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var id = new CorrelationId(4);
+ var staleCompletion = new TestResponseCompletionSource();
+ var stale = registry.Register(id, staleCompletion);
+ Assert.True(registry.TryTake(id, out var removed));
+ Assert.Same(stale, removed);
+ var replacementCompletion = new TestResponseCompletionSource();
+ var replacement = registry.Register(id, replacementCompletion);
+
+ stale.OnTimeout();
+
+ Assert.Same(replacement, registry.Take(id));
+ replacement.DoCallback(CreateResponse(replacement.Message));
+ Assert.IsType(staleCompletion.Response.Exception);
+ Assert.Same(Response.Completed, replacementCompletion.Response);
+ Assert.Equal(0, registry.Count);
+ }
+
+ [Fact]
+ public void StaleShutdownDoesNotRemoveReplacement()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var id = new CorrelationId(5);
+ var staleCompletion = new TestResponseCompletionSource();
+ var stale = registry.Register(id, staleCompletion);
+ Assert.True(registry.TryTake(id, out var removed));
+ Assert.Same(stale, removed);
+ var replacementCompletion = new TestResponseCompletionSource();
+ var replacement = registry.Register(id, replacementCompletion);
+
+ stale.OnHostShutdown();
+
+ Assert.Same(replacement, registry.Take(id));
+ replacement.DoCallback(CreateResponse(replacement.Message));
+ Assert.IsType(staleCompletion.Response.Exception);
+ Assert.Same(Response.Completed, replacementCompletion.Response);
+ Assert.Equal(0, registry.Count);
+ }
+
+ [Fact]
+ public void DuplicateRegistrationPreservesOriginalCallback()
+ {
+ using var serviceProvider = CreateServiceProvider();
+ var registry = new TestCallbackRegistry(CreateInstruments(serviceProvider));
+ var id = new CorrelationId(6);
+ var completion = new TestResponseCompletionSource();
+ var callback = registry.Register(id, completion);
+
+ var exception = Assert.Throws(() =>
+ registry.Register(id, new TestResponseCompletionSource()));
+
+ Assert.Contains(id.ToString(), exception.Message);
+ Assert.Same(callback, registry.Take(id));
+ callback.DoCallback(CreateResponse(callback.Message));
+ Assert.Same(Response.Completed, completion.Response);
+ }
+
[MethodImpl(MethodImplOptions.NoInlining)]
private static WeakReference CreateCompletedCallback(CancellationToken cancellationToken, ApplicationRequestInstruments instruments)
{
@@ -68,11 +189,11 @@ private static WeakReference CreateCompletedCallback(CancellationToken cancellat
private static CallbackData CreateCallback(
IResponseCompletionSource completion,
- Action unregister,
+ Action unregister,
ApplicationRequestInstruments instruments)
{
var shared = new SharedCallbackData(
- unregister,
+ new DelegateCallbackTarget(unregister),
logger: NullLogger.Instance,
responseTimeout: TimeSpan.FromMinutes(1),
cancelOnTimeout: false,
@@ -91,12 +212,88 @@ private static ServiceProvider CreateServiceProvider()
private static ApplicationRequestInstruments CreateInstruments(IServiceProvider serviceProvider) =>
new(new OrleansInstruments(serviceProvider.GetRequiredService()));
+ private static Message CreateResponse(Message request) => new()
+ {
+ Direction = Message.Directions.Response,
+ Id = request.Id,
+ BodyObject = Response.Completed,
+ };
+
private sealed class TestResponseCompletionSource : IResponseCompletionSource
{
- public Response Response { get; private set; } = null!;
+ private Response? _response;
+ private int _completionCount;
+
+ public Response Response => Volatile.Read(ref _response)!;
+
+ public int CompletionCount => Volatile.Read(ref _completionCount);
+
+ public void Complete(Response value)
+ {
+ Interlocked.Increment(ref _completionCount);
+ Interlocked.CompareExchange(ref _response, value, null);
+ }
+
+ public void Complete() => Complete(Response.Completed);
+ }
+
+ private sealed class DelegateCallbackTarget(Action unregister) : ICallbackDataTarget
+ {
+ public void Unregister(CallbackData callback) => unregister(callback);
+ }
+
+ private sealed class TestCallbackRegistry : ICallbackDataTarget
+ {
+ private readonly StripedCallbackDictionary _callbacks = new();
+ private readonly ApplicationRequestInstruments _instruments;
+ private readonly SharedCallbackData _sharedData;
+
+ public TestCallbackRegistry(ApplicationRequestInstruments instruments)
+ {
+ _instruments = instruments;
+ _sharedData = CreateSharedData(this);
+ }
+
+ public int Count => _callbacks.Count;
+
+ public CallbackData Register(CorrelationId id, IResponseCompletionSource completion)
+ {
+ var message = new Message { Id = id };
+ var callback = new CallbackData(_sharedData, completion, message, _instruments);
+ if (!_callbacks.TryAdd(id, callback))
+ {
+ throw new InvalidOperationException($"A callback with correlation id {id} is already registered.");
+ }
- public void Complete(Response value) => Response = value;
+ return callback;
+ }
+
+ public bool TryTake(CorrelationId id, out CallbackData? callback) =>
+ _callbacks.TryRemove(id, out callback);
+
+ public CallbackData Take(CorrelationId id)
+ {
+ Assert.True(_callbacks.TryRemove(id, out var callback));
+ return callback;
+ }
- public void Complete() => Response = Orleans.Serialization.Invocation.Response.Completed;
+ public void TryCompleteResponse(Message response)
+ {
+ if (_callbacks.TryRemove(response.Id, out var callback))
+ {
+ callback.DoCallback(response);
+ }
+ }
+
+ void ICallbackDataTarget.Unregister(CallbackData callback) =>
+ _callbacks.TryRemove(callback.Message.Id, callback);
}
+
+ private static SharedCallbackData CreateSharedData(ICallbackDataTarget target) => new(
+ target,
+ logger: NullLogger.Instance,
+ responseTimeout: TimeSpan.FromMinutes(1),
+ cancelOnTimeout: false,
+ waitForCancellationAcknowledgement: false,
+ cancellationManager: null);
}
diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs
new file mode 100644
index 00000000000..e49d12641e3
--- /dev/null
+++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs
@@ -0,0 +1,195 @@
+using Orleans.Runtime;
+using Xunit;
+
+namespace Tester;
+
+[TestSuite("BVT")]
+[TestProvider("None")]
+[TestCategory("BVT")]
+public class StripedCallbackDictionaryTests
+{
+ private static readonly Action EmptyVisitor = static (_, _) => { };
+ private static readonly Func MatchValue = static (value, expected) => value == expected;
+
+ [Fact]
+ public void CorrelationIdsDistributeAcrossStripesAtOverflowAndWithStride()
+ {
+ var start = long.MaxValue - (StripedCallbackDictionary.StripeCount / 2);
+ var consecutiveStripes = Enumerable.Range(0, StripedCallbackDictionary.StripeCount)
+ .Select(offset => new CorrelationId(unchecked(start + offset)))
+ .Select(StripedCallbackDictionary.GetStripeIndex)
+ .Distinct()
+ .Count();
+ var stridedStripes = Enumerable.Range(0, StripedCallbackDictionary.StripeCount)
+ .Select(offset => new CorrelationId(offset * StripedCallbackDictionary.StripeCount))
+ .Select(StripedCallbackDictionary.GetStripeIndex)
+ .Distinct()
+ .Count();
+
+ Assert.True(consecutiveStripes > StripedCallbackDictionary.StripeCount / 2);
+ Assert.True(stridedStripes > StripedCallbackDictionary.StripeCount / 2);
+ }
+
+ [Fact]
+ public void AddGetAndRemovePreserveValue()
+ {
+ var dictionary = new StripedCallbackDictionary();
+ var id = new CorrelationId(42);
+
+ Assert.True(dictionary.TryAdd(id, "value"));
+ Assert.False(dictionary.TryAdd(id, "duplicate"));
+ Assert.True(dictionary.TryGetValue(id, out var value));
+ Assert.Equal("value", value);
+ Assert.True(dictionary.TryRemove(id, out value));
+ Assert.Equal("value", value);
+ Assert.False(dictionary.TryGetValue(id, out _));
+ }
+
+ [Fact]
+ public void ExactRemovalDoesNotRemoveReplacement()
+ {
+ var dictionary = new StripedCallbackDictionary