From 0758a27d5b0f0e3c01a5d34c59f066886726a5b6 Mon Sep 17 00:00:00 2001 From: Robin Dadswell <19610103+RobinDadswell@users.noreply.github.com> Date: Sun, 7 Mar 2021 23:28:37 +0000 Subject: [PATCH] Generalized RateLimit logic to all indexers based on indexer id Fixes #1982 Co-authored-by: Taloth Saldono --- .../TPLTests/RateLimitServiceFixture.cs | 33 +++++++++++++++++ src/NzbDrone.Common/Http/HttpClient.cs | 2 +- src/NzbDrone.Common/Http/HttpRequest.cs | 1 + src/NzbDrone.Common/TPL/RateLimitService.cs | 36 +++++++++++++++++-- .../Download/TorrentClientBase.cs | 1 + .../Download/UsenetClientBase.cs | 1 + src/NzbDrone.Core/Indexers/HttpIndexerBase.cs | 2 ++ 7 files changed, 72 insertions(+), 4 deletions(-) diff --git a/src/NzbDrone.Common.Test/TPLTests/RateLimitServiceFixture.cs b/src/NzbDrone.Common.Test/TPLTests/RateLimitServiceFixture.cs index 919868458..a92eed1f1 100644 --- a/src/NzbDrone.Common.Test/TPLTests/RateLimitServiceFixture.cs +++ b/src/NzbDrone.Common.Test/TPLTests/RateLimitServiceFixture.cs @@ -89,5 +89,38 @@ namespace NzbDrone.Common.Test.TPLTests (GetRateLimitStore()["me"] - _epoch).Should().BeGreaterOrEqualTo(TimeSpan.FromMilliseconds(100)); } + + [Test] + public void should_extend_subkey_delay() + { + GivenExisting("me", _epoch + TimeSpan.FromMilliseconds(200)); + GivenExisting("me-sub", _epoch + TimeSpan.FromMilliseconds(300)); + + Subject.WaitAndPulse("me", "sub", TimeSpan.FromMilliseconds(100)); + + (GetRateLimitStore()["me-sub"] - _epoch).Should().BeGreaterOrEqualTo(TimeSpan.FromMilliseconds(400)); + } + + [Test] + public void should_honor_basekey_delay() + { + GivenExisting("me", _epoch + TimeSpan.FromMilliseconds(200)); + GivenExisting("me-sub", _epoch + TimeSpan.FromMilliseconds(0)); + + Subject.WaitAndPulse("me", "sub", TimeSpan.FromMilliseconds(100)); + + (GetRateLimitStore()["me-sub"] - _epoch).Should().BeGreaterOrEqualTo(TimeSpan.FromMilliseconds(200)); + } + + [Test] + public void should_not_extend_basekey_delay() + { + GivenExisting("me", _epoch + TimeSpan.FromMilliseconds(200)); + GivenExisting("me-sub", _epoch + TimeSpan.FromMilliseconds(100)); + + Subject.WaitAndPulse("me", "sub", TimeSpan.FromMilliseconds(100)); + + (GetRateLimitStore()["me"] - _epoch).Should().BeCloseTo(TimeSpan.FromMilliseconds(200)); + } } } diff --git a/src/NzbDrone.Common/Http/HttpClient.cs b/src/NzbDrone.Common/Http/HttpClient.cs index ce19b1056..5c23af7c8 100644 --- a/src/NzbDrone.Common/Http/HttpClient.cs +++ b/src/NzbDrone.Common/Http/HttpClient.cs @@ -112,7 +112,7 @@ namespace NzbDrone.Common.Http if (request.RateLimit != TimeSpan.Zero) { - _rateLimitService.WaitAndPulse(request.Url.Host, request.RateLimit); + _rateLimitService.WaitAndPulse(request.Url.Host, request.RateLimitKey, request.RateLimit); } _logger.Trace(request); diff --git a/src/NzbDrone.Common/Http/HttpRequest.cs b/src/NzbDrone.Common/Http/HttpRequest.cs index 30c554427..99bec5f8f 100644 --- a/src/NzbDrone.Common/Http/HttpRequest.cs +++ b/src/NzbDrone.Common/Http/HttpRequest.cs @@ -47,6 +47,7 @@ namespace NzbDrone.Common.Http public bool StoreResponseCookie { get; set; } public TimeSpan RequestTimeout { get; set; } public TimeSpan RateLimit { get; set; } + public string RateLimitKey { get; set; } public override string ToString() { diff --git a/src/NzbDrone.Common/TPL/RateLimitService.cs b/src/NzbDrone.Common/TPL/RateLimitService.cs index f0d30b4ff..87b0ff22f 100644 --- a/src/NzbDrone.Common/TPL/RateLimitService.cs +++ b/src/NzbDrone.Common/TPL/RateLimitService.cs @@ -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; diff --git a/src/NzbDrone.Core/Download/TorrentClientBase.cs b/src/NzbDrone.Core/Download/TorrentClientBase.cs index f359e1a51..97b8c3b14 100644 --- a/src/NzbDrone.Core/Download/TorrentClientBase.cs +++ b/src/NzbDrone.Core/Download/TorrentClientBase.cs @@ -128,6 +128,7 @@ namespace NzbDrone.Core.Download try { var request = new HttpRequest(torrentUrl); + request.RateLimitKey = remoteAlbum?.Release?.IndexerId.ToString(); request.Headers.Accept = "application/x-bittorrent"; request.AllowAutoRedirect = false; diff --git a/src/NzbDrone.Core/Download/UsenetClientBase.cs b/src/NzbDrone.Core/Download/UsenetClientBase.cs index 92d649df6..6bbd22726 100644 --- a/src/NzbDrone.Core/Download/UsenetClientBase.cs +++ b/src/NzbDrone.Core/Download/UsenetClientBase.cs @@ -45,6 +45,7 @@ namespace NzbDrone.Core.Download try { var nzbDataRequest = new HttpRequest(url); + nzbDataRequest.RateLimitKey = remoteAlbum?.Release?.IndexerId.ToString(); // TODO: Look into moving download request handling to indexer if (remoteAlbum.Release.BasicAuthString.IsNotNullOrWhiteSpace()) diff --git a/src/NzbDrone.Core/Indexers/HttpIndexerBase.cs b/src/NzbDrone.Core/Indexers/HttpIndexerBase.cs index 8be8e9aad..9b97a25c9 100644 --- a/src/NzbDrone.Core/Indexers/HttpIndexerBase.cs +++ b/src/NzbDrone.Core/Indexers/HttpIndexerBase.cs @@ -280,6 +280,8 @@ namespace NzbDrone.Core.Indexers request.HttpRequest.RateLimit = RateLimit; } + request.HttpRequest.RateLimitKey = Definition.Id.ToString(); + return new IndexerResponse(request, _httpClient.Execute(request.HttpRequest)); }