natsdotnet/src/NATS.Server/IO/OutboundBufferPool.cs

using System.Buffers;
using System.Collections.Concurrent;

namespace NATS.Server.IO;

/// <summary>
/// Tiered write buffer pool with broadcast drain capability.
/// Maintains internal pools for common sizes (512, 4096, 65536) to reduce
/// allocation overhead on the hot publish path.
/// Go reference: client.go — dynamic buffer sizing and broadcast flush coalescing for fan-out.
/// </summary>
public sealed class OutboundBufferPool
{
    private const int SmallSize  = 512;
    private const int MediumSize = 4096;
    private const int LargeSize  = 65536;

    private readonly ConcurrentBag<byte[]> _small  = new(); // 512 B
    private readonly ConcurrentBag<byte[]> _medium = new(); // 4 KiB
    private readonly ConcurrentBag<byte[]> _large  = new(); // 64 KiB

    private long _rentCount;
    private long _returnCount;
    private long _broadcastCount;

    /// <summary>Total buffer rent operations served by the pool.</summary>
    public long RentCount      => Interlocked.Read(ref _rentCount);
    /// <summary>Total buffer return operations accepted by the pool.</summary>
    public long ReturnCount    => Interlocked.Read(ref _returnCount);
    /// <summary>Total broadcast-drain operations performed.</summary>
    public long BroadcastCount => Interlocked.Read(ref _broadcastCount);

    // -----------------------------------------------------------------------
    // IMemoryOwner<byte> surface (preserves existing callers)
    // -----------------------------------------------------------------------

    /// <summary>
    /// Rents an <see cref="IMemoryOwner{T}"/> whose <c>Memory.Length</c> is at least
    /// <paramref name="size"/> bytes. Tries the internal pool first; falls back to
    /// <see cref="MemoryPool{T}.Shared"/>.
    /// </summary>
    /// <param name="size">Minimum required buffer size.</param>
    public IMemoryOwner<byte> Rent(int size)
    {
        Interlocked.Increment(ref _rentCount);

        // Try to serve from the internal pool so that Dispose() returns the
        // raw buffer back to us rather than to MemoryPool.Shared.
        if (size <= SmallSize && _small.TryTake(out var sb))
            return new PooledMemoryOwner(sb, _small);

        if (size <= MediumSize && _medium.TryTake(out var mb))
            return new PooledMemoryOwner(mb, _medium);

        if (size <= LargeSize && _large.TryTake(out var lb))
            return new PooledMemoryOwner(lb, _large);

        // Nothing cached — rent from the system pool (which may return a larger
        // buffer; that's fine, callers must honour Memory.Length, not the
        // requested size).
        int rounded = size <= SmallSize ? SmallSize
                    : size <= MediumSize ? MediumSize
                    : LargeSize;

        return MemoryPool<byte>.Shared.Rent(rounded);
    }

    // -----------------------------------------------------------------------
    // Raw byte[] surface
    // -----------------------------------------------------------------------

    /// <summary>
    /// Returns a <c>byte[]</c> from the internal pool whose length is at least
    /// <paramref name="size"/> bytes. The caller is responsible for calling
    /// <see cref="ReturnBuffer"/> when finished.
    /// </summary>
    /// <param name="size">Minimum required buffer size.</param>
    public byte[] RentBuffer(int size)
    {
        Interlocked.Increment(ref _rentCount);

        if (size <= SmallSize)
        {
            if (_small.TryTake(out var b)) return b;
            return new byte[SmallSize];
        }

        if (size <= MediumSize)
        {
            if (_medium.TryTake(out var b)) return b;
            return new byte[MediumSize];
        }

        if (_large.TryTake(out var lb)) return lb;
        return new byte[LargeSize];
    }

    /// <summary>
    /// Returns <paramref name="buffer"/> to the appropriate tier so it can be
    /// reused by a subsequent <see cref="RentBuffer"/> call.
    /// </summary>
    /// <param name="buffer">Buffer previously rented from this pool.</param>
    public void ReturnBuffer(byte[] buffer)
    {
        Interlocked.Increment(ref _returnCount);

        switch (buffer.Length)
        {
            case SmallSize:
                _small.Add(buffer);
                break;
            case MediumSize:
                _medium.Add(buffer);
                break;
            case LargeSize:
                _large.Add(buffer);
                break;
            // Buffers of unexpected sizes are simply dropped (GC reclaims them).
        }
    }

    // -----------------------------------------------------------------------
    // Broadcast drain
    // -----------------------------------------------------------------------

    /// <summary>
    /// Coalesces multiple pending payloads into a single contiguous buffer for
    /// batch writing. Copies every entry in <paramref name="pendingWrites"/>
    /// sequentially into <paramref name="destination"/> and returns the total
    /// number of bytes written.
    ///
    /// The caller must ensure <paramref name="destination"/> is large enough
    /// (use <see cref="CalculateBroadcastSize"/> to pre-check).
    ///
    /// Go reference: client.go — broadcast flush coalescing for fan-out.
    /// </summary>
    /// <param name="pendingWrites">Pending write segments to coalesce.</param>
    /// <param name="destination">Destination buffer receiving the concatenated payloads.</param>
    public int BroadcastDrain(IReadOnlyList<ReadOnlyMemory<byte>> pendingWrites, byte[] destination)
    {
        var offset = 0;
        foreach (var write in pendingWrites)
        {
            write.Span.CopyTo(destination.AsSpan(offset));
            offset += write.Length;
        }
        Interlocked.Increment(ref _broadcastCount);
        return offset;
    }

    /// <summary>
    /// Returns the total number of bytes needed to coalesce all
    /// <paramref name="pendingWrites"/> into a single buffer.
    /// </summary>
    /// <param name="pendingWrites">Pending write segments to size.</param>
    public static int CalculateBroadcastSize(IReadOnlyList<ReadOnlyMemory<byte>> pendingWrites)
    {
        var total = 0;
        foreach (var w in pendingWrites) total += w.Length;
        return total;
    }

    // -----------------------------------------------------------------------
    // Inner type: pooled IMemoryOwner<byte>
    // -----------------------------------------------------------------------

    /// <summary>
    /// Wraps a raw <c>byte[]</c> rented from an internal
    /// <see cref="ConcurrentBag{T}"/> and returns it to that bag on disposal.
    /// </summary>
    private sealed class PooledMemoryOwner : IMemoryOwner<byte>
    {
        private readonly ConcurrentBag<byte[]> _pool;
        private byte[]? _buffer;

        /// <summary>
        /// Creates a pooled memory owner backed by a reusable byte array.
        /// </summary>
        /// <param name="buffer">Rented backing buffer.</param>
        /// <param name="pool">Pool to return the buffer to on disposal.</param>
        public PooledMemoryOwner(byte[] buffer, ConcurrentBag<byte[]> pool)
        {
            _buffer = buffer;
            _pool   = pool;
        }

        /// <summary>Memory view over the currently owned buffer.</summary>
        public Memory<byte> Memory =>
            _buffer is { } b ? b.AsMemory() : Memory<byte>.Empty;

        /// <summary>Returns the owned buffer to the originating pool.</summary>
        public void Dispose()
        {
            if (Interlocked.Exchange(ref _buffer, null) is { } b)
                _pool.Add(b);
        }
    }
}