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.
This commit is contained in:
Joseph Doherty
2026-02-25 11:31:29 -05:00
parent 1a1e99f7d8
commit bc8f0e63bb
2 changed files with 351 additions and 5 deletions

View File

@@ -0,0 +1,179 @@
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);
}
}