Rendezvous/src/FinalFactory.Rendezvous.Client/SessionLeaseMaintainer.cs

using FinalFactory.Rendezvous.Contracts;

namespace FinalFactory.Rendezvous.Client;

public enum LeaseMaintenanceStopReason
{
    Cancelled = 1,
    Disposed = 2,
    LeaseLost = 3,
    Failed = 4,
}

public sealed class LeaseMaintenanceResult
{
    internal LeaseMaintenanceResult(LeaseMaintenanceStopReason reason, RendezvousErrorCode error)
    {
        Reason = reason;
        Error = error;
    }

    public LeaseMaintenanceStopReason Reason { get; }
    public RendezvousErrorCode Error { get; }
}

public sealed class SessionLeaseMaintainer : IAsyncDisposable
{
    private readonly object _gate = new();
    private readonly IRendezvousPublisherClient _publisher;
    private readonly PublishedSession _session;
    private readonly string _publisherCredential;
    private readonly IRendezvousDelay _delay;
    private readonly CancellationTokenSource _disposeCancellation = new();
    private Task<LeaseMaintenanceResult>? _activeRun;
    private Task? _disposeTask;
    private bool _disposed;

    internal SessionLeaseMaintainer(
        IRendezvousPublisherClient publisher,
        PublishedSession session,
        string publisherCredential,
        IRendezvousDelay? delay = null)
    {
        _publisher = publisher;
        _session = session;
        _publisherCredential = publisherCredential;
        _delay = delay ?? new SystemRendezvousDelay();
    }

    public event EventHandler? LeaseLost;

    public Task<LeaseMaintenanceResult> RunAsync(CancellationToken cancellationToken = default)
    {
        lock (_gate)
        {
            if (_disposed)
            {
                throw new ObjectDisposedException(nameof(SessionLeaseMaintainer));
            }
            if (_activeRun is not null)
            {
                throw new InvalidOperationException("Lease maintenance is already running.");
            }

            _activeRun = RunCoreAsync(cancellationToken);
            return _activeRun;
        }
    }

    public ValueTask DisposeAsync()
    {
        lock (_gate)
        {
            if (_disposeTask is not null)
            {
                return new(_disposeTask);
            }

            _disposed = true;
            _disposeCancellation.Cancel();
            _disposeTask = FinishDisposeAsync(_activeRun);
            return new(_disposeTask);
        }
    }

    private async Task FinishDisposeAsync(Task<LeaseMaintenanceResult>? active)
    {
        try
        {
            if (active is not null)
            {
                await active.ConfigureAwait(false);
            }
        }
        finally
        {
            _disposeCancellation.Dispose();
        }
    }

    private async Task<LeaseMaintenanceResult> RunCoreAsync(CancellationToken cancellationToken)
    {
        await Task.Yield();
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(
            cancellationToken,
            _disposeCancellation.Token);
        try
        {
            while (true)
            {
                await _delay.DelayAsync(
                    TimeSpan.FromSeconds(Math.Max(1, _session.LeaseRenewAfterSeconds)),
                    linked.Token).ConfigureAwait(false);
                RendezvousClientResult<RenewLeaseResponse> renewed = await _publisher.RenewAsync(
                    _session,
                    _publisherCredential,
                    linked.Token).ConfigureAwait(false);
                if (renewed.IsSuccess)
                {
                    continue;
                }

                if (renewed.Error is RendezvousErrorCode.NotFound
                    or RendezvousErrorCode.Expired
                    or RendezvousErrorCode.Forbidden
                    or RendezvousErrorCode.AuthenticationRequired)
                {
                    LeaseLost?.Invoke(this, EventArgs.Empty);
                    return new(LeaseMaintenanceStopReason.LeaseLost, renewed.Error);
                }

                return new(LeaseMaintenanceStopReason.Failed, renewed.Error);
            }
        }
        catch (OperationCanceledException) when (linked.IsCancellationRequested)
        {
            return new(
                _disposeCancellation.IsCancellationRequested
                    ? LeaseMaintenanceStopReason.Disposed
                    : LeaseMaintenanceStopReason.Cancelled,
                RendezvousErrorCode.None);
        }
        finally
        {
            lock (_gate)
            {
                _activeRun = null;
            }
        }
    }
}