lmxopcua/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/Cip/RequestScheduler.cs

namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Cip;

/// <summary>
/// Per-connection CIP request scheduler. Keeps exactly one request in flight on its
/// <see cref="ICipChannel"/>; requests that arrive while one is in flight queue, and when the channel
/// frees the scheduler packs as many of them (FIFO) as fit in <see cref="ICipChannel.MaxRequestBytes"/>
/// into one Multiple Service Packet. Each waiter receives its own embedded <see cref="CipReply"/>.
/// </summary>
/// <remarks>
/// One scheduler per device connection, never per tag (review finding B1). The channel is always
/// called with the scheduler's own stop token, never a waiter's, so one caller's cancellation can
/// never abort or fault a shared exchange (B2). A waiter cancelled while queued is removed and never
/// sent; one cancelled while its batch is in flight is abandoned and the batch's other waiters still
/// complete. There is deliberately no batching delay window: packing falls out of the single pump
/// being busy on the channel while requests accumulate.
/// </remarks>
public sealed class RequestScheduler : IAsyncDisposable
{
    private const byte MultipleServiceCode = 0x0A;
    private const byte EmbeddedServiceError = 0x1E;

    /// <summary>Smallest possible CIP reply: service, reserved, general status, additional-status size.</summary>
    public const int ReplyHeaderBytes = 4;

    /// <summary>The <c>expectedReplyBytes</c> default: the reply carries no data beyond its header (a
    /// write, a status-only service). Every caller whose reply carries data must pass its real bound.</summary>
    public const int StatusOnlyReply = ReplyHeaderBytes;

    /// <summary>Pass as <c>expectedReplyBytes</c> when the reply size cannot be bounded: the request is
    /// then sent alone, never packed, because its reply could fill the whole connection.</summary>
    public const int UnknownReply = -1;

    private sealed class Pending(byte[] request, bool exclusive, int expectedReplyBytes)
    {
        // Private lock object: lock(this) would let any holder of the Pending reference contend.
        private readonly object _gate = new();
        private bool _completed;
        private CancellationTokenRegistration _registration;
        private bool _hasRegistration;

        public byte[] Request { get; } = request;

        /// <summary>True when this request must never share a Multiple Service Packet with another.</summary>
        public bool Exclusive { get; } = exclusive;

        /// <summary>Expected whole-reply size (header included), or <see cref="UnknownReply"/>.</summary>
        public int ExpectedReplyBytes { get; } = expectedReplyBytes;
        public TaskCompletionSource<CipReply> Tcs { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public LinkedListNode<Pending>? Node { get; set; }

        /// <summary>Stores the waiter's registration; disposes it at once if the waiter already completed.</summary>
        public void SetRegistration(CancellationTokenRegistration registration, RequestScheduler owner)
        {
            lock (_gate)
            {
                if (!_completed)
                {
                    _registration = registration;
                    _hasRegistration = true;
                    return;
                }
            }
            registration.Dispose();
            Interlocked.Decrement(ref owner._liveRegistrations);
        }

        /// <summary>Marks the waiter complete exactly once and releases its registration (finding 10).
        /// Returns false if it was already completed.</summary>
        private bool Finish(RequestScheduler owner)
        {
            CancellationTokenRegistration registration;
            bool hasRegistration;
            lock (_gate)
            {
                if (_completed) return false;
                _completed = true;
                registration = _registration;
                hasRegistration = _hasRegistration;
                _registration = default;
                _hasRegistration = false;
            }
            if (hasRegistration)
            {
                // Safe from inside the registration's own callback: Dispose does not wait on the
                // callback that is currently executing on this thread.
                registration.Dispose();
                Interlocked.Decrement(ref owner._liveRegistrations);
            }
            return true;
        }

        public void Succeed(CipReply reply, RequestScheduler owner)
        {
            if (Finish(owner)) Tcs.TrySetResult(reply);
        }

        public void Fail(Exception ex, RequestScheduler owner)
        {
            if (Finish(owner)) Tcs.TrySetException(ex);
        }

        public void Cancel(CancellationToken ct, RequestScheduler owner)
        {
            if (Finish(owner)) Tcs.TrySetCanceled(ct);
        }
    }

    private readonly ICipChannel _channel;
    private readonly bool _allowPacking;
    private readonly LinkedList<Pending> _queue = new();
    private readonly object _lock = new();
    private readonly SemaphoreSlim _signal = new(0);
    private readonly CancellationTokenSource _stop = new();
    private readonly Task _pump;
    private bool _disposed;
    private int _liveRegistrations;
    private int _cancellationCallbacksRun;

    public RequestScheduler(ICipChannel channel, bool allowPacking)
    {
        _channel = channel ?? throw new ArgumentNullException(nameof(channel));
        _allowPacking = allowPacking;
        _pump = Task.Run(PumpAsync);
    }

    /// <summary>Test hook: waiter cancellation registrations that have not yet been disposed.</summary>
    internal int LiveRegistrations => Volatile.Read(ref _liveRegistrations);

    /// <summary>Test hook: waiter cancellation callbacks that have actually run. A disposed
    /// registration never runs, so this stays flat when a long-lived token is cancelled later.</summary>
    internal int CancellationCallbacksRun => Volatile.Read(ref _cancellationCallbacksRun);

    /// <summary>Test hook: requests still queued (not yet taken into a batch).</summary>
    internal int QueuedCount { get { lock (_lock) return _queue.Count; } }

    /// <summary>
    /// Queue one CIP request. The returned reply carries the service's own general status; a non-zero
    /// status is not an exception here. Faults with <see cref="AbTransportException"/> when the channel
    /// or the MSP envelope fails, or the scheduler is disposed; cancels only when <paramref name="ct"/> fires.
    /// </summary>
    /// <param name="request">The CIP request.</param>
    /// <param name="ct">Cancels this waiter only.</param>
    /// <param name="exclusive">When true the request is sent alone, never packed into a Multiple Service
    /// Packet: for requests whose reply can fill the whole connection.</param>
    /// <param name="expectedReplyBytes">Upper bound on this request's whole reply (header included) as
    /// it would sit inside a Multiple Service Packet reply. The packer budgets this against the
    /// connection size as well as the request bytes (#582), so a batch of small requests with large
    /// replies (STRINGs, arrays) is split instead of overflowing. Defaults to <see cref="StatusOnlyReply"/>;
    /// <see cref="UnknownReply"/> sends the request alone.</param>
    public Task<CipReply> SendAsync(byte[] request, CancellationToken ct, bool exclusive = false,
        int expectedReplyBytes = StatusOnlyReply)
    {
        ArgumentNullException.ThrowIfNull(request);
        if (ct.IsCancellationRequested) return Task.FromCanceled<CipReply>(ct);

        var p = new Pending(request, exclusive || expectedReplyBytes < 0, Math.Max(expectedReplyBytes, UnknownReply));
        lock (_lock)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            p.Node = _queue.AddLast(p);
        }

        if (ct.CanBeCanceled)
        {
            Interlocked.Increment(ref _liveRegistrations);
            var registration = ct.Register(static state =>
            {
                var (self, pending, token) = ((RequestScheduler, Pending, CancellationToken))state!;
                Interlocked.Increment(ref self._cancellationCallbacksRun);
                lock (self._lock)
                {
                    if (pending.Node?.List is not null) self._queue.Remove(pending.Node);
                }
                pending.Cancel(token, self);
            }, (this, p, ct));
            p.SetRegistration(registration, this);
        }

        lock (_lock)
        {
            if (!_disposed) _signal.Release();
        }
        return p.Tcs.Task;
    }

    /// <summary>Moves the next batch from the queue into <paramref name="batch"/>. Filling a caller-owned
    /// list means a throw part-way through still leaves every taken waiter reachable for faulting.</summary>
    private void TakeBatch(List<Pending> batch)
    {
        lock (_lock)
        {
            var sizes = new List<int>();
            var replySizes = new List<int>();
            var replyBudget = _channel.MaxRequestBytes; // connections are symmetric: one size bounds both directions
            while (_queue.First is { } n)
            {
                if (batch.Count > 0)
                {
                    if (!_allowPacking || n.Value.Exclusive) break;
                    sizes.Add(n.Value.Request.Length);
                    replySizes.Add(n.Value.ExpectedReplyBytes);
                    if (MultipleServicePacket.BuildSize(sizes) > _channel.MaxRequestBytes
                        || ReplyEnvelopeSize(replySizes) > replyBudget)
                    {
                        break;
                    }
                }
                else
                {
                    // The first request always goes, even when it alone exceeds the limit: it is sent
                    // bare, and callers that must fit (LogixTagClient) fragment before they get here.
                    sizes.Add(n.Value.Request.Length);
                    replySizes.Add(n.Value.ExpectedReplyBytes);
                }
                batch.Add(n.Value);
                _queue.RemoveFirst();
                if (n.Value.Exclusive) break;
            }
        }
    }

    /// <summary>Expected size of the MSP reply for the given embedded reply sizes: reply header, count,
    /// one offset per reply, then the replies themselves.</summary>
    internal static int ReplyEnvelopeSize(IReadOnlyList<int> replySizes)
    {
        var total = ReplyHeaderBytes + 2 + 2 * replySizes.Count;
        foreach (var r in replySizes) total += r;
        return total;
    }

    private async Task PumpAsync()
    {
        while (true)
        {
            try { await _signal.WaitAsync(_stop.Token).ConfigureAwait(false); }
            catch (OperationCanceledException) { return; }

            // A signal only means "there may be work"; drain the queue fully before waiting again,
            // so a surplus signal (e.g. from a cancelled waiter) is just one empty pass.
            while (!_stop.IsCancellationRequested)
            {
                var batch = new List<Pending>();
                try
                {
                    TakeBatch(batch);
                    if (batch.Count == 0) break;
                    await ExecuteAsync(batch).ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    // Last resort: ExecuteAsync already handles every exchange failure, so reaching here
                    // means the scheduler itself broke (e.g. MaxRequestBytes threw). Fault everything
                    // pending rather than let the pump die and leave callers hanging, then keep pumping
                    // so later requests still get an answer.
                    FaultAll(batch, ex);
                    break;
                }
            }
        }
    }

    private void FaultAll(List<Pending> batch, Exception ex)
    {
        var error = ex as AbTransportException
            ?? new AbTransportException(AbStatus.BadConnection, $"Request scheduler failed: {ex.Message}", ex);
        List<Pending> queued;
        lock (_lock)
        {
            queued = [.. _queue];
            _queue.Clear();
        }
        foreach (var b in batch) b.Fail(error, this);
        foreach (var q in queued) q.Fail(error, this);
    }

    private async Task ExecuteAsync(List<Pending> batch)
    {
        try
        {
            if (batch.Count == 1)
            {
                // Never a waiter's token (B2).
                var raw = await _channel.SendAsync(batch[0].Request, _stop.Token).ConfigureAwait(false);
                batch[0].Succeed(CipReply.Parse(raw), this);
                return;
            }

            var requests = new List<byte[]>(batch.Count);
            foreach (var b in batch) requests.Add(b.Request);
            var envelopeRaw = await _channel.SendAsync(MultipleServicePacket.Build(requests), _stop.Token).ConfigureAwait(false);
            var envelope = CipReply.Parse(envelopeRaw);
            if (envelope.Service != MultipleServiceCode)
                throw new AbTransportException(AbStatus.BadReply, $"MSP reply carries service 0x{envelope.Service:X2}.");
            if (envelope.GeneralStatus is not (0x00 or EmbeddedServiceError))
                throw new AbTransportException(envelope.Status, $"Multiple Service Packet failed: general 0x{envelope.GeneralStatus:X2}.");
            var replies = MultipleServicePacket.Parse(envelope, batch.Count);
            for (var i = 0; i < batch.Count; i++) batch[i].Succeed(replies[i], this);
        }
        catch (OperationCanceledException) when (_stop.IsCancellationRequested)
        {
            // The scheduler is being disposed: not the caller's cancellation, so not an OCE.
            var error = DisposedError();
            foreach (var b in batch) b.Fail(error, this);
        }
        catch (AbTransportException ex)
        {
            foreach (var b in batch) b.Fail(ex, this);
        }
        catch (Exception ex)
        {
            var wrapped = new AbTransportException(AbStatus.BadConnection, $"CIP exchange failed: {ex.Message}", ex);
            foreach (var b in batch) b.Fail(wrapped, this);
        }
    }

    /// <summary>
    /// The error a waiter sees when the scheduler (i.e. its connection) goes away under it. Deliberately
    /// NOT an <see cref="OperationCanceledException"/>: upstream treats OCE as "the caller unsubscribed"
    /// and swallows it, whereas this must surface as a communication failure. A waiter cancelled by its
    /// own token still gets an OCE.
    /// </summary>
    private static AbTransportException DisposedError() =>
        new(AbStatus.BadConnection, "Scheduler disposed.");

    /// <summary>Stops the pump. In-flight and queued waiters fault with
    /// <see cref="AbTransportException"/> (<see cref="AbStatus.BadConnection"/>).</summary>
    public async ValueTask DisposeAsync()
    {
        lock (_lock)
        {
            if (_disposed) return;
            _disposed = true;
        }
        _stop.Cancel();
        try { await _pump.ConfigureAwait(false); } catch { /* stopping */ }

        List<Pending> remaining;
        lock (_lock)
        {
            remaining = [.. _queue];
            _queue.Clear();
        }
        var error = DisposedError();
        foreach (var p in remaining) p.Fail(error, this);
        _signal.Dispose();
        _stop.Dispose();
    }
}