graceful shutdown

Azure · Nov 16, 2020 · dc9c668 · dc9c668
1 parent 4dbc6d1
commit dc9c668
Show file tree

Hide file tree

Showing 7 changed files with 127 additions and 24 deletions.
diff --git a/samples/ChatSample/ChatSample/Startup.cs b/samples/ChatSample/ChatSample/Startup.cs
@@ -1,6 +1,7 @@
 using System;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
+using Microsoft.AspNetCore.SignalR;
 using Microsoft.Azure.SignalR;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
@@ -18,7 +19,12 @@ public void ConfigureServices(IServiceCollection services)
                 .AddAzureSignalR(option =>
                 {
                     option.GracefulShutdown.Mode = GracefulShutdownMode.WaitForClientsClose;
-                    option.GracefulShutdown.Timeout = TimeSpan.FromSeconds(10);
+                    option.GracefulShutdown.Timeout = TimeSpan.FromSeconds(60);
+
+                    option.GracefulShutdown.Add<Chat>(async (c) =>
+                    {
+                        await c.Clients.All.SendAsync("shutdown");
+                    });
                 })
                 .AddMessagePackProtocol();
         }

diff --git a/samples/ChatSample/ChatSample/wwwroot/index.html b/samples/ChatSample/ChatSample/wwwroot/index.html
@@ -28,7 +28,17 @@ <h2 class="text-center" style="margin-top: 0; padding-top: 30px; padding-bottom:
             <div class="modal-content">
                 <div class="modal-header">
                     <div>Connection Error...</div>
-                    <div><strong style="font-size: 1.5em;">Hit Refresh/F5</strong> to rejoin. ;)</div>
+                    <div><strong style="font-size: 1.5em;">Hit Refresh/F5</strong> to rejoin. :)</div>
+                </div>
+            </div>
+        </div>
+    </div>
+    <div class="modal alert alert-success fade" id="closeModal" tabindex="-1" role="dialog" aria-labelledby="myModalLabel">
+        <div class="modal-dialog" role="document">
+            <div class="modal-content">
+                <div class="modal-header">
+                    <div>Connection Closed</div>
+                    <div><strong style="font-size: 1.5em;">Hit Refresh/F5</strong> to rejoin. :)</div>
                 </div>
             </div>
         </div>
@@ -97,9 +107,15 @@ <h2 class="text-center" style="margin-top: 0; padding-top: 30px; padding-bottom:
                     messageBox.appendChild(messageEntry);
                     messageBox.scrollTop = messageBox.scrollHeight;
                 };
+
+                var close = function (name, message) {
+                    setTimeout(() => connection.stop(), 3000);
+                }
+
                 // Create a function that the hub can call to broadcast messages.
                 connection.on('broadcastMessage', messageCallback);
                 connection.on('echo', messageCallback);
+                connection.on('exit', close);
                 connection.onclose(onConnectionError);
             }
 
@@ -138,8 +154,10 @@ <h2 class="text-center" style="margin-top: 0; padding-top: 30px; padding-bottom:
             function onConnectionError(error) {
                 if (error && error.message) {
                     console.error(error.message);
+                    var modal = document.getElementById('myModal');
+                } else {
+                    var modal = document.getElementById('closeModal');
                 }
-                var modal = document.getElementById('myModal');
                 modal.classList.add('in');
                 modal.style = 'display: block;';
             }

diff --git a/src/Microsoft.Azure.SignalR.AspNet/HubHost/ServiceHubDispatcher.cs b/src/Microsoft.Azure.SignalR.AspNet/HubHost/ServiceHubDispatcher.cs
@@ -3,9 +3,7 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Threading.Tasks;
-using Microsoft.Azure.SignalR.Protocol;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 

diff --git a/src/Microsoft.Azure.SignalR/GracefulShutdownOptions.cs b/src/Microsoft.Azure.SignalR/GracefulShutdownOptions.cs
@@ -1,20 +1,81 @@
 using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.SignalR;
 
 namespace Microsoft.Azure.SignalR
 {
     public class GracefulShutdownOptions
     {
-        /// <summary>
-        /// Specifies the timeout of a graceful shutdown process (in seconds). 
-        /// Default value is 30 seconds.
-        /// </summary>
-        public TimeSpan Timeout { get; set; } = Constants.Periods.DefaultShutdownTimeout;
+        private Dictionary<string, List<object>> _dict = new Dictionary<string, List<object>>();
 
         /// <summary>
         /// Specifies if the client-connection assigned to this server can be migrated to another server.
         /// Default value is 0.
         /// 1: Migrate client-connection if the server was shutdown gracefully.
         /// </summary>
         public GracefulShutdownMode Mode { get; set; } = GracefulShutdownMode.Off;
+
+        /// <summary>
+        /// Specifies the timeout of a graceful shutdown process (in seconds). 
+        /// Default value is 30 seconds.
+        /// </summary>
+        public TimeSpan Timeout { get; set; } = Constants.Periods.DefaultShutdownTimeout;
+
+        private List<object> GetMethods<THub>() {
+            var name = typeof(THub).Name;
+            if (!_dict.ContainsKey(typeof(THub).Name))
+            {
+                _dict.Add(name, new List<object>());
+            }
+            return _dict[name];
+        }
+
+        public void Add<THub>(Func<IHubContext<THub>, Task> func) where THub : Hub
+        {
+            GetMethods<THub>().Add(func);
+        }
+
+        public void Add<THub>(Action<IHubContext<THub>> action) where THub : Hub
+        {
+            GetMethods<THub>().Add(action);
+        }
+
+        public void Add<THub>(Func<Task> func)
+        {
+            GetMethods<THub>().Add(func);
+        }
+
+        public void Add<THub>(Action action)
+        {
+            GetMethods<THub>().Add(action);
+        }
+
+        internal async Task OnShutdown<THub>(IHubContext<THub> context) where THub : Hub
+        {
+            var name = typeof(THub).Name;
+            if (_dict.TryGetValue(name, out var methods))
+            {
+                foreach (var method in methods)
+                {
+                    if (method is Action<IHubContext<THub>> action)
+                    {
+                        action(context);
+                    }
+                    else if (method is Action action2)
+                    {
+                        action2();
+                    }
+                    else if (method is Func<IHubContext<THub>, Task> func)
+                    {
+                        await func(context);
+                    }
+                    else if (method is Func<Task> func2)
+                    {
+                        await func2();
+                    }
+                }
+            }
+        }
     }
 }
diff --git a/src/Microsoft.Azure.SignalR/HubHost/ServiceHubDispatcher.cs b/src/Microsoft.Azure.SignalR/HubHost/ServiceHubDispatcher.cs
@@ -8,6 +8,7 @@
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.SignalR;
 using Microsoft.Azure.SignalR.Protocol;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
@@ -17,6 +18,8 @@ internal class ServiceHubDispatcher<THub> where THub : Hub
     {
         private static readonly string Name = $"ServiceHubDispatcher<{typeof(THub).FullName}>";
 
+        private IHubContext<THub> Context { get; }
+
         private readonly ILoggerFactory _loggerFactory;
         private readonly ILogger<ServiceHubDispatcher<THub>> _logger;
         private readonly ServiceOptions _options;
@@ -31,6 +34,7 @@ internal class ServiceHubDispatcher<THub> where THub : Hub
 
         public ServiceHubDispatcher(
             IServiceProtocol serviceProtocol,
+            IHubContext<THub> context,
             IServiceConnectionManager<THub> serviceConnectionManager,
             IClientConnectionManager clientConnectionManager,
             IServiceEndpointManager serviceEndpointManager,
@@ -47,6 +51,8 @@ public ServiceHubDispatcher(
             _serviceEndpointManager = serviceEndpointManager;
             _options = options != null ? options.Value : throw new ArgumentNullException(nameof(options));
 
+            Context = context;
+
             _router = router ?? throw new ArgumentNullException(nameof(router));
             _loggerFactory = loggerFactory ?? throw new ArgumentNullException(nameof(loggerFactory));
             _logger = loggerFactory.CreateLogger<ServiceHubDispatcher<THub>>();
@@ -71,7 +77,8 @@ public void Start(ConnectionDelegate connectionDelegate, Action<HttpContext> con
 
         public async Task ShutdownAsync()
         {
-            if (_options.GracefulShutdown.Mode == GracefulShutdownMode.Off)
+            var options = _options.GracefulShutdown;
+            if (options.Mode == GracefulShutdownMode.Off)
             {
                 return;
             }
@@ -80,31 +87,32 @@ public async Task ShutdownAsync()
             {
                 var source = new CancellationTokenSource(_options.GracefulShutdown.Timeout);
 
-                _logger.LogInformation("[GracefulShutdown] Started.");
+                _logger.LogInformation($"[GracefulShutdown][{_hubName}] Setting server offline...");
 
                 await Task.WhenAny(
-                    OfflineAndWaitForCompletedAsync(_options.GracefulShutdown.Mode),
+                    _serviceConnectionManager.OfflineAsync(options.Mode),
                     Task.Delay(Timeout.InfiniteTimeSpan, source.Token)
                 );
 
+                _logger.LogInformation("[GracefulShutdown] Triggering shutdown hooks...");
+
+                await options.OnShutdown(Context);
+
+                _logger.LogInformation($"[GracefulShutdown][{_hubName}] Waiting client connections to complete...");
+
                 await Task.WhenAny(
-                    _serviceConnectionManager.StopAsync(),
+                    _clientConnectionManager.WhenAllCompleted(),
                     Task.Delay(Timeout.InfiniteTimeSpan, source.Token)
                 );
             }
             catch (OperationCanceledException)
             {
-                _logger.LogWarning($"[GracefulShutdown] Timeout ({_options.GracefulShutdown.Timeout.TotalMilliseconds}ms) reached, existing connections will be dropped immediately.");
+                _logger.LogWarning($"[GracefulShutdown][{_hubName}] Timeout({_options.GracefulShutdown.Timeout.TotalMilliseconds}ms) reached, existing client connections will be dropped immediately.");
             }
-        }
-
-        private async Task OfflineAndWaitForCompletedAsync(GracefulShutdownMode mode)
-        {
-            _logger.LogInformation("[GracefulShutdown] Unloading server connections.");
-            await _serviceConnectionManager.OfflineAsync(mode);
 
-            _logger.LogInformation("[GracefulShutdown] Waiting client connections to complete.");
-            await _clientConnectionManager.WhenAllCompleted();
+            _logger.LogInformation($"[GracefulShutdown][{_hubName}] Stopping server...");
+            await _serviceConnectionManager.StopAsync();
+            _logger.LogInformation($"[GracefulShutdown][{_hubName}] Finished.");
         }
 
         private IMultiEndpointServiceConnectionContainer GetMultiEndpointServiceConnectionContainer(string hub, ConnectionDelegate connectionDelegate, Action<HttpContext> contextConfig = null)

diff --git a/src/Microsoft.Azure.SignalR/ServerConnections/ServiceConnectionManager.cs b/src/Microsoft.Azure.SignalR/ServerConnections/ServiceConnectionManager.cs
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR;

diff --git a/test/Microsoft.Azure.SignalR.Tests/ServiceHubDispatcherTests.cs b/test/Microsoft.Azure.SignalR.Tests/ServiceHubDispatcherTests.cs
@@ -30,6 +30,7 @@ public async void TestShutdown()
 
             var dispatcher = new ServiceHubDispatcher<Hub>(
                 null,
+                TestHubContext<Hub>.GetInstance(),
                 serviceManager,
                 clientManager,
                 null,
@@ -48,6 +49,18 @@ public async void TestShutdown()
             Assert.Equal(1, serviceManager.OfflineIndex);
         }
 
+        private sealed class TestHubContext<THub> : IHubContext<THub> where THub : Hub
+        {
+            public IHubClients Clients => throw new NotImplementedException();
+
+            public IGroupManager Groups => throw new NotImplementedException();
+
+            public static TestHubContext<THub> GetInstance()
+            {
+                return new TestHubContext<THub>();
+            }
+        }
+
         private sealed class TestRouter : IEndpointRouter
         {
             public IEnumerable<ServiceEndpoint> GetEndpointsForBroadcast(IEnumerable<ServiceEndpoint> endpoints)