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

docs/MxAccessWorkerInstanceDesign.md

@@ -541,7 +541,11 @@ is bounded on **two** axes because no diagnostics command may be session-fatal:
   maximum less a 64 KiB envelope/reply-wrapper reserve, and the size decision
   happens inside the event queue's lock, so an event is dequeued only once it is
   known to fit. An event that does not fit stays at the head of the queue and is
-  never lost.
+  never lost. Each event's serialized size is *measured* once at enqueue, outside
+  that lock, and stored beside it: the drain only compares memoized numbers, so a
+  large drain never walks messages under the lock the STA needs to enqueue the
+  next COM callback. The memoized size cannot go stale because an enqueued event
+  is never mutated again (WRK-11).
 
 Truncation is reported in the reply's existing `DiagnosticMessage`
 ("N events returned, M remain; repeat DrainEvents for the rest") rather than in a

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; }
+    }
 }