Preserve framed reads across fake worker emission gates

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
+            await using GatedWorkerEnvelopeReader incoming = new(
-            // subscription and the gate is ready, emit an event before reading the
+                harness.ReadGatewayEnvelopeAsync, _emitGate, cancellationToken);
-            // next envelope. When _stopEmitting is set, drain the gate and read
-            // remaining envelopes (including shutdown) without emitting.
             while (!cancellationToken.IsCancellationRequested)
             {
-                // While subscribed and not stopped, try to emit gated events using
+                WorkerEnvelope? envelope = await incoming.ReadOrGateAsync(
-                // a short-timeout peek at the gate — yield to incoming envelopes to
+                    advisedServerHandle != 0 && !_stopEmitting).ConfigureAwait(false);
-                // avoid starving shutdown processing.
+                if (envelope is null)
-                while (advisedServerHandle != 0
-                    && !_stopEmitting
-                    && await _emitGate.WaitAsync(millisecondsTimeout: 0).ConfigureAwait(false))
                 {
+                    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;
                 }
 

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
+                WorkerEnvelope? envelope = await incoming.ReadOrGateAsync(
-                // at the gate so incoming envelopes (including shutdown) are never starved.
+                    advisedServerHandle != 0 && !_stopEmitting).ConfigureAwait(false);
-                while (advisedServerHandle != 0
+                if (envelope is null)
-                    && !_stopEmitting
-                    && await _emitGate.WaitAsync(millisecondsTimeout: 0).ConfigureAwait(false))
                 {
+                    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;
                 }
 

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.
+        TaskCompletionSource handshakingObserved = new(TaskCreationOptions.RunContinuationsAsynchronously);
-        _ = Task.Run(async () =>
+        TaskCompletionSource allowReady = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        workerClient.OnStateRead = state =>
         {
-            await Task.Delay(50, CancellationToken.None);
+            if (state == WorkerClientState.Handshaking)
-            workerClient.State = WorkerClientState.Ready;
+            {
-        });
+                handshakingObserved.TrySetResult();
-
+            }
-        WorkerCommandReply reply = await manager.InvokeAsync(
+        };
+        Task transition = TransitionAsync();
+        Task<WorkerCommandReply> 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;
+
         /// <inheritdoc />
         public string SessionId { get; init; } = "session-1";
 
         /// <inheritdoc />
         public int? ProcessId { get; init; } = 1234;
 
+        /// <summary>Gets or sets the optional state-observation gate.</summary>
+        public Action<WorkerClientState>? OnStateRead { get; set; }
+
         /// <inheritdoc />
-        public WorkerClientState State { get; set; } = WorkerClientState.Ready;
+        public WorkerClientState State
+        {
+            get
+            {
+                OnStateRead?.Invoke(_state);
+                return _state;
+            }
+
+            set => _state = value;
+        }
 
         /// <inheritdoc />
         public DateTimeOffset LastHeartbeatAt { get; init; } = DateTimeOffset.UtcNow;

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;
+
+/// <summary>Protects fake worker framing while the emission gate wakes its loop.</summary>
+public sealed class GatedWorkerEnvelopeReaderTests
+{
+    /// <summary>A gate wakeup must preserve a partially consumed frame and the following frame.</summary>
+    /// <returns>The asynchronous test operation.</returns>
+    [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<WorkerEnvelope?> 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());
+    }
+
+    /// <summary>Disposal or lifetime cancellation drains the retained read and gate wait.</summary>
+    /// <param name="cancelLifetime">Whether cancellation precedes disposal.</param>
+    /// <returns>The asynchronous test operation.</returns>
+    [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<WorkerEnvelope>? frameRead = null;
+        GatedWorkerEnvelopeReader selector = new(
+            token => frameRead = reader.ReadAsync(token).AsTask(), gate, lifetime.Token);
+        Task<WorkerEnvelope?> 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<OperationCanceledException>(() => frameRead);
+        await pipe.Writer.CompleteAsync();
+    }
+
+}

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;
+
+/// <summary>Coordinates a fake worker's incoming frames and event emission gate.</summary>
+internal sealed class GatedWorkerEnvelopeReader(
+    Func<CancellationToken, Task<WorkerEnvelope>> read,
+    SemaphoreSlim gate,
+    CancellationToken cancellationToken) : IAsyncDisposable
+{
+    private readonly CancellationTokenSource _lifetime = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+    private Task<WorkerEnvelope>? _pendingRead;
+    private Task? _pendingGate;
+
+    /// <summary>Reads an envelope, or returns null when emission should be checked.</summary>
+    /// <param name="watchGate">Whether the worker is currently subscribed and emitting.</param>
+    /// <returns>The next complete envelope, or null for an emission signal.</returns>
+    public async Task<WorkerEnvelope?> 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;
+    }
+
+    /// <inheritdoc />
+    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();
+        }
+    }
+}