perf(worker): memoize event frame size at enqueue; drain stops sizing under the STA's lock

src/ZB.MOM.WW.MxGateway.Worker.Tests/MxAccess/MxAccessEventQueueTests.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
+using Google.Protobuf.WellKnownTypes;
 using ZB.MOM.WW.MxGateway.Contracts.Proto;
 using ZB.MOM.WW.MxGateway.Worker.MxAccess;
 
@@ -252,6 +253,158 @@ public sealed class MxAccessEventQueueTests
         Assert.Equal(3, result.RemainingCount);
     }
 
+    /// <summary>
+    ///     Brackets the byte budget's per-event charge to the exact serialized size of the event.
+    ///     The queue measures that size once at enqueue rather than during the drain, so this pins
+    ///     the two ends of the charge against the very instance the queue holds: a budget one byte
+    ///     short must refuse the head (and leave it queued, WRK-21), and a budget of exactly the
+    ///     charge must ship it. An undercharge — sizing the event before Enqueue stamps its worker
+    ///     sequence and timestamp, say — passes the first probe and breaks the frame guarantee; an
+    ///     overcharge of even one byte fails the second. The <c>preStamped</c> case covers an event
+    ///     that arrives with those two fields already filled, which Enqueue overwrites.
+    /// </summary>
+    /// <param name="preStamped">Whether the event carries a stale sequence/timestamp on arrival.</param>
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public void Drain_ByteBudget_ChargesTheEventsExactSerializedSize(bool preStamped)
+    {
+        MxAccessEventQueue queue = new(capacity: 4);
+        MxEvent mxEvent = CreateEventWithPayload(itemHandle: 7, payloadLength: 300);
+        if (preStamped)
+        {
+            mxEvent.WorkerSequence = ulong.MaxValue;
+            mxEvent.WorkerTimestamp = Timestamp.FromDateTime(
+                new DateTime(1990, 1, 1, 0, 0, 0, DateTimeKind.Utc));
+        }
+
+        queue.Enqueue(mxEvent);
+
+        // Enqueue takes ownership but never mutates the event afterwards, so the retained reference
+        // measures exactly what the queue holds — no second drain is needed to learn the cost.
+        Assert.Equal(1UL, mxEvent.WorkerSequence);
+        int exactCost = DrainCostOf(new WorkerEvent { Event = mxEvent });
+
+        WorkerEventDrainResult refused = queue.Drain(maxEvents: 0, maxTotalBytes: exactCost - 1);
+        Assert.Empty(refused.Events);
+        Assert.True(refused.TruncatedBySize);
+        Assert.Equal(1UL, refused.OversizedHeadSequence);
+        Assert.Equal(1, queue.Count);
+
+        WorkerEventDrainResult drained = queue.Drain(maxEvents: 0, maxTotalBytes: exactCost);
+        Assert.Single(drained.Events);
+        Assert.Same(mxEvent, drained.Events[0].Event);
+        Assert.False(drained.TruncatedBySize);
+        Assert.Equal(0, queue.Count);
+    }
+
+    /// <summary>
+    ///     Pins the batch boundaries of a byte-budgeted walk over a mixed-size queue. Each batch is
+    ///     checked against the true serialized sizes of the events it returned: no batch may exceed
+    ///     the budget (an undercharge would build a reply past the frame maximum) and no batch may
+    ///     stop early — the next event, at its real size, must not have fit (an overcharge would
+    ///     ship frames smaller than the negotiated maximum allows). Both bounds come from the
+    ///     drained events themselves, so any drift between the size memoized at enqueue and the real
+    ///     one moves a boundary and fails here.
+    /// </summary>
+    [Fact]
+    public void Drain_ByteBudget_MixedSizes_KeepsBatchBoundariesOnTheRealSizes()
+    {
+        int[] payloadLengths = { 8, 512, 40, 2_048, 96, 1_200, 16, 700 };
+        const int eventCount = 240;
+
+        // Comfortably above the largest single event's cost, so the walk is never blocked by an
+        // oversized head and every stop is a genuine budget boundary.
+        const int budget = 4096;
+
+        MxAccessEventQueue queue = new(eventCount);
+        for (int index = 0; index < eventCount; index++)
+        {
+            queue.Enqueue(CreateEventWithPayload(index, payloadLengths[index % payloadLengths.Length]));
+        }
+
+        List<IReadOnlyList<WorkerEvent>> batches = new();
+        while (true)
+        {
+            WorkerEventDrainResult result = queue.Drain(maxEvents: 0, maxTotalBytes: budget);
+            if (result.Events.Count == 0)
+            {
+                Assert.Equal(0UL, result.OversizedHeadSequence);
+                break;
+            }
+
+            batches.Add(result.Events);
+            Assert.True(batches.Count <= eventCount, "Drain made no progress.");
+        }
+
+        Assert.Equal(0, queue.Count);
+        Assert.True(
+            batches.Count > 5,
+            $"Expected the byte budget to split the drain, saw {batches.Count} batches.");
+
+        ulong expectedSequence = 0;
+        for (int batchIndex = 0; batchIndex < batches.Count; batchIndex++)
+        {
+            IReadOnlyList<WorkerEvent> batch = batches[batchIndex];
+            int charged = 0;
+            foreach (WorkerEvent drained in batch)
+            {
+                charged += DrainCostOf(drained);
+                Assert.Equal(++expectedSequence, drained.Event.WorkerSequence);
+            }
+
+            Assert.True(
+                charged <= budget,
+                $"Batch {batchIndex} shipped {charged} bytes against a {budget} byte budget.");
+
+            if (batchIndex + 1 < batches.Count)
+            {
+                int nextCost = DrainCostOf(batches[batchIndex + 1][0]);
+                Assert.True(
+                    charged + nextCost > budget,
+                    $"Batch {batchIndex} stopped at {charged} bytes although the next event's {nextCost} still fit the {budget} byte budget.");
+            }
+        }
+
+        Assert.Equal((ulong)eventCount, expectedSequence);
+    }
+
+    /// <summary>
+    ///     The other half of the WRK-21 head guarantee: a head refused for its size is drained
+    ///     unchanged by a later call whose budget fits it. The size the queue charges lives with the
+    ///     event across calls, so a refused attempt must neither consume nor alter it — and the
+    ///     refusal itself is justified by the event's real serialized size.
+    /// </summary>
+    [Fact]
+    public void Drain_ByteBudget_RefusedHead_IsDrainedUnchangedOnceTheBudgetFitsIt()
+    {
+        MxAccessEventQueue queue = new(capacity: 8);
+        queue.Enqueue(CreateEventWithPayload(itemHandle: 0, payloadLength: 4096));
+        queue.Enqueue(CreateEventWithPayload(itemHandle: 1, payloadLength: 8));
+
+        WorkerEventDrainResult refused = queue.Drain(maxEvents: 0, maxTotalBytes: 1024);
+        Assert.Empty(refused.Events);
+        Assert.True(refused.TruncatedBySize);
+        Assert.Equal(1UL, refused.OversizedHeadSequence);
+        Assert.Equal(2, queue.Count);
+
+        WorkerEventDrainResult retried = queue.Drain(maxEvents: 0, maxTotalBytes: 64 * 1024);
+
+        Assert.Equal(2, retried.Events.Count);
+        Assert.Equal(1UL, retried.Events[0].Event.WorkerSequence);
+        Assert.Equal(2UL, retried.Events[1].Event.WorkerSequence);
+        Assert.False(retried.TruncatedBySize);
+        Assert.Equal(0, retried.RemainingCount);
+        Assert.Equal(0, queue.Count);
+
+        Assert.True(
+            DrainCostOf(retried.Events[0]) > 1024,
+            "The head was refused although its real cost fits the 1024-byte budget it was refused under.");
+        Assert.True(
+            DrainCostOf(retried.Events[0]) + DrainCostOf(retried.Events[1]) <= 64 * 1024,
+            "The retried batch exceeded the budget it was drained under.");
+    }
+
     /// <summary>Verifies that Enqueue is rejected after a fault is recorded manually.</summary>
     [Fact]
     public void Enqueue_AfterRecordFault_ThrowsInvalidOperationException()
@@ -515,6 +668,19 @@ public sealed class MxAccessEventQueueTests
     /// </summary>
     /// <param name="payloadLength">Length of the event's raw-status payload string.</param>
     /// <returns>The per-event byte cost.</returns>
+    /// <summary>
+    /// What the byte budget must charge for an already-stamped event: its true serialized size plus
+    /// the repeated-field allowance. Measured from the generated <c>CalculateSize()</c> so the
+    /// budget tests bound the queue's memoized size against the real one rather than against a
+    /// second copy of the queue's own arithmetic.
+    /// </summary>
+    /// <param name="workerEvent">Event as the queue holds it, sequence and timestamp stamped.</param>
+    /// <returns>The per-event byte cost.</returns>
+    private static int DrainCostOf(WorkerEvent workerEvent)
+    {
+        return workerEvent.CalculateSize() + RepeatedFieldOverheadBytes;
+    }
+
     private static int MeasureDrainCost(int payloadLength)
     {
         MxAccessEventQueue probe = new(capacity: 1);