perf(events): copy-on-write subscriber snapshot in fan-out pump

src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionEventDistributor.cs

@@ -67,11 +67,22 @@ public delegate void SubscriberOverflowHandler(bool isOnlySubscriber, bool isInt
 ///     </para>
 ///     <para>
 ///         <b>Concurrency.</b> The subscriber set is a
-///         <see cref="ConcurrentDictionary{TKey, TValue}"/> keyed by a monotonic id.
+///         <see cref="ConcurrentDictionary{TKey, TValue}"/> keyed by a monotonic id, used
-///         The pump iterates it with a snapshot-free enumerator (which never throws on
+///         for keyed add/remove only. Fan-out does NOT enumerate the dictionary: every
-///         concurrent add/remove), and <see cref="Register"/> / lease disposal mutate it
+///         mutation (<see cref="Register"/>, <see cref="RegisterWithReplay"/>, lease
-///         without any lock held across an <c>await</c>. Each subscriber channel has a
+///         disposal, overflow disconnect) happens inside the <c>_lifecycleLock</c> critical
-///         single writer — the pump — so per-channel writes never race. MXAccess parity:
+///         section and rebuilds an immutable copy-on-write <c>Subscriber[]</c> snapshot,
+///         which the pump reads once per event. This matters because
+///         <c>ConcurrentDictionary.Values</c> is a PROPERTY that acquires every internal
+///         lock and materializes a fresh <c>List</c> plus a read-only wrapper on each call
+///         — per event, on the hot fan-out path. The subscriber set is tiny (one to a
+///         handful) and mutates rarely, so paying a full array rebuild per registration to
+///         make fan-out a bare array walk is the right trade. No lock is held across an
+///         <c>await</c>. Each subscriber channel has a single writer — the pump — so
+///         per-channel writes never race. A subscriber registered after the pump captured
+///         the array for the in-flight event misses that event, which matches "late
+///         subscribers see events after they register"; the reconnect path closes that
+///         window deliberately (see <see cref="RegisterWithReplay"/>). MXAccess parity:
 ///         events are fanned in the order received; the pump never reorders or
 ///         synthesizes events.
 ///     </para>
@@ -97,6 +108,14 @@ public sealed class SessionEventDistributor : IAsyncDisposable
     private readonly CancellationTokenSource _shutdownCts = new();
     private readonly object _lifecycleLock = new();
 
+    // Copy-on-write fan-out snapshot of _subscribers.Values. Rebuilt (a whole new array)
+    // inside the _lifecycleLock section of every register/unregister; never mutated in
+    // place, so the pump can walk the array it captured with no lock and no allocation.
+    // Written with Volatile.Write / read with Volatile.Read so a reader on another core
+    // cannot observe a stale reference after the publishing store. See the type remarks
+    // for why fan-out must not touch ConcurrentDictionary.Values.
+    private Subscriber[] _subscriberSnapshot = [];
+
     // Replay ring buffer. Appended on the pump thread and queried from arbitrary
     // threads via TryGetReplayFrom, so every access is under _replayLock. Backed by a
     // fixed-size circular array preallocated to the capacity so appending a retained
@@ -268,7 +287,12 @@ public sealed class SessionEventDistributor : IAsyncDisposable
     ///     <see cref="GatewaySession.ActiveEventSubscriberCount"/>, which tracks only external
     ///     (gRPC) subscribers and excludes the internal dashboard subscriber.
     /// </summary>
-    public int SubscriberCount => _subscribers.Count;
+    /// <remarks>
+    ///     Read from the copy-on-write snapshot rather than <c>ConcurrentDictionary.Count</c>
+    ///     (which acquires every internal lock). The snapshot is rebuilt in the same
+    ///     <c>_lifecycleLock</c> section that mutates the dictionary, so the two never diverge.
+    /// </remarks>
+    public int SubscriberCount => Volatile.Read(ref _subscriberSnapshot).Length;
 
     /// <summary>
     /// Starts the background pump. Idempotent — a second call is a no-op.
@@ -332,6 +356,7 @@ public sealed class SessionEventDistributor : IAsyncDisposable
         {
             ObjectDisposedException.ThrowIf(_disposed, this);
             _subscribers[subscriber.Id] = subscriber;
+            RebuildSubscriberSnapshot();
 
             // Close the register-after-pump-completion window: if the pump already ran its
             // final CompleteAllSubscribers (source completed/faulted) but the distributor is
@@ -416,24 +441,26 @@ public sealed class SessionEventDistributor : IAsyncDisposable
     ///     <para>
     ///         <b>Why this is atomic and the handoff is correct.</b> The replay snapshot and the
     ///         subscriber registration both run inside the SAME <c>_replayLock</c> critical
-    ///         section. The pump appends each event to the replay buffer under <c>_replayLock</c>
+    ///         section. The pump appends each event to the replay buffer AND captures the
-    ///         <em>before</em> fanning it to subscribers (outside the lock). Therefore, relative
+    ///         copy-on-write subscriber array in one <c>_replayLock</c> section, then fans the
-    ///         to this method's critical section, for every event E:
+    ///         event to that captured array outside the lock. Mutual exclusion therefore places
+    ///         every event E strictly on one side of this method's critical section:
     ///     </para>
     ///     <list type="bullet">
     ///         <item>
     ///             If the pump appended E before this critical section, E is in
     ///             <paramref name="replayedEvents"/> (when newer than
-    ///             <paramref name="afterSequence"/>). The pump's fan-out of E may race the
+    ///             <paramref name="afterSequence"/>). The pump captured its subscriber array in
-    ///             registration: if it writes E to this new channel too, E's sequence is
+    ///             that same earlier section, so it cannot fan E into this not-yet-registered
-    ///             <c>&lt;= liveResumeSequence</c>, so the caller's live filter DROPS it — no
+    ///             channel. Even if it could, E's sequence is <c>&lt;= liveResumeSequence</c>, so
-    ///             duplicate.
+    ///             the caller's live filter DROPS it — no duplicate either way.
     ///         </item>
     ///         <item>
     ///             If the pump appends E after this critical section, E is NOT in the snapshot,
-    ///             but this subscriber is already registered, so the pump fans E into the live
+    ///             but this subscriber was registered — and the snapshot array republished —
-    ///             channel with sequence <c>&gt; liveResumeSequence</c> — delivered as live, no
+    ///             before that section began, so the pump's capture includes it and E is fanned
-    ///             gap.
+    ///             into the live channel with sequence <c>&gt; liveResumeSequence</c> — delivered
+    ///             as live, no gap.
     ///         </item>
     ///     </list>
     ///     <para>
@@ -508,6 +535,7 @@ public sealed class SessionEventDistributor : IAsyncDisposable
             {
                 ObjectDisposedException.ThrowIf(_disposed, this);
                 _subscribers[id] = subscriber;
+                RebuildSubscriberSnapshot();
 
                 // Same register-after-pump-completion guard as Register: a resume that races in
                 // after the source already ended still gets its retained replay batch (snapshot
@@ -591,13 +619,19 @@ public sealed class SessionEventDistributor : IAsyncDisposable
                 // Retain for replay BEFORE fan-out so a reconnecting subscriber that
                 // queries between fan-out and its own read still sees this event. Order
                 // is preserved: the pump is the single appender and events arrive in
-                // source order.
+                // source order. The same call returns the subscriber array to fan to,
-                AppendToReplayBuffer(mxEvent);
+                // captured under _replayLock — see the method for why the capture must
+                // share the append's critical section.
+                Subscriber[] subscribers = AppendToReplayBufferAndCaptureSubscribers(mxEvent);
 
-                // Enumerating a ConcurrentDictionary's Values never throws on concurrent
+                // Walk the captured copy-on-write array: no dictionary enumeration, no
-                // add/remove; a subscriber registered mid-iteration may miss this event,
+                // per-event allocation. A subscriber registered after this capture misses
-                // which matches "late subscribers see events after they register".
+                // this event, which matches "late subscribers see events after they
-                foreach (Subscriber subscriber in _subscribers.Values)
+                // register". A subscriber unregistered after the capture is still written
+                // to — the identical window the previous ConcurrentDictionary.Values
+                // enumeration had (that property also materialized its list up front), so
+                // the race and its outcome are unchanged by the copy-on-write array.
+                foreach (Subscriber subscriber in subscribers)
                 {
                     // Non-blocking write: TryWrite never blocks the pump on a slow reader.
                     // A false return means this subscriber's bounded channel is full — the
@@ -669,7 +703,11 @@ public sealed class SessionEventDistributor : IAsyncDisposable
         // remove it from the fan-out set. Its gRPC reader's MoveNextAsync then throws the
         // SessionManagerException, which EventStreamService surfaces to the client exactly as
         // the pre-epic per-RPC overflow did. The pump and every other subscriber are untouched.
-        if (_subscribers.TryRemove(subscriber.Id, out _))
+        //
+        // The removal takes _lifecycleLock because it must republish the fan-out snapshot; the
+        // pump holds no other lock here (fan-out runs outside _replayLock), so this cannot
+        // invert the _replayLock-then-_lifecycleLock order RegisterWithReplay uses.
+        if (RemoveSubscriber(subscriber))
         {
             subscriber.Channel.Writer.TryComplete(new SessionManagerException(
                 SessionManagerErrorCode.EventQueueOverflow,
@@ -699,12 +737,36 @@ public sealed class SessionEventDistributor : IAsyncDisposable
 
     private void Unregister(Subscriber subscriber)
     {
-        if (_subscribers.TryRemove(subscriber.Id, out _))
+        if (RemoveSubscriber(subscriber))
         {
             subscriber.Channel.Writer.TryComplete();
         }
     }
 
+    // Removes a subscriber from the fan-out set and republishes the copy-on-write snapshot.
+    // Returns true only for the caller that actually removed it, so the channel is completed
+    // exactly once however many disposal/overflow paths race. Completing the channel is left to
+    // that caller and happens OUTSIDE the lock: this lock guards set membership only.
+    private bool RemoveSubscriber(Subscriber subscriber)
+    {
+        lock (_lifecycleLock)
+        {
+            if (!_subscribers.TryRemove(subscriber.Id, out _))
+            {
+                return false;
+            }
+
+            RebuildSubscriberSnapshot();
+            return true;
+        }
+    }
+
+    // Republishes the fan-out array from the current dictionary contents. MUST be called with
+    // _lifecycleLock held — holding that lock across the dictionary mutation and this rebuild is
+    // what keeps the array and the dictionary from diverging.
+    private void RebuildSubscriberSnapshot()
+        => Volatile.Write(ref _subscriberSnapshot, [.. _subscribers.Values]);
+
     /// <summary>
     ///     Returns the retained events with <see cref="MxEvent.WorkerSequence"/> strictly
     ///     greater than <paramref name="afterSequence"/>, in ascending sequence order, so a
@@ -791,7 +853,16 @@ public sealed class SessionEventDistributor : IAsyncDisposable
         }
     }
 
-    private void AppendToReplayBuffer(MxEvent mxEvent)
+    // Appends an event to the replay ring AND captures the fan-out array the pump will write it
+    // to, in ONE _replayLock section. The capture must share the append's critical section, not
+    // follow it: RegisterWithReplay snapshots the ring and registers under the same _replayLock,
+    // so mutual exclusion is what puts each event strictly before or strictly after a resume —
+    // "replayed, not fanned" or "fanned, not replayed", never neither. Capturing after the lock
+    // released would let a resume interleave between the append and the capture, replaying
+    // nothing for the event and fanning it to a stale array that omits the new subscriber: a
+    // silently dropped event. Returns the array; the pump fans OUTSIDE the lock so a slow
+    // reader can never stall replay.
+    private Subscriber[] AppendToReplayBufferAndCaptureSubscribers(MxEvent mxEvent)
     {
         lock (_replayLock)
         {
@@ -805,7 +876,7 @@ public sealed class SessionEventDistributor : IAsyncDisposable
             // can still report a gap) but keep no events.
             if (_replayBufferCapacity == 0)
             {
-                return;
+                return Volatile.Read(ref _subscriberSnapshot);
             }
 
             // Append at the logical tail. When the ring is full the oldest entry is
@@ -824,6 +895,7 @@ public sealed class SessionEventDistributor : IAsyncDisposable
             }
 
             EvictAged();
+            return Volatile.Read(ref _subscriberSnapshot);
         }
     }
 

src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionEventDistributorTests.cs

@@ -1057,6 +1057,63 @@ public sealed class SessionEventDistributorTests
         Assert.False(lateCts.IsCancellationRequested);
     }
 
+    /// <summary>
+    ///     Guards the copy-on-write fan-out snapshot: registrations and unregistrations churn on
+    ///     another thread while the pump is actively fanning events, and the stable subscriber
+    ///     must still receive every event exactly once and in order. The pump captures the
+    ///     subscriber array once per event instead of enumerating the dictionary, so a mutation
+    ///     racing the fan-out must never drop, duplicate, or reorder an event for a subscriber
+    ///     registered throughout — nor leave the array and the dictionary disagreeing on the
+    ///     subscriber count once the churn stops.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task RegistrationChurnDuringFanOut_StableSubscriberStillReceivesEveryEventInOrder()
+    {
+        // Below the 64-event per-subscriber queue capacity, so the stable subscriber cannot
+        // overflow and be disconnected while the writes race the churn — the assertion stays
+        // deterministic no matter how the threads interleave.
+        const int EventCount = 50;
+
+        Channel<MxEvent> source = Channel.CreateUnbounded<MxEvent>();
+        await using SessionEventDistributor distributor = CreateDistributor(source.Reader);
+        await distributor.StartAsync(CancellationToken.None);
+
+        using IEventSubscriberLease stable = distributor.Register();
+
+        using CancellationTokenSource churnCts = new();
+        Task churn = Task.Run(async () =>
+        {
+            while (!churnCts.IsCancellationRequested)
+            {
+                // Register then immediately unregister: every iteration rebuilds the fan-out
+                // snapshot twice, maximizing the chance of landing inside a fan-out pass.
+                distributor.Register().Dispose();
+                await Task.Yield();
+            }
+        });
+
+        for (ulong sequence = 1; sequence <= EventCount; sequence++)
+        {
+            source.Writer.TryWrite(Event(sequence));
+        }
+
+        List<ulong> received = [];
+        for (int i = 0; i < EventCount; i++)
+        {
+            received.Add((await ReadOneAsync(stable.Reader)).WorkerSequence);
+        }
+
+        await churnCts.CancelAsync();
+        await churn.WaitAsync(ReadTimeout);
+
+        Assert.Equal(Enumerable.Range(1, EventCount).Select(sequence => (ulong)sequence), received);
+
+        // Only the stable subscriber remains: the snapshot the count is read from tracked every
+        // add and remove the churn performed.
+        Assert.Equal(1, distributor.SubscriberCount);
+    }
+
     private static SessionEventDistributor CreateDistributor(ChannelReader<MxEvent> source)
         => CreateDistributor(source, replayBufferCapacity: 1024, replayRetentionSeconds: 300);