mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Dashboard/DashboardSnapshotFeed.cs

using System.Runtime.CompilerServices;
using System.Threading.Channels;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;

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

/// <summary>
/// Fans one <see cref="IDashboardSnapshotService.WatchSnapshotsAsync"/> enumeration out to
/// every dashboard circuit. The underlying watch is not multicast — each enumeration owns a
/// timer and builds its own snapshot per tick — so subscribing per page would multiply the
/// snapshot cost by the number of open pages.
/// </summary>
/// <remarks>
/// <para>
/// The pump is idle-gated: it starts when the first subscriber arrives and is cancelled and
/// awaited when the last one leaves, so an unwatched gateway runs no timer and builds no
/// snapshots. Successive pumps are chained through <c>_pumpTask</c>, so a rapid
/// unsubscribe/resubscribe restarts a fresh pump without ever running two enumerations at
/// once.
/// </para>
/// <para>
/// Every subscriber is tagged with the pump generation it joined under, and a dying pump only
/// ever detaches its own generation. A pump ends its generation the instant its source fails
/// or completes — before the (possibly slow) enumerator disposal — so a subscriber arriving
/// while a pump unwinds starts a fresh generation instead of silently attaching to a dead
/// pump that is about to detach everybody and leave nobody watching.
/// </para>
/// </remarks>
public sealed class DashboardSnapshotFeed : IDashboardSnapshotFeed
{
    /// <summary>Generation value meaning "no pump is accepting subscribers".</summary>
    private const long NoGeneration = 0;

    private readonly IDashboardSnapshotService _snapshotService;
    private readonly ILogger<DashboardSnapshotFeed> _logger;
    private readonly object _gate = new();
    private readonly List<Subscription> _subscribers = [];

    /// <summary>
    /// The most recent pump, completed while idle. A starting pump awaits its predecessor
    /// before enumerating, which is what guarantees a single live enumeration.
    /// </summary>
    private Task _pumpTask = Task.CompletedTask;

    /// <summary>Cancellation for the live pump; null when no generation is accepting subscribers.</summary>
    private CancellationTokenSource? _pumpCancellation;

    /// <summary>The generation new subscribers join, or <see cref="NoGeneration"/> when no pump is live.</summary>
    private long _generation = NoGeneration;

    /// <summary>Last generation handed out; only ever incremented under <c>_gate</c>.</summary>
    private long _lastGeneration = NoGeneration;

    /// <summary>Initializes a new instance of the <see cref="DashboardSnapshotFeed"/> class.</summary>
    /// <param name="snapshotService">Snapshot source to multicast.</param>
    /// <param name="logger">Optional logger for pump faults.</param>
    public DashboardSnapshotFeed(
        IDashboardSnapshotService snapshotService,
        ILogger<DashboardSnapshotFeed>? logger = null)
    {
        _snapshotService = snapshotService ?? throw new ArgumentNullException(nameof(snapshotService));
        _logger = logger ?? NullLogger<DashboardSnapshotFeed>.Instance;
    }

    /// <inheritdoc />
    public async IAsyncEnumerable<DashboardSnapshot> WatchAsync(
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        // Capacity 1 + DropOldest: a viewer only ever wants the latest snapshot, so a
        // circuit that renders slowly neither buffers without bound nor blocks the pump
        // (TryWrite always succeeds) — it just skips the snapshots it was too slow for.
        Channel<DashboardSnapshot> channel = Channel.CreateBounded<DashboardSnapshot>(
            new BoundedChannelOptions(1)
            {
                FullMode = BoundedChannelFullMode.DropOldest,
                SingleReader = true,
                SingleWriter = false,
            });

        Subscription subscription = Subscribe(channel);
        try
        {
            await foreach (DashboardSnapshot snapshot in channel.Reader
                .ReadAllAsync(cancellationToken)
                .ConfigureAwait(false))
            {
                yield return snapshot;
            }
        }
        finally
        {
            // Untokened on purpose: teardown must run to completion even when this
            // subscriber is unwinding because its own token fired.
            await UnsubscribeAsync(subscription).ConfigureAwait(false);
        }
    }

    private Subscription Subscribe(Channel<DashboardSnapshot> channel)
    {
        lock (_gate)
        {
            // A live generation is joined; otherwise this subscriber starts one. Keying on
            // "is a generation live" rather than "is this the first subscriber" is what makes
            // a subscriber arriving while a pump unwinds start a fresh pump for itself.
            long generation = _pumpCancellation is null ? StartPumpLocked() : _generation;
            Subscription subscription = new(channel, generation);
            _subscribers.Add(subscription);
            return subscription;
        }
    }

    private async Task UnsubscribeAsync(Subscription subscription)
    {
        CancellationTokenSource? cancellation;
        Task pump;
        lock (_gate)
        {
            if (!_subscribers.Remove(subscription) || _subscribers.Count != 0)
            {
                // Either the pump already detached this subscription (it completed or
                // faulted), or other viewers are still watching.
                return;
            }

            cancellation = _pumpCancellation;
            _pumpCancellation = null;
            _generation = NoGeneration;
            pump = _pumpTask;
        }

        try
        {
            cancellation?.Cancel();
        }
        catch (ObjectDisposedException)
        {
            // The pump ended on its own and disposed its cancellation source first.
        }

        try
        {
            await pump.ConfigureAwait(false);
        }
        catch (Exception)
        {
            // A pump fault has already been reported to the subscribers it had; the
            // unsubscribing caller is only waiting for the enumeration to stop.
        }
    }

    /// <summary>
    /// Starts a pump generation. Must be called while holding <c>_gate</c>; the caller adds
    /// the subscribers that belong to the returned generation.
    /// </summary>
    /// <returns>The new generation identifier.</returns>
    private long StartPumpLocked()
    {
        long generation = ++_lastGeneration;
        CancellationTokenSource cancellation = new();
        Task previous = _pumpTask;
        _generation = generation;
        _pumpCancellation = cancellation;

        // Task.Run, not a direct call: an async iterator runs synchronously up to its
        // first suspension, and the first pull of the underlying watch can read the API
        // key table. That must not run on the subscribing circuit's thread, let alone
        // while this lock is held.
        _pumpTask = Task.Run(() => PumpAsync(generation, previous, cancellation, cancellation.Token));
        return generation;
    }

    private async Task PumpAsync(
        long generation,
        Task previous,
        CancellationTokenSource cancellation,
        CancellationToken cancellationToken)
    {
        try
        {
            // Never overlap with the enumeration this pump replaces.
            await previous.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

            // Enumerated by hand rather than with await foreach so the generation can be
            // ended the moment the source fails or completes — await foreach would run the
            // enumerator's disposal first, and a subscriber arriving during that disposal
            // would join a generation that is already doomed.
            IAsyncEnumerator<DashboardSnapshot> snapshots = _snapshotService
                .WatchSnapshotsAsync(cancellationToken)
                .GetAsyncEnumerator(cancellationToken);
            try
            {
                while (true)
                {
                    bool moved;
                    try
                    {
                        moved = await snapshots.MoveNextAsync().ConfigureAwait(false);
                    }
                    catch
                    {
                        EndGeneration(generation);
                        throw;
                    }

                    if (!moved)
                    {
                        EndGeneration(generation);
                        break;
                    }

                    Broadcast(generation, snapshots.Current);
                }
            }
            finally
            {
                await snapshots.DisposeAsync().ConfigureAwait(false);
            }

            // The source completed on its own; hand the completion to this generation's
            // subscribers and re-arm so the next one starts a fresh enumeration.
            Reset(generation, error: null);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            // The production DashboardSnapshotService swallows cancellation and yield-breaks,
            // so normal teardown exits through the fall-through above (with an ownership-checked
            // Reset that finds no subscribers); an implementation that propagates the token
            // instead exits here. Both shapes end the generation exactly once.
            EndGeneration(generation);
        }
        catch (Exception error)
        {
            _logger.LogWarning(error, "Dashboard snapshot feed stopped; the next subscriber restarts it.");
            Reset(generation, error);
        }
        finally
        {
            cancellation.Dispose();
        }
    }

    private void Broadcast(long generation, DashboardSnapshot snapshot)
    {
        lock (_gate)
        {
            foreach (Subscription subscriber in _subscribers)
            {
                if (subscriber.Generation != generation)
                {
                    continue;
                }

                // Bounded/DropOldest: always accepted unless the channel is completed.
                subscriber.Channel.Writer.TryWrite(snapshot);
            }
        }
    }

    /// <summary>
    /// Stops handing <paramref name="generation"/> to new subscribers. Called the instant a
    /// pump's source fails or completes, before its enumerator is disposed.
    /// </summary>
    /// <param name="generation">The generation that has ended.</param>
    private void EndGeneration(long generation)
    {
        lock (_gate)
        {
            EndGenerationLocked(generation);
        }
    }

    /// <summary>Clears the live-pump state if <paramref name="generation"/> still owns it.</summary>
    /// <param name="generation">The generation that has ended.</param>
    private void EndGenerationLocked(long generation)
    {
        if (_generation != generation)
        {
            return;
        }

        _generation = NoGeneration;
        _pumpCancellation = null;
    }

    /// <summary>
    /// Detaches the subscribers of a finished generation and re-arms the feed. Subscribers of
    /// any other generation are left alone — they belong to a pump that is still running (or
    /// about to), so a dying pump must not take them down with it.
    /// </summary>
    /// <param name="generation">The generation whose subscribers are being detached.</param>
    /// <param name="error">Failure to surface, or null when the source completed cleanly.</param>
    private void Reset(long generation, Exception? error)
    {
        List<Channel<DashboardSnapshot>> detached = [];
        lock (_gate)
        {
            for (int index = _subscribers.Count - 1; index >= 0; index--)
            {
                if (_subscribers[index].Generation != generation)
                {
                    continue;
                }

                detached.Add(_subscribers[index].Channel);
                _subscribers.RemoveAt(index);
            }

            EndGenerationLocked(generation);

            if (_subscribers.Count > 0 && _pumpCancellation is null)
            {
                // Belt and braces: subscribers left with no live pump would be frozen for
                // good, because only a subscriber that finds no generation starts one.
                long restarted = StartPumpLocked();
                foreach (Subscription subscriber in _subscribers)
                {
                    subscriber.Generation = restarted;
                }
            }
        }

        foreach (Channel<DashboardSnapshot> channel in detached)
        {
            channel.Writer.TryComplete(error);
        }
    }

    /// <summary>One viewer's delivery channel plus the pump generation serving it.</summary>
    /// <param name="channel">Delivery channel for this viewer.</param>
    /// <param name="generation">Pump generation this viewer joined under.</param>
    private sealed class Subscription(Channel<DashboardSnapshot> channel, long generation)
    {
        /// <summary>Gets the viewer's delivery channel.</summary>
        public Channel<DashboardSnapshot> Channel { get; } = channel;

        /// <summary>Gets or sets the pump generation currently serving this viewer.</summary>
        public long Generation { get; set; } = generation;
    }
}