using JobTrackerApi.Data; using JobTrackerApi.Models; using Microsoft.EntityFrameworkCore; namespace JobTrackerApi.Services; public sealed record CreateUserOperation( string TaskType, string IdempotencyKey, string EntitlementDecision, string PrivacyPolicy, string? SubjectType = null, string? SubjectId = null, int Priority = 0, int MaxAttempts = 3, DateTime? DeadlineAtUtc = null); public sealed record UserOperationCreation(UserOperation Operation, bool Created); public sealed record UserOperationLease(Guid OperationId, string OwnerUserId, string LeaseToken, string TaskType, string PrivacyPolicy, string? SubjectType, string? SubjectId, int AttemptCount, DateTime? DeadlineAtUtc); public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timeProvider) { private DateTime UtcNow => timeProvider.GetUtcNow().UtcDateTime; public Task GetAsync(Guid operationId, CancellationToken cancellationToken) { EnsureOwnerScope(); return db.UserOperations.AsNoTracking().FirstOrDefaultAsync(operation => operation.Id == operationId, cancellationToken); } public Task> ListAsync(int limit, CancellationToken cancellationToken) { EnsureOwnerScope(); if (limit is < 1 or > 100) throw new ArgumentOutOfRangeException(nameof(limit)); return db.UserOperations.AsNoTracking() .OrderByDescending(operation => operation.CreatedAtUtc) .Take(limit) .ToListAsync(cancellationToken); } public Task FindByIdempotencyAsync(string taskType, string idempotencyKey, CancellationToken cancellationToken) { EnsureOwnerScope(); return db.UserOperations.AsNoTracking().FirstOrDefaultAsync( operation => operation.TaskType == taskType && operation.IdempotencyKey == idempotencyKey, cancellationToken); } public async Task CreateAsync(CreateUserOperation request, CancellationToken cancellationToken) { var owner = db.CurrentUserId ?? throw new InvalidOperationException("Operation creation requires an authenticated owner scope."); Validate(request); var existing = await db.UserOperations.FirstOrDefaultAsync( operation => operation.TaskType == request.TaskType && operation.IdempotencyKey == request.IdempotencyKey, cancellationToken); if (existing is not null) return new UserOperationCreation(existing, false); var now = UtcNow; var operation = new UserOperation { Id = Guid.NewGuid(), OwnerUserId = owner, TaskType = request.TaskType, IdempotencyKey = request.IdempotencyKey, EntitlementDecision = request.EntitlementDecision, PrivacyPolicy = request.PrivacyPolicy, SubjectType = request.SubjectType, SubjectId = request.SubjectId, Priority = request.Priority, MaxAttempts = request.MaxAttempts, CreatedAtUtc = now, AvailableAtUtc = now, DeadlineAtUtc = request.DeadlineAtUtc, }; db.UserOperations.Add(operation); try { await db.SaveChangesAsync(cancellationToken); return new UserOperationCreation(operation, true); } catch (DbUpdateException) { db.Entry(operation).State = EntityState.Detached; existing = await db.UserOperations.FirstOrDefaultAsync( item => item.TaskType == request.TaskType && item.IdempotencyKey == request.IdempotencyKey, cancellationToken); if (existing is not null) return new UserOperationCreation(existing, false); throw; } } public async Task ClaimNextAsync( TimeSpan leaseDuration, CancellationToken cancellationToken, IReadOnlyCollection? allowedTaskTypes = null) { if (db.CurrentUserId is not null) throw new InvalidOperationException("Worker claims require a neutral background scope."); ValidateLeaseDuration(leaseDuration); var now = UtcNow; await RecoverExpiredLeasesAsync(now, cancellationToken); var candidates = db.UserOperations.IgnoreQueryFilters().AsNoTracking() .Where(operation => (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry) && operation.AvailableAtUtc <= now && operation.CancellationRequestedAtUtc == null && operation.AttemptCount < operation.MaxAttempts && (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now)); if (allowedTaskTypes is { Count: > 0 }) candidates = candidates.Where(operation => allowedTaskTypes.Contains(operation.TaskType)); var candidateIds = await candidates .OrderByDescending(operation => operation.Priority) .ThenBy(operation => operation.CreatedAtUtc) .Select(operation => operation.Id) .Take(16) .ToListAsync(cancellationToken); foreach (var candidateId in candidateIds) { var leaseToken = Guid.NewGuid().ToString("N"); var affected = await db.UserOperations.IgnoreQueryFilters() .Where(operation => operation.Id == candidateId && (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry) && operation.AvailableAtUtc <= now && operation.CancellationRequestedAtUtc == null && operation.AttemptCount < operation.MaxAttempts && (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now)) .ExecuteUpdateAsync(setters => setters .SetProperty(operation => operation.Status, OperationStatuses.Running) .SetProperty(operation => operation.LeaseToken, leaseToken) .SetProperty(operation => operation.LeaseExpiresAtUtc, now.Add(leaseDuration)) .SetProperty(operation => operation.LastHeartbeatAtUtc, now) .SetProperty(operation => operation.StartedAtUtc, operation => operation.StartedAtUtc ?? now) .SetProperty(operation => operation.AttemptCount, operation => operation.AttemptCount + 1), cancellationToken); if (affected != 1) continue; var claimed = await db.UserOperations.IgnoreQueryFilters().AsNoTracking() .SingleAsync(operation => operation.Id == candidateId && operation.LeaseToken == leaseToken, cancellationToken); return new UserOperationLease(claimed.Id, claimed.OwnerUserId, leaseToken, claimed.TaskType, claimed.PrivacyPolicy, claimed.SubjectType, claimed.SubjectId, claimed.AttemptCount, claimed.DeadlineAtUtc); } return null; } public Task HeartbeatAsync(Guid operationId, string leaseToken, TimeSpan leaseDuration, string? progressStage, int? progressPercent, CancellationToken cancellationToken) { EnsureOwnerScope(); ValidateLeaseDuration(leaseDuration); if (progressPercent is < 0 or > 100) throw new ArgumentOutOfRangeException(nameof(progressPercent)); ValidateOptional(progressStage, 64, nameof(progressStage)); var now = UtcNow; return db.UserOperations .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running && operation.LeaseToken == leaseToken) .ExecuteUpdateAsync(setters => setters .SetProperty(operation => operation.LeaseExpiresAtUtc, now.Add(leaseDuration)) .SetProperty(operation => operation.LastHeartbeatAtUtc, now) .SetProperty(operation => operation.ProgressStage, progressStage) .SetProperty(operation => operation.ProgressPercent, progressPercent), cancellationToken); } public Task CompleteAsync(Guid operationId, string leaseToken, string? resultReference, CancellationToken cancellationToken) => CompleteAsync(operationId, leaseToken, resultReference, null, null, null, cancellationToken); public async Task CompleteAsync( Guid operationId, string leaseToken, string? resultReference, string? provider, string? model, string? completionStage, CancellationToken cancellationToken) { EnsureOwnerScope(); ValidateOptional(resultReference, 256, nameof(resultReference)); ValidateOptional(provider, 128, nameof(provider)); ValidateOptional(model, 128, nameof(model)); ValidateOptional(completionStage, 64, nameof(completionStage)); var now = UtcNow; await using var transaction = await BeginTransactionAsync(cancellationToken); var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken); if (operation is null) return 0; var affected = await db.UserOperations .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running && operation.LeaseToken == leaseToken && operation.CancellationRequestedAtUtc == null) .ExecuteUpdateAsync(setters => setters .SetProperty(operation => operation.Status, OperationStatuses.Succeeded) .SetProperty(operation => operation.ResultReference, resultReference) .SetProperty(operation => operation.Provider, provider) .SetProperty(operation => operation.Model, model) .SetProperty(operation => operation.ProgressStage, completionStage) .SetProperty(operation => operation.CompletedAtUtc, now) .SetProperty(operation => operation.ProgressPercent, 100) .SetProperty(operation => operation.LeaseToken, (string?)null) .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken); if (affected == 1) { db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Succeeded, now)); await db.SaveChangesAsync(cancellationToken); if (transaction is not null) await transaction.CommitAsync(cancellationToken); } return affected; } public Task FailAsync(Guid operationId, string leaseToken, bool retryable, string category, string message, TimeSpan retryDelay, CancellationToken cancellationToken) => FailAsync(operationId, leaseToken, retryable, category, message, retryDelay, null, null, null, cancellationToken); public async Task FailAsync( Guid operationId, string leaseToken, bool retryable, string category, string message, TimeSpan retryDelay, string? provider, string? model, string? progressStage, CancellationToken cancellationToken) { EnsureOwnerScope(); ValidateRequired(category, 64, nameof(category)); ValidateRequired(message, 512, nameof(message)); ValidateOptional(provider, 128, nameof(provider)); ValidateOptional(model, 128, nameof(model)); ValidateOptional(progressStage, 64, nameof(progressStage)); if (retryDelay < TimeSpan.Zero || retryDelay > TimeSpan.FromHours(1)) throw new ArgumentOutOfRangeException(nameof(retryDelay)); await using var transaction = await BeginTransactionAsync(cancellationToken); var operation = await db.UserOperations.AsNoTracking() .FirstOrDefaultAsync(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.LeaseToken == leaseToken, cancellationToken); if (operation is null) return false; var now = UtcNow; var canRetry = retryable && operation.AttemptCount < operation.MaxAttempts && (operation.DeadlineAtUtc is null || operation.DeadlineAtUtc > now); var retryAt = now.Add(retryDelay); var affected = await db.UserOperations .Where(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.LeaseToken == leaseToken) .ExecuteUpdateAsync(setters => setters .SetProperty(item => item.Status, canRetry ? OperationStatuses.WaitingForRetry : OperationStatuses.Failed) .SetProperty(item => item.AvailableAtUtc, item => canRetry ? retryAt : item.AvailableAtUtc) .SetProperty(item => item.CompletedAtUtc, canRetry ? null : now) .SetProperty(item => item.FailureCategory, category) .SetProperty(item => item.FailureMessage, message) .SetProperty(item => item.Provider, provider) .SetProperty(item => item.Model, model) .SetProperty(item => item.ProgressStage, progressStage) .SetProperty(item => item.LeaseToken, (string?)null) .SetProperty(item => item.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken); if (affected == 1 && !canRetry) { db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Failed, now)); await db.SaveChangesAsync(cancellationToken); } if (affected == 1 && transaction is not null) await transaction.CommitAsync(cancellationToken); return affected == 1; } public async Task RequestCancellationAsync(Guid operationId, CancellationToken cancellationToken) { EnsureOwnerScope(); var now = UtcNow; await using var transaction = await BeginTransactionAsync(cancellationToken); var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken); if (operation is null || OperationStatuses.IsTerminal(operation.Status)) return false; var cancelled = await db.UserOperations .Where(item => item.Id == operationId && (item.Status == OperationStatuses.Queued || item.Status == OperationStatuses.WaitingForRetry || item.Status == OperationStatuses.WaitingForExternalFallback)) .ExecuteUpdateAsync(setters => setters .SetProperty(item => item.Status, OperationStatuses.Cancelled) .SetProperty(item => item.CompletedAtUtc, now), cancellationToken); if (cancelled == 1) { db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Cancelled, now)); await db.SaveChangesAsync(cancellationToken); if (transaction is not null) await transaction.CommitAsync(cancellationToken); return true; } var requested = await db.UserOperations .Where(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.CancellationRequestedAtUtc == null) .ExecuteUpdateAsync(setters => setters.SetProperty(item => item.CancellationRequestedAtUtc, now), cancellationToken); if (requested == 1 && transaction is not null) await transaction.CommitAsync(cancellationToken); return requested == 1; } public async Task AcknowledgeCancellationAsync(Guid operationId, string leaseToken, CancellationToken cancellationToken) { EnsureOwnerScope(); var now = UtcNow; await using var transaction = await BeginTransactionAsync(cancellationToken); var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken); if (operation is null) return 0; var affected = await db.UserOperations .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running && operation.LeaseToken == leaseToken && operation.CancellationRequestedAtUtc != null) .ExecuteUpdateAsync(setters => setters .SetProperty(operation => operation.Status, OperationStatuses.Cancelled) .SetProperty(operation => operation.CompletedAtUtc, now) .SetProperty(operation => operation.LeaseToken, (string?)null) .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken); if (affected == 1) { db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Cancelled, now)); await db.SaveChangesAsync(cancellationToken); if (transaction is not null) await transaction.CommitAsync(cancellationToken); } return affected; } public async Task RetryAsync(Guid operationId, CancellationToken cancellationToken) { EnsureOwnerScope(); await using var transaction = await BeginTransactionAsync(cancellationToken); var operation = await db.UserOperations.FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken); if (operation is null || operation.Status is not (OperationStatuses.Failed or OperationStatuses.Cancelled)) return false; operation.Status = OperationStatuses.Queued; operation.AttemptCount = 0; operation.AvailableAtUtc = UtcNow; operation.StartedAtUtc = null; operation.CompletedAtUtc = null; operation.CancellationRequestedAtUtc = null; operation.FailureCategory = null; operation.FailureMessage = null; operation.ResultReference = null; operation.Provider = null; operation.Model = null; operation.ProgressStage = null; operation.ProgressPercent = null; var existingNotification = await db.UserNotifications.FirstOrDefaultAsync(item => item.OperationId == operationId, cancellationToken); if (existingNotification is not null) db.UserNotifications.Remove(existingNotification); await db.SaveChangesAsync(cancellationToken); if (transaction is not null) await transaction.CommitAsync(cancellationToken); return true; } private async Task RecoverExpiredLeasesAsync(DateTime now, CancellationToken cancellationToken) { var cancelled = await db.UserOperations.IgnoreQueryFilters().AsNoTracking() .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now && operation.CancellationRequestedAtUtc != null) .ToListAsync(cancellationToken); foreach (var operation in cancelled) await FinalizeRecoveredAsync(operation, RecoveryTerminal.CancelledLease, OperationStatuses.Cancelled, "cancelled", "The operation was cancelled.", now, cancellationToken); var exhausted = await db.UserOperations.IgnoreQueryFilters().AsNoTracking() .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now && operation.CancellationRequestedAtUtc == null && (operation.AttemptCount >= operation.MaxAttempts || operation.DeadlineAtUtc <= now)) .ToListAsync(cancellationToken); foreach (var operation in exhausted) await FinalizeRecoveredAsync(operation, RecoveryTerminal.ExpiredLease, OperationStatuses.Failed, "lease_expired", "The operation could not be recovered after its final worker attempt.", now, cancellationToken); var deadlineExpired = await db.UserOperations.IgnoreQueryFilters().AsNoTracking() .Where(operation => (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry || operation.Status == OperationStatuses.WaitingForExternalFallback) && operation.DeadlineAtUtc <= now) .ToListAsync(cancellationToken); foreach (var operation in deadlineExpired) await FinalizeRecoveredAsync(operation, RecoveryTerminal.QueuedDeadline, OperationStatuses.Failed, "deadline_exceeded", "The operation deadline elapsed before work could complete.", now, cancellationToken); await db.UserOperations.IgnoreQueryFilters() .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now && operation.CancellationRequestedAtUtc == null && operation.AttemptCount < operation.MaxAttempts && (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now)) .ExecuteUpdateAsync(setters => setters .SetProperty(operation => operation.Status, OperationStatuses.WaitingForRetry) .SetProperty(operation => operation.AvailableAtUtc, now) .SetProperty(operation => operation.FailureCategory, "lease_expired") .SetProperty(operation => operation.FailureMessage, "The worker stopped before completing this operation; it will be retried.") .SetProperty(operation => operation.LeaseToken, (string?)null) .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken); } private async Task FinalizeRecoveredAsync(UserOperation operation, RecoveryTerminal reason, string status, string category, string message, DateTime now, CancellationToken cancellationToken) { await using var transaction = await BeginTransactionAsync(cancellationToken); var query = db.UserOperations.IgnoreQueryFilters().Where(item => item.Id == operation.Id); query = reason switch { RecoveryTerminal.CancelledLease => query.Where(item => item.Status == OperationStatuses.Running && item.LeaseExpiresAtUtc <= now && item.CancellationRequestedAtUtc != null), RecoveryTerminal.ExpiredLease => query.Where(item => item.Status == OperationStatuses.Running && item.LeaseExpiresAtUtc <= now && item.CancellationRequestedAtUtc == null && (item.AttemptCount >= item.MaxAttempts || item.DeadlineAtUtc <= now)), _ => query.Where(item => (item.Status == OperationStatuses.Queued || item.Status == OperationStatuses.WaitingForRetry || item.Status == OperationStatuses.WaitingForExternalFallback) && item.DeadlineAtUtc <= now), }; var affected = await query.ExecuteUpdateAsync(setters => setters .SetProperty(item => item.Status, status) .SetProperty(item => item.CompletedAtUtc, now) .SetProperty(item => item.FailureCategory, category) .SetProperty(item => item.FailureMessage, message) .SetProperty(item => item.LeaseToken, (string?)null) .SetProperty(item => item.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken); if (affected == 1) { db.UserNotifications.Add(CreateTerminalNotification(operation, status, now)); await db.SaveChangesAsync(cancellationToken); if (transaction is not null) await transaction.CommitAsync(cancellationToken); } } private async Task BeginTransactionAsync(CancellationToken cancellationToken) { if (!db.Database.IsRelational()) return null; return await db.Database.BeginTransactionAsync(cancellationToken); } private static UserNotification CreateTerminalNotification(UserOperation operation, string status, DateTime now) { var (kind, title, message) = status switch { OperationStatuses.Succeeded => ("operation_succeeded", "Operation completed", "Your background operation completed."), OperationStatuses.Cancelled => ("operation_cancelled", "Operation cancelled", "Your background operation was cancelled."), _ => ("operation_failed", "Operation failed", "A background operation failed. Review it for details."), }; return new UserNotification { Id = Guid.NewGuid(), OwnerUserId = operation.OwnerUserId, OperationId = operation.Id, Kind = kind, Title = title, Message = message, CreatedAtUtc = now, }; } private enum RecoveryTerminal { CancelledLease, ExpiredLease, QueuedDeadline } private static void Validate(CreateUserOperation request) { ValidateRequired(request.TaskType, 64, nameof(request.TaskType)); ValidateRequired(request.IdempotencyKey, 128, nameof(request.IdempotencyKey)); ValidateRequired(request.EntitlementDecision, 32, nameof(request.EntitlementDecision)); ValidateRequired(request.PrivacyPolicy, 32, nameof(request.PrivacyPolicy)); ValidateOptional(request.SubjectType, 64, nameof(request.SubjectType)); ValidateOptional(request.SubjectId, 128, nameof(request.SubjectId)); if (request.MaxAttempts is < 1 or > 10) throw new ArgumentOutOfRangeException(nameof(request.MaxAttempts)); if (request.DeadlineAtUtc is { Kind: not DateTimeKind.Utc }) throw new ArgumentException("Operation deadlines must be UTC.", nameof(request.DeadlineAtUtc)); } private void EnsureOwnerScope() { if (db.CurrentUserId is null) throw new InvalidOperationException("Operation mutation requires an explicit owner scope."); } private static void ValidateLeaseDuration(TimeSpan leaseDuration) { if (leaseDuration < TimeSpan.FromSeconds(5) || leaseDuration > TimeSpan.FromMinutes(30)) throw new ArgumentOutOfRangeException(nameof(leaseDuration)); } private static void ValidateRequired(string value, int maxLength, string name) { if (string.IsNullOrWhiteSpace(value) || value.Length > maxLength) throw new ArgumentException($"{name} is required and limited to {maxLength} characters.", name); } private static void ValidateOptional(string? value, int maxLength, string name) { if (value?.Length > maxLength) throw new ArgumentException($"{name} is limited to {maxLength} characters.", name); } }