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);

src/ZB.MOM.WW.MxGateway.Worker/MxAccess/MxAccessEventQueue.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Globalization;
 using System.Threading;
 using System.Threading.Tasks;
+using Google.Protobuf;
 using Google.Protobuf.WellKnownTypes;
 using ZB.MOM.WW.MxGateway.Contracts.Proto;
 
@@ -21,6 +22,11 @@ namespace ZB.MOM.WW.MxGateway.Worker.MxAccess;
 ///     MxAccessAlarmEventSink, AlarmCommandHandler) build a new event per
 ///     Enqueue via the mapper and satisfy this; the value cache stores its own
 ///     independent snapshot (see <see cref="MxAccessValueCache.Set"/>).
+///     <para>
+///         The byte-budgeted <see cref="Drain(uint, int)"/> relies on the same invariant: each
+///         event's serialized size is measured once at enqueue and stored beside it, which is
+///         only sound because an enqueued event is never mutated again.
+///     </para>
 /// </remarks>
 public sealed class MxAccessEventQueue
 {
@@ -63,8 +69,16 @@ public sealed class MxAccessEventQueue
     // edge can never push the packed reply past the frame maximum.
     private const int RepeatedFieldOverheadBytes = 8;
 
+    // Every field the enqueue-time size reconstruction charges by hand — WorkerEvent.event (field
+    // 1), MxEvent.worker_sequence (9) and MxEvent.worker_timestamp (10) — has a field number below
+    // 16, so protobuf encodes each of their tags in a single byte. Renumbering any of them above 15
+    // would widen its tag and make this estimate one byte light per field, which the
+    // RepeatedFieldOverheadBytes slack absorbs — so the frame stays safe — but the constant should
+    // be revisited with the .proto.
+    private const int SingleByteFieldTagBytes = 1;
+
     private readonly int capacity;
-    private readonly Queue<WorkerEvent> events;
+    private readonly Queue<QueuedEvent> events;
     private readonly object syncRoot = new();
 
     // Wake signal for the worker's event drain loop (see WorkerPipeSession.RunEventDrainLoopAsync).
@@ -110,7 +124,7 @@ public sealed class MxAccessEventQueue
         }
 
         this.capacity = capacity;
-        events = new Queue<WorkerEvent>(capacity);
+        events = new Queue<QueuedEvent>(capacity);
     }
 
     /// <summary>
@@ -210,6 +224,28 @@ public sealed class MxAccessEventQueue
             throw new ArgumentNullException(nameof(mxEvent));
         }
 
+        // Size the event here, outside syncRoot, and carry the result into the queue. The
+        // byte-budgeted Drain used to call CalculateSize() on every candidate while holding the
+        // lock, so a large gateway-pulled drain kept the STA's next Enqueue blocked for the whole
+        // sizing walk — exactly when a burst of COM callbacks needs the lock. Measuring on the
+        // enqueuing thread is safe on both counts: that thread owns this event exclusively at this
+        // point, and WRK-11 (see the class remarks) makes the hand-off final — nothing mutates an
+        // event after Enqueue, so a memoized size can never go stale.
+        //
+        // Enqueue itself stamps the worker sequence and worker timestamp, and those must stay under
+        // the lock (the sequence has to be assigned in queue order), so they are excluded from this
+        // measurement and charged back by CalculateWorkerEventSize with O(1) arithmetic instead of a
+        // second walk of the message. Clearing them first is observationally a no-op — both are
+        // overwritten unconditionally below — and it makes this measurement exactly "the event minus
+        // the two stamped fields" even if a caller pre-filled them.
+        //
+        // CalculateSize() on a valid protobuf message does not throw, so moving it ahead of the
+        // fault and overflow checks introduces no new failure before them; an event those checks
+        // reject simply discards the size computed here.
+        mxEvent.WorkerSequence = 0UL;
+        mxEvent.WorkerTimestamp = null;
+        int unstampedEventSize = mxEvent.CalculateSize();
+
         try
         {
             lock (syncRoot)
@@ -237,7 +273,9 @@ public sealed class MxAccessEventQueue
                 {
                     Event = mxEvent,
                 };
-                events.Enqueue(workerEvent);
+                events.Enqueue(new QueuedEvent(
+                    workerEvent,
+                    CalculateWorkerEventSize(mxEvent, unstampedEventSize)));
             }
         }
         finally
@@ -295,7 +333,7 @@ public sealed class MxAccessEventQueue
                 return false;
             }
 
-            workerEvent = events.Dequeue();
+            workerEvent = events.Dequeue().Event;
             return true;
         }
     }
@@ -320,7 +358,7 @@ public sealed class MxAccessEventQueue
             List<WorkerEvent> drained = new(drainCount);
             for (int index = 0; index < drainCount; index++)
             {
-                drained.Add(events.Dequeue());
+                drained.Add(events.Dequeue().Event);
             }
 
             return drained;
@@ -334,11 +372,13 @@ public sealed class MxAccessEventQueue
     /// <remarks>
     ///     The size decision happens inside the queue lock, so an event is dequeued only once it is
     ///     known to fit: an event that does not fit stays at the head for the next call and is never
-    ///     lost (WRK-21). Per-event cost is <c>WorkerEvent.CalculateSize()</c> — which already
-    ///     includes the event's own tag and length prefix, the same shape the reply's
-    ///     <c>events</c> repeated field packs it into — plus <see cref="RepeatedFieldOverheadBytes"/>
-    ///     of pure slack, so the running total stays strictly ahead of the true serialized size and
-    ///     the estimate errs on the safe side.
+    ///     lost (WRK-21). Per-event cost is the serialized <c>WorkerEvent</c> size measured once at
+    ///     enqueue time — which already includes the event's own tag and length prefix, the same
+    ///     shape the reply's <c>events</c> repeated field packs it into — plus
+    ///     <see cref="RepeatedFieldOverheadBytes"/> of pure slack, so the running total stays
+    ///     strictly ahead of the true serialized size and the estimate errs on the safe side. Only
+    ///     the memoized number is read here: this loop no longer walks a message under the lock the
+    ///     STA needs for <see cref="Enqueue"/>, so a large drain cannot stall COM callbacks.
     /// </remarks>
     /// <param name="maxEvents">Maximum number of events to drain; 0 means "no count limit".</param>
     /// <param name="maxTotalBytes">Byte budget for the drained events' estimated serialized size.</param>
@@ -357,8 +397,8 @@ public sealed class MxAccessEventQueue
 
             while (drained.Count < countLimit && events.Count > 0)
             {
-                WorkerEvent head = events.Peek();
-                int cost = head.CalculateSize() + RepeatedFieldOverheadBytes;
+                QueuedEvent head = events.Peek();
+                int cost = head.Size + RepeatedFieldOverheadBytes;
                 if (cost > remainingBudget)
                 {
                     truncatedBySize = true;
@@ -367,14 +407,14 @@ public sealed class MxAccessEventQueue
                         // The head alone cannot fit this budget, so repeating the call will not
                         // move it either. Report its sequence instead of silently stalling; the
                         // events that did fit are still returned.
-                        oversizedHeadSequence = head.Event?.WorkerSequence ?? 0;
+                        oversizedHeadSequence = head.Event.Event?.WorkerSequence ?? 0;
                     }
 
                     break;
                 }
 
                 remainingBudget -= cost;
-                drained.Add(events.Dequeue());
+                drained.Add(events.Dequeue().Event);
             }
 
             return new WorkerEventDrainResult(
@@ -447,6 +487,24 @@ public sealed class MxAccessEventQueue
         }
     }
 
+    // Serialized size of the WorkerEvent wrapping mxEvent, rebuilt from the size measured before the
+    // lock plus the exact cost of the two fields Enqueue stamps inside it. It mirrors the generated
+    // CalculateSize() instead of calling it, because calling it again under the lock is the cost
+    // this memoization exists to remove: WorkerEvent carries the single length-delimited `event`
+    // field, MxEvent charges worker_sequence as tag + varint and worker_timestamp as tag +
+    // length-delimited message, and both are always serialized here — the sequence is
+    // pre-incremented, so it is never the proto3 default that generated code would skip, and the
+    // timestamp is never null. Exactness matters: Drain must charge the same bytes it charged when
+    // it sized events itself, so an unchanged budget still yields unchanged batches.
+    private static int CalculateWorkerEventSize(MxEvent mxEvent, int unstampedEventSize)
+    {
+        int eventSize = unstampedEventSize
+            + SingleByteFieldTagBytes + CodedOutputStream.ComputeUInt64Size(mxEvent.WorkerSequence)
+            + SingleByteFieldTagBytes + CodedOutputStream.ComputeMessageSize(mxEvent.WorkerTimestamp);
+
+        return SingleByteFieldTagBytes + CodedOutputStream.ComputeLengthSize(eventSize) + eventSize;
+    }
+
     private WorkerFault CreateOverflowFault()
     {
         string message = $"MXAccess outbound event queue reached capacity {capacity}.";
@@ -461,4 +519,24 @@ public sealed class MxAccessEventQueue
             },
         };
     }
+
+    // A queued event and the serialized size measured for it at enqueue time. A struct, so carrying
+    // the size costs the queue one extra field per slot rather than an allocation per event on the
+    // STA's path. Nothing outside the queue sees it: every drain hands back the WorkerEvent alone,
+    // exactly as before.
+    private readonly struct QueuedEvent
+    {
+        public QueuedEvent(WorkerEvent workerEvent, int size)
+        {
+            Event = workerEvent;
+            Size = size;
+        }
+
+        // The queued event, returned unchanged by TryDequeue and both Drain overloads.
+        public WorkerEvent Event { get; }
+
+        // Serialized size of Event, valid for the queue's whole lifetime because an enqueued event
+        // is never mutated again (WRK-11).
+        public int Size { get; }
+    }
 }