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>

Author: ReubenBond

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -28,6 +28,7 @@ internal partial class OutsideRuntimeClient : IRuntimeClient, IDisposable, IClus
 
         private readonly StripedCallbackDictionary<CallbackData> 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);
 

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<HostedClient>();
+        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<GrainLocator>());
+        private GrainLocator GrainLocator => this.grainLocator;
 
-        private List<IIncomingGrainCallFilter> GrainCallFilters
-            => this.grainCallFilters ??= new List<IIncomingGrainCallFilter>(this.ServiceProvider.GetServices<IIncomingGrainCallFilter>());
+        private List<IIncomingGrainCallFilter> GrainCallFilters => this.grainCallFilters;
 
-        private MessageCenter MessageCenter => this.messageCenter ?? (this.messageCenter = this.ServiceProvider.GetRequiredService<MessageCenter>());
+        private MessageCenter MessageCenter => this.messageCenter;
 
-        public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime ?? (this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService<IGrainReferenceRuntime>());
+        public IGrainReferenceRuntime GrainReferenceRuntime => this.grainReferenceRuntime;
+
+        internal void ConsumeServices()
+        {
+            this.grainLocator = this.ServiceProvider.GetRequiredService<GrainLocator>();
+            this.grainCallFilters = new List<IIncomingGrainCallFilter>(this.ServiceProvider.GetServices<IIncomingGrainCallFilter>());
+            this.messageCenter = this.ServiceProvider.GetRequiredService<MessageCenter>();
+            this.grainReferenceRuntime = this.ServiceProvider.GetRequiredService<IGrainReferenceRuntime>();
+            this.hostedClient = this.ServiceProvider.GetRequiredService<HostedClient>();
+            _cancellationManager = this.ServiceProvider.GetRequiredService<IGrainCallCancellationManager>();
+            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<IGrainCallCancellationManager>();
-            sharedCallbackData.CancellationManager = _cancellationManager;
-            systemSharedCallbackData.CancellationManager = _cancellationManager;
+            ConsumeServices();
             lifecycle.Subscribe<InsideRuntimeClient>(ServiceLifecycleStage.RuntimeInitialize, OnRuntimeInitializeStart, OnRuntimeInitializeStop);
         }