ShrinkNetwork/Runtime/Core/ShrinkNetworkWorkQueue.cs

#nullable enable
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork
{
    public enum ShrinkNetworkQueueResult { Completed, Rejected, Replaced, Canceled }

    public sealed class ShrinkNetworkQueueDiagnostics
    {
        public int PendingCount { get; internal set; }
        public long PendingBytes { get; internal set; }
        public long Rejected { get; internal set; }
        public long Replaced { get; internal set; }
        public long Completed { get; internal set; }
        public double OldestWaitMilliseconds { get; internal set; }
        public double LastWaitMilliseconds { get; internal set; }
    }

    /// <summary>Caller-pumped serial execution, fair across session/channel partitions. Only queued state is replaceable.</summary>
    public sealed class ShrinkNetworkWorkQueue : IDisposable
    {
        private sealed class Work
        {
            public Func<UniTask> Callback = null!;
            public UniTaskCompletionSource<ShrinkNetworkQueueResult> Completion = new();
            public string? StateKey;
            public int Bytes;
            public long Enqueued = Stopwatch.GetTimestamp();
            public CancellationToken Cancellation;
        }
        private sealed class Partition
        {
            public readonly LinkedList<Work> Items = new();
            public readonly Dictionary<string, LinkedListNode<Work>> States = new(StringComparer.Ordinal);
        }
        private readonly object _gate = new();
        private readonly Dictionary<(long, string), Partition> _partitions = new();
        private readonly Queue<(long, string)> _ready = new();
        private readonly int _capacity;
        private readonly long _byteCapacity;
        private readonly int _partitionCapacity;
        private int _count, _pumping;
        private long _bytes, _rejected, _replaced, _completed;
        private double _lastWait;
        private bool _disposed;

        public int PendingCount { get { lock (_gate) return _count; } }

        public ShrinkNetworkWorkQueue(int capacity, long byteCapacity, int perPartitionCapacity = int.MaxValue)
        {
            if (capacity <= 0 || byteCapacity <= 0 || perPartitionCapacity <= 0) throw new ArgumentOutOfRangeException(nameof(capacity));
            _capacity = capacity; _byteCapacity = byteCapacity; _partitionCapacity = perPartitionCapacity;
        }

        /// <summary>A nonempty stateKey explicitly permits replacing an unsent state in this session/channel. Never use for RPC or snapshot fragments.</summary>
        public UniTask<ShrinkNetworkQueueResult> EnqueueAsync(long sessionId, string channel, int byteCount,
            Func<UniTask> callback, string? stateKey = null, CancellationToken cancellationToken = default)
        {
            if (callback == null) throw new ArgumentNullException(nameof(callback));
            if (channel == null) throw new ArgumentNullException(nameof(channel));
            if (byteCount < 0) throw new ArgumentOutOfRangeException(nameof(byteCount));
            Work? replaced = null;
            Work item;
            lock (_gate)
            {
                if (_disposed || cancellationToken.IsCancellationRequested) return UniTask.FromResult(ShrinkNetworkQueueResult.Canceled);
                var key = (sessionId, channel);
                _partitions.TryGetValue(key, out var partition);
                LinkedListNode<Work>? old = null;
                if (!string.IsNullOrEmpty(stateKey)) partition?.States.TryGetValue(stateKey!, out old);
                var nextBytes = _bytes - (old?.Value.Bytes ?? 0) + byteCount;
                if (nextBytes > _byteCapacity || (old == null && (_count >= _capacity || (partition?.Items.Count ?? 0) >= _partitionCapacity)))
                { _rejected++; return UniTask.FromResult(ShrinkNetworkQueueResult.Rejected); }
                item = new Work { Callback = callback, Bytes = byteCount, StateKey = string.IsNullOrEmpty(stateKey) ? null : stateKey, Cancellation = cancellationToken };
                if (partition == null) { partition = new Partition(); _partitions.Add(key, partition); _ready.Enqueue(key); }
                if (old != null)
                {
                    replaced = old.Value;
                    // Move a replacement to the tail: later state must not jump ahead of intervening reliable operations.
                    partition.Items.Remove(old); _replaced++;
                }
                else _count++;
                var node = partition.Items.AddLast(item);
                if (item.StateKey != null) partition.States[item.StateKey] = node;
                _bytes = nextBytes;
            }
            replaced?.Completion.TrySetResult(ShrinkNetworkQueueResult.Replaced);
            return item.Completion.Task;
        }

        public async UniTask<int> PumpAsync(int maxItems, long maxBytes = long.MaxValue, TimeSpan? timeBudget = null)
        {
            if (maxItems <= 0 || maxBytes <= 0 || (timeBudget.HasValue && timeBudget.Value <= TimeSpan.Zero)) throw new ArgumentOutOfRangeException(nameof(maxItems));
            if (Interlocked.Exchange(ref _pumping, 1) != 0) return 0;
            var started = Stopwatch.GetTimestamp();
            var processed = 0;
            long bytes = 0;
            try
            {
                while (processed < maxItems && (!timeBudget.HasValue || Elapsed(started) < timeBudget.Value.TotalMilliseconds))
                {
                    Work item;
                    lock (_gate)
                    {
                        if (_disposed || _ready.Count == 0) break;
                        var key = _ready.Peek();
                        var partition = _partitions[key];
                        item = partition.Items.First!.Value;
                        // Allow one oversized item so a byte budget cannot permanently starve a valid packet.
                        if (processed > 0 && item.Bytes > maxBytes - bytes) break;
                        _ready.Dequeue(); partition.Items.RemoveFirst();
                        if (item.StateKey != null) partition.States.Remove(item.StateKey);
                        if (partition.Items.Count == 0) _partitions.Remove(key); else _ready.Enqueue(key);
                        _count--; _bytes -= item.Bytes; _lastWait = Elapsed(item.Enqueued);
                    }
                    try
                    {
                        if (item.Cancellation.IsCancellationRequested) item.Completion.TrySetResult(ShrinkNetworkQueueResult.Canceled);
                        else { await item.Callback(); item.Completion.TrySetResult(ShrinkNetworkQueueResult.Completed); lock (_gate) _completed++; }
                    }
                    catch (OperationCanceledException) { item.Completion.TrySetResult(ShrinkNetworkQueueResult.Canceled); }
                    catch (Exception ex) { item.Completion.TrySetException(ex); }
                    processed++; bytes += item.Bytes;
                }
                return processed;
            }
            finally { Volatile.Write(ref _pumping, 0); }
        }

        public ShrinkNetworkQueueDiagnostics CaptureDiagnostics()
        {
            lock (_gate)
            {
                long oldest = Stopwatch.GetTimestamp();
                foreach (var partition in _partitions.Values)
                    if (partition.Items.First != null) oldest = Math.Min(oldest, partition.Items.First.Value.Enqueued);
                return new ShrinkNetworkQueueDiagnostics { PendingCount = _count, PendingBytes = _bytes, Rejected = _rejected, Replaced = _replaced,
                    Completed = _completed, LastWaitMilliseconds = _lastWait, OldestWaitMilliseconds = _count == 0 ? 0 : Elapsed(oldest) };
            }
        }
        private static double Elapsed(long start) => (Stopwatch.GetTimestamp() - start) * 1000d / Stopwatch.Frequency;
        public void Dispose()
        {
            List<Work> canceled = new();
            lock (_gate)
            {
                if (_disposed) return;
                _disposed = true;
                foreach (var partition in _partitions.Values) canceled.AddRange(partition.Items);
                _partitions.Clear(); _ready.Clear(); _count = 0; _bytes = 0;
            }
            foreach (var item in canceled) item.Completion.TrySetResult(ShrinkNetworkQueueResult.Canceled);
        }
    }
}