[Mocha] Fix transports not stopping when the host shuts down

src/Mocha/src/Mocha.Transport.Postgres/PostgresMessagingTransport.cs

@@ -419,7 +419,7 @@ private async Task RecoverFromEvictionAsync(CancellationToken cancellationToken)
     }
 
     /// <inheritdoc />
-    protected override async ValueTask OnBeforeStopAsync(CancellationToken cancellationToken)
+    protected override async ValueTask OnAfterStopAsync(CancellationToken cancellationToken)
     {
         await _backgroundTasks.DisposeAsync();
 

src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs

@@ -288,9 +288,8 @@ public async ValueTask StartAsync(IMessagingRuntimeContext context, Cancellation
     /// Stops this endpoint, ceasing message consumption and releasing runtime resources.
     /// </summary>
     /// <remarks>
-    /// Delegates to <see cref="OnStopAsync"/> for transport-specific shutdown and clears
-    /// the runtime state. This method is idempotent; calling it on an already-stopped
-    /// endpoint is a no-op.
+    /// Delegates to <see cref="OnStopAsync"/> for transport-specific shutdown. This method is
+    /// idempotent; calling it on an already-stopped endpoint is a no-op.
     /// </remarks>
     /// <param name="context">The messaging runtime context.</param>
     /// <param name="cancellationToken">Token to signal cancellation of the stop operation.</param>
@@ -303,7 +302,6 @@ public async ValueTask StopAsync(IMessagingRuntimeContext context, CancellationT
 
         await OnStopAsync(context, cancellationToken);
 
-        _runtimeState = null;
         IsStarted = false;
     }
 

src/Mocha/src/Mocha/Runtime/MessagingRuntime.cs

@@ -36,6 +36,9 @@ public sealed class MessagingRuntime(
     IEndpointRouter endpointRouter,
     IFeatureCollection features) : IMessagingRuntime, IAsyncDisposable
 {
+    private bool _stopped;
+    private bool _disposed;
+
     /// <inheritdoc />
     public IServiceProvider Services => services;
 
@@ -122,8 +125,14 @@ public MessageType GetMessageType(Type type)
     /// Starts all registered transports and their receive endpoints, enabling message consumption.
     /// </summary>
     /// <param name="cancellationToken">A token to cancel the startup sequence.</param>
+    /// <exception cref="InvalidOperationException">Thrown if the runtime has been stopped.</exception>
     public async ValueTask StartAsync(CancellationToken cancellationToken)
     {
+        if (_stopped)
+        {
+            throw ThrowHelper.RuntimeStopped();
+        }
+
         if (IsStarted)
         {
             return;
@@ -137,14 +146,53 @@ public async ValueTask StartAsync(CancellationToken cancellationToken)
         IsStarted = true;
     }
 
+    /// <summary>
+    /// Stops the started transports so their receive endpoints stop consuming messages. A stopped
+    /// runtime cannot be started again.
+    /// </summary>
+    /// <param name="cancellationToken">A token to cancel the shutdown operations.</param>
+    public async ValueTask StopAsync(CancellationToken cancellationToken)
+    {
+        _stopped = true;
+
+        foreach (var transport in transports)
+        {
+            if (transport.IsStarted)
+            {
+                await transport.StopAsync(this, cancellationToken);
+            }
+        }
+
+        IsStarted = false;
+    }
+
     /// <inheritdoc />
     public async ValueTask DisposeAsync()
     {
-        foreach (var consumer in consumers)
+        if (_disposed)
         {
-            await consumer.DisposeAsync();
+            return;
         }
 
-        _changeTokenSource?.Dispose();
+        _disposed = true;
+
+        try
+        {
+            await StopAsync(CancellationToken.None);
+        }
+        finally
+        {
+            foreach (var consumer in consumers)
+            {
+                await consumer.DisposeAsync();
+            }
+
+            foreach (var transport in transports)
+            {
+                await transport.DisposeAsync();
+            }
+
+            _changeTokenSource?.Dispose();
+        }
     }
 }

src/Mocha/src/Mocha/Runtime/MessagingRuntimeHostedService.cs

@@ -3,7 +3,7 @@
 namespace Mocha;
 
 /// <summary>
-/// Hosted service that automatically starts the messaging runtime when the host starts.
+/// Hosted service that starts and stops the messaging runtime with the host.
 /// </summary>
 internal sealed class MessagingRuntimeHostedService(IMessagingRuntime runtime) : IHostedService
 {
@@ -16,6 +16,6 @@ public async Task StartAsync(CancellationToken cancellationToken)
 
     public async Task StopAsync(CancellationToken cancellationToken)
     {
-        await _runtime.DisposeAsync();
+        await _runtime.StopAsync(cancellationToken);
     }
 }

src/Mocha/src/Mocha/ThrowHelper.cs

@@ -110,6 +110,9 @@ public static Exception TransportAlreadyStarted()
     public static Exception TransportNotStarted()
         => new InvalidOperationException("Transport is not started");
 
+    public static Exception RuntimeStopped()
+        => new InvalidOperationException("Messaging runtime is stopped and cannot be started again");
+
     public static Exception EndpointConfigurationFailed()
         => new InvalidOperationException("Failed to create endpoint configuration");
 

src/Mocha/src/Mocha/Transport/MessagingTransport.cs

@@ -245,6 +245,8 @@ public async ValueTask StopAsync(IMessagingRuntimeContext context, CancellationT
             await endpoint.StopAsync(context, cancellationToken);
         }
 
+        await OnAfterStopAsync(cancellationToken);
+
         IsStarted = false;
     }
 
@@ -266,6 +268,13 @@ protected virtual ValueTask OnBeforeStartAsync(
     /// <param name="cancellationToken">A token to cancel the pre-stop operation.</param>
     protected virtual ValueTask OnBeforeStopAsync(CancellationToken cancellationToken) => ValueTask.CompletedTask;
 
+    /// <summary>
+    /// Called after receive endpoints are stopped, allowing derived transports to release resources
+    /// that the endpoints use while stopping.
+    /// </summary>
+    /// <param name="cancellationToken">A token to cancel the post-stop operation.</param>
+    protected virtual ValueTask OnAfterStopAsync(CancellationToken cancellationToken) => ValueTask.CompletedTask;
+
     /// <summary>
     /// Creates the transport-specific configuration from the setup context during the bus build phase.
     /// </summary>

src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxIntegrationTests.cs

@@ -196,7 +196,10 @@ public async Task Outbox_Should_ResumeProcessing_When_WorkerRestartedAfterInterr
             await db.Database.EnsureCreatedAsync(TestContext.Current.CancellationToken);
         }
 
-        var hostedServices = provider.GetServices<IHostedService>().ToList();
+        // the runtime is started above and stays running, only the outbox worker is restarted
+        var hostedServices = provider.GetServices<IHostedService>()
+            .Where(s => s is not MessagingRuntimeHostedService)
+            .ToList();
         foreach (var svc in hostedServices)
         {
             await svc.StartAsync(TestContext.Current.CancellationToken);

src/Mocha/test/Mocha.Tests/Runtime/MessagingRuntimeHostedServiceTests.cs

@@ -0,0 +1,93 @@
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Mocha.Transport.InMemory;
+
+namespace Mocha.Tests;
+
+public sealed class MessagingRuntimeHostedServiceTests
+{
+    [Fact]
+    public async Task StopAsync_Should_StopTransportsAndEndpoints_When_HostStops()
+    {
+        // arrange
+        await using var provider = CreateProvider();
+        var hostedService = GetHostedService(provider);
+        var runtime = GetRuntime(provider);
+        await hostedService.StartAsync(TestContext.Current.CancellationToken);
+
+        // act
+        await hostedService.StopAsync(TestContext.Current.CancellationToken);
+
+        // assert
+        Assert.False(runtime.IsStarted);
+        Assert.All(runtime.Transports, transport => Assert.False(transport.IsStarted));
+        Assert.All(
+            runtime.Transports.SelectMany(transport => transport.ReceiveEndpoints),
+            endpoint => Assert.False(endpoint.IsStarted));
+    }
+
+    [Fact]
+    public async Task StartAsync_Should_Throw_When_HostWasStopped()
+    {
+        // arrange
+        await using var provider = CreateProvider();
+        var hostedService = GetHostedService(provider);
+        await hostedService.StartAsync(TestContext.Current.CancellationToken);
+        await hostedService.StopAsync(TestContext.Current.CancellationToken);
+
+        // act
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(
+            () => hostedService.StartAsync(TestContext.Current.CancellationToken));
+
+        // assert
+        Assert.Equal("Messaging runtime is stopped and cannot be started again", exception.Message);
+    }
+
+    [Fact]
+    public async Task DisposeAsync_Should_StopTransports_When_ServiceProviderIsDisposed()
+    {
+        // arrange
+        var provider = CreateProvider();
+        var runtime = GetRuntime(provider);
+        await GetHostedService(provider).StartAsync(TestContext.Current.CancellationToken);
+
+        // act
+        await provider.DisposeAsync();
+
+        // assert
+        Assert.False(runtime.IsStarted);
+        Assert.All(runtime.Transports, transport => Assert.False(transport.IsStarted));
+    }
+
+    private static ServiceProvider CreateProvider()
+        => new ServiceCollection()
+            .AddMessageBus()
+            .AddEventHandler<EventHandler>()
+            .AddBatchHandler<BatchHandler>(o => o.MaxBatchSize = 1)
+            .AddInMemory(t =>
+            {
+                t.Queue("events").Handler<EventHandler>();
+                t.Queue("batches").Handler<BatchHandler>();
+            })
+            .Services.BuildServiceProvider();
+
+    private static MessagingRuntimeHostedService GetHostedService(IServiceProvider provider)
+        => provider.GetServices<IHostedService>().OfType<MessagingRuntimeHostedService>().Single();
+
+    private static MessagingRuntime GetRuntime(IServiceProvider provider)
+        => (MessagingRuntime)provider.GetRequiredService<IMessagingRuntime>();
+
+    public sealed record TestEvent(string Name);
+
+    public sealed class EventHandler : IEventHandler<TestEvent>
+    {
+        public ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken)
+            => ValueTask.CompletedTask;
+    }
+
+    public sealed class BatchHandler : IBatchEventHandler<TestEvent>
+    {
+        public ValueTask HandleAsync(IMessageBatch<TestEvent> messages, CancellationToken cancellationToken)
+            => ValueTask.CompletedTask;
+    }
+}

src/Mocha/test/Mocha.Transport.Postgres.Tests/Behaviors/ShutdownTests.cs

@@ -0,0 +1,72 @@
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Mocha.Transport.Postgres.Tests.Helpers;
+using Npgsql;
+
+namespace Mocha.Transport.Postgres.Tests.Behaviors;
+
+[Collection("Postgres")]
+public sealed class ShutdownTests(PostgresFixture fixture)
+{
+    [Fact]
+    public async Task StopAsync_Should_UnregisterConsumer_When_HostStops()
+    {
+        // arrange
+        await using var db = await fixture.CreateDatabaseAsync("shutdown_unregister");
+        await using var provider = CreateProvider(db.ConnectionString);
+        var services = provider.GetServices<IHostedService>().ToArray();
+        await StartAsync(services);
+        var consumersBeforeStop = await CountConsumersAsync(db.ConnectionString);
+
+        // act
+        await StopAsync(services);
+
+        // assert
+        Assert.Equal(1, consumersBeforeStop);
+        Assert.Equal(0, await CountConsumersAsync(db.ConnectionString));
+    }
+
+    private static ServiceProvider CreateProvider(string connectionString)
+        => new ServiceCollection()
+            .AddMessageBus()
+            .AddEventHandler<NoOpHandler>()
+            .AddPostgres(t =>
+            {
+                t.ConnectionString(connectionString);
+                t.Endpoint("shutdown").Handler<NoOpHandler>();
+            })
+            .Services.BuildServiceProvider();
+
+    private static async Task StartAsync(IHostedService[] services)
+    {
+        foreach (var service in services)
+        {
+            await service.StartAsync(TestContext.Current.CancellationToken);
+        }
+    }
+
+    private static async Task StopAsync(IHostedService[] services)
+    {
+        foreach (var service in services.Reverse())
+        {
+            await service.StopAsync(TestContext.Current.CancellationToken);
+        }
+    }
+
+    private static async Task<long> CountConsumersAsync(string connectionString)
+    {
+        await using var connection = new NpgsqlConnection(connectionString);
+        await connection.OpenAsync(TestContext.Current.CancellationToken);
+        await using var command = connection.CreateCommand();
+        command.CommandText = "SELECT count(*) FROM mocha_consumers";
+        return (long)(await command.ExecuteScalarAsync(TestContext.Current.CancellationToken))!;
+    }
+
+    public sealed class TestEvent;
+
+    public sealed class NoOpHandler : IEventHandler<TestEvent>
+    {
+        public ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken)
+            => ValueTask.CompletedTask;
+    }
+}