Add main logic for Client Invocation

build/dependencies.props

@@ -8,6 +8,7 @@
     <MicrosoftAspNetCoreHttpConnectionsClientPackage3_0Version>3.0.0</MicrosoftAspNetCoreHttpConnectionsClientPackage3_0Version>
     <MicrosoftAspNetCoreHttpConnectionsClientPackage3_1Version>3.1.9</MicrosoftAspNetCoreHttpConnectionsClientPackage3_1Version>
     <MicrosoftAspNetCoreHttpConnectionsClientPackage5_0Version>5.0.1</MicrosoftAspNetCoreHttpConnectionsClientPackage5_0Version>
+    <MicrosoftAspNetCoreHttpConnectionsClientPackage7_0Version>7.0.0-preview.7.22376.6</MicrosoftAspNetCoreHttpConnectionsClientPackage7_0Version>
     <MicrosoftAspNetCoreHttpConnectionsClientPackage6_0Version>6.0.10</MicrosoftAspNetCoreHttpConnectionsClientPackage6_0Version>
     <MicrosoftAspNetCoreHttpConnectionsCommonPackageVersion>1.0.4</MicrosoftAspNetCoreHttpConnectionsCommonPackageVersion>
     <MicrosoftAspNetCoreHttpConnectionsCommonPackage3_0Version>3.0.0</MicrosoftAspNetCoreHttpConnectionsCommonPackage3_0Version>

src/Microsoft.Azure.SignalR.Common/ClientInvocation/ICallerClientResultsManager.cs

@@ -17,10 +17,9 @@ internal interface ICallerClientResultsManager : IClientResultsManager
         /// <typeparam name="T"></typeparam>
         /// <param name="connectionId"></param>
         /// <param name="invocationId"></param>
-        /// <param name="instanceId"> The InstanceId of target client the caller server knows when this method is called. If the target client is managed by the caller server, the caller server knows the InstanceId of target client and this parameter is not null. Otherwise, this parameter is null. </param>
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
-        Task<T> AddInvocation<T>(string connectionId, string invocationId, string instanceId, CancellationToken cancellationToken);
+        Task<T> AddInvocation<T>(string connectionId, string invocationId, CancellationToken cancellationToken);
 
         void AddServiceMapping(ServiceMappingMessage serviceMappingMessage);
 
@@ -29,5 +28,7 @@ internal interface ICallerClientResultsManager : IClientResultsManager
         bool TryCompleteResult(string connectionId, ClientCompletionMessage message);
 
         bool TryCompleteResult(string connectionId, ErrorCompletionMessage message);
+
+        bool RemoveInvocation(string invocationId);
     }
 }

src/Microsoft.Azure.SignalR.Common/ServiceConnections/ServiceConnectionBase.cs

@@ -295,7 +295,7 @@ protected Task OnServiceErrorAsync(ServiceErrorMessage serviceErrorMessage)
             return Task.CompletedTask;
         }
 
-        protected Task OnPingMessageAsync(PingMessage pingMessage)
+        protected virtual Task OnPingMessageAsync(PingMessage pingMessage)
         {
             if (RuntimeServicePingMessage.TryGetOffline(pingMessage, out var instanceId))
             {
@@ -550,7 +550,7 @@ private async Task ProcessIncomingAsync(ConnectionContext connection)
             }
         }
 
-        private Task DispatchMessageAsync(ServiceMessage message)
+        protected virtual Task DispatchMessageAsync(ServiceMessage message)
         {
             return message switch
             {

src/Microsoft.Azure.SignalR/ClientInvocation/CallerClientResultsManager.cs

@@ -31,7 +31,7 @@ public string GenerateInvocationId(string connectionId)
             return $"{connectionId}-{_clientResultManagerId}-{Interlocked.Increment(ref _lastInvocationId)}";
         }
 
-        public Task<T> AddInvocation<T>(string connectionId, string invocationId, string instanceId, CancellationToken cancellationToken)
+        public Task<T> AddInvocation<T>(string connectionId, string invocationId, CancellationToken cancellationToken)
         {
             var tcs = new TaskCompletionSourceWithCancellation<T>(
                 cancellationToken,
@@ -53,7 +53,7 @@ public Task<T> AddInvocation<T>(string connectionId, string invocationId, string
                         {
                             tcs.TrySetException(new Exception(completionMessage.Error));
                         }
-                    }) { RouterInstanceId = instanceId }
+                    }) { RouterInstanceId = null }
             );
             Debug.Assert(result);
 
@@ -70,14 +70,6 @@ public void AddServiceMapping(ServiceMappingMessage serviceMappingMessage)
                 {
                     invocation.RouterInstanceId = serviceMappingMessage.InstanceId;
                 }
-                else
-                {
-                    // do nothing
-                }
-            }
-            else
-            {
-                // do nothing
             }
         }
 
@@ -174,6 +166,8 @@ public bool TryGetInvocationReturnType(string invocationId, out Type type)
             return false;
         }
 
+        public bool RemoveInvocation(string invocationId) => _pendingInvocations.TryRemove(invocationId, out _);
+
         // Unused, here to honor the IInvocationBinder interface but should never be called
         public IReadOnlyList<Type> GetParameterTypes(string methodName) => throw new NotImplementedException();
 

src/Microsoft.Azure.SignalR/DependencyInjectionExtensions.cs

@@ -88,6 +88,11 @@ private static ISignalRServerBuilder AddAzureSignalRCore(this ISignalRServerBuil
                 .AddSingleton<IClientConnectionFactory, ClientConnectionFactory>()
                 .AddSingleton<IHostedService, HeartBeat>()
                 .AddSingleton<IAccessKeySynchronizer, AccessKeySynchronizer>()
+#if NET7_0_OR_GREATER
+                .AddSingleton<IClientInvocationManager, ClientInvocationManager>()
+#else
+                .AddSingleton<IClientInvocationManager, DummyClientInvocationManager>()
+#endif
                 .AddSingleton(typeof(NegotiateHandler<>));
 
             // If a custom router is added, do not add the default router

src/Microsoft.Azure.SignalR/HubHost/ServiceHubDispatcher.cs

@@ -30,6 +30,7 @@ internal class ServiceHubDispatcher<THub> where THub : Hub
         private readonly IEndpointRouter _router;
         private readonly string _hubName;
         private readonly IServiceEventHandler _serviceEventHandler;
+        private readonly IClientInvocationManager _clientInvocationManager;
 
         protected readonly IServerNameProvider _nameProvider;
 
@@ -45,6 +46,7 @@ public ServiceHubDispatcher(
             IServerNameProvider nameProvider,
             ServerLifetimeManager serverLifetimeManager,
             IClientConnectionFactory clientConnectionFactory,
+            IClientInvocationManager clientInvocationManager,
             IServiceEventHandler serviceEventHandler)
         {
             _serviceProtocol = serviceProtocol;
@@ -62,6 +64,7 @@ public ServiceHubDispatcher(
             _nameProvider = nameProvider;
             _hubName = typeof(THub).Name;
             _serviceEventHandler = serviceEventHandler;
+            _clientInvocationManager = clientInvocationManager;
 
             serverLifetimeManager?.Register(ShutdownAsync);
         }
@@ -150,7 +153,8 @@ internal virtual ServiceConnectionFactory GetServiceConnectionFactory(
                 connectionDelegate,
                 _clientConnectionFactory,
                 _nameProvider,
-                _serviceEventHandler)
+                _serviceEventHandler,
+                _clientInvocationManager)
             {
                 ConfigureContext = contextConfig,
                 ShutdownMode = _options.GracefulShutdown.Mode

src/Microsoft.Azure.SignalR/HubHost/ServiceLifetimeManager.cs

@@ -3,10 +3,12 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR;
+using Microsoft.AspNetCore.SignalR.Protocol;
 using Microsoft.Azure.SignalR.Common;
 using Microsoft.Azure.SignalR.Protocol;
 using Microsoft.Extensions.Logging;
@@ -18,8 +20,9 @@ internal class ServiceLifetimeManager<THub> : ServiceLifetimeManagerBase<THub> w
     {
         private const string MarkerNotConfiguredError =
             "'AddAzureSignalR(...)' was called without a matching call to 'IApplicationBuilder.UseAzureSignalR(...)'.";
-
+        private readonly IClientInvocationManager _clientInvocationManager;
         private readonly IClientConnectionManager _clientConnectionManager;
+        private readonly string _callerId;
 
         public ServiceLifetimeManager(
             IServiceConnectionManager<THub> serviceConnectionManager,
@@ -29,12 +32,15 @@ public ServiceLifetimeManager(
             AzureSignalRMarkerService marker,
             IOptions<HubOptions> globalHubOptions,
             IOptions<HubOptions<THub>> hubOptions,
-            IBlazorDetector blazorDetector)
+            IBlazorDetector blazorDetector,
+            IServerNameProvider nameProvider,
+            IClientInvocationManager clientInvocationManager)
             : base(
                   serviceConnectionManager,
                   protocolResolver,
                   globalHubOptions,
-                  hubOptions, logger)
+                  hubOptions,
+                  logger)
         {
             // after core 3.0 UseAzureSignalR() is not required.
 #if NETSTANDARD2_0
@@ -43,12 +49,19 @@ public ServiceLifetimeManager(
                 throw new InvalidOperationException(MarkerNotConfiguredError);
             }
 #endif
-            _clientConnectionManager = clientConnectionManager;
-
             if (hubOptions.Value.SupportedProtocols != null && hubOptions.Value.SupportedProtocols.Any(x => x.Equals(Constants.Protocol.BlazorPack, StringComparison.OrdinalIgnoreCase)))
             {
                 blazorDetector?.TrySetBlazor(typeof(THub).Name, true);
             }
+
+            if (nameProvider == null)
+            {
+                throw new ArgumentNullException(nameof(nameProvider));
+            }
+            _callerId = nameProvider.GetName();
+
+            _clientInvocationManager = clientInvocationManager ?? throw new ArgumentNullException(nameof(clientInvocationManager));
+            _clientConnectionManager = clientConnectionManager ?? throw new ArgumentNullException(nameof(clientConnectionManager));
         }
 
         public override Task OnConnectedAsync(HubConnectionContext connection)
@@ -103,6 +116,85 @@ public override async Task SendConnectionAsync(string connectionId, string metho
             }
         }
 
+#if NET7_0_OR_GREATER
+        public override async Task<T> InvokeConnectionAsync<T>(string connectionId, string methodName, object[] args, CancellationToken cancellationToken = default)
+        {
+            if (IsInvalidArgument(connectionId))
+            {
+                throw new ArgumentNullException(nameof(connectionId));
+            }
+
+            if (IsInvalidArgument(methodName))
+            {
+                throw new ArgumentNullException(nameof(methodName));
+            }
+
+            var invocationId = _clientInvocationManager.Caller.GenerateInvocationId(connectionId);
+            var task = _clientInvocationManager.Caller.AddInvocation<T>(connectionId, invocationId, cancellationToken);
+
+            // Exception handling follows https://source.dot.net/#Microsoft.AspNetCore.SignalR.Core/DefaultHubLifetimeManager.cs,349
+            try
+            {
+                var message = AppendMessageTracingId(new ClientInvocationMessage(invocationId, connectionId, _callerId, SerializeAllProtocols(methodName, args, invocationId)));
+                await WriteAsync(message);
+                return await task;
+            }
+            catch
+            {
+                _clientInvocationManager.Caller.RemoveInvocation(invocationId);
+                throw;
+            }
+        }
+
+        public override async Task SetConnectionResultAsync(string connectionId, CompletionMessage result)
+        {
+            if (IsInvalidArgument(connectionId))
+            {
+                throw new ArgumentException(NullOrEmptyStringErrorMessage, nameof(connectionId));
+            }
+            if (_clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var clientConnectionContext))
+            {
+                // Determine which manager (Caller / Router) the `result` belongs to.
+                // `TryCompletionResult` returns false when the corresponding invocation is not existing.
+                IClientResultsManager clientResultsManager = null;
+                if (_clientInvocationManager.Caller.TryCompleteResult(connectionId, result))
+                {
+                    clientResultsManager = _clientInvocationManager.Caller;
+                }
+                if (_clientInvocationManager.Router.TryCompleteResult(connectionId, result))
+                {
+                    clientResultsManager = _clientInvocationManager.Router;
+                }
+
+                // Block unknown `results` which belongs to neither Caller nor Router
+                if (clientResultsManager != null)
+                {
+                    var protocol = clientConnectionContext.Protocol;
+                    // For router server, it should send a ClientCompletionMessage with accurate payload content, which is necessary for the caller server.
+                    // For caller server, the only purpose of sending ClientCompletionMessage is to inform service to cleanup the invocation, which means only InvocationId and ConnectionId make sense. To avoid serialization for useless payload, we assign payload with empty bytes.
+                    var payload = clientResultsManager == _clientInvocationManager.Router
+                                ? SerializeCompletionMessage(result, protocol)
+                                : Array.Empty<byte>();
+
+                    var message = AppendMessageTracingId(new ClientCompletionMessage(result.InvocationId, connectionId, _callerId, protocol, payload));
+                    await WriteAsync(message);
+                }
+            }
+        }
+
+        public override bool TryGetReturnType(string invocationId, [NotNullWhen(true)] out Type type)
+        {
+            if (_clientInvocationManager.Router.ContainsInvocation(invocationId))
+            {
+                return _clientInvocationManager.Router.TryGetInvocationReturnType(invocationId, out type);
+            }
+            else
+            {
+                return _clientInvocationManager.Caller.TryGetInvocationReturnType(invocationId, out type);
+            }
+        }
+#endif
+
         private MultiConnectionDataMessage CreateMessage(string connectionId, string methodName, object[] args, ClientConnectionContext serviceConnectionContext)
         {
             IDictionary<string, ReadOnlyMemory<byte>> payloads;

src/Microsoft.Azure.SignalR/HubHost/ServiceLifetimeManagerBase.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Diagnostics.CodeAnalysis;
 using Microsoft.AspNetCore.SignalR;
 using Microsoft.AspNetCore.SignalR.Protocol;
 using Microsoft.Azure.SignalR.Protocol;
@@ -283,10 +284,18 @@ protected static bool IsInvalidArgument(IReadOnlyList<object> list)
             return list == null;
         }
 
-        protected IDictionary<string, ReadOnlyMemory<byte>> SerializeAllProtocols(string method, object[] args)
+        protected IDictionary<string, ReadOnlyMemory<byte>> SerializeAllProtocols(string method, object[] args, string invocationId = null)
         {
             var payloads = new Dictionary<string, ReadOnlyMemory<byte>>();
-            var message = new InvocationMessage(method, args);
+            InvocationMessage message;
+            if (invocationId == null)
+            {
+                message = new InvocationMessage(method, args);
+            }
+            else
+            {
+                message = new InvocationMessage(invocationId, method, args);
+            }
             var serializedHubMessages = _messageSerializer.SerializeMessage(message);
             foreach (var serializedMessage in serializedHubMessages)
             {
@@ -298,6 +307,9 @@ protected IDictionary<string, ReadOnlyMemory<byte>> SerializeAllProtocols(string
         protected ReadOnlyMemory<byte> SerializeProtocol(string protocol, string method, object[] args) =>
             _messageSerializer.SerializeMessage(protocol, new InvocationMessage(method, args));
 
+        protected ReadOnlyMemory<byte> SerializeCompletionMessage(CompletionMessage message, string protocol) =>
+            _messageSerializer.SerializeMessage(protocol, message);
+
         protected virtual T AppendMessageTracingId<T>(T message) where T : ServiceMessage, IMessageWithTracingId
         {
             return message.WithTracingId();