Skip to content

Commit

Permalink
New: Round-robin over available Download Client instead of the first …
Browse files Browse the repository at this point in the history
…enabled one.
  • Loading branch information
Taloth committed Nov 8, 2017
1 parent 0688340 commit c8695ba
Show file tree
Hide file tree
Showing 3 changed files with 222 additions and 2 deletions.
185 changes: 185 additions & 0 deletions src/NzbDrone.Core.Test/Download/DownloadClientProviderFixture.cs
@@ -0,0 +1,185 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using FizzWare.NBuilder;
using FluentAssertions;
using Moq;
using NUnit.Framework;
using NzbDrone.Core.Download;
using NzbDrone.Core.Indexers;
using NzbDrone.Core.Test.Framework;

namespace NzbDrone.Core.Test.Download
{
[TestFixture]
public class DownloadClientProviderFixture : CoreTest<DownloadClientProvider>
{
private List<IDownloadClient> _downloadClients;
private List<DownloadClientStatus> _blockedProviders;
private int _nextId;

[SetUp]
public void SetUp()
{
_downloadClients = new List<IDownloadClient>();
_blockedProviders = new List<DownloadClientStatus>();
_nextId = 1;

Mocker.GetMock<IDownloadClientFactory>()
.Setup(v => v.GetAvailableProviders())
.Returns(_downloadClients);

Mocker.GetMock<IDownloadClientStatusService>()
.Setup(v => v.GetBlockedProviders())
.Returns(_blockedProviders);
}

private Mock<IDownloadClient> WithUsenetClient()
{
var mock = new Mock<IDownloadClient>(MockBehavior.Default);
mock.SetupGet(s => s.Definition)
.Returns(Builder<DownloadClientDefinition>
.CreateNew()
.With(v => v.Id = _nextId++)
.Build());

_downloadClients.Add(mock.Object);

mock.SetupGet(v => v.Protocol).Returns(DownloadProtocol.Usenet);

return mock;
}

private Mock<IDownloadClient> WithTorrentClient()
{
var mock = new Mock<IDownloadClient>(MockBehavior.Default);
mock.SetupGet(s => s.Definition)
.Returns(Builder<DownloadClientDefinition>
.CreateNew()
.With(v => v.Id = _nextId++)
.Build());

_downloadClients.Add(mock.Object);

mock.SetupGet(v => v.Protocol).Returns(DownloadProtocol.Torrent);

return mock;
}

private void GivenBlockedClient(int id)
{
_blockedProviders.Add(new DownloadClientStatus
{
ProviderId = id,
DisabledTill = DateTime.UtcNow.AddHours(3)
});
}

[Test]
public void should_roundrobin_over_usenet_client()
{
WithUsenetClient();
WithUsenetClient();
WithUsenetClient();
WithTorrentClient();

var client1 = Subject.GetDownloadClient(DownloadProtocol.Usenet);
var client2 = Subject.GetDownloadClient(DownloadProtocol.Usenet);
var client3 = Subject.GetDownloadClient(DownloadProtocol.Usenet);
var client4 = Subject.GetDownloadClient(DownloadProtocol.Usenet);
var client5 = Subject.GetDownloadClient(DownloadProtocol.Usenet);

client1.Definition.Id.Should().Be(1);
client2.Definition.Id.Should().Be(2);
client3.Definition.Id.Should().Be(3);
client4.Definition.Id.Should().Be(1);
client5.Definition.Id.Should().Be(2);
}

[Test]
public void should_roundrobin_over_torrent_client()
{
WithUsenetClient();
WithTorrentClient();
WithTorrentClient();
WithTorrentClient();

var client1 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client2 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client3 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client4 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client5 = Subject.GetDownloadClient(DownloadProtocol.Torrent);

client1.Definition.Id.Should().Be(2);
client2.Definition.Id.Should().Be(3);
client3.Definition.Id.Should().Be(4);
client4.Definition.Id.Should().Be(2);
client5.Definition.Id.Should().Be(3);
}

[Test]
public void should_roundrobin_over_protocol_separately()
{
WithUsenetClient();
WithTorrentClient();
WithTorrentClient();

var client1 = Subject.GetDownloadClient(DownloadProtocol.Usenet);
var client2 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client3 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client4 = Subject.GetDownloadClient(DownloadProtocol.Torrent);

client1.Definition.Id.Should().Be(1);
client2.Definition.Id.Should().Be(2);
client3.Definition.Id.Should().Be(3);
client4.Definition.Id.Should().Be(2);
}

[Test]
public void should_skip_blocked_torrent_client()
{
WithUsenetClient();
WithTorrentClient();
WithTorrentClient();
WithTorrentClient();

GivenBlockedClient(3);

var client1 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client2 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client3 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client4 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client5 = Subject.GetDownloadClient(DownloadProtocol.Torrent);

client1.Definition.Id.Should().Be(2);
client2.Definition.Id.Should().Be(4);
client3.Definition.Id.Should().Be(2);
client4.Definition.Id.Should().Be(4);
}

[Test]
public void should_not_skip_blocked_torrent_client_if_all_blocked()
{
WithUsenetClient();
WithTorrentClient();
WithTorrentClient();
WithTorrentClient();

GivenBlockedClient(2);
GivenBlockedClient(3);
GivenBlockedClient(4);

var client1 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client2 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client3 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client4 = Subject.GetDownloadClient(DownloadProtocol.Torrent);
var client5 = Subject.GetDownloadClient(DownloadProtocol.Torrent);

client1.Definition.Id.Should().Be(2);
client2.Definition.Id.Should().Be(3);
client3.Definition.Id.Should().Be(4);
client4.Definition.Id.Should().Be(2);
}
}
}
1 change: 1 addition & 0 deletions src/NzbDrone.Core.Test/NzbDrone.Core.Test.csproj
Expand Up @@ -171,6 +171,7 @@
<Compile Include="DiskSpace\DiskSpaceServiceFixture.cs" />
<Compile Include="Download\CompletedDownloadServiceFixture.cs" />
<Compile Include="Download\DownloadApprovedReportsTests\DownloadApprovedFixture.cs" />
<Compile Include="Download\DownloadClientProviderFixture.cs" />
<Compile Include="Download\DownloadClientTests\Blackhole\ScanWatchFolderFixture.cs" />
<Compile Include="Download\DownloadClientTests\Blackhole\TorrentBlackholeFixture.cs" />
<Compile Include="Download\DownloadClientTests\Blackhole\UsenetBlackholeFixture.cs" />
Expand Down
38 changes: 36 additions & 2 deletions src/NzbDrone.Core/Download/DownloadClientProvider.cs
@@ -1,6 +1,8 @@
using System.Linq;
using System.Collections.Generic;
using NzbDrone.Core.Indexers;
using NzbDrone.Common.Cache;
using NLog;

namespace NzbDrone.Core.Download
{
Expand All @@ -13,16 +15,48 @@ public interface IProvideDownloadClient

public class DownloadClientProvider : IProvideDownloadClient
{
private readonly Logger _logger;
private readonly IDownloadClientFactory _downloadClientFactory;
private readonly IDownloadClientStatusService _downloadClientStatusService;
private readonly ICached<int> _lastUsedDownloadClient;

public DownloadClientProvider(IDownloadClientFactory downloadClientFactory)
public DownloadClientProvider(IDownloadClientStatusService downloadClientStatusService, IDownloadClientFactory downloadClientFactory, ICacheManager cacheManager, Logger logger)
{
_logger = logger;
_downloadClientFactory = downloadClientFactory;
_downloadClientStatusService = downloadClientStatusService;
_lastUsedDownloadClient = cacheManager.GetCache<int>(GetType(), "lastDownloadClientId");
}

public IDownloadClient GetDownloadClient(DownloadProtocol downloadProtocol)
{
return _downloadClientFactory.GetAvailableProviders().FirstOrDefault(v => v.Protocol == downloadProtocol);
var availableProviders = _downloadClientFactory.GetAvailableProviders().Where(v => v.Protocol == downloadProtocol).ToList();

if (!availableProviders.Any()) return null;

var blockedProviders = new HashSet<int>(_downloadClientStatusService.GetBlockedProviders().Select(v => v.ProviderId));

if (blockedProviders.Any())
{
var nonBlockedProviders = availableProviders.Where(v => !blockedProviders.Contains(v.Definition.Id)).ToList();

if (nonBlockedProviders.Any())
{
availableProviders = nonBlockedProviders;
}
else
{
_logger.Trace("No non-blocked Download Client available, retrying blocked one.");
}
}

var lastId = _lastUsedDownloadClient.Find(downloadProtocol.ToString());

var provider = availableProviders.FirstOrDefault(v => v.Definition.Id > lastId) ?? availableProviders.First();

_lastUsedDownloadClient.Set(downloadProtocol.ToString(), provider.Definition.Id);

return provider;
}

public IEnumerable<IDownloadClient> GetDownloadClients()
Expand Down

0 comments on commit c8695ba

Please sign in to comment.