mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/GatewaySessionWorkerExitTests.cs

using Microsoft.Extensions.Options;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Metrics;
using ZB.MOM.WW.MxGateway.Server.Sessions;
using ZB.MOM.WW.MxGateway.Server.Workers;
using ZB.MOM.WW.MxGateway.Tests.Gateway.Workers.Fakes;

namespace ZB.MOM.WW.MxGateway.Tests.Gateway.Sessions;

/// <summary>
///     A worker that goes away on its own — the client sent the <c>ShutdownWorker</c> control
///     command and the worker exited cleanly, or the worker crashed / its pipe broke — must be
///     noticed by its <see cref="GatewaySession"/>. The real <see cref="WorkerClient"/> reaches
///     <c>Closed</c> / <c>Faulted</c>, but nothing used to tell the session, which kept reporting
///     <c>Ready</c> (live windev finding) until its lease expired.
/// </summary>
public sealed class GatewaySessionWorkerExitTests
{
    private static readonly TimeSpan TestTimeout = TimeSpan.FromSeconds(5);

    /// <summary>
    ///     After the client's <c>ShutdownWorker</c> command the fake worker replies OK, sends the
    ///     shutdown ack and the worker client closes. The session must fault (the worker is gone
    ///     outside a gateway-initiated close), stay observable with a fault reason, and then be
    ///     reaped by the ordinary faulted-session sweep — which is where the session is accounted
    ///     closed, exactly once.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ClientShutdownWorker_WorkerExits_SessionFaultsAndSweepReapsItOnce()
    {
        await using FakeWorkerHarness fakeWorker = await FakeWorkerHarness.CreateConnectedPairAsync();
        using GatewayMetrics metrics = new();
        WorkerClient client = fakeWorker.CreateClient(metrics: metrics);
        await StartClientAsync(fakeWorker, client);
        SessionManager manager = CreateManager(client, metrics);
        GatewaySession session = await OpenSessionAsync(manager);
        using IEventSubscriberLease stream = session.AttachEventSubscriber(maxSubscribers: 1);

        Task<WorkerCommandReply> invokeTask = manager.InvokeAsync(session.SessionId, CreateShutdownWorkerCommand(), CancellationToken.None);
        await fakeWorker.RespondToControlCommandAsync().WaitAsync(TestTimeout);
        WorkerCommandReply reply = await invokeTask.WaitAsync(TestTimeout);
        Assert.Equal(ProtocolStatusCode.Ok, reply.Reply.ProtocolStatus.Code);

        await WaitUntilAsync(() => session.State == SessionState.Faulted);

        Assert.Equal(WorkerClientState.Closed, client.State);
        Assert.Contains("worker", session.FinalFault, StringComparison.OrdinalIgnoreCase);
        Assert.True(manager.TryGetSession(session.SessionId, out _));
        Assert.Equal(1, metrics.GetSnapshot().OpenSessions);
        Assert.Equal(0, metrics.GetSnapshot().SessionsClosed);

        int closedCount = await manager.CloseExpiredLeasesAsync(DateTimeOffset.UtcNow, CancellationToken.None);

        Assert.Equal(1, closedCount);
        Assert.Equal(SessionState.Closed, session.State);
        Assert.False(manager.TryGetSession(session.SessionId, out _));
        GatewayMetricsSnapshot snapshot = metrics.GetSnapshot();
        Assert.Equal(1, snapshot.SessionsClosed);
        Assert.Equal(0, snapshot.OpenSessions);
    }

    /// <summary>
    ///     The other way a worker goes away: its pipe breaks (crash). The worker client faults, and
    ///     the session must follow it to <c>Faulted</c> rather than stay <c>Ready</c>.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WorkerPipeBreaks_SessionFaults()
    {
        await using FakeWorkerHarness fakeWorker = await FakeWorkerHarness.CreateConnectedPairAsync();
        using GatewayMetrics metrics = new();
        WorkerClient client = fakeWorker.CreateClient(metrics: metrics);
        await StartClientAsync(fakeWorker, client);
        SessionManager manager = CreateManager(client, metrics);
        GatewaySession session = await OpenSessionAsync(manager);
        using IEventSubscriberLease stream = session.AttachEventSubscriber(maxSubscribers: 1);

        await fakeWorker.DisposeWorkerSideAsync();

        await WaitUntilAsync(() => session.State == SessionState.Faulted);
        Assert.Equal(WorkerClientState.Faulted, client.State);

        await manager.CloseExpiredLeasesAsync(DateTimeOffset.UtcNow, CancellationToken.None);
        Assert.False(manager.TryGetSession(session.SessionId, out _));
        Assert.Equal(1, metrics.GetSnapshot().SessionsClosed);
        Assert.Equal(0, metrics.GetSnapshot().OpenSessions);
    }

    private static async Task StartClientAsync(FakeWorkerHarness fakeWorker, WorkerClient client)
    {
        Task startTask = client.StartAsync(CancellationToken.None);
        await fakeWorker.CompleteStartupAsync();
        await startTask.WaitAsync(TestTimeout);
        Assert.Equal(WorkerClientState.Ready, client.State);
    }

    private static SessionManager CreateManager(IWorkerClient client, GatewayMetrics metrics)
    {
        return new SessionManager(
            new SessionRegistry(),
            new SingleClientFactory(client),
            Options.Create(new GatewayOptions()),
            metrics);
    }

    private static Task<GatewaySession> OpenSessionAsync(SessionManager manager)
    {
        return manager.OpenSessionAsync(
            new SessionOpenRequest(RequestedBackend: null, ClientSessionName: "worker-exit", ClientCorrelationId: null, CommandTimeout: null),
            "client-1",
            ownerKeyId: null,
            CancellationToken.None);
    }

    private static WorkerCommand CreateShutdownWorkerCommand()
    {
        return new WorkerCommand
        {
            Command = new MxCommand
            {
                Kind = MxCommandKind.ShutdownWorker,
                ShutdownWorker = new ShutdownWorkerCommand(),
            },
        };
    }

    private static async Task WaitUntilAsync(Func<bool> condition)
    {
        DateTime deadline = DateTime.UtcNow + TestTimeout;
        while (!condition())
        {
            if (DateTime.UtcNow > deadline)
            {
                Assert.Fail("Condition was not met within the test timeout.");
            }

            await Task.Delay(10);
        }
    }

    private sealed class SingleClientFactory(IWorkerClient client) : ISessionWorkerClientFactory
    {
        /// <inheritdoc />
        public Task<IWorkerClient> CreateAsync(GatewaySession session, CancellationToken cancellationToken)
        {
            return Task.FromResult(client);
        }
    }
}