Update DI for alt stats container (#8472)

NuGet · Apr 1, 2021 · 91c68fe · 91c68fe
1 parent d772bac
commit 91c68fe
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 25 deletions.
diff --git a/src/NuGetGallery/App_Start/DefaultDependenciesModule.cs b/src/NuGetGallery/App_Start/DefaultDependenciesModule.cs
@@ -707,27 +707,41 @@ private static void RegisterDeleteAccountService(ContainerBuilder builder, Confi
 
         private static void RegisterStatisticsServices(ContainerBuilder builder, IGalleryConfigurationService configuration, ITelemetryService telemetryService)
         {
+            // when running on Windows Azure, download counts come from the downloads.v1.json blob
+            builder.Register(c => new SimpleBlobStorageConfiguration(configuration.Current.AzureStorage_Statistics_ConnectionString, configuration.Current.AzureStorageReadAccessGeoRedundant))
+                .SingleInstance()
+                .Keyed<IBlobStorageConfiguration>(BindingKeys.PrimaryStatisticsKey);
+
+            builder.Register(c => new SimpleBlobStorageConfiguration(configuration.Current.AzureStorage_Statistics_ConnectionString_Alternate, configuration.Current.AzureStorageReadAccessGeoRedundant))
+                .SingleInstance()
+                .Keyed<IBlobStorageConfiguration>(BindingKeys.AlternateStatisticsKey);
+
             // when running on Windows Azure, we use a back-end job to calculate stats totals and store in the blobs
-            builder.RegisterInstance(new JsonAggregateStatsService(configuration.Current.AzureStorage_Statistics_ConnectionString, configuration.Current.AzureStorageReadAccessGeoRedundant))
+            builder.Register(c =>
+            {
+                var primaryConfiguration = c.ResolveKeyed<IBlobStorageConfiguration>(BindingKeys.PrimaryStatisticsKey);
+                var alternateConfiguration = c.ResolveKeyed<IBlobStorageConfiguration>(BindingKeys.AlternateStatisticsKey);
+                var featureFlagService = c.Resolve<IFeatureFlagService>();
+                var jsonAggregateStatsService = new JsonAggregateStatsService(featureFlagService, primaryConfiguration, alternateConfiguration);
+                return jsonAggregateStatsService;
+            })
                 .AsSelf()
                 .As<IAggregateStatsService>()
                 .SingleInstance();
 
             // when running on Windows Azure, pull the statistics from the warehouse via storage
-            builder.RegisterInstance(new CloudReportService(configuration.Current.AzureStorage_Statistics_ConnectionString, configuration.Current.AzureStorageReadAccessGeoRedundant))
+            builder.Register(c =>
+            {
+                var primaryConfiguration = c.ResolveKeyed<IBlobStorageConfiguration>(BindingKeys.PrimaryStatisticsKey);
+                var alternateConfiguration = c.ResolveKeyed<IBlobStorageConfiguration>(BindingKeys.AlternateStatisticsKey);
+                var featureFlagService = c.Resolve<IFeatureFlagService>();
+                var cloudReportService = new CloudReportService(featureFlagService, primaryConfiguration, alternateConfiguration);
+                return cloudReportService;
+            })
                 .AsSelf()
                 .As<IReportService>()
                 .SingleInstance();
 
-            // when running on Windows Azure, download counts come from the downloads.v1.json blob
-            builder.Register(c => new SimpleBlobStorageConfiguration(configuration.Current.AzureStorage_Statistics_ConnectionString, configuration.Current.AzureStorageReadAccessGeoRedundant))
-                .SingleInstance()
-                .Keyed<IBlobStorageConfiguration>(BindingKeys.PrimaryStatisticsKey);
-
-            builder.Register(c => new SimpleBlobStorageConfiguration(configuration.Current.AzureStorage_Statistics_ConnectionString_Alternate, configuration.Current.AzureStorageReadAccessGeoRedundant))
-                .SingleInstance()
-                .Keyed<IBlobStorageConfiguration>(BindingKeys.AlternateStatisticsKey);
-
             builder.Register(c =>
             {
                 var primaryConfiguration = c.ResolveKeyed<IBlobStorageConfiguration>(BindingKeys.PrimaryStatisticsKey);

diff --git a/src/NuGetGallery/Services/CloudReportService.cs b/src/NuGetGallery/Services/CloudReportService.cs
@@ -1,23 +1,30 @@
 // 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;
 using Microsoft.WindowsAzure.Storage;
 using Microsoft.WindowsAzure.Storage.Blob;
 using Microsoft.WindowsAzure.Storage.RetryPolicies;
+using NuGetGallery.Services;
 
 namespace NuGetGallery
 {
     public class CloudReportService : IReportService
     {
         private const string _statsContainerName = "nuget-cdnstats";
-        private readonly string _connectionString;
-        private readonly bool _readAccessGeoRedundant;
+        private readonly IFeatureFlagService _featureFlagService;
+        private readonly IBlobStorageConfiguration _primaryStorageConfiguration;
+        private readonly IBlobStorageConfiguration _alternateBlobStorageConfiguration;
 
-        public CloudReportService(string connectionString, bool readAccessGeoRedundant)
+        public CloudReportService(
+            IFeatureFlagService featureFlagService,
+            IBlobStorageConfiguration primaryBlobStorageConfiguration,
+            IBlobStorageConfiguration alternateBlobStorageConfiguration)
         {
-            _connectionString = connectionString;
-            _readAccessGeoRedundant = readAccessGeoRedundant;
+            _featureFlagService = featureFlagService ?? throw new ArgumentNullException(nameof(featureFlagService));
+            _primaryStorageConfiguration = primaryBlobStorageConfiguration ?? throw new ArgumentNullException(nameof(primaryBlobStorageConfiguration));
+            _alternateBlobStorageConfiguration = alternateBlobStorageConfiguration;
         }
 
         public async Task<StatisticsReport> Load(string reportName)
@@ -42,10 +49,19 @@ public async Task<StatisticsReport> Load(string reportName)
 
         private CloudBlobContainer GetCloudBlobContainer()
         {
-            var storageAccount = CloudStorageAccount.Parse(_connectionString);
+            var connectionString = _primaryStorageConfiguration.ConnectionString;
+            var readAccessGeoRedundant = _primaryStorageConfiguration.ReadAccessGeoRedundant;
+
+            if(_alternateBlobStorageConfiguration != null && _featureFlagService.IsAlternateStatisticsSourceEnabled())
+            {
+                connectionString = _alternateBlobStorageConfiguration.ConnectionString;
+                readAccessGeoRedundant = _alternateBlobStorageConfiguration.ReadAccessGeoRedundant;
+            }
+
+            var storageAccount = CloudStorageAccount.Parse(connectionString);
             var blobClient = storageAccount.CreateCloudBlobClient();
 
-            if (_readAccessGeoRedundant)
+            if (readAccessGeoRedundant)
             {
                 blobClient.DefaultRequestOptions.LocationMode = LocationMode.PrimaryThenSecondary;
             }

diff --git a/src/NuGetGallery/Services/JsonAggregateStatsService.cs b/src/NuGetGallery/Services/JsonAggregateStatsService.cs
@@ -1,31 +1,47 @@
 // 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;
 using Microsoft.WindowsAzure.Storage;
 using Microsoft.WindowsAzure.Storage.Blob;
 using Microsoft.WindowsAzure.Storage.RetryPolicies;
 using Newtonsoft.Json;
+using NuGetGallery.Services;
 
 namespace NuGetGallery
 {
     public class JsonAggregateStatsService : IAggregateStatsService
     {
-        private readonly string _connectionString;
-        private readonly bool _readAccessGeoRedundant;
+        private readonly IFeatureFlagService _featureFlagService;
+        private readonly IBlobStorageConfiguration _primaryStorageConfiguration;
+        private readonly IBlobStorageConfiguration _alternateBlobStorageConfiguration;
 
-        public JsonAggregateStatsService(string connectionString, bool readAccessGeoRedundant)
+        public JsonAggregateStatsService(
+            IFeatureFlagService featureFlagService,
+            IBlobStorageConfiguration primaryBlobStorageConfiguration,
+            IBlobStorageConfiguration alternateBlobStorageConfiguration)
         {
-            _connectionString = connectionString;
-            _readAccessGeoRedundant = readAccessGeoRedundant;
+            _featureFlagService = featureFlagService ?? throw new ArgumentNullException(nameof(featureFlagService));
+            _primaryStorageConfiguration = primaryBlobStorageConfiguration ?? throw new ArgumentNullException(nameof(primaryBlobStorageConfiguration));
+            _alternateBlobStorageConfiguration = alternateBlobStorageConfiguration;
         }
 
         public async Task<AggregateStats> GetAggregateStats()
         {
-            CloudStorageAccount storageAccount = CloudStorageAccount.Parse(_connectionString);
+            var connectionString = _primaryStorageConfiguration.ConnectionString;
+            var readAccessGeoRedundant = _primaryStorageConfiguration.ReadAccessGeoRedundant;
+
+            if (_alternateBlobStorageConfiguration != null && _featureFlagService.IsAlternateStatisticsSourceEnabled())
+            {
+                connectionString = _alternateBlobStorageConfiguration.ConnectionString;
+                readAccessGeoRedundant = _alternateBlobStorageConfiguration.ReadAccessGeoRedundant;
+            }
+
+            CloudStorageAccount storageAccount = CloudStorageAccount.Parse(connectionString);
             CloudBlobClient blobClient = storageAccount.CreateCloudBlobClient();
 
-            if (_readAccessGeoRedundant)
+            if (readAccessGeoRedundant)
             {
                 blobClient.DefaultRequestOptions.LocationMode = LocationMode.PrimaryThenSecondary;
             }