Preserve framed reads across fake worker emission gates

2026-10-13 09:42:09 -04:00
parent 19920b13da
commit 42c120aa83
5 changed files with 213 additions and 50 deletions
@@ -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();
+    }
+
+}