|
|
|
@ -2,12 +2,14 @@
|
|
|
|
|
using System.Collections.Concurrent;
|
|
|
|
|
using NLog;
|
|
|
|
|
using NzbDrone.Common.Cache;
|
|
|
|
|
using NzbDrone.Common.Extensions;
|
|
|
|
|
|
|
|
|
|
namespace NzbDrone.Common.TPL
|
|
|
|
|
{
|
|
|
|
|
public interface IRateLimitService
|
|
|
|
|
{
|
|
|
|
|
void WaitAndPulse(string key, TimeSpan interval);
|
|
|
|
|
void WaitAndPulse(string key, string subKey, TimeSpan interval);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
public class RateLimitService : IRateLimitService
|
|
|
|
@ -23,9 +25,37 @@ namespace NzbDrone.Common.TPL
|
|
|
|
|
|
|
|
|
|
public void WaitAndPulse(string key, TimeSpan interval)
|
|
|
|
|
{
|
|
|
|
|
var waitUntil = _rateLimitStore.AddOrUpdate(key,
|
|
|
|
|
(s) => DateTime.UtcNow + interval,
|
|
|
|
|
(s, i) => new DateTime(Math.Max(DateTime.UtcNow.Ticks, i.Ticks), DateTimeKind.Utc) + interval);
|
|
|
|
|
WaitAndPulse(key, null, interval);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
public void WaitAndPulse(string key, string subKey, TimeSpan interval)
|
|
|
|
|
{
|
|
|
|
|
var waitUntil = DateTime.UtcNow.Add(interval);
|
|
|
|
|
|
|
|
|
|
if (subKey.IsNotNullOrWhiteSpace())
|
|
|
|
|
{
|
|
|
|
|
// Expand the base key timer, but don't extend it beyond now+interval.
|
|
|
|
|
var baseUntil = _rateLimitStore.AddOrUpdate(key,
|
|
|
|
|
(s) => waitUntil,
|
|
|
|
|
(s, i) => new DateTime(Math.Max(waitUntil.Ticks, i.Ticks), DateTimeKind.Utc));
|
|
|
|
|
|
|
|
|
|
if (baseUntil > waitUntil)
|
|
|
|
|
{
|
|
|
|
|
waitUntil = baseUntil;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Wait for the full key
|
|
|
|
|
var combinedKey = key + "-" + subKey;
|
|
|
|
|
waitUntil = _rateLimitStore.AddOrUpdate(combinedKey,
|
|
|
|
|
(s) => waitUntil,
|
|
|
|
|
(s, i) => new DateTime(Math.Max(waitUntil.Ticks, i.Add(interval).Ticks), DateTimeKind.Utc));
|
|
|
|
|
}
|
|
|
|
|
else
|
|
|
|
|
{
|
|
|
|
|
waitUntil = _rateLimitStore.AddOrUpdate(key,
|
|
|
|
|
(s) => waitUntil,
|
|
|
|
|
(s, i) => new DateTime(Math.Max(waitUntil.Ticks, i.Add(interval).Ticks), DateTimeKind.Utc));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
waitUntil -= interval;
|
|
|
|
|
|
|
|
|
|