private IAlgorithm GetTokenBucketProcessor(string storageType, int limitNumber)
{
var tokenBucketRules = new TokenBucketRule[]
{
new TokenBucketRule(limitNumber, 100, TimeSpan.FromSeconds(1))
{
Id = Guid.NewGuid().ToString(),
LockSeconds = 1,
ExtractTarget = (request) =>
{
return((request as SimulationRequest).RequestResource);
},
CheckRuleMatching = (request) =>
{
return(true);
},
}
};
if (storageType == "redis")
{
return(new RedisTokenBucketAlgorithm(tokenBucketRules, redisClient));
}
else
{
return(new InProcessTokenBucketAlgorithm(tokenBucketRules));
}
}
private IAlgorithm GetAlgorithm(string storageType, int capacity, int inflowQuantity, TimeSpan inflowUnit, int lockSeconds, StartTimeType startTimeType = StartTimeType.FromCurrent)
{
var tokenBucketRules = new TokenBucketRule[]
{
new TokenBucketRule(capacity, inflowQuantity, inflowUnit)
{
Id = Guid.NewGuid().ToString(),
LockSeconds = lockSeconds,
StartTimeType = startTimeType,
ExtractTarget = (request) =>
{
return((request as SimulationRequest).RequestResource);
},
CheckRuleMatching = (request) =>
{
return(true);
},
ExtractTargetAsync = (request) =>
{
return(Task.FromResult((request as SimulationRequest).RequestResource));
},
CheckRuleMatchingAsync = (request) =>
{
return(Task.FromResult(true));
},
}
};
if (storageType == "redis")
{
var redisClient = StackExchange.Redis.ConnectionMultiplexer.Connect("127.0.0.1");
return(new RedisTokenBucketAlgorithm(tokenBucketRules, redisClient));
}
else
{
return(new InProcessTokenBucketAlgorithm(tokenBucketRules));
}
}
/// <summary>
/// Decrease the count value of the rate limit target for token bucket algorithm.
/// </summary>
/// <param name="target">The target</param>
/// <param name="amount">The amount of decrease</param>
/// <param name="currentRule">The rate limit rule</param>
/// <returns>Amount of token in the bucket</returns>
public Tuple <bool, long> InnerCheckSingleRule(string target, long amount, TokenBucketRule currentRule)
{
bool locked = CheckLocked(target);
if (locked)
{
return(new Tuple <bool, long>(true, -1));
}
var inflowUnit = currentRule.InflowUnit.TotalMilliseconds;
lock (target)
{
var currentTime = _timeProvider.GetCurrentLocalTime();
long bucketAmount = 0;
var result = _cache.GetCacheItem(target);
if (result == null)
{
var startTime = AlgorithmStartTime.ToSpecifiedTypeTime(currentTime, TimeSpan.FromMilliseconds(inflowUnit), currentRule.StartTimeType);
bucketAmount = currentRule.Capacity - amount;
_cache.Add(target, new CountValue(bucketAmount)
{
LastFlowTime = startTime
}, DateTimeOffset.MaxValue);
return(new Tuple <bool, long>(false, bucketAmount));
}
var countValue = (CountValue)result.Value;
var lastTime = countValue.LastFlowTime;
var lastTimeChanged = false;
var pastTime = currentTime - lastTime;
var pastTimeMilliseconds = pastTime.TotalMilliseconds;
// Debug.WriteLine(currentTime.ToString("mm:ss.fff") + "," + lastTime.ToString("mm:ss.fff") + "," + pastTimeMilliseconds);
if (pastTimeMilliseconds < inflowUnit)
{
bucketAmount = countValue.Value - amount;
}
else
{
var pastInflowUnitQuantity = (int)(pastTimeMilliseconds / inflowUnit);
lastTime = lastTime.AddMilliseconds(pastInflowUnitQuantity * inflowUnit);
lastTimeChanged = true;
var pastInflowQuantity = currentRule.InflowQuantityPerUnit * pastInflowUnitQuantity;
bucketAmount = (countValue.Value < 0 ? 0 : countValue.Value) + pastInflowQuantity - amount;
}
if (bucketAmount < 0)
{
if (currentRule.LockSeconds > 0)
{
TryLock(target, currentTime, TimeSpan.FromSeconds(currentRule.LockSeconds));
}
return(new Tuple <bool, long>(true, bucketAmount));
}
if (bucketAmount >= currentRule.Capacity)
{
bucketAmount = currentRule.Capacity - amount;
}
countValue.Value = bucketAmount;
if (lastTimeChanged)
{
countValue.LastFlowTime = lastTime;
}
return(new Tuple <bool, long>(false, countValue.Value));
}
}
