AyCode.Core/AyCode.Services.Server.Tests/SignalRs/TestMultiSegmentProtocol.cs

using System.Buffers;
using System.Diagnostics.CodeAnalysis;
using System.IO.Pipelines;
using AyCode.Core.Serializers.Binaries;
using AyCode.Services.SignalRs;
using Microsoft.AspNetCore.SignalR;
using Microsoft.AspNetCore.SignalR.Protocol;

namespace AyCode.Services.Server.Tests.SignalRs;

/// <summary>
/// Test protocol that simulates production Kestrel transport behavior:
///
/// WRITE SIDE (GetMessageBytesMultiSegment):
///   WriteMessage → SlabTransportWriter (MemoryManager-backed, TryGetArray fails)
///   → BWO takes owned-buffer fallback path (same as production Kestrel PipeWriter)
///   → Automatic byte-by-byte comparison against ArrayBufferWriter reference
///
/// READ SIDE (TryParseMessage):
///   Bytes → Pipe(SlabSimulatingPool, 256B segments) → multi-segment ReadOnlySequence
///   → SequenceBinaryInput cross-boundary reads (same as production Kestrel PipeReader)
///
/// Production: SignalR → WriteMessage(PipeWriter/Kestrel) → 4096B slab segments → TryParseMessage
/// Test:       SignalR → WriteMessage(SlabTransportWriter) → 256B slab segments → TryParseMessage
/// </summary>
internal class TestMultiSegmentProtocol : AyCodeBinaryHubProtocol
{
    private const int SegmentSize = 256;
    private readonly BinaryProtocolMode _mode;

    public TestMultiSegmentProtocol(BinaryProtocolMode mode = BinaryProtocolMode.Bytes)
        : base(AcBinarySerializerOptions.Default, mode)
    {
        _mode = mode;
        Options.BufferWriterChunkSize = SegmentSize;
    }

    /// <summary>
    /// Write side: WriteMessage → SlabTransportWriter (MemoryManager-backed Memory).
    /// Forces BWO into owned-buffer fallback path — same code path as production Kestrel.
    /// Validates output against ArrayBufferWriter reference on every call.
    /// </summary>
    public ReadOnlyMemory<byte> GetMessageBytesMultiSegment(HubMessage message)
    {
        if (_mode == BinaryProtocolMode.AsyncSegment)
            return GetMessageBytesAsyncSegment(message);

        // ── Transport-double path (production simulation) ──────────────────
        var transport = new SlabTransportWriter(SegmentSize);
        WriteMessage(message, transport);
        var transportBytes = transport.ToArray();

        // ── Reference path: ArrayBufferWriter with large capacity (no resize → lengthSpan stays valid) ──
        // NOTE: base.GetMessageBytes uses capacity = chunkSize + 4 = 260, which causes resize
        //       on any non-trivial message → invalidates the back-patched lengthSpan.
        //       We use a large capacity to ensure no resize ever happens.
        var referenceCapacity = Math.Max(transportBytes.Length + 256, 65536);
        var reference = new ArrayBufferWriter<byte>(referenceCapacity);
        WriteMessage(message, reference);
        var referenceBytes = reference.WrittenSpan;

        // ── Validate: transport output must match reference byte-for-byte ──
        ValidateAgainstReference(transportBytes, referenceBytes);

        return transportBytes;
    }

    /// <summary>
    /// AsyncSegment write side: uses a real PipeWriter transport so chunked protocol path activates.
    /// </summary>
    private ReadOnlyMemory<byte> GetMessageBytesAsyncSegment(HubMessage message)
    {
        using var transport = new AsyncSegmentPipeTransportWriter(SegmentSize);
        WriteMessage(message, transport.Writer);
        transport.CompleteWriter();
        var bytes = transport.DrainAllAsync().GetAwaiter().GetResult();
        return bytes;
    }

    /// <summary>
    /// Read side: fill Pipe with 256-byte slab segments → multi-segment ReadOnlySequence.
    /// Same as production Kestrel PipeReader delivering slab-sized segments.
    /// </summary>
    public override bool TryParseMessage(ref ReadOnlySequence<byte> input, IInvocationBinder binder,
        [NotNullWhen(true)] out HubMessage? message)
    {
        var bytes = input.ToArray();
        var pipe = new Pipe(new PipeOptions(pool: new SlabSimulatingPool(SegmentSize)));
        var writer = pipe.Writer;

        // Write in chunks — GetMemory may return less than requested (like Kestrel slab)
        var remaining = bytes.AsSpan();
        while (remaining.Length > 0)
        {
            var mem = writer.GetMemory(Math.Min(SegmentSize, remaining.Length));
            var chunk = Math.Min(mem.Length, remaining.Length);
            remaining[..chunk].CopyTo(mem.Span);
            writer.Advance(chunk);
            remaining = remaining[chunk..];
        }

        writer.Complete();
        pipe.Reader.TryRead(out var result);
        var seq = result.Buffer;

        // Assert multi-segment: if payload > SegmentSize, sequence must NOT be single-segment
        if (bytes.Length > SegmentSize)
        {
            AssertMultiSegment(seq, bytes.Length);
        }

        var success = base.TryParseMessage(ref seq, binder, out message);
        pipe.Reader.Complete();
        return success;
    }

    #region Validation

    /// <summary>
    /// Byte-by-byte comparison of transport-double output vs ArrayBufferWriter reference.
    /// Catches any BWO owned-buffer path bug (position drift, data corruption, length mismatch).
    /// </summary>
    private static void ValidateAgainstReference(byte[] transportBytes, ReadOnlySpan<byte> referenceBytes)
    {
        var refArray = referenceBytes.ToArray();

        // ── Total byte count ──
        if (transportBytes.Length != refArray.Length)
        {
            // Also check outer length prefix for diagnostics
            var transportPrefix = transportBytes.Length >= 4 ? BitConverter.ToInt32(transportBytes, 0) : -1;
            var referencePrefix = refArray.Length >= 4 ? BitConverter.ToInt32(refArray, 0) : -1;

            throw new InvalidOperationException(
                $"[TRANSPORT_DOUBLE] Total byte count mismatch: transport={transportBytes.Length}, reference={refArray.Length}. " +
                $"Diff={transportBytes.Length - refArray.Length} bytes. " +
                $"Outer length prefix: transport={transportPrefix}, reference={referencePrefix}. " +
                $"BWO GetTotalPosition drift detected.");
        }

        // ── Outer length prefix check (first 4 bytes) ──
        if (transportBytes.Length >= 4)
        {
            var transportPrefix = BitConverter.ToInt32(transportBytes, 0);
            var referencePrefix = BitConverter.ToInt32(refArray, 0);
            if (transportPrefix != referencePrefix)
            {
                throw new InvalidOperationException(
                    $"[TRANSPORT_DOUBLE] Outer length prefix mismatch: transport={transportPrefix}, reference={referencePrefix}. " +
                    $"Total bytes match ({transportBytes.Length}) but prefix differs. " +
                    $"BWO owned-buffer path likely has position drift.");
            }
        }

        // ── Byte-by-byte content ──
        var minLen = Math.Min(transportBytes.Length, refArray.Length);
        for (int i = 0; i < minLen; i++)
        {
            if (transportBytes[i] != refArray[i])
            {
                var start = Math.Max(0, i - 8);
                var end = Math.Min(minLen, i + 16);
                var refHex = Convert.ToHexString(refArray.AsSpan(start, end - start));
                var trnHex = Convert.ToHexString(transportBytes.AsSpan(start, end - start));
                throw new InvalidOperationException(
                    $"[TRANSPORT_DOUBLE] Content mismatch at byte {i}/{minLen}: " +
                    $"ref={refHex} transport={trnHex}. " +
                    $"BWO owned-buffer path writing corrupt data.");
            }
        }
    }

    private static void AssertMultiSegment(ReadOnlySequence<byte> seq, int totalLength)
    {
        if (seq.IsSingleSegment)
        {
            throw new InvalidOperationException(
                $"[MULTI_SEGMENT] Expected multi-segment sequence for {totalLength} bytes " +
                $"(> {SegmentSize}B segment size), but got single segment. " +
                $"SlabSimulatingPool or Pipe configuration is wrong.");
        }

        // Count segments
        var segmentCount = 0;
        foreach (var _ in seq)
            segmentCount++;

        if (segmentCount < 2)
        {
            throw new InvalidOperationException(
                $"[MULTI_SEGMENT] Expected >= 2 segments for {totalLength} bytes, got {segmentCount}.");
        }
    }

    #endregion

    /// <summary>
    /// MemoryPool that returns <paramref name="segmentSize"/>-byte blocks at random offsets
    /// within a larger backing array — simulating Kestrel's slab allocator where segments
    /// share a large slab and have non-zero offsets.
    /// </summary>
    private sealed class SlabSimulatingPool(int segmentSize) : MemoryPool<byte>
    {
        private readonly Random _rng = new(42); // deterministic seed for reproducibility

        public override int MaxBufferSize => segmentSize;

        public override IMemoryOwner<byte> Rent(int minBufferSize = -1)
        {
            var size = Math.Max(minBufferSize, segmentSize);
            var offset = _rng.Next(0, segmentSize);       // random slab offset
            var jitter = _rng.Next(-segmentSize / 4, segmentSize / 4 + 1); // ±25% size variance
            var actualSize = Math.Max(16, size + jitter);
            var array = new byte[actualSize + offset];
            return new Owner(array, offset, actualSize);
        }

        protected override void Dispose(bool disposing) { }

        private sealed class Owner(byte[] array, int offset, int length) : IMemoryOwner<byte>
        {
            public Memory<byte> Memory { get; } = array.AsMemory(offset, length);
            public void Dispose() { }
        }
    }
}