git.stella-ops.org/src/StellaOps.Scanner.WebService/Services/ScanProgressStream.cs

using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Runtime.CompilerServices;
using System.Threading.Channels;
using StellaOps.Scanner.WebService.Domain;

namespace StellaOps.Scanner.WebService.Services;

public interface IScanProgressPublisher
{
    ScanProgressEvent Publish(
        ScanId scanId,
        string state,
        string? message = null,
        IReadOnlyDictionary<string, object?>? data = null,
        string? correlationId = null);
}

public interface IScanProgressReader
{
    bool Exists(ScanId scanId);

    IAsyncEnumerable<ScanProgressEvent> SubscribeAsync(ScanId scanId, CancellationToken cancellationToken);
}

public sealed class ScanProgressStream : IScanProgressPublisher, IScanProgressReader
{
    private sealed class ProgressChannel
    {
        private readonly List<ScanProgressEvent> history = new();
        private readonly Channel<ScanProgressEvent> channel = Channel.CreateUnbounded<ScanProgressEvent>(new UnboundedChannelOptions
        {
            AllowSynchronousContinuations = true,
            SingleReader = false,
            SingleWriter = false
        });

        public int Sequence { get; private set; }

        public ScanProgressEvent Append(ScanProgressEvent progressEvent)
        {
            history.Add(progressEvent);
            channel.Writer.TryWrite(progressEvent);
            return progressEvent;
        }

        public IReadOnlyList<ScanProgressEvent> Snapshot()
        {
            return history.Count == 0
                ? Array.Empty<ScanProgressEvent>()
                : history.ToArray();
        }

        public ChannelReader<ScanProgressEvent> Reader => channel.Reader;

        public int NextSequence() => ++Sequence;
    }

    private static readonly IReadOnlyDictionary<string, object?> EmptyData =
        new ReadOnlyDictionary<string, object?>(new SortedDictionary<string, object?>(StringComparer.OrdinalIgnoreCase));

    private readonly ConcurrentDictionary<string, ProgressChannel> channels = new(StringComparer.OrdinalIgnoreCase);
    private readonly TimeProvider timeProvider;

    public ScanProgressStream(TimeProvider timeProvider)
    {
        this.timeProvider = timeProvider ?? throw new ArgumentNullException(nameof(timeProvider));
    }

    public bool Exists(ScanId scanId)
        => channels.ContainsKey(scanId.Value);

    public ScanProgressEvent Publish(
        ScanId scanId,
        string state,
        string? message = null,
        IReadOnlyDictionary<string, object?>? data = null,
        string? correlationId = null)
    {
        var channel = channels.GetOrAdd(scanId.Value, _ => new ProgressChannel());

        ScanProgressEvent progressEvent;
        lock (channel)
        {
            var sequence = channel.NextSequence();
            var correlation = correlationId ?? $"{scanId.Value}:{sequence:D4}";
            progressEvent = new ScanProgressEvent(
                scanId,
                sequence,
                timeProvider.GetUtcNow(),
                state,
                message,
                correlation,
                NormalizePayload(data));

            channel.Append(progressEvent);
        }

        return progressEvent;
    }

    public async IAsyncEnumerable<ScanProgressEvent> SubscribeAsync(
        ScanId scanId,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        if (!channels.TryGetValue(scanId.Value, out var channel))
        {
            yield break;
        }

        IReadOnlyList<ScanProgressEvent> snapshot;
        lock (channel)
        {
            snapshot = channel.Snapshot();
        }

        foreach (var progressEvent in snapshot)
        {
            yield return progressEvent;
        }

        var reader = channel.Reader;
        while (await reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
        {
            while (reader.TryRead(out var progressEvent))
            {
                yield return progressEvent;
            }
        }
    }

    private static IReadOnlyDictionary<string, object?> NormalizePayload(IReadOnlyDictionary<string, object?>? data)
    {
        if (data is null || data.Count == 0)
        {
            return EmptyData;
        }

        var sorted = new SortedDictionary<string, object?>(StringComparer.OrdinalIgnoreCase);
        foreach (var pair in data)
        {
            sorted[pair.Key] = pair.Value;
        }

        return sorted.Count == 0
            ? EmptyData
            : new ReadOnlyDictionary<string, object?>(sorted);
    }
}