ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.IntegrationTests/GracefulStopWithOpenSyncStreamTests.cs

using System.Diagnostics;
using System.Net;
using Akka.Actor;
using Akka.Cluster;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Hosting.Server;
using Microsoft.AspNetCore.Hosting.Server.Features;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.AspNetCore.Server.Kestrel.Core;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Xunit.Abstractions;
using ZB.MOM.WW.LocalDb;
using ZB.MOM.WW.LocalDb.Replication;
using ZB.MOM.WW.ScadaBridge.ClusterInfrastructure;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Host;
using ZB.MOM.WW.ScadaBridge.Host.Actors;
using AkkaCluster = Akka.Cluster.Cluster;

namespace ZB.MOM.WW.ScadaBridge.IntegrationTests;

/// <summary>
/// ScadaBridge#140 (the ScadaBridge half of lmxopcua#906): a graceful stop of a site node whose
/// peer holds an open LocalDb sync stream must leave the Akka cluster FIRST and finish promptly,
/// not wait out Kestrel's graceful-drain timeout.
/// </summary>
/// <remarks>
/// <para>
/// The stopping node is a real <see cref="WebApplication"/> — the hosting model whose hosted-service
/// order caused the bug (the web server stops before the Akka service) — with a real Kestrel h2c
/// listener, the real <c>MapZbLocalDbSync</c> behind the real <see cref="LocalDbSyncAuthInterceptor"/>,
/// the production <see cref="AkkaHostedService"/> forming a real single-member Akka cluster, and the
/// production <see cref="ShutdownStreamCancellationInterceptor"/>. The peer is a real LocalDb
/// <c>SyncBackgroundService</c> dialling it over loopback and holding its stream open.
/// </para>
/// <para>
/// The web server's <see cref="IServer"/> is wrapped to timestamp the moment its <c>StopAsync</c> (the
/// drain) begins, so "left the cluster before the drain" is asserted as an ORDER, not inferred from
/// a total duration.
/// </para>
/// <para>
/// The second test is the positive control: the same stack as it was before the fix — the Akka service
/// seen by the host as a plain <see cref="IHostedService"/> (so its shutdown runs only in
/// <c>StopAsync</c>) and no stream interceptor — reproduces the bug: the stop waits out the whole
/// shutdown timeout and the cluster is left only after the drain. Without it, a fast stop in the first
/// test could be a harness that never held a stream open at all.
/// </para>
/// </remarks>
[Collection("LocalDbSitePairConvergence")]
public sealed class GracefulStopWithOpenSyncStreamTests
{
    private const string ApiKey = "scadabridge-140-graceful-stop-key";

    /// <summary>The host's ShutdownTimeout in these tests — Kestrel's drain budget (30 s in production).</summary>
    private static readonly TimeSpan ShutdownTimeout = TimeSpan.FromSeconds(8);

    private readonly ITestOutputHelper _output;

    static GracefulStopWithOpenSyncStreamTests() =>
        AppContext.SetSwitch("System.Net.Http.SocketsHttpHandler.Http2UnencryptedSupport", true);

    /// <summary>Creates the test class.</summary>
    /// <param name="output">xunit output, for the measured timings.</param>
    public GracefulStopWithOpenSyncStreamTests(ITestOutputHelper output) => _output = output;

    /// <summary>With the fix: the node leaves the cluster before the drain begins and stops well inside the budget.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task GracefulStop_WithOpenSyncStream_LeavesClusterBeforeDrain_AndFinishesPromptly()
    {
        var result = await RunStopScenarioAsync(withFix: true);

        // The peer's stream was open and was ended by the stop (not by a peer that never connected).
        Assert.True(result.StreamWasOpen, "the peer's sync stream was never open on the stopping node");

        // Left the cluster, and did so before the web server began draining.
        Assert.NotNull(result.MemberRemovedAt);
        Assert.NotNull(result.ServerStopBeganAt);
        Assert.True(result.MemberRemovedAt < result.ServerStopBeganAt,
            $"the node must leave the cluster BEFORE Kestrel drains (removed at {result.MemberRemovedAt?.TotalMilliseconds:0} ms, " +
            $"drain began at {result.ServerStopBeganAt?.TotalMilliseconds:0} ms)");

        // Finished well inside the drain budget: the open sync stream did not hold the stop.
        Assert.True(result.StopElapsed < TimeSpan.FromSeconds(4),
            $"stop took {result.StopElapsed.TotalMilliseconds:0} ms against a {ShutdownTimeout.TotalSeconds:0} s drain budget");

        // A planned stop is not a self-down: no exit code 75.
        Assert.False(result.TerminatedOutsideHostShutdown);
    }

    /// <summary>Positive control: the pre-fix wiring reproduces the late leave and the full-budget stop.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task PositiveControl_WithoutTheFix_StopWaitsOutTheDrain_AndLeavesTheClusterLast()
    {
        var result = await RunStopScenarioAsync(withFix: false);

        Assert.True(result.StreamWasOpen);
        Assert.True(result.StopElapsed >= ShutdownTimeout - TimeSpan.FromSeconds(1),
            $"stop took {result.StopElapsed.TotalMilliseconds:0} ms; expected it to wait out the {ShutdownTimeout.TotalSeconds:0} s drain");
        Assert.NotNull(result.MemberRemovedAt);
        Assert.True(result.MemberRemovedAt > result.ServerStopBeganAt);
        Assert.True(result.MemberRemovedAt >= ShutdownTimeout - TimeSpan.FromSeconds(1));
        Assert.False(result.TerminatedOutsideHostShutdown);
    }

    private sealed record StopResult(
        TimeSpan StopElapsed,
        TimeSpan? MemberRemovedAt,
        TimeSpan? ServerStopBeganAt,
        bool StreamWasOpen,
        bool TerminatedOutsideHostShutdown);

    private async Task<StopResult> RunStopScenarioAsync(bool withFix)
    {
        var pathA = Path.Combine(Path.GetTempPath(), $"sb140-A-{Guid.NewGuid():N}.db");
        var pathB = Path.Combine(Path.GetTempPath(), $"sb140-B-{Guid.NewGuid():N}.db");
        await using var dbA = BuildDatabaseProvider(pathA, "node-a");
        await using var dbB = BuildDatabaseProvider(pathB, "node-b");
        var clock = new Stopwatch();
        var probe = new ServerStopProbe(clock);
        IHost? initiator = null;
        WebApplication? passive = null;

        try
        {
            passive = BuildStoppingNode(dbB.GetRequiredService<ILocalDb>(), probe, withFix);
            await passive.StartAsync();

            var akka = passive.Services.GetRequiredService<AkkaHostedService>();
            var cluster = AkkaCluster.Get(akka.ActorSystem!);
            await WaitUntilAsync(() => cluster.State.Members.Any(m => m.Status == MemberStatus.Up),
                "the node to come Up as its own seed");

            TimeSpan? memberRemovedAt = null;
            cluster.RegisterOnMemberRemoved(() => memberRemovedAt ??= clock.Elapsed);

            var address = passive.Services.GetRequiredService<IServer>()
                .Features.Get<IServerAddressesFeature>()!.Addresses.Single();
            initiator = await StartInitiatorAsync(dbA.GetRequiredService<ILocalDb>(), address);

            // The peer's stream is open and holding a session on the stopping node.
            var status = passive.Services.GetRequiredService<ISyncStatus>();
            await WaitUntilAsync(() => status.Connected, "the peer's sync stream to be open on the stopping node");
            var streamWasOpen = status.Connected;

            clock.Start();
            await passive.StopAsync(CancellationToken.None);
            var elapsed = clock.Elapsed;

            _output.WriteLine(
                $"withFix={withFix}: stop {elapsed.TotalMilliseconds:0} ms; member removed at " +
                $"{memberRemovedAt?.TotalMilliseconds:0} ms; web-server drain began at " +
                $"{probe.BeganAt?.TotalMilliseconds:0} ms (drain budget {ShutdownTimeout.TotalSeconds:0} s)");
            return new StopResult(elapsed, memberRemovedAt, probe.BeganAt, streamWasOpen,
                akka.TerminatedOutsideHostShutdown);
        }
        finally
        {
            if (initiator is not null)
            {
                try { await initiator.StopAsync(TimeSpan.FromSeconds(5)); } catch { /* teardown */ }
                initiator.Dispose();
            }

            if (passive is not null)
                await passive.DisposeAsync();

            Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
            foreach (var path in new[] { pathA, pathB })
                foreach (var suffix in new[] { "", "-wal", "-shm" })
                    try { File.Delete(path + suffix); } catch { /* best effort */ }
        }
    }

    private static WebApplication BuildStoppingNode(ILocalDb db, ServerStopProbe probe, bool withFix)
    {
        var builder = WebApplication.CreateBuilder();
        builder.Environment.ApplicationName = typeof(GracefulStopWithOpenSyncStreamTests).Assembly.GetName().Name!;
        builder.Logging.ClearProviders();
        builder.WebHost.ConfigureKestrel(k =>
            k.Listen(IPAddress.Loopback, 0, o => o.Protocols = HttpProtocols.Http2));
        builder.Services.Configure<HostOptions>(o => o.ShutdownTimeout = ShutdownTimeout);

        // The production site interceptor chain: the real auth gate, then (with the fix) the
        // stream-cancellation interceptor innermost, exactly as Program.cs orders them.
        builder.Services.AddGrpc(o =>
        {
            o.Interceptors.Add<LocalDbSyncAuthInterceptor>();
            if (withFix)
                o.Interceptors.Add<ShutdownStreamCancellationInterceptor>();
        });
        builder.Services.AddSingleton(db);
        builder.Services.AddZbLocalDbReplication(ReplicationConfig(peerAddress: null));

        // The production AkkaHostedService, registered the way Program/SiteServiceRegistration do.
        // Role "StopProbe" (neither Central nor Site) makes StartAsync register only the dead-letter
        // monitor; the cluster formation, CoordinatedShutdown and the WhenTerminated watchdog are
        // the production code paths.
        var port = FreePort();
        builder.Services.AddSingleton(sp => new AkkaHostedService(
            sp,
            Options.Create(new NodeOptions { Role = "StopProbe", NodeHostname = "127.0.0.1", RemotingPort = port }),
            Options.Create(new ClusterOptions
            {
                SeedNodes = new List<string> { $"akka.tcp://scadabridge@127.0.0.1:{port}" },
                AllowSingleNodeCluster = true,
                MinNrOfMembers = 1,
            }),
            Options.Create(new CommunicationOptions()),
            sp.GetRequiredService<ILogger<AkkaHostedService>>(),
            sp.GetRequiredService<IHostApplicationLifetime>(),
            sp.GetRequiredService<IHostLifetime>()));
        if (withFix)
            builder.Services.AddHostedService(sp => sp.GetRequiredService<AkkaHostedService>());
        else
            builder.Services.AddHostedService(sp => new PreFixHostedService(sp.GetRequiredService<AkkaHostedService>()));

        // Wrap the web server so the start of its drain is timestamped.
        var serverDescriptor = builder.Services.Last(d => d.ServiceType == typeof(IServer));
        builder.Services.Remove(serverDescriptor);
        builder.Services.AddSingleton<IServer>(sp =>
        {
            var inner = serverDescriptor.ImplementationFactory is { } factory
                ? (IServer)factory(sp)
                : (IServer)ActivatorUtilities.CreateInstance(sp, serverDescriptor.ImplementationType!);
            return probe.Wrap(inner);
        });

        var app = builder.Build();
        app.MapZbLocalDbSync();
        return app;
    }

    private static async Task<IHost> StartInitiatorAsync(ILocalDb db, string peerAddress)
    {
        var config = ReplicationConfig(peerAddress);
        var host = new HostBuilder()
            .ConfigureServices(services =>
            {
                services.AddLogging();
                services.AddSingleton(db);
                services.AddZbLocalDbReplication(config);
            })
            .Build();
        await host.StartAsync();
        return host;
    }

    private static IConfiguration ReplicationConfig(string? peerAddress)
    {
        var values = new Dictionary<string, string?>
        {
            ["LocalDb:Replication:FlushInterval"] = "00:00:00.050",
            ["LocalDb:Replication:ReconnectBackoffMax"] = "00:00:02",
            ["LocalDb:Replication:ApiKey"] = ApiKey,
        };
        if (peerAddress is not null)
            values["LocalDb:Replication:PeerAddress"] = peerAddress;
        return new ConfigurationBuilder().AddInMemoryCollection(values).Build();
    }

    private static ServiceProvider BuildDatabaseProvider(string path, string nodeName)
    {
        string Absent() => Path.Combine(Path.GetTempPath(), $"absent-{Guid.NewGuid():N}.db");
        var config = new ConfigurationBuilder()
            .AddInMemoryCollection(new Dictionary<string, string?>
            {
                ["LocalDb:Path"] = path,
                ["ScadaBridge:Node:NodeName"] = nodeName,
                ["LocalDb:Replication:ApiKey"] = ApiKey,
                // Legacy migrators pointed at paths that do not exist, so they no-op.
                ["ScadaBridge:OperationTracking:ConnectionString"] = $"Data Source={Absent()}",
                ["ScadaBridge:SiteEventLog:DatabasePath"] = Absent(),
                ["ScadaBridge:StoreAndForward:SqliteDbPath"] = Absent(),
                ["ScadaBridge:Database:SiteDbPath"] = Absent(),
            })
            .Build();

        return new ServiceCollection()
            .AddZbLocalDb(config, db => SiteLocalDbSetup.OnReady(db, config))
            .BuildServiceProvider();
    }

    private static int FreePort()
    {
        var l = new System.Net.Sockets.TcpListener(IPAddress.Loopback, 0);
        l.Start();
        var p = ((IPEndPoint)l.LocalEndpoint).Port;
        l.Stop();
        return p;
    }

    private static async Task WaitUntilAsync(Func<bool> condition, string because)
    {
        var deadline = DateTime.UtcNow + TimeSpan.FromSeconds(30);
        while (DateTime.UtcNow < deadline)
        {
            if (condition())
                return;
            await Task.Delay(50);
        }

        Assert.Fail($"Timed out waiting for: {because}");
    }

    /// <summary>
    /// The pre-fix registration shape: the host sees the Akka service only as an
    /// <see cref="IHostedService"/>, so its shutdown runs in <c>StopAsync</c>, after Kestrel's.
    /// </summary>
    private sealed class PreFixHostedService(AkkaHostedService inner) : IHostedService
    {
        public Task StartAsync(CancellationToken cancellationToken) => inner.StartAsync(cancellationToken);

        public Task StopAsync(CancellationToken cancellationToken) => inner.StopAsync(cancellationToken);
    }

    /// <summary>Records when the wrapped web server's <c>StopAsync</c> (its graceful drain) begins.</summary>
    private sealed class ServerStopProbe(Stopwatch clock)
    {
        public TimeSpan? BeganAt { get; private set; }

        private void MarkBegan() => BeganAt ??= clock.Elapsed;

        public IServer Wrap(IServer inner) => new RecordingServer(inner, this);

        private sealed class RecordingServer(IServer inner, ServerStopProbe probe) : IServer
        {
            public IFeatureCollection Features => inner.Features;

            public Task StartAsync<TContext>(IHttpApplication<TContext> application, CancellationToken cancellationToken)
                where TContext : notnull
                => inner.StartAsync(application, cancellationToken);

            public Task StopAsync(CancellationToken cancellationToken)
            {
                probe.MarkBegan();
                return inner.StopAsync(cancellationToken);
            }

            public void Dispose() => inner.Dispose();
        }
    }
}