natsdotnet/tests/NATS.Server.Tests/IO/DynamicBufferPoolTests.cs

using System.Text;
using NATS.Server.IO;
using Shouldly;

// Go reference: client.go — dynamic buffer sizing and broadcast flush coalescing for fan-out.

namespace NATS.Server.Tests.IO;

public class DynamicBufferPoolTests
{
    // -----------------------------------------------------------------------
    // Rent (IMemoryOwner<byte>)
    // -----------------------------------------------------------------------

    [Fact]
    public void Rent_returns_buffer_of_requested_size_or_larger()
    {
        // Go ref: client.go — dynamic buffer sizing (512 → 65536).
        var pool = new OutboundBufferPool();
        using var owner = pool.Rent(100);
        owner.Memory.Length.ShouldBeGreaterThanOrEqualTo(100);
    }

    // -----------------------------------------------------------------------
    // RentBuffer — tier sizing
    // -----------------------------------------------------------------------

    [Fact]
    public void RentBuffer_returns_small_buffer()
    {
        // Go ref: client.go — initial 512 B write buffer per connection.
        var pool = new OutboundBufferPool();
        var buf = pool.RentBuffer(100);
        buf.Length.ShouldBeGreaterThanOrEqualTo(512);
        pool.ReturnBuffer(buf);
    }

    [Fact]
    public void RentBuffer_returns_medium_buffer()
    {
        // Go ref: client.go — 4 KiB write buffer growth step.
        var pool = new OutboundBufferPool();
        var buf = pool.RentBuffer(1000);
        buf.Length.ShouldBeGreaterThanOrEqualTo(4096);
        pool.ReturnBuffer(buf);
    }

    [Fact]
    public void RentBuffer_returns_large_buffer()
    {
        // Go ref: client.go — max 64 KiB write buffer per connection.
        var pool = new OutboundBufferPool();
        var buf = pool.RentBuffer(10000);
        buf.Length.ShouldBeGreaterThanOrEqualTo(65536);
        pool.ReturnBuffer(buf);
    }

    // -----------------------------------------------------------------------
    // ReturnBuffer + reuse
    // -----------------------------------------------------------------------

    [Fact]
    public void ReturnBuffer_and_reuse()
    {
        // Verifies that a returned buffer is available for reuse on the next
        // RentBuffer call of the same tier.
        // Go ref: client.go — buffer pooling to avoid GC pressure.
        var pool = new OutboundBufferPool();

        var first = pool.RentBuffer(100);   // small tier → 512 B
        first.Length.ShouldBe(512);
        pool.ReturnBuffer(first);

        var second = pool.RentBuffer(100);  // should reuse the returned buffer
        second.Length.ShouldBe(512);
        // ReferenceEquals confirms the exact same array instance was reused.
        ReferenceEquals(first, second).ShouldBeTrue();
        pool.ReturnBuffer(second);
    }

    // -----------------------------------------------------------------------
    // BroadcastDrain — coalescing
    // -----------------------------------------------------------------------

    [Fact]
    public void BroadcastDrain_coalesces_writes()
    {
        // Go ref: client.go — broadcast flush for fan-out publish.
        var pool = new OutboundBufferPool();

        var p1 = Encoding.UTF8.GetBytes("Hello");
        var p2 = Encoding.UTF8.GetBytes(", ");
        var p3 = Encoding.UTF8.GetBytes("World");

        IReadOnlyList<ReadOnlyMemory<byte>> pending =
        [
            p1.AsMemory(),
            p2.AsMemory(),
            p3.AsMemory(),
        ];

        var dest = new byte[OutboundBufferPool.CalculateBroadcastSize(pending)];
        pool.BroadcastDrain(pending, dest);

        Encoding.UTF8.GetString(dest).ShouldBe("Hello, World");
    }

    [Fact]
    public void BroadcastDrain_returns_correct_byte_count()
    {
        // Go ref: client.go — total bytes written during coalesced drain.
        var pool = new OutboundBufferPool();

        IReadOnlyList<ReadOnlyMemory<byte>> pending =
        [
            new byte[10].AsMemory(),
            new byte[20].AsMemory(),
            new byte[30].AsMemory(),
        ];

        var dest = new byte[60];
        var written = pool.BroadcastDrain(pending, dest);

        written.ShouldBe(60);
    }

    // -----------------------------------------------------------------------
    // CalculateBroadcastSize
    // -----------------------------------------------------------------------

    [Fact]
    public void CalculateBroadcastSize_sums_all_writes()
    {
        // Go ref: client.go — pre-check buffer capacity before coalesced drain.
        IReadOnlyList<ReadOnlyMemory<byte>> pending =
        [
            new byte[7].AsMemory(),
            new byte[13].AsMemory(),
        ];

        OutboundBufferPool.CalculateBroadcastSize(pending).ShouldBe(20);
    }

    // -----------------------------------------------------------------------
    // Stats counters
    // -----------------------------------------------------------------------

    [Fact]
    public void RentCount_increments()
    {
        // Go ref: client.go — observability for buffer allocation rate.
        var pool = new OutboundBufferPool();

        pool.RentCount.ShouldBe(0L);

        using var _ = pool.Rent(100);
        pool.RentBuffer(200);

        pool.RentCount.ShouldBe(2L);
    }

    [Fact]
    public void BroadcastCount_increments()
    {
        // Go ref: client.go — observability for fan-out drain operations.
        var pool = new OutboundBufferPool();

        pool.BroadcastCount.ShouldBe(0L);

        IReadOnlyList<ReadOnlyMemory<byte>> pending = [new byte[4].AsMemory()];
        var dest = new byte[4];

        pool.BroadcastDrain(pending, dest);
        pool.BroadcastDrain(pending, dest);
        pool.BroadcastDrain(pending, dest);

        pool.BroadcastCount.ShouldBe(3L);
    }
}