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>
/// The pump is idle-gated: it starts when the subscriber count goes 0 → 1 and is cancelled
/// and awaited when it goes 1 → 0, 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.
/// </remarks>
public sealed class DashboardSnapshotFeed : IDashboardSnapshotFeed
{
    private readonly IDashboardSnapshotService _snapshotService;
    private readonly ILogger<DashboardSnapshotFeed> _logger;
    private readonly object _gate = new();
    private readonly List<Channel<DashboardSnapshot>> _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 pump is running or one is being torn down.</summary>
    private CancellationTokenSource? _pumpCancellation;

    /// <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,
            });

        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(channel).ConfigureAwait(false);
        }
    }

    private void Subscribe(Channel<DashboardSnapshot> channel)
    {
        lock (_gate)
        {
            _subscribers.Add(channel);
            if (_subscribers.Count != 1)
            {
                return;
            }

            CancellationTokenSource cancellation = new();
            Task previous = _pumpTask;
            _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(previous, cancellation, cancellation.Token));
        }
    }

    private async Task UnsubscribeAsync(Channel<DashboardSnapshot> channel)
    {
        CancellationTokenSource? cancellation;
        Task pump;
        lock (_gate)
        {
            if (!_subscribers.Remove(channel) || _subscribers.Count != 0)
            {
                // Either the pump already dropped this channel (it completed or faulted
                // and reset itself), or other viewers are still watching.
                return;
            }

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

        try
        {
            cancellation?.Cancel();
        }
        catch (ObjectDisposedException)
        {
            // The pump reset itself and disposed its own 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.
        }
    }

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

            await foreach (DashboardSnapshot snapshot in _snapshotService
                .WatchSnapshotsAsync(cancellationToken)
                .ConfigureAwait(false))
            {
                Broadcast(snapshot);
            }

            // The source completed on its own; hand the completion to the subscribers
            // and re-arm so the next one starts a fresh enumeration.
            Reset(cancellation, error: null);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            // Last subscriber left: the unsubscribing caller already detached the
            // channels and cleared the pump state.
        }
        catch (Exception error)
        {
            _logger.LogWarning(error, "Dashboard snapshot feed stopped; the next subscriber restarts it.");
            Reset(cancellation, error);
        }
        finally
        {
            cancellation.Dispose();
        }
    }

    private void Broadcast(DashboardSnapshot snapshot)
    {
        lock (_gate)
        {
            foreach (Channel<DashboardSnapshot> subscriber in _subscribers)
            {
                // Bounded/DropOldest: always accepted unless the channel is completed.
                subscriber.Writer.TryWrite(snapshot);
            }
        }
    }

    /// <summary>
    /// Detaches every subscriber and clears the pump state so the next subscriber starts a
    /// new enumeration. The detached subscribers observe <paramref name="error"/> (or a
    /// clean end of stream) from their own <c>WatchAsync</c>.
    /// </summary>
    /// <param name="cancellation">The calling pump's cancellation source, used as its ownership token.</param>
    /// <param name="error">Failure to surface, or null when the source completed cleanly.</param>
    private void Reset(CancellationTokenSource cancellation, Exception? error)
    {
        Channel<DashboardSnapshot>[] detached;
        lock (_gate)
        {
            if (!ReferenceEquals(_pumpCancellation, cancellation))
            {
                // A newer pump (or an in-flight teardown) owns the state now: its
                // subscribers must not be detached by this pump's exit.
                return;
            }

            detached = _subscribers.ToArray();
            _subscribers.Clear();
            _pumpCancellation = null;
        }

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