Use async scope to make sure scoped IAsyncDisposable instances get disposed (#6650) (#6658)

* AcceptanceTest to verify scoped async disposable support

* use CreateAsyncScope

* Add Assert

* Extend test for singletons

* Owned builder should be disposed using async disposable

* Async dispose token registration

* Dispose stopping token source

* Fix leaking token source

* Enable nullable in RunningEndpointInstance to indicate case when builder is null

* Make the externally managed mode more visible in the code

* Coverage of externally and internally managed mode

* Rename tests

Co-authored-by: Simon <[email protected]>

Co-authored-by: Simon <[email protected]>

Author: Daniel Marbach

src/NServiceBus.AcceptanceTests/Core/DependencyInjection/When_registering_async_disposables_externally_managed.cs

@@ -0,0 +1,114 @@
+namespace NServiceBus.AcceptanceTests.Core.DependencyInjection
+{
+    using System;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using EndpointTemplates;
+    using Microsoft.Extensions.DependencyInjection;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class When_registering_async_disposables_externally_managed : NServiceBusAcceptanceTest
+    {
+        [Test]
+        public async Task Should_dispose()
+        {
+            ServiceProvider serviceProvider = null;
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddSingleton<SingletonAsyncDisposable>();
+            serviceCollection.AddScoped<ScopedAsyncDisposable>();
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<EndpointWithAsyncDisposable>(b =>
+                {
+                    IStartableEndpointWithExternallyManagedContainer configuredEndpoint = null;
+
+                    b.ToCreateInstance(
+                        config =>
+                        {
+                            configuredEndpoint =
+                                EndpointWithExternallyManagedContainer.Create(config, serviceCollection);
+                            return Task.FromResult(configuredEndpoint);
+                        },
+                        configured =>
+                        {
+                            serviceProvider = serviceCollection.BuildServiceProvider();
+                            return configured.Start(serviceProvider);
+                        });
+                    b.When(e => e.SendLocal(new SomeMessage()));
+                })
+                .Done(c => c.ScopedAsyncDisposableDisposed)
+                .Run(TimeSpan.FromSeconds(10));
+
+            await serviceProvider.DisposeAsync();
+
+            Assert.That(context.ScopedAsyncDisposableDisposed, Is.True, "Scoped AsyncDisposable wasn't disposed as it should have been.");
+            Assert.That(context.SingletonAsyncDisposableDisposed, Is.True, "Singleton AsyncDisposable wasn't disposed as it should have been.");
+        }
+
+        class Context : ScenarioContext
+        {
+            public bool ScopedAsyncDisposableDisposed { get; set; }
+            public bool SingletonAsyncDisposableDisposed { get; set; }
+        }
+
+        public class EndpointWithAsyncDisposable : EndpointConfigurationBuilder
+        {
+            public EndpointWithAsyncDisposable() =>
+                EndpointSetup<DefaultServer>();
+
+            class HandlerWithAsyncDisposable : IHandleMessages<SomeMessage>
+            {
+                public HandlerWithAsyncDisposable(Context context, ScopedAsyncDisposable scopedAsyncDisposable, SingletonAsyncDisposable singletonAsyncDisposable)
+                {
+                    testContext = context;
+                    this.scopedAsyncDisposable = scopedAsyncDisposable;
+                    this.singletonAsyncDisposable = singletonAsyncDisposable;
+                }
+
+                public Task Handle(SomeMessage message, IMessageHandlerContext context)
+                {
+                    scopedAsyncDisposable.Initialize(testContext);
+                    singletonAsyncDisposable.Initialize(testContext);
+                    return Task.CompletedTask;
+                }
+
+                readonly Context testContext;
+                readonly ScopedAsyncDisposable scopedAsyncDisposable;
+                readonly SingletonAsyncDisposable singletonAsyncDisposable;
+            }
+        }
+
+        public class SomeMessage : IMessage
+        {
+        }
+
+        class SingletonAsyncDisposable : IAsyncDisposable
+        {
+            // This method is here to make the code being used in the handler to not trigger compiler warnings
+            public void Initialize(Context scenarioContext) => context = scenarioContext;
+
+            public ValueTask DisposeAsync()
+            {
+                context.SingletonAsyncDisposableDisposed = true;
+                return new ValueTask();
+            }
+
+            Context context;
+        }
+
+        class ScopedAsyncDisposable : IAsyncDisposable
+        {
+            // This method is here to make the code being used in the handler to not trigger compiler warnings
+            public void Initialize(Context scenarioContext) => context = scenarioContext;
+
+            public ValueTask DisposeAsync()
+            {
+                context.ScopedAsyncDisposableDisposed = true;
+                return new ValueTask();
+            }
+
+            Context context;
+        }
+    }
+}

src/NServiceBus.AcceptanceTests/Core/DependencyInjection/When_registering_async_disposables_internally_managed.cs

@@ -0,0 +1,100 @@
+namespace NServiceBus.AcceptanceTests.Core.DependencyInjection
+{
+    using System;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using EndpointTemplates;
+    using Microsoft.Extensions.DependencyInjection;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class When_registering_async_disposables_internally_managed : NServiceBusAcceptanceTest
+    {
+        [Test]
+        public async Task Should_dispose()
+        {
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<EndpointWithAsyncDisposable>(b =>
+                {
+                    b.When(e => e.SendLocal(new SomeMessage()));
+                })
+                .Done(c => c.ScopedAsyncDisposableDisposed)
+                .Run(TimeSpan.FromSeconds(10));
+
+            Assert.That(context.ScopedAsyncDisposableDisposed, Is.True, "Scoped AsyncDisposable wasn't disposed as it should have been.");
+            Assert.That(context.SingletonAsyncDisposableDisposed, Is.True, "Singleton AsyncDisposable wasn't disposed as it should have been.");
+        }
+
+        class Context : ScenarioContext
+        {
+            public bool ScopedAsyncDisposableDisposed { get; set; }
+            public bool SingletonAsyncDisposableDisposed { get; set; }
+        }
+
+        public class EndpointWithAsyncDisposable : EndpointConfigurationBuilder
+        {
+            public EndpointWithAsyncDisposable() =>
+                EndpointSetup<DefaultServer>(c =>
+                {
+                    c.RegisterComponents(s =>
+                    {
+                        s.AddScoped<ScopedAsyncDisposable>();
+                        s.AddSingleton<SingletonAsyncDisposable>();
+                    });
+                });
+
+            class HandlerWithAsyncDisposable : IHandleMessages<SomeMessage>
+            {
+                public HandlerWithAsyncDisposable(Context context, ScopedAsyncDisposable scopedAsyncDisposable, SingletonAsyncDisposable singletonAsyncDisposable)
+                {
+                    testContext = context;
+                    this.scopedAsyncDisposable = scopedAsyncDisposable;
+                    this.singletonAsyncDisposable = singletonAsyncDisposable;
+                }
+
+                public Task Handle(SomeMessage message, IMessageHandlerContext context)
+                {
+                    scopedAsyncDisposable.Initialize(testContext);
+                    singletonAsyncDisposable.Initialize(testContext);
+                    return Task.CompletedTask;
+                }
+
+                readonly Context testContext;
+                readonly ScopedAsyncDisposable scopedAsyncDisposable;
+                readonly SingletonAsyncDisposable singletonAsyncDisposable;
+            }
+        }
+
+        public class SomeMessage : IMessage
+        {
+        }
+
+        class SingletonAsyncDisposable : IAsyncDisposable
+        {
+            // This method is here to make the code being used in the handler to not trigger compiler warnings
+            public void Initialize(Context scenarioContext) => context = scenarioContext;
+
+            public ValueTask DisposeAsync()
+            {
+                context.SingletonAsyncDisposableDisposed = true;
+                return new ValueTask();
+            }
+
+            Context context;
+        }
+
+        class ScopedAsyncDisposable : IAsyncDisposable
+        {
+            // This method is here to make the code being used in the handler to not trigger compiler warnings
+            public void Initialize(Context scenarioContext) => context = scenarioContext;
+
+            public ValueTask DisposeAsync()
+            {
+                context.ScopedAsyncDisposableDisposed = true;
+                return new ValueTask();
+            }
+
+            Context context;
+        }
+    }
+}

src/NServiceBus.Core/EndpointInstanceExtensions.cs

@@ -17,7 +17,10 @@ public static class EndpointInstanceExtensions
         /// <param name="endpoint">The endpoint to stop.</param>
         /// <param name="gracefulStopTimeout">The length of time granted to gracefully complete processing.</param>
         [SuppressMessage("Code", "PS0018:A task-returning method should have a CancellationToken parameter unless it has a parameter implementing ICancellableContext", Justification = "Convenience method wrapping the CancellationToken overload.")]
-        public static Task Stop(this IEndpointInstance endpoint, TimeSpan gracefulStopTimeout) =>
-            endpoint.Stop(new CancellationTokenSource(gracefulStopTimeout).Token);
+        public static async Task Stop(this IEndpointInstance endpoint, TimeSpan gracefulStopTimeout)
+        {
+            using var cancellationTokenSource = new CancellationTokenSource(gracefulStopTimeout);
+            await endpoint.Stop(cancellationTokenSource.Token).ConfigureAwait(false);
+        }
     }
 }

src/NServiceBus.Core/Hosting/ExternallyManagedContainerHost.cs

@@ -34,13 +34,13 @@ public ExternallyManagedContainerHost(EndpointCreator endpointCreator, HostingCo
 
         internal Lazy<IServiceProvider> Builder { get; private set; }
 
-        public async Task<IEndpointInstance> Start(IServiceProvider externalBuilder, CancellationToken cancellationToken = default)
+        public async Task<IEndpointInstance> Start(IServiceProvider serviceProvider, CancellationToken cancellationToken = default)
         {
-            objectBuilder = externalBuilder;
+            objectBuilder = serviceProvider;
 
-            var startableEndpoint = endpointCreator.CreateStartableEndpoint(externalBuilder, hostingComponent);
+            var startableEndpoint = endpointCreator.CreateStartableEndpoint(serviceProvider, hostingComponent);
 
-            hostingComponent.RegisterBuilder(externalBuilder);
+            hostingComponent.RegisterServiceProvider(serviceProvider);
 
             await hostingComponent.RunInstallers(cancellationToken).ConfigureAwait(false);
 

src/NServiceBus.Core/Hosting/HostCreator.cs

@@ -59,7 +59,7 @@ public static async Task<IStartableEndpoint> CreateWithInternallyManagedContaine
 
             var serviceProvider = serviceCollection.BuildServiceProvider();
             var startableEndpoint = endpointCreator.CreateStartableEndpoint(serviceProvider, hostingComponent);
-            hostingComponent.RegisterBuilder(serviceProvider);
+            hostingComponent.RegisterServiceProvider(serviceProvider);
 
             await hostingComponent.RunInstallers(cancellationToken).ConfigureAwait(false);
 
@@ -76,4 +76,4 @@ static void CheckIfSettingsWhereUsedToCreateAnotherEndpoint(SettingsHolder setti
             settings.Set("UsedToCreateEndpoint", true);
         }
     }
-}
+}

src/NServiceBus.Core/Hosting/HostingComponent.cs

@@ -12,13 +12,13 @@
 
     partial class HostingComponent
     {
-        public HostingComponent(Configuration configuration, bool shouldDisposeBuilder)
+        public HostingComponent(Configuration configuration, bool shouldDisposeServiceProvider)
         {
             this.configuration = configuration;
-            this.shouldDisposeBuilder = shouldDisposeBuilder;
+            this.shouldDisposeServiceProvider = shouldDisposeServiceProvider;
         }
 
-        public static HostingComponent Initialize(Configuration configuration, IServiceCollection serviceCollection, bool shouldDisposeBuilder)
+        public static HostingComponent Initialize(Configuration configuration, IServiceCollection serviceCollection, bool shouldDisposeServiceProvider)
         {
             serviceCollection.ConfigureComponent(() => configuration.HostInformation, DependencyLifecycle.SingleInstance);
             serviceCollection.ConfigureComponent(() => configuration.CriticalError, DependencyLifecycle.SingleInstance);
@@ -53,13 +53,10 @@ public static HostingComponent Initialize(Configuration configuration, IServiceC
                 PathToExe = PathUtilities.SanitizedPath(Environment.CommandLine)
             });
 
-            return new HostingComponent(configuration, shouldDisposeBuilder);
+            return new HostingComponent(configuration, shouldDisposeServiceProvider);
         }
 
-        public void RegisterBuilder(IServiceProvider objectBuilder)
-        {
-            builder = objectBuilder;
-        }
+        public void RegisterServiceProvider(IServiceProvider serviceProvider) => this.serviceProvider = serviceProvider;
 
         // This can't happen at start due to an old "feature" that allowed users to
         // run installers by "just creating the endpoint". See https://docs.particular.net/nservicebus/operations/installers#running-installers for more details.
@@ -72,7 +69,7 @@ public async Task RunInstallers(CancellationToken cancellationToken = default)
 
             var installationUserName = GetInstallationUserName();
 
-            foreach (var installer in builder.GetServices<INeedToInstallSomething>())
+            foreach (var installer in serviceProvider.GetServices<INeedToInstallSomething>())
             {
                 await installer.Install(installationUserName, cancellationToken).ConfigureAwait(false);
             }
@@ -91,11 +88,11 @@ public async Task<IEndpointInstance> Start(IStartableEndpoint startableEndpoint,
             return endpointInstance;
         }
 
-        public void Stop()
+        public async Task Stop(CancellationToken cancellationToken = default)
         {
-            if (shouldDisposeBuilder)
+            if (shouldDisposeServiceProvider && serviceProvider is IAsyncDisposable asyncDisposableBuilder)
             {
-                (builder as IDisposable)?.Dispose();
+                await asyncDisposableBuilder.DisposeAsync().ConfigureAwait(false);
             }
         }
 
@@ -115,7 +112,7 @@ string GetInstallationUserName()
         }
 
         readonly Configuration configuration;
-        bool shouldDisposeBuilder;
-        IServiceProvider builder;
+        bool shouldDisposeServiceProvider;
+        IServiceProvider serviceProvider;
     }
 }

src/NServiceBus.Core/Pipeline/MainPipelineExecutor.cs

@@ -25,7 +25,8 @@ public async Task Invoke(MessageContext messageContext, CancellationToken cancel
 
             using var activity = activityFactory.StartIncomingActivity(messageContext);
 
-            using (var childScope = rootBuilder.CreateScope())
+            var childScope = rootBuilder.CreateAsyncScope();
+            await using (childScope.ConfigureAwait(false))
             {
                 var message = new IncomingMessage(messageContext.NativeMessageId, messageContext.Headers, messageContext.Body);
 

src/NServiceBus.Core/Recoverability/RecoverabilityPipelineExecutor.cs

@@ -31,7 +31,8 @@ public RecoverabilityPipelineExecutor(
 
         public async Task<ErrorHandleResult> Invoke(ErrorContext errorContext, CancellationToken cancellationToken = default)
         {
-            using (var childScope = serviceProvider.CreateScope())
+            var childScope = serviceProvider.CreateAsyncScope();
+            await using (childScope.ConfigureAwait(false))
             {
                 var rootContext = new RootContext(childScope.ServiceProvider, messageOperations, pipelineCache, cancellationToken);
                 rootContext.Extensions.Merge(errorContext.Extensions);

src/NServiceBus.Core/Unicast/RunningEndpointInstance.cs

@@ -63,15 +63,20 @@ public async Task Stop(CancellationToken cancellationToken = default)
                 finally
                 {
                     settings.Clear();
-                    hostingComponent.Stop();
+                    await hostingComponent.Stop(cancellationToken).ConfigureAwait(false);
                     status = Status.Stopped;
                     Log.Info("Shutdown complete.");
                 }
             }
             finally
             {
                 stopSemaphore.Release();
+#if NET
+                await tokenRegistration.DisposeAsync().ConfigureAwait(false);
+#else
                 tokenRegistration.Dispose();
+#endif
+                stoppingTokenSource.Dispose();
             }
         }