fix(clients): reject outstanding requests to dropped clients

src/Orleans.Core.Abstractions/IDs/ClientGrainId.cs

@@ -59,6 +59,29 @@ public static bool TryParse(GrainId grainId, out ClientGrainId clientId)
             return true;
         }
 
+        /// <summary>
+        /// Checks if the provided <see cref="GrainId"/> points to the same client as this <see cref="ClientGrainId"/>.
+        /// </summary>
+        /// <param name="other">The <see cref="GrainId"/> to compare.</param>
+        /// <returns><see langword="true"/> if the provided <see cref="GrainId"/> corresponds to the same client, otherwise <see langword="false"/>.</returns>
+        public bool IsClientEqual(GrainId other)
+        {
+            if (!GrainId.Type.Equals(other.Type))
+            {
+                return false;
+            }
+
+            // Strip the observer id, if present, using span-based operations to avoid allocations.
+            var key = other.Key.AsSpan();
+            if (key.IndexOf((byte)ObserverGrainId.SegmentSeparator) is int index && index >= 0)
+            {
+                key = key[..index];
+            }
+
+            // Compare the stripped key with the current GrainId's key.
+            return GrainId.Key.AsSpan().SequenceEqual(key);
+        }
+
         /// <inheritdoc/>
         public override bool Equals(object? obj) => obj is ClientGrainId clientId && GrainId.Equals(clientId.GrainId);
 

src/Orleans.Core/Diagnostics/EventSourceEvents.cs

@@ -52,7 +52,7 @@ public void OnTargetSiloFail(Message message)
         /// Indicates that a request completed.
         /// </summary>
         [NonEvent]
-        public void DoCallback(Message message)
+        public void OnResponse(Message message)
         {
             if (this.IsEnabled())
             {

src/Orleans.Core/Messaging/Message.cs

@@ -9,13 +9,27 @@ namespace Orleans.Runtime
     [Id(101)]
     internal sealed class Message : ISpanFormattable
     {
+        private const string GatewayRequestOwnerHeader = "#orleans.gateway.request-owner";
+        private const string GatewayRequestOwnerSiloHeader = "#orleans.gateway.request-owner-silo";
+        private const string GatewayResponseTargetHeader = "#orleans.gateway.response-target";
+        private const string GatewayRequestTimeoutHeader = "#orleans.gateway.request-timeout";
+
         public const int LENGTH_HEADER_SIZE = 8;
         public const int LENGTH_META_HEADER = 4;
         internal const int MaxCacheInvalidationHeaderEntries = 16;
 
         [NonSerialized]
         private short _retryCount;
 
+        [NonSerialized]
+        private bool _hasGatewayRequestSource;
+
+        [NonSerialized]
+        private SiloAddress? _gatewayRequestSource;
+
+        [NonSerialized]
+        private bool _hasTrustedGatewayResponseTarget;
+
         public CoarseStopwatch _timeToExpiry;
 
         public object? BodyObject { get; set; }
@@ -257,6 +271,179 @@ public Dictionary<string, object>? RequestContextData
             }
         }
 
+        internal static bool IsGatewayRequestContextHeader(string key)
+            => key is GatewayRequestOwnerHeader
+                or GatewayRequestOwnerSiloHeader
+                or GatewayResponseTargetHeader
+                or GatewayRequestTimeoutHeader;
+
+        internal void SetGatewayRequestTimeout(TimeSpan timeout)
+        {
+            var context = RequestContextData ??= [];
+            context[GatewayRequestTimeoutHeader] = timeout;
+        }
+
+        internal TimeSpan? GetGatewayRequestTimeout()
+            => RequestContextData is { } context
+                && context.TryGetValue(GatewayRequestTimeoutHeader, out var value)
+                && value is TimeSpan timeout
+                    ? timeout
+                    : null;
+
+        internal void ClearGatewayRequestOwner()
+        {
+            ClearGatewayRequestRouting();
+            if (RequestContextData is { } context)
+            {
+                context.Remove(GatewayRequestTimeoutHeader);
+                if (context.Count == 0)
+                {
+                    RequestContextData = null;
+                }
+            }
+        }
+
+        internal void ClearGatewayRequestRouting()
+        {
+            if (RequestContextData is { } context)
+            {
+                context.Remove(GatewayRequestOwnerHeader);
+                context.Remove(GatewayRequestOwnerSiloHeader);
+                context.Remove(GatewayResponseTargetHeader);
+                if (context.Count == 0)
+                {
+                    RequestContextData = null;
+                }
+            }
+
+            _gatewayRequestSource = null;
+            _hasGatewayRequestSource = false;
+            _hasTrustedGatewayResponseTarget = false;
+        }
+
+        internal void SetGatewayRequestOwner(SiloAddress ownerGateway, SiloAddress ownerSilo)
+        {
+            var context = RequestContextData ??= [];
+            if (!_hasGatewayRequestSource)
+            {
+                _gatewayRequestSource = SendingSilo;
+                _hasGatewayRequestSource = true;
+            }
+
+            context.Remove(GatewayRequestOwnerHeader);
+            context.Remove(GatewayRequestOwnerSiloHeader);
+            context.Remove(GatewayResponseTargetHeader);
+            context[GatewayRequestOwnerHeader] = ownerGateway;
+            context[GatewayRequestOwnerSiloHeader] = ownerSilo;
+            if (_gatewayRequestSource is { } responseTarget)
+            {
+                context[GatewayResponseTargetHeader] = responseTarget;
+            }
+
+            SendingSilo = ownerGateway;
+        }
+
+        internal void RestoreGatewayRequestSource()
+        {
+            if (RequestContextData is not { } context
+                || !context.Remove(GatewayRequestOwnerHeader))
+            {
+                return;
+            }
+
+            SendingSilo = context.Remove(GatewayResponseTargetHeader, out var targetValue)
+                && targetValue is SiloAddress responseTarget
+                    ? responseTarget
+                    : null;
+            context.Remove(GatewayRequestOwnerSiloHeader);
+            _gatewayRequestSource = SendingSilo;
+
+            if (context.Count == 0)
+            {
+                RequestContextData = null;
+            }
+        }
+
+        internal void ApplyGatewayRequestOwner(Message request)
+        {
+            if (request.RequestContextData is not { } requestContext
+                || !requestContext.TryGetValue(GatewayRequestOwnerHeader, out var ownerValue)
+                || ownerValue is not SiloAddress ownerGateway)
+            {
+                return;
+            }
+
+            var responseContext = RequestContextData ??= [];
+            responseContext[GatewayRequestOwnerHeader] = ownerGateway;
+            if (requestContext.TryGetValue(GatewayRequestOwnerSiloHeader, out var ownerSiloValue)
+                && ownerSiloValue is SiloAddress ownerSilo)
+            {
+                responseContext[GatewayRequestOwnerSiloHeader] = ownerSilo;
+            }
+            if (requestContext.TryGetValue(GatewayResponseTargetHeader, out var targetValue)
+                && targetValue is SiloAddress responseTarget)
+            {
+                responseContext[GatewayResponseTargetHeader] = responseTarget;
+            }
+            else
+            {
+                responseContext.Remove(GatewayResponseTargetHeader);
+            }
+
+            TargetSilo = ownerGateway;
+        }
+
+        internal bool TryGetGatewayRequestOwner(out SiloAddress ownerGateway, out SiloAddress ownerSilo)
+        {
+            ownerGateway = default!;
+            ownerSilo = default!;
+            if (RequestContextData is not { } context
+                || !context.TryGetValue(GatewayRequestOwnerHeader, out var ownerValue)
+                || ownerValue is not SiloAddress gateway
+                || !context.TryGetValue(GatewayRequestOwnerSiloHeader, out var ownerSiloValue)
+                || ownerSiloValue is not SiloAddress silo)
+            {
+                return false;
+            }
+
+            ownerGateway = gateway;
+            ownerSilo = silo;
+            return true;
+        }
+
+        internal void RestoreGatewayResponseTarget(bool preserveRoute = false)
+        {
+            if (RequestContextData is not { } context)
+            {
+                return;
+            }
+
+            context.Remove(GatewayRequestOwnerHeader);
+            context.Remove(GatewayRequestOwnerSiloHeader);
+            TargetSilo = context.Remove(GatewayResponseTargetHeader, out var targetValue)
+                && targetValue is SiloAddress responseTarget
+                    ? responseTarget
+                    : null;
+            _hasTrustedGatewayResponseTarget = preserveRoute && TargetSilo is not null;
+
+            if (context.Count == 0)
+            {
+                RequestContextData = null;
+            }
+        }
+
+        internal bool TryTakeTrustedGatewayResponseTarget(out SiloAddress target)
+        {
+            target = TargetSilo!;
+            if (!_hasTrustedGatewayResponseTarget || target is null)
+            {
+                return false;
+            }
+
+            _hasTrustedGatewayResponseTarget = false;
+            return true;
+        }
+
         public GrainInterfaceType InterfaceType
         {
             get => _interfaceType;

src/Orleans.Core/Messaging/MessageFactory.cs

@@ -41,6 +41,7 @@ public Message CreateMessage(object? body, InvokeMethodOptions options)
                 BodyObject = body,
                 RequestContextData = RequestContextExtensions.Export(_deepCopier),
             };
+            message.ClearGatewayRequestOwner();
 
             _messagingTrace.OnCreateMessage(message);
             return message;
@@ -69,6 +70,7 @@ public Message CreateResponseMessage(Message request)
                 TimeToLive = request.TimeToLive,
                 RequestContextData = RequestContextExtensions.Export(_deepCopier),
             };
+            response.ApplyGatewayRequestOwner(request);
 
             _messagingTrace.OnCreateMessage(response);
             return response;

src/Orleans.Core/OrleansContracts.txt

@@ -16,6 +16,8 @@ interface [GrainInterfaceType("Orleans.ClientObservers.IClientGatewayObserver")]
 interface [GrainInterfaceType("Orleans.IMembershipTableSystemTarget")] Orleans.IMembershipTableSystemTarget [Version(0)]
 
 interface [GrainInterfaceType("Orleans.ISiloControl")] Orleans.ISiloControl [Version(0)]
+  16D39D91: CompleteGatewayRequest(Orleans.Runtime.GrainId, Orleans.Runtime.GrainId, Orleans.Runtime.CorrelationId) -> Task
+  B99FB859: DropDisconnectedClients(bool) -> Task
   45D07D09: ForceActivationCollection(System.TimeSpan) -> Task
   F388CED1: ForceGarbageCollection() -> Task
   0C7DBD0C: ForceRuntimeStatisticsCollection() -> Task
@@ -58,6 +60,7 @@ interface [GrainInterfaceType("Orleans.Runtime.IGrainCallCancellationExtension")
   FA239824: CancelRequestAsync(Orleans.Runtime.GrainId, Orleans.Runtime.CorrelationId) -> ValueTask
 
 interface [GrainInterfaceType("Orleans.Runtime.IManagementGrain")] Orleans.Runtime.IManagementGrain [Version(0)]
+  101564A8: DropDisconnectedClients(bool) -> Task
   329F9A1B: ForceActivationCollection(Orleans.Runtime.SiloAddress[], System.TimeSpan) -> Task
   54E6D1D1: ForceActivationCollection(System.TimeSpan) -> Task
   5922EB76: ForceGarbageCollection(Orleans.Runtime.SiloAddress[]) -> Task

src/Orleans.Core/Runtime/CallbackData.cs

@@ -211,7 +211,7 @@ public void DoCallback(Message response)
                 return;
             }
 
-            OrleansCallBackDataEvent.Instance.DoCallback(this.Message);
+            OrleansCallBackDataEvent.Instance.OnResponse(this.Message);
 
             this.stopwatch.Stop();
             DisposeCancellationRegistration();
@@ -221,6 +221,8 @@ public void DoCallback(Message response)
             ResponseCallback(response, this.context);
         }
 
+        public void OnResponse(Message response) => DoCallback(response);
+
         private bool TryComplete() => (Interlocked.Or(ref _state, StateCompleted) & StateCompleted) == 0;
 
         private void DisposeCancellationRegistration()

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -281,11 +281,16 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp
                 message.TargetSilo = systemTargetGrainId.GetSiloAddress();
             }
 
+            var responseTimeout = request.GetDefaultResponseTimeout() ?? this.sharedCallbackData.ResponseTimeout;
+            if (targetGrainId.IsClient())
+            {
+                message.SetGatewayRequestTimeout(responseTimeout);
+            }
+
             if (this.clientMessagingOptions.DropExpiredMessages && message.IsExpirableMessage())
             {
                 // don't set expiration for system target messages.
-                var ttl = request.GetDefaultResponseTimeout() ?? this.clientMessagingOptions.ResponseTimeout;
-                message.TimeToLive = ttl;
+                message.TimeToLive = responseTimeout;
             }
 
             if (!oneWay)

src/Orleans.Core/Runtime/RequestContextExtensions.cs

@@ -18,11 +18,20 @@ public static class RequestContextExtensions
         /// <param name="contextData">The context data.</param>
         public static void Import(Dictionary<string, object>? contextData)
         {
-            var values = contextData switch
+            Dictionary<string, object>? values = null;
+            if (contextData is { Count: > 0 })
             {
-                { Count: > 0 } => contextData.ToDictionary(kvp => kvp.Key, kvp => kvp.Value),
-                _ => null,
-            };
+                foreach (var (key, value) in contextData)
+                {
+                    if (Message.IsGatewayRequestContextHeader(key))
+                    {
+                        continue;
+                    }
+
+                    values ??= new(contextData.Count);
+                    values.Add(key, value);
+                }
+            }
 
             RequestContext.CallContextData.Value = new RequestContext.ContextProperties
             {

src/Orleans.Core/SystemTargetInterfaces/IManagementGrain.cs

@@ -141,6 +141,13 @@ public interface IManagementGrain : IGrainWithIntegerKey, IVersionManager
         /// <param name="hostsIds">The hosts to invoke the operation on.</param>
         /// <returns>A task representing the work performed.</returns>
         ValueTask ResetGrainCallFrequencies(SiloAddress[]? hostsIds = null);
+
+        /// <summary>
+        /// Instructs all gateways to drop defunct (disconnected and expired) clients.
+        /// </summary>
+        /// <param name="excludeRecent">If true, only clients that have been disconnected for longer than the configured client expiration time will be dropped.</param>
+        /// <returns>A task representing the work performed.</returns>
+        Task DropDisconnectedClients(bool excludeRecent);
     }
 
     /// <summary>

src/Orleans.Core/SystemTargetInterfaces/ISiloControl.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
+using Orleans.Concurrency;
 using Orleans.Providers;
 using Orleans.Runtime;
 
@@ -25,5 +26,8 @@ internal interface ISiloControl : ISystemTarget, IVersionManager
 
         Task<object?> SendControlCommandToProvider<T>(string providerName, int command, object? arg) where T : IControllable;
         Task<List<GrainId>> GetActiveGrains(GrainType grainType);
+        [OneWay, AlwaysInterleave]
+        Task CompleteGatewayRequest(GrainId clientId, GrainId sourceId, CorrelationId correlationId);
+        Task DropDisconnectedClients(bool excludeRecent);
     }
 }