From 28464dacb0d59c752ef86697fd9d8aecce891507 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Wed, 29 Apr 2026 17:32:36 -0700 Subject: [PATCH 01/18] perf(runtime): reduce callback tracking contention Distribute callbacks across striped dictionaries using correlation-id bits so concurrent request registration and completion contend on independent locks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Orleans.Core/Messaging/CorrelationId.cs | 2 +- src/Orleans.Core/Messaging/MessageFactory.cs | 3 +- .../Messaging/StripedCallbackDictionary.cs | 255 ++++++++++++++++++ .../Runtime/OutsideRuntimeClient.cs | 110 +++----- .../Core/InsideRuntimeClient.cs | 221 +++++---------- 5 files changed, 368 insertions(+), 223 deletions(-) create mode 100644 src/Orleans.Core/Messaging/StripedCallbackDictionary.cs diff --git a/src/Orleans.Core/Messaging/CorrelationId.cs b/src/Orleans.Core/Messaging/CorrelationId.cs index b40167a4fca..c3a7b055dd3 100644 --- a/src/Orleans.Core/Messaging/CorrelationId.cs +++ b/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() => id.GetHashCode(); + public override int GetHashCode() => HashCode.Combine(id); public override bool Equals(object? obj) => obj is CorrelationId correlationId && Equals(correlationId); diff --git a/src/Orleans.Core/Messaging/MessageFactory.cs b/src/Orleans.Core/Messaging/MessageFactory.cs index 1f815944ea9..3cf1e429c26 100644 --- a/src/Orleans.Core/Messaging/MessageFactory.cs +++ b/src/Orleans.Core/Messaging/MessageFactory.cs @@ -49,7 +49,8 @@ public Message CreateMessage(object? body, InvokeMethodOptions options) private CorrelationId GetNextCorrelationId() { var id = _seed ^ Interlocked.Increment(ref _nextId); - return new CorrelationId(unchecked((long)id)); + var stripeIndex = StripedCallbackDictionary.GetCurrentThreadStripeIndex(); + return StripedCallbackDictionary.CreateCorrelationId(unchecked((long)id), stripeIndex); } public Message CreateResponseMessage(Message request) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs new file mode 100644 index 00000000000..3e131d8dda9 --- /dev/null +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -0,0 +1,255 @@ +#nullable enable +using System; +using System.Collections; +using System.Collections.Generic; +using System.Runtime.CompilerServices; + +namespace Orleans.Runtime; + +/// +/// A striped dictionary that distributes entries across multiple internal dictionaries +/// to reduce lock contention. The stripe is determined by bits embedded in the CorrelationId. +/// +/// The type of values stored in the dictionary. +internal sealed class StripedCallbackDictionary : IEnumerable> +{ + /// + /// The number of bits used to identify the stripe (stored in the upper bits of the CorrelationId). + /// + public const int StripeBits = 7; + + /// + /// The number of stripes (must be a power of 2). + /// + public const int StripeCount = 1 << StripeBits; // 128 stripes + + /// + /// Mask to extract the stripe index from the upper bits. + /// + private const long StripeMask = (long)(StripeCount - 1) << (64 - StripeBits); + + /// + /// The shift amount to move the stripe bits to the lowest position. + /// + private const int StripeShift = 64 - StripeBits; + + private readonly Stripe[] _stripes; + + public StripedCallbackDictionary() + { + _stripes = new Stripe[StripeCount]; + for (int i = 0; i < StripeCount; i++) + { + _stripes[i] = new Stripe(); + } + } + + /// + /// Encodes a stripe index into the upper bits of a base value to create a CorrelationId. + /// + /// The base value (e.g., from an incrementing counter XORed with a seed). + /// The stripe index (typically derived from thread id). + /// A CorrelationId with the stripe encoded in the upper bits. + [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); + } + + /// + /// Extracts the stripe index from a CorrelationId. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static int GetStripeIndex(CorrelationId correlationId) + { + return (int)((correlationId.ToInt64() & StripeMask) >>> StripeShift); + } + + /// + /// Gets the stripe index for the current thread. Use this when creating new CorrelationIds. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static int GetCurrentThreadStripeIndex() + { + return Environment.CurrentManagedThreadId & (StripeCount - 1); + } + + /// + /// Gets the stripe for the given correlation 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 key, TValue value) + { + var stripe = GetStripe(key); + lock (stripe.Lock) + { + return stripe.Dictionary.TryAdd(key, value); + } + } + + /// + /// Attempts to get the value associated with the specified key. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public bool TryGetValue(CorrelationId key, out TValue? value) + { + var stripe = GetStripe(key); + lock (stripe.Lock) + { + return stripe.Dictionary.TryGetValue(key, out value); + } + } + + /// + /// Attempts to remove the value with the specified key. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public bool TryRemove(CorrelationId key, out TValue? value) + { + var stripe = GetStripe(key); + lock (stripe.Lock) + { + return stripe.Dictionary.Remove(key, out value); + } + } + + /// + /// 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, bool> predicate) + { + int count = 0; + foreach (var stripe in _stripes) + { + lock (stripe.Lock) + { + foreach (var kvp in stripe.Dictionary) + { + if (predicate(kvp)) + { + count++; + } + } + } + } + return count; + } + + /// + /// Returns an enumerator that iterates through all items in all stripes. + /// Note: This takes a snapshot of each stripe under its lock. + /// + public Enumerator GetEnumerator() => new(this); + + IEnumerator> IEnumerable>.GetEnumerator() => GetEnumerator(); + + IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); + + private sealed class Stripe + { + public readonly object Lock = new(); + public readonly Dictionary Dictionary = new(); + } + + public struct Enumerator : IEnumerator> + { + private readonly StripedCallbackDictionary _dictionary; + private int _stripeIndex; + private List>? _currentSnapshot; + private int _snapshotIndex; + + internal Enumerator(StripedCallbackDictionary dictionary) + { + _dictionary = dictionary; + _stripeIndex = -1; + _currentSnapshot = null; + _snapshotIndex = -1; + } + + public KeyValuePair Current => _currentSnapshot![_snapshotIndex]; + + object IEnumerator.Current => Current; + + public bool MoveNext() + { + while (true) + { + // Try to advance within current snapshot + if (_currentSnapshot != null) + { + _snapshotIndex++; + if (_snapshotIndex < _currentSnapshot.Count) + { + return true; + } + } + + // Move to next stripe + _stripeIndex++; + if (_stripeIndex >= _dictionary._stripes.Length) + { + _currentSnapshot = null; + return false; + } + + // Take a snapshot of the next stripe + var stripe = _dictionary._stripes[_stripeIndex]; + lock (stripe.Lock) + { + if (stripe.Dictionary.Count > 0) + { + _currentSnapshot = new List>(stripe.Dictionary); + _snapshotIndex = -1; + } + else + { + _currentSnapshot = null; + } + } + } + } + + public void Reset() + { + _stripeIndex = -1; + _currentSnapshot = null; + _snapshotIndex = -1; + } + + public void Dispose() + { + _currentSnapshot = null; + } + } +} diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index bb86127f633..38f7a2a1c30 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -16,6 +16,7 @@ using Orleans.Serialization.Invocation; using static Orleans.Internal.StandardExtensions; +#nullable disable namespace Orleans { internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClusterConnectionStatusListener @@ -25,44 +26,43 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus private readonly ILogger logger; private readonly ClientMessagingOptions clientMessagingOptions; - private readonly ConcurrentDictionary callbacks; - private InvokableObjectManager? localObjects; - private int _isStopping; + private readonly StripedCallbackDictionary callbacks; + private InvokableObjectManager localObjects; private bool disposing; private bool disposed; private readonly MessagingTrace messagingTrace; private readonly InterfaceToImplementationMappingCache _interfaceToImplementationMapping; private readonly ApplicationRequestInstruments _applicationRequestInstruments; - private IGrainCallCancellationManager? _cancellationManager; - private IClusterConnectionStatusObserver[]? _statusObservers; + private IGrainCallCancellationManager _cancellationManager; + private IClusterConnectionStatusObserver[] _statusObservers; - public IInternalGrainFactory InternalGrainFactory { get; private set; } = null!; + public IInternalGrainFactory InternalGrainFactory { get; private set; } - private ClientClusterManifestProvider? _manifestProvider; - private MessageFactory? messageFactory; + private ClientClusterManifestProvider _manifestProvider; + private MessageFactory messageFactory; private readonly LocalClientDetails _localClientDetails; private readonly ILoggerFactory loggerFactory; private readonly SharedCallbackData sharedCallbackData; private readonly PeriodicTimer callbackTimer; - private Task? callbackTimerTask; + private Task callbackTimerTask; public GrainAddress CurrentActivationAddress { get; private set; - } = null!; - public ClientGatewayObserver? gatewayObserver { get; private set; } + } + public ClientGatewayObserver gatewayObserver { get; private set; } public string CurrentActivationIdentity { get { return CurrentActivationAddress.ToString(); } } - public IGrainReferenceRuntime GrainReferenceRuntime { get; private set; } = null!; + public IGrainReferenceRuntime GrainReferenceRuntime { get; private set; } - internal ClientMessageCenter? MessageCenter { get; private set; } + internal ClientMessageCenter MessageCenter { get; private set; } [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "MessageCenter is IDisposable but cannot call Dispose yet as it lives past the end of this method call.")] @@ -72,7 +72,7 @@ public OutsideRuntimeClient( IOptions clientMessagingOptions, MessagingTrace messagingTrace, IServiceProvider serviceProvider, - [FromKeyedServices(TimeProviderNames.Messaging)] TimeProvider timeProvider, + TimeProvider timeProvider, InterfaceToImplementationMappingCache interfaceToImplementationMapping, OrleansInstruments orleansInstruments) { @@ -84,7 +84,7 @@ public OutsideRuntimeClient( this.loggerFactory = loggerFactory; this.messagingTrace = messagingTrace; this.logger = loggerFactory.CreateLogger(); - callbacks = new ConcurrentDictionary(); + callbacks = new StripedCallbackDictionary(); this.clientMessagingOptions = clientMessagingOptions.Value; var period = Max( TimeSpan.FromMilliseconds(1), @@ -155,13 +155,8 @@ public async Task StartAsync(CancellationToken cancellationToken) public async Task StopAsync(CancellationToken cancellationToken) { - Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); - // Fault callbacks before any cancellation-sensitive waits. Completing them can resume code - // which issues follow-up calls, so request admission must already be closed. - BreakOutstandingMessages(); - if (this.callbackTimerTask is { } task) { await task.WaitAsync(cancellationToken); @@ -194,20 +189,20 @@ await ExecuteWithRetries( MessageCenter = ActivatorUtilities.CreateInstance(this.ServiceProvider); MessageCenter.RegisterLocalMessageHandler(this.HandleMessage); await ExecuteWithRetries( - async () => await MessageCenter!.StartAsync(cancellationToken), + async () => await MessageCenter.StartAsync(cancellationToken), retryFilter, cancellationToken); CurrentActivationAddress = GrainAddress.NewActivationAddress(MessageCenter.MyAddress, _localClientDetails.ClientId.GrainId); this.gatewayObserver = new ClientGatewayObserver(gatewayManager); - this.InternalGrainFactory.CreateObjectReference(this.gatewayObserver!); + this.InternalGrainFactory.CreateObjectReference(this.gatewayObserver); await ExecuteWithRetries( - _manifestProvider!.StartAsync, + _manifestProvider.StartAsync, retryFilter, cancellationToken); - static async Task ExecuteWithRetries(Func task, IClientConnectionRetryFilter? retryFilter, CancellationToken cancellationToken) + static async Task ExecuteWithRetries(Func task, IClientConnectionRetryFilter retryFilter, CancellationToken cancellationToken) { do { @@ -241,7 +236,7 @@ private void HandleMessage(Message message) case Message.Directions.OneWay: case Message.Directions.Request: { - this.localObjects!.Dispatch(message); + this.localObjects.Dispatch(message); break; } default: @@ -253,19 +248,19 @@ private void HandleMessage(Message message) public void SendResponse(Message request, Response response) { ThrowIfDisposed(); - var message = this.messageFactory!.CreateResponseMessage(request); + var message = this.messageFactory.CreateResponseMessage(request); OrleansOutsideRuntimeClientEvent.Instance.SendResponse(message); message.BodyObject = response; - MessageCenter!.SendMessage(message); + MessageCenter.SendMessage(message); } - public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options) + public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource context, InvokeMethodOptions options) { ThrowIfDisposed(); var cancellationToken = request.GetCancellationToken(); cancellationToken.ThrowIfCancellationRequested(); - var message = this.messageFactory!.CreateMessage(request, options); + var message = this.messageFactory.CreateMessage(request, options); OrleansOutsideRuntimeClientEvent.Instance.SendRequest(message); message.InterfaceType = target.InterfaceType; @@ -290,33 +285,17 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp if (!oneWay) { - var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments); - if (Volatile.Read(ref _isStopping) != 0) - { - callbackData.OnHostShutdown(); - return; - } - - callbacks.TryAdd(message.Id, callbackData); + var callbackData = new CallbackData(this.sharedCallbackData, context, message, _applicationRequestInstruments); callbackData.SubscribeForCancellation(cancellationToken); - - if (Volatile.Read(ref _isStopping) != 0) - { - callbackData.OnHostShutdown(); - return; - } + callbacks.TryAdd(message.Id, callbackData); } else { context?.Complete(); - if (Volatile.Read(ref _isStopping) != 0) - { - return; - } } LogSendingMessage(logger, message); - MessageCenter!.SendMessage(message); + MessageCenter.SendMessage(message); } public void ReceiveResponse(Message response) @@ -327,12 +306,12 @@ public void ReceiveResponse(Message response) if (response.Result is Message.ResponseTypes.Status) { - var status = (StatusResponse)response.BodyObject!; + var status = (StatusResponse)response.BodyObject; callbacks.TryGetValue(response.Id, out var callback); var request = callback?.Message; if (request is not null) { - callback!.OnStatusUpdate(status); + callback.OnStatusUpdate(status); if (status.Diagnostics != null && status.Diagnostics.Count > 0) { LogReceivedStatusUpdateForPendingRequest(logger, request, new(status.Diagnostics)); @@ -360,14 +339,14 @@ public void ReceiveResponse(Message response) return; } - CallbackData? callbackData; + CallbackData callbackData; var found = callbacks.TryRemove(response.Id, out callbackData); if (found) { // We need to import the RequestContext here as well. // Unfortunately, it is not enough, since CallContext.LogicalGetData will not flow "up" from task completion source into the resolved task. // RequestContextExtensions.Import(response.RequestContextData); - callbackData!.DoCallback(response); + callbackData.DoCallback(response); } else { @@ -404,7 +383,7 @@ public IAddressable CreateObjectReference(IAddressable obj) : ObserverGrainId.Create(_localClientDetails.ClientId); var reference = this.InternalGrainFactory.GetGrain(observerId.GrainId); - if (!localObjects!.TryRegister(obj, observerId)) + if (!localObjects.TryRegister(obj, observerId)) { throw new ArgumentException($"Failed to add new observer {reference} to localObjects collection.", "reference"); } @@ -424,7 +403,7 @@ public void DeleteObjectReference(IAddressable obj) throw new ArgumentException($"Reference {reference.GrainId} is not an observer reference"); } - if (!localObjects!.TryDeregister(observerId)) + if (!localObjects.TryDeregister(observerId)) { throw new ArgumentException("Reference is not associated with a local object.", "reference"); } @@ -434,10 +413,8 @@ public void Dispose() { if (this.disposing) return; this.disposing = true; - Volatile.Write(ref _isStopping, 1); Utils.SafeExecute(() => this.callbackTimer.Dispose()); - BreakOutstandingMessages(); Utils.SafeExecute(() => MessageCenter?.Dispose()); @@ -456,28 +433,13 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) } } - private void BreakOutstandingMessages() - { - foreach (var (_, callback) in callbacks) - { - try - { - callback.OnHostShutdown(); - } - catch (Exception exception) - { - LogErrorWhileProcessingCallbackExpiry(logger, exception); - } - } - } - public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) - => this.callbacks.Count(c => c.Value.Message.InterfaceType == grainInterfaceType); + => this.callbacks.CountWhere(c => c.Value.Message.InterfaceType == grainInterfaceType); /// public void NotifyClusterConnectionLost() { - foreach (var observer in _statusObservers!) + foreach (var observer in _statusObservers) { try { @@ -493,7 +455,7 @@ public void NotifyClusterConnectionLost() /// public void NotifyGatewayCountChanged(int currentNumberOfGateways, int previousNumberOfGateways) { - foreach (var observer in _statusObservers!) + foreach (var observer in _statusObservers) { try { diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index d605549c37e..d163eda1082 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -1,8 +1,6 @@ using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics; -using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.DependencyInjection; @@ -20,6 +18,7 @@ using Orleans.Storage; using static Orleans.Internal.StandardExtensions; +#nullable disable namespace Orleans.Runtime { /// @@ -31,25 +30,24 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa private readonly ILogger invokeExceptionLogger; private readonly ILoggerFactory loggerFactory; private readonly SiloMessagingOptions messagingOptions; - private readonly ConcurrentDictionary<(GrainId, CorrelationId), CallbackData> callbacks; + private readonly StripedCallbackDictionary callbacks; private readonly InterfaceToImplementationMappingCache interfaceToImplementationMapping; private readonly SharedCallbackData sharedCallbackData; private readonly SharedCallbackData systemSharedCallbackData; private readonly PeriodicTimer callbackTimer; - private int _isStopping; - private GrainLocator grainLocator = null!; - private MessageCenter messageCenter = null!; - private List grainCallFilters = null!; + private GrainLocator grainLocator; + private MessageCenter messageCenter; + private List grainCallFilters; private readonly DeepCopier _deepCopier; private readonly ApplicationRequestInstruments _applicationRequestInstruments; - private IGrainCallCancellationManager _cancellationManager = null!; - private HostedClient hostedClient = null!; + private IGrainCallCancellationManager _cancellationManager; + private HostedClient hostedClient; - private HostedClient HostedClient => this.hostedClient; + private HostedClient HostedClient => this.hostedClient ??= this.ServiceProvider.GetRequiredService(); private readonly MessageFactory messageFactory; - private IGrainReferenceRuntime grainReferenceRuntime = null!; - private Task? callbackTimerTask; + private IGrainReferenceRuntime grainReferenceRuntime; + private Task callbackTimerTask; private readonly MessagingTrace messagingTrace; private readonly DeepCopier responseCopier; @@ -64,7 +62,7 @@ public InsideRuntimeClient( GrainInterfaceTypeResolver interfaceIdResolver, GrainInterfaceTypeToGrainTypeResolver interfaceToTypeResolver, DeepCopier deepCopier, - [FromKeyedServices(TimeProviderNames.Messaging)] TimeProvider timeProvider, + TimeProvider timeProvider, InterfaceToImplementationMappingCache interfaceToImplementationMapping, OrleansInstruments orleansInstruments) { @@ -74,7 +72,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,20 +86,20 @@ public InsideRuntimeClient( var callbackDataLogger = loggerFactory.CreateLogger(); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), + msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, this.messagingOptions.WaitForCancellationAcknowledgement, - cancellationManager: null!); + cancellationManager: null); this.systemSharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), + msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, waitForCancellationAcknowledgement: this.messagingOptions.WaitForCancellationAcknowledgement, - cancellationManager: null!); + cancellationManager: null); } public IServiceProvider ServiceProvider { get; } @@ -112,30 +110,20 @@ public InsideRuntimeClient( public GrainFactory ConcreteGrainFactory { get; } - private GrainLocator GrainLocator => this.grainLocator; + private GrainLocator GrainLocator + => this.grainLocator ?? (this.grainLocator = this.ServiceProvider.GetRequiredService()); - private List GrainCallFilters => this.grainCallFilters; + private List GrainCallFilters + => this.grainCallFilters ??= new List(this.ServiceProvider.GetServices()); - private MessageCenter MessageCenter => this.messageCenter; + private MessageCenter MessageCenter => this.messageCenter ?? (this.messageCenter = this.ServiceProvider.GetRequiredService()); - public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime; - - internal void ConsumeServices() - { - this.grainLocator = this.ServiceProvider.GetRequiredService(); - this.grainCallFilters = new List(this.ServiceProvider.GetServices()); - this.messageCenter = this.ServiceProvider.GetRequiredService(); - this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService(); - this.hostedClient = this.ServiceProvider.GetRequiredService(); - _cancellationManager = this.ServiceProvider.GetRequiredService(); - sharedCallbackData.CancellationManager = _cancellationManager; - systemSharedCallbackData.CancellationManager = _cancellationManager; - } + public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime ?? (this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService()); public void SendRequest( GrainReference target, IInvokable request, - IResponseCompletionSource? context, + IResponseCompletionSource context, InvokeMethodOptions options) { var cancellationToken = request.GetCancellationToken(); @@ -149,7 +137,7 @@ public void SendRequest( if (message.SendingSilo == null) message.SendingSilo = MySilo; - IGrainContext? sendingActivation = RuntimeContext.Current; + IGrainContext sendingActivation = RuntimeContext.Current; if (sendingActivation == null) { @@ -182,37 +170,18 @@ public void SendRequest( } var oneWay = (options & InvokeMethodOptions.OneWay) != 0; - CallbackData? callbackData = null; if (!oneWay) { Debug.Assert(context is not null); // Register a callback for the request. - callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); - if (Volatile.Read(ref _isStopping) != 0) - { - callbackData.OnHostShutdown(); - return; - } - - callbacks.TryAdd((message.SendingGrain, message.Id), callbackData); + var callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); + callbacks.TryAdd(message.Id, callbackData); callbackData.SubscribeForCancellation(cancellationToken); } else { context?.Complete(); - if (Volatile.Read(ref _isStopping) != 0) - { - return; - } - } - - // Completing callbacks during shutdown can resume application code which issues follow-up - // calls. Reject those calls so that they cannot outlive the shutdown callback sweep. - if (Volatile.Read(ref _isStopping) != 0) - { - callbackData?.OnHostShutdown(); - return; } this.messagingTrace.OnSendRequest(message); @@ -236,9 +205,9 @@ public void SendResponse(Message request, Response response) /// /// UnRegister a callback. /// - private void UnregisterCallback(GrainId grainId, CorrelationId correlationId) + private void UnregisterCallback(CorrelationId correlationId) { - callbacks.TryRemove((grainId, correlationId), out _); + callbacks.TryRemove(correlationId, out _); } public void SniffIncomingMessage(Message message) @@ -315,13 +284,12 @@ public async Task Invoke(IGrainContext target, Message message) { var invoker = new GrainMethodInvoker(message, target, invokable, GrainCallFilters, this.interfaceToImplementationMapping, this.responseCopier); await invoker.Invoke(); - response = invoker.Response!; + response = invoker.Response; } else { response = await invokable.Invoke(); - // The copier preserves the null state of its input. - response = this.responseCopier.Copy(response)!; + response = this.responseCopier.Copy(response); } invokable.Dispose(); @@ -382,8 +350,7 @@ private void SafeSendResponse(Message message, Response response) { try { - // The copier preserves the null state of its input. - SendResponse(message, (Response)this._deepCopier.Copy(response)!); + SendResponse(message, (Response)this._deepCopier.Copy(response)); } catch (Exception exc) { @@ -415,17 +382,9 @@ private void SafeSendExceptionResponse(Message message, Exception ex) public void ReceiveResponse(Message message) { OrleansInsideRuntimeClientEvent.Instance.ReceiveResponse(message); - - var result = message.Result; - if (result != Message.ResponseTypes.Rejection && result != Message.ResponseTypes.Status) + if (message.Result is Message.ResponseTypes.Rejection) { - ProcessResponseCallback(message); - return; - } - - if (result is Message.ResponseTypes.Rejection) - { - if (!message.TargetSilo!.Matches(this.MySilo)) + if (!message.TargetSilo.Matches(this.MySilo)) { // gatewayed message - gateway back to sender LogTraceNoCallbackForRejection(this.logger, message); @@ -434,7 +393,7 @@ public void ReceiveResponse(Message message) } LogHandleMessage(this.logger, message); - var rejection = (RejectionResponse)message.BodyObject!; + var rejection = (RejectionResponse)message.BodyObject; switch (rejection.RejectionType) { case Message.RejectionTypes.Overloaded: @@ -457,18 +416,46 @@ public void ReceiveResponse(Message message) LogErrorUnsupportedRejectionType(this.logger, rejection.RejectionType); break; } - - ProcessResponseCallback(message); } - else + else if (message.Result == Message.ResponseTypes.Status) { - ProcessStatusResponse(message); + var status = (StatusResponse)message.BodyObject; + callbacks.TryGetValue(message.Id, out var callback); + var request = callback?.Message; + if (request is not null) + { + callback.OnStatusUpdate(status); + if (status.Diagnostics != null && status.Diagnostics.Count > 0) + { + LogInformationReceivedStatusUpdate(this.logger, request, status.Diagnostics); + } + } + else + { + if (messagingOptions.CancelUnknownRequestOnStatusUpdate) + { + // Cancel the call since the caller has abandoned it. + // Note that the target and sender arguments are swapped because this is a response to the original request. + _cancellationManager.SignalCancellation( + message.SendingSilo, + targetGrainId: message.SendingGrain, + sendingGrainId: message.TargetGrain, + messageId: message.Id); + } + + if (status.Diagnostics != null && status.Diagnostics.Count > 0 && logger.IsEnabled(LogLevel.Debug)) + { + var diagnosticsString = string.Join("\n", status.Diagnostics); + LogDebugReceivedStatusUpdateUnknownRequest(this.logger, message, diagnosticsString); + } + } + + return; } - } - private void ProcessResponseCallback(Message message) - { - if (callbacks.TryRemove((message.TargetGrain, message.Id), out var callbackData)) + CallbackData callbackData; + bool found = callbacks.TryRemove(message.Id, out callbackData); + if (found) { // 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. @@ -480,45 +467,6 @@ private void ProcessResponseCallback(Message message) } } - private void ProcessStatusResponse(Message message) - { - var status = (StatusResponse)message.BodyObject!; - callbacks.TryGetValue((message.TargetGrain, message.Id), out var callback); - var request = callback?.Message; - if (request is not null) - { - callback!.OnStatusUpdate(status); - if (status.Diagnostics is { Count: > 0 }) - { - LogInformationReceivedStatusUpdate(this.logger, request, status.Diagnostics); - } - - return; - } - - HandleUnknownStatusUpdate(message, status); - } - - private void HandleUnknownStatusUpdate(Message message, StatusResponse status) - { - if (messagingOptions.CancelUnknownRequestOnStatusUpdate) - { - // Cancel the call since the caller has abandoned it. - // Note that the target and sender arguments are swapped because this is a response to the original request. - _cancellationManager.SignalCancellation( - message.SendingSilo, - targetGrainId: message.SendingGrain, - sendingGrainId: message.TargetGrain, - messageId: message.Id); - } - - if (status.Diagnostics is { Count: > 0 } && logger.IsEnabled(LogLevel.Debug)) - { - var diagnosticsString = string.Join("\n", status.Diagnostics); - LogDebugReceivedStatusUpdateUnknownRequest(this.logger, message, diagnosticsString); - } - } - public string CurrentActivationIdentity => RuntimeContext.Current?.Address.ToString() ?? this.HostedClient.ToString(); public TimeProvider TimeProvider { get; } @@ -549,36 +497,13 @@ public void DeleteObjectReference(IAddressable obj) private async Task OnRuntimeInitializeStop(CancellationToken tc) { - Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); - // Once the silo is shutting down it can no longer receive responses, so any requests which - // are still outstanding will never complete. Fault them now so that in-flight grain calls - // observe a terminal result instead of hanging forever, which would otherwise deadlock grain - // deactivation and host disposal during an ungraceful shutdown. This must happen before - // waiting for the timer task since that wait observes the shutdown cancellation token. - BreakOutstandingMessages(); - if (this.callbackTimerTask is { } task) { await task.WaitAsync(tc); } } - private void BreakOutstandingMessages() - { - foreach (var (_, callback) in callbacks) - { - try - { - callback.OnHostShutdown(); - } - catch (Exception exception) - { - LogWarningWhileProcessingCallbackExpiry(this.logger, exception); - } - } - } - private Task OnRuntimeInitializeStart(CancellationToken tc) { var stopWatch = ValueStopwatch.StartNew(); @@ -611,12 +536,14 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) public void Participate(ISiloLifecycle lifecycle) { - ConsumeServices(); + _cancellationManager = this.ServiceProvider.GetRequiredService(); + sharedCallbackData.CancellationManager = _cancellationManager; + systemSharedCallbackData.CancellationManager = _cancellationManager; lifecycle.Subscribe(ServiceLifecycleStage.RuntimeInitialize, OnRuntimeInitializeStart, OnRuntimeInitializeStop); } public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) - => this.callbacks.Count(c => c.Value.Message.InterfaceType == grainInterfaceType); + => this.callbacks.CountWhere(c => c.Value.Message.InterfaceType == grainInterfaceType); private async Task MonitorCallbackExpiry() { From 3fa1619bff9935aa089c8fd4983735dc4d0d9682 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Tue, 18 Aug 2026 04:35:56 -0700 Subject: [PATCH 02/18] fix(runtime): preserve callback shutdown semantics Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Messaging/StripedCallbackDictionary.cs | 29 +++++++++++++++---- .../Runtime/OutsideRuntimeClient.cs | 7 ++++- .../Core/InsideRuntimeClient.cs | 7 ++++- 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 3e131d8dda9..57335b3593f 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -1,5 +1,6 @@ #nullable enable using System; +using System.Buffers; using System.Collections; using System.Collections.Generic; using System.Runtime.CompilerServices; @@ -186,7 +187,8 @@ public struct Enumerator : IEnumerator> { private readonly StripedCallbackDictionary _dictionary; private int _stripeIndex; - private List>? _currentSnapshot; + private KeyValuePair[]? _currentSnapshot; + private int _snapshotCount; private int _snapshotIndex; internal Enumerator(StripedCallbackDictionary dictionary) @@ -194,6 +196,7 @@ internal Enumerator(StripedCallbackDictionary dictionary) _dictionary = dictionary; _stripeIndex = -1; _currentSnapshot = null; + _snapshotCount = 0; _snapshotIndex = -1; } @@ -209,10 +212,12 @@ public bool MoveNext() if (_currentSnapshot != null) { _snapshotIndex++; - if (_snapshotIndex < _currentSnapshot.Count) + if (_snapshotIndex < _snapshotCount) { return true; } + + ReturnSnapshot(); } // Move to next stripe @@ -229,7 +234,12 @@ public bool MoveNext() { if (stripe.Dictionary.Count > 0) { - _currentSnapshot = new List>(stripe.Dictionary); + _currentSnapshot = ArrayPool>.Shared.Rent(stripe.Dictionary.Count); + _snapshotCount = 0; + foreach (var pair in stripe.Dictionary) + { + _currentSnapshot[_snapshotCount++] = pair; + } _snapshotIndex = -1; } else @@ -243,13 +253,20 @@ public bool MoveNext() public void Reset() { _stripeIndex = -1; - _currentSnapshot = null; + ReturnSnapshot(); _snapshotIndex = -1; } - public void Dispose() + public void Dispose() => ReturnSnapshot(); + + private void ReturnSnapshot() { - _currentSnapshot = null; + if (_currentSnapshot is { } snapshot) + { + ArrayPool>.Shared.Return(snapshot, clearArray: true); + _currentSnapshot = null; + _snapshotCount = 0; + } } } } diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 38f7a2a1c30..162465ae240 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -72,7 +72,7 @@ public OutsideRuntimeClient( IOptions clientMessagingOptions, MessagingTrace messagingTrace, IServiceProvider serviceProvider, - TimeProvider timeProvider, + [FromKeyedServices(TimeProviderNames.Messaging)] TimeProvider timeProvider, InterfaceToImplementationMappingCache interfaceToImplementationMapping, OrleansInstruments orleansInstruments) { @@ -155,6 +155,11 @@ public async Task StartAsync(CancellationToken cancellationToken) public async Task StopAsync(CancellationToken cancellationToken) { + foreach (var (_, callback) in callbacks) + { + callback.OnHostShutdown(); + } + this.callbackTimer.Dispose(); if (this.callbackTimerTask is { } task) diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index d163eda1082..995447c2c8b 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -62,7 +62,7 @@ public InsideRuntimeClient( GrainInterfaceTypeResolver interfaceIdResolver, GrainInterfaceTypeToGrainTypeResolver interfaceToTypeResolver, DeepCopier deepCopier, - TimeProvider timeProvider, + [FromKeyedServices(TimeProviderNames.Messaging)] TimeProvider timeProvider, InterfaceToImplementationMappingCache interfaceToImplementationMapping, OrleansInstruments orleansInstruments) { @@ -497,6 +497,11 @@ public void DeleteObjectReference(IAddressable obj) private async Task OnRuntimeInitializeStop(CancellationToken tc) { + foreach (var (_, callback) in callbacks) + { + callback.OnHostShutdown(); + } + this.callbackTimer.Dispose(); if (this.callbackTimerTask is { } task) { From 6c287f2632ce20313c6e38086289248d7b8c92b4 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Tue, 18 Aug 2026 05:00:37 -0700 Subject: [PATCH 03/18] test(runtime): cover striped callback dictionary Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../StripedCallbackDictionaryTests.cs | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs new file mode 100644 index 00000000000..674edadab6d --- /dev/null +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -0,0 +1,63 @@ +using Orleans.Runtime; +using Xunit; + +namespace Tester; + +[TestSuite("BVT")] +[TestProvider("None")] +[TestCategory("BVT")] +public class StripedCallbackDictionaryTests +{ + [Fact] + public void CorrelationIdsRetainStripeIndex() + { + for (var stripe = 0; stripe < StripedCallbackDictionary.StripeCount; stripe++) + { + var id = StripedCallbackDictionary.CreateCorrelationId(42, stripe); + Assert.Equal(stripe, StripedCallbackDictionary.GetStripeIndex(id)); + } + } + + [Fact] + public void AddGetAndRemovePreserveValue() + { + var dictionary = new StripedCallbackDictionary(); + var id = StripedCallbackDictionary.CreateCorrelationId(42, 7); + + 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 EnumerationReturnsSnapshotValues() + { + var dictionary = new StripedCallbackDictionary(); + for (var i = 0; i < 32; i++) + { + var id = StripedCallbackDictionary.CreateCorrelationId(i, i); + Assert.True(dictionary.TryAdd(id, i)); + } + + Assert.Equal(Enumerable.Range(0, 32), dictionary.Select(pair => pair.Value).Order()); + } + + [Fact] + public void ConcurrentOperationsPreserveAllEntries() + { + var dictionary = new StripedCallbackDictionary(); + + Parallel.For(0, 10_000, i => + { + var id = StripedCallbackDictionary.CreateCorrelationId(i, i); + Assert.True(dictionary.TryAdd(id, i)); + }); + + Assert.Equal(10_000, dictionary.Count); + Assert.Equal(10_000, dictionary.CountWhere(static pair => pair.Value >= 0)); + } +} From cffc9bf1f17684d5631111aab57048d996128ee2 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 20 Aug 2026 07:46:05 -0700 Subject: [PATCH 04/18] fix(runtime): preserve callback shutdown boundaries Close callback registration before shutdown sweeps and eagerly capture silo services so late responses cannot resolve disposed providers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Runtime/OutsideRuntimeClient.cs | 43 +++++++++-- .../Core/InsideRuntimeClient.cs | 71 ++++++++++++++----- 2 files changed, 92 insertions(+), 22 deletions(-) diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 162465ae240..63b148ed193 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -28,6 +28,7 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus private readonly StripedCallbackDictionary callbacks; private InvokableObjectManager localObjects; + private int _isStopping; private bool disposing; private bool disposed; @@ -155,12 +156,9 @@ public async Task StartAsync(CancellationToken cancellationToken) public async Task StopAsync(CancellationToken cancellationToken) { - foreach (var (_, callback) in callbacks) - { - callback.OnHostShutdown(); - } - + Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); + BreakOutstandingMessages(); if (this.callbackTimerTask is { } task) { @@ -291,12 +289,28 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp if (!oneWay) { var callbackData = new CallbackData(this.sharedCallbackData, context, message, _applicationRequestInstruments); - callbackData.SubscribeForCancellation(cancellationToken); + if (Volatile.Read(ref _isStopping) != 0) + { + callbackData.OnHostShutdown(); + return; + } + callbacks.TryAdd(message.Id, callbackData); + callbackData.SubscribeForCancellation(cancellationToken); + + if (Volatile.Read(ref _isStopping) != 0) + { + callbackData.OnHostShutdown(); + return; + } } else { context?.Complete(); + if (Volatile.Read(ref _isStopping) != 0) + { + return; + } } LogSendingMessage(logger, message); @@ -418,8 +432,10 @@ public void Dispose() { if (this.disposing) return; this.disposing = true; + Volatile.Write(ref _isStopping, 1); Utils.SafeExecute(() => this.callbackTimer.Dispose()); + BreakOutstandingMessages(); Utils.SafeExecute(() => MessageCenter?.Dispose()); @@ -438,6 +454,21 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) } } + private void BreakOutstandingMessages() + { + foreach (var (_, callback) in callbacks) + { + try + { + callback.OnHostShutdown(); + } + catch (Exception exception) + { + LogErrorWhileProcessingCallbackExpiry(logger, exception); + } + } + } + public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) => this.callbacks.CountWhere(c => c.Value.Message.InterfaceType == grainInterfaceType); diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index 995447c2c8b..1bac2d9d53c 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -35,6 +35,7 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa private readonly SharedCallbackData sharedCallbackData; private readonly SharedCallbackData systemSharedCallbackData; private readonly PeriodicTimer callbackTimer; + private int _isStopping; private GrainLocator grainLocator; private MessageCenter messageCenter; @@ -44,7 +45,7 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa private IGrainCallCancellationManager _cancellationManager; private HostedClient hostedClient; - private HostedClient HostedClient => this.hostedClient ??= this.ServiceProvider.GetRequiredService(); + private HostedClient HostedClient => this.hostedClient; private readonly MessageFactory messageFactory; private IGrainReferenceRuntime grainReferenceRuntime; private Task callbackTimerTask; @@ -110,15 +111,25 @@ public InsideRuntimeClient( public GrainFactory ConcreteGrainFactory { get; } - private GrainLocator GrainLocator - => this.grainLocator ?? (this.grainLocator = this.ServiceProvider.GetRequiredService()); + private GrainLocator GrainLocator => this.grainLocator; - private List GrainCallFilters - => this.grainCallFilters ??= new List(this.ServiceProvider.GetServices()); + private List GrainCallFilters => this.grainCallFilters; - private MessageCenter MessageCenter => this.messageCenter ?? (this.messageCenter = this.ServiceProvider.GetRequiredService()); + private MessageCenter MessageCenter => this.messageCenter; - public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime ?? (this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService()); + public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime; + + internal void ConsumeServices() + { + this.grainLocator = this.ServiceProvider.GetRequiredService(); + this.grainCallFilters = new List(this.ServiceProvider.GetServices()); + this.messageCenter = this.ServiceProvider.GetRequiredService(); + this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService(); + this.hostedClient = this.ServiceProvider.GetRequiredService(); + _cancellationManager = this.ServiceProvider.GetRequiredService(); + sharedCallbackData.CancellationManager = _cancellationManager; + systemSharedCallbackData.CancellationManager = _cancellationManager; + } public void SendRequest( GrainReference target, @@ -170,18 +181,35 @@ public void SendRequest( } var oneWay = (options & InvokeMethodOptions.OneWay) != 0; + CallbackData callbackData = null; if (!oneWay) { Debug.Assert(context is not null); // Register a callback for the request. - var callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); + callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); + if (Volatile.Read(ref _isStopping) != 0) + { + callbackData.OnHostShutdown(); + return; + } + callbacks.TryAdd(message.Id, callbackData); callbackData.SubscribeForCancellation(cancellationToken); } else { context?.Complete(); + if (Volatile.Read(ref _isStopping) != 0) + { + return; + } + } + + if (Volatile.Read(ref _isStopping) != 0) + { + callbackData?.OnHostShutdown(); + return; } this.messagingTrace.OnSendRequest(message); @@ -497,18 +525,31 @@ public void DeleteObjectReference(IAddressable obj) private async Task OnRuntimeInitializeStop(CancellationToken tc) { - foreach (var (_, callback) in callbacks) - { - callback.OnHostShutdown(); - } - + Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); + BreakOutstandingMessages(); + if (this.callbackTimerTask is { } task) { await task.WaitAsync(tc); } } + private void BreakOutstandingMessages() + { + foreach (var (_, callback) in callbacks) + { + try + { + callback.OnHostShutdown(); + } + catch (Exception exception) + { + LogWarningWhileProcessingCallbackExpiry(this.logger, exception); + } + } + } + private Task OnRuntimeInitializeStart(CancellationToken tc) { var stopWatch = ValueStopwatch.StartNew(); @@ -541,9 +582,7 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) public void Participate(ISiloLifecycle lifecycle) { - _cancellationManager = this.ServiceProvider.GetRequiredService(); - sharedCallbackData.CancellationManager = _cancellationManager; - systemSharedCallbackData.CancellationManager = _cancellationManager; + ConsumeServices(); lifecycle.Subscribe(ServiceLifecycleStage.RuntimeInitialize, OnRuntimeInitializeStart, OnRuntimeInitializeStop); } From a123c103e82448521a8c267ff895d99b355e948f Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 20 Aug 2026 08:55:46 -0700 Subject: [PATCH 05/18] fix(runtime): retain nullable callback invariants Port striped callback storage onto the current nullable-safe runtime clients and clear pooled snapshots only when their entries contain references. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Messaging/StripedCallbackDictionary.cs | 10 +- .../Runtime/OutsideRuntimeClient.cs | 65 +++++---- .../Core/InsideRuntimeClient.cs | 138 +++++++++++------- 3 files changed, 123 insertions(+), 90 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 57335b3593f..e283892294a 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -3,6 +3,7 @@ using System.Buffers; using System.Collections; using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; namespace Orleans.Runtime; @@ -13,6 +14,7 @@ namespace Orleans.Runtime; /// /// The type of values stored in the dictionary. internal sealed class StripedCallbackDictionary : IEnumerable> + where TValue : notnull { /// /// The number of bits used to identify the stripe (stored in the upper bits of the CorrelationId). @@ -104,7 +106,7 @@ public bool TryAdd(CorrelationId key, TValue value) /// Attempts to get the value associated with the specified key. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryGetValue(CorrelationId key, out TValue? value) + public bool TryGetValue(CorrelationId key, [NotNullWhen(true)] out TValue? value) { var stripe = GetStripe(key); lock (stripe.Lock) @@ -117,7 +119,7 @@ public bool TryGetValue(CorrelationId key, out TValue? value) /// Attempts to remove the value with the specified key. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryRemove(CorrelationId key, out TValue? value) + public bool TryRemove(CorrelationId key, [NotNullWhen(true)] out TValue? value) { var stripe = GetStripe(key); lock (stripe.Lock) @@ -263,7 +265,9 @@ private void ReturnSnapshot() { if (_currentSnapshot is { } snapshot) { - ArrayPool>.Shared.Return(snapshot, clearArray: true); + ArrayPool>.Shared.Return( + snapshot, + clearArray: RuntimeHelpers.IsReferenceOrContainsReferences>()); _currentSnapshot = null; _snapshotCount = 0; } diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 63b148ed193..5663eacdec2 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -1,5 +1,4 @@ using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Threading; @@ -16,7 +15,6 @@ using Orleans.Serialization.Invocation; using static Orleans.Internal.StandardExtensions; -#nullable disable namespace Orleans { internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClusterConnectionStatusListener @@ -27,7 +25,7 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus private readonly ClientMessagingOptions clientMessagingOptions; private readonly StripedCallbackDictionary callbacks; - private InvokableObjectManager localObjects; + private InvokableObjectManager? localObjects; private int _isStopping; private bool disposing; private bool disposed; @@ -35,35 +33,35 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus private readonly MessagingTrace messagingTrace; private readonly InterfaceToImplementationMappingCache _interfaceToImplementationMapping; private readonly ApplicationRequestInstruments _applicationRequestInstruments; - private IGrainCallCancellationManager _cancellationManager; - private IClusterConnectionStatusObserver[] _statusObservers; + private IGrainCallCancellationManager? _cancellationManager; + private IClusterConnectionStatusObserver[]? _statusObservers; - public IInternalGrainFactory InternalGrainFactory { get; private set; } + public IInternalGrainFactory InternalGrainFactory { get; private set; } = null!; - private ClientClusterManifestProvider _manifestProvider; - private MessageFactory messageFactory; + private ClientClusterManifestProvider? _manifestProvider; + private MessageFactory? messageFactory; private readonly LocalClientDetails _localClientDetails; private readonly ILoggerFactory loggerFactory; private readonly SharedCallbackData sharedCallbackData; private readonly PeriodicTimer callbackTimer; - private Task callbackTimerTask; + private Task? callbackTimerTask; public GrainAddress CurrentActivationAddress { get; private set; - } - public ClientGatewayObserver gatewayObserver { get; private set; } + } = null!; + public ClientGatewayObserver? gatewayObserver { get; private set; } public string CurrentActivationIdentity { get { return CurrentActivationAddress.ToString(); } } - public IGrainReferenceRuntime GrainReferenceRuntime { get; private set; } + public IGrainReferenceRuntime GrainReferenceRuntime { get; private set; } = null!; - internal ClientMessageCenter MessageCenter { get; private set; } + internal ClientMessageCenter? MessageCenter { get; private set; } [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "MessageCenter is IDisposable but cannot call Dispose yet as it lives past the end of this method call.")] @@ -158,6 +156,9 @@ public async Task StopAsync(CancellationToken cancellationToken) { Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); + + // Fault callbacks before any cancellation-sensitive waits. Completing them can resume code + // which issues follow-up calls, so request admission must already be closed. BreakOutstandingMessages(); if (this.callbackTimerTask is { } task) @@ -192,20 +193,20 @@ await ExecuteWithRetries( MessageCenter = ActivatorUtilities.CreateInstance(this.ServiceProvider); MessageCenter.RegisterLocalMessageHandler(this.HandleMessage); await ExecuteWithRetries( - async () => await MessageCenter.StartAsync(cancellationToken), + async () => await MessageCenter!.StartAsync(cancellationToken), retryFilter, cancellationToken); CurrentActivationAddress = GrainAddress.NewActivationAddress(MessageCenter.MyAddress, _localClientDetails.ClientId.GrainId); this.gatewayObserver = new ClientGatewayObserver(gatewayManager); - this.InternalGrainFactory.CreateObjectReference(this.gatewayObserver); + this.InternalGrainFactory.CreateObjectReference(this.gatewayObserver!); await ExecuteWithRetries( - _manifestProvider.StartAsync, + _manifestProvider!.StartAsync, retryFilter, cancellationToken); - static async Task ExecuteWithRetries(Func task, IClientConnectionRetryFilter retryFilter, CancellationToken cancellationToken) + static async Task ExecuteWithRetries(Func task, IClientConnectionRetryFilter? retryFilter, CancellationToken cancellationToken) { do { @@ -239,7 +240,7 @@ private void HandleMessage(Message message) case Message.Directions.OneWay: case Message.Directions.Request: { - this.localObjects.Dispatch(message); + this.localObjects!.Dispatch(message); break; } default: @@ -251,19 +252,19 @@ private void HandleMessage(Message message) public void SendResponse(Message request, Response response) { ThrowIfDisposed(); - var message = this.messageFactory.CreateResponseMessage(request); + var message = this.messageFactory!.CreateResponseMessage(request); OrleansOutsideRuntimeClientEvent.Instance.SendResponse(message); message.BodyObject = response; - MessageCenter.SendMessage(message); + MessageCenter!.SendMessage(message); } - public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource context, InvokeMethodOptions options) + public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options) { ThrowIfDisposed(); var cancellationToken = request.GetCancellationToken(); cancellationToken.ThrowIfCancellationRequested(); - var message = this.messageFactory.CreateMessage(request, options); + var message = this.messageFactory!.CreateMessage(request, options); OrleansOutsideRuntimeClientEvent.Instance.SendRequest(message); message.InterfaceType = target.InterfaceType; @@ -288,7 +289,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp if (!oneWay) { - var callbackData = new CallbackData(this.sharedCallbackData, context, message, _applicationRequestInstruments); + var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments); if (Volatile.Read(ref _isStopping) != 0) { callbackData.OnHostShutdown(); @@ -314,7 +315,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp } LogSendingMessage(logger, message); - MessageCenter.SendMessage(message); + MessageCenter!.SendMessage(message); } public void ReceiveResponse(Message response) @@ -325,12 +326,12 @@ public void ReceiveResponse(Message response) if (response.Result is Message.ResponseTypes.Status) { - var status = (StatusResponse)response.BodyObject; + var status = (StatusResponse)response.BodyObject!; callbacks.TryGetValue(response.Id, out var callback); var request = callback?.Message; if (request is not null) { - callback.OnStatusUpdate(status); + callback!.OnStatusUpdate(status); if (status.Diagnostics != null && status.Diagnostics.Count > 0) { LogReceivedStatusUpdateForPendingRequest(logger, request, new(status.Diagnostics)); @@ -358,14 +359,14 @@ public void ReceiveResponse(Message response) return; } - CallbackData callbackData; + CallbackData? callbackData; var found = callbacks.TryRemove(response.Id, out callbackData); if (found) { // We need to import the RequestContext here as well. // Unfortunately, it is not enough, since CallContext.LogicalGetData will not flow "up" from task completion source into the resolved task. // RequestContextExtensions.Import(response.RequestContextData); - callbackData.DoCallback(response); + callbackData!.DoCallback(response); } else { @@ -402,7 +403,7 @@ public IAddressable CreateObjectReference(IAddressable obj) : ObserverGrainId.Create(_localClientDetails.ClientId); var reference = this.InternalGrainFactory.GetGrain(observerId.GrainId); - if (!localObjects.TryRegister(obj, observerId)) + if (!localObjects!.TryRegister(obj, observerId)) { throw new ArgumentException($"Failed to add new observer {reference} to localObjects collection.", "reference"); } @@ -422,7 +423,7 @@ public void DeleteObjectReference(IAddressable obj) throw new ArgumentException($"Reference {reference.GrainId} is not an observer reference"); } - if (!localObjects.TryDeregister(observerId)) + if (!localObjects!.TryDeregister(observerId)) { throw new ArgumentException("Reference is not associated with a local object.", "reference"); } @@ -475,7 +476,7 @@ public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) /// public void NotifyClusterConnectionLost() { - foreach (var observer in _statusObservers) + foreach (var observer in _statusObservers!) { try { @@ -491,7 +492,7 @@ public void NotifyClusterConnectionLost() /// public void NotifyGatewayCountChanged(int currentNumberOfGateways, int previousNumberOfGateways) { - foreach (var observer in _statusObservers) + foreach (var observer in _statusObservers!) { try { diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index 1bac2d9d53c..03554b05833 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Generic; using System.Diagnostics; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.DependencyInjection; @@ -18,7 +19,6 @@ using Orleans.Storage; using static Orleans.Internal.StandardExtensions; -#nullable disable namespace Orleans.Runtime { /// @@ -37,18 +37,18 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa private readonly PeriodicTimer callbackTimer; private int _isStopping; - private GrainLocator grainLocator; - private MessageCenter messageCenter; - private List grainCallFilters; + private GrainLocator grainLocator = null!; + private MessageCenter messageCenter = null!; + private List grainCallFilters = null!; private readonly DeepCopier _deepCopier; private readonly ApplicationRequestInstruments _applicationRequestInstruments; - private IGrainCallCancellationManager _cancellationManager; - private HostedClient hostedClient; + private IGrainCallCancellationManager _cancellationManager = null!; + private HostedClient hostedClient = null!; private HostedClient HostedClient => this.hostedClient; private readonly MessageFactory messageFactory; - private IGrainReferenceRuntime grainReferenceRuntime; - private Task callbackTimerTask; + private IGrainReferenceRuntime grainReferenceRuntime = null!; + private Task? callbackTimerTask; private readonly MessagingTrace messagingTrace; private readonly DeepCopier responseCopier; @@ -92,7 +92,7 @@ public InsideRuntimeClient( this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, this.messagingOptions.WaitForCancellationAcknowledgement, - cancellationManager: null); + cancellationManager: null!); this.systemSharedCallbackData = new SharedCallbackData( msg => this.UnregisterCallback(msg.Id), @@ -100,7 +100,7 @@ public InsideRuntimeClient( this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, waitForCancellationAcknowledgement: this.messagingOptions.WaitForCancellationAcknowledgement, - cancellationManager: null); + cancellationManager: null!); } public IServiceProvider ServiceProvider { get; } @@ -134,7 +134,7 @@ internal void ConsumeServices() public void SendRequest( GrainReference target, IInvokable request, - IResponseCompletionSource context, + IResponseCompletionSource? context, InvokeMethodOptions options) { var cancellationToken = request.GetCancellationToken(); @@ -148,7 +148,7 @@ public void SendRequest( if (message.SendingSilo == null) message.SendingSilo = MySilo; - IGrainContext sendingActivation = RuntimeContext.Current; + IGrainContext? sendingActivation = RuntimeContext.Current; if (sendingActivation == null) { @@ -181,7 +181,7 @@ public void SendRequest( } var oneWay = (options & InvokeMethodOptions.OneWay) != 0; - CallbackData callbackData = null; + CallbackData? callbackData = null; if (!oneWay) { Debug.Assert(context is not null); @@ -206,6 +206,8 @@ public void SendRequest( } } + // Completing callbacks during shutdown can resume application code which issues follow-up + // calls. Reject those calls so that they cannot outlive the shutdown callback sweep. if (Volatile.Read(ref _isStopping) != 0) { callbackData?.OnHostShutdown(); @@ -312,12 +314,13 @@ public async Task Invoke(IGrainContext target, Message message) { var invoker = new GrainMethodInvoker(message, target, invokable, GrainCallFilters, this.interfaceToImplementationMapping, this.responseCopier); await invoker.Invoke(); - response = invoker.Response; + response = invoker.Response!; } else { response = await invokable.Invoke(); - response = this.responseCopier.Copy(response); + // The copier preserves the null state of its input. + response = this.responseCopier.Copy(response)!; } invokable.Dispose(); @@ -378,7 +381,8 @@ private void SafeSendResponse(Message message, Response response) { try { - SendResponse(message, (Response)this._deepCopier.Copy(response)); + // The copier preserves the null state of its input. + SendResponse(message, (Response)this._deepCopier.Copy(response)!); } catch (Exception exc) { @@ -410,9 +414,17 @@ private void SafeSendExceptionResponse(Message message, Exception ex) public void ReceiveResponse(Message message) { OrleansInsideRuntimeClientEvent.Instance.ReceiveResponse(message); - if (message.Result is Message.ResponseTypes.Rejection) + + var result = message.Result; + if (result != Message.ResponseTypes.Rejection && result != Message.ResponseTypes.Status) + { + ProcessResponseCallback(message); + return; + } + + if (result is Message.ResponseTypes.Rejection) { - if (!message.TargetSilo.Matches(this.MySilo)) + if (!message.TargetSilo!.Matches(this.MySilo)) { // gatewayed message - gateway back to sender LogTraceNoCallbackForRejection(this.logger, message); @@ -421,7 +433,7 @@ public void ReceiveResponse(Message message) } LogHandleMessage(this.logger, message); - var rejection = (RejectionResponse)message.BodyObject; + var rejection = (RejectionResponse)message.BodyObject!; switch (rejection.RejectionType) { case Message.RejectionTypes.Overloaded: @@ -444,46 +456,18 @@ public void ReceiveResponse(Message message) LogErrorUnsupportedRejectionType(this.logger, rejection.RejectionType); break; } + + ProcessResponseCallback(message); } - else if (message.Result == Message.ResponseTypes.Status) + else { - var status = (StatusResponse)message.BodyObject; - callbacks.TryGetValue(message.Id, out var callback); - var request = callback?.Message; - if (request is not null) - { - callback.OnStatusUpdate(status); - if (status.Diagnostics != null && status.Diagnostics.Count > 0) - { - LogInformationReceivedStatusUpdate(this.logger, request, status.Diagnostics); - } - } - else - { - if (messagingOptions.CancelUnknownRequestOnStatusUpdate) - { - // Cancel the call since the caller has abandoned it. - // Note that the target and sender arguments are swapped because this is a response to the original request. - _cancellationManager.SignalCancellation( - message.SendingSilo, - targetGrainId: message.SendingGrain, - sendingGrainId: message.TargetGrain, - messageId: message.Id); - } - - if (status.Diagnostics != null && status.Diagnostics.Count > 0 && logger.IsEnabled(LogLevel.Debug)) - { - var diagnosticsString = string.Join("\n", status.Diagnostics); - LogDebugReceivedStatusUpdateUnknownRequest(this.logger, message, diagnosticsString); - } - } - - return; + ProcessStatusResponse(message); } + } - CallbackData callbackData; - bool found = callbacks.TryRemove(message.Id, out callbackData); - if (found) + private void ProcessResponseCallback(Message message) + { + 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. @@ -495,6 +479,45 @@ public void ReceiveResponse(Message message) } } + private void ProcessStatusResponse(Message message) + { + var status = (StatusResponse)message.BodyObject!; + callbacks.TryGetValue(message.Id, out var callback); + var request = callback?.Message; + if (request is not null) + { + callback!.OnStatusUpdate(status); + if (status.Diagnostics is { Count: > 0 }) + { + LogInformationReceivedStatusUpdate(this.logger, request, status.Diagnostics); + } + + return; + } + + HandleUnknownStatusUpdate(message, status); + } + + private void HandleUnknownStatusUpdate(Message message, StatusResponse status) + { + if (messagingOptions.CancelUnknownRequestOnStatusUpdate) + { + // Cancel the call since the caller has abandoned it. + // Note that the target and sender arguments are swapped because this is a response to the original request. + _cancellationManager.SignalCancellation( + message.SendingSilo, + targetGrainId: message.SendingGrain, + sendingGrainId: message.TargetGrain, + messageId: message.Id); + } + + if (status.Diagnostics is { Count: > 0 } && logger.IsEnabled(LogLevel.Debug)) + { + var diagnosticsString = string.Join("\n", status.Diagnostics); + LogDebugReceivedStatusUpdateUnknownRequest(this.logger, message, diagnosticsString); + } + } + public string CurrentActivationIdentity => RuntimeContext.Current?.Address.ToString() ?? this.HostedClient.ToString(); public TimeProvider TimeProvider { get; } @@ -527,6 +550,11 @@ private async Task OnRuntimeInitializeStop(CancellationToken tc) { Volatile.Write(ref _isStopping, 1); this.callbackTimer.Dispose(); + // Once the silo is shutting down it can no longer receive responses, so any requests which + // are still outstanding will never complete. Fault them now so that in-flight grain calls + // observe a terminal result instead of hanging forever, which would otherwise deadlock grain + // deactivation and host disposal during an ungraceful shutdown. This must happen before + // waiting for the timer task since that wait observes the shutdown cancellation token. BreakOutstandingMessages(); if (this.callbackTimerTask is { } task) From e2450e24282da623a6ed3ba94afbb2f58eec1a38 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 21 Aug 2026 12:40:30 -0700 Subject: [PATCH 06/18] fix(runtime): prevent pooled callback snapshot reuse --- src/Orleans.Core/Messaging/StripedCallbackDictionary.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index e283892294a..87e4e4a9254 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -185,7 +185,7 @@ private sealed class Stripe public readonly Dictionary Dictionary = new(); } - public struct Enumerator : IEnumerator> + public sealed class Enumerator : IEnumerator> { private readonly StripedCallbackDictionary _dictionary; private int _stripeIndex; From f58d1db760216ac7eba23b0ab8d0ba8c9cc46f8b Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Sun, 23 Aug 2026 04:07:59 -0700 Subject: [PATCH 07/18] fix(runtime): preserve callback ownership --- src/Orleans.Core/Messaging/CorrelationId.cs | 2 +- src/Orleans.Core/Messaging/MessageFactory.cs | 3 +- .../Messaging/StripedCallbackDictionary.cs | 103 +++++---------- .../Runtime/OutsideRuntimeClient.cs | 22 ++-- .../Core/InsideRuntimeClient.cs | 24 ++-- .../StripedCallbackDictionaryTests.cs | 118 +++++++++++++++--- 6 files changed, 158 insertions(+), 114 deletions(-) diff --git a/src/Orleans.Core/Messaging/CorrelationId.cs b/src/Orleans.Core/Messaging/CorrelationId.cs index c3a7b055dd3..b40167a4fca 100644 --- a/src/Orleans.Core/Messaging/CorrelationId.cs +++ b/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); diff --git a/src/Orleans.Core/Messaging/MessageFactory.cs b/src/Orleans.Core/Messaging/MessageFactory.cs index 3cf1e429c26..1f815944ea9 100644 --- a/src/Orleans.Core/Messaging/MessageFactory.cs +++ b/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.GetCurrentThreadStripeIndex(); - return StripedCallbackDictionary.CreateCorrelationId(unchecked((long)id), stripeIndex); + return new CorrelationId(unchecked((long)id)); } public Message CreateResponseMessage(Message request) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 87e4e4a9254..180d1e654fd 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -10,31 +10,20 @@ namespace Orleans.Runtime; /// /// 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. /// /// The type of values stored in the dictionary. -internal sealed class StripedCallbackDictionary : IEnumerable> +internal sealed class StripedCallbackDictionary : IEnumerable where TValue : notnull { - /// - /// The number of bits used to identify the stripe (stored in the upper bits of the CorrelationId). - /// - public const int StripeBits = 7; - - /// - /// The number of stripes (must be a power of 2). - /// - public const int StripeCount = 1 << StripeBits; // 128 stripes - - /// - /// Mask to extract the stripe index from the upper bits. - /// - 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; /// - /// The shift amount to move the stripe bits to the lowest position. + /// The number of stripes. /// - private const int StripeShift = 64 - StripeBits; + public const int StripeCount = 1 << StripeBits; private readonly Stripe[] _stripes; @@ -48,40 +37,14 @@ public StripedCallbackDictionary() } /// - /// Encodes a stripe index into the upper bits of a base value to create a CorrelationId. - /// - /// The base value (e.g., from an incrementing counter XORed with a seed). - /// The stripe index (typically derived from thread id). - /// A CorrelationId with the stripe encoded in the upper bits. - [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); - } - - /// - /// Extracts the stripe index from a CorrelationId. + /// Computes the stripe index for a correlation id. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] public static int GetStripeIndex(CorrelationId correlationId) - { - return (int)((correlationId.ToInt64() & StripeMask) >>> StripeShift); - } + => (int)(unchecked((ulong)correlationId.ToInt64() * HashFactor) >> (64 - StripeBits)); /// - /// Gets the stripe index for the current thread. Use this when creating new CorrelationIds. - /// - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public static int GetCurrentThreadStripeIndex() - { - return Environment.CurrentManagedThreadId & (StripeCount - 1); - } - - /// - /// Gets the stripe for the given correlation id. + /// Gets the stripe for the given callback id. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] private Stripe GetStripe(CorrelationId correlationId) @@ -93,12 +56,12 @@ private Stripe GetStripe(CorrelationId correlationId) /// Attempts to add the specified key and value to the dictionary. /// [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); } } @@ -106,12 +69,12 @@ public bool TryAdd(CorrelationId key, TValue value) /// Attempts to get the value associated with the specified key. /// [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); } } @@ -119,12 +82,12 @@ public bool TryGetValue(CorrelationId key, [NotNullWhen(true)] out TValue? value /// Attempts to remove the value with the specified key. /// [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 /// /// Counts items matching a predicate across all stripes. /// - public int CountWhere(Func, bool> predicate) + public int CountWhere(Func 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, bool> predicate) /// public Enumerator GetEnumerator() => new(this); - IEnumerator> IEnumerable>.GetEnumerator() => GetEnumerator(); + IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); private sealed class Stripe { public readonly object Lock = new(); - public readonly Dictionary Dictionary = new(); + public readonly Dictionary Dictionary = new(); } - public sealed class Enumerator : IEnumerator> + private readonly record struct CallbackKey(GrainId Owner, CorrelationId Id); + + public sealed class Enumerator : IEnumerator { private readonly StripedCallbackDictionary _dictionary; private int _stripeIndex; - private KeyValuePair[]? _currentSnapshot; + private TValue[]? _currentSnapshot; private int _snapshotCount; private int _snapshotIndex; @@ -202,7 +167,7 @@ internal Enumerator(StripedCallbackDictionary dictionary) _snapshotIndex = -1; } - public KeyValuePair 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>.Shared.Rent(stripe.Dictionary.Count); + _currentSnapshot = ArrayPool.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>.Shared.Return( + ArrayPool.Shared.Return( snapshot, - clearArray: RuntimeHelpers.IsReferenceOrContainsReferences>()); + clearArray: RuntimeHelpers.IsReferenceOrContainsReferences()); _currentSnapshot = null; _snapshotCount = 0; } diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 5663eacdec2..1ae4ebb5a0b 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/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(), 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); /// 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) { diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index 03554b05833..9c898cc0829 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -87,7 +87,7 @@ public InsideRuntimeClient( var callbackDataLogger = loggerFactory.CreateLogger(); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.Id), + msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), callbackDataLogger, this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, @@ -95,7 +95,7 @@ public InsideRuntimeClient( cancellationManager: null!); this.systemSharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.Id), + msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), callbackDataLogger, this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, @@ -194,7 +194,7 @@ public void SendRequest( return; } - callbacks.TryAdd(message.Id, callbackData); + callbacks.TryAdd(message.SendingGrain, message.Id, callbackData); callbackData.SubscribeForCancellation(cancellationToken); } else @@ -235,9 +235,9 @@ public void SendResponse(Message request, Response response) /// /// UnRegister a callback. /// - private void UnregisterCallback(CorrelationId correlationId) + private void UnregisterCallback(GrainId owner, CorrelationId correlationId) { - callbacks.TryRemove(correlationId, out _); + callbacks.TryRemove(owner, correlationId, out _); } public void SniffIncomingMessage(Message message) @@ -467,7 +467,7 @@ public void ReceiveResponse(Message message) private void ProcessResponseCallback(Message message) { - if (callbacks.TryRemove(message.Id, out var callbackData)) + if (callbacks.TryRemove(message.TargetGrain, 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. @@ -482,7 +482,7 @@ private void ProcessResponseCallback(Message message) private void ProcessStatusResponse(Message message) { var status = (StatusResponse)message.BodyObject!; - callbacks.TryGetValue(message.Id, out var callback); + callbacks.TryGetValue(message.TargetGrain, message.Id, out var callback); var request = callback?.Message; if (request is not null) { @@ -565,7 +565,7 @@ private async Task OnRuntimeInitializeStop(CancellationToken tc) private void BreakOutstandingMessages() { - foreach (var (_, callback) in callbacks) + foreach (var callback in callbacks) { try { @@ -601,9 +601,9 @@ 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(); } } } @@ -615,7 +615,7 @@ public void Participate(ISiloLifecycle lifecycle) } public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) - => this.callbacks.CountWhere(c => c.Value.Message.InterfaceType == grainInterfaceType); + => this.callbacks.CountWhere(c => c.Message.InterfaceType == grainInterfaceType); private async Task MonitorCallbackExpiry() { @@ -624,7 +624,7 @@ private async Task MonitorCallbackExpiry() try { var currentStopwatchTicks = ValueStopwatch.GetTimestamp(); - foreach (var (_, callback) in callbacks) + foreach (var callback in callbacks) { if (callback.IsCompleted) { diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index 674edadab6d..68e3ae00282 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -8,29 +8,54 @@ namespace Tester; [TestCategory("BVT")] public class StripedCallbackDictionaryTests { + private static readonly GrainId Owner = GrainId.Create("test", "owner"); + [Fact] - public void CorrelationIdsRetainStripeIndex() + public void CorrelationIdsDistributeAcrossStripesAtOverflowAndWithStride() { - for (var stripe = 0; stripe < StripedCallbackDictionary.StripeCount; stripe++) - { - var id = StripedCallbackDictionary.CreateCorrelationId(42, stripe); - Assert.Equal(stripe, StripedCallbackDictionary.GetStripeIndex(id)); - } + 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 = StripedCallbackDictionary.CreateCorrelationId(42, 7); + 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.True(dictionary.TryAdd(Owner, id, "value")); + Assert.False(dictionary.TryAdd(Owner, id, "duplicate")); + Assert.True(dictionary.TryGetValue(Owner, id, out var value)); Assert.Equal("value", value); - Assert.True(dictionary.TryRemove(id, out value)); + Assert.True(dictionary.TryRemove(Owner, id, out value)); + Assert.Equal("value", value); + Assert.False(dictionary.TryGetValue(Owner, id, out _)); + } + + [Fact] + public void CallbackOwnerIsPartOfTheKey() + { + var dictionary = new StripedCallbackDictionary(); + var otherOwner = GrainId.Create("test", "other-owner"); + var id = new CorrelationId(42); + + Assert.True(dictionary.TryAdd(Owner, id, "value")); + Assert.False(dictionary.TryGetValue(otherOwner, id, out _)); + Assert.False(dictionary.TryRemove(otherOwner, id, out _)); + Assert.True(dictionary.TryGetValue(Owner, id, out var value)); Assert.Equal("value", value); - Assert.False(dictionary.TryGetValue(id, out _)); } [Fact] @@ -39,25 +64,80 @@ public void EnumerationReturnsSnapshotValues() var dictionary = new StripedCallbackDictionary(); for (var i = 0; i < 32; i++) { - var id = StripedCallbackDictionary.CreateCorrelationId(i, i); - Assert.True(dictionary.TryAdd(id, i)); + var id = new CorrelationId(i); + Assert.True(dictionary.TryAdd(Owner, id, i)); } - Assert.Equal(Enumerable.Range(0, 32), dictionary.Select(pair => pair.Value).Order()); + Assert.Equal(Enumerable.Range(0, 32), dictionary.Order()); } [Fact] - public void ConcurrentOperationsPreserveAllEntries() + public void ConcurrentOperationsPreserveCountAndValues() { var dictionary = new StripedCallbackDictionary(); Parallel.For(0, 10_000, i => { - var id = StripedCallbackDictionary.CreateCorrelationId(i, i); - Assert.True(dictionary.TryAdd(id, i)); + var id = new CorrelationId(i); + Assert.True(dictionary.TryAdd(Owner, id, i)); + Assert.True(dictionary.TryGetValue(Owner, id, out var value)); + Assert.Equal(i, value); }); Assert.Equal(10_000, dictionary.Count); - Assert.Equal(10_000, dictionary.CountWhere(static pair => pair.Value >= 0)); + Assert.Equal(10_000, dictionary.CountWhere(static value => value >= 0)); + + Parallel.For(0, 10_000, i => + { + Assert.True(dictionary.TryRemove(Owner, 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(Owner, new CorrelationId(i), i)); + } + + Parallel.Invoke( + () => Parallel.For(0, count, i => + { + if (dictionary.TryGetValue(Owner, new CorrelationId(i), out var value)) + { + Assert.Equal(i, value); + } + }), + () => Parallel.For(0, count, i => + { + Assert.True(dictionary.TryRemove(Owner, new CorrelationId(i), out var value)); + Assert.Equal(i, value); + })); + + Assert.Equal(0, dictionary.Count); + } + + [Fact] + public void DisposingPartialEnumerationLeavesDictionaryUsable() + { + var dictionary = new StripedCallbackDictionary(); + for (var i = 0; i < 32; i++) + { + Assert.True(dictionary.TryAdd(Owner, new CorrelationId(i), i)); + } + + using (var enumerator = dictionary.GetEnumerator()) + { + Assert.True(enumerator.MoveNext()); + } + + Assert.True(dictionary.TryRemove(Owner, new CorrelationId(0), out var value)); + Assert.Equal(0, value); } } From bf1cb5a055e1431c0fa344741f99f21c93389f7c Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Sun, 23 Aug 2026 04:56:20 -0700 Subject: [PATCH 08/18] perf(runtime): remove callback scan allocation --- .../Messaging/StripedCallbackDictionary.cs | 120 +++++------------- .../Runtime/OutsideRuntimeClient.cs | 18 +-- .../Core/InsideRuntimeClient.cs | 16 +-- .../StripedCallbackDictionaryTests.cs | 31 ++++- 4 files changed, 72 insertions(+), 113 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 180d1e654fd..4ab6cc1278f 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -1,7 +1,6 @@ #nullable enable using System; using System.Buffers; -using System.Collections; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; @@ -13,7 +12,7 @@ namespace Orleans.Runtime; /// to reduce lock contention by hashing correlation ids across stripes. /// /// The type of values stored in the dictionary. -internal sealed class StripedCallbackDictionary : IEnumerable +internal sealed class StripedCallbackDictionary where TValue : notnull { private const int StripeBits = 7; @@ -133,109 +132,52 @@ public int CountWhere(Func predicate) } /// - /// Returns an enumerator that iterates through all items in all stripes. - /// Note: This takes a snapshot of each stripe under its lock. + /// Visits a snapshot of the values in each stripe. /// - public Enumerator GetEnumerator() => new(this); - - IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); - - IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); - - private sealed class Stripe - { - public readonly object Lock = new(); - public readonly Dictionary Dictionary = new(); - } - - private readonly record struct CallbackKey(GrainId Owner, CorrelationId Id); - - public sealed class Enumerator : IEnumerator + public void ForEach(TState state, Action action) { - private readonly StripedCallbackDictionary _dictionary; - private int _stripeIndex; - private TValue[]? _currentSnapshot; - private int _snapshotCount; - private int _snapshotIndex; - - internal Enumerator(StripedCallbackDictionary dictionary) - { - _dictionary = dictionary; - _stripeIndex = -1; - _currentSnapshot = null; - _snapshotCount = 0; - _snapshotIndex = -1; - } - - public TValue Current => _currentSnapshot![_snapshotIndex]; - - object IEnumerator.Current => Current; - - public bool MoveNext() + foreach (var stripe in _stripes) { - while (true) + TValue[]? snapshot = null; + var snapshotCount = 0; + try { - // Try to advance within current snapshot - if (_currentSnapshot != null) + lock (stripe.Lock) { - _snapshotIndex++; - if (_snapshotIndex < _snapshotCount) + if (stripe.Dictionary.Count == 0) { - return true; + continue; } - ReturnSnapshot(); + snapshot = ArrayPool.Shared.Rent(stripe.Dictionary.Count); + foreach (var value in stripe.Dictionary.Values) + { + snapshot[snapshotCount++] = value; + } } - // Move to next stripe - _stripeIndex++; - if (_stripeIndex >= _dictionary._stripes.Length) + for (var i = 0; i < snapshotCount; i++) { - _currentSnapshot = null; - return false; + action(snapshot[i], state); } - - // Take a snapshot of the next stripe - var stripe = _dictionary._stripes[_stripeIndex]; - lock (stripe.Lock) + } + finally + { + if (snapshot is not null) { - if (stripe.Dictionary.Count > 0) - { - _currentSnapshot = ArrayPool.Shared.Rent(stripe.Dictionary.Count); - _snapshotCount = 0; - foreach (var value in stripe.Dictionary.Values) - { - _currentSnapshot[_snapshotCount++] = value; - } - _snapshotIndex = -1; - } - else - { - _currentSnapshot = null; - } + ArrayPool.Shared.Return( + snapshot, + clearArray: RuntimeHelpers.IsReferenceOrContainsReferences()); } } } + } - public void Reset() - { - _stripeIndex = -1; - ReturnSnapshot(); - _snapshotIndex = -1; - } - - public void Dispose() => ReturnSnapshot(); - - private void ReturnSnapshot() - { - if (_currentSnapshot is { } snapshot) - { - ArrayPool.Shared.Return( - snapshot, - clearArray: RuntimeHelpers.IsReferenceOrContainsReferences()); - _currentSnapshot = null; - _snapshotCount = 0; - } - } + private sealed class Stripe + { + public readonly object Lock = new(); + public readonly Dictionary Dictionary = new(); } + + private readonly record struct CallbackKey(GrainId Owner, CorrelationId Id); } diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 1ae4ebb5a0b..2739f64274c 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -446,18 +446,18 @@ public void Dispose() public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) { - foreach (var callback in callbacks) + callbacks.ForEach(deadSilo, static (callback, deadSilo) => { if (deadSilo.Equals(callback.Message.TargetSilo)) { callback.OnTargetSiloFail(); } - } + }); } private void BreakOutstandingMessages() { - foreach (var callback in callbacks) + callbacks.ForEach(this, static (callback, self) => { try { @@ -465,9 +465,9 @@ private void BreakOutstandingMessages() } catch (Exception exception) { - LogErrorWhileProcessingCallbackExpiry(logger, exception); + LogErrorWhileProcessingCallbackExpiry(self.logger, exception); } - } + }); } public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) @@ -515,18 +515,18 @@ private async Task MonitorCallbackExpiry() try { var currentStopwatchTicks = ValueStopwatch.GetTimestamp(); - foreach (var callback in callbacks) + callbacks.ForEach((Self: this, CurrentStopwatchTicks: currentStopwatchTicks), static (callback, state) => { if (callback.IsCompleted) { - continue; + return; } - if (callback.IsExpired(currentStopwatchTicks)) + if (callback.IsExpired(state.CurrentStopwatchTicks)) { callback.OnTimeout(); } - } + }); } catch (Exception ex) { diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index 9c898cc0829..edb06397b64 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -565,7 +565,7 @@ private async Task OnRuntimeInitializeStop(CancellationToken tc) private void BreakOutstandingMessages() { - foreach (var callback in callbacks) + callbacks.ForEach(this, static (callback, self) => { try { @@ -573,9 +573,9 @@ private void BreakOutstandingMessages() } catch (Exception exception) { - LogWarningWhileProcessingCallbackExpiry(this.logger, exception); + LogWarningWhileProcessingCallbackExpiry(self.logger, exception); } - } + }); } private Task OnRuntimeInitializeStart(CancellationToken tc) @@ -599,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.Message.TargetSilo)) { callback.OnTargetSiloFail(); } - } + }); } public void Participate(ISiloLifecycle lifecycle) @@ -624,18 +624,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/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index 68e3ae00282..e239e9f3197 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -9,6 +9,7 @@ namespace Tester; public class StripedCallbackDictionaryTests { private static readonly GrainId Owner = GrainId.Create("test", "owner"); + private static readonly Action EmptyVisitor = static (_, _) => { }; [Fact] public void CorrelationIdsDistributeAcrossStripesAtOverflowAndWithStride() @@ -68,7 +69,10 @@ public void EnumerationReturnsSnapshotValues() Assert.True(dictionary.TryAdd(Owner, id, i)); } - Assert.Equal(Enumerable.Range(0, 32), dictionary.Order()); + var values = new List(); + dictionary.ForEach(values, static (value, values) => values.Add(value)); + + Assert.Equal(Enumerable.Range(0, 32), values.Order()); } [Fact] @@ -124,7 +128,7 @@ public void ConcurrentLookupAndRemovalRemainConsistent() } [Fact] - public void DisposingPartialEnumerationLeavesDictionaryUsable() + public void SnapshotVisitorAllowsValuesToRemoveThemselves() { var dictionary = new StripedCallbackDictionary(); for (var i = 0; i < 32; i++) @@ -132,12 +136,25 @@ public void DisposingPartialEnumerationLeavesDictionaryUsable() Assert.True(dictionary.TryAdd(Owner, new CorrelationId(i), i)); } - using (var enumerator = dictionary.GetEnumerator()) + dictionary.ForEach((Dictionary: dictionary, Owner), static (value, state) => { - Assert.True(enumerator.MoveNext()); - } + Assert.True(state.Dictionary.TryRemove(state.Owner, 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.True(dictionary.TryRemove(Owner, new CorrelationId(0), out var value)); - Assert.Equal(0, value); + Assert.Equal(0, allocated); } } From 7b435bd4b333a8376b9bb31d98c2e0292fe1f36a Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 09:39:27 -0700 Subject: [PATCH 09/18] perf(runtime): preserve client callback fast path --- .../Messaging/StripedCallbackDictionary.cs | 4 ++ .../Runtime/OutsideRuntimeClient.cs | 41 ++++++++++--------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 4ab6cc1278f..27a2702d0af 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -175,7 +175,11 @@ public void ForEach(TState state, Action action) 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/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 2739f64274c..bb86127f633 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Threading; @@ -24,7 +25,7 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus private readonly ILogger logger; private readonly ClientMessagingOptions clientMessagingOptions; - private readonly StripedCallbackDictionary callbacks; + private readonly ConcurrentDictionary callbacks; private InvokableObjectManager? localObjects; private int _isStopping; private bool disposing; @@ -83,7 +84,7 @@ public OutsideRuntimeClient( this.loggerFactory = loggerFactory; this.messagingTrace = messagingTrace; this.logger = loggerFactory.CreateLogger(); - callbacks = new StripedCallbackDictionary(); + callbacks = new ConcurrentDictionary(); this.clientMessagingOptions = clientMessagingOptions.Value; var period = Max( TimeSpan.FromMilliseconds(1), @@ -92,7 +93,7 @@ public OutsideRuntimeClient( TimeSpan.FromSeconds(1))); this.callbackTimer = new PeriodicTimer(period, timeProvider); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), + msg => this.UnregisterCallback(msg.Id), this.loggerFactory.CreateLogger(), this.clientMessagingOptions.ResponseTimeout, this.clientMessagingOptions.CancelRequestOnTimeout, @@ -296,7 +297,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp return; } - callbacks.TryAdd(message.SendingGrain, message.Id, callbackData); + callbacks.TryAdd(message.Id, callbackData); callbackData.SubscribeForCancellation(cancellationToken); if (Volatile.Read(ref _isStopping) != 0) @@ -327,7 +328,7 @@ public void ReceiveResponse(Message response) if (response.Result is Message.ResponseTypes.Status) { var status = (StatusResponse)response.BodyObject!; - callbacks.TryGetValue(response.TargetGrain, response.Id, out var callback); + callbacks.TryGetValue(response.Id, out var callback); var request = callback?.Message; if (request is not null) { @@ -360,7 +361,7 @@ public void ReceiveResponse(Message response) } CallbackData? callbackData; - var found = callbacks.TryRemove(response.TargetGrain, response.Id, out callbackData); + var found = callbacks.TryRemove(response.Id, out callbackData); if (found) { // We need to import the RequestContext here as well. @@ -374,9 +375,9 @@ public void ReceiveResponse(Message response) } } - private void UnregisterCallback(GrainId owner, CorrelationId id) + private void UnregisterCallback(CorrelationId id) { - callbacks.TryRemove(owner, id, out _); + callbacks.TryRemove(id, out _); } private void ConstructorReset() @@ -446,18 +447,18 @@ public void Dispose() public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) { - callbacks.ForEach(deadSilo, static (callback, deadSilo) => + foreach (var callback in callbacks) { - if (deadSilo.Equals(callback.Message.TargetSilo)) + if (deadSilo.Equals(callback.Value.Message.TargetSilo)) { - callback.OnTargetSiloFail(); + callback.Value.OnTargetSiloFail(); } - }); + } } private void BreakOutstandingMessages() { - callbacks.ForEach(this, static (callback, self) => + foreach (var (_, callback) in callbacks) { try { @@ -465,13 +466,13 @@ private void BreakOutstandingMessages() } catch (Exception exception) { - LogErrorWhileProcessingCallbackExpiry(self.logger, exception); + LogErrorWhileProcessingCallbackExpiry(logger, exception); } - }); + } } public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) - => this.callbacks.CountWhere(c => c.Message.InterfaceType == grainInterfaceType); + => this.callbacks.Count(c => c.Value.Message.InterfaceType == grainInterfaceType); /// public void NotifyClusterConnectionLost() @@ -515,18 +516,18 @@ private async Task MonitorCallbackExpiry() try { var currentStopwatchTicks = ValueStopwatch.GetTimestamp(); - callbacks.ForEach((Self: this, CurrentStopwatchTicks: currentStopwatchTicks), static (callback, state) => + foreach (var (_, callback) in callbacks) { if (callback.IsCompleted) { - return; + continue; } - if (callback.IsExpired(state.CurrentStopwatchTicks)) + if (callback.IsExpired(currentStopwatchTicks)) { callback.OnTimeout(); } - }); + } } catch (Exception ex) { From 1c7e0e15a19d746c37cf26737740a96382a31b18 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 17:10:18 -0700 Subject: [PATCH 10/18] perf(runtime): key silo callbacks by correlation id --- .../Messaging/StripedCallbackDictionary.cs | 16 +++---- .../Core/InsideRuntimeClient.cs | 15 ++++--- .../StripedCallbackDictionaryTests.cs | 45 +++++++------------ 3 files changed, 30 insertions(+), 46 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 27a2702d0af..a3daa79591f 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -55,12 +55,12 @@ private Stripe GetStripe(CorrelationId correlationId) /// Attempts to add the specified key and value to the dictionary. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryAdd(GrainId owner, CorrelationId id, TValue value) + public bool TryAdd(CorrelationId id, TValue value) { var stripe = GetStripe(id); lock (stripe.Lock) { - return stripe.Dictionary.TryAdd(new(owner, id), value); + return stripe.Dictionary.TryAdd(id, value); } } @@ -68,12 +68,12 @@ public bool TryAdd(GrainId owner, CorrelationId id, TValue value) /// Attempts to get the value associated with the specified key. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryGetValue(GrainId owner, CorrelationId id, [NotNullWhen(true)] out TValue? value) + public bool TryGetValue(CorrelationId id, [NotNullWhen(true)] out TValue? value) { var stripe = GetStripe(id); lock (stripe.Lock) { - return stripe.Dictionary.TryGetValue(new(owner, id), out value); + return stripe.Dictionary.TryGetValue(id, out value); } } @@ -81,12 +81,12 @@ public bool TryGetValue(GrainId owner, CorrelationId id, [NotNullWhen(true)] out /// Attempts to remove the value with the specified key. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryRemove(GrainId owner, CorrelationId id, [NotNullWhen(true)] out TValue? value) + public bool TryRemove(CorrelationId id, [NotNullWhen(true)] out TValue? value) { var stripe = GetStripe(id); lock (stripe.Lock) { - return stripe.Dictionary.Remove(new(owner, id), out value); + return stripe.Dictionary.Remove(id, out value); } } @@ -180,8 +180,6 @@ private sealed class Stripe #else public readonly object Lock = new(); #endif - public readonly Dictionary Dictionary = new(); + public readonly Dictionary Dictionary = new(); } - - private readonly record struct CallbackKey(GrainId Owner, CorrelationId Id); } diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index edb06397b64..b105ce0de7b 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -30,6 +30,7 @@ internal sealed partial class InsideRuntimeClient : IRuntimeClient, ILifecyclePa private readonly ILogger invokeExceptionLogger; private readonly ILoggerFactory loggerFactory; private readonly SiloMessagingOptions messagingOptions; + // 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; @@ -87,7 +88,7 @@ public InsideRuntimeClient( var callbackDataLogger = loggerFactory.CreateLogger(); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), + msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, @@ -95,7 +96,7 @@ public InsideRuntimeClient( cancellationManager: null!); this.systemSharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.SendingGrain, msg.Id), + msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, @@ -194,7 +195,7 @@ public void SendRequest( return; } - callbacks.TryAdd(message.SendingGrain, message.Id, callbackData); + callbacks.TryAdd(message.Id, callbackData); callbackData.SubscribeForCancellation(cancellationToken); } else @@ -235,9 +236,9 @@ public void SendResponse(Message request, Response response) /// /// UnRegister a callback. /// - private void UnregisterCallback(GrainId owner, CorrelationId correlationId) + private void UnregisterCallback(CorrelationId correlationId) { - callbacks.TryRemove(owner, correlationId, out _); + callbacks.TryRemove(correlationId, out _); } public void SniffIncomingMessage(Message message) @@ -467,7 +468,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. @@ -482,7 +483,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) { diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index e239e9f3197..dfd1b291060 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -8,7 +8,6 @@ namespace Tester; [TestCategory("BVT")] public class StripedCallbackDictionaryTests { - private static readonly GrainId Owner = GrainId.Create("test", "owner"); private static readonly Action EmptyVisitor = static (_, _) => { }; [Fact] @@ -36,27 +35,13 @@ public void AddGetAndRemovePreserveValue() var dictionary = new StripedCallbackDictionary(); var id = new CorrelationId(42); - Assert.True(dictionary.TryAdd(Owner, id, "value")); - Assert.False(dictionary.TryAdd(Owner, id, "duplicate")); - Assert.True(dictionary.TryGetValue(Owner, id, out var value)); + 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(Owner, id, out value)); - Assert.Equal("value", value); - Assert.False(dictionary.TryGetValue(Owner, id, out _)); - } - - [Fact] - public void CallbackOwnerIsPartOfTheKey() - { - var dictionary = new StripedCallbackDictionary(); - var otherOwner = GrainId.Create("test", "other-owner"); - var id = new CorrelationId(42); - - Assert.True(dictionary.TryAdd(Owner, id, "value")); - Assert.False(dictionary.TryGetValue(otherOwner, id, out _)); - Assert.False(dictionary.TryRemove(otherOwner, id, out _)); - Assert.True(dictionary.TryGetValue(Owner, id, out var value)); + Assert.True(dictionary.TryRemove(id, out value)); Assert.Equal("value", value); + Assert.False(dictionary.TryGetValue(id, out _)); } [Fact] @@ -66,7 +51,7 @@ public void EnumerationReturnsSnapshotValues() for (var i = 0; i < 32; i++) { var id = new CorrelationId(i); - Assert.True(dictionary.TryAdd(Owner, id, i)); + Assert.True(dictionary.TryAdd(id, i)); } var values = new List(); @@ -83,8 +68,8 @@ public void ConcurrentOperationsPreserveCountAndValues() Parallel.For(0, 10_000, i => { var id = new CorrelationId(i); - Assert.True(dictionary.TryAdd(Owner, id, i)); - Assert.True(dictionary.TryGetValue(Owner, id, out var value)); + Assert.True(dictionary.TryAdd(id, i)); + Assert.True(dictionary.TryGetValue(id, out var value)); Assert.Equal(i, value); }); @@ -93,7 +78,7 @@ public void ConcurrentOperationsPreserveCountAndValues() Parallel.For(0, 10_000, i => { - Assert.True(dictionary.TryRemove(Owner, new CorrelationId(i), out var value)); + Assert.True(dictionary.TryRemove(new CorrelationId(i), out var value)); Assert.Equal(i, value); }); @@ -107,20 +92,20 @@ public void ConcurrentLookupAndRemovalRemainConsistent() var dictionary = new StripedCallbackDictionary(); for (var i = 0; i < count; i++) { - Assert.True(dictionary.TryAdd(Owner, new CorrelationId(i), i)); + Assert.True(dictionary.TryAdd(new CorrelationId(i), i)); } Parallel.Invoke( () => Parallel.For(0, count, i => { - if (dictionary.TryGetValue(Owner, new CorrelationId(i), out var value)) + if (dictionary.TryGetValue(new CorrelationId(i), out var value)) { Assert.Equal(i, value); } }), () => Parallel.For(0, count, i => { - Assert.True(dictionary.TryRemove(Owner, new CorrelationId(i), out var value)); + Assert.True(dictionary.TryRemove(new CorrelationId(i), out var value)); Assert.Equal(i, value); })); @@ -133,12 +118,12 @@ public void SnapshotVisitorAllowsValuesToRemoveThemselves() var dictionary = new StripedCallbackDictionary(); for (var i = 0; i < 32; i++) { - Assert.True(dictionary.TryAdd(Owner, new CorrelationId(i), i)); + Assert.True(dictionary.TryAdd(new CorrelationId(i), i)); } - dictionary.ForEach((Dictionary: dictionary, Owner), static (value, state) => + dictionary.ForEach(dictionary, static (value, dictionary) => { - Assert.True(state.Dictionary.TryRemove(state.Owner, new CorrelationId(value), out var removed)); + Assert.True(dictionary.TryRemove(new CorrelationId(value), out var removed)); Assert.Equal(value, removed); }); From f7c990e2c59a07984258fc3900a3bec4786ad4be Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 21:10:50 -0700 Subject: [PATCH 11/18] perf(runtime): remove callback count allocation --- .../Messaging/StripedCallbackDictionary.cs | 8 +++++++- src/Orleans.Runtime/Core/InsideRuntimeClient.cs | 4 +++- .../StripedCallbackDictionaryTests.cs | 16 ++++++++++++++++ 3 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index a3daa79591f..9afdba06d84 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -113,6 +113,12 @@ public int 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) @@ -121,7 +127,7 @@ public int CountWhere(Func predicate) { foreach (var value in stripe.Dictionary.Values) { - if (predicate(value)) + if (predicate(value, state)) { count++; } diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index b105ce0de7b..c46d7d05b89 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -616,7 +616,9 @@ public void Participate(ISiloLifecycle lifecycle) } public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) - => this.callbacks.CountWhere(c => c.Message.InterfaceType == grainInterfaceType); + => this.callbacks.CountWhere( + grainInterfaceType, + static (callback, grainInterfaceType) => callback.Message.InterfaceType == grainInterfaceType); private async Task MonitorCallbackExpiry() { diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index dfd1b291060..e12b6b7867a 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -9,6 +9,7 @@ namespace Tester; public class StripedCallbackDictionaryTests { private static readonly Action EmptyVisitor = static (_, _) => { }; + private static readonly Func MatchValue = static (value, expected) => value == expected; [Fact] public void CorrelationIdsDistributeAcrossStripesAtOverflowAndWithStride() @@ -142,4 +143,19 @@ public void EmptyVisitorDoesNotAllocate() 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); + } } From 4bd2adfb627f6fde3031d9e751a75fe0a0fabd20 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 18:06:21 -0700 Subject: [PATCH 12/18] fix(runtime): remove callbacks by exact instance --- .../Messaging/StripedCallbackDictionary.cs | 18 +++++++++++++++++ .../StripedCallbackDictionaryTests.cs | 20 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index 9afdba06d84..e35ad042cd9 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -90,6 +90,24 @@ public bool TryRemove(CorrelationId id, [NotNullWhen(true)] out TValue? value) } } + /// + /// Attempts to remove the value with the specified key if it is the expected instance. + /// + [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) || !ReferenceEquals(value, expected)) + { + return false; + } + + return stripe.Dictionary.Remove(id); + } + } + /// /// Gets the approximate total count of items across all stripes. /// diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index e12b6b7867a..dd8e77129af 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -45,6 +45,26 @@ public void AddGetAndRemovePreserveValue() 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 EnumerationReturnsSnapshotValues() { From e41af20c115590fddad6cb0e8dd08ef72cdec128 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 18:07:14 -0700 Subject: [PATCH 13/18] fix(runtime): give callbacks unregister ownership --- src/Orleans.Core/Runtime/CallbackData.cs | 16 ++++++++++---- .../Runtime/OutsideRuntimeClient.cs | 17 ++++++++------- .../Runtime/SharedCallbackData.cs | 3 --- .../Core/InsideRuntimeClient.cs | 21 ++++++++----------- .../CallbackDataTests.cs | 10 ++++++--- 5 files changed, 37 insertions(+), 30 deletions(-) diff --git a/src/Orleans.Core/Runtime/CallbackData.cs b/src/Orleans.Core/Runtime/CallbackData.cs index bfbfb7a28f4..f96a0537fb0 100644 --- a/src/Orleans.Core/Runtime/CallbackData.cs +++ b/src/Orleans.Core/Runtime/CallbackData.cs @@ -6,6 +6,11 @@ namespace Orleans.Runtime { + internal interface ICallbackDataTarget + { + void Unregister(CallbackData callback); + } + internal sealed partial class CallbackData { private const int StateNone = 0; @@ -14,6 +19,7 @@ internal sealed partial class CallbackData private const int StateCancellationRegistrationPublished = 4; private readonly SharedCallbackData shared; + private readonly ICallbackDataTarget target; private readonly IResponseCompletionSource context; private readonly ApplicationRequestInstruments _applicationRequestInstruments; private int _state; @@ -23,11 +29,13 @@ internal sealed partial class CallbackData public CallbackData( SharedCallbackData shared, + ICallbackDataTarget target, IResponseCompletionSource ctx, Message msg, ApplicationRequestInstruments applicationRequestInstruments) { this.shared = shared; + this.target = target; this.context = ctx; this.Message = msg; _applicationRequestInstruments = applicationRequestInstruments; @@ -129,7 +137,7 @@ private void OnCancellation(CancellationToken cancellationToken) stopwatch.Stop(); SignalCancellation(); - shared.Unregister(Message); + target.Unregister(this); _applicationRequestInstruments.OnAppRequestsEnd((long)stopwatch.Elapsed.TotalMilliseconds); _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType()); OrleansCallBackDataEvent.Instance.OnCanceled(Message); @@ -150,7 +158,7 @@ public void OnTimeout() SignalCancellation(); } - this.shared.Unregister(this.Message); + this.target.Unregister(this); DisposeCancellationRegistration(); _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds); _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType()); @@ -175,7 +183,7 @@ public void OnTargetSiloFail() } this.stopwatch.Stop(); - this.shared.Unregister(this.Message); + this.target.Unregister(this); DisposeCancellationRegistration(); _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds); @@ -195,7 +203,7 @@ public void OnHostShutdown() } this.stopwatch.Stop(); - this.shared.Unregister(this.Message); + this.target.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..75e8033668e 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,6 @@ public OutsideRuntimeClient( TimeSpan.FromSeconds(1))); this.callbackTimer = new PeriodicTimer(period, timeProvider); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.Id), this.loggerFactory.CreateLogger(), this.clientMessagingOptions.ResponseTimeout, this.clientMessagingOptions.CancelRequestOnTimeout, @@ -290,14 +289,18 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp if (!oneWay) { - var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments); + var callbackData = new CallbackData(this.sharedCallbackData, this, context!, message, _applicationRequestInstruments); if (Volatile.Read(ref _isStopping) != 0) { callbackData.OnHostShutdown(); 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 +378,8 @@ public void ReceiveResponse(Message response) } } - private void UnregisterCallback(CorrelationId id) - { - callbacks.TryRemove(id, out _); - } + void ICallbackDataTarget.Unregister(CallbackData callback) => + callbacks.TryRemove(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..a4689beb1ee 100644 --- a/src/Orleans.Core/Runtime/SharedCallbackData.cs +++ b/src/Orleans.Core/Runtime/SharedCallbackData.cs @@ -6,20 +6,17 @@ namespace Orleans.Runtime; internal sealed class SharedCallbackData { - public readonly Action Unregister; public readonly ILogger Logger; private TimeSpan _responseTimeout; public long ResponseTimeoutStopwatchTicks; public SharedCallbackData( - Action unregister, ILogger logger, TimeSpan responseTimeout, bool cancelOnTimeout, bool waitForCancellationAcknowledgement, IGrainCallCancellationManager? cancellationManager) { - Unregister = unregister; Logger = logger; ResponseTimeout = responseTimeout; CancelRequestOnTimeout = cancelOnTimeout; diff --git a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs index c46d7d05b89..40c60d1ae79 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -24,7 +24,7 @@ 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; @@ -88,7 +88,6 @@ public InsideRuntimeClient( var callbackDataLogger = loggerFactory.CreateLogger(); this.sharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, @@ -96,7 +95,6 @@ public InsideRuntimeClient( cancellationManager: null!); this.systemSharedCallbackData = new SharedCallbackData( - msg => this.UnregisterCallback(msg.Id), callbackDataLogger, this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, @@ -188,14 +186,18 @@ public void SendRequest( Debug.Assert(context is not null); // Register a callback for the request. - callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); + callbackData = new CallbackData(sharedData, this, context, message, _applicationRequestInstruments); if (Volatile.Read(ref _isStopping) != 0) { callbackData.OnHostShutdown(); 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); } else @@ -233,13 +235,8 @@ public void SendResponse(Message request, Response response) this.MessageCenter.SendResponse(request, response); } - /// - /// UnRegister a callback. - /// - private void UnregisterCallback(CorrelationId correlationId) - { - callbacks.TryRemove(correlationId, out _); - } + void ICallbackDataTarget.Unregister(CallbackData callback) => + callbacks.TryRemove(callback.Message.Id, callback); public void SniffIncomingMessage(Message message) { diff --git a/test/Orleans.Runtime.Tests/CallbackDataTests.cs b/test/Orleans.Runtime.Tests/CallbackDataTests.cs index dedcc46b756..8193eeab9c1 100644 --- a/test/Orleans.Runtime.Tests/CallbackDataTests.cs +++ b/test/Orleans.Runtime.Tests/CallbackDataTests.cs @@ -68,17 +68,16 @@ private static WeakReference CreateCompletedCallback(CancellationToken cancellat private static CallbackData CreateCallback( IResponseCompletionSource completion, - Action unregister, + Action unregister, ApplicationRequestInstruments instruments) { var shared = new SharedCallbackData( - unregister, logger: NullLogger.Instance, responseTimeout: TimeSpan.FromMinutes(1), cancelOnTimeout: false, waitForCancellationAcknowledgement: false, cancellationManager: null); - return new CallbackData(shared, completion, new Message(), instruments); + return new CallbackData(shared, new DelegateCallbackTarget(unregister), completion, new Message(), instruments); } private static ServiceProvider CreateServiceProvider() @@ -99,4 +98,9 @@ private sealed class TestResponseCompletionSource : IResponseCompletionSource public void Complete() => Response = Orleans.Serialization.Invocation.Response.Completed; } + + private sealed class DelegateCallbackTarget(Action unregister) : ICallbackDataTarget + { + public void Unregister(CallbackData callback) => unregister(callback); + } } From d10ef9762f2841f075bae5b6959457fe506c559f Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 18:08:16 -0700 Subject: [PATCH 14/18] test(runtime): cover callback ownership races --- .../CallbackDataTests.cs | 201 +++++++++++++++++- 1 file changed, 192 insertions(+), 9 deletions(-) diff --git a/test/Orleans.Runtime.Tests/CallbackDataTests.cs b/test/Orleans.Runtime.Tests/CallbackDataTests.cs index 8193eeab9c1..c28d47f5484 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) { @@ -90,17 +211,79 @@ 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 void Complete(Response value) => Response = value; + public int CompletionCount => Volatile.Read(ref _completionCount); - public void Complete() => Response = Orleans.Serialization.Invocation.Response.Completed; + 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(ApplicationRequestInstruments instruments) : ICallbackDataTarget + { + private readonly StripedCallbackDictionary _callbacks = new(); + + public int Count => _callbacks.Count; + + public CallbackData Register(CorrelationId id, IResponseCompletionSource completion) + { + var message = new Message { Id = id }; + var callback = new CallbackData(CreateSharedData(), this, completion, message, instruments); + if (!_callbacks.TryAdd(id, callback)) + { + throw new InvalidOperationException($"A callback with correlation id {id} is already registered."); + } + + 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 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() => new( + logger: NullLogger.Instance, + responseTimeout: TimeSpan.FromMinutes(1), + cancelOnTimeout: false, + waitForCancellationAcknowledgement: false, + cancellationManager: null); } From ae789d5fc424953f597200c7ec70ed3903e1741c Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 18:14:16 -0700 Subject: [PATCH 15/18] fix(runtime): preserve client callback ownership on net8 --- src/Orleans.Core/Runtime/OutsideRuntimeClient.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs index 75e8033668e..909232a10d2 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -379,7 +379,8 @@ public void ReceiveResponse(Message response) } void ICallbackDataTarget.Unregister(CallbackData callback) => - callbacks.TryRemove(KeyValuePair.Create(callback.Message.Id, callback)); + ((ICollection>)callbacks) + .Remove(KeyValuePair.Create(callback.Message.Id, callback)); private void ConstructorReset() { From 570da7ef31ea6fe84ce45afa1756c52210e8098d Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 27 Aug 2026 18:16:40 -0700 Subject: [PATCH 16/18] fix(tests): preserve nullable callback flow --- test/Orleans.Runtime.Tests/CallbackDataTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Orleans.Runtime.Tests/CallbackDataTests.cs b/test/Orleans.Runtime.Tests/CallbackDataTests.cs index c28d47f5484..e7fb4e20320 100644 --- a/test/Orleans.Runtime.Tests/CallbackDataTests.cs +++ b/test/Orleans.Runtime.Tests/CallbackDataTests.cs @@ -259,7 +259,7 @@ public CallbackData Register(CorrelationId id, IResponseCompletionSource complet return callback; } - public bool TryTake(CorrelationId id, out CallbackData callback) => + public bool TryTake(CorrelationId id, out CallbackData? callback) => _callbacks.TryRemove(id, out callback); public CallbackData Take(CorrelationId id) From 92004f0378471a16fe9fb7a492486cd12a54fed1 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 28 Aug 2026 07:02:46 -0700 Subject: [PATCH 17/18] perf(runtime): share callback ownership target --- src/Orleans.Core/Runtime/CallbackData.cs | 12 ++++-------- .../Runtime/OutsideRuntimeClient.cs | 3 ++- src/Orleans.Core/Runtime/SharedCallbackData.cs | 5 +++++ .../Core/InsideRuntimeClient.cs | 4 +++- .../Orleans.Runtime.Tests/CallbackDataTests.cs | 18 ++++++++++++++---- 5 files changed, 28 insertions(+), 14 deletions(-) diff --git a/src/Orleans.Core/Runtime/CallbackData.cs b/src/Orleans.Core/Runtime/CallbackData.cs index f96a0537fb0..3383c0aba6a 100644 --- a/src/Orleans.Core/Runtime/CallbackData.cs +++ b/src/Orleans.Core/Runtime/CallbackData.cs @@ -17,9 +17,7 @@ internal sealed partial class CallbackData private const int StateCompleted = 1; private const int StateCancellationRegistrationPending = 2; private const int StateCancellationRegistrationPublished = 4; - private readonly SharedCallbackData shared; - private readonly ICallbackDataTarget target; private readonly IResponseCompletionSource context; private readonly ApplicationRequestInstruments _applicationRequestInstruments; private int _state; @@ -29,13 +27,11 @@ internal sealed partial class CallbackData public CallbackData( SharedCallbackData shared, - ICallbackDataTarget target, IResponseCompletionSource ctx, Message msg, ApplicationRequestInstruments applicationRequestInstruments) { this.shared = shared; - this.target = target; this.context = ctx; this.Message = msg; _applicationRequestInstruments = applicationRequestInstruments; @@ -137,7 +133,7 @@ private void OnCancellation(CancellationToken cancellationToken) stopwatch.Stop(); SignalCancellation(); - target.Unregister(this); + shared.Unregister(this); _applicationRequestInstruments.OnAppRequestsEnd((long)stopwatch.Elapsed.TotalMilliseconds); _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType()); OrleansCallBackDataEvent.Instance.OnCanceled(Message); @@ -158,7 +154,7 @@ public void OnTimeout() SignalCancellation(); } - this.target.Unregister(this); + this.shared.Unregister(this); DisposeCancellationRegistration(); _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds); _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType()); @@ -183,7 +179,7 @@ public void OnTargetSiloFail() } this.stopwatch.Stop(); - this.target.Unregister(this); + this.shared.Unregister(this); DisposeCancellationRegistration(); _applicationRequestInstruments.OnAppRequestsEnd((long)this.stopwatch.Elapsed.TotalMilliseconds); @@ -203,7 +199,7 @@ public void OnHostShutdown() } this.stopwatch.Stop(); - this.target.Unregister(this); + 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 909232a10d2..545233b4f3d 100644 --- a/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs +++ b/src/Orleans.Core/Runtime/OutsideRuntimeClient.cs @@ -93,6 +93,7 @@ public OutsideRuntimeClient( TimeSpan.FromSeconds(1))); this.callbackTimer = new PeriodicTimer(period, timeProvider); this.sharedCallbackData = new SharedCallbackData( + this, this.loggerFactory.CreateLogger(), this.clientMessagingOptions.ResponseTimeout, this.clientMessagingOptions.CancelRequestOnTimeout, @@ -289,7 +290,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp if (!oneWay) { - var callbackData = new CallbackData(this.sharedCallbackData, this, context!, message, _applicationRequestInstruments); + var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments); if (Volatile.Read(ref _isStopping) != 0) { callbackData.OnHostShutdown(); diff --git a/src/Orleans.Core/Runtime/SharedCallbackData.cs b/src/Orleans.Core/Runtime/SharedCallbackData.cs index a4689beb1ee..f83d57a8ad2 100644 --- a/src/Orleans.Core/Runtime/SharedCallbackData.cs +++ b/src/Orleans.Core/Runtime/SharedCallbackData.cs @@ -6,17 +6,20 @@ namespace Orleans.Runtime; internal sealed class SharedCallbackData { + private readonly ICallbackDataTarget _target; public readonly ILogger Logger; private TimeSpan _responseTimeout; public long ResponseTimeoutStopwatchTicks; public SharedCallbackData( + ICallbackDataTarget target, ILogger logger, TimeSpan responseTimeout, bool cancelOnTimeout, bool waitForCancellationAcknowledgement, IGrainCallCancellationManager? cancellationManager) { + _target = target; Logger = logger; ResponseTimeout = responseTimeout; CancelRequestOnTimeout = cancelOnTimeout; @@ -39,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 40c60d1ae79..478e6ca9255 100644 --- a/src/Orleans.Runtime/Core/InsideRuntimeClient.cs +++ b/src/Orleans.Runtime/Core/InsideRuntimeClient.cs @@ -88,6 +88,7 @@ public InsideRuntimeClient( var callbackDataLogger = loggerFactory.CreateLogger(); this.sharedCallbackData = new SharedCallbackData( + this, callbackDataLogger, this.messagingOptions.ResponseTimeout, this.messagingOptions.CancelRequestOnTimeout, @@ -95,6 +96,7 @@ public InsideRuntimeClient( cancellationManager: null!); this.systemSharedCallbackData = new SharedCallbackData( + this, callbackDataLogger, this.messagingOptions.SystemResponseTimeout, cancelOnTimeout: false, @@ -186,7 +188,7 @@ public void SendRequest( Debug.Assert(context is not null); // Register a callback for the request. - callbackData = new CallbackData(sharedData, this, context, message, _applicationRequestInstruments); + callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments); if (Volatile.Read(ref _isStopping) != 0) { callbackData.OnHostShutdown(); diff --git a/test/Orleans.Runtime.Tests/CallbackDataTests.cs b/test/Orleans.Runtime.Tests/CallbackDataTests.cs index e7fb4e20320..8d5da45c884 100644 --- a/test/Orleans.Runtime.Tests/CallbackDataTests.cs +++ b/test/Orleans.Runtime.Tests/CallbackDataTests.cs @@ -193,12 +193,13 @@ private static CallbackData CreateCallback( ApplicationRequestInstruments instruments) { var shared = new SharedCallbackData( + new DelegateCallbackTarget(unregister), logger: NullLogger.Instance, responseTimeout: TimeSpan.FromMinutes(1), cancelOnTimeout: false, waitForCancellationAcknowledgement: false, cancellationManager: null); - return new CallbackData(shared, new DelegateCallbackTarget(unregister), completion, new Message(), instruments); + return new CallbackData(shared, completion, new Message(), instruments); } private static ServiceProvider CreateServiceProvider() @@ -241,16 +242,24 @@ private sealed class DelegateCallbackTarget(Action unregister) : I public void Unregister(CallbackData callback) => unregister(callback); } - private sealed class TestCallbackRegistry(ApplicationRequestInstruments instruments) : ICallbackDataTarget + 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(CreateSharedData(), this, completion, message, instruments); + 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."); @@ -280,7 +289,8 @@ void ICallbackDataTarget.Unregister(CallbackData callback) => _callbacks.TryRemove(callback.Message.Id, callback); } - private static SharedCallbackData CreateSharedData() => new( + private static SharedCallbackData CreateSharedData(ICallbackDataTarget target) => new( + target, logger: NullLogger.Instance, responseTimeout: TimeSpan.FromMinutes(1), cancelOnTimeout: false, From a6ef79f58e7abde8ce070525fdd6abda1a47319e Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Fri, 28 Aug 2026 07:35:37 -0700 Subject: [PATCH 18/18] fix(runtime): restore exact value removal --- .../Messaging/StripedCallbackDictionary.cs | 7 +++++-- .../StripedCallbackDictionaryTests.cs | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs index e35ad042cd9..84b669382a5 100644 --- a/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs +++ b/src/Orleans.Core/Messaging/StripedCallbackDictionary.cs @@ -91,7 +91,7 @@ public bool TryRemove(CorrelationId id, [NotNullWhen(true)] out TValue? value) } /// - /// Attempts to remove the value with the specified key if it is the expected instance. + /// 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) @@ -99,7 +99,10 @@ public bool TryRemove(CorrelationId id, TValue expected) var stripe = GetStripe(id); lock (stripe.Lock) { - if (!stripe.Dictionary.TryGetValue(id, out var value) || !ReferenceEquals(value, expected)) + if (!stripe.Dictionary.TryGetValue(id, out var value) + || (typeof(TValue).IsValueType + ? !EqualityComparer.Default.Equals(value, expected) + : !ReferenceEquals(value, expected))) { return false; } diff --git a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs index dd8e77129af..e49d12641e3 100644 --- a/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs +++ b/test/Orleans.Runtime.Tests/StripedCallbackDictionaryTests.cs @@ -65,6 +65,20 @@ public void ExactRemovalDoesNotRemoveReplacement() 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() {