[Revalidation] Improve throughput by batching revalidations (#667)

Increase the throughput of the revalidation job by batching revalidations. The revalidation job does a bunch of work before enqueueing revalidations as it must:

1. Ensure the revalidation job hasn't been killswitched
2. Verify that the ingestion pipeline is healthy
3. Verify that the desired package event rate hasn't been reached

This change amortizes that work by enqueueing revalidations in batches. Also, the job now takes into account how long it spent enqueueing revalidations when deciding how long to sleep for.

Addresses https://github.com/NuGet/Engineering/issues/1877

Author: loic-sharma

src/NuGet.Services.Revalidate/Configuration/RevalidationQueueConfiguration.cs

@@ -8,19 +8,13 @@ namespace NuGet.Services.Revalidate
     public class RevalidationQueueConfiguration
     {
         /// <summary>
-        /// If non-null, this skips revalidations of packages with more than this many versions.
+        /// The maximum number of revalidations that should be returned by <see cref="IRevalidationQueue.NextAsync"/>.
         /// </summary>
-        public int? MaximumPackageVersions { get; set; }
+        public int MaxBatchSize { get; set; } = 64;
 
         /// <summary>
-        /// The maximum times that the <see cref="RevalidationQueue"/> should look for a revalidation
-        /// before giving up.
-        /// </summary>
-        public int MaximumAttempts { get; set; } = 5;
-
-        /// <summary>
-        /// The time to sleep after an initialized revalidation is deemed completed.
+        /// If non-null, this skips revalidations of packages with more than this many versions.
         /// </summary>
-        public TimeSpan SleepBetweenAttempts { get; set; } = TimeSpan.FromSeconds(5);
+        public int? MaximumPackageVersions { get; set; }
     }
 }

src/NuGet.Services.Revalidate/NuGet.Services.Revalidate.csproj

@@ -69,7 +69,7 @@
     <Compile Include="Initialization\PackageRevalidationInserter.cs" />
     <Compile Include="Services\RevalidationOperation.cs" />
     <Compile Include="Services\RevalidationQueue.cs" />
-    <Compile Include="Services\RevalidationResult.cs" />
+    <Compile Include="Services\StartRevalidationStatus.cs" />
     <Compile Include="Services\RevalidationService.cs" />
     <Compile Include="Services\RevalidationJobStateService.cs" />
     <Compile Include="Services\PackageRevalidationStateService.cs" />
@@ -82,6 +82,7 @@
     <Compile Include="Program.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
     <Compile Include="Properties\AssemblyInfo.*.cs" />
+    <Compile Include="Services\StartRevalidationResult.cs" />
   </ItemGroup>
   <ItemGroup>
     <None Include="App.config" />

src/NuGet.Services.Revalidate/Services/IPackageRevalidationStateService.cs

@@ -34,10 +34,10 @@ public interface IPackageRevalidationStateService
         Task<int> CountRevalidationsEnqueuedInPastHourAsync();
 
         /// <summary>
-        /// Update the package revalidation and mark is as enqueued.
+        /// Update the package revalidations and mark them as enqueued.
         /// </summary>
-        /// <param name="revalidation">The revalidation to update.</param>
-        /// <returns>A task that completes once the revalidation has been updated.</returns>
-        Task MarkPackageRevalidationAsEnqueuedAsync(PackageRevalidation revalidation);
+        /// <param name="revalidations">The revalidations to update.</param>
+        /// <returns>A task that completes once the revalidations have been updated.</returns>
+        Task MarkPackageRevalidationsAsEnqueuedAsync(IReadOnlyList<PackageRevalidation> revalidations);
     }
 }

src/NuGet.Services.Revalidate/Services/IRevalidationQueue.cs

@@ -1,6 +1,7 @@
 // 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.Collections.Generic;
 using System.Threading.Tasks;
 using NuGet.Services.Validation;
 
@@ -9,9 +10,9 @@ namespace NuGet.Services.Revalidate
     public interface IRevalidationQueue
     {
         /// <summary>
-        /// Fetch the next package to revalidate.
+        /// Fetch the next packages to revalidate.
         /// </summary>
-        /// <returns>The next package to revalidate, or null if there are no packages to revalidate at this time.</returns>
-        Task<PackageRevalidation> NextOrNullAsync();
+        /// <returns>The next packages to revalidate, or an empty list if there are no packages to revalidate at this time.</returns>
+        Task<IReadOnlyList<PackageRevalidation>> NextAsync();
     }
 }

src/NuGet.Services.Revalidate/Services/IRevalidationStarter.cs

@@ -20,6 +20,6 @@ public interface IRevalidationStarter
         /// 4. A revalidation could not be found at this time
         /// </summary>
         /// <returns>The result of the revalidation attempt.</returns>
-        Task<RevalidationResult> StartNextRevalidationAsync();
+        Task<StartRevalidationResult> StartNextRevalidationsAsync();
     }
 }

src/NuGet.Services.Revalidate/Services/IRevalidationThrottler.cs

@@ -1,6 +1,7 @@
 // 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.Threading.Tasks;
 
 namespace NuGet.Services.Revalidate
@@ -16,8 +17,10 @@ public interface IRevalidationThrottler
         /// <summary>
         /// Delay the current task to achieve the desired revalidation rate.
         /// </summary>
+        /// <param name="started">The number of revalidations started.</param>
+        /// <param name="startDuration">How long it took it took to start the revalidations.</param>
         /// <returns>Delay the task to ensure the desired revalidation rate.</returns>
-        Task DelayUntilNextRevalidationAsync();
+        Task DelayUntilNextRevalidationAsync(int started, TimeSpan startDuration);
 
         /// <summary>
         /// Delay the current task until when a revalidation can be retried.

src/NuGet.Services.Revalidate/Services/ITelemetryService.cs

@@ -1,12 +1,15 @@
 // 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.Logging;
 
 namespace NuGet.Services.Revalidate
 {
     public interface ITelemetryService
     {
+        IDisposable TrackFindNextRevalidations();
+
         DurationMetric<StartNextRevalidationOperation> TrackStartNextRevalidationOperation();
 
         void TrackPackageRevalidationMarkedAsCompleted(string packageId, string normalizedVersion);

src/NuGet.Services.Revalidate/Services/PackageRevalidationStateService.cs

@@ -67,21 +67,21 @@ public async Task<int> CountRevalidationsEnqueuedInPastHourAsync()
                 .CountAsync();
         }
 
-        public async Task MarkPackageRevalidationAsEnqueuedAsync(PackageRevalidation revalidation)
+        public async Task MarkPackageRevalidationsAsEnqueuedAsync(IReadOnlyList<PackageRevalidation> revalidations)
         {
             try
             {
-                revalidation.Enqueued = DateTime.UtcNow;
+                var enqueueTime = DateTime.UtcNow;
+                foreach (var revalidation in revalidations)
+                {
+                    revalidation.Enqueued = enqueueTime;
+                }
 
                 await _context.SaveChangesAsync();
             }
             catch (DbUpdateConcurrencyException)
             {
-                _logger.LogWarning(
-                    "Failed to update revalidation as enqueued for {PackageId} {PackageNormalizedVersion}",
-                    revalidation.PackageId,
-                    revalidation.PackageNormalizedVersion);
-
+                _logger.LogWarning("Failed to update revalidations as enqueued");
                 throw;
             }
         }

src/NuGet.Services.Revalidate/Services/RevalidationOperation.cs

@@ -6,8 +6,13 @@ namespace NuGet.Services.Revalidate
     public class StartNextRevalidationOperation
     {
         /// <summary>
-        /// The result of attempting to start the next revalidation.
+        /// The result of attempting to start the next revalidations.
         /// </summary>
-        public RevalidationResult Result { get; set; }
+        public StartRevalidationStatus Result { get; set; }
+
+        /// <summary>
+        /// The number of revalidations started.
+        /// </summary>
+        public int Started { get; set; }
     }
 }

src/NuGet.Services.Revalidate/Services/RevalidationQueue.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;
+using System.Collections.Generic;
 using System.Data.Entity;
 using System.Data.Entity.Infrastructure;
 using System.Linq;
@@ -37,19 +38,17 @@ public RevalidationQueue(
             _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         }
 
-        public async Task<PackageRevalidation> NextOrNullAsync()
+        public async Task<IReadOnlyList<PackageRevalidation>> NextAsync()
         {
-            for (var i = 0; i < _config.MaximumAttempts; i++)
+            // Find the next package to revalidate. We will skip packages if:
+            //   1. The package has more than "MaximumPackageVersions" versions
+            //   2. The package has already been enqueued for revalidation
+            //   3. The package's revalidation was completed by an external factory (like manual admin revalidation)
+            List<PackageRevalidation> next;
+            using (_telemetry.TrackFindNextRevalidations())
             {
-                _logger.LogInformation(
-                    "Attempting to find the next revalidation. Try {Attempt} of {MaxAttempts}",
-                    i + 1,
-                    _config.MaximumAttempts);
-
-                // Find the next package to revalidate. We will skip packages if:
-                //   1. The package has more than "MaximumPackageVersions" versions
-                //   2. The package has already been enqueued for revalidation
-                //   3. The package's revalidation was completed by an external factory (like manual admin revalidation)
+                _logger.LogInformation("Finding the next packages to revalidate...");
+
                 IQueryable<PackageRevalidation> query = _validationContext.PackageRevalidations;
 
                 if (_config.MaximumPackageVersions.HasValue)
@@ -61,82 +60,130 @@ public async Task<PackageRevalidation> NextOrNullAsync()
                         .Any(g => g.Key == r.PackageId));
                 }
 
-                var next = await query
+                next = await query
                     .Where(r => r.Enqueued == null)
                     .Where(r => r.Completed == false)
                     .OrderBy(r => r.Key)
-                    .FirstOrDefaultAsync();
+                    .Take(_config.MaxBatchSize)
+                    .ToListAsync();
+            }
 
-                if (next == null)
-                {
-                    _logger.LogWarning("Could not find any incomplete revalidations");
-                    return null;
-                }
+            _logger.LogInformation("Found {Revalidations} packages to revalidate", next.Count);
 
-                // Don't revalidate packages that already have a repository signature or that no longer exist.
-                if (await HasRepositorySignature(next) || await IsDeleted(next))
-                {
-                    await MarkAsCompleted(next);
-                    await Task.Delay(_config.SleepBetweenAttempts);
+            // Return all the revalidations that aren't already completed.
+            return await FilterCompletedRevalidationsAsync(next);
+        }
 
-                    continue;
-                }
+        private async Task<IReadOnlyList<PackageRevalidation>> FilterCompletedRevalidationsAsync(IReadOnlyList<PackageRevalidation> revalidations)
+        {
+            if (!revalidations.Any())
+            {
+                return revalidations;
+            }
 
-                _logger.LogInformation(
-                    "Found revalidation for {PackageId} {PackageNormalizedVersion} after {Attempt} attempts",
-                    next.PackageId,
-                    next.PackageNormalizedVersion,
-                    i + 1);
+            var completed = new List<PackageRevalidation>();
+            var uncompleted = revalidations.ToDictionary(
+                r => $"{r.PackageId}/{r.PackageNormalizedVersion}",
+                r => r);
 
-                return next;
-            }
+            // Packages that already have a repository signature do not need to be revalidated.
+            _logger.LogInformation("Finding revalidations that can be skipped because their packages are already repository signed...");
+
+            var hasRepositorySignatures = await _validationContext.PackageSigningStates
+                .Select(s => new {
+                    IdAndVersion = s.PackageId + "/" + s.PackageNormalizedVersion,
+                    s.PackageSignatures
+                })
+                .Where(s => uncompleted.Keys.Contains(s.IdAndVersion))
+                .Where(s => s.PackageSignatures.Any(sig => sig.Type == PackageSignatureType.Repository))
+                .Select(s => s.IdAndVersion)
+                .ToListAsync();
 
             _logger.LogInformation(
-                "Did not find any revalidations after {MaxAttempts}. Retry later...",
-                _config.MaximumAttempts);
+                "Found {RevalidationCount} revalidations that can be skipped because their packages are already repository signed",
+                hasRepositorySignatures.Count);
 
-            return null;
-        }
+            foreach (var idAndVersion in hasRepositorySignatures)
+            {
+                completed.Add(uncompleted[idAndVersion]);
+                uncompleted.Remove(idAndVersion);
+            }
 
-        private Task<bool> HasRepositorySignature(PackageRevalidation revalidation)
-        {
-            return _validationContext.PackageSigningStates
-                .Where(s => s.PackageId == revalidation.PackageId)
-                .Where(s => s.PackageNormalizedVersion == revalidation.PackageNormalizedVersion)
-                .Where(s => s.PackageSignatures.Any(sig => sig.Type == PackageSignatureType.Repository))
-                .AnyAsync();
-        }
+            // Packages that are no longer available should not be revalidated.
+            _logger.LogInformation("Finding revalidations' package statuses...");
 
-        private async Task<bool> IsDeleted(PackageRevalidation revalidation)
-        {
-            var packageStatus = await _galleryContext.Set<Package>()
-                .Where(p => p.PackageRegistration.Id == revalidation.PackageId)
-                .Where(p => p.NormalizedVersion == revalidation.PackageNormalizedVersion)
-                .Select(p => (PackageStatus?)p.PackageStatusKey)
-                .FirstOrDefaultAsync();
+            var packageStatuses = await _galleryContext.Set<Package>()
+                .Select(p => new
+                {
+                    Identity = p.PackageRegistration.Id + "/" + p.NormalizedVersion,
+                    p.PackageStatusKey
+                })
+                .Where(p => uncompleted.Keys.Contains(p.Identity))
+                .ToDictionaryAsync(
+                    p => p.Identity,
+                    p => p.PackageStatusKey);
 
-            return (packageStatus == null || packageStatus == PackageStatus.Deleted);
-        }
+            _logger.LogInformation("Found {PackageStatusCount} revalidations' package statuses", packageStatuses.Count);
+
+            foreach (var key in uncompleted.Keys.ToList())
+            {
+                // Packages that are hard deleted won't have a status.
+                if (!packageStatuses.TryGetValue(key, out var status) || status == PackageStatus.Deleted)
+                {
+                    completed.Add(uncompleted[key]);
+                    uncompleted.Remove(key);
+                    continue;
+                }
+            }
 
-        private async Task MarkAsCompleted(PackageRevalidation revalidation)
-        {
             _logger.LogInformation(
-                "Marking package revalidation as completed as it has a repository signature or is deleted for {PackageId} {PackageNormalizedVersion}",
-                revalidation.PackageId,
-                revalidation.PackageNormalizedVersion);
+                "Found {CompletedRevalidations} revalidations that can be skipped. There are {UncompletedRevalidations} " +
+                "revalidations remaining in this batch",
+                completed.Count,
+                uncompleted.Count);
 
+            // Update revalidations that were determined to be completed and return the remaining revalidations.
+            if (completed.Any())
+            {
+                await MarkRevalidationsAsCompletedAsync(completed);
+            }
+
+            return uncompleted.Values.ToList();
+        }
+
+        private async Task MarkRevalidationsAsCompletedAsync(IReadOnlyList<PackageRevalidation> revalidations)
+        {
             try
             {
-                revalidation.Completed = true;
+                foreach (var revalidation in revalidations)
+                {
+                    _logger.LogInformation(
+                        "Marking package {PackageId} {PackageNormalizedVersion} revalidation as completed as the package is unavailable or the package is already repository signed...",
+                        revalidation.PackageId,
+                        revalidation.PackageNormalizedVersion);
+
+                    revalidation.Completed = true;
+                }
 
                 await _validationContext.SaveChangesAsync();
 
-                _telemetry.TrackPackageRevalidationMarkedAsCompleted(revalidation.PackageId, revalidation.PackageNormalizedVersion);
+                foreach (var revalidation in revalidations)
+                {
+                    _logger.LogInformation(
+                        "Marked package {PackageId} {PackageNormalizedVersion} revalidation as completed",
+                        revalidation.PackageId,
+                        revalidation.PackageNormalizedVersion);
+
+                    _telemetry.TrackPackageRevalidationMarkedAsCompleted(revalidation.PackageId, revalidation.PackageNormalizedVersion);
+                }
             }
-            catch (DbUpdateConcurrencyException)
+            catch (DbUpdateConcurrencyException e)
             {
-                // Swallow concurrency exceptions. The package will be marked as completed
-                // on the next iteration of "NextOrNullAsync".
+                _logger.LogError(
+                    0,
+                    e,
+                    "Failed to mark package revalidations as completed. " +
+                    $"These revalidations will be marked as completed on the next iteration of {nameof(NextAsync)}...");
             }
         }
     }