fix(messaging): preserve pooled callback lifecycle

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: ReubenBond

src/Orleans.Core/Messaging/Message.cs

@@ -296,7 +296,10 @@ internal void AddToCacheInvalidationHeader(GrainAddress invalidAddress, GrainAdd
                     // Another thread initialized it, add to the existing list
                     lock (_cacheInvalidationHeader)
                     {
-                        _cacheInvalidationHeader.Add(grainAddressCacheUpdate);
+                        if (_cacheInvalidationHeader.Count < MaxCacheInvalidationHeaderEntries)
+                        {
+                            _cacheInvalidationHeader.Add(grainAddressCacheUpdate);
+                        }
                     }
                 }
                 else
@@ -308,7 +311,10 @@ internal void AddToCacheInvalidationHeader(GrainAddress invalidAddress, GrainAdd
             {
                 lock (_cacheInvalidationHeader)
                 {
-                    _cacheInvalidationHeader.Add(grainAddressCacheUpdate);
+                    if (_cacheInvalidationHeader.Count < MaxCacheInvalidationHeaderEntries)
+                    {
+                        _cacheInvalidationHeader.Add(grainAddressCacheUpdate);
+                    }
                 }
             }
         }

src/Orleans.Core/Networking/Connection.cs

@@ -482,6 +482,7 @@ private bool HandleReceiveMessageFailure(Message message, Exception exception)
 
                     // Send the error response and continue processing the next message.
                     this.Send(response);
+                    message.ReleaseDropped("ReceiveMessageDeserializationFailure");
                 }
                 else if (message.Direction == Message.Directions.Response)
                 {

src/Orleans.Core/Runtime/CallbackData.cs

@@ -14,6 +14,7 @@ internal sealed partial class CallbackData
         private int completed;
         private StatusResponse? lastKnownStatus;
         private ValueStopwatch stopwatch;
+        private CancellationToken _cancellationToken;
         private CancellationTokenRegistration _cancellationTokenRegistration;
 
         public CallbackData(
@@ -33,7 +34,7 @@ public CallbackData(
 
         public Message Message { get; } // might hold metadata used by response pipeline
 
-        public bool IsCompleted => this.completed == 1;
+        public bool IsCompleted => Volatile.Read(ref completed) != 0;
 
         public void SubscribeForCancellation(CancellationToken cancellationToken)
         {
@@ -42,12 +43,26 @@ public void SubscribeForCancellation(CancellationToken cancellationToken)
                 return;
             }
 
-            cancellationToken.ThrowIfCancellationRequested();
-            _cancellationTokenRegistration = cancellationToken.UnsafeRegister(static arg =>
+            if (IsCompleted)
+            {
+                return;
+            }
+
+            _cancellationToken = cancellationToken;
+            var registration = cancellationToken.UnsafeRegister(static arg =>
             {
                 var callbackData = (CallbackData)arg!;
                 callbackData.OnCancellation();
             }, this);
+
+            if (IsCompleted)
+            {
+                registration.Dispose();
+            }
+            else
+            {
+                _cancellationTokenRegistration = registration;
+            }
         }
 
         private void SignalCancellation()
@@ -112,9 +127,9 @@ private void OnCancellation()
             SignalCancellation();
             shared.Unregister(Message);
             _applicationRequestInstruments.OnAppRequestsEnd((long)stopwatch.Elapsed.TotalMilliseconds);
-            _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType());
+            _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType());
             OrleansCallBackDataEvent.Instance.OnCanceled(Message);
-            context.Complete(Response.FromException(new OperationCanceledException(_cancellationTokenRegistration.Token)));
+            context.Complete(Response.FromException(new OperationCanceledException(_cancellationToken)));
             _cancellationTokenRegistration.Dispose();
             ReleaseRequest("CallbackData.OnCancellation");
         }

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -155,6 +155,12 @@ public async Task StartAsync(CancellationToken cancellationToken)
 
         public async Task StopAsync(CancellationToken cancellationToken)
         {
+            disposing = true;
+            foreach (var callback in callbacks.Values)
+            {
+                callback.OnHostShutdown();
+            }
+
             this.callbackTimer.Dispose();
 
             if (this.callbackTimerTask is { } task)
@@ -505,7 +511,7 @@ private async Task MonitorCallbackExpiry()
 
         private void ThrowIfDisposed()
         {
-            if (disposed)
+            if (disposing || disposed)
             {
                 ThrowObjectDisposedException();
             }

src/Orleans.Runtime/Core/InsideRuntimeClient.cs

@@ -113,10 +113,10 @@ public InsideRuntimeClient(
         public GrainFactory ConcreteGrainFactory { get; }
 
         private GrainLocator GrainLocator
-            => this.grainLocator ?? (this.grainLocator = this.ServiceProvider.GetRequiredService<GrainLocator>());
+            => this.grainLocator;
 
         private List<IIncomingGrainCallFilter> GrainCallFilters
-            => this.grainCallFilters ??= new List<IIncomingGrainCallFilter>(this.ServiceProvider.GetServices<IIncomingGrainCallFilter>());
+            => this.grainCallFilters;
 
         private MessageCenter MessageCenter => this.messageCenter ?? (this.messageCenter = this.ServiceProvider.GetRequiredService<MessageCenter>());
 
@@ -510,6 +510,11 @@ public void DeleteObjectReference(IAddressable obj)
 
         private async Task OnRuntimeInitializeStop(CancellationToken tc)
         {
+            foreach (var callback in callbacks.Values)
+            {
+                callback.OnHostShutdown();
+            }
+
             this.callbackTimer.Dispose();
             if (this.callbackTimerTask is { } task)
             {
@@ -549,6 +554,8 @@ public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo)
 
         public void Participate(ISiloLifecycle lifecycle)
         {
+            grainLocator = ServiceProvider.GetRequiredService<GrainLocator>();
+            grainCallFilters = new List<IIncomingGrainCallFilter>(ServiceProvider.GetServices<IIncomingGrainCallFilter>());
             _cancellationManager = this.ServiceProvider.GetRequiredService<IGrainCallCancellationManager>();
             sharedCallbackData.CancellationManager = _cancellationManager;
             systemSharedCallbackData.CancellationManager = _cancellationManager;

test/Orleans.Runtime.Tests/CallbackDataTests.cs

@@ -78,7 +78,11 @@ private static CallbackData CreateCallback(
             cancelOnTimeout: false,
             waitForCancellationAcknowledgement: false,
             cancellationManager: null);
-        return new CallbackData(shared, completion, new Message(), instruments);
+        var message = new Message();
+        message.InitializeRefCount();
+        var callback = new CallbackData(shared, completion, message, instruments);
+        message.Release();
+        return callback;
     }
 
     private static ServiceProvider CreateServiceProvider()

test/Orleans.Runtime.Tests/InsideRuntimeClientTests.cs

@@ -36,6 +36,7 @@ public async Task LateRejectionAfterSiloDisposalDoesNotResolveServices()
                 RejectionInfo = "The outbound queue is stopped",
             },
         };
+        rejection.InitializeRefCount();
 
         await silo.StopSiloAsync(stopGracefully: false);
         await silo.DisposeAsync();