Rendezvous/tests/FinalFactory.Rendezvous.Tests/Deployment/GracefulDrainServiceTests.cs

using System.Diagnostics;
using FinalFactory.Rendezvous.Server.Deployment;
using FinalFactory.Rendezvous.Server.State;
using FinalFactory.Rendezvous.Tests.State;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;

namespace FinalFactory.Rendezvous.Tests.Deployment;

public sealed class GracefulDrainServiceTests
{
    [Fact]
    public async Task ShutdownRejectsNewWorkAndClearsStateAfterBoundedAttemptDeadline()
    {
        EphemeralStoreOptions stateOptions = new()
        {
            GracefulDrainLifetime = TimeSpan.FromSeconds(1),
        };
        EphemeralStateFixture fixture = new(stateOptions);
        StoredListing listing = fixture.CreateVisibleListing(out _);
        StoreResult<StoredJoinAttempt> attempt = fixture.Store.CreateJoinAttempt(
            fixture.AttemptCommand(listing));
        Assert.True(attempt.Succeeded);
        using FakeApplicationLifetime lifetime = new();
        using GracefulDrainService service = CreateService(fixture.Store, lifetime);
        await service.StartAsync(CancellationToken.None);

        Task stopping = Task.Run(lifetime.StopApplication);
        await WaitUntilAsync(() => fixture.Store.IsDraining, TimeSpan.FromSeconds(1));
        StoreResult<StoredListing> rejected = fixture.Store.CreateListing(fixture.ListingCommand());
        long sweepsAfterAdmissionCheck = fixture.Store.MaintenanceSweepCount;
        Stopwatch elapsed = Stopwatch.StartNew();
        await stopping;
        await service.StopAsync(CancellationToken.None);

        Assert.Equal(StoreResultCode.Draining, rejected.Code);
        Assert.Equal(sweepsAfterAdmissionCheck, fixture.Store.MaintenanceSweepCount);
        Assert.InRange(elapsed.Elapsed, TimeSpan.FromMilliseconds(850), TimeSpan.FromSeconds(2));
        Assert.False(fixture.Store.IsAvailable);
        Assert.Equal(0, fixture.Store.GetSnapshot().ActiveJoinAttempts);
        Assert.Equal(0, fixture.Store.GetSnapshot().ActiveListings);
    }

    [Fact]
    public async Task ShutdownWithoutAttemptsStopsAfterTheConfiguredMinimumOnly()
    {
        EphemeralStateFixture fixture = new(new EphemeralStoreOptions
        {
            GracefulDrainLifetime = TimeSpan.FromSeconds(2),
        });
        fixture.CreateVisibleListing(out _);
        using FakeApplicationLifetime lifetime = new();
        DeploymentOptions options = new()
        {
            DrainDeadlineSeconds = 2,
            MinimumDrainSeconds = 0,
        };
        using GracefulDrainService service = new(
            fixture.Store,
            lifetime,
            Options.Create(options),
            NullLogger<GracefulDrainService>.Instance);
        await service.StartAsync(CancellationToken.None);

        Stopwatch elapsed = Stopwatch.StartNew();
        await service.StopAsync(CancellationToken.None);

        Assert.True(elapsed.Elapsed < TimeSpan.FromMilliseconds(500));
        Assert.False(fixture.Store.IsAvailable);
        Assert.Equal(0, fixture.Store.GetSnapshot().ActiveListings);
    }

    [Fact]
    public async Task ShutdownStopsWhenTheLastAttemptExpiresWithoutAFullStateSweep()
    {
        EphemeralStateFixture fixture = new(new EphemeralStoreOptions
        {
            JoinAttemptLifetime = TimeSpan.FromMilliseconds(100),
            ConnectionTicketLifetime = TimeSpan.FromMilliseconds(50),
            GracefulDrainLifetime = TimeSpan.FromSeconds(2),
        });
        StoredListing listing = fixture.CreateVisibleListing(out _);
        Assert.True(fixture.Store.CreateJoinAttempt(fixture.AttemptCommand(listing)).Succeeded);
        using FakeApplicationLifetime lifetime = new();
        using GracefulDrainService service = new(
            fixture.Store,
            lifetime,
            Options.Create(new DeploymentOptions
            {
                DrainDeadlineSeconds = 2,
                MinimumDrainSeconds = 0,
            }),
            NullLogger<GracefulDrainService>.Instance);
        await service.StartAsync(CancellationToken.None);

        Task stopping = Task.Run(lifetime.StopApplication);
        await WaitUntilAsync(() => fixture.Store.IsDraining, TimeSpan.FromSeconds(1));
        long sweepsBeforeExpiry = fixture.Store.MaintenanceSweepCount;
        fixture.Clock.Advance(TimeSpan.FromMilliseconds(150));
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(1));
        await stopping.WaitAsync(timeout.Token);

        Assert.Equal(sweepsBeforeExpiry, fixture.Store.MaintenanceSweepCount);
        Assert.False(fixture.Store.IsAvailable);
    }

    private static GracefulDrainService CreateService(
        InMemoryEphemeralRendezvousStore store,
        IHostApplicationLifetime lifetime) => new(
            store,
            lifetime,
            Options.Create(new DeploymentOptions
            {
                DrainDeadlineSeconds = 1,
                MinimumDrainSeconds = 0,
            }),
            NullLogger<GracefulDrainService>.Instance);

    private static async Task WaitUntilAsync(Func<bool> predicate, TimeSpan timeout)
    {
        Stopwatch elapsed = Stopwatch.StartNew();
        while (!predicate())
        {
            Assert.True(elapsed.Elapsed < timeout, "The service did not enter drain in time.");
            await Task.Delay(10);
        }
    }

    private sealed class FakeApplicationLifetime : IHostApplicationLifetime, IDisposable
    {
        private readonly CancellationTokenSource _started = new();
        private readonly CancellationTokenSource _stopping = new();
        private readonly CancellationTokenSource _stopped = new();

        public CancellationToken ApplicationStarted => _started.Token;
        public CancellationToken ApplicationStopping => _stopping.Token;
        public CancellationToken ApplicationStopped => _stopped.Token;

        public void StopApplication() => _stopping.Cancel();

        public void Dispose()
        {
            _started.Dispose();
            _stopping.Dispose();
            _stopped.Dispose();
        }
    }
}