using System.Diagnostics; using System.Text.Json; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.DependencyInjection; using Tiku.Application.Jobs; using Tiku.Application.Security; using Tiku.Domain.Common; using Tiku.Domain.Operations; using Tiku.Infrastructure.Observability; namespace Tiku.Infrastructure.Jobs; internal sealed partial class BackgroundJobService { public async Task ProcessPendingAsync( string workerId, int batchSize, bool includeImmediateJobs = true, CancellationToken cancellationToken = default) { var now = DateTimeOffset.UtcNow; var leaseExpiresAt = now.Add(LeaseDuration); var claimedIds = await dbContext.Database.SqlQuery($""" UPDATE background_jobs AS job SET status = 'processing', locked_by = {workerId}, lock_expires_at = {leaseExpiresAt}, started_at = COALESCE(started_at, {now}), updated_at = {now} WHERE job.id IN ( SELECT candidate.id FROM background_jobs AS candidate WHERE ( (candidate.status = 'pending' AND ({includeImmediateJobs} OR candidate.run_after IS NOT NULL) AND (candidate.run_after IS NULL OR candidate.run_after <= {now})) OR (candidate.status = 'processing' AND candidate.lock_expires_at <= {now}) ) ORDER BY candidate.created_at, candidate.id FOR UPDATE SKIP LOCKED LIMIT {Math.Clamp(batchSize, 1, 100)} ) RETURNING job.id AS "Value" """) .ToArrayAsync(cancellationToken); var processed = 0; dbContext.ChangeTracker.Clear(); foreach (var jobId in claimedIds) { cancellationToken.ThrowIfCancellationRequested(); var job = await dbContext.BackgroundJobs.SingleAsync(value => value.Id == jobId, cancellationToken); if (await ProcessJobAsync(job, workerId, true, cancellationToken)) processed++; dbContext.ChangeTracker.Clear(); } return processed; } public async Task ProcessRequestedAsync( Guid jobId, Guid tenantId, string jobType, string workerId, CancellationToken cancellationToken = default) { var normalizedJobType = NormalizeJobType(jobType); var claimed = await dbContext.BackgroundJobs .Where(item => item.Id == jobId && item.TenantId == tenantId && item.JobType == normalizedJobType && item.Status == BackgroundJobStatus.Pending && item.RunAfter == null) .ExecuteUpdateAsync(setters => setters .SetProperty(item => item.Status, BackgroundJobStatus.Processing) .SetProperty(item => item.LockedBy, workerId) .SetProperty(item => item.LockExpiresAt, DateTimeOffset.UtcNow.Add(LeaseDuration)) .SetProperty(item => item.StartedAt, DateTimeOffset.UtcNow), cancellationToken); if (claimed == 0) return false; dbContext.ChangeTracker.Clear(); var job = await dbContext.BackgroundJobs.SingleOrDefaultAsync( item => item.Id == jobId && item.TenantId == tenantId, cancellationToken); if (job is null) throw new InvalidOperationException("The requested background job does not exist in the target tenant."); if (!string.Equals(job.JobType, normalizedJobType, StringComparison.Ordinal)) throw new InvalidOperationException("The requested background job type does not match the persisted job."); return await ProcessJobAsync(job, workerId, true, cancellationToken); } private async Task ProcessJobAsync( BackgroundJob job, string workerId, bool alreadyClaimed, CancellationToken cancellationToken) { var startedTimestamp = Stopwatch.GetTimestamp(); if ((!alreadyClaimed && job.Status != BackgroundJobStatus.Pending) || (alreadyClaimed && (job.Status != BackgroundJobStatus.Processing || job.LockedBy != workerId))) return false; await dbContext.Entry(job).ReloadAsync(cancellationToken); if (job.CancellationRequestedAt.HasValue) { job.CompletedAt = DateTimeOffset.UtcNow; await CompleteAsync( job, workerId, BackgroundJobStatus.Cancelled, job.Result, null, cancellationToken); return true; } if (job.JobType is not ("asset_security_scan" or "tenant_export") && !(await featureAccessService.EvaluateAsync( job.TenantId, ResolveRequiredFeature(job.JobType, job.Payload), FeatureAccessOperation.Write, cancellationToken)).Allowed) { await CompleteAsync( job, workerId, BackgroundJobStatus.Failed, JsonDefaults.Object(), "Tenant feature entitlement was revoked before job execution.", cancellationToken); return true; } if (!alreadyClaimed) { var now = DateTimeOffset.UtcNow; job.Status = BackgroundJobStatus.Processing; job.LockedBy = workerId; job.LockExpiresAt = now.Add(LeaseDuration); job.StartedAt = now; await dbContext.SaveChangesAsync(cancellationToken); } try { var handlerResult = await tenantExecutionScope.ExecuteAsync( new SystemScopeRequest( job.TenantId, SystemScopeCallerType.Worker, workerId, $"Background job {job.JobType}", job.Id.ToString("N")), async (provider, token) => { var registry = provider.GetRequiredService(); var handler = registry.GetRequired(job.JobType); return await handler.HandleAsync( new BackgroundJobExecutionContext(job.Id, job.TenantId, job.JobType, job.Payload), token); }, cancellationToken); var cancellationRequested = await dbContext.BackgroundJobs.AsNoTracking() .Where(item => item.Id == job.Id) .Select(item => item.CancellationRequestedAt != null) .SingleAsync(cancellationToken); job.Status = cancellationRequested ? BackgroundJobStatus.Cancelled : BackgroundJobStatus.Succeeded; job.CompletedAt = DateTimeOffset.UtcNow; job.LastError = null; job.Result = handlerResult.Result; job.OutputAssetId = handlerResult.OutputAssetId; } catch (Exception exception) when (exception is not OperationCanceledException) { job.RetryCount++; job.LastError = exception is BackgroundJobHandlerException handlerException ? $"{handlerException.Code}: {handlerException.Message}" : exception.Message; job.Status = job.RetryCount > job.MaxRetries ? BackgroundJobStatus.Failed : BackgroundJobStatus.Pending; job.RunAfter = job.Status == BackgroundJobStatus.Pending ? DateTimeOffset.UtcNow.AddSeconds(Math.Min(300, 10 * job.RetryCount)) : null; try { await tenantExecutionScope.ExecuteAsync( new SystemScopeRequest( job.TenantId, SystemScopeCallerType.Worker, workerId, $"Background job failure compensation {job.JobType}", job.Id.ToString("N")), async (provider, token) => { var handler = provider.GetRequiredService() .GetRequired(job.JobType); await handler.HandleFailureAsync( new BackgroundJobExecutionContext(job.Id, job.TenantId, job.JobType, job.Payload), exception, token); }, cancellationToken); } catch (Exception compensationException) when (compensationException is not OperationCanceledException) { job.LastError = $"{job.LastError} Compensation failed: {compensationException.Message}"; } } finally { await CompleteAsync(job, workerId, job.Status, job.Result, job.LastError, cancellationToken); WorkerTelemetry.RecordJob(job.JobType, job.Status.ToString(), Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds); } return true; } private async Task CompleteAsync( BackgroundJob job, string workerId, BackgroundJobStatus status, JsonElement result, string? lastError, CancellationToken cancellationToken) { await dbContext.BackgroundJobs .Where(value => value.Id == job.Id && value.LockedBy == workerId) .ExecuteUpdateAsync(setters => setters .SetProperty(value => value.Status, status) .SetProperty(value => value.RetryCount, job.RetryCount) .SetProperty(value => value.RunAfter, job.RunAfter) .SetProperty(value => value.CompletedAt, job.CompletedAt) .SetProperty(value => value.LastError, lastError) .SetProperty(value => value.OutputAssetId, job.OutputAssetId) .SetProperty(value => value.Result, result) .SetProperty(value => value.LockedBy, (string?)null) .SetProperty(value => value.LockExpiresAt, (DateTimeOffset?)null), cancellationToken); } }