From 136e8b1a8a732b20e13d62557cde88ffb903efae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Sharma?= Date: Fri, 1 May 2020 15:40:39 -0700 Subject: [PATCH] [Package Renames] Add feature flag to disable popularity transfers (#773) Adds the feature flag `Search.TransferPopularity` that, if disabled, removes all popularity transfers on the search index. Note that download overrides are still applied even if the popularity transfer feature flag is disabled. Addresses https://github.com/NuGet/NuGetGallery/issues/7944 The feature flag staleness monitoring is tracked by https://github.com/NuGet/NuGetGallery/issues/7966 --- src/NuGet.Jobs.Db2AzureSearch/Job.cs | 1 - .../UpdateDownloadsCommand.cs | 18 ++++- .../AzureSearchJob.cs | 10 +++ .../NewPackageRegistrationProducer.cs | 18 ++++- .../DependencyInjectionExtensions.cs | 5 ++ .../FeatureFlagService.cs | 25 +++++++ .../IFeatureFlagService.cs | 10 +++ .../NuGet.Services.AzureSearch.csproj | 2 + .../DependencyInjectionExtensions.cs | 68 ++++++++++++++++++ src/NuGet.Services.V3/V3TelemetryService.cs | 10 ++- .../PopularityTransferIntegrationTests.cs | 63 ++++++++++++++++- .../UpdateDownloadsCommandFacts.cs | 70 +++++++++++++++++++ .../NewPackageRegistrationProducerFacts.cs | 68 ++++++++++++++++++ 13 files changed, 363 insertions(+), 5 deletions(-) create mode 100644 src/NuGet.Services.AzureSearch/FeatureFlagService.cs create mode 100644 src/NuGet.Services.AzureSearch/IFeatureFlagService.cs diff --git a/src/NuGet.Jobs.Db2AzureSearch/Job.cs b/src/NuGet.Jobs.Db2AzureSearch/Job.cs index baef374ed..7ee5090af 100644 --- a/src/NuGet.Jobs.Db2AzureSearch/Job.cs +++ b/src/NuGet.Jobs.Db2AzureSearch/Job.cs @@ -3,7 +3,6 @@ using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; -using Microsoft.Extensions.Options; using NuGet.Services.AzureSearch; using NuGet.Services.AzureSearch.AuxiliaryFiles; using NuGet.Services.AzureSearch.Db2AzureSearch; diff --git a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/UpdateDownloadsCommand.cs b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/UpdateDownloadsCommand.cs index dbb5229a9..526100852 100644 --- a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/UpdateDownloadsCommand.cs +++ b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/UpdateDownloadsCommand.cs @@ -37,6 +37,7 @@ public class UpdateDownloadsCommand : IAzureSearchCommand private readonly ISearchIndexActionBuilder _indexActionBuilder; private readonly Func _batchPusherFactory; private readonly ISystemTime _systemTime; + private readonly IFeatureFlagService _featureFlags; private readonly IOptionsSnapshot _options; private readonly IAzureSearchTelemetryService _telemetryService; private readonly ILogger _logger; @@ -53,6 +54,7 @@ public UpdateDownloadsCommand( ISearchIndexActionBuilder indexActionBuilder, Func batchPusherFactory, ISystemTime systemTime, + IFeatureFlagService featureFlags, IOptionsSnapshot options, IAzureSearchTelemetryService telemetryService, ILogger logger) @@ -67,6 +69,7 @@ public UpdateDownloadsCommand( _indexActionBuilder = indexActionBuilder ?? throw new ArgumentNullException(nameof(indexActionBuilder)); _batchPusherFactory = batchPusherFactory ?? throw new ArgumentNullException(nameof(batchPusherFactory)); _systemTime = systemTime ?? throw new ArgumentNullException(nameof(systemTime)); + _featureFlags = featureFlags ?? throw new ArgumentNullException(nameof(featureFlags)); _options = options ?? throw new ArgumentNullException(nameof(options)); _telemetryService = telemetryService ?? throw new ArgumentNullException(nameof(telemetryService)); _logger = logger ?? throw new ArgumentNullException(nameof(logger)); @@ -132,7 +135,7 @@ private async Task PushIndexChangesAsync() // The "new" data is the latest popularity transfers data from the database. _logger.LogInformation("Fetching new popularity transfer data from database."); - var newTransfers = await _databaseFetcher.GetPackageIdToPopularityTransfersAsync(); + var newTransfers = await GetPopularityTransfersAsync(); _logger.LogInformation("Fetching new download overrides from blob storage."); var downloadOverrides = await _auxiliaryFileClient.LoadDownloadOverridesAsync(); @@ -171,6 +174,19 @@ await _popularityTransferDataClient.ReplaceLatestIndexedAsync( return true; } + private async Task>> GetPopularityTransfersAsync() + { + if (!_featureFlags.IsPopularityTransferEnabled()) + { + _logger.LogWarning( + "Popularity transfers feature flag is disabled. " + + "All popularity transfers will be removed."); + return new SortedDictionary>(StringComparer.OrdinalIgnoreCase); + } + + return await _databaseFetcher.GetPackageIdToPopularityTransfersAsync(); + } + private void ApplyDownloadTransfers( DownloadData newData, SortedDictionary> oldTransfers, diff --git a/src/NuGet.Services.AzureSearch/AzureSearchJob.cs b/src/NuGet.Services.AzureSearch/AzureSearchJob.cs index bd90ee19d..97223bbd5 100644 --- a/src/NuGet.Services.AzureSearch/AzureSearchJob.cs +++ b/src/NuGet.Services.AzureSearch/AzureSearchJob.cs @@ -8,16 +8,22 @@ using Microsoft.Extensions.DependencyInjection; using Microsoft.Rest; using NuGet.Jobs; +using NuGet.Jobs.Validation; namespace NuGet.Services.AzureSearch { public abstract class AzureSearchJob : JsonConfigurationJob where T : IAzureSearchCommand { + private const string FeatureFlagConfigurationSectionName = "FeatureFlags"; + public override async Task Run() { ServicePointManager.DefaultConnectionLimit = 64; ServicePointManager.MaxServicePointIdleTime = 10000; + var featureFlagRefresher = _serviceProvider.GetRequiredService(); + await featureFlagRefresher.StartIfConfiguredAsync(); + var tracingInterceptor = _serviceProvider.GetRequiredService(); try { @@ -30,6 +36,8 @@ public override async Task Run() { ServiceClientTracing.RemoveTracingInterceptor(tracingInterceptor); } + + await featureFlagRefresher.StopAndWaitAsync(); } protected override void ConfigureAutofacServices(ContainerBuilder containerBuilder) @@ -40,6 +48,8 @@ protected override void ConfigureAutofacServices(ContainerBuilder containerBuild protected override void ConfigureJobServices(IServiceCollection services, IConfigurationRoot configurationRoot) { services.AddAzureSearch(GlobalTelemetryDimensions); + + services.Configure(configurationRoot.GetSection(FeatureFlagConfigurationSectionName)); } } } diff --git a/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs b/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs index 1b847380b..a2fb11de1 100644 --- a/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs +++ b/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs @@ -23,6 +23,7 @@ public class NewPackageRegistrationProducer : INewPackageRegistrationProducer private readonly IAuxiliaryFileClient _auxiliaryFileClient; private readonly IDatabaseAuxiliaryDataFetcher _databaseFetcher; private readonly IDownloadTransferrer _downloadTransferrer; + private readonly IFeatureFlagService _featureFlags; private readonly IOptionsSnapshot _options; private readonly IOptionsSnapshot _developmentOptions; private readonly ILogger _logger; @@ -32,6 +33,7 @@ public NewPackageRegistrationProducer( IAuxiliaryFileClient auxiliaryFileClient, IDatabaseAuxiliaryDataFetcher databaseFetcher, IDownloadTransferrer downloadTransferrer, + IFeatureFlagService featureFlags, IOptionsSnapshot options, IOptionsSnapshot developmentOptions, ILogger logger) @@ -40,6 +42,7 @@ public NewPackageRegistrationProducer( _auxiliaryFileClient = auxiliaryFileClient ?? throw new ArgumentNullException(nameof(auxiliaryFileClient)); _databaseFetcher = databaseFetcher ?? throw new ArgumentNullException(nameof(databaseFetcher)); _downloadTransferrer = downloadTransferrer ?? throw new ArgumentNullException(nameof(downloadTransferrer)); + _featureFlags = featureFlags ?? throw new ArgumentNullException(nameof(featureFlags)); _options = options ?? throw new ArgumentNullException(nameof(options)); _developmentOptions = developmentOptions ?? throw new ArgumentNullException(nameof(developmentOptions)); _logger = logger ?? throw new ArgumentNullException(nameof(logger)); @@ -64,7 +67,7 @@ public async Task ProduceWorkAsync( // auxiliary file. var downloads = await _auxiliaryFileClient.LoadDownloadDataAsync(); - var popularityTransfers = await _databaseFetcher.GetPackageIdToPopularityTransfersAsync(); + var popularityTransfers = await GetPopularityTransfersAsync(); var downloadOverrides = await _auxiliaryFileClient.LoadDownloadOverridesAsync(); // Apply changes from popularity transfers and download overrides. @@ -160,6 +163,19 @@ private bool ShouldWait(ConcurrentBag allWork, bool log) return false; } + private async Task>> GetPopularityTransfersAsync() + { + if (!_featureFlags.IsPopularityTransferEnabled()) + { + _logger.LogWarning( + "Popularity transfers feature flag is disabled. " + + "Popularity transfers will be ignored."); + return new SortedDictionary>(StringComparer.OrdinalIgnoreCase); + } + + return await _databaseFetcher.GetPackageIdToPopularityTransfersAsync(); + } + private Dictionary GetTransferredDownloads( DownloadData downloads, SortedDictionary> popularityTransfers, diff --git a/src/NuGet.Services.AzureSearch/DependencyInjectionExtensions.cs b/src/NuGet.Services.AzureSearch/DependencyInjectionExtensions.cs index 839aa9fcb..ff19bd972 100644 --- a/src/NuGet.Services.AzureSearch/DependencyInjectionExtensions.cs +++ b/src/NuGet.Services.AzureSearch/DependencyInjectionExtensions.cs @@ -29,6 +29,8 @@ public static class DependencyInjectionExtensions { public static ContainerBuilder AddAzureSearch(this ContainerBuilder containerBuilder) { + containerBuilder.AddFeatureFlags(); + /// Here, we register services that depend on an interface that there are multiple implementations. /// There are multiple implementations of . @@ -226,6 +228,9 @@ public static IServiceCollection AddAzureSearch( { services.AddV3(telemetryGlobalDimensions); + services.AddFeatureFlags(); + services.AddTransient(); + services .AddTransient(p => { diff --git a/src/NuGet.Services.AzureSearch/FeatureFlagService.cs b/src/NuGet.Services.AzureSearch/FeatureFlagService.cs new file mode 100644 index 000000000..8669c4f48 --- /dev/null +++ b/src/NuGet.Services.AzureSearch/FeatureFlagService.cs @@ -0,0 +1,25 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. + +using System; +using NuGet.Services.FeatureFlags; + +namespace NuGet.Services.AzureSearch +{ + public class FeatureFlagService : IFeatureFlagService + { + private const string SearchPrefix = "Search."; + + private readonly IFeatureFlagClient _featureFlagClient; + + public FeatureFlagService(IFeatureFlagClient featureFlagClient) + { + _featureFlagClient = featureFlagClient ?? throw new ArgumentNullException(nameof(featureFlagClient)); + } + + public bool IsPopularityTransferEnabled() + { + return _featureFlagClient.IsEnabled(SearchPrefix + "TransferPopularity", defaultValue: true); + } + } +} diff --git a/src/NuGet.Services.AzureSearch/IFeatureFlagService.cs b/src/NuGet.Services.AzureSearch/IFeatureFlagService.cs new file mode 100644 index 000000000..5c038196d --- /dev/null +++ b/src/NuGet.Services.AzureSearch/IFeatureFlagService.cs @@ -0,0 +1,10 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. + +namespace NuGet.Services.AzureSearch +{ + public interface IFeatureFlagService + { + bool IsPopularityTransferEnabled(); + } +} diff --git a/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj b/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj index 91b53a4bb..d71b8baf9 100644 --- a/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj +++ b/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj @@ -81,11 +81,13 @@ + + diff --git a/src/NuGet.Services.V3/DependencyInjectionExtensions.cs b/src/NuGet.Services.V3/DependencyInjectionExtensions.cs index aba36dd59..13c95122b 100644 --- a/src/NuGet.Services.V3/DependencyInjectionExtensions.cs +++ b/src/NuGet.Services.V3/DependencyInjectionExtensions.cs @@ -1,22 +1,35 @@ // Copyright (c) .NET Foundation. All rights reserved. // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. +using System; using System.Collections.Generic; using System.Net; using System.Net.Http; +using Autofac; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Options; +using Microsoft.WindowsAzure.Storage.Blob; +using Microsoft.WindowsAzure.Storage.RetryPolicies; using NuGet.Jobs.Validation; using NuGet.Protocol.Catalog; using NuGet.Protocol.Registration; +using NuGet.Services.FeatureFlags; using NuGet.Services.Logging; using NuGet.Services.Metadata.Catalog; +using NuGetGallery; using NuGetGallery.Diagnostics; +using NuGetGallery.Features; namespace NuGet.Services.V3 { public static class DependencyInjectionExtensions { + private const string FeatureFlagBindingKey = nameof(FeatureFlagBindingKey); + + private static readonly TimeSpan FeatureFlagServerTimeout = TimeSpan.FromSeconds(30); + private static readonly TimeSpan FeatureFlagMaxExecutionTime = TimeSpan.FromMinutes(10); + public static IServiceCollection AddV3(this IServiceCollection services, IDictionary telemetryGlobalDimensions) { services @@ -50,5 +63,60 @@ public static IServiceCollection AddV3(this IServiceCollection services, IDictio return services; } + + public static void AddFeatureFlags(this ContainerBuilder containerBuilder) + { + containerBuilder + .Register(c => + { + var options = c.Resolve>(); + var requestOptions = new BlobRequestOptions + { + ServerTimeout = FeatureFlagServerTimeout, + MaximumExecutionTime = FeatureFlagMaxExecutionTime, + LocationMode = LocationMode.PrimaryThenSecondary, + RetryPolicy = new ExponentialRetry(), + }; + + return new CloudBlobClientWrapper( + options.Value.ConnectionString, + requestOptions); + }) + .Keyed(FeatureFlagBindingKey); + + containerBuilder + .Register(c => new CloudBlobCoreFileStorageService( + c.ResolveKeyed(FeatureFlagBindingKey), + c.Resolve(), + c.Resolve())) + .Keyed(FeatureFlagBindingKey); + + containerBuilder + .Register(c => new FeatureFlagFileStorageService( + c.ResolveKeyed(FeatureFlagBindingKey))) + .As(); + } + + public static IServiceCollection AddFeatureFlags(this IServiceCollection services) + { + services + .AddTransient(p => + { + var options = p.GetRequiredService>(); + return new FeatureFlagOptions + { + RefreshInterval = options.Value.RefreshInternal, + }; + }); + + services.AddTransient(); + services.AddTransient(); + services.AddTransient(); + + services.AddSingleton(); + services.AddSingleton(); + + return services; + } } } diff --git a/src/NuGet.Services.V3/V3TelemetryService.cs b/src/NuGet.Services.V3/V3TelemetryService.cs index d91d092e9..89a35cac5 100644 --- a/src/NuGet.Services.V3/V3TelemetryService.cs +++ b/src/NuGet.Services.V3/V3TelemetryService.cs @@ -3,11 +3,12 @@ using System; using System.Collections.Generic; +using NuGet.Services.FeatureFlags; using NuGet.Services.Logging; namespace NuGet.Services.V3 { - public class V3TelemetryService : IV3TelemetryService + public class V3TelemetryService : IV3TelemetryService, IFeatureFlagTelemetryService { private const string Prefix = "V3."; @@ -27,5 +28,12 @@ public IDisposable TrackCatalogLeafDownloadBatch(int count) { "Count", count.ToString() }, }); } + + public void TrackFeatureFlagStaleness(TimeSpan staleness) + { + _telemetryClient.TrackMetric( + Prefix + "FeatureFlagStalenessSeconds", + staleness.TotalSeconds); + } } } diff --git a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Integration/PopularityTransferIntegrationTests.cs b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Integration/PopularityTransferIntegrationTests.cs index 8e4e7f54d..af5be8327 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Integration/PopularityTransferIntegrationTests.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Integration/PopularityTransferIntegrationTests.cs @@ -1,4 +1,7 @@ -using System; +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. + +using System; using System.Collections.Generic; using System.Linq; using System.Threading.Tasks; @@ -19,6 +22,7 @@ public class PopularityTransferIntegrationTests private readonly InMemoryCloudBlobContainer _storageContainer; private readonly Mock _searchOperations; + private readonly Mock _featureFlags; private readonly Auxiliary2AzureSearchConfiguration _config; private readonly AzureSearchJobDevelopmentConfiguration _developmentConfig; private readonly Mock _telemetry; @@ -30,6 +34,7 @@ public class PopularityTransferIntegrationTests public PopularityTransferIntegrationTests(ITestOutputHelper output) { + _featureFlags = new Mock(); _telemetry = new Mock(); _config = new Auxiliary2AzureSearchConfiguration @@ -147,6 +152,8 @@ public PopularityTransferIntegrationTests(ITestOutputHelper output) var time = new Mock(); + _featureFlags.Setup(x => x.IsPopularityTransferEnabled()).Returns(true); + _target = new UpdateDownloadsCommand( auxiliaryFileClient, databaseFetcher.Object, @@ -158,6 +165,7 @@ public PopularityTransferIntegrationTests(ITestOutputHelper output) searchIndexActionBuilder, batchPusherFactory, time.Object, + _featureFlags.Object, options.Object, _telemetry.Object, output.GetLogger()); @@ -349,6 +357,59 @@ public async Task ReverseTransferChangesDownloads() VerifyUpdateDownloadCountAction("B", 10, actions[7]); } + [Fact] + public async Task DisablingPopularityTransferFeatureRemovesTransfers() + { + SetExcludedPackagesJson("{}"); + + AddVersionList("A", "1.0.0"); + AddVersionList("B", "1.0.0"); + AddVersionList("C", "1.0.0"); + + SetOldDownloadsJson(@" +{ + ""A"": { ""1.0.0"": 100 }, + ""B"": { ""1.0.0"": 20 }, + ""C"": { ""1.0.0"": 1 } +}"); + SetNewDownloadsJson(@" +[ + [ ""A"", [ ""1.0.0"", 100 ] ], + [ ""B"", [ ""1.0.0"", 20 ] ], + [ ""C"", [ ""1.0.0"", 1 ] ] +]"); + + // Old: A -> B rename + // New: A -> B rename + SetOldPopularityTransfersJson(@"{ ""A"": [ ""B"" ] }"); + AddNewPopularityTransfer("A", "B"); + SetDownloadOverrides(@"{ ""C"": 5 }"); + + _config.Scoring.PopularityTransfer = 0.5; + _featureFlags + .Setup(x => x.IsPopularityTransferEnabled()) + .Returns(false); + + await _target.ExecuteAsync(); + + Assert.NotNull(_indexedBatch); + var actions = _indexedBatch.Actions.OrderBy(x => x.Document.Key).ToList(); + Assert.Equal(12, actions.Count); + + VerifyUpdateDownloadCountAction("A", 100, actions[0]); + VerifyUpdateDownloadCountAction("A", 100, actions[1]); + VerifyUpdateDownloadCountAction("A", 100, actions[2]); + VerifyUpdateDownloadCountAction("A", 100, actions[3]); + VerifyUpdateDownloadCountAction("B", 20, actions[4]); + VerifyUpdateDownloadCountAction("B", 20, actions[5]); + VerifyUpdateDownloadCountAction("B", 20, actions[6]); + VerifyUpdateDownloadCountAction("B", 20, actions[7]); + VerifyUpdateDownloadCountAction("C", 5, actions[8]); + VerifyUpdateDownloadCountAction("C", 5, actions[9]); + VerifyUpdateDownloadCountAction("C", 5, actions[10]); + VerifyUpdateDownloadCountAction("C", 5, actions[11]); + } + private void SetOldDownloadsJson(string json) { _storageContainer.Blobs["downloads/downloads.v2.json"] = new InMemoryCloudBlob(json); diff --git a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/UpdateDownloadsCommandFacts.cs b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/UpdateDownloadsCommandFacts.cs index 57ebac98c..dd17c787b 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/UpdateDownloadsCommandFacts.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/UpdateDownloadsCommandFacts.cs @@ -232,6 +232,71 @@ public async Task AppliesTransferChanges() Times.Once); } + [Fact] + public async Task DisablesPopularityTransfers() + { + var downloadChanges = new SortedDictionary(StringComparer.OrdinalIgnoreCase) + { + { "Package1", 5 } + }; + + DownloadSetComparer + .Setup(c => c.Compare(It.IsAny(), It.IsAny())) + .Returns((oldData, newData) => + { + return downloadChanges; + }); + + OldTransfers["Package1"] = new SortedSet(StringComparer.OrdinalIgnoreCase) + { + "Package2" + }; + + NewTransfers["Package1"] = new SortedSet(StringComparer.OrdinalIgnoreCase) + { + "Package2" + }; + + DownloadOverrides["Package1"] = 100; + + FeatureFlags + .Setup(x => x.IsPopularityTransferEnabled()) + .Returns(false); + + await Target.ExecuteAsync(); + + PopularityTransferDataClient + .Verify( + c => c.ReadLatestIndexedAsync(), + Times.Once); + DatabaseFetcher + .Verify( + d => d.GetPackageIdToPopularityTransfersAsync(), + Times.Never); + AuxiliaryFileClient + .Verify( + a => a.LoadDownloadOverridesAsync(), + Times.Once); + + // The popularity transfers should not be given to the download transferrer. + DownloadTransferrer + .Verify( + x => x.UpdateDownloadTransfers( + NewDownloadData, + downloadChanges, + OldTransfers, + It.Is>>(d => d.Count == 0), + DownloadOverrides), + Times.Once); + + // Popularity transfers auxiliary file should be empty. + PopularityTransferDataClient.Verify( + c => c.ReplaceLatestIndexedAsync( + It.Is>>(d => d.Count == 0), + It.IsAny()), + Times.Once); + } + [Fact] public async Task TransferChangesOverrideDownloadChanges() { @@ -321,6 +386,7 @@ public Facts(ITestOutputHelper output) IndexActionBuilder = new Mock(); BatchPusher = new Mock(); SystemTime = new Mock(); + FeatureFlags = new Mock(); Options = new Mock>(); TelemetryService = new Mock(); Logger = output.GetLogger(); @@ -409,6 +475,8 @@ public Facts(ITestOutputHelper output) CurrentBatch = new ConcurrentBag(); }); + FeatureFlags.Setup(x => x.IsPopularityTransferEnabled()).Returns(true); + Target = new UpdateDownloadsCommand( AuxiliaryFileClient.Object, DatabaseFetcher.Object, @@ -420,6 +488,7 @@ public Facts(ITestOutputHelper output) IndexActionBuilder.Object, () => BatchPusher.Object, SystemTime.Object, + FeatureFlags.Object, Options.Object, TelemetryService.Object, Logger); @@ -435,6 +504,7 @@ public Facts(ITestOutputHelper output) public Mock IndexActionBuilder { get; } public Mock BatchPusher { get; } public Mock SystemTime { get; } + public Mock FeatureFlags { get; } public Mock> Options { get; } public Mock TelemetryService { get; } public RecordingLogger Logger { get; } diff --git a/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs b/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs index d40d22584..c8e7b1174 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs @@ -39,7 +39,9 @@ public class ProduceWorkAsync private readonly Mock _auxiliaryFileClient; private readonly Mock _databaseFetcher; private readonly Mock _downloadTransferrer; + private readonly Mock _featureFlags; private readonly DownloadData _downloads; + private readonly Dictionary _downloadOverrides; private readonly SortedDictionary> _popularityTransfers; private readonly SortedDictionary _transferChanges; private HashSet _excludedPackages; @@ -70,6 +72,10 @@ public ProduceWorkAsync(ITestOutputHelper output) _auxiliaryFileClient .Setup(x => x.LoadDownloadDataAsync()) .ReturnsAsync(() => _downloads); + _downloadOverrides = new Dictionary(); + _auxiliaryFileClient + .Setup(x => x.LoadDownloadOverridesAsync()) + .ReturnsAsync(() => _downloadOverrides); _popularityTransfers = new SortedDictionary>(); _databaseFetcher = new Mock(); @@ -86,6 +92,11 @@ public ProduceWorkAsync(ITestOutputHelper output) It.IsAny>())) .Returns(_transferChanges); + _featureFlags = new Mock(); + _featureFlags + .Setup(x => x.IsPopularityTransferEnabled()) + .Returns(true); + _entitiesContextFactory .Setup(x => x.CreateAsync(It.IsAny())) .ReturnsAsync(() => _entitiesContext.Object); @@ -107,6 +118,7 @@ public ProduceWorkAsync(ITestOutputHelper output) _auxiliaryFileClient.Object, _databaseFetcher.Object, _downloadTransferrer.Object, + _featureFlags.Object, _options.Object, _developmentOptions.Object, _logger); @@ -523,8 +535,21 @@ public async Task AppliesDownloadTransfers() _transferChanges["b"] = 66; _transferChanges["C"] = 123; + _downloadOverrides["C"] = 5; + var result = await _target.ProduceWorkAsync(_work, _token); + _auxiliaryFileClient.Verify(x => x.LoadDownloadOverridesAsync(), Times.Once); + _databaseFetcher.Verify(x => x.GetPackageIdToPopularityTransfersAsync(), Times.Once); + + _downloadTransferrer + .Verify( + x => x.InitializeDownloadTransfers( + _downloads, + _popularityTransfers, + _downloadOverrides), + Times.Once); + // Documents should have overriden downloads. var work = _work.Reverse().ToList(); Assert.Equal(3, work.Count); @@ -553,6 +578,49 @@ public async Task AppliesDownloadTransfers() Assert.Equal(3, result.Downloads["C"]["6.0.0"]); } + [Fact] + public async Task DisablesPopularityTransfers() + { + _packageRegistrations.Add(new PackageRegistration + { + Key = 1, + Id = "A", + Packages = new[] + { + new Package { Version = "1.0.0" }, + }, + }); + _downloads.SetDownloadCount("A", "1.0.0", 100); + _popularityTransfers["A"] = new SortedSet { "A" }; + _downloadOverrides["A"] = 5; + + InitializePackagesFromPackageRegistrations(); + + _featureFlags + .Setup(x => x.IsPopularityTransferEnabled()) + .Returns(false); + + var result = await _target.ProduceWorkAsync(_work, _token); + + // The popularity transfers should not be loaded from the database. + _databaseFetcher + .Verify( + x => x.GetPackageIdToPopularityTransfersAsync(), + Times.Never); + + // Popularity transfers should not be passed to the download transferrer. + _downloadTransferrer + .Verify( + x => x.InitializeDownloadTransfers( + _downloads, + It.Is>>(data => data.Count == 0), + _downloadOverrides), + Times.Once); + + // There should be no popularity transfers. + Assert.Empty(result.PopularityTransfers); + } + [Fact] public async Task IgnoresDownloadTransfersForNonexistentPackages() {