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

using System.Collections.Concurrent;
using System.Threading.Channels;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Server.Sessions;

/// <summary>
///     Invoked by the pump (on the pump thread) when a subscriber's bounded channel is full
///     and the event cannot be written. The handler applies policy side-effects only:
///     it records the overflow metric and, in the legacy single-subscriber FailFast case,
///     faults the owning session. The handler MUST NOT complete the subscriber's channel —
///     the distributor performs the disconnect and channel-completion unconditionally,
///     regardless of what the handler does.
/// </summary>
/// <param name="isOnlySubscriber">
///     <see langword="true"/> when FailFast is allowed to fault the whole session for this
///     overflow. This is gated on the SESSION MODE, not a live count: it is
///     <see langword="true"/> only for an external subscriber in single-subscriber mode
///     (<c>AllowMultipleEventSubscribers == false</c>), where at most one external subscriber
///     can ever exist. In multi-subscriber mode it is always <see langword="false"/>, so
///     FailFast degrades to a per-subscriber disconnect and one slow consumer never faults a
///     session shared by others; gating on the fixed mode also removes the race where a
///     concurrent registration could make a count snapshot falsely report a sole subscriber.
///     Always <see langword="false"/> for internal subscribers (the dashboard mirror) so a
///     slow/broken dashboard can never fault the session.
/// </param>
/// <param name="isInternal">
///     <see langword="true"/> when the overflowing subscriber is the gateway-owned internal
///     dashboard mirror subscriber. The handler uses this to choose the correct metric label
///     (<c>"dashboard-mirror"</c> vs <c>"grpc-event-stream"</c>).
/// </param>
public delegate void SubscriberOverflowHandler(bool isOnlySubscriber, bool isInternal);

/// <summary>
///     Per-session event pump and fan-out. A single background task drains the
///     session's event source <em>exactly once</em> and fans each event out to
///     every currently-registered subscriber's own bounded channel.
/// </summary>
/// <remarks>
///     <para>
///         The bounded replay ring buffer is wired into <c>GatewaySession</c> and
///         <c>EventStreamService</c>. The per-subscriber backpressure-isolation policy is
///         implemented here: a slow subscriber overflows only its own bounded channel and the
///         pump applies the policy to that subscriber alone (see
///         <see cref="SubscriberOverflowHandler"/> and <c>OnSubscriberOverflow</c>), leaving
///         the pump, the session, and other subscribers running. The FailFast-faults-session
///         decision is mode-gated: it fires only in single-subscriber mode
///         (<c>singleSubscriberMode</c>), so multi-subscriber FailFast always degrades to
///         a per-subscriber disconnect — see <c>OnSubscriberOverflow</c>. The ring buffer supports capacity
///         eviction (oldest entry dropped when the count exceeds
///         <c>replayBufferCapacity</c>) and age eviction (entries older than
///         <c>replayRetentionSeconds</c> dropped on the next append or query), and is
///         queried via <see cref="TryGetReplayFrom"/> by reconnecting subscribers.
///     </para>
///     <para>
///         <b>Source seam.</b> The event source is injected as a
///         <see cref="Func{T, TResult}"/> producing an
///         <see cref="IAsyncEnumerable{T}"/> of already-mapped public
///         <see cref="MxEvent"/>s, given a <see cref="CancellationToken"/>. This is the
///         cleanest seam: it can pass
///         <c>ct =&gt; session.ReadEventsAsync(ct).Select(mapper.MapEvent)</c> (or a
///         channel reader's <c>ReadAllAsync</c>), while unit tests pass a plain
///         channel reader's <c>ReadAllAsync</c> with no real session. The pump owns the
///         single consumption of this enumerable; fan-out happens on the public
///         <see cref="MxEvent"/> after mapping, mirroring today's
///         <c>EventStreamService.ProduceEventsAsync</c> ordering.
///     </para>
///     <para>
///         <b>Concurrency.</b> The subscriber set is a
///         <see cref="ConcurrentDictionary{TKey, TValue}"/> keyed by a monotonic id.
///         The pump iterates it with a snapshot-free enumerator (which never throws on
///         concurrent add/remove), and <see cref="Register"/> / lease disposal mutate it
///         without any lock held across an <c>await</c>. Each subscriber channel has a
///         single writer — the pump — so per-channel writes never race. MXAccess parity:
///         events are fanned in the order received; the pump never reorders or
///         synthesizes events.
///     </para>
/// </remarks>
public sealed class SessionEventDistributor : IAsyncDisposable
{
    /// <summary>
    ///     Bounded wait for the pump to stop during disposal. A source factory that
    ///     ignores cancellation must not hang dispose forever; after this window the
    ///     pump is abandoned and subscribers are completed anyway.
    /// </summary>
    private static readonly TimeSpan DefaultShutdownTimeout = TimeSpan.FromSeconds(5);

    private readonly string _sessionId;
    private readonly Func<CancellationToken, IAsyncEnumerable<MxEvent>> _eventSourceFactory;
    private readonly int _subscriberQueueCapacity;
    private readonly bool _singleSubscriberMode;
    private readonly SubscriberOverflowHandler? _overflowHandler;
    private readonly TimeSpan _shutdownTimeout;
    private readonly ILogger<SessionEventDistributor> _logger;
    private readonly TimeProvider _timeProvider;
    private readonly ConcurrentDictionary<long, Subscriber> _subscribers = new();
    private readonly CancellationTokenSource _shutdownCts = new();
    private readonly object _lifecycleLock = new();

    // 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
    // event allocates no node (the LinkedList this replaced allocated one node per
    // event on the fan-out hot path). Events are kept in ascending WorkerSequence order
    // (the pump fans in source order): _replayHead is the logical front (oldest retained
    // event) and _replayCount entries follow it, wrapping modulo the array length. The
    // logical entry at position i is _replayBuffer[(_replayHead + i) % Length]. Capacity
    // == 0 disables retention (the array is empty and never indexed); RetentionSeconds
    // <= 0 disables age-based eviction.
    private readonly int _replayBufferCapacity;
    private readonly TimeSpan _replayRetention;
    private readonly bool _ageEvictionEnabled;
    private readonly ReplayEntry[] _replayBuffer;
    private int _replayHead;
    private int _replayCount;
    private readonly object _replayLock = new();
    private bool _anyEventSeen;
    private ulong _highestSequenceSeen;

    private long _nextSubscriberId;
    private Task? _pumpTask;
    private bool _started;
    private bool _disposed;

    // Set once the pump has run its final CompleteAllSubscribers sweep — the event source
    // completed or faulted and the pump exited. Guarded by _lifecycleLock together with the
    // subscriber add. A subscriber that registers AFTER this point but BEFORE DisposeAsync
    // (the source ended but the session is not yet torn down) would otherwise be added with a
    // channel the now-exited pump never completes, hanging its reader forever. The register
    // paths complete such a late registrant's channel immediately with the same terminal
    // state. _completionError carries the terminal exception (source fault) or null (graceful
    // source completion), mirroring what the final CompleteAllSubscribers passed.
    private bool _completed;
    private Exception? _completionError;

    /// <summary>
    /// Initializes a per-session event distributor.
    /// </summary>
    /// <param name="sessionId">Owning session id, used only for logging context.</param>
    /// <param name="eventSourceFactory">
    ///     Factory producing the session's event stream given a cancellation token.
    ///     The pump consumes this exactly once. See the type remarks for the seam this
    ///     plugs into.
    /// </param>
    /// <param name="subscriberQueueCapacity">
    ///     Bounded capacity of each per-subscriber channel. Mirrors the gRPC event-stream
    ///     queue capacity shape used today.
    /// </param>
    /// <param name="logger">Logger for pump lifecycle diagnostics.</param>
    /// <param name="overflowHandler">
    ///     Optional per-subscriber backpressure handler invoked when a subscriber's bounded
    ///     channel is full. See the primary constructor overload for the full contract.
    /// </param>
    /// <param name="singleSubscriberMode">
    ///     <see langword="true"/> when the owning session is in single-subscriber mode. See
    ///     the primary constructor overload for the full contract.
    /// </param>
    /// <remarks>
    ///     This overload disables the replay ring buffer (capacity 0). Use the overload
    ///     taking replay parameters to retain events for reconnect/reattach replay.
    ///     Kept <c>internal</c> so production wiring cannot accidentally use
    ///     the no-replay path; tests reach it via <c>InternalsVisibleTo</c>.
    /// </remarks>
    internal SessionEventDistributor(
        string sessionId,
        Func<CancellationToken, IAsyncEnumerable<MxEvent>> eventSourceFactory,
        int subscriberQueueCapacity,
        ILogger<SessionEventDistributor> logger,
        SubscriberOverflowHandler? overflowHandler = null,
        bool singleSubscriberMode = true)
        : this(
            sessionId,
            eventSourceFactory,
            subscriberQueueCapacity,
            replayBufferCapacity: 0,
            replayRetentionSeconds: 0,
            logger,
            TimeProvider.System,
            overflowHandler,
            singleSubscriberMode)
    {
    }

    /// <summary>
    /// Initializes a per-session event distributor with a bounded replay ring buffer.
    /// </summary>
    /// <param name="sessionId">Owning session id, used only for logging context.</param>
    /// <param name="eventSourceFactory">
    ///     Factory producing the session's event stream given a cancellation token.
    ///     The pump consumes this exactly once. See the type remarks for the seam this
    ///     plugs into.
    /// </param>
    /// <param name="subscriberQueueCapacity">
    ///     Bounded capacity of each per-subscriber channel. Mirrors the gRPC event-stream
    ///     queue capacity shape used today.
    /// </param>
    /// <param name="replayBufferCapacity">
    ///     Maximum number of events retained for replay. The oldest retained event is
    ///     evicted once this count is exceeded. <c>0</c> disables retention entirely.
    /// </param>
    /// <param name="replayRetentionSeconds">
    ///     Maximum age, in seconds, of a retained event. Entries older than this are
    ///     evicted regardless of capacity. <c>0</c> (or less) disables age-based eviction.
    /// </param>
    /// <param name="logger">Logger for pump lifecycle diagnostics.</param>
    /// <param name="timeProvider">
    ///     Clock used to timestamp and age-evict replay entries. Inject a fake to make
    ///     age-eviction deterministic in tests.
    /// </param>
    /// <param name="overflowHandler">
    ///     Optional per-subscriber backpressure handler invoked when a subscriber's bounded
    ///     channel is full. It records the overflow metric and, for the legacy
    ///     single-subscriber FailFast case, faults the owning session. The distributor always
    ///     disconnects the offending subscriber with an overflow fault regardless of the
    ///     handler. When <see langword="null"/> (unit/skeleton use) the offending subscriber is
    ///     still disconnected but no metric/fault side effect runs.
    /// </param>
    /// <param name="singleSubscriberMode">
    ///     <see langword="true"/> when the owning session is in single-subscriber mode
    ///     (<c>AllowMultipleEventSubscribers == false</c>). This gates the FailFast
    ///     session-fault decision in <c>OnSubscriberOverflow</c>: an external subscriber that
    ///     overflows reports <c>isOnlySubscriber == true</c> (legacy FailFast faults the
    ///     session) ONLY in single-subscriber mode. In multi-subscriber mode it is always
    ///     <see langword="false"/>, so FailFast degrades to a per-subscriber disconnect and a
    ///     transient registration race can never falsely fault a shared session. Defaults to
    ///     <see langword="true"/> so existing call sites and unit tests keep legacy
    ///     single-subscriber FailFast behavior.
    /// </param>
    public SessionEventDistributor(
        string sessionId,
        Func<CancellationToken, IAsyncEnumerable<MxEvent>> eventSourceFactory,
        int subscriberQueueCapacity,
        int replayBufferCapacity,
        double replayRetentionSeconds,
        ILogger<SessionEventDistributor> logger,
        TimeProvider timeProvider,
        SubscriberOverflowHandler? overflowHandler = null,
        bool singleSubscriberMode = true)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);
        ArgumentNullException.ThrowIfNull(eventSourceFactory);
        ArgumentOutOfRangeException.ThrowIfLessThan(subscriberQueueCapacity, 1);
        ArgumentOutOfRangeException.ThrowIfNegative(replayBufferCapacity);
        ArgumentOutOfRangeException.ThrowIfNegative(replayRetentionSeconds);
        ArgumentNullException.ThrowIfNull(logger);
        ArgumentNullException.ThrowIfNull(timeProvider);

        _sessionId = sessionId;
        _eventSourceFactory = eventSourceFactory;
        _subscriberQueueCapacity = subscriberQueueCapacity;
        _singleSubscriberMode = singleSubscriberMode;
        _overflowHandler = overflowHandler;
        _shutdownTimeout = DefaultShutdownTimeout;
        _replayBufferCapacity = replayBufferCapacity;
        _replayBuffer = new ReplayEntry[replayBufferCapacity];
        _ageEvictionEnabled = replayRetentionSeconds > 0;
        _replayRetention = _ageEvictionEnabled
            ? TimeSpan.FromSeconds(replayRetentionSeconds)
            : TimeSpan.Zero;
        _logger = logger;
        _timeProvider = timeProvider;
    }

    /// <summary>
    ///     Gets the count of currently-registered subscribers. This count INCLUDES internal
    ///     subscribers (e.g. the gateway-owned dashboard mirror registered via
    ///     <c>Register(isInternal: true)</c>), and therefore differs from
    ///     <see cref="GatewaySession.ActiveEventSubscriberCount"/>, which tracks only external
    ///     (gRPC) subscribers and excludes the internal dashboard subscriber.
    /// </summary>
    public int SubscriberCount => _subscribers.Count;

    /// <summary>
    /// Starts the background pump. Idempotent — a second call is a no-op.
    /// </summary>
    /// <param name="cancellationToken">Token observed only while starting.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public Task StartAsync(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();

        lock (_lifecycleLock)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            if (_started)
            {
                return Task.CompletedTask;
            }

            _started = true;
            _pumpTask = Task.Run(() => PumpAsync(_shutdownCts.Token), CancellationToken.None);
        }

        return Task.CompletedTask;
    }

    /// <summary>
    ///     Registers a new subscriber and returns its lease. The lease exposes the
    ///     subscriber's <see cref="ChannelReader{T}"/> and, when disposed, unregisters the
    ///     subscriber and completes its channel without disturbing the pump or other
    ///     subscribers.
    /// </summary>
    /// <param name="isInternal">
    ///     <see langword="true"/> for a gateway-owned internal subscriber (the
    ///     session's dashboard mirror) that must NOT participate in the single-subscriber
    ///     overflow accounting. An internal subscriber is excluded from the
    ///     <c>isOnlySubscriber</c> count, so a lone external gRPC subscriber still reports
    ///     <c>isOnlySubscriber == true</c> (preserving legacy FailFast session-fault
    ///     behavior) even while the dashboard subscriber is attached; and an internal
    ///     subscriber that itself overflows always reports <c>isOnlySubscriber == false</c>,
    ///     so a slow/broken dashboard can never fault the session — it is merely
    ///     disconnected from the mirror. Defaults to <see langword="false"/> (external
    ///     subscriber) so every existing call site is unchanged.
    /// </param>
    /// <returns>The lease for the newly-registered subscriber.</returns>
    public IEventSubscriberLease Register(bool isInternal = false)
    {
        Channel<MxEvent> channel = CreateSubscriberChannel();
        long id = Interlocked.Increment(ref _nextSubscriberId);
        Subscriber subscriber = new(id, channel, isInternal);
        return RegisterSubscriber(subscriber);
    }

    private IEventSubscriberLease RegisterSubscriber(Subscriber subscriber)
    {
        // The disposed check AND the map add happen under the same lock with no await
        // in between. DisposeAsync sets _disposed=true under this same lock before it
        // calls CompleteAllSubscribers, so once disposal has begun no further subscriber
        // can be added — closing the Register-after-DisposeAsync window that would
        // otherwise leave a subscriber's channel never completed.
        lock (_lifecycleLock)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            _subscribers[subscriber.Id] = subscriber;

            // Close the register-after-pump-completion window: if the pump already ran its
            // final CompleteAllSubscribers (source completed/faulted) but the distributor is
            // not yet disposed, no further completion sweep will run, so complete this late
            // registrant's channel now with the same terminal state instead of leaving its
            // reader hanging.
            if (_completed)
            {
                subscriber.Channel.Writer.TryComplete(_completionError);
            }
        }

        return new SubscriberLease(this, subscriber);
    }

    // Creates a per-subscriber bounded channel. The pump is the single writer; readers are
    // single-consumer (one gRPC stream / dashboard subscriber). Synchronous continuations are
    // disabled so a slow reader can never stall the pump on its completion.
    //
    // The pump MUST stay non-blocking: it writes with the non-blocking TryWrite so one slow
    // reader can never stall the single pump that feeds every subscriber. FullMode is
    // deliberately Wait — NOT because the pump ever blocks (it never calls the blocking
    // WriteAsync overload), but because Wait is the only BoundedChannelFullMode under which
    // TryWrite returns false when the channel is full. That false return IS the overflow signal
    // the pump needs to apply the per-subscriber backpressure policy. The Drop* modes would
    // make TryWrite silently succeed-and-drop, hiding overflow and re-introducing silent data
    // loss. So: Wait mode + TryWrite = a non-blocking pump that still detects a full channel.
    private Channel<MxEvent> CreateSubscriberChannel()
        => Channel.CreateBounded<MxEvent>(
            new BoundedChannelOptions(_subscriberQueueCapacity)
            {
                SingleReader = true,
                SingleWriter = true,
                FullMode = BoundedChannelFullMode.Wait,
                AllowSynchronousContinuations = false,
            });

    /// <summary>
    ///     Atomically snapshots the replay ring for events newer than
    ///     <paramref name="afterSequence"/> AND registers a live subscriber, so the
    ///     replay→live handoff has no gap and no duplicate (reconnect/resume).
    /// </summary>
    /// <param name="afterSequence">
    ///     The last worker sequence the reconnecting client already observed. Replay returns
    ///     events strictly newer than this; the live channel is filtered (by the caller) to
    ///     events strictly newer than the last replayed sequence.
    /// </param>
    /// <param name="replayedEvents">
    ///     The retained events newer than <paramref name="afterSequence"/>, in ascending
    ///     sequence order. Never null; empty when nothing newer is retained.
    /// </param>
    /// <param name="gap">
    ///     <see langword="true"/> when events between <paramref name="afterSequence"/> and the
    ///     oldest retained event were already evicted (capacity/age), so the client missed
    ///     events that can no longer be replayed and must re-snapshot. Mirrors
    ///     <see cref="TryGetReplayFrom"/> gap semantics.
    /// </param>
    /// <param name="oldestAvailableSequence">
    ///     The oldest worker sequence still retained and replayable. <c>0</c> when nothing is
    ///     retained. Meaningful to the caller only when <paramref name="gap"/> is
    ///     <see langword="true"/> (it populates the ReplayGap sentinel's
    ///     <c>oldest_available_sequence</c>).
    /// </param>
    /// <param name="liveResumeSequence">
    ///     The worker sequence the live channel must resume strictly after: the highest
    ///     replayed sequence, or <paramref name="afterSequence"/> when nothing was replayed.
    ///     The caller MUST apply this as the per-subscriber live filter so any event that was
    ///     both replayed here and subsequently fanned into this subscriber's live channel is
    ///     dropped exactly once (no duplicate), while every newer event is delivered (no gap).
    /// </param>
    /// <param name="isInternal">
    ///     <see langword="true"/> for a gateway-owned internal subscriber. See
    ///     <see cref="Register"/>.
    /// </param>
    /// <returns>The lease for the newly-registered subscriber.</returns>
    /// <remarks>
    ///     <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>
    ///         <em>before</em> fanning it to subscribers (outside the lock). Therefore, relative
    ///         to this method's critical section, for every event E:
    ///     </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
    ///             registration: if it writes E to this new channel too, E's sequence is
    ///             <c>&lt;= liveResumeSequence</c>, so the caller's live filter DROPS it — no
    ///             duplicate.
    ///         </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
    ///             channel with sequence <c>&gt; liveResumeSequence</c> — delivered as live, no
    ///             gap.
    ///         </item>
    ///     </list>
    ///     <para>
    ///         Lock ordering: this is the only path that holds both <c>_replayLock</c> and
    ///         <c>_lifecycleLock</c>; it always takes <c>_replayLock</c> first then
    ///         <c>_lifecycleLock</c>. No other path acquires both, so there is no inversion.
    ///     </para>
    /// </remarks>
    public IEventSubscriberLease RegisterWithReplay(
        ulong afterSequence,
        out IReadOnlyList<MxEvent> replayedEvents,
        out bool gap,
        out ulong oldestAvailableSequence,
        out ulong liveResumeSequence,
        bool isInternal = false)
    {
        Channel<MxEvent> channel = CreateSubscriberChannel();
        long id = Interlocked.Increment(ref _nextSubscriberId);
        Subscriber subscriber = new(id, channel, isInternal);

        // Snapshot replay AND register under a single _replayLock section so the live channel
        // begins exactly where the replay snapshot ends — see the remarks for the no-gap /
        // no-duplicate argument. _lifecycleLock is nested inside (consistent ordering) only to
        // honor the disposed check and the same add semantics as Register.
        lock (_replayLock)
        {
            EvictAged();

            List<MxEvent> newer = [];
            ulong highestReplayed = afterSequence;

            if (_replayCount == 0)
            {
                gap = _anyEventSeen && afterSequence < _highestSequenceSeen;
                oldestAvailableSequence = 0; // meaningful only when gap == true; 0 here since nothing is retained
            }
            else
            {
                ulong oldestRetained = ReplayEntryAt(0).Event.WorkerSequence;
                gap = oldestRetained > 0 && afterSequence < oldestRetained - 1;
                // Per the contract on OldestAvailableSequence: meaningful only when gap == true.
                oldestAvailableSequence = gap ? oldestRetained : 0;

                for (int i = 0; i < _replayCount; i++)
                {
                    MxEvent retained = ReplayEntryAt(i).Event;
                    if (retained.WorkerSequence > afterSequence)
                    {
                        newer.Add(retained);
                        highestReplayed = retained.WorkerSequence;
                    }
                }
            }

            replayedEvents = newer;
            liveResumeSequence = highestReplayed;

            lock (_lifecycleLock)
            {
                ObjectDisposedException.ThrowIf(_disposed, this);
                _subscribers[id] = subscriber;

                // 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
                // above), but its live channel must be completed now since the pump is gone.
                if (_completed)
                {
                    subscriber.Channel.Writer.TryComplete(_completionError);
                }
            }
        }

        return new SubscriberLease(this, subscriber);
    }

    /// <summary>
    /// Stops the pump and completes all subscriber channels. Idempotent.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async ValueTask DisposeAsync()
    {
        Task? pumpTask;
        lock (_lifecycleLock)
        {
            if (_disposed)
            {
                return;
            }

            _disposed = true;
            pumpTask = _pumpTask;
        }

        // Signal the pump to stop. It must not block on a non-reading subscriber:
        // it writes with non-blocking TryWrite, so cancellation tears it down promptly.
        await _shutdownCts.CancelAsync().ConfigureAwait(false);

        if (pumpTask is not null)
        {
            // Bound the wait: a source factory that ignores cancellation would otherwise
            // hang dispose forever. If the pump does not stop in time we log and proceed
            // to complete subscribers anyway; DisposeAsync must not throw on this path.
            Task completed = await Task.WhenAny(pumpTask, Task.Delay(_shutdownTimeout)).ConfigureAwait(false);
            if (!ReferenceEquals(completed, pumpTask))
            {
                _logger.LogWarning(
                    "Event distributor pump did not stop within {ShutdownTimeoutSeconds}s for session {SessionId}; completing subscribers and abandoning the pump.",
                    _shutdownTimeout.TotalSeconds,
                    _sessionId);
            }
            else
            {
                try
                {
                    await pumpTask.ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                }
                catch (Exception exception)
                {
                    _logger.LogDebug(
                        exception,
                        "Event distributor pump faulted during shutdown for session {SessionId}.",
                        _sessionId);
                }
            }
        }

        CompleteAllSubscribers(error: null);
        _shutdownCts.Dispose();
    }

    private async Task PumpAsync(CancellationToken cancellationToken)
    {
        try
        {
            await foreach (MxEvent mxEvent in _eventSourceFactory(cancellationToken)
                .WithCancellation(cancellationToken)
                .ConfigureAwait(false))
            {
                // 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.
                AppendToReplayBuffer(mxEvent);

                // Enumerating a ConcurrentDictionary's Values never throws on concurrent
                // add/remove; a subscriber registered mid-iteration may miss this event,
                // which matches "late subscribers see events after they register".
                foreach (Subscriber subscriber in _subscribers.Values)
                {
                    // Non-blocking write: TryWrite never blocks the pump on a slow reader.
                    // A false return means this subscriber's bounded channel is full — the
                    // per-subscriber overflow signal. We apply the backpressure policy to
                    // THIS subscriber only; the pump, the session, and every other subscriber
                    // keep running. Logs identifiers (worker sequence, subscriber id, session)
                    // only, never the event payload or tag values.
                    if (!subscriber.Channel.Writer.TryWrite(mxEvent))
                    {
                        OnSubscriberOverflow(subscriber, mxEvent.WorkerSequence);
                    }
                }
            }

            CompleteAllSubscribers(error: null);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            // Shutdown path: DisposeAsync completes subscribers.
        }
        catch (Exception exception)
        {
            // Unexpected source fault (not the shutdown-cancellation path above) — visible
            // by default so an event stream silently dying is not lost in Debug noise.
            _logger.LogError(
                exception,
                "Event distributor source faulted for session {SessionId}.",
                _sessionId);
            CompleteAllSubscribers(exception);
        }
    }

    // Applies the per-subscriber backpressure policy when a subscriber's bounded channel is
    // full. Runs on the pump thread. The offending subscriber is ALWAYS disconnected with an
    // overflow fault and unregistered, so it can never wedge the pump again; the overflow
    // handler decides the observable side effects (overflow metric, and — for legacy
    // single-subscriber FailFast — faulting the owning session). Multi-subscriber FailFast
    // intentionally degrades to a plain disconnect (see SubscriberOverflowHandler docs): one
    // slow consumer must not fault a session shared by other healthy subscribers.
    private void OnSubscriberOverflow(Subscriber subscriber, ulong workerSequence)
    {
        // Decide whether FailFast may fault the whole session for this overflow. This is the
        // "isOnlySubscriber" signal the legacy single-subscriber FailFast path keys on.
        bool isOnlySubscriber = !subscriber.IsInternal && _singleSubscriberMode;

        _logger.LogDebug(
            "Event distributor disconnecting subscriber {SubscriberId} in session {SessionId} after queue overflow (worker sequence {WorkerSequence}).",
            subscriber.Id,
            _sessionId,
            workerSequence);

        // Observability + session-fault decision. Errors here must not stall the pump or
        // leave the subscriber attached, so the disconnect below runs regardless.
        // Pass subscriber.IsInternal so the handler can choose the correct metric label.
        try
        {
            _overflowHandler?.Invoke(isOnlySubscriber, subscriber.IsInternal);
        }
        catch (Exception exception)
        {
            _logger.LogError(
                exception,
                "Event distributor overflow handler threw for session {SessionId}; disconnecting subscriber {SubscriberId} anyway.",
                _sessionId,
                subscriber.Id);
        }

        // Disconnect ONLY this subscriber: complete its channel with the overflow fault and
        // 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 _))
        {
            subscriber.Channel.Writer.TryComplete(new SessionManagerException(
                SessionManagerErrorCode.EventQueueOverflow,
                $"Session {_sessionId} event stream queue overflowed."));
        }
    }

    private void CompleteAllSubscribers(Exception? error)
    {
        // Record the terminal state AND complete the current subscribers under _lifecycleLock
        // so this serializes with the subscriber-add in Register/RegisterWithReplay: a
        // subscriber added before this runs is in the map and completed by the loop; one that
        // races in afterward sees _completed and completes its own channel in the register
        // path. Exactly one of the two completes each subscriber. TryComplete is non-blocking
        // and (channels use AllowSynchronousContinuations=false) runs no continuation inline,
        // so holding the lock across the loop cannot stall or re-enter.
        lock (_lifecycleLock)
        {
            _completed = true;
            _completionError = error;
            foreach (Subscriber subscriber in _subscribers.Values)
            {
                subscriber.Channel.Writer.TryComplete(error);
            }
        }
    }

    private void Unregister(Subscriber subscriber)
    {
        if (_subscribers.TryRemove(subscriber.Id, out _))
        {
            subscriber.Channel.Writer.TryComplete();
        }
    }

    /// <summary>
    ///     Returns the retained events with <see cref="MxEvent.WorkerSequence"/> strictly
    ///     greater than <paramref name="afterSequence"/>, in ascending sequence order, so a
    ///     reconnecting or reattaching subscriber can replay what it missed.
    /// </summary>
    /// <param name="afterSequence">
    ///     The last worker sequence the caller already observed. Only events newer than this
    ///     are returned.
    /// </param>
    /// <param name="events">
    ///     The retained events newer than <paramref name="afterSequence"/>, in order. Never
    ///     null; empty when nothing newer is retained.
    /// </param>
    /// <param name="gap">
    ///     <see langword="true"/> when events between <paramref name="afterSequence"/> and the
    ///     oldest retained event were already evicted (by capacity or age), meaning the caller
    ///     missed events that can no longer be replayed and must re-snapshot. When
    ///     <see langword="true"/>, whatever IS still retained is still returned via
    ///     <paramref name="events"/>.
    /// </param>
    /// <returns>
    ///     Always <see langword="true"/> — the out parameters fully describe the result. The
    ///     return value exists for a fluent call shape and future extension.
    /// </returns>
    /// <remarks>
    ///     <para>Gap semantics, by buffer state:</para>
    ///     <list type="bullet">
    ///         <item>
    ///             Buffer non-empty: <paramref name="gap"/> is <see langword="true"/> iff
    ///             <paramref name="afterSequence"/> is below the oldest retained sequence minus
    ///             one (i.e. at least one event newer than <paramref name="afterSequence"/> but
    ///             older than the oldest retained was evicted). When
    ///             <paramref name="afterSequence"/> equals or exceeds the newest retained
    ///             sequence the caller is fully caught up: empty list, no gap.
    ///         </item>
    ///         <item>
    ///             Buffer empty (retention disabled, nothing seen yet, or everything evicted):
    ///             empty list, and <paramref name="gap"/> is <see langword="true"/> iff
    ///             <paramref name="afterSequence"/> is below the highest sequence ever seen —
    ///             i.e. the caller is behind but nothing is retained to replay. If no event has
    ///             ever been seen, or the caller is already at/ahead of the highest seen, there
    ///             is nothing to miss: no gap.
    ///         </item>
    ///     </list>
    /// </remarks>
    public bool TryGetReplayFrom(ulong afterSequence, out IReadOnlyList<MxEvent> events, out bool gap)
    {
        lock (_replayLock)
        {
            EvictAged();

            if (_replayCount == 0)
            {
                events = [];
                // Nothing retained. The caller missed events only if it is behind the
                // highest sequence ever seen (and we have seen at least one event).
                gap = _anyEventSeen && afterSequence < _highestSequenceSeen;
                return true;
            }

            ulong oldestRetained = ReplayEntryAt(0).Event.WorkerSequence;

            // A gap exists when at least one event newer than afterSequence was evicted,
            // i.e. afterSequence sits below the oldest-retained-minus-one boundary.
            // Written as (oldestRetained > 0 && afterSequence < oldestRetained - 1) to
            // avoid wrapping when afterSequence == ulong.MaxValue (afterSequence + 1
            // would overflow to 0, falsely reporting a gap).
            gap = oldestRetained > 0 && afterSequence < oldestRetained - 1;

            // O(n) scan over the retained buffer — acceptable because TryGetReplayFrom
            // is only called on subscriber reconnect, never on the hot fan-out path.
            List<MxEvent> newer = [];
            for (int i = 0; i < _replayCount; i++)
            {
                MxEvent retained = ReplayEntryAt(i).Event;
                if (retained.WorkerSequence > afterSequence)
                {
                    newer.Add(retained);
                }
            }

            events = newer;
            return true;
        }
    }

    private void AppendToReplayBuffer(MxEvent mxEvent)
    {
        lock (_replayLock)
        {
            _anyEventSeen = true;
            if (mxEvent.WorkerSequence > _highestSequenceSeen)
            {
                _highestSequenceSeen = mxEvent.WorkerSequence;
            }

            // Capacity 0 disables retention: track the highest-seen sequence (so replay
            // can still report a gap) but keep no events.
            if (_replayBufferCapacity == 0)
            {
                return;
            }

            // Append at the logical tail. When the ring is full the oldest entry is
            // overwritten in place (its slot becomes the new tail) and the head advances,
            // so the newest _replayBufferCapacity events are retained with no allocation.
            ReplayEntry entry = new(mxEvent, _timeProvider.GetUtcNow());
            if (_replayCount < _replayBufferCapacity)
            {
                _replayBuffer[(_replayHead + _replayCount) % _replayBufferCapacity] = entry;
                _replayCount++;
            }
            else
            {
                _replayBuffer[_replayHead] = entry;
                _replayHead = (_replayHead + 1) % _replayBufferCapacity;
            }

            EvictAged();
        }
    }

    // Returns the logical entry at position i (0 == oldest retained). Must be called
    // under _replayLock with 0 <= i < _replayCount (so the array length is nonzero).
    private ReplayEntry ReplayEntryAt(int i) => _replayBuffer[(_replayHead + i) % _replayBuffer.Length];

    // Must be called under _replayLock. Drops entries older than the retention window
    // by advancing the head past them (no dealloc; slots are reused on the next append).
    private void EvictAged()
    {
        if (!_ageEvictionEnabled || _replayCount == 0)
        {
            return;
        }

        DateTimeOffset cutoff = _timeProvider.GetUtcNow() - _replayRetention;
        while (_replayCount > 0 && _replayBuffer[_replayHead].RetainedAt < cutoff)
        {
            _replayHead = (_replayHead + 1) % _replayBuffer.Length;
            _replayCount--;
        }
    }

    private readonly record struct ReplayEntry(MxEvent Event, DateTimeOffset RetainedAt);

    private sealed class Subscriber(long id, Channel<MxEvent> channel, bool isInternal)
    {
        /// <summary>Gets the subscriber's monotonic id, assigned at registration.</summary>
        public long Id { get; } = id;

        /// <summary>Gets the subscriber's own bounded event channel, written by the pump.</summary>
        public Channel<MxEvent> Channel { get; } = channel;

        /// <summary>
        ///     Gets a value indicating whether this is the gateway-owned internal dashboard
        ///     subscriber. Excluded from the single-subscriber overflow accounting so it cannot
        ///     fault the session.
        /// </summary>
        public bool IsInternal { get; } = isInternal;
    }

    private sealed class SubscriberLease(SessionEventDistributor distributor, Subscriber subscriber)
        : IEventSubscriberLease
    {
        private int _leaseDisposed;

        /// <inheritdoc />
        public ChannelReader<MxEvent> Reader => subscriber.Channel.Reader;

        /// <summary>
        ///     Unregisters the subscriber from the distributor and completes its channel.
        ///     Safe to call more than once; only the first call takes effect.
        /// </summary>
        public void Dispose()
        {
            // Atomic check-and-set so concurrent Dispose calls unregister at most once.
            if (Interlocked.Exchange(ref _leaseDisposed, 1) == 0)
            {
                distributor.Unregister(subscriber);
            }
        }
    }
}