diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndMultiSubscriberTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndMultiSubscriberTests.cs index 276bccb..4dfceed 100644 --- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndMultiSubscriberTests.cs +++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndMultiSubscriberTests.cs @@ -685,19 +685,19 @@ public sealed class GatewayEndToEndMultiSubscriberTests int advisedItemHandle = 0; int emittedCount = 0; - // Read envelopes one at a time. Between each envelope, if we have a - // subscription and the gate is ready, emit an event before reading the - // next envelope. When _stopEmitting is set, drain the gate and read - // remaining envelopes (including shutdown) without emitting. + await using GatedWorkerEnvelopeReader incoming = new( + harness.ReadGatewayEnvelopeAsync, _emitGate, cancellationToken); while (!cancellationToken.IsCancellationRequested) { - // While subscribed and not stopped, try to emit gated events using - // a short-timeout peek at the gate — yield to incoming envelopes to - // avoid starving shutdown processing. - while (advisedServerHandle != 0 - && !_stopEmitting - && await _emitGate.WaitAsync(millisecondsTimeout: 0).ConfigureAwait(false)) + WorkerEnvelope? envelope = await incoming.ReadOrGateAsync( + advisedServerHandle != 0 && !_stopEmitting).ConfigureAwait(false); + if (envelope is null) { + if (_stopEmitting) + { + continue; + } + int index = ++emittedCount; await harness.EmitEventAsync( MxEventFamily.OnDataChange, @@ -714,20 +714,6 @@ public sealed class GatewayEndToEndMultiSubscriberTests }; mxEvent.OnDataChange = new OnDataChangeEvent(); }).ConfigureAwait(false); - } - - // Use a short timeout so the emit loop above is re-evaluated - // periodically — but long enough not to spam. - WorkerEnvelope? envelope = null; - try - { - using CancellationTokenSource readCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - readCts.CancelAfter(TimeSpan.FromMilliseconds(50)); - envelope = await harness.ReadGatewayEnvelopeAsync(readCts.Token).ConfigureAwait(false); - } - catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) - { - // Timed out waiting for an envelope — loop back to check gate / emit. continue; } diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndReconnectReplayTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndReconnectReplayTests.cs index df3c53a..b23eea1 100644 --- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndReconnectReplayTests.cs +++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/GatewayEndToEndReconnectReplayTests.cs @@ -697,14 +697,19 @@ public sealed class GatewayEndToEndReconnectReplayTests int advisedItemHandle = 0; int emittedCount = 0; + await using GatedWorkerEnvelopeReader incoming = new( + harness.ReadGatewayEnvelopeAsync, _emitGate, cancellationToken); while (!cancellationToken.IsCancellationRequested) { - // While subscribed and not stopped, emit gated events using a non-blocking peek - // at the gate so incoming envelopes (including shutdown) are never starved. - while (advisedServerHandle != 0 - && !_stopEmitting - && await _emitGate.WaitAsync(millisecondsTimeout: 0).ConfigureAwait(false)) + WorkerEnvelope? envelope = await incoming.ReadOrGateAsync( + advisedServerHandle != 0 && !_stopEmitting).ConfigureAwait(false); + if (envelope is null) { + if (_stopEmitting) + { + continue; + } + int index = ++emittedCount; await harness.EmitEventAsync( MxEventFamily.OnDataChange, @@ -721,18 +726,6 @@ public sealed class GatewayEndToEndReconnectReplayTests }; mxEvent.OnDataChange = new OnDataChangeEvent(); }).ConfigureAwait(false); - } - - WorkerEnvelope? envelope; - try - { - using CancellationTokenSource readCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - readCts.CancelAfter(TimeSpan.FromMilliseconds(50)); - envelope = await harness.ReadGatewayEnvelopeAsync(readCts.Token).ConfigureAwait(false); - } - catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) - { - // Timed out waiting for an envelope — loop back to check the gate / emit. continue; } diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs index 0a04f13..5528e9d 100644 --- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs +++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs @@ -478,24 +478,42 @@ public sealed class SessionManagerTests [Fact] public async Task InvokeAsync_WhenWorkerHandshakingThenReadyWithinTimeout_Succeeds() { + FakeTimeProvider clock = new(new DateTimeOffset(2026, 1, 1, 0, 0, 0, TimeSpan.Zero)); FakeWorkerClient workerClient = new() { State = WorkerClientState.Handshaking }; SessionManager manager = CreateManager( new FakeSessionWorkerClientFactory(workerClient), - options: CreateOptions(workerReadyWaitTimeoutMs: 500)); + options: CreateOptions(workerReadyWaitTimeoutMs: 500), + timeProvider: clock); GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None); Assert.Equal(SessionState.Ready, session.State); - // Flip the worker to Ready shortly after the invoke starts waiting. - _ = Task.Run(async () => + TaskCompletionSource handshakingObserved = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource allowReady = new(TaskCreationOptions.RunContinuationsAsynchronously); + workerClient.OnStateRead = state => { - await Task.Delay(50, CancellationToken.None); - workerClient.State = WorkerClientState.Ready; - }); - - WorkerCommandReply reply = await manager.InvokeAsync( + if (state == WorkerClientState.Handshaking) + { + handshakingObserved.TrySetResult(); + } + }; + Task transition = TransitionAsync(); + Task invoke = manager.InvokeAsync( session.SessionId, CreateCommand(MxCommandKind.Ping), CancellationToken.None); + await handshakingObserved.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.False(invoke.IsCompleted); + Assert.Equal(0, workerClient.InvokeCount); + allowReady.SetResult(); + await transition.WaitAsync(TimeSpan.FromSeconds(5)); + clock.Advance(TimeSpan.FromMilliseconds(50)); + WorkerCommandReply reply = await invoke.WaitAsync(TimeSpan.FromSeconds(5)); + + async Task TransitionAsync() + { + await allowReady.Task; + workerClient.State = WorkerClientState.Ready; + } Assert.NotNull(reply); Assert.Equal(1, workerClient.InvokeCount); @@ -1406,14 +1424,28 @@ public sealed class SessionManagerTests private sealed class FakeWorkerClient : IWorkerClient { + private WorkerClientState _state = WorkerClientState.Ready; + /// public string SessionId { get; init; } = "session-1"; /// public int? ProcessId { get; init; } = 1234; + /// Gets or sets the optional state-observation gate. + public Action? OnStateRead { get; set; } + /// - public WorkerClientState State { get; set; } = WorkerClientState.Ready; + public WorkerClientState State + { + get + { + OnStateRead?.Invoke(_state); + return _state; + } + + set => _state = value; + } /// public DateTimeOffset LastHeartbeatAt { get; init; } = DateTimeOffset.UtcNow; diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/GatedWorkerEnvelopeReaderTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/GatedWorkerEnvelopeReaderTests.cs new file mode 100644 index 0000000..5ab097a --- /dev/null +++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/GatedWorkerEnvelopeReaderTests.cs @@ -0,0 +1,83 @@ +using System.IO.Pipelines; +using ZB.MOM.WW.MxGateway.Contracts; +using ZB.MOM.WW.MxGateway.Contracts.Proto; +using ZB.MOM.WW.MxGateway.Server.Workers; +using ZB.MOM.WW.MxGateway.Tests.TestSupport; + +namespace ZB.MOM.WW.MxGateway.Tests.Gateway.Workers; + +/// Protects fake worker framing while the emission gate wakes its loop. +public sealed class GatedWorkerEnvelopeReaderTests +{ + /// A gate wakeup must preserve a partially consumed frame and the following frame. + /// The asynchronous test operation. + [Fact] + public async Task ReadOrGateAsync_WhenGateOpensDuringSegmentedFrame_PreservesBothFrames() + { + WorkerFrameProtocolOptions options = new("segmented-session"); + WorkerEnvelope first = new() + { + SessionId = "segmented-session", + ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion, + Sequence = 1, + WorkerShutdown = new WorkerShutdown(), + }; + WorkerEnvelope second = first.Clone(); + second.Sequence = 2; + await using MemoryStream encoded = new(); + WorkerFrameWriter writer = new(encoded, options); + await writer.WriteAsync(first); + await writer.WriteAsync(second); + byte[] bytes = encoded.ToArray(); + Pipe pipe = new(); + await using Stream input = pipe.Reader.AsStream(); + WorkerFrameReader reader = new(input, options); + using SemaphoreSlim gate = new(0); + using CancellationTokenSource lifetime = new(TimeSpan.FromSeconds(5)); + await using GatedWorkerEnvelopeReader selector = new(token => reader.ReadAsync(token).AsTask(), gate, lifetime.Token); + + // The length and first payload byte are consumed before the remaining bytes exist. + await pipe.Writer.WriteAsync(bytes.AsMemory(0, 5)); + Task pending = selector.ReadOrGateAsync(); + Assert.False(pending.IsCompleted); + gate.Release(); + Assert.Null(await pending); + await pipe.Writer.WriteAsync(bytes.AsMemory(5)); + await pipe.Writer.CompleteAsync(); + + Assert.Equal(first, await selector.ReadOrGateAsync()); + Assert.Equal(second, await selector.ReadOrGateAsync()); + } + + /// Disposal or lifetime cancellation drains the retained read and gate wait. + /// Whether cancellation precedes disposal. + /// The asynchronous test operation. + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task DisposeAsync_WithIncompleteFrame_CancelsRetainedRead(bool cancelLifetime) + { + Pipe pipe = new(); + await using Stream input = pipe.Reader.AsStream(); + WorkerFrameReader reader = new(input, new WorkerFrameProtocolOptions("cancel-session")); + using SemaphoreSlim gate = new(0); + using CancellationTokenSource lifetime = new(); + Task? frameRead = null; + GatedWorkerEnvelopeReader selector = new( + token => frameRead = reader.ReadAsync(token).AsTask(), gate, lifetime.Token); + Task pending = selector.ReadOrGateAsync(); + gate.Release(); + Assert.Null(await pending.WaitAsync(TimeSpan.FromSeconds(5))); + Assert.NotNull(frameRead); + Assert.False(frameRead.IsCompleted); + if (cancelLifetime) + { + await lifetime.CancelAsync(); + } + + await selector.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(5)); + await Assert.ThrowsAnyAsync(() => frameRead); + await pipe.Writer.CompleteAsync(); + } + +} diff --git a/src/ZB.MOM.WW.MxGateway.Tests/TestSupport/GatedWorkerEnvelopeReader.cs b/src/ZB.MOM.WW.MxGateway.Tests/TestSupport/GatedWorkerEnvelopeReader.cs new file mode 100644 index 0000000..352bbce --- /dev/null +++ b/src/ZB.MOM.WW.MxGateway.Tests/TestSupport/GatedWorkerEnvelopeReader.cs @@ -0,0 +1,69 @@ +using ZB.MOM.WW.MxGateway.Contracts.Proto; + +namespace ZB.MOM.WW.MxGateway.Tests.TestSupport; + +/// Coordinates a fake worker's incoming frames and event emission gate. +internal sealed class GatedWorkerEnvelopeReader( + Func> read, + SemaphoreSlim gate, + CancellationToken cancellationToken) : IAsyncDisposable +{ + private readonly CancellationTokenSource _lifetime = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + private Task? _pendingRead; + private Task? _pendingGate; + + /// Reads an envelope, or returns null when emission should be checked. + /// Whether the worker is currently subscribed and emitting. + /// The next complete envelope, or null for an emission signal. + public async Task ReadOrGateAsync(bool watchGate = true) + { + // A gate wakeup must never cancel or replace a partially consumed frame. + _pendingRead ??= read(_lifetime.Token); + if (watchGate) + { + _pendingGate ??= gate.WaitAsync(_lifetime.Token); + await Task.WhenAny(_pendingRead, _pendingGate).ConfigureAwait(false); + if (!_pendingRead.IsCompleted) + { + await _pendingGate.ConfigureAwait(false); + _pendingGate = null; + return null; + } + } + + WorkerEnvelope envelope = await _pendingRead.ConfigureAwait(false); + _pendingRead = null; + return envelope; + } + + /// + public async ValueTask DisposeAsync() + { + await _lifetime.CancelAsync().ConfigureAwait(false); + try + { + if (_pendingRead is not null) + { + await _pendingRead.ConfigureAwait(false); + } + } + catch (OperationCanceledException) when (_lifetime.IsCancellationRequested) + { + } + finally + { + try + { + if (_pendingGate is not null) + { + await _pendingGate.ConfigureAwait(false); + } + } + catch (OperationCanceledException) when (_lifetime.IsCancellationRequested) + { + } + + _lifetime.Dispose(); + } + } +}