Implement the retry consumer on EF Core

src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs

@@ -43,7 +43,7 @@ public Task StoreFailedErrorImport(FailedErrorImport failure) =>
             }
 
             var failedAt = timeProvider.GetUtcNow().UtcDateTime;
-            var headersJson = JsonSerializer.Serialize(failure.Message.Headers, HeadersJsonContext.Default.DictionaryStringString);
+            var headersJson = MessageHeaders.Write(failure.Message.Headers);
             byte[] storedBody = storeExternally ? [] : body;
 
             await dbContext.UpsertAsync([uniqueMessageId], () => new FailedErrorImportEntity
@@ -143,7 +143,7 @@ async Task<List<FailedErrorImportEntity>> ReadBatch(int offset, CancellationToke
 
     async Task<FailedTransportMessage> ToTransportMessage(FailedErrorImportEntity import, CancellationToken cancellationToken)
     {
-        var headers = JsonSerializer.Deserialize(import.HeadersJson, HeadersJsonContext.Default.DictionaryStringString) ?? [];
+        var headers = MessageHeaders.Read(import.HeadersJson);
 
         var body = import.BodyStoredExternally
             ? await ReadExternalBody(import.UniqueMessageId, cancellationToken)

src/ServiceControl.Persistence.EFCore/Implementation/FailedMessageViewMapper.cs

@@ -6,6 +6,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 using ServiceControl.MessageFailures.Api;
 using ServiceControl.Operations;
 using ServiceControl.Persistence.EFCore.Entities;
+using ServiceControl.Persistence.EFCore.Infrastructure;
 using ServiceControl.Persistence.EFCore.Implementation.UnitOfWork;
 
 static class FailedMessageViewMapper
@@ -147,6 +148,5 @@ static ExceptionDetails ToExceptionDetails(this FailedMessageEntity entity, Dict
                 HostId = entity.ReceivingEndpointHostId ?? Guid.Empty
             };
 
-    static Dictionary<string, string> ReadHeaders(this FailedMessageEntity entity) =>
-        JsonSerializer.Deserialize(entity.HeadersJson, HeadersJsonContext.Default.DictionaryStringString) ?? [];
+    static Dictionary<string, string> ReadHeaders(this FailedMessageEntity entity) => MessageHeaders.Read(entity.HeadersJson);
 }

src/ServiceControl.Persistence.EFCore/Implementation/RetryStagingStore.cs

@@ -1,34 +1,239 @@
 namespace ServiceControl.Persistence.EFCore.Implementation;
 
-public class RetryStagingStore : IRetryStagingStore
+using Microsoft.EntityFrameworkCore;
+using Microsoft.Extensions.DependencyInjection;
+using ServiceControl.MessageFailures;
+using ServiceControl.Persistence.EFCore.DbContexts;
+using ServiceControl.Persistence.EFCore.Entities;
+using ServiceControl.Persistence.EFCore.Infrastructure;
+
+public class RetryStagingStore(IServiceScopeFactory scopeFactory, TimeProvider timeProvider) : DataStoreBase(scopeFactory), IRetryStagingStore
 {
     public Task<RetryBatch?> GetStagingBatch() =>
-        throw new NotImplementedException();
+        ExecuteWithDbContext(async dbContext =>
+        {
+            var batch = await dbContext.RetryBatches
+                .AsNoTracking()
+                .Where(batch => batch.Status == RetryBatchStatus.Staging)
+                .OrderBy(batch => batch.StartTime)
+                .ThenBy(batch => batch.Id)
+                .FirstOrDefaultAsync();
+
+            return batch?.ToRetryBatch(await CountMessages(dbContext, batch.Id));
+        });
+
+    public Task<StagingMessage[]> GetMessagesToStage(string batchId)
+    {
+        var batch = ParseBatchId(batchId);
+
+        return ExecuteWithDbContext(async dbContext =>
+        {
+            // A message claimed by an earlier batch is not claimed by this one, and a claim whose
+            // message is gone drops out of the join, which is what leaves it out of the staging.
+            var rows = await dbContext.FailedMessageRetries
+                .AsNoTracking()
+                .Where(retry => retry.RetryBatchId == batch)
+                .Join(dbContext.FailedMessages.AsNoTracking(),
+                    retry => retry.UniqueMessageId,
+                    message => message.UniqueMessageId,
+                    (retry, message) => new
+                    {
+                        message.UniqueMessageId,
+                        message.MessageId,
+                        message.FailingEndpointAddress,
+                        message.HeadersJson,
+                        retry.StageAttempts
+                    })
+                .ToListAsync();
+
+            return rows.Select(row => new StagingMessage(
+                row.UniqueMessageId.ToString(),
+                row.UniqueMessageId.ToString(),
+                row.MessageId!,
+                row.FailingEndpointAddress,
+                MessageHeaders.Read(row.HeadersJson),
+                row.StageAttempts)).ToArray();
+        });
+    }
+
+    public Task MarkBatchAsForwarding(string batchId, string stagingId, IReadOnlyCollection<string> stagedMessageIds)
+    {
+        var batch = ParseBatchId(batchId);
+        var staged = ParseMessageIds(stagedMessageIds);
+        var now = timeProvider.GetUtcNow().UtcDateTime;
+
+        return ExecuteWithDbContext(dbContext =>
+            InTransaction(dbContext, async () =>
+            {
+                await dbContext.RetryBatches
+                    .Where(row => row.Id == batch)
+                    .ExecuteUpdateAsync(setters => setters
+                        .SetProperty(row => row.Status, RetryBatchStatus.Forwarding)
+                        .SetProperty(row => row.StagingId, stagingId));
 
-    public Task<StagingMessage[]> GetMessagesToStage(string batchId) =>
-        throw new NotImplementedException();
+                // The batch keeps only what it staged, so its message count is what the forwarder is
+                // told to expect. The claims dropped here are of messages that no longer exist.
+                await dbContext.FailedMessageRetries
+                    .Where(row => row.RetryBatchId == batch && !staged.Contains(row.UniqueMessageId))
+                    .ExecuteDeleteAsync();
 
-    public Task MarkBatchAsForwarding(string batchId, string stagingId, IReadOnlyCollection<string> stagedMessageIds) =>
-        throw new NotImplementedException();
+                await dbContext.FailedMessages
+                    .Where(row => staged.Contains(row.UniqueMessageId))
+                    .ExecuteUpdateAsync(setters => setters
+                        .SetProperty(row => row.Status, FailedMessageStatus.RetryIssued)
+                        .SetProperty(row => row.StatusChangedAt, now)
+                        .SetProperty(row => row.LastModified, now));
 
-    public Task DiscardBatch(string batchId) =>
-        throw new NotImplementedException();
+                await PointForwarderAt(dbContext, batch);
+            }));
+    }
+
+    public Task DiscardBatch(string batchId)
+    {
+        var batch = ParseBatchId(batchId);
+
+        return ExecuteWithDbContext(dbContext =>
+            InTransaction(dbContext, async () =>
+            {
+                // Nothing was staged, so every claim of this batch is of a message that is gone.
+                await dbContext.FailedMessageRetries
+                    .Where(row => row.RetryBatchId == batch)
+                    .ExecuteDeleteAsync();
+
+                await dbContext.RetryBatches
+                    .Where(row => row.Id == batch)
+                    .ExecuteDeleteAsync();
+            }));
+    }
 
     public Task<string?> GetForwardingBatchId() =>
-        throw new NotImplementedException();
+        ExecuteWithDbContext(async dbContext =>
+        {
+            var nowForwarding = await dbContext.RetryBatchNowForwarding
+                .AsNoTracking()
+                .SingleOrDefaultAsync();
+
+            return nowForwarding?.RetryBatchId.ToString();
+        });
+
+    public Task<RetryBatch?> GetBatch(string batchId, CancellationToken cancellationToken)
+    {
+        var batch = ParseBatchId(batchId);
+
+        return ExecuteWithDbContext(async dbContext =>
+        {
+            var entity = await dbContext.RetryBatches
+                .AsNoTracking()
+                .SingleOrDefaultAsync(row => row.Id == batch, cancellationToken);
+
+            return entity?.ToRetryBatch(await CountMessages(dbContext, batch, cancellationToken));
+        });
+    }
+
+    public Task CompleteForwarding(string batchId)
+    {
+        var batch = ParseBatchId(batchId);
+
+        return ExecuteWithDbContext(dbContext =>
+            InTransaction(dbContext, async () =>
+            {
+                // The claims outlive the batch: they are what stops a message being staged again
+                // before its retry is confirmed.
+                await dbContext.RetryBatches
+                    .Where(row => row.Id == batch)
+                    .ExecuteDeleteAsync();
+
+                await dbContext.RetryBatchNowForwarding
+                    .Where(row => row.RetryBatchId == batch)
+                    .ExecuteDeleteAsync();
+            }));
+    }
+
+    public Task RecordStagingFailure(IReadOnlyCollection<string> uniqueMessageIds)
+    {
+        var failed = ParseMessageIds(uniqueMessageIds);
+
+        return ExecuteWithDbContext(dbContext => dbContext.FailedMessageRetries
+            .Where(row => failed.Contains(row.UniqueMessageId))
+            .ExecuteUpdateAsync(setters => setters.SetProperty(row => row.StageAttempts, 1)));
+    }
+
+    public Task IncrementStagingAttempts(string uniqueMessageId)
+    {
+        if (!Guid.TryParse(uniqueMessageId, out var message))
+        {
+            return Task.CompletedTask;
+        }
+
+        return ExecuteWithDbContext(dbContext => dbContext.FailedMessageRetries
+            .Where(row => row.UniqueMessageId == message)
+            .ExecuteUpdateAsync(setters => setters.SetProperty(row => row.StageAttempts, row => row.StageAttempts + 1)));
+    }
+
+    public Task RemoveFromBatch(string uniqueMessageId)
+    {
+        if (!Guid.TryParse(uniqueMessageId, out var message))
+        {
+            return Task.CompletedTask;
+        }
+
+        return ExecuteWithDbContext(dbContext => dbContext.FailedMessageRetries
+            .Where(row => row.UniqueMessageId == message)
+            .ExecuteDeleteAsync());
+    }
+
+    static async Task PointForwarderAt(ServiceControlDbContext dbContext, Guid batch)
+    {
+        var updated = await dbContext.RetryBatchNowForwarding
+            .Where(row => row.Id == RetryBatchNowForwardingEntity.SingleRowId)
+            .ExecuteUpdateAsync(setters => setters.SetProperty(row => row.RetryBatchId, batch));
+
+        if (updated == 0)
+        {
+            dbContext.RetryBatchNowForwarding.Add(new RetryBatchNowForwardingEntity { RetryBatchId = batch });
+            await dbContext.SaveChangesAsync();
+        }
+    }
+
+    static Task<int> CountMessages(ServiceControlDbContext dbContext, Guid batch, CancellationToken cancellationToken = default) =>
+        dbContext.FailedMessageRetries
+            .AsNoTracking()
+            .CountAsync(retry => retry.RetryBatchId == batch, cancellationToken);
+
+    static Task InTransaction(ServiceControlDbContext dbContext, Func<Task> operations) =>
+        dbContext.Database.CreateExecutionStrategy().ExecuteAsync(async () =>
+        {
+            await using var transaction = await dbContext.Database.BeginTransactionAsync();
+
+            await operations();
 
-    public Task<RetryBatch?> GetBatch(string batchId, CancellationToken cancellationToken) =>
-        throw new NotImplementedException();
+            await transaction.CommitAsync();
+        });
 
-    public Task CompleteForwarding(string batchId) =>
-        throw new NotImplementedException();
+    /// <summary>
+    /// Message ids reach this from the API, so an id that is not a message id cannot match a stored
+    /// message and is left out rather than thrown at.
+    /// </summary>
+    static HashSet<Guid> ParseMessageIds(IReadOnlyCollection<string> uniqueMessageIds)
+    {
+        var parsed = new HashSet<Guid>(uniqueMessageIds.Count);
 
-    public Task RecordStagingFailure(IReadOnlyCollection<string> uniqueMessageIds) =>
-        throw new NotImplementedException();
+        foreach (var uniqueMessageId in uniqueMessageIds)
+        {
+            if (Guid.TryParse(uniqueMessageId, out var message))
+            {
+                parsed.Add(message);
+            }
+        }
 
-    public Task IncrementStagingAttempts(string uniqueMessageId) =>
-        throw new NotImplementedException();
+        return parsed;
+    }
 
-    public Task RemoveFromBatch(string uniqueMessageId) =>
-        throw new NotImplementedException();
+    /// <summary>
+    /// Batch ids only ever come from CreateBatch, so anything else is a programming error.
+    /// </summary>
+    static Guid ParseBatchId(string batchId) =>
+        Guid.TryParse(batchId, out var parsed)
+            ? parsed
+            : throw new ArgumentException($"'{batchId}' is not a retry batch id issued by this store.", nameof(batchId));
 }

src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFRecoverabilityIngestionUnitOfWork.cs

@@ -36,7 +36,7 @@ public Task RecordFailedProcessingAttempt(MessageContext context,
             AttemptedAt = processingAttempt.AttemptedAt,
             TimeOfFailure = processingAttempt.FailureDetails.TimeOfFailure,
             Groups = groups,
-            HeadersJson = JsonSerializer.Serialize(processingAttempt.Headers, HeadersJsonContext.Default.DictionaryStringString),
+            HeadersJson = MessageHeaders.Write(processingAttempt.Headers),
             MessageId = processingAttempt.MessageId,
             MessageType = GetMetadata<string>(processingAttempt, "MessageType"),
             TimeSent = GetMetadata<DateTime?>(processingAttempt, "TimeSent"),

src/ServiceControl.Persistence.EFCore/Infrastructure/MessageHeaders.cs

@@ -0,0 +1,18 @@
+namespace ServiceControl.Persistence.EFCore.Infrastructure;
+
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+// The headers of a failed message are stored verbatim as the HeadersJson column.
+static class MessageHeaders
+{
+    public static string Write(Dictionary<string, string> headers) =>
+        JsonSerializer.Serialize(headers, HeadersJsonContext.Default.DictionaryStringString);
+
+    public static Dictionary<string, string> Read(string headersJson) =>
+        JsonSerializer.Deserialize(headersJson, HeadersJsonContext.Default.DictionaryStringString) ?? [];
+}
+
+// Source generated serialization, which keeps the reflection-based serializer off the ingestion hot path.
+[JsonSerializable(typeof(Dictionary<string, string>))]
+partial class HeadersJsonContext : JsonSerializerContext;

src/ServiceControl.Persistence.Tests.PostgreSql/ServiceControl.Persistence.Tests.PostgreSql.csproj

@@ -35,7 +35,6 @@
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\EditHandlerAuditTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\EditMessageTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\RetryConfirmationProcessorTests.cs" />
-    <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\RetryStagingStoreTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\AuditServiceMetadataTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\BrokerMetadataTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\EndpointsTests.cs" />

src/ServiceControl.Persistence.Tests.SqlServer/ServiceControl.Persistence.Tests.SqlServer.csproj

@@ -35,7 +35,6 @@
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\EditHandlerAuditTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\EditMessageTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\RetryConfirmationProcessorTests.cs" />
-    <Compile Remove="..\ServiceControl.Persistence.Tests\Recoverability\RetryStagingStoreTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\AuditServiceMetadataTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\BrokerMetadataTests.cs" />
     <Compile Remove="..\ServiceControl.Persistence.Tests\Throughput\EndpointsTests.cs" />

src/ServiceControl.Persistence.Tests/EFCore/PersistenceTestsContext.cs

@@ -8,7 +8,7 @@ namespace ServiceControl.Persistence.Tests;
 using System.Threading.Tasks;
 using EFCore.DbContexts;
 using EFCore.Entities;
-using EFCore.Implementation.UnitOfWork;
+using EFCore.Infrastructure;
 using MessageFailures;
 using Microsoft.Extensions.DependencyInjection;
 using NServiceBus;
@@ -41,7 +41,7 @@ static async Task InsertFailedMessagesDirect(IServiceProvider serviceProvider, F
                 FirstTimeOfFailure = ordered.Min(pa => pa.FailureDetails.TimeOfFailure),
                 LastTimeOfFailure = ordered.Max(pa => pa.FailureDetails.TimeOfFailure),
                 LastAttemptedAt = attempt.AttemptedAt,
-                HeadersJson = JsonSerializer.Serialize(attempt.Headers, HeadersJsonContext.Default.DictionaryStringString),
+                HeadersJson = MessageHeaders.Write(attempt.Headers),
                 MessageId = attempt.MessageId,
                 MessageType = GetMetadata<string>(attempt, "MessageType"),
                 TimeSent = GetMetadata<DateTime?>(attempt, "TimeSent"),