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>
     <MicrosoftAspNetCoreHttpConnectionsCommonPackageVersion>1.0.4</MicrosoftAspNetCoreHttpConnectionsCommonPackageVersion>
     <MicrosoftAspNetCoreHttpConnectionsCommonPackage3_0Version>3.0.0</MicrosoftAspNetCoreHttpConnectionsCommonPackage3_0Version>
     <MicrosoftAspNetCoreHttpConnectionsCommonPackage3_1Version>3.1.9</MicrosoftAspNetCoreHttpConnectionsCommonPackage3_1Version>

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

@@ -42,6 +42,7 @@ public ServiceConnection(
             ILoggerFactory loggerFactory,
             IServiceMessageHandler serviceMessageHandler,
             IServiceEventHandler serviceEventHandler,
+            IClientInvocationManager clientInvocationManager = null,
             ServiceConnectionType connectionType = ServiceConnectionType.Default)
             : base(
                   serviceProtocol,
@@ -51,6 +52,7 @@ public ServiceConnection(
                   serviceMessageHandler,
                   serviceEventHandler,
                   connectionType,
+                  clientInvocationManager,
                   loggerFactory?.CreateLogger<ServiceConnection>())
         {
             _connectionFactory = connectionFactory;

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

@@ -41,6 +41,7 @@ public IServiceConnection Create(HubServiceEndpoint endpoint, IServiceMessageHan
                 _logger,
                 serviceMessageHandler,
                 _serviceEventHandler,
+                null,
                 type);
         }
     }

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

@@ -7,7 +7,6 @@
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
-
 using Microsoft.AspNetCore.Connections;
 using Microsoft.Azure.SignalR.Common;
 using Microsoft.Azure.SignalR.Protocol;
@@ -74,6 +73,8 @@ internal abstract class ServiceConnectionBase : IServiceConnection
 
         protected IServiceProtocol ServiceProtocol { get; }
 
+        protected IClientInvocationManager ClientInvocationManager { get;  }
+
         private ConnectionContext _connectionContext;
 
         public event Action<StatusChange> ConnectionStatusChanged;
@@ -111,6 +112,7 @@ protected ServiceConnectionBase(
             IServiceMessageHandler serviceMessageHandler,
             IServiceEventHandler serviceEventHandler,
             ServiceConnectionType connectionType,
+            IClientInvocationManager clientInvocationManager,
             ILogger logger,
             GracefulShutdownMode mode = GracefulShutdownMode.Off)
         {
@@ -132,6 +134,7 @@ protected ServiceConnectionBase(
             Logger = logger ?? throw new ArgumentNullException(nameof(logger));
             _serviceMessageHandler = serviceMessageHandler;
             _serviceEventHandler = serviceEventHandler;
+            ClientInvocationManager = clientInvocationManager ?? throw new ArgumentNullException(nameof(clientInvocationManager));
         }
 
         /// <summary>
@@ -300,6 +303,9 @@ protected Task OnPingMessageAsync(PingMessage pingMessage)
             if (RuntimeServicePingMessage.TryGetOffline(pingMessage, out var instanceId))
             {
                 Log.ReceivedInstanceOfflinePing(Logger, instanceId);
+#if NET7_0_OR_GREATER
+                ClientInvocationManager.Caller.CleanupInvocationsByInstance(instanceId);
+#endif
                 return CleanupClientConnections(instanceId);
             }
             if (RuntimeServicePingMessage.IsFinAck(pingMessage))
@@ -310,6 +316,30 @@ protected Task OnPingMessageAsync(PingMessage pingMessage)
             return _serviceMessageHandler.HandlePingAsync(pingMessage);
         }
 
+        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;
+        }
+
         protected Task OnAckMessageAsync(AckMessage ackMessage)
         {
             _serviceMessageHandler.HandleAck(ackMessage);
@@ -562,6 +592,10 @@ private Task DispatchMessageAsync(ServiceMessage message)
                 AckMessage ackMessage => OnAckMessageAsync(ackMessage),
                 ServiceEventMessage eventMessage => OnEventMessageAsync(eventMessage),
                 AccessKeyResponseMessage keyMessage => OnAccessKeyMessageAsync(keyMessage),
+                ClientInvocationMessage clientInvocationMessage => OnClientInvocationAsync(clientInvocationMessage),
+                ServiceMappingMessage serviceMappingMessage => OnServiceMappingAsync(serviceMappingMessage),
+                ClientCompletionMessage clientCompletionMessage => OnClientCompletionAsync(clientCompletionMessage),
+                ErrorCompletionMessage errorCompletionMessage => OnErrorCompletionAsync(errorCompletionMessage),
                 _ => Task.CompletedTask,
             };
         }

src/Microsoft.Azure.SignalR.Management/DependencyInjectionExtensions.cs

@@ -87,7 +87,7 @@ private static IServiceCollection AddSignalRServiceCore(this IServiceCollection
             // Remove the JsonHubProtocol and add new one.
             // On .NET Standard 2.0, registering multiple hub protocols with the same name is forbidden.
             services.Replace(ServiceDescriptor.Singleton<IHubProtocol>(sp =>
-            {
+                {
                 var objectSerializer = sp.GetRequiredService<IOptions<ServiceManagerOptions>>().Value.ObjectSerializer;
                 return objectSerializer != null ? new JsonObjectSerializerHubProtocol(objectSerializer) : new JsonHubProtocol();
             }));

src/Microsoft.Azure.SignalR.Management/WebsocketsHubLifetimeManager.cs

@@ -16,8 +16,8 @@ internal class WebSocketsHubLifetimeManager<THub> : ServiceLifetimeManagerBase<T
         private readonly IOptions<ServiceManagerOptions> _serviceManagerOptions;
 
         public WebSocketsHubLifetimeManager(IServiceConnectionManager<THub> serviceConnectionManager, IHubProtocolResolver protocolResolver,
-            IOptions<HubOptions> globalHubOptions, IOptions<HubOptions<THub>> hubOptions, ILoggerFactory loggerFactory, IOptions<ServiceManagerOptions> serviceManagerOptions) :
-            base(serviceConnectionManager, protocolResolver, globalHubOptions, hubOptions, loggerFactory?.CreateLogger(nameof(WebSocketsHubLifetimeManager<Hub>)))
+            IOptions<HubOptions> globalHubOptions, IOptions<HubOptions<THub>> hubOptions, IClientInvocationManager clientInvocationManager, IClientConnectionManager clientConnectionManager, ILoggerFactory loggerFactory, IOptions<ServiceManagerOptions> serviceManagerOptions) :
+            base(serviceConnectionManager, protocolResolver, globalHubOptions, hubOptions, new DefaultServerNameProvider(), clientInvocationManager, clientConnectionManager, loggerFactory?.CreateLogger(nameof(WebSocketsHubLifetimeManager<Hub>)))
         {
             _serviceManagerOptions = serviceManagerOptions ?? throw new ArgumentNullException(nameof(serviceManagerOptions));
         }

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

@@ -19,8 +19,6 @@ internal class ServiceLifetimeManager<THub> : ServiceLifetimeManagerBase<THub> w
         private const string MarkerNotConfiguredError =
             "'AddAzureSignalR(...)' was called without a matching call to 'IApplicationBuilder.UseAzureSignalR(...)'.";
 
-        private readonly IClientConnectionManager _clientConnectionManager;
-
         public ServiceLifetimeManager(
             IServiceConnectionManager<THub> serviceConnectionManager,
             IClientConnectionManager clientConnectionManager,
@@ -29,12 +27,18 @@ 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,
+                  nameProvider,
+                  clientInvocationManager,
+                  clientConnectionManager,
+                  logger)
         {
             // after core 3.0 UseAzureSignalR() is not required.
 #if NETSTANDARD2_0
@@ -43,8 +47,6 @@ 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);
@@ -74,7 +76,7 @@ public override async Task SendConnectionAsync(string connectionId, string metho
                 throw new ArgumentException(NullOrEmptyStringErrorMessage, nameof(methodName));
             }
 
-            if (_clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var serviceConnectionContext))
+            if (ClientConnectionManager.ClientConnections.TryGetValue(connectionId, out var serviceConnectionContext))
             {
                 var message = CreateMessage(connectionId, methodName, args, serviceConnectionContext);
                 var messageWithTracingId = (IMessageWithTracingId)message;