public void Create_TokenBucket()
{
var options = new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 10, TimeSpan.FromMinutes(1), 1, true);
var partition = RateLimitPartition.CreateTokenBucketLimiter(1, key => options);
var factoryProperty = typeof(RateLimitPartition <int>).GetField("Factory", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance) !;
var factory = (Func <int, RateLimiter>)factoryProperty.GetValue(partition);
var limiter = factory(1);
var tokenBucketLimiter = Assert.IsType <TokenBucketRateLimiter>(limiter);
Assert.Equal(options.TokenLimit, tokenBucketLimiter.GetAvailablePermits());
// TODO: Check other properties when ReplenshingRateLimiter is merged
// TODO: Check that autoReplenishment: true got changed to false
}
public async Task CorrectRetryMetadataWithNonZeroAvailableItems()
{
var options = new TokenBucketRateLimiterOptions(3, 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(3);
Assert.False(failedLease.IsAcquired);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata));
Assert.Equal(options.ReplenishmentPeriod.Ticks * 2, typedMetadata.Ticks);
}
/// <summary>
/// Defines a partition with a <see cref="TokenBucketRateLimiter"/> with the given <see cref="TokenBucketRateLimiterOptions"/>.
/// </summary>
/// <remarks>
/// Set <see cref="TokenBucketRateLimiterOptions.AutoReplenishment"/> to <see langword="false"/> to save an allocation. This method will create a new options type and set <see cref="TokenBucketRateLimiterOptions.AutoReplenishment"/> to <see langword="false"/> otherwise.
/// </remarks>
/// <typeparam name="TKey">The type to distinguish partitions with.</typeparam>
/// <param name="partitionKey">The specific key for this partition.</param>
/// <param name="factory">The function called when a rate limiter for the given <paramref name="partitionKey"/> is needed. This can return the same instance of <see cref="TokenBucketRateLimiterOptions"/> across different calls.</param>
/// <returns></returns>
public static RateLimitPartition <TKey> GetTokenBucketLimiter <TKey>(
TKey partitionKey,
Func <TKey, TokenBucketRateLimiterOptions> factory)
{
return(Get(partitionKey, key =>
{
TokenBucketRateLimiterOptions options = factory(key);
// We don't want individual TokenBucketRateLimiters to have timers. We will instead have our own internal Timer handling all of them
if (options.AutoReplenishment is true)
{
options = new TokenBucketRateLimiterOptions(options.TokenLimit, options.QueueProcessingOrder, options.QueueLimit,
options.ReplenishmentPeriod, options.TokensPerPeriod, autoReplenishment: false);
}
return new TokenBucketRateLimiter(options);
}));
}
public async Task CorrectRetryMetadataWithLargeTokensPerPeriod()
{
var options = new TokenBucketRateLimiterOptions(2, QueueProcessingOrder.OldestFirst, 1,
TimeSpan.FromSeconds(20), 100, autoReplenishment: false);
var limiter = new TokenBucketRateLimiter(options);
using var lease = limiter.Acquire(2);
// Queue item which changes the retry after time for failed items
var wait = limiter.WaitAsync(1);
Assert.False(wait.IsCompleted);
var failedLease = await limiter.WaitAsync(2);
Assert.False(failedLease.IsAcquired);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata));
Assert.Equal(options.ReplenishmentPeriod, typedMetadata);
}
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).DefaultTimeout();
Assert.False(failedLease.IsAcquired);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter.Name, out var metadata));
Assert.Equal(options.ReplenishmentPeriod * 2, metadata);
Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata));
Assert.Equal(options.ReplenishmentPeriod * 2, typedMetadata);
Assert.Collection(failedLease.MetadataNames, item => item.Equals(MetadataName.RetryAfter.Name));
}
/// <summary>
/// Adds a new <see cref="TokenBucketRateLimiter"/> with the given <see cref="TokenBucketRateLimiterOptions"/> to the <see cref="RateLimiterOptions"/>.
/// </summary>
/// <param name="options">The <see cref="RateLimiterOptions"/> to add a limiter to.</param>
/// <param name="policyName">The name that will be associated with the limiter.</param>
/// <param name="configureOptions">A callback to configure the <see cref="TokenBucketRateLimiterOptions"/> to be used for the limiter.</param>
/// <returns>This <see cref="RateLimiterOptions"/>.</returns>
public static RateLimiterOptions AddTokenBucketLimiter(this RateLimiterOptions options, string policyName, Action <TokenBucketRateLimiterOptions> configureOptions)
{
ArgumentNullException.ThrowIfNull(configureOptions);
var key = new PolicyNameKey()
{
PolicyName = policyName
};
var tokenBucketRateLimiterOptions = new TokenBucketRateLimiterOptions();
configureOptions.Invoke(tokenBucketRateLimiterOptions);
// Saves an allocation in GetTokenBucketLimiter, which would have created a new set of options if this was true.
tokenBucketRateLimiterOptions.AutoReplenishment = false;
return(options.AddPolicy(policyName, context =>
{
return RateLimitPartition.GetTokenBucketLimiter(key,
_ => tokenBucketRateLimiterOptions);
}));
}
public void Create_TokenBucket()
{
var options = new TokenBucketRateLimiterOptions
{
TokenLimit = 1,
QueueProcessingOrder = QueueProcessingOrder.OldestFirst,
QueueLimit = 10,
ReplenishmentPeriod = TimeSpan.FromMinutes(1),
TokensPerPeriod = 1,
AutoReplenishment = true
};
var partition = RateLimitPartition.GetTokenBucketLimiter(1, key => options);
var limiter = partition.Factory(1);
var tokenBucketLimiter = Assert.IsType <TokenBucketRateLimiter>(limiter);
Assert.Equal(options.TokenLimit, tokenBucketLimiter.GetStatistics().CurrentAvailablePermits);
Assert.Equal(options.ReplenishmentPeriod, tokenBucketLimiter.ReplenishmentPeriod);
Assert.False(tokenBucketLimiter.IsAutoReplenishing);
}
