feat: add dynamic write buffer pooling with broadcast drain (Gap 5.6)

Enhances OutboundBufferPool with tiered internal pools (512/4096/65536),
RentBuffer/ReturnBuffer raw-array surface, BroadcastDrain coalescing for
fan-out publish, and Interlocked stats counters (RentCount, ReturnCount,
BroadcastCount). Adds 10 DynamicBufferPoolTests covering all new paths.

src/NATS.Server/IO/OutboundBufferPool.cs

@@ -1,15 +1,182 @@
 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;
+
+    public long RentCount      => Interlocked.Read(ref _rentCount);
+    public long ReturnCount    => Interlocked.Read(ref _returnCount);
+    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>
     public IMemoryOwner<byte> Rent(int size)
     {
-        if (size <= 512)
-            return MemoryPool<byte>.Shared.Rent(512);
-        if (size <= 4096)
-            return MemoryPool<byte>.Shared.Rent(4096);
-        return MemoryPool<byte>.Shared.Rent(64 * 1024);
+        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>
+    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>
+    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>
+    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>
+    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;
+
+        public PooledMemoryOwner(byte[] buffer, ConcurrentBag<byte[]> pool)
+        {
+            _buffer = buffer;
+            _pool   = pool;
+        }
+
+        public Memory<byte> Memory =>
+            _buffer is { } b ? b.AsMemory() : Memory<byte>.Empty;
+
+        public void Dispose()
+        {
+            if (Interlocked.Exchange(ref _buffer, null) is { } b)
+                _pool.Add(b);
+        }
     }
 }