fix(worker): enable safe multi-instance job processing

Tiku.Infrastructure/Jobs/Processor/BackgroundJobService.Processor.cs

@@ -17,6 +17,30 @@ internal sealed partial class BackgroundJobService
         int batchSize,
         bool includeImmediateJobs = true,
         CancellationToken cancellationToken = default)
+    {
+        var processed = 0;
+        var limit = Math.Clamp(batchSize, 1, 100);
+        for (var index = 0; index < limit; index++)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            var jobId = await TryClaimNextAsync(workerId, includeImmediateJobs, cancellationToken);
+            if (!jobId.HasValue) break;
+
+            unitOfWork.ChangeTracker.Clear();
+            var job = await jobsOperationsPersistence.BackgroundJobs.SingleAsync(
+                value => value.Id == jobId.Value,
+                cancellationToken);
+            if (await ProcessJobAsync(job, workerId, true, cancellationToken)) processed++;
+            unitOfWork.ChangeTracker.Clear();
+        }
+
+        return processed;
+    }
+
+    private async Task<Guid?> TryClaimNextAsync(
+        string workerId,
+        bool includeImmediateJobs,
+        CancellationToken cancellationToken)
     {
         var now = DateTimeOffset.UtcNow;
         var leaseExpiresAt = now.Add(LeaseDuration);
@@ -27,7 +51,7 @@ internal sealed partial class BackgroundJobService
                                                                       lock_expires_at = {leaseExpiresAt},
                                                                       started_at = COALESCE(started_at, {now}),
                                                                       updated_at = {now}
-                                                                  WHERE job.id IN (
+                                                                  WHERE job.id = (
                                                                       SELECT candidate.id
                                                                       FROM background_jobs AS candidate
                                                                       WHERE (
@@ -37,23 +61,12 @@ internal sealed partial class BackgroundJobService
                                                                       )
                                                                       ORDER BY candidate.created_at, candidate.id
                                                                       FOR UPDATE SKIP LOCKED
-                                                                      LIMIT {Math.Clamp(batchSize, 1, 100)}
+                                                                      LIMIT 1
                                                                   )
                                                                   RETURNING job.id AS "Value"
                                                                   """)
             .ToArrayAsync(cancellationToken);
-
-        var processed = 0;
-        unitOfWork.ChangeTracker.Clear();
-        foreach (var jobId in claimedIds)
-        {
-            cancellationToken.ThrowIfCancellationRequested();
-            var job = await jobsOperationsPersistence.BackgroundJobs.SingleAsync(value => value.Id == jobId, cancellationToken);
-            if (await ProcessJobAsync(job, workerId, true, cancellationToken)) processed++;
-            unitOfWork.ChangeTracker.Clear();
-        }
-
-        return processed;
+        return claimedIds.Length == 0 ? null : claimedIds[0];
     }
 
     public async Task<bool> ProcessRequestedAsync(
@@ -97,17 +110,18 @@ internal sealed partial class BackgroundJobService
             (alreadyClaimed && (job.Status != BackgroundJobStatus.Processing || job.LockedBy != workerId)))
             return false;
         await unitOfWork.Entry(job).ReloadAsync(cancellationToken);
+        if (alreadyClaimed && (job.Status != BackgroundJobStatus.Processing || job.LockedBy != workerId))
+            return false;
         if (job.CancellationRequestedAt.HasValue)
         {
             job.CompletedAt = DateTimeOffset.UtcNow;
-            await CompleteAsync(
+            return await CompleteAsync(
                 job,
                 workerId,
                 BackgroundJobStatus.Cancelled,
                 job.Result,
                 null,
-                cancellationToken);
-            return true;
+                cancellationToken) > 0;
         }
 
         if (job.JobType is not ("asset_security_scan" or "tenant_export") && !(await featureAccessService.EvaluateAsync(
@@ -116,14 +130,13 @@ internal sealed partial class BackgroundJobService
                 FeatureAccessOperation.Write,
                 cancellationToken)).Allowed)
         {
-            await CompleteAsync(
+            return await CompleteAsync(
                 job,
                 workerId,
                 BackgroundJobStatus.Failed,
                 JsonDefaults.Object(),
                 "Tenant feature entitlement was revoked before job execution.",
-                cancellationToken);
-            return true;
+                cancellationToken) > 0;
         }
 
         if (!alreadyClaimed)
@@ -136,6 +149,13 @@ internal sealed partial class BackgroundJobService
             await unitOfWork.SaveChangesAsync(cancellationToken);
         }
 
+        if (!job.LockExpiresAt.HasValue) return false;
+
+        await using var lease = leaseManager.Start(
+            job.Id,
+            workerId,
+            job.LockExpiresAt.Value,
+            cancellationToken);
         try
         {
             var handlerResult = await tenantExecutionScope.ExecuteAsync(
@@ -150,7 +170,13 @@ internal sealed partial class BackgroundJobService
                         new BackgroundJobExecutionContext(job.Id, job.TenantId, job.JobType, job.Payload),
                         token);
                 },
-                cancellationToken);
+                lease.ExecutionToken);
+            if (lease.OwnershipLost)
+            {
+                RecordLeaseLost(job, startedTimestamp);
+                return false;
+            }
+
             var cancellationRequested = await jobsOperationsPersistence.BackgroundJobs.AsNoTracking()
                 .Where(item => item.Id == job.Id)
                 .Select(item => item.CancellationRequestedAt != null)
@@ -161,6 +187,16 @@ internal sealed partial class BackgroundJobService
             job.Result = handlerResult.Result;
             job.OutputAssetId = handlerResult.OutputAssetId;
         }
+        catch (OperationCanceledException) when (lease.OwnershipLost)
+        {
+            RecordLeaseLost(job, startedTimestamp);
+            return false;
+        }
+        catch (Exception) when (lease.OwnershipLost)
+        {
+            RecordLeaseLost(job, startedTimestamp);
+            return false;
+        }
         catch (Exception exception) when (exception is not OperationCanceledException)
         {
             job.RetryCount++;
@@ -188,24 +224,39 @@ internal sealed partial class BackgroundJobService
                             exception,
                             token);
                     },
-                    cancellationToken);
+                    lease.ExecutionToken);
             }
             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);
+            catch (OperationCanceledException) when (lease.OwnershipLost)
+            {
+                RecordLeaseLost(job, startedTimestamp);
+                return false;
+            }
         }
 
+        await lease.StopAsync();
+        if (lease.OwnershipLost)
+        {
+            RecordLeaseLost(job, startedTimestamp);
+            return false;
+        }
+
+        var completed = await CompleteAsync(job, workerId, job.Status, job.Result, job.LastError, cancellationToken);
+        if (completed == 0)
+        {
+            RecordLeaseLost(job, startedTimestamp);
+            return false;
+        }
+
+        WorkerTelemetry.RecordJob(job.JobType, job.Status.ToString(),
+            Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
         return true;
     }
 
-    private async Task CompleteAsync(
+    private async Task<int> CompleteAsync(
         BackgroundJob job,
         string workerId,
         BackgroundJobStatus status,
@@ -213,7 +264,7 @@ internal sealed partial class BackgroundJobService
         string? lastError,
         CancellationToken cancellationToken)
     {
-        await jobsOperationsPersistence.BackgroundJobs
+        return await jobsOperationsPersistence.BackgroundJobs
             .Where(value => value.Id == job.Id && value.LockedBy == workerId)
             .ExecuteUpdateAsync(setters => setters
                     .SetProperty(value => value.Status, status)
@@ -227,4 +278,10 @@ internal sealed partial class BackgroundJobService
                     .SetProperty(value => value.LockExpiresAt, (DateTimeOffset?)null),
                 cancellationToken);
     }
+
+    private static void RecordLeaseLost(BackgroundJob job, long startedTimestamp)
+    {
+        WorkerTelemetry.RecordJob(job.JobType, "LeaseLost",
+            Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
+    }
 }