feat: add token bucket rate limiter for consumers (Gap 3.13)

Implements TokenBucketRateLimiter with refill-over-time semantics,
TryConsume/EstimateWait/WaitForTokensAsync API, and dynamic rate updates.
12 tests covering all behaviors including SW004-suppressed refill timing test.
This commit is contained in:
Joseph Doherty
2026-02-25 11:15:58 -05:00
parent 778687cf6f
commit aad9cf17e4
2 changed files with 339 additions and 0 deletions

View File

@@ -0,0 +1,130 @@
// Go: consumer.go (rateLimitBps config, rate limiting in consumer delivery)
namespace NATS.Server.JetStream.Consumers;
/// <summary>
/// Token bucket rate limiter for consumer message delivery.
/// Tokens refill at a configurable rate (bytes per second).
/// Each message consumes tokens equal to its payload size.
/// Go reference: consumer.go rate limiting via rateLimitBps config.
/// </summary>
public sealed class TokenBucketRateLimiter
{
private readonly object _lock = new();
private double _tokens;
private double _maxTokens;
private double _refillRate; // tokens per millisecond
private DateTime _lastRefill;
/// <summary>
/// Creates a rate limiter with the specified rate in bytes per second.
/// </summary>
/// <param name="bytesPerSecond">Maximum bytes per second. 0 = unlimited.</param>
/// <param name="burstSize">Maximum burst size in bytes. Defaults to 2x rate.</param>
public TokenBucketRateLimiter(long bytesPerSecond, long burstSize = 0)
{
BytesPerSecond = bytesPerSecond;
_refillRate = bytesPerSecond / 1000.0; // tokens per ms
_maxTokens = burstSize > 0 ? burstSize : bytesPerSecond * 2;
_tokens = _maxTokens; // Start full
_lastRefill = DateTime.UtcNow;
}
/// <summary>Configured rate in bytes per second.</summary>
public long BytesPerSecond { get; private set; }
/// <summary>Current available tokens (approximate).</summary>
public double AvailableTokens
{
get
{
lock (_lock)
{
Refill();
return _tokens;
}
}
}
/// <summary>
/// Tries to consume the specified number of tokens (bytes).
/// Returns true if tokens were available (message can be sent).
/// Returns false if not enough tokens (caller should wait).
/// </summary>
public bool TryConsume(long bytes)
{
if (BytesPerSecond <= 0) return true; // Unlimited
lock (_lock)
{
Refill();
if (_tokens >= bytes)
{
_tokens -= bytes;
return true;
}
return false;
}
}
/// <summary>
/// Returns the estimated wait time until enough tokens are available.
/// </summary>
public TimeSpan EstimateWait(long bytes)
{
if (BytesPerSecond <= 0) return TimeSpan.Zero;
lock (_lock)
{
Refill();
if (_tokens >= bytes) return TimeSpan.Zero;
var deficit = bytes - _tokens;
var waitMs = deficit / _refillRate;
return TimeSpan.FromMilliseconds(waitMs);
}
}
/// <summary>
/// Waits until enough tokens are available, then consumes them.
/// </summary>
public async ValueTask WaitForTokensAsync(long bytes, CancellationToken ct = default)
{
if (BytesPerSecond <= 0) return;
while (!ct.IsCancellationRequested)
{
if (TryConsume(bytes)) return;
var wait = EstimateWait(bytes);
if (wait > TimeSpan.Zero)
await Task.Delay(wait, ct).ConfigureAwait(false);
}
ct.ThrowIfCancellationRequested();
}
/// <summary>
/// Updates the rate dynamically.
/// Go reference: consumer.go — rate can change on config update.
/// </summary>
public void UpdateRate(long bytesPerSecond, long burstSize = 0)
{
lock (_lock)
{
BytesPerSecond = bytesPerSecond;
_refillRate = bytesPerSecond / 1000.0;
_maxTokens = burstSize > 0 ? burstSize : bytesPerSecond * 2;
_tokens = Math.Min(_tokens, _maxTokens);
}
}
private void Refill()
{
var now = DateTime.UtcNow;
var elapsed = (now - _lastRefill).TotalMilliseconds;
if (elapsed <= 0) return;
_tokens = Math.Min(_maxTokens, _tokens + elapsed * _refillRate);
_lastRefill = now;
}
}