Skip to content

Push the cancellation token through the retry stores #5732

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
johnsimons merged 1 commit into from
Aug 11, 2026
Merged

Push the cancellation token through the retry stores #5732

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
61 changes: 31 additions & 30 deletions src/ServiceControl.Persistence.EFCore/Implementation/RetryBatchStore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@ public class RetryBatchStore(IServiceScopeFactory scopeFactory, IRetryBatchSqlDi
public Task<string> CreateBatch(string retrySessionId, string requestId, RetryType retryType,
string[] failedMessageRetryIds, string originator, DateTime startTime, DateTime? last = null,
string? batchName = null, string? classifier = null,
string? initiatedById = null, string? initiatedByName = null, string? operationId = null) =>
string? initiatedById = null, string? initiatedByName = null, string? operationId = null,
CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var batch = new RetryBatchEntity
Expand All @@ -36,15 +37,15 @@ public Task<string> CreateBatch(string retrySessionId, string requestId, RetryTy

dbContext.RetryBatches.Add(batch);

await dbContext.SaveChangesAsync();
await dbContext.SaveChangesAsync(cancellationToken);

return batch.Id.ToString();
});

/// <summary>
/// Claims the messages for the batch. A message already claimed by another batch keeps that claim, so the batch it is staged with is whichever one got there first.
/// </summary>
public Task AssignMessagesToBatch(string batchId, string[] messageIds) =>
public Task AssignMessagesToBatch(string batchId, string[] messageIds, CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var batch = ParseBatchId(batchId);
Expand Down Expand Up @@ -75,21 +76,21 @@ public Task AssignMessagesToBatch(string batchId, string[] messageIds) =>

await strategy.ExecuteAsync(async () =>
{
await using var transaction = await dbContext.Database.BeginTransactionAsync();
await using var transaction = await dbContext.Database.BeginTransactionAsync(cancellationToken);

await dialect.InsertMissingRetryClaims(dbContext, claims, CancellationToken.None);

await transaction.CommitAsync();
await transaction.CommitAsync(cancellationToken);
});
});

public Task MoveBatchToStaging(string batchId)
public Task MoveBatchToStaging(string batchId, CancellationToken cancellationToken = default)
{
var batch = ParseBatchId(batchId);

return ExecuteWithDbContext(dbContext => dbContext.RetryBatches
.Where(row => row.Id == batch)
.ExecuteUpdateAsync(setters => setters.SetProperty(row => row.Status, RetryBatchStatus.Staging)));
.ExecuteUpdateAsync(setters => setters.SetProperty(row => row.Status, RetryBatchStatus.Staging), cancellationToken));
}

// Batch ids only ever come from CreateBatch, so anything else is a programming error.
Expand All @@ -98,22 +99,22 @@ static Guid ParseBatchId(string batchId) =>
? parsed
: throw new ArgumentException($"'{batchId}' is not a retry batch id issued by this store.", nameof(batchId));

public Task<QueryResult<IList<RetryBatch>>> GetOrphanedBatches(string retrySessionId) =>
public Task<QueryResult<IList<RetryBatch>>> GetOrphanedBatches(string retrySessionId, CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var orphaned = await dbContext.RetryBatches
.AsNoTracking()
.Where(batch => batch.Status == RetryBatchStatus.MarkingDocuments && batch.RetrySessionId != retrySessionId)
.ToListAsync();
.ToListAsync(cancellationToken);

var messageCounts = await CountMessages(dbContext, [.. orphaned.Select(batch => batch.Id)]);
var messageCounts = await CountMessages(dbContext, [.. orphaned.Select(batch => batch.Id)], cancellationToken);

IList<RetryBatch> batches = [.. orphaned.Select(batch => batch.ToRetryBatch(messageCounts.GetValueOrDefault(batch.Id)))];

return new QueryResult<IList<RetryBatch>>(batches, new QueryStatsInfo(string.Empty, batches.Count, false));
});

public Task<IList<RetryBatchGroup>> GetAvailableBatchGroups() =>
public Task<IList<RetryBatchGroup>> GetAvailableBatchGroups(CancellationToken cancellationToken = default) =>
ExecuteWithDbContext<IList<RetryBatchGroup>>(async dbContext =>
{
var groups = await dbContext.RetryBatches
Expand All @@ -132,7 +133,7 @@ public Task<IList<RetryBatchGroup>> GetAvailableBatchGroups() =>
Originator = group.Max(batch => batch.Originator),
Classifier = group.Max(batch => batch.Classifier)
})
.ToListAsync();
.ToListAsync(cancellationToken);

return [.. groups.Select(group => new RetryBatchGroup
{
Expand All @@ -148,12 +149,12 @@ public Task<IList<RetryBatchGroup>> GetAvailableBatchGroups() =>
})];
});

public Task<ForwardingRetryBatch?> GetCurrentForwardingBatch() =>
public Task<ForwardingRetryBatch?> GetCurrentForwardingBatch(CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var nowForwarding = await dbContext.RetryBatchNowForwarding
.AsNoTracking()
.SingleOrDefaultAsync();
.SingleOrDefaultAsync(cancellationToken);

if (nowForwarding == null)
{
Expand All @@ -164,45 +165,45 @@ public Task<IList<RetryBatchGroup>> GetAvailableBatchGroups() =>
.AsNoTracking()
.Where(batch => batch.Id == nowForwarding.RetryBatchId)
.Select(batch => new ForwardingRetryBatch(batch.RequestId, batch.RetryType, batch.Originator!, batch.Classifier!))
.SingleOrDefaultAsync();
.SingleOrDefaultAsync(cancellationToken);
});

public Task ForEachUnresolvedMessage(Func<string, DateTime, Task> callback) =>
ForEach(Unresolved, callback);
public Task ForEachUnresolvedMessage(Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken = default) =>
ForEach(Unresolved, callback, cancellationToken);

public Task ForEachUnresolvedMessageForEndpoint(string endpoint, Func<string, DateTime, Task> callback) =>
public Task ForEachUnresolvedMessageForEndpoint(string endpoint, Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken = default) =>
ForEach(dbContext => Unresolved(dbContext)
.Where(message => message.ReceivingEndpointName == endpoint), callback);
.Where(message => message.ReceivingEndpointName == endpoint), callback, cancellationToken);

public Task ForEachMessageForQueueAddress(string failedQueueAddress, FailedMessageStatus status, Func<string, DateTime, Task> callback) =>
public Task ForEachMessageForQueueAddress(string failedQueueAddress, FailedMessageStatus status, Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken = default) =>
ForEach(dbContext => Unresolved(dbContext)
.Where(message => message.FailingEndpointAddress == failedQueueAddress && message.Status == status), callback);
.Where(message => message.FailingEndpointAddress == failedQueueAddress && message.Status == status), callback, cancellationToken);

public Task ForEachUnresolvedMessageInGroup(string groupId, Func<string, DateTime, Task> callback) =>
public Task ForEachUnresolvedMessageInGroup(string groupId, Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken = default) =>
ForEach(dbContext => Unresolved(dbContext)
.Where(message => dbContext.FailedMessageGroups.Any(group => group.GroupId == groupId && group.FailedMessageUniqueId == message.UniqueMessageId)), callback);
.Where(message => dbContext.FailedMessageGroups.Any(group => group.GroupId == groupId && group.FailedMessageUniqueId == message.UniqueMessageId)), callback, cancellationToken);

Task ForEach(Func<ServiceControlDbContext, IQueryable<FailedMessageEntity>> query, Func<string, DateTime, Task> callback) =>
ExecuteWithDbContext(dbContext => Stream(query(dbContext), callback));
Task ForEach(Func<ServiceControlDbContext, IQueryable<FailedMessageEntity>> query, Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken) =>
ExecuteWithDbContext(dbContext => Stream(query(dbContext), callback, cancellationToken));

static IQueryable<FailedMessageEntity> Unresolved(ServiceControlDbContext dbContext) =>
dbContext.FailedMessages
.AsNoTracking()
.Where(message => message.Status == FailedMessageStatus.Unresolved);

static async Task Stream(IQueryable<FailedMessageEntity> messages, Func<string, DateTime, Task> callback)
static async Task Stream(IQueryable<FailedMessageEntity> messages, Func<string, DateTime, CancellationToken, Task> callback, CancellationToken cancellationToken)
{
var rows = messages
.Select(message => new { message.UniqueMessageId, message.LastTimeOfFailure })
.AsAsyncEnumerable();

await foreach (var row in rows)
await foreach (var row in rows.WithCancellation(cancellationToken))
{
await callback(row.UniqueMessageId.ToString(), row.LastTimeOfFailure);
await callback(row.UniqueMessageId.ToString(), row.LastTimeOfFailure, cancellationToken);
}
}

static async Task<Dictionary<Guid, int>> CountMessages(ServiceControlDbContext dbContext, Guid[] batchIds)
static async Task<Dictionary<Guid, int>> CountMessages(ServiceControlDbContext dbContext, Guid[] batchIds, CancellationToken cancellationToken)
{
if (batchIds.Length == 0)
{
Expand All @@ -214,6 +215,6 @@ static async Task<Dictionary<Guid, int>> CountMessages(ServiceControlDbContext d
.Where(retry => batchIds.Contains(retry.RetryBatchId))
.GroupBy(retry => retry.RetryBatchId)
.Select(group => new { RetryBatchId = group.Key, MessageCount = group.Count() })
.ToDictionaryAsync(row => row.RetryBatchId, row => row.MessageCount);
.ToDictionaryAsync(row => row.RetryBatchId, row => row.MessageCount, cancellationToken);
}
}
35 changes: 18 additions & 17 deletions src/ServiceControl.Persistence.EFCore/Implementation/RetryHistoryDataStore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;

public class RetryHistoryDataStore(IServiceScopeFactory scopeFactory) : DataStoreBase(scopeFactory), IRetryHistoryDataStore
{
public Task<RetryHistory> GetRetryHistory() =>
public Task<RetryHistory> GetRetryHistory(CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var historicOperations = await dbContext.HistoricRetryOperations
Expand All @@ -25,7 +25,7 @@ public Task<RetryHistory> GetRetryHistory() =>
Failed = operation.Failed,
NumberOfMessagesProcessed = operation.NumberOfMessagesProcessed
})
.ToListAsync();
.ToListAsync(cancellationToken);

var unacknowledgedOperations = await dbContext.UnacknowledgedRetryOperations
.AsNoTracking()
Expand All @@ -41,7 +41,7 @@ public Task<RetryHistory> GetRetryHistory() =>
Failed = operation.Failed,
NumberOfMessagesProcessed = operation.NumberOfMessagesProcessed
})
.ToListAsync();
.ToListAsync(cancellationToken);

return new RetryHistory
{
Expand All @@ -51,14 +51,15 @@ public Task<RetryHistory> GetRetryHistory() =>
});

public Task RecordRetryOperationCompleted(string requestId, RetryType retryType, DateTime startTime, DateTime completionTime,
string originator, string classifier, bool messageFailed, int numberOfMessagesProcessed, DateTime lastProcessed, int retryHistoryDepth) =>
string originator, string classifier, bool messageFailed, int numberOfMessagesProcessed, DateTime lastProcessed, int retryHistoryDepth,
CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var strategy = dbContext.Database.CreateExecutionStrategy();

await strategy.ExecuteAsync(async () =>
{
await using var transaction = await dbContext.Database.BeginTransactionAsync();
await using var transaction = await dbContext.Database.BeginTransactionAsync(cancellationToken);

dbContext.HistoricRetryOperations.Add(new HistoricRetryOperationEntity
{
Expand All @@ -74,24 +75,24 @@ await strategy.ExecuteAsync(async () =>
if (NeedsAcknowledgement(retryType))
{
await RecordUnacknowledged(dbContext, requestId, retryType, startTime, completionTime,
originator, classifier, messageFailed, numberOfMessagesProcessed, lastProcessed);
originator, classifier, messageFailed, numberOfMessagesProcessed, lastProcessed, cancellationToken);
}

await dbContext.SaveChangesAsync();
await dbContext.SaveChangesAsync(cancellationToken);

// After the insert, so the operation just recorded competes for a place in the history.
await TrimHistory(dbContext, retryHistoryDepth);
await TrimHistory(dbContext, retryHistoryDepth, cancellationToken);

await transaction.CommitAsync();
await transaction.CommitAsync(cancellationToken);
});
});

public Task<bool> AcknowledgeRetryGroup(string groupId) =>
public Task<bool> AcknowledgeRetryGroup(string groupId, CancellationToken cancellationToken = default) =>
ExecuteWithDbContext(async dbContext =>
{
var acknowledged = await dbContext.UnacknowledgedRetryOperations
.Where(operation => operation.RequestId == groupId && operation.RetryType == RetryType.FailureGroup)
.ExecuteDeleteAsync();
.ExecuteDeleteAsync(cancellationToken);

return acknowledged > 0;
});
Expand All @@ -101,10 +102,10 @@ static bool NeedsAcknowledgement(RetryType retryType) =>

static async Task RecordUnacknowledged(ServiceControlDbContext dbContext, string requestId, RetryType retryType,
DateTime startTime, DateTime completionTime, string originator, string classifier, bool messageFailed,
int numberOfMessagesProcessed, DateTime lastProcessed)
int numberOfMessagesProcessed, DateTime lastProcessed, CancellationToken cancellationToken)
{
var unacknowledged = await dbContext.UnacknowledgedRetryOperations
.SingleOrDefaultAsync(operation => operation.RequestId == requestId && operation.RetryType == retryType);
.SingleOrDefaultAsync(operation => operation.RequestId == requestId && operation.RetryType == retryType, cancellationToken);

if (unacknowledged == null)
{
Expand All @@ -121,11 +122,11 @@ static async Task RecordUnacknowledged(ServiceControlDbContext dbContext, string
unacknowledged.NumberOfMessagesProcessed = numberOfMessagesProcessed;
}

static async Task TrimHistory(ServiceControlDbContext dbContext, int retryHistoryDepth)
static async Task TrimHistory(ServiceControlDbContext dbContext, int retryHistoryDepth, CancellationToken cancellationToken)
{
if (retryHistoryDepth <= 0)
{
await dbContext.HistoricRetryOperations.ExecuteDeleteAsync();
await dbContext.HistoricRetryOperations.ExecuteDeleteAsync(cancellationToken);
return;
}

Expand All @@ -136,7 +137,7 @@ static async Task TrimHistory(ServiceControlDbContext dbContext, int retryHistor
.ThenByDescending(operation => operation.Id)
.Skip(retryHistoryDepth - 1)
.Select(operation => new { operation.CompletionTime, operation.Id })
.FirstOrDefaultAsync();
.FirstOrDefaultAsync(cancellationToken);

if (cutoff == null)
{
Expand All @@ -146,6 +147,6 @@ static async Task TrimHistory(ServiceControlDbContext dbContext, int retryHistor
await dbContext.HistoricRetryOperations
.Where(operation => operation.CompletionTime < cutoff.CompletionTime
|| (operation.CompletionTime == cutoff.CompletionTime && operation.Id < cutoff.Id))
.ExecuteDeleteAsync();
.ExecuteDeleteAsync(cancellationToken);
}
}
Loading
Loading