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

using System.Buffers;
using System.Runtime.InteropServices;

namespace AyCode.Services.Server.Tests.SignalRs;

/// <summary>
/// Custom IBufferWriter that simulates Kestrel's PinnedBlockMemoryPool transport behavior:
///
/// Production: Kestrel PipeWriter → PinnedBlockMemoryPool → MemoryManager-backed Memory
///             → TryGetArray FAILS → BWO takes owned-buffer fallback path
///
/// Test:       SlabTransportWriter → SlabMemoryManager-backed Memory
///             → TryGetArray FAILS → BWO takes owned-buffer fallback path (same as production)
///
/// Key behaviors:
/// - GetMemory returns MemoryManager-backed Memory (TryGetArray always fails)
/// - Fixed-size slab segments with random offsets (simulates slab allocator)
/// - After Advance, remaining slab space is reused (offset grows)
/// - GetMemory may return less than sizeHint (remaining slab space)
/// - Deterministic via seeded Random for reproducible tests
/// </summary>
internal sealed class SlabTransportWriter : IBufferWriter<byte>
{
    private readonly int _slabSize;
    private readonly Random _rng;

    // Committed data tracking
    private readonly List<CommittedSegment> _segments = new();
    private int _totalCommitted;

    // Current slab state
    private byte[]? _currentSlab;
    private int _writePos;   // current write position within slab
    private int _slabEnd;    // end of usable area in slab

    public SlabTransportWriter(int slabSize = 256, int seed = 42)
    {
        _slabSize = slabSize;
        _rng = new Random(seed);
    }

    /// <summary>Total bytes committed via Advance.</summary>
    public int WrittenCount => _totalCommitted;

    /// <summary>Number of slab segments allocated.</summary>
    public int SlabCount { get; private set; }

    public void Advance(int count)
    {
        if (_currentSlab == null)
            throw new InvalidOperationException("Call GetMemory/GetSpan before Advance");
        if (count < 0 || _writePos + count > _slabEnd)
            throw new InvalidOperationException(
                $"Advance({count}) invalid: writePos={_writePos}, slabEnd={_slabEnd}, remaining={_slabEnd - _writePos}");

        _segments.Add(new CommittedSegment(_currentSlab, _writePos, count));
        _writePos += count;
        _totalCommitted += count;
    }

    /// <summary>
    /// Returns MemoryManager-backed Memory so that TryGetArray ALWAYS fails.
    /// May return fewer bytes than sizeHint (remaining slab space) — legal per IBufferWriter contract.
    /// This forces BWO to rent from ArrayPool (owned-buffer path), matching production behavior.
    /// </summary>
    public Memory<byte> GetMemory(int sizeHint = 0)
    {
        sizeHint = Math.Max(1, sizeHint);
        var remaining = _currentSlab != null ? _slabEnd - _writePos : 0;

        if (remaining <= 0)
        {
            // Allocate new slab — at least sizeHint to avoid starving the caller
            AllocateNewSlab(Math.Max(sizeHint, _slabSize));
            remaining = _slabEnd - _writePos;
        }

        // Return MemoryManager-backed Memory: TryGetArray will fail
        return new SlabMemoryManager(_currentSlab!, _writePos, remaining).Memory;
    }

    /// <summary>
    /// Returns Span with at least sizeHint bytes.
    /// Used by FlushOwnedBuffer: _writer.GetSpan(bytesInChunk) must be large enough for CopyTo.
    /// Allocates new slab if remaining space is insufficient.
    /// </summary>
    public Span<byte> GetSpan(int sizeHint = 0)
    {
        sizeHint = Math.Max(1, sizeHint);
        var remaining = _currentSlab != null ? _slabEnd - _writePos : 0;

        if (remaining < sizeHint)
        {
            AllocateNewSlab(Math.Max(sizeHint, _slabSize));
            remaining = _slabEnd - _writePos;
        }

        return _currentSlab.AsSpan(_writePos, remaining);
    }

    private void AllocateNewSlab(int minSize)
    {
        // Random offset within slab — simulates Kestrel slab allocator non-zero offsets
        var offset = _rng.Next(0, Math.Max(1, _slabSize / 4));
        // ±12% size jitter for variety
        var jitter = _rng.Next(-_slabSize / 8, _slabSize / 8 + 1);
        var actualSize = Math.Max(minSize, _slabSize + jitter);
        _currentSlab = new byte[actualSize + offset];
        _writePos = offset;
        _slabEnd = offset + actualSize;
        SlabCount++;
    }

    /// <summary>
    /// Get all committed bytes as a contiguous array.
    /// </summary>
    public byte[] ToArray()
    {
        var result = new byte[_totalCommitted];
        var pos = 0;
        foreach (var seg in _segments)
        {
            Buffer.BlockCopy(seg.Array, seg.Offset, result, pos, seg.Length);
            pos += seg.Length;
        }
        return result;
    }

    /// <summary>
    /// Build a multi-segment ReadOnlySequence from committed data, splitting at slab boundaries.
    /// Each committed segment becomes a separate ReadOnlySequence segment.
    /// </summary>
    public ReadOnlySequence<byte> ToReadOnlySequence()
    {
        if (_segments.Count == 0)
            return ReadOnlySequence<byte>.Empty;

        if (_segments.Count == 1)
        {
            var seg = _segments[0];
            return new ReadOnlySequence<byte>(seg.Array, seg.Offset, seg.Length);
        }

        // Build linked segment list
        var first = new MemorySegment(new ReadOnlyMemory<byte>(_segments[0].Array, _segments[0].Offset, _segments[0].Length));
        var current = first;
        for (var i = 1; i < _segments.Count; i++)
        {
            var seg = _segments[i];
            current = current.Append(new ReadOnlyMemory<byte>(seg.Array, seg.Offset, seg.Length));
        }

        return new ReadOnlySequence<byte>(first, 0, current, current.Memory.Length);
    }

    /// <summary>
    /// Verify that TryGetArray fails on our Memory (sanity check for test correctness).
    /// </summary>
    public static void VerifyMemoryManagerBacked()
    {
        var writer = new SlabTransportWriter(64);
        var mem = writer.GetMemory(16);
        if (MemoryMarshal.TryGetArray(mem, out ArraySegment<byte> _))
            throw new InvalidOperationException(
                "SlabTransportWriter.GetMemory returned array-backed Memory — TryGetArray should fail!");
    }

    private record struct CommittedSegment(byte[] Array, int Offset, int Length);

    /// <summary>
    /// MemoryManager that wraps an array region but returns Memory where TryGetArray fails.
    /// This is the key difference from production Kestrel PinnedBlockMemoryPool.
    /// </summary>
    private sealed class SlabMemoryManager : MemoryManager<byte>
    {
        private readonly byte[] _array;
        private readonly int _offset;
        private readonly int _length;

        public SlabMemoryManager(byte[] array, int offset, int length)
        {
            _array = array;
            _offset = offset;
            _length = length;
        }

        public override Span<byte> GetSpan() => _array.AsSpan(_offset, _length);

        public override MemoryHandle Pin(int elementIndex = 0)
            => throw new NotSupportedException("SlabMemoryManager does not support pinning");

        public override void Unpin()
            => throw new NotSupportedException("SlabMemoryManager does not support pinning");

        protected override void Dispose(bool disposing) { }
    }

    /// <summary>
    /// ReadOnlySequenceSegment for building multi-segment sequences from committed data.
    /// </summary>
    private sealed class MemorySegment : ReadOnlySequenceSegment<byte>
    {
        public MemorySegment(ReadOnlyMemory<byte> memory)
        {
            Memory = memory;
        }

        public MemorySegment Append(ReadOnlyMemory<byte> memory)
        {
            var next = new MemorySegment(memory) { RunningIndex = RunningIndex + Memory.Length };
            Next = next;
            return next;
        }
    }
}