Rendezvous/src/FinalFactory.Rendezvous.Server/Deployment/GracefulDrainService.cs

using System.Diagnostics;
using FinalFactory.Rendezvous.Server.State;
using Microsoft.Extensions.Options;

namespace FinalFactory.Rendezvous.Server.Deployment;

internal sealed partial class GracefulDrainService : IHostedService, IDisposable
{
    private static readonly TimeSpan PollInterval = TimeSpan.FromMilliseconds(50);
    private readonly InMemoryEphemeralRendezvousStore _store;
    private readonly IHostApplicationLifetime _lifetime;
    private readonly DeploymentOptions _options;
    private readonly ILogger<GracefulDrainService> _logger;
    private readonly object _gate = new();
    private CancellationTokenRegistration _stoppingRegistration;
    private Task? _drainTask;

    public GracefulDrainService(
        InMemoryEphemeralRendezvousStore store,
        IHostApplicationLifetime lifetime,
        IOptions<DeploymentOptions> options,
        ILogger<GracefulDrainService> logger)
    {
        _store = store;
        _lifetime = lifetime;
        _options = options.Value;
        _logger = logger;
    }

    public Task StartAsync(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        _stoppingRegistration = _lifetime.ApplicationStopping.Register(
            () => EnsureDrainAsync().GetAwaiter().GetResult());
        return Task.CompletedTask;
    }

    public Task StopAsync(CancellationToken cancellationToken)
    {
        // ApplicationStopping callbacks run before hosted services and listeners
        // stop. StopAsync is the idempotent fallback for directly driven hosts.
        _ = cancellationToken;
        return EnsureDrainAsync();
    }

    public void Dispose() => _stoppingRegistration.Dispose();

    private Task EnsureDrainAsync()
    {
        lock (_gate)
        {
            return _drainTask ??= DrainAsync();
        }
    }

    private async Task DrainAsync()
    {
        _store.BeginDrain(CancellationToken.None);
        TimeSpan deadline = TimeSpan.FromSeconds(_options.DrainDeadlineSeconds);
        TimeSpan minimum = TimeSpan.FromSeconds(_options.MinimumDrainSeconds);
        long startedAt = Stopwatch.GetTimestamp();
        LogDrainStarted(_logger, _options.DrainDeadlineSeconds);
        try
        {
            while (Stopwatch.GetElapsedTime(startedAt) < deadline)
            {
                TimeSpan elapsed = Stopwatch.GetElapsedTime(startedAt);
                if (elapsed >= minimum && _store.GetActiveJoinAttemptCountForDrain() == 0)
                {
                    break;
                }

                TimeSpan remaining = deadline - elapsed;
                await Task.Delay(
                    remaining < PollInterval ? remaining : PollInterval,
                    CancellationToken.None).ConfigureAwait(false);
            }
        }
        finally
        {
            _store.MarkUnavailable();
            double elapsedMilliseconds = Stopwatch.GetElapsedTime(startedAt).TotalMilliseconds;
            LogDrainFinished(_logger, elapsedMilliseconds);
        }
    }

    [LoggerMessage(
        EventId = 1,
        Level = LogLevel.Information,
        Message = "Graceful drain started with a {DrainDeadlineSeconds}-second deadline")]
    private static partial void LogDrainStarted(ILogger logger, int drainDeadlineSeconds);

    [LoggerMessage(
        EventId = 2,
        Level = LogLevel.Information,
        Message = "Graceful drain finished after {ElapsedMilliseconds:F0} ms; ephemeral state was cleared")]
    private static partial void LogDrainFinished(ILogger logger, double elapsedMilliseconds);
}