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(); + var id = new CorrelationId(42); + var stale = new object(); + var replacement = new object(); + + Assert.True(dictionary.TryAdd(id, stale)); + Assert.True(dictionary.TryRemove(id, out var removed)); + Assert.Same(stale, removed); + Assert.True(dictionary.TryAdd(id, replacement)); + + Assert.False(dictionary.TryRemove(id, stale)); + Assert.True(dictionary.TryGetValue(id, out var current)); + Assert.Same(replacement, current); + Assert.True(dictionary.TryRemove(id, replacement)); + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void ExactRemovalSupportsValueTypes() + { + var dictionary = new StripedCallbackDictionary(); + var id = new CorrelationId(42); + + Assert.True(dictionary.TryAdd(id, 1)); + Assert.False(dictionary.TryRemove(id, 2)); + Assert.True(dictionary.TryGetValue(id, out var current)); + Assert.Equal(1, current); + Assert.True(dictionary.TryRemove(id, 1)); + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void EnumerationReturnsSnapshotValues() + { + var dictionary = new StripedCallbackDictionary(); + for (var i = 0; i < 32; i++) + { + var id = new CorrelationId(i); + Assert.True(dictionary.TryAdd(id, i)); + } + + var values = new List(); + dictionary.ForEach(values, static (value, values) => values.Add(value)); + + Assert.Equal(Enumerable.Range(0, 32), values.Order()); + } + + [Fact] + public void ConcurrentOperationsPreserveCountAndValues() + { + var dictionary = new StripedCallbackDictionary(); + + Parallel.For(0, 10_000, i => + { + var id = new CorrelationId(i); + Assert.True(dictionary.TryAdd(id, i)); + Assert.True(dictionary.TryGetValue(id, out var value)); + Assert.Equal(i, value); + }); + + Assert.Equal(10_000, dictionary.Count); + Assert.Equal(10_000, dictionary.CountWhere(static value => value >= 0)); + + Parallel.For(0, 10_000, i => + { + Assert.True(dictionary.TryRemove(new CorrelationId(i), out var value)); + Assert.Equal(i, value); + }); + + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void ConcurrentLookupAndRemovalRemainConsistent() + { + const int count = 10_000; + var dictionary = new StripedCallbackDictionary(); + for (var i = 0; i < count; i++) + { + Assert.True(dictionary.TryAdd(new CorrelationId(i), i)); + } + + Parallel.Invoke( + () => Parallel.For(0, count, i => + { + if (dictionary.TryGetValue(new CorrelationId(i), out var value)) + { + Assert.Equal(i, value); + } + }), + () => Parallel.For(0, count, i => + { + Assert.True(dictionary.TryRemove(new CorrelationId(i), out var value)); + Assert.Equal(i, value); + })); + + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void SnapshotVisitorAllowsValuesToRemoveThemselves() + { + var dictionary = new StripedCallbackDictionary(); + for (var i = 0; i < 32; i++) + { + Assert.True(dictionary.TryAdd(new CorrelationId(i), i)); + } + + dictionary.ForEach(dictionary, static (value, dictionary) => + { + Assert.True(dictionary.TryRemove(new CorrelationId(value), out var removed)); + Assert.Equal(value, removed); + }); + + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void EmptyVisitorDoesNotAllocate() + { + var dictionary = new StripedCallbackDictionary(); + dictionary.ForEach((object?)null, EmptyVisitor); + + var allocatedBefore = GC.GetAllocatedBytesForCurrentThread(); + dictionary.ForEach((object?)null, EmptyVisitor); + var allocated = GC.GetAllocatedBytesForCurrentThread() - allocatedBefore; + + Assert.Equal(0, allocated); + } + + [Fact] + public void StatefulCountDoesNotAllocate() + { + var dictionary = new StripedCallbackDictionary(); + Assert.True(dictionary.TryAdd(new CorrelationId(42), 42)); + Assert.Equal(1, dictionary.CountWhere(42, MatchValue)); + + var allocatedBefore = GC.GetAllocatedBytesForCurrentThread(); + var count = dictionary.CountWhere(42, MatchValue); + var allocated = GC.GetAllocatedBytesForCurrentThread() - allocatedBefore; + + Assert.Equal(1, count); + Assert.Equal(0, allocated); + } +}