diff --git a/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs b/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs index c817000ee7..1f3fe12169 100644 --- a/src/Microsoft.AspNetCore.Sockets/ConnectionManager.cs +++ b/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 logger) + public ConnectionManager(ILogger logger, IApplicationLifetime appLifetime) { _logger = logger; + + appLifetime.ApplicationStarted.Register(() => Start()); + appLifetime.ApplicationStopping.Register(() => CloseConnections()); } public void Start() diff --git a/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj b/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj index 786e01c81c..6bd6ad9bca 100644 --- a/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj +++ b/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj @@ -13,7 +13,7 @@ - + diff --git a/src/Microsoft.AspNetCore.Sockets/SocketsApplicationLifetimeService.cs b/src/Microsoft.AspNetCore.Sockets/SocketsApplicationLifetimeService.cs deleted file mode 100644 index fbeae6694c..0000000000 --- a/src/Microsoft.AspNetCore.Sockets/SocketsApplicationLifetimeService.cs +++ /dev/null @@ -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; - } - } -} diff --git a/src/Microsoft.AspNetCore.Sockets/SocketsDependencyInjectionExtensions.cs b/src/Microsoft.AspNetCore.Sockets/SocketsDependencyInjectionExtensions.cs index 85c1c2efd0..104f630042 100644 --- a/src/Microsoft.AspNetCore.Sockets/SocketsDependencyInjectionExtensions.cs +++ b/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(); - services.TryAddEnumerable(ServiceDescriptor.Singleton()); return services; } } diff --git a/test/Microsoft.AspNetCore.Sockets.Tests/ApplicationLifetime.cs b/test/Microsoft.AspNetCore.Sockets.Tests/ApplicationLifetime.cs new file mode 100644 index 0000000000..b112ea6aa4 --- /dev/null +++ b/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() + { + } + } +} diff --git a/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs b/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs index 6f054633d7..9f0c56a1a1 100644 --- a/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs +++ b/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(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(new LoggerFactory()), lifetime); } } } diff --git a/test/Microsoft.AspNetCore.Sockets.Tests/HttpConnectionDispatcherTests.cs b/test/Microsoft.AspNetCore.Sockets.Tests/HttpConnectionDispatcherTests.cs index 5b7c44bfea..b6bb73d717 100644 --- a/test/Microsoft.AspNetCore.Sockets.Tests/HttpConnectionDispatcherTests.cs +++ b/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(new LoggerFactory())); + return new ConnectionManager(new Logger(new LoggerFactory()), new EmptyApplicationLifetime()); } private string GetContentAsString(Stream body)