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;
///
/// A worker that goes away on its own — the client sent the ShutdownWorker control
/// command and the worker exited cleanly, or the worker crashed / its pipe broke — must be
/// noticed by its . The real reaches
/// Closed / Faulted, but nothing used to tell the session, which kept reporting
/// Ready (live windev finding) until its lease expired.
///
public sealed class GatewaySessionWorkerExitTests
{
private static readonly TimeSpan TestTimeout = TimeSpan.FromSeconds(5);
///
/// After the client's ShutdownWorker 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.
///
/// A task that represents the asynchronous operation.
[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 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);
}
///
/// The other way a worker goes away: its pipe breaks (crash). The worker client faults, and
/// the session must follow it to Faulted rather than stay Ready.
///
/// A task that represents the asynchronous operation.
[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 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 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
{
///
public Task CreateAsync(GatewaySession session, CancellationToken cancellationToken)
{
return Task.FromResult(client);
}
}
}