Skip to content

Add Complete() method to IAuditIngestionUnitOfWork #5918

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rbev merged 1 commit into from
Sep 25, 2026
Merged

Add Complete() method to IAuditIngestionUnitOfWork #5918

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWork.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<byte> body, CancellationToken cancellationToken = default)
{
if (!body.IsEmpty)
Expand Down
30 changes: 28 additions & 2 deletions src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditIngestionUnitOfWork.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
}
}
8 changes: 4 additions & 4 deletions src/ServiceControl.Audit.Persistence.Tests.RavenDB/RetentionTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/ServiceControl.Audit.Persistence.Tests.RavenDB/SagaDetailsIndexTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
4 changes: 2 additions & 2 deletions src/ServiceControl.Audit.Persistence.Tests/AuditCountingTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
37 changes: 19 additions & 18 deletions src/ServiceControl.Audit.Persistence.Tests/AuditTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -77,19 +77,20 @@ 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);
var processedMessage = MakeMessage();

await unitOfWork.RecordProcessedMessage(processedMessage, body);

await unitOfWork.DisposeAsync();
await unitOfWork.Complete();

var bodyId = GetBodyId(processedMessage);

Expand Down Expand Up @@ -118,15 +119,15 @@ 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);
var processedMessage = MakeMessage();

await unitOfWork.RecordProcessedMessage(processedMessage, body);

await unitOfWork.DisposeAsync();
await unitOfWork.Complete();

var bodyId = GetBodyId(processedMessage);

Expand All @@ -140,15 +141,15 @@ 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);
var processedMessage = MakeMessage();

await unitOfWork.RecordProcessedMessage(processedMessage, body);

await unitOfWork.DisposeAsync();
await unitOfWork.Complete();

var bodyId = GetBodyId(processedMessage);

Expand All @@ -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;
Expand All @@ -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();

Expand All @@ -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();

Expand All @@ -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;
Expand All @@ -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();

Expand Down Expand Up @@ -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();
}

Expand Down
4 changes: 2 additions & 2 deletions src/ServiceControl.Audit.Persistence.Tests/SagaHistoryTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
5 changes: 5 additions & 0 deletions src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWork.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,5 +10,10 @@ public interface IAuditIngestionUnitOfWork : IAsyncDisposable
{
Task RecordProcessedMessage(ProcessedMessage processedMessage, ReadOnlyMemory<byte> body = default, CancellationToken cancellationToken = default);
Task RecordSagaSnapshot(SagaSnapshot sagaSnapshot, CancellationToken cancellationToken = default);

/// <summary>
/// Signals that all intended records have been added and the unit of work should commit.
/// </summary>
Task Complete(CancellationToken cancellationToken = default);
}
}
2 changes: 2 additions & 0 deletions src/ServiceControl.Audit/Auditing/AuditPersister.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageCo

storedContexts.Add(context);
}

await unitOfWork.Complete(cancellationToken);
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
Expand Down
Loading