mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Workers/WorkerFrameReader.cs

using System.Buffers;
using System.Buffers.Binary;
using Google.Protobuf;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Server.Workers;

/// <summary>
/// Reads length-prefixed WorkerEnvelope protobuf frames from a stream.
/// </summary>
/// <remarks>
///     <see cref="ReadAsync"/> is not reentrant: the reader keeps a per-instance length-prefix scratch
///     buffer, so exactly one consumer may be inside a read at a time. That matches how the reader is
///     used — a single read loop per <c>WorkerClient</c>, with handshake reads completing before the
///     loop starts.
/// </remarks>
public sealed class WorkerFrameReader
{
    private readonly WorkerFrameProtocolOptions _options;
    private readonly Stream _stream;

    // Reused across frames rather than allocated per read. Safe because ReadAsync is
    // single-consumer by construction; the prefix is fully overwritten by every read.
    private readonly byte[] _lengthPrefix = new byte[sizeof(uint)];

    /// <summary>
    /// Initializes a new instance of <see cref="WorkerFrameReader"/>.
    /// </summary>
    /// <param name="stream">Stream to read frames from.</param>
    /// <param name="options">Frame protocol options.</param>
    public WorkerFrameReader(
        Stream stream,
        WorkerFrameProtocolOptions options)
    {
        _stream = stream ?? throw new ArgumentNullException(nameof(stream));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    /// <summary>
    /// Reads a worker envelope frame from the stream asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Parsed worker envelope.</returns>
    public async ValueTask<WorkerEnvelope> ReadAsync(CancellationToken cancellationToken = default)
    {
        await ReadExactlyOrThrowAsync(_lengthPrefix, cancellationToken).ConfigureAwait(false);

        uint payloadLength = BinaryPrimitives.ReadUInt32LittleEndian(_lengthPrefix);
        if (payloadLength == 0)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MalformedLength,
                "Worker frame payload length must be greater than zero.");
        }

        if (payloadLength > _options.MaxMessageBytes)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MessageTooLarge,
                $"Worker frame payload length {payloadLength} exceeds the configured maximum of {_options.MaxMessageBytes} bytes.");
        }

        // Rent the payload buffer from the shared pool rather than allocating a fresh byte[] per
        // frame; large event frames near the cap would otherwise allocate an LOH buffer each time.
        // ParseFrom copies whatever it needs into the parsed message, so the rented buffer
        // can be returned as soon as parsing completes without the envelope aliasing it. The rented
        // buffer may be larger than requested, so the read and parse are bounded to length.
        int length = checked((int)payloadLength);
        byte[] payload = ArrayPool<byte>.Shared.Rent(length);
        WorkerEnvelope envelope;
        try
        {
            await ReadExactlyOrThrowAsync(new Memory<byte>(payload, 0, length), cancellationToken)
                .ConfigureAwait(false);

            try
            {
                envelope = WorkerEnvelope.Parser.ParseFrom(payload, 0, length);
            }
            catch (InvalidProtocolBufferException exception)
            {
                throw new WorkerFrameProtocolException(
                    WorkerFrameProtocolErrorCode.InvalidEnvelope,
                    "Worker frame payload is not a valid WorkerEnvelope protobuf message.",
                    exception);
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(payload);
        }

        WorkerEnvelopeValidator.Validate(envelope, _options);

        return envelope;
    }

    private async ValueTask ReadExactlyOrThrowAsync(
        Memory<byte> buffer,
        CancellationToken cancellationToken)
    {
        try
        {
            await _stream.ReadExactlyAsync(buffer, cancellationToken).ConfigureAwait(false);
        }
        catch (EndOfStreamException exception)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.EndOfStream,
                "Worker frame ended before the expected number of bytes were read.",
                exception);
        }
    }
}