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/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/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,82 @@ 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);
+
+            // 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);
+                if (_clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var clientConnectionContext))
+                {
+                    var instanceId = clientConnectionContext.InstanceId;
+                    var task = _clientInvocationManager.Caller.AddInvocation<T>(connectionId, invocationId, instanceId, cancellationToken);
+                    return await task;
+                }
+            }
+            catch
+            {
+                // Use `TryCompleteResult` to remove the failed invoation.
+                // If the invocation was not added by caller, it will be ignored.
+                // The content of error message is useless and will not be exposed to client.
+                _clientInvocationManager.Caller.TryCompleteResult(connectionId, CompletionMessage.WithError(invocationId, "Invocation Failed"));
+                throw;
+            }
+
+            // when condition `_clientConnectionManager.ClientConnections.TryGetValue` is false
+            throw new InvalidOperationException($"ConnectionId {connectionId} is invalid.");
+        }
+
+        // Only route server will reach here
+        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))
+            {
+                // Block unknown `results` which belongs to neither Caller nor Router
+                // `TryCompletionResult` returns false when the corresponding invocation is not existing.
+                if (!_clientInvocationManager.Caller.TryCompleteResult(connectionId, result) 
+                    && !_clientInvocationManager.Router.TryCompleteResult(connectionId, result))
+                {
+                    return; 
+                }
+
+                var protocol = clientConnectionContext.Protocol;
+                var message = AppendMessageTracingId(new ClientCompletionMessage(result.InvocationId, connectionId, _callerId, protocol, SerializeCompletionMessage(result, protocol)));
+                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();

src/Microsoft.Azure.SignalR/ServerConnections/ServiceConnection.cs

@@ -41,6 +41,8 @@ internal partial class ServiceConnection : ServiceConnectionBase
 
         private readonly ConnectionDelegate _connectionDelegate;
 
+        private readonly IClientInvocationManager _clientInvocationManager;
+
         public Action<HttpContext> ConfigureContext { get; set; }
 
         public ServiceConnection(IServiceProtocol serviceProtocol,
@@ -54,6 +56,7 @@ public ServiceConnection(IServiceProtocol serviceProtocol,
                                  HubServiceEndpoint endpoint,
                                  IServiceMessageHandler serviceMessageHandler,
                                  IServiceEventHandler serviceEventHandler,
+                                 IClientInvocationManager clientInvocationManager,
                                  ServiceConnectionType connectionType = ServiceConnectionType.Default,
                                  GracefulShutdownMode mode = GracefulShutdownMode.Off,
                                  int closeTimeOutMilliseconds = DefaultCloseTimeoutMilliseconds
@@ -64,6 +67,7 @@ public ServiceConnection(IServiceProtocol serviceProtocol,
             _connectionDelegate = connectionDelegate;
             _clientConnectionFactory = clientConnectionFactory;
             _closeTimeOutMilliseconds = closeTimeOutMilliseconds;
+            _clientInvocationManager = clientInvocationManager;
         }
 
         protected override Task<ConnectionContext> CreateConnection(string target = null)
@@ -188,7 +192,31 @@ protected override async Task OnClientMessageAsync(ConnectionDataMessage connect
             }
         }
 
-        private async Task ProcessClientConnectionAsync(ClientConnectionContext connection)
+        protected override Task DispatchMessageAsync(ServiceMessage message)
+        {
+            return message switch
+            {
+                PingMessage pingMessage => OnPingMessageAsync(pingMessage),
+                ClientInvocationMessage clientInvocationMessage => OnClientInvocationAsync(clientInvocationMessage),
+                ServiceMappingMessage serviceMappingMessage => OnServiceMappingAsync(serviceMappingMessage),
+                ClientCompletionMessage clientCompletionMessage => OnClientCompletionAsync(clientCompletionMessage),
+                ErrorCompletionMessage errorCompletionMessage => OnErrorCompletionAsync(errorCompletionMessage),
+                _ => base.DispatchMessageAsync(message)
+            };
+        }
+
+        protected override Task OnPingMessageAsync(PingMessage pingMessage)
+        {
+#if NET7_0_OR_GREATER
+            if (RuntimeServicePingMessage.TryGetOffline(pingMessage, out var instanceId))
+            {
+                _clientInvocationManager.Caller.CleanupInvocationsByInstance(instanceId);
+            }
+#endif
+            return base.OnPingMessageAsync(pingMessage);
+        }
+
+    private async Task ProcessClientConnectionAsync(ClientConnectionContext connection)
         {
             try
             {
@@ -434,7 +462,34 @@ private ClientConnectionContext RemoveClientConnection(string connectionId)
         {
             _connectionIds.TryRemove(connectionId, out _);
             _clientConnectionManager.TryRemoveClientConnection(connectionId, out var connection);
+#if NET7_0_OR_GREATER
+            _clientInvocationManager.Router.CleanupInvocationsByConnection(connectionId);
+#endif
             return connection;
         }
+
+        private Task OnClientInvocationAsync(ClientInvocationMessage message)
+        {
+            _clientInvocationManager.Router.AddInvocation(message.ConnectionId, message.InvocationId, message.CallerServerId, default);
+            return Task.CompletedTask;
+        }
+
+        private Task OnServiceMappingAsync(ServiceMappingMessage message)
+        {
+            _clientInvocationManager.Caller.AddServiceMapping(message);
+            return Task.CompletedTask;
+        }
+
+        private Task OnClientCompletionAsync(ClientCompletionMessage clientCompletionMessage)
+        {
+            _clientInvocationManager.Caller.TryCompleteResult(clientCompletionMessage.ConnectionId, clientCompletionMessage);
+            return Task.CompletedTask;
+        }
+
+        private Task OnErrorCompletionAsync(ErrorCompletionMessage errorCompletionMessage)
+        {
+            _clientInvocationManager.Caller.TryCompleteResult(errorCompletionMessage.ConnectionId, errorCompletionMessage);
+            return Task.CompletedTask;
+        }
     }
 }

src/Microsoft.Azure.SignalR/ServerConnections/ServiceConnectionFactory.cs

@@ -16,6 +16,7 @@ internal class ServiceConnectionFactory : IServiceConnectionFactory
         private readonly IClientConnectionFactory _clientConnectionFactory;
         private readonly IServerNameProvider _nameProvider;
         private readonly IServiceEventHandler _serviceEventHandler;
+        private readonly IClientInvocationManager _clientInvocationManager;
 
         public GracefulShutdownMode ShutdownMode { get; set; } = GracefulShutdownMode.Off;
 
@@ -29,7 +30,8 @@ public ServiceConnectionFactory(
             ConnectionDelegate connectionDelegate,
             IClientConnectionFactory clientConnectionFactory,
             IServerNameProvider nameProvider,
-            IServiceEventHandler serviceEventHandler)
+            IServiceEventHandler serviceEventHandler,
+            IClientInvocationManager clientInvocationManager)
         {
             _serviceProtocol = serviceProtocol;
             _clientConnectionManager = clientConnectionManager;
@@ -39,6 +41,7 @@ public ServiceConnectionFactory(
             _clientConnectionFactory = clientConnectionFactory;
             _nameProvider = nameProvider;
             _serviceEventHandler = serviceEventHandler;
+            _clientInvocationManager = clientInvocationManager;
         }
 
         public virtual IServiceConnection Create(HubServiceEndpoint endpoint, IServiceMessageHandler serviceMessageHandler, ServiceConnectionType type)
@@ -55,6 +58,7 @@ public virtual IServiceConnection Create(HubServiceEndpoint endpoint, IServiceMe
                 endpoint,
                 serviceMessageHandler,
                 _serviceEventHandler,
+                _clientInvocationManager,
                 type,
                 ShutdownMode
             )