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:
130
src/NATS.Server/JetStream/Consumers/TokenBucketRateLimiter.cs
Normal file
130
src/NATS.Server/JetStream/Consumers/TokenBucketRateLimiter.cs
Normal 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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user