Use IApplicationLifetime instead of IHostedService (#757)

src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs

@@ -8,6 +8,7 @@ using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Channels;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Sockets.Internal;
 using Microsoft.Extensions.Logging;
 
@@ -21,9 +22,12 @@ namespace Microsoft.AspNetCore.Sockets
         private object _executionLock = new object();
         private bool _disposed;
 
-        public ConnectionManager(ILogger<ConnectionManager> logger)
+        public ConnectionManager(ILogger<ConnectionManager> logger, IApplicationLifetime appLifetime)
         {
             _logger = logger;
+
+            appLifetime.ApplicationStarted.Register(() => Start());
+            appLifetime.ApplicationStopping.Register(() => CloseConnections());
         }
 
         public void Start()

src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj

@@ -13,7 +13,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\Microsoft.AspNetCore.Sockets.Abstractions\Microsoft.AspNetCore.Sockets.Abstractions.csproj" />
-    <PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" Version="$(AspNetCoreVersion)" />
+    <PackageReference Include="Microsoft.AspNetCore.Hosting.Abstractions" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="System.Threading.Tasks.Channels" Version="$(CoreFxLabsVersion)" />

src/Microsoft.AspNetCore.Sockets/SocketsApplicationLifetimeService.cs

@@ -1,31 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
-
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Extensions.Hosting;
-
-namespace Microsoft.AspNetCore.Sockets
-{
-    public class SocketsApplicationLifetimeService : IHostedService
-    {
-        private readonly ConnectionManager _connectionManager;
-
-        public SocketsApplicationLifetimeService(ConnectionManager connectionManager)
-        {
-            _connectionManager = connectionManager;
-        }
-
-        public Task StartAsync(CancellationToken token)
-        {
-            _connectionManager.Start();
-            return Task.CompletedTask;
-        }
-
-        public Task StopAsync(CancellationToken token)
-        {
-            _connectionManager.CloseConnections();
-            return Task.CompletedTask;
-        }
-    }
-}

src/Microsoft.AspNetCore.Sockets/SocketsDependencyInjectionExtensions.cs

@@ -3,7 +3,6 @@
 
 using Microsoft.AspNetCore.Sockets;
 using Microsoft.Extensions.DependencyInjection.Extensions;
-using Microsoft.Extensions.Hosting;
 
 namespace Microsoft.Extensions.DependencyInjection
 {
@@ -12,7 +11,6 @@ namespace Microsoft.Extensions.DependencyInjection
         public static IServiceCollection AddSocketsCore(this IServiceCollection services)
         {
             services.TryAddSingleton<ConnectionManager>();
-            services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, SocketsApplicationLifetimeService>());
             return services;
         }
     }

test/Microsoft.AspNetCore.Sockets.Tests/ApplicationLifetime.cs

@@ -0,0 +1,44 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Threading;
+using Microsoft.AspNetCore.Hosting;
+
+namespace Microsoft.AspNetCore.Sockets.Tests
+{
+    public class TestApplicationLifetime : IApplicationLifetime
+    {
+        private readonly CancellationTokenSource _startedSource = new CancellationTokenSource();
+        private readonly CancellationTokenSource _stoppingSource = new CancellationTokenSource();
+        private readonly CancellationTokenSource _stoppedSource = new CancellationTokenSource();
+
+        public CancellationToken ApplicationStarted => _startedSource.Token;
+
+        public CancellationToken ApplicationStopping => _stoppingSource.Token;
+
+        public CancellationToken ApplicationStopped => _stoppedSource.Token;
+
+        public void StopApplication()
+        {
+            _stoppingSource.Cancel(throwOnFirstException: false);
+        }
+
+        public void Start()
+        {
+            _startedSource.Cancel(throwOnFirstException: false);
+        }
+    }
+
+    public class EmptyApplicationLifetime : IApplicationLifetime
+	{
+        public CancellationToken ApplicationStarted => CancellationToken.None;
+
+        public CancellationToken ApplicationStopping => CancellationToken.None;
+
+        public CancellationToken ApplicationStopped => CancellationToken.None;
+
+        public void StopApplication()
+        {
+        }
+    }
+}

test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.SignalR.Tests.Common;
 using Microsoft.Extensions.Logging;
 using Xunit;
@@ -183,9 +184,26 @@ namespace Microsoft.AspNetCore.Sockets.Tests
             connectionManager.Scan();
         }
 
-        private static ConnectionManager CreateConnectionManager()
+        [Fact]
+        public async Task ApplicationLifetimeIsHookedUp()
         {
-            return new ConnectionManager(new Logger<ConnectionManager>(new LoggerFactory()));
+            var appLifetime = new TestApplicationLifetime();
+            var connectionManager = CreateConnectionManager(appLifetime);
+
+            appLifetime.Start();
+
+            var connection = connectionManager.CreateConnection();
+
+            appLifetime.StopApplication();
+
+            // Connection should be disposed so this should complete immediately
+            Assert.False(await connection.Application.Out.WaitToWriteAsync().OrTimeout());
+        }
+
+        private static ConnectionManager CreateConnectionManager(IApplicationLifetime lifetime = null)
+        {
+            lifetime = lifetime ?? new EmptyApplicationLifetime();
+            return new ConnectionManager(new Logger<ConnectionManager>(new LoggerFactory()), lifetime);
         }
     }
 }

test/Microsoft.AspNetCore.Sockets.Tests/HttpConnectionDispatcherTests.cs

@@ -1049,7 +1049,7 @@ namespace Microsoft.AspNetCore.Sockets.Tests
 
         private static ConnectionManager CreateConnectionManager()
         {
-            return new ConnectionManager(new Logger<ConnectionManager>(new LoggerFactory()));
+            return new ConnectionManager(new Logger<ConnectionManager>(new LoggerFactory()), new EmptyApplicationLifetime());
         }
 
         private string GetContentAsString(Stream body)