[Revalidation] Improve package registration finder

src/NuGet.Services.Revalidate/Initialization/IPackageFinder.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using NuGet.Versioning;
 
 namespace NuGet.Services.Revalidate
@@ -41,7 +42,7 @@ public interface IPackageFinder
         /// <param name="setName">The name of this set of packages.</param>
         /// <param name="packageRegistrationKeys">The set of package registration keys.</param>
         /// <returns>Information about each package registration, if it exists in the database.</returns>
-        List<PackageRegistrationInformation> FindPackageRegistrationInformation(string setName, HashSet<int> packageRegistrationKeys);
+        Task<List<PackageRegistrationInformation>> FindPackageRegistrationInformationAsync(string setName, HashSet<int> packageRegistrationKeys);
 
         /// <summary>
         /// Find versions that are appropriate for revalidations.

src/NuGet.Services.Revalidate/Initialization/InitializationManager.cs

@@ -122,7 +122,7 @@ private async Task ClearPackageRevalidationStateAsync()
 
         private async Task InitializePackageSetAsync(string setName, HashSet<int> packageRegistrationKeys)
         {
-            var packageInformations = _packageFinder.FindPackageRegistrationInformation(setName, packageRegistrationKeys);
+            var packageInformations = await _packageFinder.FindPackageRegistrationInformationAsync(setName, packageRegistrationKeys);
 
             var chunks = packageInformations
                 .OrderByDescending(p => p.Downloads)

src/NuGet.Services.Revalidate/Initialization/PackageFinder.cs

@@ -7,6 +7,8 @@
 using System.Linq;
 using System.Linq.Expressions;
 using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Newtonsoft.Json;
 using NuGet.Versioning;
@@ -29,15 +31,18 @@ public class PackageFinder : IPackageFinder
         private static string MicrosoftAccountName = "Microsoft";
 
         private readonly IGalleryContext _galleryContext;
+        private readonly IServiceScopeFactory _scopeFactory;
         private readonly InitializationConfiguration _config;
         private readonly ILogger<PackageFinder> _logger;
 
         public PackageFinder(
             IGalleryContext galleryContext,
+            IServiceScopeFactory scopeFactory,
             InitializationConfiguration config,
             ILogger<PackageFinder> logger)
         {
             _galleryContext = galleryContext ?? throw new ArgumentNullException(nameof(galleryContext));
+            _scopeFactory = scopeFactory ?? throw new ArgumentNullException(nameof(scopeFactory));
             _config = config ?? throw new ArgumentNullException(nameof(config));
             _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         }
@@ -105,7 +110,7 @@ public HashSet<int> FindAllPackages(HashSet<int> except)
             return FindRegistrationKeys(RemainingSetName, r => !except.Contains(r.Key));
         }
 
-        public List<PackageRegistrationInformation> FindPackageRegistrationInformation(string setName, HashSet<int> packageRegistrationKeys)
+        public async Task<List<PackageRegistrationInformation>> FindPackageRegistrationInformationAsync(string setName, HashSet<int> packageRegistrationKeys)
         {
             // Fetch the packages' information in batches.
             var batches = packageRegistrationKeys.Batch(BatchSize);
@@ -121,23 +126,38 @@ public List<PackageRegistrationInformation> FindPackageRegistrationInformation(s
 
                 var batch = batches[batchIndex];
 
-                var packages = _galleryContext.PackageRegistrations
-                    .Where(r => batch.Contains(r.Key))
-                    .Select(r => new PackageRegistrationInformation
-                    {
-                        Key = r.Key,
-                        Id = r.Id,
-                        Downloads = r.DownloadCount,
-                        Versions = r.Packages.Count(),
-                    });
+                using (var scope = _scopeFactory.CreateScope())
+                {
+                    var scopedContext = scope.ServiceProvider.GetRequiredService<IGalleryContext>();
+                    var packages = scopedContext
+                        .PackageRegistrations
+                        .Where(r => batch.Contains(r.Key))
+                        .Select(r => new PackageRegistrationInformation
+                        {
+                            Key = r.Key,
+                            Id = r.Id,
+                            Downloads = r.DownloadCount,
+                            Versions = r.Packages.Count(),
+                        });
 
-                result.AddRange(packages);
+                    result.AddRange(packages);
+                }
 
                 _logger.LogInformation(
                     "Fetched batch {Batch} of {BatchesCount} of package informations for package set {SetName}",
                     batchIndex + 1,
                     batches.Count,
                     setName);
+
+                if (batchIndex < batches.Count - 1)
+                {
+                    _logger.LogInformation(
+                        "Sleeping for {SleepDuration} before fetching the next batch for package set {SetName}...",
+                        _config.SleepDurationBetweenBatches,
+                        setName);
+
+                    await Task.Delay(_config.SleepDurationBetweenBatches);
+                }
             }
 
             return result;

tests/NuGet.Services.Revalidate.Tests/Initializer/InitializationManagerFacts.cs

@@ -39,8 +39,8 @@ public async Task RemovesPreviousRevalidations()
                 _packageFinder.Setup(f => f.FindDependencyPackages(It.IsAny<HashSet<int>>())).Returns(new HashSet<int>());
                 _packageFinder.Setup(f => f.FindAllPackages(It.IsAny<HashSet<int>>())).Returns(new HashSet<int>());
 
-                _packageFinder.Setup(f => f.FindPackageRegistrationInformation(It.IsAny<string>(), It.IsAny<HashSet<int>>()))
-                    .Returns(new List<PackageRegistrationInformation>());
+                _packageFinder.Setup(f => f.FindPackageRegistrationInformationAsync(It.IsAny<string>(), It.IsAny<HashSet<int>>()))
+                    .ReturnsAsync(new List<PackageRegistrationInformation>());
 
                 var firstRemove = true;
 
@@ -339,20 +339,20 @@ List<PackageRegistrationInformation> RegistrationInformation(IEnumerable<Package
                 }
 
                 _packageFinder
-                    .Setup(f => f.FindPackageRegistrationInformation(PackageFinder.MicrosoftSetName, It.IsAny<HashSet<int>>()))
-                    .Returns(RegistrationInformation(microsoftPackages));
+                    .Setup(f => f.FindPackageRegistrationInformationAsync(PackageFinder.MicrosoftSetName, It.IsAny<HashSet<int>>()))
+                    .ReturnsAsync(RegistrationInformation(microsoftPackages));
 
                 _packageFinder
-                    .Setup(f => f.FindPackageRegistrationInformation(PackageFinder.PreinstalledSetName, It.IsAny<HashSet<int>>()))
-                    .Returns(RegistrationInformation(preinstalledPackages));
+                    .Setup(f => f.FindPackageRegistrationInformationAsync(PackageFinder.PreinstalledSetName, It.IsAny<HashSet<int>>()))
+                    .ReturnsAsync(RegistrationInformation(preinstalledPackages));
 
                 _packageFinder
-                    .Setup(f => f.FindPackageRegistrationInformation(PackageFinder.DependencySetName, It.IsAny<HashSet<int>>()))
-                    .Returns(RegistrationInformation(dependencyPackages));
+                    .Setup(f => f.FindPackageRegistrationInformationAsync(PackageFinder.DependencySetName, It.IsAny<HashSet<int>>()))
+                    .ReturnsAsync(RegistrationInformation(dependencyPackages));
 
                 _packageFinder
-                    .Setup(f => f.FindPackageRegistrationInformation(PackageFinder.RemainingSetName, It.IsAny<HashSet<int>>()))
-                    .Returns(RegistrationInformation(remainingPackages));
+                    .Setup(f => f.FindPackageRegistrationInformationAsync(PackageFinder.RemainingSetName, It.IsAny<HashSet<int>>()))
+                    .ReturnsAsync(RegistrationInformation(remainingPackages));
 
                 // Build the list of versions for each version of packages.
                 var versions = microsoftPackages

tests/NuGet.Services.Revalidate.Tests/Initializer/PackageFinderFacts.cs

@@ -3,6 +3,8 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Moq;
 using NuGet.Versioning;
@@ -207,7 +209,7 @@ public void ExcludesIdsFromExceptParameter()
         public class TheFindPackageRegistrationInformationMethod : FactsBase
         {
             [Fact]
-            public void FindsRegistrationInformation()
+            public async Task FindsRegistrationInformation()
             {
                 // Arrange
                 // Package registration 3 isn't in the input set and should be ignored.
@@ -249,7 +251,7 @@ public void FindsRegistrationInformation()
                     });
 
                 // Act and assert
-                var actual = _target.FindPackageRegistrationInformation("Name", new HashSet<int> { 1, 2 });
+                var actual = await _target.FindPackageRegistrationInformationAsync("Name", new HashSet<int> { 1, 2 });
 
                 Assert.Equal(2, actual.Count);
                 Assert.Equal(1, actual[0].Key);
@@ -264,7 +266,7 @@ public void FindsRegistrationInformation()
             }
 
             [Fact]
-            public void PackageCountDoesntFilterPackagesThatWontBeRevalidated()
+            public async Task PackageCountDoesntFilterPackagesThatWontBeRevalidated()
             {
                 // Arrange
                 // Package registration 3 isn't in the input set and should be ignored.
@@ -294,7 +296,7 @@ public void PackageCountDoesntFilterPackagesThatWontBeRevalidated()
                     });
 
                 // Act & Assert
-                var actual = _target.FindPackageRegistrationInformation("Name", new HashSet<int> { 1 });
+                var actual = await _target.FindPackageRegistrationInformationAsync("Name", new HashSet<int> { 1 });
 
                 Assert.Single(actual);
                 Assert.Equal(1, actual[0].Key);
@@ -464,17 +466,27 @@ public void SkipsPackagesThatAreNotAvailableOrDeleted()
         public class FactsBase
         {
             public readonly Mock<IEntitiesContext> _context;
+            public readonly Mock<IServiceScopeFactory> _scopeFactory;
 
             public readonly InitializationConfiguration _config;
             public readonly PackageFinder _target;
 
             public FactsBase()
             {
                 _context = new Mock<IEntitiesContext>();
+                _scopeFactory = new Mock<IServiceScopeFactory>();
                 _config = new InitializationConfiguration();
 
+                var scope = new Mock<IServiceScope>();
+                var serviceProvider = new Mock<IServiceProvider>();
+
+                _scopeFactory.Setup(s => s.CreateScope()).Returns(scope.Object);
+                scope.Setup(s => s.ServiceProvider).Returns(serviceProvider.Object);
+                serviceProvider.Setup(p => p.GetService(typeof(IEntitiesContext))).Returns(_context.Object);
+
                 _target = new PackageFinder(
                     _context.Object,
+                    _scopeFactory.Object,
                     _config,
                     Mock.Of<ILogger<PackageFinder>>());
             }

tests/Validation.PackageSigning.ScanAndSign.Tests/ScanAndSignProcessorFacts.cs

@@ -311,9 +311,12 @@ public async Task EnqueuesScanAndSignEvenIfRepositorySigningIsDisabled()
                         _request.NupkgUrl,
                         _config.V3ServiceIndexUrl,
                         It.Is<List<string>>(l =>
-                            l.Count() == 2 &&
-                            l.Contains("Billy") &&
-                            l.Contains("Bob"))),
+                            // Ensure that the owners are lexicographically ordered.
+                            l.Count() == 4 &&
+                            l[0] == "Annie" &&
+                            l[1] == "Bob" &&
+                            l[2] == "zack" &&
+                            l[3] == "Zorro")),
                     Times.Once);
 
             _validatorStateServiceMock