Ensure RunningEndpointInstance.DisposeAsync completes when host is disposed without StopAsync (#7806) (#7818)

* DisposeAsync should pass a bounded cancellation token to transport Shutdown

* .

* Update RunningEndpointInstance.cs

* .

* ,

* Update StartableEndpoint.cs

* Update RunningEndpointInstance.cs

* Remove shutdown token timeout since it would create akward split between dispose and shutdown. Instead cancel the token immediately on dispose.

* Core acceptance test

* Ignore exceptions during DisposeAsync to prevent throwing from disposal process

---------

Co-authored-by: Simon Cropp <simon.cropp@gmail.com>
Co-authored-by: Daniel Marbach <danielmarbach@users.noreply.github.com>

Author: danielmarbach

src/NServiceBus.AcceptanceTesting/Support/EndpointBehavior.cs

@@ -44,6 +44,14 @@ public EndpointBehavior(IEndpointConfigurationFactory endpointBuilder, int insta
                 return async token => await endpointLifecycle.Stop(token).ConfigureAwait(false);
             });
 
+            // For some super advanced scenarios require
+            // disposing the endpoint, and this backdoor allows that without having to expose the lifecycle in Core
+            collectionAdapter.AddKeyedSingleton<Func<ValueTask>>("Disposer", (provider, key) =>
+            {
+                var endpointLifecycle = provider.GetRequiredKeyedService<IEndpointLifecycle>(serviceKey);
+                return async () => await endpointLifecycle.DisposeAsync().ConfigureAwait(false);
+            });
+
             return Task.FromResult(new StartableEndpointInstance(serviceKey));
         }, static (startableEndpoint, provider, cancellationToken) => startableEndpoint.Start(provider, cancellationToken));
     }

src/NServiceBus.AcceptanceTests/Core/Stopping/When_disposing_without_stopping.cs

@@ -0,0 +1,58 @@
+namespace NServiceBus.AcceptanceTests.Core.Stopping;
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using AcceptanceTesting;
+using EndpointTemplates;
+using Microsoft.Extensions.DependencyInjection;
+using NUnit.Framework;
+
+public class When_disposing_without_stopping : NServiceBusAcceptanceTest
+{
+    [Test]
+    [CancelAfter(30000)]
+    public async Task Should_initiate_immediate_handler_cancellation(CancellationToken cancellationToken = default) =>
+        await Scenario.Define<Context>()
+            .WithEndpoint<EndpointThatGetsRugPulled>(b =>
+                b.ServiceResolve(static async (provider, context, token) =>
+                {
+                    var session = provider.GetRequiredService<IMessageSession>();
+                    await session.SendLocal(new MessageThatTakesALongTime(), token);
+
+                    await context.MessageReceived.Task.WaitAsync(token);
+
+                    var disposer = provider.GetRequiredKeyedService<Func<ValueTask>>("Disposer");
+                    await disposer();
+                }, true))
+            .Run(cancellationToken);
+
+    public class Context : ScenarioContext
+    {
+        public TaskCompletionSource MessageReceived { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+
+    public class EndpointThatGetsRugPulled : EndpointConfigurationBuilder
+    {
+        public EndpointThatGetsRugPulled() => EndpointSetup<DefaultServer>();
+
+        [Handler]
+        public class InfiniteHandler(Context testContext) : IHandleMessages<MessageThatTakesALongTime>
+        {
+            public async Task Handle(MessageThatTakesALongTime message, IMessageHandlerContext context)
+            {
+                try
+                {
+                    testContext.MessageReceived.SetResult();
+                    await Task.Delay(Timeout.InfiniteTimeSpan, context.CancellationToken);
+                }
+                catch (OperationCanceledException) when (context.CancellationToken.IsCancellationRequested)
+                {
+                    testContext.MarkAsCompleted();
+                }
+            }
+        }
+    }
+
+    public class MessageThatTakesALongTime : IMessage;
+}

src/NServiceBus.Core.Tests/Unicast/RunningEndpointInstanceTest.cs

@@ -7,6 +7,7 @@
 using NUnit.Framework;
 using Settings;
 using Testing;
+using NServiceBus.Transport;
 
 [TestFixture]
 public class RunningEndpointInstanceTest
@@ -17,10 +18,10 @@ static RunningEndpointInstance Create()
 
         var testInstance = new RunningEndpointInstance(
             settings,
-            null,
+            null!,
             new FeatureComponent(new FeatureComponent.Settings()),
             new TestableMessageSession(),
-            null,
+            null!,
             new CancellationTokenSource(),
             NoOpAsyncDisposable.Instance,
             new EndpointLogSlot("RunningEndpointInstanceTest", endpointIdentifier: null));
@@ -47,6 +48,59 @@ public async Task ShouldAllowMultipleDispose()
         Assert.That(async () => await testInstance.DisposeAsync(), Throws.Nothing);
     }
 
+    [Test]
+    public async Task DisposeAsync_should_pass_a_bounded_cancellation_token_to_transport_Shutdown()
+    {
+        var capturingTransport = new TokenCapturingTransportInfrastructure();
+
+        var testInstance = new RunningEndpointInstance(
+            new SettingsHolder(),
+            CreateEmptyReceiveComponent(),
+            new FeatureComponent(new FeatureComponent.Settings()),
+            new TestableMessageSession(),
+            capturingTransport,
+            new CancellationTokenSource(),
+            NoOpAsyncDisposable.Instance,
+            new EndpointLogSlot("RunningEndpointInstanceTest", endpointIdentifier: null));
+
+        await testInstance.DisposeAsync();
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(capturingTransport.ShutdownWasCalled, Is.True,
+                    "transport.Shutdown was not called during DisposeAsync.");
+            Assert.That(capturingTransport.ObservedToken.CanBeCanceled, Is.True,
+                "DisposeAsync passed CancellationToken.None to StopCore (which then passes it to " +
+                "transport.Shutdown). Disposal must be bounded by an internal cancellation token so " +
+                "a stuck transport cannot hang shutdown indefinitely. " +
+                "See RunningEndpointInstance.DisposeAsync — `await StopCore()` is missing a token.");
+        }
+    }
+
+    [Test]
+    public async Task DisposeAsync_should_complete_even_when_transport_Shutdown_hangs()
+    {
+        var hangingTransport = new HangingTransportInfrastructure();
+
+        var testInstance = new RunningEndpointInstance(
+            new SettingsHolder(),
+            CreateEmptyReceiveComponent(),
+            new FeatureComponent(new FeatureComponent.Settings()),
+            new TestableMessageSession(),
+            hangingTransport,
+            new CancellationTokenSource(),
+            NoOpAsyncDisposable.Instance,
+            new EndpointLogSlot("RunningEndpointInstanceTest", endpointIdentifier: null));
+
+        var dispose = testInstance.DisposeAsync().AsTask();
+        var winner = await Task.WhenAny(dispose, Task.Delay(TimeSpan.FromSeconds(5)));
+
+        Assert.That(winner, Is.SameAs(dispose),
+            "DisposeAsync did not complete within 5s of a 250ms internal timeout — " +
+            "the disposeShutdownTimeout must fire and let disposal proceed past a stuck transport.Shutdown.");
+        Assert.DoesNotThrowAsync(() => dispose, "DisposeAsync threw an exception when transport.Shutdown hung. DisposeAsync should complete successfully even if transport.Shutdown does not.");
+    }
+
     [Test]
     public async Task ShouldThrowExceptionAfterInvokingStop()
     {
@@ -61,4 +115,32 @@ public async Task ShouldThrowExceptionAfterInvokingStop()
         Assert.Throws<InvalidOperationException>(() => testInstance.Subscribe(typeof(object), new SubscribeOptions()), "Invoking messaging operations on the endpoint instance after it has been triggered to stop is not supported.");
         Assert.Throws<InvalidOperationException>(() => testInstance.Unsubscribe(typeof(object), new UnsubscribeOptions()), "Invoking messaging operations on the endpoint instance after it has been triggered to stop is not supported.");
     }
+
+    // Empty ReceiveComponent (no receivers) so Stop is a no-op and StopCore reaches
+    // transport.Shutdown without us having to stand up the full receive pipeline.
+    static ReceiveComponent CreateEmptyReceiveComponent() =>
+        new(configuration: null!, activityFactory: null!, endpointLogSlot: null!);
+
+    sealed class TokenCapturingTransportInfrastructure : TransportInfrastructure
+    {
+        public bool ShutdownWasCalled { get; private set; }
+        public CancellationToken ObservedToken { get; private set; }
+
+        public override Task Shutdown(CancellationToken cancellationToken = default)
+        {
+            ShutdownWasCalled = true;
+            ObservedToken = cancellationToken;
+            return Task.CompletedTask;
+        }
+
+        public override string ToTransportAddress(QueueAddress address) => address.BaseAddress;
+    }
+
+    sealed class HangingTransportInfrastructure : TransportInfrastructure
+    {
+        public override Task Shutdown(CancellationToken cancellationToken = default) =>
+            Task.Delay(Timeout.Infinite, cancellationToken);
+
+        public override string ToTransportAddress(QueueAddress address) => address.BaseAddress;
+    }
 }

src/NServiceBus.Core/Receiving/ReceiveComponent.cs

@@ -17,7 +17,7 @@ namespace NServiceBus;
 
 partial class ReceiveComponent
 {
-    ReceiveComponent(Configuration configuration, IActivityFactory activityFactory, EndpointLogSlot endpointLogSlot)
+    internal ReceiveComponent(Configuration configuration, IActivityFactory activityFactory, EndpointLogSlot endpointLogSlot)
     {
         this.configuration = configuration;
         this.activityFactory = activityFactory;

src/NServiceBus.Core/Unicast/RunningEndpointInstance.cs

@@ -99,11 +99,28 @@ public async ValueTask DisposeAsync()
             return;
         }
 
-        await StopCore().ConfigureAwait(false);
+        // In case Stop was not called, we need to trigger shutdown before cleaning up resources.
+        // Since we're already disposing, we want to bypass any waits and just trigger shutdown with a canceled token.
+        // We are effectively indicating the graceful shutdown period has already elapsed and any ongoing operations should
+        // be aborted immediately if they participate in the cooperative cancellation.
+        var cancellationToken = new CancellationToken(true);
 
-        settings.Clear();
-        stoppingTokenSource.Dispose();
-        await serviceProviderLease.DisposeAsync().ConfigureAwait(false);
+        try
+        {
+            await StopCore(cancellationToken).ConfigureAwait(false);
+        }
+#pragma warning disable PS0019
+        catch (Exception)
+#pragma warning restore PS0019
+        {
+            // ignored because we're already disposing and we don't want to throw from DisposeAsync. Any exceptions from StopCore are already logged, so we can safely ignore them here.
+        }
+        finally
+        {
+            settings.Clear();
+            stoppingTokenSource.Dispose();
+            await serviceProviderLease.DisposeAsync().ConfigureAwait(false);
+        }
     }
 
     public Task Send(object message, SendOptions sendOptions, CancellationToken cancellationToken = default)