/// <summary>
/// Initializes the <see cref="TokenBucketRateLimiter"/>.
/// </summary>
/// <param name="options">Options to specify the behavior of the <see cref="TokenBucketRateLimiter"/>.</param>
public TokenBucketRateLimiter(TokenBucketRateLimiterOptions options)
{
if (options is null)
{
throw new ArgumentNullException(nameof(options));
}
if (options.TokenLimit <= 0 || options.TokensPerPeriod <= 0)
{
throw new ArgumentException($"Both {nameof(options.TokenLimit)} and {nameof(options.TokensPerPeriod)} must be set to values greater than 0.", nameof(options));
}
if (options.QueueLimit < 0)
{
throw new ArgumentException($"{nameof(options.QueueLimit)} must be set to a value greater than or equal to 0.", nameof(options));
}
if (options.ReplenishmentPeriod < TimeSpan.Zero)
{
throw new ArgumentException($"{nameof(options.ReplenishmentPeriod)} must be set to a value greater than or equal to TimeSpan.Zero.", nameof(options));
}
_options = new TokenBucketRateLimiterOptions
{
TokenLimit = options.TokenLimit,
QueueProcessingOrder = options.QueueProcessingOrder,
QueueLimit = options.QueueLimit,
ReplenishmentPeriod = options.ReplenishmentPeriod,
TokensPerPeriod = options.TokensPerPeriod,
AutoReplenishment = options.AutoReplenishment
};
_tokenCount = options.TokenLimit;
_idleSince = _lastReplenishmentTick = Stopwatch.GetTimestamp();
if (_options.AutoReplenishment)
{
_renewTimer = new Timer(Replenish, this, _options.ReplenishmentPeriod, _options.ReplenishmentPeriod);
}
}
/// <summary>
/// Initializes the <see cref="TokenBucketRateLimiter"/>.
/// </summary>
/// <param name="options">Options to specify the behavior of the <see cref="TokenBucketRateLimiter"/>.</param>
public TokenBucketRateLimiter(TokenBucketRateLimiterOptions options !!)
{
// Used in tests that test behavior with specific time intervals
internal void ReplenishInternal(uint nowTicks)
{
bool wrapped = false;
// (uint)TickCount will wrap every ~50 days, we can detect that by checking if the new ticks is less than the last replenishment
if (nowTicks < _lastReplenishmentTick)
{
wrapped = true;
}
// method is re-entrant (from Timer), lock to avoid multiple simultaneous replenishes
lock (Lock)
{
// Fix the wrapping by using a long and adding uint.MaxValue in the wrapped case
long nonWrappedTicks = wrapped ? (long)nowTicks + uint.MaxValue : nowTicks;
if (nonWrappedTicks - _lastReplenishmentTick < _options.ReplenishmentPeriod.TotalMilliseconds)
{
return;
}
_lastReplenishmentTick = nowTicks;
int availablePermits = _tokenCount;
TokenBucketRateLimiterOptions options = _options;
int maxPermits = options.TokenLimit;
int resourcesToAdd;
if (availablePermits < maxPermits)
{
resourcesToAdd = Math.Min(options.TokensPerPeriod, maxPermits - availablePermits);
}
else
{
// All tokens available, nothing to do
return;
}
// Process queued requests
Deque <RequestRegistration> queue = _queue;
_tokenCount += resourcesToAdd;
Debug.Assert(_tokenCount <= _options.TokenLimit);
while (queue.Count > 0)
{
RequestRegistration nextPendingRequest =
options.QueueProcessingOrder == QueueProcessingOrder.OldestFirst
? queue.PeekHead()
: queue.PeekTail();
if (_tokenCount >= nextPendingRequest.Count)
{
// Request can be fulfilled
nextPendingRequest =
options.QueueProcessingOrder == QueueProcessingOrder.OldestFirst
? queue.DequeueHead()
: queue.DequeueTail();
_queueCount -= nextPendingRequest.Count;
_tokenCount -= nextPendingRequest.Count;
Debug.Assert(_queueCount >= 0);
Debug.Assert(_tokenCount >= 0);
if (!nextPendingRequest.Tcs.TrySetResult(SuccessfulLease))
{
// Queued item was canceled so add count back
_tokenCount += nextPendingRequest.Count;
}
nextPendingRequest.CancellationTokenRegistration.Dispose();
}
else
{
// Request cannot be fulfilled
break;
}
}
}
}
