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,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();
+        }
+    }
+}