Add scheduling support to the Mocha in-memory transport

src/Mocha/src/Mocha.Transport.InMemory/MessageBusBuilderExtensions.cs

@@ -1,4 +1,9 @@
-namespace Mocha.Transport.InMemory;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.Logging;
+using Mocha.Scheduling;
+
+namespace Mocha.Transport.InMemory;
 
 /// <summary>
 /// Extension methods for registering the in-memory messaging transport on an <see cref="IMessageBusHostBuilder"/>.
@@ -19,6 +24,22 @@ public static IMessageBusHostBuilder AddInMemory(
 
         busBuilder.ConfigureMessageBus(b => b.AddTransport(transport));
 
+        busBuilder.Services.TryAddSingleton<InMemoryScheduledMessageStore>();
+        busBuilder.Services.TryAddSingleton<IScheduledMessageStore>(
+            sp => sp.GetRequiredService<InMemoryScheduledMessageStore>());
+        busBuilder.Services.TryAddSingleton(sp => new InMemoryScheduledMessageWorker(
+            sp,
+            sp.GetRequiredService<IMessagingRuntime>(),
+            sp.GetRequiredService<IMessagingPools>(),
+            sp.GetRequiredService<ISchedulerSignal>(),
+            sp.GetRequiredService<InMemoryScheduledMessageStore>(),
+            sp.GetService<TimeProvider>() ?? TimeProvider.System,
+            sp.GetRequiredService<ILogger<InMemoryScheduledMessageWorker>>()));
+        busBuilder.Services.AddHostedService(
+            sp => sp.GetRequiredService<InMemoryScheduledMessageWorker>());
+
+        busBuilder.UseSchedulerCore();
+
         return busBuilder;
     }
 

src/Mocha/src/Mocha.Transport.InMemory/Scheduling/InMemoryScheduledMessageStore.cs

@@ -0,0 +1,115 @@
+using System.Diagnostics.CodeAnalysis;
+using Mocha.Middlewares;
+using Mocha.Scheduling;
+
+namespace Mocha.Transport.InMemory;
+
+/// <summary>
+/// In-memory implementation of <see cref="IScheduledMessageStore"/> that holds scheduled message
+/// envelopes in process memory until they are due for delivery. State is not durable: scheduled
+/// messages are lost when the process stops.
+/// </summary>
+public sealed class InMemoryScheduledMessageStore(ISchedulerSignal signal) : IScheduledMessageStore
+{
+    private const string TokenPrefix = "in-memory-scheduler:";
+#if NET9_0_OR_GREATER
+    private readonly Lock _lock = new();
+#else
+    private readonly object _lock = new();
+#endif
+    private readonly Dictionary<Guid, Entry> _entries = [];
+
+    /// <inheritdoc />
+    public ValueTask<string> PersistAsync(
+        MessageEnvelope envelope,
+        DateTimeOffset scheduledTime,
+        CancellationToken cancellationToken)
+    {
+        var id = Guid.NewGuid();
+
+        var stored = new MessageEnvelope(envelope) { Body = envelope.Body.ToArray() };
+
+        lock (_lock)
+        {
+            _entries[id] = new Entry(stored, scheduledTime);
+        }
+
+        signal.Notify(scheduledTime);
+
+        return new ValueTask<string>($"{TokenPrefix}{id:D}");
+    }
+
+    /// <inheritdoc />
+    public ValueTask<bool> CancelAsync(string token, CancellationToken cancellationToken)
+    {
+        var value = token.StartsWith(TokenPrefix, StringComparison.Ordinal)
+            ? token[TokenPrefix.Length..]
+            : token;
+
+        if (!Guid.TryParse(value, out var id))
+        {
+            return new ValueTask<bool>(false);
+        }
+
+        lock (_lock)
+        {
+            return new ValueTask<bool>(_entries.Remove(id));
+        }
+    }
+
+    /// <summary>
+    /// Removes and returns the earliest entry whose scheduled time is at or before <paramref name="now"/>.
+    /// </summary>
+    public bool TryTakeDue(DateTimeOffset now, [NotNullWhen(true)] out MessageEnvelope? envelope)
+    {
+        lock (_lock)
+        {
+            var foundId = Guid.Empty;
+            var foundTime = DateTimeOffset.MaxValue;
+            MessageEnvelope? foundEnvelope = null;
+
+            foreach (var (id, entry) in _entries)
+            {
+                if (entry.ScheduledTime <= now && entry.ScheduledTime < foundTime)
+                {
+                    foundId = id;
+                    foundTime = entry.ScheduledTime;
+                    foundEnvelope = entry.Envelope;
+                }
+            }
+
+            if (foundEnvelope is null)
+            {
+                envelope = null;
+                return false;
+            }
+
+            _entries.Remove(foundId);
+            envelope = foundEnvelope;
+            return true;
+        }
+    }
+
+    /// <summary>
+    /// Returns the earliest scheduled time across all pending entries, or <c>null</c> when empty.
+    /// </summary>
+    public DateTimeOffset? NextDueTime()
+    {
+        lock (_lock)
+        {
+            DateTimeOffset? next = null;
+
+            foreach (var entry in _entries.Values)
+            {
+                if (next is null || entry.ScheduledTime < next)
+                {
+                    next = entry.ScheduledTime;
+                }
+            }
+
+            return next;
+        }
+    }
+
+    private readonly record struct Entry(MessageEnvelope Envelope, DateTimeOffset ScheduledTime);
+}

src/Mocha/src/Mocha.Transport.InMemory/Scheduling/InMemoryScheduledMessageWorker.cs

@@ -0,0 +1,152 @@
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+using Mocha.Middlewares;
+using Mocha.Scheduling;
+using Mocha.Threading;
+
+namespace Mocha.Transport.InMemory;
+
+/// <summary>
+/// Background service that delivers due in-memory scheduled messages. It sleeps on the scheduler
+/// signal until the next message is due, then dispatches it through the normal endpoint pipeline.
+/// Delivery is best-effort: a message that fails to dispatch is logged and dropped.
+/// </summary>
+internal sealed class InMemoryScheduledMessageWorker(
+    IServiceProvider services,
+    IMessagingRuntime runtime,
+    IMessagingPools pools,
+    ISchedulerSignal signal,
+    InMemoryScheduledMessageStore store,
+    TimeProvider timeProvider,
+    ILogger<InMemoryScheduledMessageWorker> logger) : IHostedService, IAsyncDisposable
+{
+#if NET9_0_OR_GREATER
+    private readonly Lock _lock = new();
+#else
+    private readonly object _lock = new();
+#endif
+    private ContinuousTask? _task;
+
+    public Task StartAsync(CancellationToken cancellationToken)
+    {
+        lock (_lock)
+        {
+            _task ??= new ContinuousTask(ProcessAsync);
+        }
+
+        return Task.CompletedTask;
+    }
+
+    public async Task StopAsync(CancellationToken cancellationToken)
+    {
+        ContinuousTask? task;
+
+        lock (_lock)
+        {
+            task = _task;
+            _task = null;
+        }
+
+        if (task is not null)
+        {
+            await task.DisposeAsync();
+        }
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await StopAsync(CancellationToken.None);
+    }
+
+    private async Task ProcessAsync(CancellationToken cancellationToken)
+    {
+        while (!cancellationToken.IsCancellationRequested)
+        {
+            if (store.TryTakeDue(timeProvider.GetUtcNow(), out var envelope))
+            {
+                await DispatchAsync(envelope, cancellationToken);
+                continue;
+            }
+
+            await signal.WaitUntilAsync(store.NextDueTime() ?? DateTimeOffset.MaxValue, cancellationToken);
+        }
+    }
+
+    private async Task DispatchAsync(MessageEnvelope envelope, CancellationToken cancellationToken)
+    {
+        var messageType = runtime.GetMessageType(envelope.MessageType);
+        var isReply = envelope.Headers?.IsReply() ?? false;
+        var endpoint = isReply
+            ? GetReplyEndpoint(envelope.DestinationAddress)
+            : GetSendEndpoint(envelope.DestinationAddress);
+
+        if (messageType is null || endpoint is null)
+        {
+            logger.CouldNotDispatchScheduledMessage(envelope.MessageId);
+            return;
+        }
+
+        var context = pools.DispatchContext.Get();
+        try
+        {
+            await using var scope = services.CreateAsyncScope();
+            context.Initialize(scope.ServiceProvider, endpoint, runtime, messageType, cancellationToken);
+            context.SkipScheduler();
+
+            // Unlike the Postgres dispatcher, this worker does not call SkipOutbox(): the in-memory
+            // transport is not paired with an outbox, and the re-dispatched envelope carries no
+            // headers, so the outbox middleware passes through. Avoids coupling to Mocha.Outbox.
+            context.Envelope = envelope;
+            await endpoint.ExecuteAsync(context);
+        }
+        catch (Exception ex) when (ex is not OperationCanceledException)
+        {
+            logger.ScheduledMessageDispatchFailed(envelope.MessageId, ex);
+        }
+        finally
+        {
+            pools.DispatchContext.Return(context);
+        }
+    }
+
+    private DispatchEndpoint? GetSendEndpoint(string? address)
+    {
+        try
+        {
+            return Uri.TryCreate(address, UriKind.Absolute, out var uri)
+                ? runtime.GetDispatchEndpoint(uri)
+                : null;
+        }
+        catch
+        {
+            return null;
+        }
+    }
+
+    private DispatchEndpoint? GetReplyEndpoint(string? address)
+    {
+        try
+        {
+            return Uri.TryCreate(address, UriKind.Absolute, out var uri)
+                ? runtime.GetTransport(uri)?.ReplyDispatchEndpoint
+                : null;
+        }
+        catch
+        {
+            return null;
+        }
+    }
+}
+
+internal static partial class InMemoryScheduledMessageWorkerLog
+{
+    [LoggerMessage(1, LogLevel.Warning, "Could not dispatch scheduled message {MessageId}")]
+    public static partial void CouldNotDispatchScheduledMessage(this ILogger logger, string? messageId);
+
+    [LoggerMessage(2, LogLevel.Error, "Scheduled message {MessageId} dispatch failed")]
+    public static partial void ScheduledMessageDispatchFailed(
+        this ILogger logger,
+        string? messageId,
+        Exception exception);
+}

src/Mocha/src/Mocha/Scheduling/SchedulerCoreServiceCollectionExtensions.cs

@@ -16,6 +16,7 @@ public static class SchedulerCoreServiceCollectionExtensions
     /// Adds <see cref="ISchedulerSignal"/> as a singleton and configures the dispatch pipeline
     /// to persist outgoing messages with a <see cref="Middlewares.IDispatchContext.ScheduledTime"/>
     /// through <see cref="IScheduledMessageStore"/> instead of dispatching them directly.
+    /// Calling this method more than once is safe: the dispatch middleware is installed at most once.
     /// </remarks>
     /// <param name="builder">The message bus host builder to configure.</param>
     /// <returns>The same <paramref name="builder"/> instance for chaining.</returns>
@@ -24,8 +25,20 @@ public static IMessageBusHostBuilder UseSchedulerCore(this IMessageBusHostBuilde
         builder.Services.TryAddSingleton<ISchedulerSignal>(sp =>
             new MessageBusSchedulerSignal(sp.GetService<TimeProvider>() ?? TimeProvider.System));
 
+        if (builder.Services.Any(d => d.ServiceType == typeof(SchedulerCoreMarker)))
+        {
+            return builder;
+        }
+
+        builder.Services.AddSingleton<SchedulerCoreMarker>();
         builder.ConfigureMessageBus(x => x.UseDispatch(DispatchSchedulingMiddleware.Create(), after: "Serialization"));
 
         return builder;
     }
 }
+
+/// <summary>
+/// Marker service used to detect whether <see cref="SchedulerCoreServiceCollectionExtensions.UseSchedulerCore"/>
+/// has already been called, preventing the scheduling dispatch middleware from being installed twice.
+/// </summary>
+internal sealed class SchedulerCoreMarker;

src/Mocha/test/Mocha.Sagas.Tests/IntegrationTests.cs

@@ -1,5 +1,6 @@
 using System.Collections.Concurrent;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Mocha.Transport.InMemory;
 
 namespace Mocha.Sagas.Tests;
@@ -19,6 +20,12 @@ private static async Task<ServiceProvider> CreateBusAsync(Action<IMessageBusHost
         var provider = services.BuildServiceProvider();
         var runtime = (MessagingRuntime)provider.GetRequiredService<IMessagingRuntime>();
         await runtime.StartAsync(CancellationToken.None);
+
+        foreach (var hostedService in provider.GetServices<IHostedService>())
+        {
+            await hostedService.StartAsync(CancellationToken.None);
+        }
+
         return provider;
     }
 
@@ -123,13 +130,28 @@ public async Task Saga_Should_TimeoutWithCustomResponse()
         var recorded = Assert.Single(recorder.Messages);
         var sagaId = Assert.IsType<TriggerEvent>(recorded).CorrelationId!.Value;
 
+        // wait until the saga instance is persisted before sending the timeout, so the
+        // SagaTimedOutEvent is applied to the stored instance
+        var runtime = provider.GetRequiredService<IMessagingRuntime>();
+        var sagaName = runtime.Naming.GetSagaName(typeof(TimeoutWithResponseSaga));
+        var persistDeadline = DateTime.UtcNow + s_timeout;
+        while (storage.Load<TimeoutState>(sagaName, sagaId) is null && DateTime.UtcNow < persistDeadline)
+        {
+            await Task.Delay(50, TestContext.Current.CancellationToken);
+        }
+
+        Assert.NotNull(storage.Load<TimeoutState>(sagaName, sagaId));
+
         // simulate timeout by sending SagaTimedOutEvent with the saga ID
         await bus.SendAsync(new SagaTimedOutEvent(sagaId), CancellationToken.None);
 
-        // allow time for final state processing
-        await Task.Delay(500, TestContext.Current.CancellationToken);
+        // wait for the saga to reach its final state and be deleted from the store
+        var deadline = DateTime.UtcNow + s_timeout;
+        while (storage.Load<TimeoutState>(sagaName, sagaId) is not null && DateTime.UtcNow < deadline)
+        {
+            await Task.Delay(50, TestContext.Current.CancellationToken);
+        }
 
-        // assert - saga should be deleted from store after reaching final state
         Assert.Equal(0, storage.Count);
     }