Add registration updater to update all of the hives for a single pack…

…age ID (#723)

Progress on NuGet/NuGetGallery#7739

src/NuGet.Jobs.Catalog2Registration/IRegistrationUpdater.cs

@@ -0,0 +1,18 @@
+// 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.Protocol.Catalog;
+using NuGet.Services.Metadata.Catalog;
+
+namespace NuGet.Jobs.Catalog2Registration
+{
+    public interface IRegistrationUpdater
+    {
+        Task UpdateAsync(
+            string id,
+            IReadOnlyList<CatalogCommitItem> entries,
+            IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf> entryToLeaf);
+    }
+}

src/NuGet.Jobs.Catalog2Registration/NuGet.Jobs.Catalog2Registration.csproj

@@ -45,6 +45,7 @@
     <Reference Include="System.Xml" />
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="RegistrationUpdater.cs" />
     <Compile Include="Hives\CatalogCommit.cs" />
     <Compile Include="Hives\Bookkeeping\LeafInfo.cs" />
     <Compile Include="Hives\Bookkeeping\IndexInfo.cs" />
@@ -54,6 +55,7 @@
     <Compile Include="Hives\HiveMergeResult.cs" />
     <Compile Include="Hives\HiveStorage.cs" />
     <Compile Include="Hives\HiveType.cs" />
+    <Compile Include="IRegistrationUpdater.cs" />
     <Compile Include="Schema\IEntityBuilder.cs" />
     <Compile Include="Hives\IHiveMerger.cs" />
     <Compile Include="Hives\IHiveStorage.cs" />

src/NuGet.Jobs.Catalog2Registration/RegistrationUpdater.cs

@@ -0,0 +1,96 @@
+// 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.Concurrent;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using NuGet.Protocol.Catalog;
+using NuGet.Services.Metadata.Catalog;
+using NuGet.Services.Metadata.Catalog.Helpers;
+
+namespace NuGet.Jobs.Catalog2Registration
+{
+    public class RegistrationUpdater : IRegistrationUpdater
+    {
+        private static readonly Dictionary<HiveType, IReadOnlyList<HiveType>> HiveToReplicaHives = new Dictionary<HiveType, IReadOnlyList<HiveType>>
+        {
+            {
+                HiveType.Legacy,
+                new List<HiveType> { HiveType.Gzipped }
+            },
+            {
+                HiveType.SemVer2,
+                new List<HiveType>()
+            },
+        };
+
+        private readonly IHiveUpdater _hiveUpdater;
+        private readonly IOptionsSnapshot<Catalog2RegistrationConfiguration> _options;
+        private readonly ILogger<RegistrationUpdater> _logger;
+
+        public RegistrationUpdater(
+            IHiveUpdater hiveUpdater,
+            IOptionsSnapshot<Catalog2RegistrationConfiguration> options,
+            ILogger<RegistrationUpdater> logger)
+        {
+            _hiveUpdater = hiveUpdater ?? throw new ArgumentNullException(nameof(hiveUpdater));
+            _options = options ?? throw new ArgumentNullException(nameof(options));
+            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+
+            if (_options.Value.MaxConcurrentHivesPerId <= 0)
+            {
+                throw new ArgumentOutOfRangeException(
+                    nameof(options),
+                    $"The {nameof(Catalog2RegistrationConfiguration.MaxConcurrentHivesPerId)} must be greater than zero.");
+            }
+        }
+
+        public async Task UpdateAsync(
+            string id,
+            IReadOnlyList<CatalogCommitItem> entries,
+            IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf> entryToLeaf)
+        {
+            var registrationCommitTimestamp = DateTimeOffset.UtcNow;
+            var hives = new ConcurrentBag<HiveType>(HiveToReplicaHives.Keys);
+            await ParallelAsync.Repeat(
+                async () =>
+                {
+                    await Task.Yield();
+                    while (hives.TryTake(out var hive))
+                    {
+                        await ProcessHiveAsync(hive, id, entries, entryToLeaf, registrationCommitTimestamp);
+                    }
+                },
+                _options.Value.MaxConcurrentHivesPerId);
+        }
+
+        private async Task ProcessHiveAsync(
+            HiveType hive,
+            string id,
+            IReadOnlyList<CatalogCommitItem> entries,
+            IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf> entryToLeaf,
+            DateTimeOffset registrationCommitTimestamp)
+        {
+            var registrationCommitId = Guid.NewGuid().ToString();
+            var registrationCommit = new CatalogCommit(registrationCommitId, registrationCommitTimestamp);
+            using (_logger.BeginScope(
+                "Processing package {PackageId}, " +
+                "hive {Hive}, " +
+                "replica hives {ReplicaHives}, " +
+                "registration commit ID {RegistrationCommitId}, " +
+                "registration commit timestamp {RegistrationCommitTimestamp:O}.",
+                id,
+                hive,
+                HiveToReplicaHives[hive],
+                registrationCommitId,
+                registrationCommitTimestamp))
+            {
+                _logger.LogInformation("Processing {Count} catalog commit items.", entries.Count);
+                await _hiveUpdater.UpdateAsync(hive, HiveToReplicaHives[hive], id, entries, entryToLeaf, registrationCommit);
+            }
+        }
+    }
+}

tests/NuGet.Jobs.Catalog2Registration.Tests/NuGet.Jobs.Catalog2Registration.Tests.csproj

@@ -64,6 +64,7 @@
     <Compile Include="Hives\HiveUpdaterFacts.cs" />
     <Compile Include="Hives\RegistrationUrlBuilderFacts.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
+    <Compile Include="RegistrationUpdaterFacts.cs" />
     <Compile Include="Schema\EntityBuilderFacts.cs" />
   </ItemGroup>
   <ItemGroup>

tests/NuGet.Jobs.Catalog2Registration.Tests/RegistrationUpdaterFacts.cs

@@ -0,0 +1,116 @@
+// 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.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Options;
+using Moq;
+using NuGet.Protocol.Catalog;
+using NuGet.Services;
+using NuGet.Services.Metadata.Catalog;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace NuGet.Jobs.Catalog2Registration
+{
+    public class RegistrationUpdaterFacts
+    {
+        public class UpdateAsync : Facts
+        {
+            public UpdateAsync(ITestOutputHelper output) : base(output)
+            {
+            }
+
+            [Fact]
+            public async Task UsesDifferentCommitIdButSameCommitTimestamp()
+            {
+                var commits = new ConcurrentBag<CatalogCommit>();
+                HiveUpdater
+                    .Setup(x => x.UpdateAsync(
+                        It.IsAny<HiveType>(),
+                        It.IsAny<IReadOnlyList<HiveType>>(),
+                        It.IsAny<string>(),
+                        It.IsAny<IReadOnlyList<CatalogCommitItem>>(),
+                        It.IsAny<IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf>>(),
+                        It.IsAny<CatalogCommit>()))
+                    .Returns(Task.CompletedTask)
+                    .Callback<HiveType, IReadOnlyList<HiveType>, string, IReadOnlyList<CatalogCommitItem>, IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf>, CatalogCommit>(
+                        (h, r, i, e, l, c) => commits.Add(c));
+
+                await Target.UpdateAsync(Id, Entries, EntryToLeaf);
+
+                Assert.Equal(2, commits.Count);
+                Assert.Single(commits.Select(x => x.Timestamp).Distinct());
+                Assert.Equal(2, commits.Select(x => x.Id).Distinct().Count());
+            }
+
+            [Fact]
+            public async Task UsesProperReplicaHives()
+            {
+                await Target.UpdateAsync(Id, Entries, EntryToLeaf);
+
+                HiveUpdater.Verify(
+                    x => x.UpdateAsync(
+                        It.IsAny<HiveType>(),
+                        It.IsAny<IReadOnlyList<HiveType>>(),
+                        It.IsAny<string>(),
+                        It.IsAny<IReadOnlyList<CatalogCommitItem>>(),
+                        It.IsAny<IReadOnlyDictionary<CatalogCommitItem, PackageDetailsCatalogLeaf>>(),
+                        It.IsAny<CatalogCommit>()),
+                    Times.Exactly(2));
+                HiveUpdater.Verify(
+                    x => x.UpdateAsync(
+                        HiveType.Legacy,
+                        It.Is<IReadOnlyList<HiveType>>(r => r.Count == 1 && r[0] == HiveType.Gzipped),
+                        Id,
+                        Entries,
+                        EntryToLeaf,
+                        It.IsAny<CatalogCommit>()),
+                    Times.Once);
+                HiveUpdater.Verify(
+                    x => x.UpdateAsync(
+                        HiveType.SemVer2,
+                        It.Is<IReadOnlyList<HiveType>>(r => r.Count == 0),
+                        Id,
+                        Entries,
+                        EntryToLeaf,
+                        It.IsAny<CatalogCommit>()),
+                    Times.Once);
+            }
+        }
+
+        public abstract class Facts
+        {
+            public Facts(ITestOutputHelper output)
+            {
+                HiveUpdater = new Mock<IHiveUpdater>();
+                Options = new Mock<IOptionsSnapshot<Catalog2RegistrationConfiguration>>();
+                Logger = output.GetLogger<RegistrationUpdater>();
+
+                Config = new Catalog2RegistrationConfiguration();
+                Config.MaxConcurrentHivesPerId = 1;
+                Id = "NuGet.Versioning";
+                Entries = new List<CatalogCommitItem>();
+                EntryToLeaf = new Dictionary<CatalogCommitItem, PackageDetailsCatalogLeaf>();
+
+                Options.Setup(x => x.Value).Returns(() => Config);
+
+                Target = new RegistrationUpdater(
+                    HiveUpdater.Object,
+                    Options.Object,
+                    Logger);
+            }
+
+            public Mock<IHiveUpdater> HiveUpdater { get; }
+            public Mock<IOptionsSnapshot<Catalog2RegistrationConfiguration>> Options { get; }
+            public RecordingLogger<RegistrationUpdater> Logger { get; }
+            public Catalog2RegistrationConfiguration Config { get; }
+            public string Id { get; }
+            public List<CatalogCommitItem> Entries { get; }
+            public Dictionary<CatalogCommitItem, PackageDetailsCatalogLeaf> EntryToLeaf { get; }
+            public RegistrationUpdater Target { get; }
+        }
+    }
+}