lmxopcua/src/Core/ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol/ScriptWorkerFramePipe.cs

using System.Buffers;
using System.Buffers.Binary;
using System.Collections.Concurrent;
using Google.Protobuf;
using ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol.V1;

namespace ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol;

/// <summary>
///     A buffered reader of <see cref="ScriptWorkerFrameCodec"/> frames (lmxopcua#789). One read from the pipe pulls in
///     every frame the peer has written so far, and they are then parsed from the buffer without another read, so a
///     batch of results costs one wake-up instead of two reads per frame. Same validation as
///     <see cref="ScriptWorkerFrameCodec.ReadAsync"/>: zero-length, oversize, truncated, unparseable, wrong-version and
///     bodiless frames are a <see cref="ScriptWorkerProtocolException"/>. Not thread-safe: one reader per stream.
/// </summary>
public sealed class ScriptWorkerFrameReader
{
    private const int InitialBuffer = 64 * 1024;

    private readonly Stream _stream;
    private byte[] _buffer = new byte[InitialBuffer];
    private int _start;
    private int _end;

    /// <summary>Initializes the reader.</summary>
    /// <param name="stream">The pipe.</param>
    /// <param name="maxFrameBytes">The frame cap.</param>
    public ScriptWorkerFrameReader(Stream stream, int maxFrameBytes = ScriptWorkerFrameCodec.DefaultMaxFrameBytes)
    {
        ArgumentNullException.ThrowIfNull(stream);
        _stream = stream;
        MaxFrameBytes = maxFrameBytes;
    }

    /// <summary>The negotiated frame cap.</summary>
    public int MaxFrameBytes { get; set; }

    /// <summary>Reads the next frame.</summary>
    /// <param name="ct">Cancellation.</param>
    /// <returns>The envelope, or null on a clean end of stream at a frame boundary.</returns>
    public async ValueTask<Envelope?> ReadAsync(CancellationToken ct = default)
    {
        while (true)
        {
            if (TryTake(out var envelope, out var need)) return envelope;
            Prepare(need);
            var n = await _stream.ReadAsync(_buffer.AsMemory(_end), ct).ConfigureAwait(false);
            if (n == 0) return EndOfStream();
            _end += n;
        }
    }

    /// <summary>Reads the next frame, blocking (for a dedicated reader thread).</summary>
    /// <returns>The envelope, or null on a clean end of stream at a frame boundary.</returns>
    public Envelope? Read()
    {
        while (true)
        {
            if (TryTake(out var envelope, out var need)) return envelope;
            Prepare(need);
            var n = _stream.Read(_buffer, _end, _buffer.Length - _end);
            if (n == 0) return EndOfStream();
            _end += n;
        }
    }

    private bool TryTake(out Envelope? envelope, out int need)
    {
        envelope = null;
        var available = _end - _start;
        if (available < 4)
        {
            need = 4;
            return false;
        }
        var size = BinaryPrimitives.ReadInt32LittleEndian(_buffer.AsSpan(_start, 4));
        ScriptWorkerFrameCodec.CheckLength(size, MaxFrameBytes);
        need = size + 4;
        if (available < need) return false;
        envelope = ScriptWorkerFrameCodec.ParseBody(_buffer.AsSpan(_start + 4, size));
        _start += need;
        if (_start == _end) _start = _end = 0;
        return true;
    }

    /// <summary>Makes room for a frame of <paramref name="need"/> bytes (prefix included) starting at <c>_start</c>.</summary>
    private void Prepare(int need)
    {
        if (_start + need <= _buffer.Length && _end < _buffer.Length) return;
        var pending = _end - _start;
        if (need > _buffer.Length)
        {
            var bigger = new byte[Math.Max(need, _buffer.Length * 2)];
            Buffer.BlockCopy(_buffer, _start, bigger, 0, pending);
            _buffer = bigger;
        }
        else if (_start > 0)
        {
            Buffer.BlockCopy(_buffer, _start, _buffer, 0, pending);
        }
        _start = 0;
        _end = pending;
    }

    private Envelope? EndOfStream()
    {
        var pending = _end - _start;
        if (pending == 0) return null;
        if (pending < 4) throw new ScriptWorkerProtocolException("truncated frame header");
        var size = BinaryPrimitives.ReadInt32LittleEndian(_buffer.AsSpan(_start, 4));
        throw new ScriptWorkerProtocolException($"truncated frame: {pending - 4} of {size} bytes");
    }
}

/// <summary>
///     A combining frame writer (lmxopcua#789): the thread that queues a frame writes it itself when no write is in
///     progress, so the common case needs no hand-off to a writer task (no thread-pool wake-up, no spin-wait), and
///     frames queued while a write is in progress are written by that writer in the same pass. Within a pass consecutive
///     <see cref="EvalBatch"/> items are merged into one batch and consecutive <see cref="EvalResults"/> items into one
///     (design §4.5 opportunistic batching), every frame is serialised into one buffer, and the buffer goes out in one
///     write. Frame order is preserved.
/// </summary>
/// <remarks>
///     <para>Asynchronous mode: a write that cannot complete at once (a full pipe) continues on the I/O completion, and
///     frames queued meanwhile wait for it. Synchronous mode (<see cref="ScriptWorkerFrameCodec.UseBlockingIo"/>, so
///     the socket under a Unix pipe never turns non-blocking): the draining thread blocks while the pipe is full. Only
///     one thread drains at a time — every other producer just queues and returns — and a peer that stops reading is
///     killed by its supervisor (the host's heartbeat check), which fails the write and releases that thread.</para>
///     <para>A failed write completes the writer and reports the exception once through <c>onFault</c>.</para>
/// </remarks>
public sealed class ScriptWorkerFrameWriter
{
    private readonly Stream _stream;
    private readonly Action<Exception> _onFault;
    private readonly ConcurrentQueue<Envelope> _queue = new();
    private readonly ArrayBufferWriter<byte> _pass = new(16 * 1024);
    private readonly bool _synchronous;
    private int _draining;
    private int _completed;
    private int _faulted;
    private long _seq;

    /// <summary>Initializes the writer.</summary>
    /// <param name="stream">The pipe.</param>
    /// <param name="onFault">Called once when a write fails (or a frame is refused); the writer is then completed.</param>
    /// <param name="maxFrameBytes">The frame cap.</param>
    /// <param name="firstSeq">The last sequence number already used (the handshake's).</param>
    /// <param name="synchronous">Write with blocking calls (see the remarks).</param>
    public ScriptWorkerFrameWriter(Stream stream, Action<Exception> onFault,
        int maxFrameBytes = ScriptWorkerFrameCodec.DefaultMaxFrameBytes, long firstSeq = 0, bool synchronous = false)
    {
        _synchronous = synchronous;
        ArgumentNullException.ThrowIfNull(stream);
        ArgumentNullException.ThrowIfNull(onFault);
        _stream = stream;
        _onFault = onFault;
        MaxFrameBytes = maxFrameBytes;
        _seq = firstSeq;
    }

    /// <summary>The negotiated frame cap.</summary>
    public int MaxFrameBytes { get; set; }

    /// <summary>Whether <see cref="Complete"/> has been called or a write failed.</summary>
    public bool IsCompleted => Volatile.Read(ref _completed) != 0;

    /// <summary>
    ///     Queues a frame and, when no write is in progress, writes it on the calling thread. Returns false once the
    ///     writer is completed; the frame is then dropped.
    /// </summary>
    /// <param name="envelope">The frame (its version and sequence number are set here).</param>
    /// <returns>Whether the frame was accepted.</returns>
    public bool TryWrite(Envelope envelope)
    {
        ArgumentNullException.ThrowIfNull(envelope);
        // A full fence before the completion check, so a caller that registered state (a pending evaluation) before
        // calling here and a completer that completes and then sweeps that state cannot both miss each other.
        Interlocked.MemoryBarrier();
        if (Volatile.Read(ref _completed) != 0) return false;
        envelope.ProtocolVersion = ScriptWorkerFrameCodec.ProtocolVersion;
        _queue.Enqueue(envelope);
        if (Interlocked.CompareExchange(ref _draining, 1, 0) == 0) _ = DrainAsync();
        return true;
    }

    /// <summary>Stops accepting frames. Frames already queued are still written by a pass in progress.</summary>
    public void Complete() => Interlocked.Exchange(ref _completed, 1);

    private async Task DrainAsync()
    {
        try
        {
            while (true)
            {
                FillPass();
                if (_pass.WrittenCount > 0)
                {
                    if (_synchronous) _stream.Write(_pass.WrittenSpan);
                    else await _stream.WriteAsync(_pass.WrittenMemory).ConfigureAwait(false);
                    _pass.ResetWrittenCount();
                    // Synchronous on purpose: a pipe has no write buffer to flush, and PipeStream does not override
                    // FlushAsync, so the base Stream.FlushAsync would queue a thread-pool work item per pass.
                    _stream.Flush();
                }

                Volatile.Write(ref _draining, 0);
                // A frame queued after the last TryDequeue but before the flag dropped saw _draining == 1 and left it
                // to us: take the flag back and go round again, unless another thread already has.
                if (_queue.IsEmpty || Interlocked.CompareExchange(ref _draining, 1, 0) != 0) return;
            }
        }
        catch (Exception ex)
        {
            Interlocked.Exchange(ref _completed, 1);
            _pass.ResetWrittenCount();
            while (_queue.TryDequeue(out _)) { }
            Volatile.Write(ref _draining, 0);
            if (Interlocked.Exchange(ref _faulted, 1) == 0) _onFault(ex);
        }
    }

    private void FillPass()
    {
        EvalBatch? evals = null;
        EvalResults? results = null;
        while (_queue.TryDequeue(out var envelope))
        {
            switch (envelope.BodyCase)
            {
                case Envelope.BodyOneofCase.Eval:
                    if (results is not null) { Append(new Envelope { Results = results }); results = null; }
                    (evals ??= new EvalBatch()).Items.AddRange(envelope.Eval.Items);
                    break;
                case Envelope.BodyOneofCase.Results:
                    if (evals is not null) { Append(new Envelope { Eval = evals }); evals = null; }
                    results ??= new EvalResults();
                    results.Items.AddRange(envelope.Results.Items);
                    results.StartedEvalIds.AddRange(envelope.Results.StartedEvalIds);
                    results.DeferredEvalIds.AddRange(envelope.Results.DeferredEvalIds);
                    break;
                default:
                    if (evals is not null) { Append(new Envelope { Eval = evals }); evals = null; }
                    if (results is not null) { Append(new Envelope { Results = results }); results = null; }
                    Append(envelope);
                    break;
            }
        }
        if (evals is not null) Append(new Envelope { Eval = evals });
        if (results is not null) Append(new Envelope { Results = results });
    }

    private void Append(Envelope envelope)
    {
        envelope.ProtocolVersion = ScriptWorkerFrameCodec.ProtocolVersion;
        envelope.Seq = (ulong)Interlocked.Increment(ref _seq);
        var size = envelope.CalculateSize();
        if (size == 0) throw new ScriptWorkerProtocolException("refusing to write an empty frame");
        if (size > MaxFrameBytes)
            throw new ScriptWorkerProtocolException($"frame of {size} bytes exceeds the {MaxFrameBytes}-byte cap");
        var span = _pass.GetSpan(size + 4);
        BinaryPrimitives.WriteInt32LittleEndian(span, size);
        envelope.WriteTo(span.Slice(4, size));
        _pass.Advance(size + 4);
    }
}