From 113a15bd45094b96898cdd294eae51f61142b739 Mon Sep 17 00:00:00 2001 From: Rhys Bevilaqua Date: Thu, 24 Sep 2026 13:25:32 +0800 Subject: [PATCH] Add Complete() to IAuditIngestionUnitOfWork to support transaction rollback --- .../InMemoryAuditIngestionUnitOfWork.cs | 3 ++ .../RavenAuditIngestionUnitOfWork.cs | 30 ++++++++++++++- .../RetentionTests.cs | 8 ++-- .../SagaDetailsIndexTests.cs | 4 +- .../AuditCountingTests.cs | 4 +- .../AuditTests.cs | 37 ++++++++++--------- .../SagaHistoryTests.cs | 4 +- .../UnitOfWork/IAuditIngestionUnitOfWork.cs | 5 +++ .../Auditing/AuditPersister.cs | 2 + 9 files changed, 67 insertions(+), 30 deletions(-) diff --git a/src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWork.cs b/src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWork.cs index bc51d76815..649b916220 100644 --- a/src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWork.cs +++ b/src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWork.cs @@ -15,6 +15,9 @@ class InMemoryAuditIngestionUnitOfWork( { public ValueTask DisposeAsync() => ValueTask.CompletedTask; + // InMemory writes are immediate, so Complete has nothing to commit. + public Task Complete(CancellationToken cancellationToken = default) => Task.CompletedTask; + public async Task RecordProcessedMessage(ProcessedMessage processedMessage, ReadOnlyMemory body, CancellationToken cancellationToken = default) { if (!body.IsEmpty) diff --git a/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditIngestionUnitOfWork.cs b/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditIngestionUnitOfWork.cs index 664dd7020b..18eb78e7f7 100644 --- a/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditIngestionUnitOfWork.cs +++ b/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditIngestionUnitOfWork.cs @@ -49,10 +49,36 @@ MetadataAsDictionary GetExpirationMetadata() => public Task RecordSagaSnapshot(SagaSnapshot sagaSnapshot, CancellationToken cancellationToken = default) => bulkInsert.StoreAsync(sagaSnapshot, GetExpirationMetadata()); - public async ValueTask DisposeAsync() + bool completed; + + public async Task Complete(CancellationToken cancellationToken = default) { + // Closing the bulk insert flushes its remaining buffered documents. await bulkInsert.DisposeAsync(); - timedCancellationSource.Dispose(); + completed = true; + } + + public async ValueTask DisposeAsync() + { + try + { + if (!completed) + { + // Bulk inserts are not atomic; abort prevents further writes, not earlier ones. + try + { + await bulkInsert.AbortAsync(); + } + finally + { + await bulkInsert.DisposeAsync(); + } + } + } + finally + { + timedCancellationSource.Dispose(); + } } } } \ No newline at end of file diff --git a/src/ServiceControl.Audit.Persistence.Tests.RavenDB/RetentionTests.cs b/src/ServiceControl.Audit.Persistence.Tests.RavenDB/RetentionTests.cs index 2d850fa6cd..b020506078 100644 --- a/src/ServiceControl.Audit.Persistence.Tests.RavenDB/RetentionTests.cs +++ b/src/ServiceControl.Audit.Persistence.Tests.RavenDB/RetentionTests.cs @@ -134,23 +134,23 @@ ProcessedMessage MakeMessage( async Task IngestProcessedMessagesAudits(params ProcessedMessage[] processedMessages) { - var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); + await using var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); foreach (var processedMessage in processedMessages) { await unitOfWork.RecordProcessedMessage(processedMessage); } - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } async Task IngestSagaAudits(params SagaSnapshot[] snapshots) { - var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); + await using var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); foreach (var snapshot in snapshots) { await unitOfWork.RecordSagaSnapshot(snapshot); } - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } } diff --git a/src/ServiceControl.Audit.Persistence.Tests.RavenDB/SagaDetailsIndexTests.cs b/src/ServiceControl.Audit.Persistence.Tests.RavenDB/SagaDetailsIndexTests.cs index 71ec7400d2..5f9007740f 100644 --- a/src/ServiceControl.Audit.Persistence.Tests.RavenDB/SagaDetailsIndexTests.cs +++ b/src/ServiceControl.Audit.Persistence.Tests.RavenDB/SagaDetailsIndexTests.cs @@ -109,13 +109,13 @@ await IngestSagaAudits(new SagaSnapshot async Task IngestSagaAudits(params SagaSnapshot[] snapshots) { - var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); + await using var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); foreach (var snapshot in snapshots) { await unitOfWork.RecordSagaSnapshot(snapshot); } - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } } \ No newline at end of file diff --git a/src/ServiceControl.Audit.Persistence.Tests/AuditCountingTests.cs b/src/ServiceControl.Audit.Persistence.Tests/AuditCountingTests.cs index 6486fa0fdc..c7bddaf10a 100644 --- a/src/ServiceControl.Audit.Persistence.Tests/AuditCountingTests.cs +++ b/src/ServiceControl.Audit.Persistence.Tests/AuditCountingTests.cs @@ -137,12 +137,12 @@ static ProcessedMessage MakeMessage(string processingEndpoint, string sendingEnd async Task IngestProcessedMessagesAudits(params ProcessedMessage[] processedMessages) { - var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); + await using var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); foreach (var processedMessage in processedMessages) { await unitOfWork.RecordProcessedMessage(processedMessage); } - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } } diff --git a/src/ServiceControl.Audit.Persistence.Tests/AuditTests.cs b/src/ServiceControl.Audit.Persistence.Tests/AuditTests.cs index 96d2fc6b37..b9a1e8b04c 100644 --- a/src/ServiceControl.Audit.Persistence.Tests/AuditTests.cs +++ b/src/ServiceControl.Audit.Persistence.Tests/AuditTests.cs @@ -28,7 +28,7 @@ public async Task Basic_Roundtrip() await IngestProcessedMessagesAudits( message - ); + ); var queryResult = await DataStore.QueryMessages("MyMessageId", new PagingInfo(), new SortInfo("Id", "asc"), cancellationToken: TestContext.CurrentContext.CancellationToken); @@ -77,11 +77,12 @@ await IngestProcessedMessagesAudits( Assert.That(queryResult.Results, Has.Count.EqualTo(2)); } + [Test] public async Task Can_roundtrip_message_body() { string expectedContentType = "text/plain"; - var unitOfWork = await StartAuditUnitOfWork(1); + await using var unitOfWork = await StartAuditUnitOfWork(1); var body = new byte[100]; Random.Shared.NextBytes(body); @@ -89,7 +90,7 @@ public async Task Can_roundtrip_message_body() await unitOfWork.RecordProcessedMessage(processedMessage, body); - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); var bodyId = GetBodyId(processedMessage); @@ -118,7 +119,7 @@ public async Task Can_roundtrip_message_body() [Test] public async Task Message_body_validator_is_stable_across_reads() { - var unitOfWork = await StartAuditUnitOfWork(1); + await using var unitOfWork = await StartAuditUnitOfWork(1); var body = new byte[100]; Random.Shared.NextBytes(body); @@ -126,7 +127,7 @@ public async Task Message_body_validator_is_stable_across_reads() await unitOfWork.RecordProcessedMessage(processedMessage, body); - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); var bodyId = GetBodyId(processedMessage); @@ -140,7 +141,7 @@ public async Task Message_body_validator_is_stable_across_reads() [Test] public async Task Does_respect_max_message_body() { - var unitOfWork = await StartAuditUnitOfWork(1); + await using var unitOfWork = await StartAuditUnitOfWork(1); var body = new byte[MAX_BODY_SIZE + 1000]; Random.Shared.NextBytes(body); @@ -148,7 +149,7 @@ public async Task Does_respect_max_message_body() await unitOfWork.RecordProcessedMessage(processedMessage, body); - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); var bodyId = GetBodyId(processedMessage); @@ -160,13 +161,12 @@ public async Task Does_respect_max_message_body() Assert.That(retrievedMessage.Found, Is.True); Assert.That(retrievedMessage.HasContent, Is.False); } - } [Test] public async Task Deduplicates_messages_in_same_batch() { - var unitOfWork = await StartAuditUnitOfWork(1); + await using var unitOfWork = await StartAuditUnitOfWork(1); var messageId = "duplicatedId"; var processingEndpoint = "endpoint"; var processingStarted = DateTimeOffset.UtcNow; @@ -176,7 +176,7 @@ public async Task Deduplicates_messages_in_same_batch() await unitOfWork.RecordProcessedMessage(processedMessage); await unitOfWork.RecordProcessedMessage(duplicatedMessage); - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); @@ -193,14 +193,14 @@ public async Task Deduplicates_messages_in_different_batches() var processingStarted = DateTimeOffset.UtcNow; var processedMessage = MakeMessage(messageId: messageId, processingEndpoint: processingEndpoint, processingStarted: processingStarted); - var unitOfWork1 = await StartAuditUnitOfWork(1); + await using var unitOfWork1 = await StartAuditUnitOfWork(1); await unitOfWork1.RecordProcessedMessage(processedMessage); - await unitOfWork1.DisposeAsync(); + await unitOfWork1.Complete(); var duplicatedMessage = MakeMessage(messageId: messageId, processingEndpoint: processingEndpoint, processingStarted: processingStarted); - var unitOfWork2 = await StartAuditUnitOfWork(1); + await using var unitOfWork2 = await StartAuditUnitOfWork(1); await unitOfWork2.RecordProcessedMessage(duplicatedMessage); - await unitOfWork2.DisposeAsync(); + await unitOfWork2.Complete(); await configuration.CompleteDBOperation(); @@ -212,7 +212,7 @@ public async Task Deduplicates_messages_in_different_batches() [Test] public async Task Does_not_deduplicate_with_different_processing_started_header() { - var unitOfWork = await StartAuditUnitOfWork(1); + await using var unitOfWork = await StartAuditUnitOfWork(1); var messageId = "duplicatedId"; var processingEndpoint = "endpoint"; var processingStarted = DateTimeOffset.UtcNow; @@ -223,7 +223,7 @@ public async Task Does_not_deduplicate_with_different_processing_started_header( await unitOfWork.RecordProcessedMessage(processedMessage); await unitOfWork.RecordProcessedMessage(duplicatedMessage); - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); @@ -294,12 +294,13 @@ ProcessedMessage MakeMessage( async Task IngestProcessedMessagesAudits(params ProcessedMessage[] processedMessages) { - var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); + await using var unitOfWork = await StartAuditUnitOfWork(processedMessages.Length); foreach (var processedMessage in processedMessages) { await unitOfWork.RecordProcessedMessage(processedMessage); } - await unitOfWork.DisposeAsync(); + + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } diff --git a/src/ServiceControl.Audit.Persistence.Tests/SagaHistoryTests.cs b/src/ServiceControl.Audit.Persistence.Tests/SagaHistoryTests.cs index 197821b3cf..dcdfb69a16 100644 --- a/src/ServiceControl.Audit.Persistence.Tests/SagaHistoryTests.cs +++ b/src/ServiceControl.Audit.Persistence.Tests/SagaHistoryTests.cs @@ -61,12 +61,12 @@ await IngestSagaAudits( async Task IngestSagaAudits(params SagaSnapshot[] snapshots) { - var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); + await using var unitOfWork = await StartAuditUnitOfWork(snapshots.Length); foreach (var snapshot in snapshots) { await unitOfWork.RecordSagaSnapshot(snapshot); } - await unitOfWork.DisposeAsync(); + await unitOfWork.Complete(); await configuration.CompleteDBOperation(); } } diff --git a/src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWork.cs b/src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWork.cs index 01b18c6415..3b6a4b9e87 100644 --- a/src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWork.cs +++ b/src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWork.cs @@ -10,5 +10,10 @@ public interface IAuditIngestionUnitOfWork : IAsyncDisposable { Task RecordProcessedMessage(ProcessedMessage processedMessage, ReadOnlyMemory body = default, CancellationToken cancellationToken = default); Task RecordSagaSnapshot(SagaSnapshot sagaSnapshot, CancellationToken cancellationToken = default); + + /// + /// Signals that all intended records have been added and the unit of work should commit. + /// + Task Complete(CancellationToken cancellationToken = default); } } \ No newline at end of file diff --git a/src/ServiceControl.Audit/Auditing/AuditPersister.cs b/src/ServiceControl.Audit/Auditing/AuditPersister.cs index d82b1fa3ab..f6089e2f28 100644 --- a/src/ServiceControl.Audit/Auditing/AuditPersister.cs +++ b/src/ServiceControl.Audit/Auditing/AuditPersister.cs @@ -57,6 +57,8 @@ public async Task> Persist(IReadOnlyList