tiku-backend.net/Tiku.Infrastructure/Jobs/Processor/BackgroundJobLeaseManager.cs

using Microsoft.Extensions.Logging;
using Npgsql;
using Tiku.Infrastructure.Observability;

namespace Tiku.Infrastructure.Jobs;

internal sealed record BackgroundJobLeaseTiming(
    TimeSpan LeaseDuration,
    TimeSpan RenewalInterval,
    TimeSpan RetryInterval)
{
    internal static readonly BackgroundJobLeaseTiming Default = new(
        TimeSpan.FromMinutes(5),
        TimeSpan.FromMinutes(1),
        TimeSpan.FromSeconds(10));
}

internal sealed class BackgroundJobLeaseManager(
    NpgsqlDataSource dataSource,
    BackgroundJobLeaseTiming timing,
    TimeProvider timeProvider,
    ILogger<BackgroundJobLeaseManager> logger)
{
    internal BackgroundJobLease Start(
        Guid jobId,
        string workerId,
        DateTimeOffset leaseExpiresAt,
        CancellationToken executionCancellationToken)
    {
        return new BackgroundJobLease(
            jobId,
            workerId,
            leaseExpiresAt,
            dataSource,
            timing,
            timeProvider,
            logger,
            executionCancellationToken);
    }
}

internal sealed class BackgroundJobLease : IAsyncDisposable
{
    private readonly Guid jobId;
    private readonly string workerId;
    private readonly NpgsqlDataSource dataSource;
    private readonly BackgroundJobLeaseTiming timing;
    private readonly TimeProvider timeProvider;
    private readonly ILogger logger;
    private readonly CancellationTokenSource executionCancellation;
    private readonly CancellationTokenSource renewalCancellation = new();
    private readonly Task renewalTask;
    private DateTimeOffset confirmedExpiresAt;
    private int ownershipLost;
    private int stopped;

    internal BackgroundJobLease(
        Guid jobId,
        string workerId,
        DateTimeOffset leaseExpiresAt,
        NpgsqlDataSource dataSource,
        BackgroundJobLeaseTiming timing,
        TimeProvider timeProvider,
        ILogger logger,
        CancellationToken executionCancellationToken)
    {
        this.jobId = jobId;
        this.workerId = workerId;
        this.dataSource = dataSource;
        this.timing = timing;
        this.timeProvider = timeProvider;
        this.logger = logger;
        confirmedExpiresAt = leaseExpiresAt;
        executionCancellation = CancellationTokenSource.CreateLinkedTokenSource(executionCancellationToken);
        renewalTask = RenewAsync();
    }

    internal CancellationToken ExecutionToken => executionCancellation.Token;
    internal bool OwnershipLost => Volatile.Read(ref ownershipLost) != 0;

    internal async Task StopAsync()
    {
        if (Interlocked.Exchange(ref stopped, 1) != 0) return;

        await renewalCancellation.CancelAsync();
        try
        {
            await renewalTask;
        }
        catch (OperationCanceledException) when (renewalCancellation.IsCancellationRequested)
        {
        }
    }

    public async ValueTask DisposeAsync()
    {
        await StopAsync();
        renewalCancellation.Dispose();
        executionCancellation.Dispose();
    }

    private async Task RenewAsync()
    {
        var delay = timing.RenewalInterval;
        while (!renewalCancellation.IsCancellationRequested)
        {
            await Task.Delay(delay, timeProvider, renewalCancellation.Token);
            try
            {
                var now = timeProvider.GetUtcNow();
                var nextExpiresAt = now.Add(timing.LeaseDuration);
                await using var connection = await dataSource.OpenConnectionAsync(renewalCancellation.Token);
                await using var command = connection.CreateCommand();
                command.CommandText = """
                                      UPDATE background_jobs
                                      SET lock_expires_at = @lock_expires_at,
                                          updated_at = @updated_at
                                      WHERE id = @id
                                        AND locked_by = @worker_id
                                        AND status = 'processing'
                                      """;
                command.Parameters.AddWithValue("lock_expires_at", nextExpiresAt);
                command.Parameters.AddWithValue("updated_at", now);
                command.Parameters.AddWithValue("id", jobId);
                command.Parameters.AddWithValue("worker_id", workerId);
                var renewed = await command.ExecuteNonQueryAsync(renewalCancellation.Token);
                if (renewed == 0)
                {
                    LoseOwnership("The job is no longer owned by this worker.");
                    return;
                }

                confirmedExpiresAt = nextExpiresAt;
                delay = timing.RenewalInterval;
                WorkerTelemetry.RecordJobLeaseRenewal("succeeded");
            }
            catch (OperationCanceledException) when (renewalCancellation.IsCancellationRequested)
            {
                return;
            }
            catch (Exception exception)
            {
                WorkerTelemetry.RecordJobLeaseRenewal("failed");
                logger.LogWarning(exception, "Failed to renew background job {JobId} lease for {WorkerId}.", jobId,
                    workerId);
                if (timeProvider.GetUtcNow() >= confirmedExpiresAt)
                {
                    LoseOwnership("The last confirmed lease expired before it could be renewed.");
                    return;
                }

                delay = timing.RetryInterval;
            }
        }
    }

    private void LoseOwnership(string reason)
    {
        if (Interlocked.Exchange(ref ownershipLost, 1) != 0) return;

        WorkerTelemetry.RecordJobLeaseRenewal("lost");
        logger.LogError("Background job {JobId} lease was lost by {WorkerId}. {Reason}", jobId, workerId, reason);
        executionCancellation.Cancel();
    }
}