lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.Host.IntegrationTests/LocalDb/GracefulStopWithOpenSyncStreamTests.cs

using System.Diagnostics;
using System.Net;
using Akka.Actor;
using Akka.Cluster;
using Akka.Configuration;
using Akka.Hosting;
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.DependencyInjection.Extensions;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Shouldly;
using Xunit;
using ZB.MOM.WW.LocalDb;
using ZB.MOM.WW.LocalDb.Replication;
using ZB.MOM.WW.OtOpcUa.Cluster;
using ZB.MOM.WW.OtOpcUa.Host.Configuration;
using AkkaCluster = Akka.Cluster.Cluster;

namespace ZB.MOM.WW.OtOpcUa.Host.IntegrationTests.LocalDb;

/// <summary>
///     lmxopcua#906 — a graceful stop of a 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 passive node is a real <see cref="WebApplication"/> — the hosting model whose hosted-service
///         order caused the bug (the web server stops before Akka) — with a real Kestrel h2c listener, the
///         real <c>MapZbLocalDbSync</c> behind the real auth interceptor, a real single-member Akka cluster
///         and the production <see cref="GracefulClusterLeaveRegistration.AddOtOpcUaGracefulClusterLeave"/>
///         + <see cref="ShutdownStreamCancellationInterceptor"/>. The peer is a real
///         <c>SyncBackgroundService</c> dialling it over loopback.
///     </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 WITHOUT the fix reproduces #906 — 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("LocalDbPairConvergence")]
public sealed class GracefulStopWithOpenSyncStreamTests
{
    private const string ApiKey = "otopcua-906-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);

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

    [Fact]
    public async Task GracefulStop_WithOpenSyncStream_LeavesClusterBeforeDrain_AndFinishesPromptly()
    {
        var result = await RunStopScenarioAsync(withFix: true);

        // Left the cluster, and did so before the web server began draining.
        result.MemberRemovedAt.ShouldNotBeNull("the stopping node never left the Akka cluster");
        result.MemberRemovedAt.Value.ShouldBeLessThan(result.ServerStopBeganAt!.Value,
            "the node must leave the cluster BEFORE Kestrel drains, so the peer promotes promptly");

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

    [Fact]
    public async Task PositiveControl_WithoutTheFix_StopWaitsOutTheDrain_AndLeavesTheClusterLast()
    {
        var result = await RunStopScenarioAsync(withFix: false);

        // The #906 shape: Kestrel waits for the open stream until the shutdown timeout aborts it,
        // and only then does Akka's own hosted service leave the cluster.
        result.StopElapsed.ShouldBeGreaterThanOrEqualTo(ShutdownTimeout - TimeSpan.FromSeconds(1));
        result.MemberRemovedAt.ShouldNotBeNull();
        result.MemberRemovedAt.Value.ShouldBeGreaterThan(result.ServerStopBeganAt!.Value);
        result.MemberRemovedAt.Value.ShouldBeGreaterThanOrEqualTo(ShutdownTimeout - TimeSpan.FromSeconds(1));
    }

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

    private static async Task<StopResult> RunStopScenarioAsync(bool withFix)
    {
        var ct = TestContext.Current.CancellationToken;
        var pathA = Path.Combine(Path.GetTempPath(), $"otopcua-906-A-{Guid.NewGuid():N}.db");
        var pathB = Path.Combine(Path.GetTempPath(), $"otopcua-906-B-{Guid.NewGuid():N}.db");
        await using var dbA = BuildDatabaseProvider(pathA);
        await using var dbB = BuildDatabaseProvider(pathB);
        var clock = new Stopwatch();
        var serverStopBegan = new ServerStopProbe(clock);
        IHost? initiator = null;
        WebApplication? passive = null;

        try
        {
            passive = BuildPassiveNode(dbB.GetRequiredService<ILocalDb>(), serverStopBegan, withFix);
            await passive.StartAsync(ct);

            var system = passive.Services.GetRequiredService<ActorSystem>();
            var cluster = AkkaCluster.Get(system);
            cluster.Join(cluster.SelfAddress);
            await WaitUntilAsync(() => cluster.State.Members.Any(m => m.Status == MemberStatus.Up),
                "the passive 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, ct);

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

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

            TestContext.Current.TestOutputHelper?.WriteLine(
                $"withFix={withFix}: stop {elapsed.TotalMilliseconds:0} ms; member removed at " +
                $"{memberRemovedAt?.TotalMilliseconds:0} ms; web-server drain began at " +
                $"{serverStopBegan.BeganAt?.TotalMilliseconds:0} ms (drain budget {ShutdownTimeout.TotalSeconds:0} s)");
            return new StopResult(elapsed, memberRemovedAt, serverStopBegan.BeganAt);
        }
        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 BuildPassiveNode(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);

        builder.Services.AddGrpc(o =>
        {
            o.Interceptors.Add<LocalDbSyncAuthInterceptor>();
            if (withFix)
                o.Interceptors.Add<ShutdownStreamCancellationInterceptor>();
        });
        builder.Services.AddSingleton(db);
        builder.Services.AddZbLocalDbReplication(ReplicationConfig(peerAddress: null));
        if (!withFix)
        {
            // From LocalDb 0.3.8 the library itself ends open sync streams on ApplicationStopping
            // (scadaproj#57), which would hide the #906 drain in this control. Re-register the
            // service WITHOUT the host lifetime, i.e. the 0.3.7 behaviour, so the control still
            // reproduces the bug the host-side fix guards against.
            builder.Services.Replace(ServiceDescriptor.Singleton(sp => new LocalDbSyncService(
                sp.GetRequiredService<ILocalDb>(),
                sp.GetRequiredService<IOptions<ReplicationOptions>>(),
                sp.GetRequiredService<ILoggerFactory>(),
                sp.GetService<LocalDbMetrics>(),
                sp.GetService<ISyncStatus>(),
                lifetime: null)));
        }

        builder.Services.AddAkka("otopcua906", (ab, _) => ab.AddHocon(
            ConfigurationFactory.ParseString(
                    "akka.remote.dot-netty.tcp.hostname = \"127.0.0.1\"\n"
                    + "akka.remote.dot-netty.tcp.public-hostname = \"127.0.0.1\"\n"
                    + "akka.remote.dot-netty.tcp.port = 0\n"
                    + "akka.cluster.seed-nodes = []\n"
                    + "akka.loglevel = \"WARNING\"\n")
                .WithFallback(HoconLoader.LoadBaseConfig()),
            HoconAddMode.Prepend));
        if (withFix)
            builder.Services.AddOtOpcUaGracefulClusterLeave();

        // 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, CancellationToken ct)
    {
        var config = ReplicationConfig(peerAddress);
        var host = new HostBuilder()
            .ConfigureServices(services =>
            {
                services.AddLogging();
                services.AddSingleton(db);
                services.AddZbLocalDbReplication(config);
            })
            .Build();
        await host.StartAsync(ct);
        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)
    {
        var config = new ConfigurationBuilder()
            .AddInMemoryCollection(new Dictionary<string, string?>
            {
                ["LocalDb:Path"] = path,
                ["LocalDb:Replication:ApiKey"] = ApiKey,
            })
            .Build();
        return new ServiceCollection()
            .AddZbLocalDb(config, db => LocalDbSetup.OnReady(db, config))
            .BuildServiceProvider();
    }

    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, TestContext.Current.CancellationToken);
        }

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

    /// <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();
        }
    }
}