public async Task RetryMetadataOnFailedWaitAsync()
{
var options = new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.OldestFirst, 1,
TimeSpan.FromSeconds(20), 1, autoReplenishment: false);
var limiter = new TokenBucketRateLimiter(options);
using var lease = limiter.Acquire(2);
var failedLease = await limiter.WaitAsync(2);
Assert.False(failedLease.IsAcquired);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter.Name, out var metadata));
var metaDataTime = Assert.IsType <TimeSpan>(metadata);
Assert.Equal(options.ReplenishmentPeriod.Ticks * 2, metaDataTime.Ticks);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata));
Assert.Equal(options.ReplenishmentPeriod.Ticks * 2, typedMetadata.Ticks);
Assert.Collection(failedLease.MetadataNames, item => item.Equals(MetadataName.RetryAfter.Name));
}
public override async Task DropsOldestWhenQueuingMoreThanLimit_NewestFirst()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.NewestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
var lease = limiter.Acquire(1);
var wait = limiter.WaitAndAcquireAsync(1);
Assert.False(wait.IsCompleted);
var wait2 = limiter.WaitAndAcquireAsync(1);
var lease1 = await wait;
Assert.False(lease1.IsAcquired);
Assert.False(wait2.IsCompleted);
limiter.TryReplenish();
lease = await wait2;
Assert.True(lease.IsAcquired);
}
public override async Task CannotAcquireResourcesWithAcquireWithQueuedItemsIfOldestFirst()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.OldestFirst, 3,
TimeSpan.Zero, 2, autoReplenishment: false));
var lease = limiter.Acquire(1);
Assert.True(lease.IsAcquired);
var wait = limiter.WaitAsync(2);
Assert.False(wait.IsCompleted);
lease = limiter.Acquire(1);
Assert.False(lease.IsAcquired);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
}
public override async Task CanCancelWaitAsyncBeforeQueuing()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
var lease = limiter.Acquire(1);
Assert.True(lease.IsAcquired);
var cts = new CancellationTokenSource();
cts.Cancel();
var ex = await Assert.ThrowsAsync <TaskCanceledException>(() => limiter.WaitAsync(1, cts.Token).AsTask());
Assert.Equal(cts.Token, ex.CancellationToken);
lease.Dispose();
Assert.True(limiter.TryReplenish());
Assert.Equal(1, limiter.GetAvailablePermits());
}
public async Task ReplenishWorksWhenTicksWrap()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(10, QueueProcessingOrder.OldestFirst, 2,
TimeSpan.FromMilliseconds(2), 1, autoReplenishment: false));
var lease = limiter.Acquire(10);
Assert.True(lease.IsAcquired);
var wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
var replenishInternalMethod = typeof(TokenBucketRateLimiter).GetMethod("ReplenishInternal", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance) !;
// This will set the last tick to the max value
replenishInternalMethod.Invoke(limiter, new object[] { uint.MaxValue });
lease = await wait;
Assert.True(lease.IsAcquired);
wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
// ticks wrapped, should replenish
replenishInternalMethod.Invoke(limiter, new object[] { 2U });
lease = await wait;
Assert.True(lease.IsAcquired);
replenishInternalMethod.Invoke(limiter, new object[] { uint.MaxValue });
wait = limiter.WaitAsync(2);
Assert.False(wait.IsCompleted);
// ticks wrapped, but only 1 millisecond passed, make sure the wrapping behaves correctly and replenish doesn't happen
replenishInternalMethod.Invoke(limiter, new object[] { 1U });
Assert.False(wait.IsCompleted);
Assert.Equal(1, limiter.GetAvailablePermits());
}
public override async Task QueueAvailableAfterQueueLimitHitAndResources_BecomeAvailable()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
var lease = limiter.Acquire(1);
var wait = limiter.WaitAsync(1);
var failedLease = await limiter.WaitAsync(1);
Assert.False(failedLease.IsAcquired);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
}
public override async Task DropsRequestedLeaseIfPermitCountGreaterThanQueueLimitAndNoAvailability_NewestFirst()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.NewestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
var lease = limiter.Acquire(2);
Assert.True(lease.IsAcquired);
// Fill queue
var wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
var lease1 = await limiter.WaitAsync(2);
Assert.False(lease1.IsAcquired);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
}
public async Task ReplenishWorksWhenTicksWrap()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(10, QueueProcessingOrder.OldestFirst, 2,
TimeSpan.FromMilliseconds(2), 1, autoReplenishment: false));
var lease = limiter.Acquire(10);
Assert.True(lease.IsAcquired);
var wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
// This will set the last tick to the max value
limiter.ReplenishInternal(uint.MaxValue);
lease = await wait.DefaultTimeout();
Assert.True(lease.IsAcquired);
wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
// ticks wrapped, should replenish
limiter.ReplenishInternal(2);
lease = await wait.DefaultTimeout();
Assert.True(lease.IsAcquired);
limiter.ReplenishInternal(uint.MaxValue);
wait = limiter.WaitAsync(2);
Assert.False(wait.IsCompleted);
// ticks wrapped, but only 1 millisecond passed, make sure the wrapping behaves correctly and replenish doesn't happen
limiter.ReplenishInternal(1);
Assert.False(wait.IsCompleted);
Assert.Equal(1, limiter.GetAvailablePermits());
}
public override async Task CancelUpdatesQueueLimit()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
var lease = limiter.Acquire(1);
Assert.True(lease.IsAcquired);
var cts = new CancellationTokenSource();
var wait = limiter.WaitAsync(1, cts.Token);
cts.Cancel();
var ex = await Assert.ThrowsAsync <TaskCanceledException>(() => wait.AsTask());
Assert.Equal(cts.Token, ex.CancellationToken);
wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
}
public override async Task CanAcquireResourcesWithWaitAsyncWithQueuedItemsIfNewestFirst()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.NewestFirst, 2,
TimeSpan.Zero, 2, autoReplenishment: false));
var lease = limiter.Acquire(1);
Assert.True(lease.IsAcquired);
var wait = limiter.WaitAsync(2);
Assert.False(wait.IsCompleted);
Assert.Equal(1, limiter.GetAvailablePermits());
lease = await limiter.WaitAsync(1);
Assert.True(lease.IsAcquired);
Assert.False(wait.IsCompleted);
limiter.TryReplenish();
lease = await wait;
Assert.True(lease.IsAcquired);
}
public override async Task CanDequeueMultipleResourcesAtOnce()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.OldestFirst, 2,
TimeSpan.Zero, 2, autoReplenishment: false));
using var lease = await limiter.WaitAsync(2);
Assert.True(lease.IsAcquired);
var wait1 = limiter.WaitAsync(1);
var wait2 = limiter.WaitAsync(1);
Assert.False(wait1.IsCompleted);
Assert.False(wait2.IsCompleted);
lease.Dispose();
Assert.True(limiter.TryReplenish());
var lease1 = await wait1;
var lease2 = await wait2;
Assert.True(lease1.IsAcquired);
Assert.True(lease2.IsAcquired);
}
public override async Task ThrowsWhenWaitingForLessThanZero()
{
var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.NewestFirst, 1,
TimeSpan.Zero, 1, autoReplenishment: false));
await Assert.ThrowsAsync <ArgumentOutOfRangeException>(async() => await limiter.WaitAsync(-1));
}
/// <summary>
/// Constructor.
/// </summary>
/// <param name="localBucket">The local <see cref="TokenBucketRateLimiter"/>.</param>
/// <param name="globalBucket">The shared global <see cref="TokenBucketRateLimiter"/>.</param>
/// <exception cref="ArgumentNullException"><paramref name="localBucket"/> or <paramref name="globalBucket"/> is null.</exception>
public DualTokenBucketRateLimiter(TokenBucketRateLimiter localBucket, TokenBucketRateLimiter globalBucket)
{
this.LocalBucket = localBucket ?? throw new ArgumentNullException(nameof(localBucket));
this.GlobalBucket = globalBucket ?? throw new ArgumentNullException(nameof(globalBucket));
}
