From f51fcadeb18c0013cc25b43547c43d66639a6d48 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Wed, 2 Nov 2016 23:14:12 -0700
Subject: [PATCH] Dispose the ConnectionManager on application shutdown

---
 .../ConnectionManager.cs                      | 25 ++++++++++++++++++-
 .../HttpDispatcherAppBuilderExtensions.cs     |  8 ++++++
 src/Microsoft.AspNetCore.Sockets/project.json | 23 +++++++++--------
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs b/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs
index 72944f272d..1d128c2ab5 100644
--- a/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs
+++ b/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs
@@ -5,7 +5,7 @@ using Channels;
 
 namespace Microsoft.AspNetCore.Sockets
 {
-    public class ConnectionManager
+    public class ConnectionManager : IDisposable
     {
         private ConcurrentDictionary<string, ConnectionState> _connections = new ConcurrentDictionary<string, ConnectionState>();
         private Timer _timer;
@@ -92,5 +92,28 @@ namespace Microsoft.AspNetCore.Sockets
                 }
             }
         }
+
+        public void Dispose()
+        {
+            // Stop firing the timer
+            _timer.Dispose();
+
+            foreach (var c in _connections)
+            {
+                ConnectionState s;
+                if (_connections.TryRemove(c.Key, out s))
+                {
+                    // Longpolling connections should do this
+                    if (s.Close != null)
+                    {
+                        s.Close();
+                    }
+                    else
+                    {
+                        s.Connection.Channel.Dispose();
+                    }
+                }
+            }
+        }
     }
 }
diff --git a/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs b/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs
index b15d74a9b4..b6f24205e8 100644
--- a/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs
+++ b/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs
@@ -1,8 +1,10 @@
 using System;
 using Channels;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.AspNetCore.Sockets;
 using Microsoft.AspNetCore.Sockets.Routing;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Microsoft.AspNetCore.Builder
 {
@@ -12,7 +14,13 @@ namespace Microsoft.AspNetCore.Builder
         {
             var manager = new ConnectionManager();
             var factory = new ChannelFactory();
+
             var dispatcher = new HttpConnectionDispatcher(manager, factory);
+
+            // Dispose the connection manager when application shutdown is triggered
+            var lifetime = app.ApplicationServices.GetRequiredService<IApplicationLifetime>();
+            lifetime.ApplicationStopping.Register(state => ((IDisposable)state).Dispose(), manager);
+
             var routes = new RouteBuilder(app);
 
             callback(new SocketRouteBuilder(routes, dispatcher));
diff --git a/src/Microsoft.AspNetCore.Sockets/project.json b/src/Microsoft.AspNetCore.Sockets/project.json
index 455ae26c5c..72175223a5 100644
--- a/src/Microsoft.AspNetCore.Sockets/project.json
+++ b/src/Microsoft.AspNetCore.Sockets/project.json
@@ -1,13 +1,14 @@
 {
-    "version": "0.1.0-*",
-    "dependencies": {
-        "Channels": "0.2.0-beta-*",
-        "Microsoft.AspNetCore.Routing": "1.1.0-*",
-        "Microsoft.AspNetCore.WebSockets": "0.2.0-*",
-        "NETStandard.Library": "1.6.1-*"
-    },
-    "frameworks": {
-        "netstandard1.3": {},
-        "net46": {}
-    }
+  "version": "0.1.0-*",
+  "dependencies": {
+    "Channels": "0.2.0-beta-*",
+    "Microsoft.AspNetCore.Hosting.Abstractions": "1.1.0-*",
+    "Microsoft.AspNetCore.Routing": "1.1.0-*",
+    "Microsoft.AspNetCore.WebSockets": "0.2.0-*",
+    "NETStandard.Library": "1.6.1-*"
+  },
+  "frameworks": {
+    "netstandard1.3": {},
+    "net46": {}
+  }
 }