Speed up bulk insertions (#530)

The initialization phase of the job has to insert 1.5 million records into the validation DB in batches of 1,000 records. On DEV, each batch was taking 30 - 120 seconds using Entity Framework. This fix uses `SqlBulkCopy` to speed up batches to be sub-second.

Author: loic-sharma

src/NuGet.Jobs.Common/DelegateSqlConnectionFactory.cs

@@ -0,0 +1,52 @@
+// 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.Data.SqlClient;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using NuGet.Jobs.Configuration;
+
+namespace NuGet.Jobs
+{
+    public class DelegateSqlConnectionFactory<TbDbConfiguration> : ISqlConnectionFactory<TbDbConfiguration>
+        where TbDbConfiguration : IDbConfiguration
+    {
+        private readonly Func<Task<SqlConnection>> _connectionFunc;
+        private readonly ILogger<DelegateSqlConnectionFactory<TbDbConfiguration>> _logger;
+
+        public DelegateSqlConnectionFactory(Func<Task<SqlConnection>> connectionFunc, ILogger<DelegateSqlConnectionFactory<TbDbConfiguration>> logger)
+        {
+            _connectionFunc = connectionFunc ?? throw new ArgumentNullException(nameof(connectionFunc));
+            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+        }
+
+        public Task<SqlConnection> CreateAsync() => _connectionFunc();
+
+        public async Task<SqlConnection> OpenAsync()
+        {
+            SqlConnection connection = null;
+
+            try
+            {
+                _logger.LogDebug("Opening SQL connection...");
+
+                connection = await _connectionFunc();
+
+                await connection.OpenAsync();
+
+                _logger.LogDebug("Opened SQL connection");
+
+                return connection;
+            }
+            catch (Exception e)
+            {
+                _logger.LogError(0, e, "Unable to open SQL connection due to exception");
+
+                connection?.Dispose();
+
+                throw;
+            }
+        }
+    }
+}

src/NuGet.Jobs.Common/ISqlConnectionFactory.cs

@@ -0,0 +1,38 @@
+// 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.Data.SqlClient;
+using System.Threading.Tasks;
+using NuGet.Jobs.Configuration;
+
+namespace NuGet.Jobs
+{
+    /// <summary>
+    /// A factory to create and open <see cref="SqlConnection"/>s.
+    /// </summary>
+    public interface ISqlConnectionFactory
+    {
+        /// <summary>
+        /// Create an unopened SQL connection.
+        /// </summary>
+        /// <returns>The unopened SQL connection.</returns>
+        Task<SqlConnection> CreateAsync();
+
+        /// <summary>
+        /// Create and then open a SQL connection.
+        /// </summary>
+        /// <returns>A task that creates and then opens a SQL connection.</returns>
+        Task<SqlConnection> OpenAsync();
+    }
+
+    /// <summary>
+    /// A factory to create and open <see cref="SqlConnection"/>s for a specific
+    /// <see cref="TDbConfiguration"/>. This type can be used to avoid Dependency
+    /// Injection key bindings.
+    /// </summary>
+    /// <typeparam name="TDbConfiguration">The configuration used to create the connection.</typeparam>
+    public interface ISqlConnectionFactory<TDbConfiguration> : ISqlConnectionFactory
+        where TDbConfiguration : IDbConfiguration
+    {
+    }
+}

src/NuGet.Jobs.Common/JobBase.cs

@@ -16,6 +16,8 @@
 
 namespace NuGet.Jobs
 {
+    using ICoreSqlConnectionFactory = NuGet.Services.Sql.ISqlConnectionFactory;
+
     public abstract class JobBase
     {
         private readonly EventSource _jobEventSource;
@@ -29,7 +31,7 @@ protected JobBase(EventSource jobEventSource)
         {
             JobName = GetType().ToString();
             _jobEventSource = jobEventSource;
-            SqlConnectionFactories = new Dictionary<string, ISqlConnectionFactory>();
+            SqlConnectionFactories = new Dictionary<string, ICoreSqlConnectionFactory>();
         }
 
         public string JobName { get; private set; }
@@ -38,7 +40,7 @@ protected JobBase(EventSource jobEventSource)
 
         protected ILogger Logger { get; private set; }
 
-        private Dictionary<string, ISqlConnectionFactory> SqlConnectionFactories { get; }
+        private Dictionary<string, ICoreSqlConnectionFactory> SqlConnectionFactories { get; }
 
         public void SetLogger(ILoggerFactory loggerFactory, ILogger logger)
         {
@@ -56,11 +58,10 @@ public void SetLogger(ILoggerFactory loggerFactory, ILogger logger)
         /// </summary>
         public abstract Task Run();
 
-
         /// <summary>
         /// Test connection early to fail fast, and log connection diagnostics.
         /// </summary>
-        private async Task TestConnection(string name, ISqlConnectionFactory connectionFactory)
+        private async Task TestConnection(string name, ICoreSqlConnectionFactory connectionFactory)
         {
             try
             {
@@ -98,17 +99,17 @@ public SqlConnectionStringBuilder GetDatabaseRegistration<T>()
         /// </summary>
         /// <returns>ConnectionStringBuilder, used for diagnostics.</returns>
         public SqlConnectionStringBuilder RegisterDatabase<T>(
-            IServiceProvider serviceProvider,
+            IServiceProvider services,
             bool testConnection = true)
             where T : IDbConfiguration
         {
-            if (serviceProvider == null)
+            if (services == null)
             {
-                throw new ArgumentNullException(nameof(serviceProvider));
+                throw new ArgumentNullException(nameof(services));
             }
 
-            var secretInjector = serviceProvider.GetRequiredService<ISecretInjector>();
-            var connectionString = serviceProvider.GetRequiredService<IOptionsSnapshot<T>>().Value.ConnectionString;
+            var secretInjector = services.GetRequiredService<ISecretInjector>();
+            var connectionString = services.GetRequiredService<IOptionsSnapshot<T>>().Value.ConnectionString;
             var connectionFactory = new AzureSqlConnectionFactory(connectionString, secretInjector);
 
             return RegisterDatabase(GetDatabaseKey<T>(), connectionString, testConnection, secretInjector);
@@ -167,13 +168,13 @@ private SqlConnectionStringBuilder RegisterDatabase(
             return connectionFactory.SqlConnectionStringBuilder;
         }
 
-        private ISqlConnectionFactory GetSqlConnectionFactory<T>()
+        private ICoreSqlConnectionFactory GetSqlConnectionFactory<T>()
             where T : IDbConfiguration
         {
             return GetSqlConnectionFactory(GetDatabaseKey<T>());
         }
 
-        private ISqlConnectionFactory GetSqlConnectionFactory(string name)
+        private ICoreSqlConnectionFactory GetSqlConnectionFactory(string name)
         {
             if (!SqlConnectionFactories.ContainsKey(name))
             {

src/NuGet.Jobs.Common/JsonConfigurationJob.cs

@@ -117,6 +117,27 @@ protected virtual void ConfigureDefaultJobServices(IServiceCollection services,
 
             services.AddSingleton(new TelemetryClient());
             services.AddTransient<ITelemetryClient, TelemetryClientWrapper>();
+
+            services.AddScoped<ISqlConnectionFactory<GalleryDbConfiguration>>(p =>
+            {
+                return new DelegateSqlConnectionFactory<GalleryDbConfiguration>(
+                    CreateSqlConnectionAsync<GalleryDbConfiguration>,
+                    p.GetRequiredService<ILogger<DelegateSqlConnectionFactory<GalleryDbConfiguration>>>());
+            });
+
+            services.AddScoped<ISqlConnectionFactory<StatisticsDbConfiguration>>(p =>
+            {
+                return new DelegateSqlConnectionFactory<StatisticsDbConfiguration>(
+                    CreateSqlConnectionAsync<StatisticsDbConfiguration>,
+                    p.GetRequiredService<ILogger<DelegateSqlConnectionFactory<StatisticsDbConfiguration>>>());
+            });
+
+            services.AddScoped<ISqlConnectionFactory<ValidationDbConfiguration>>(p =>
+            {
+                return new DelegateSqlConnectionFactory<ValidationDbConfiguration>(
+                    CreateSqlConnectionAsync<ValidationDbConfiguration>,
+                    p.GetRequiredService<ILogger<DelegateSqlConnectionFactory<ValidationDbConfiguration>>>());
+            });
         }
 
         private void ConfigureLibraries(IServiceCollection services)

src/NuGet.Jobs.Common/NuGet.Jobs.Common.csproj

@@ -50,8 +50,10 @@
     <Compile Include="Configuration\StatisticsDbConfiguration.cs" />
     <Compile Include="Configuration\ValidationDbConfiguration.cs" />
     <Compile Include="Configuration\ValidationStorageConfiguration.cs" />
+    <Compile Include="DelegateSqlConnectionFactory.cs" />
     <Compile Include="Extensions\LoggerExtensions.cs" />
     <Compile Include="Extensions\XElementExtensions.cs" />
+    <Compile Include="ISqlConnectionFactory.cs" />
     <Compile Include="JsonConfigurationJob.cs" />
     <Compile Include="SecretReader\ISecretReaderFactory.cs" />
     <Compile Include="SecretReader\SecretReaderFactory.cs" />

src/NuGet.Services.Revalidate/Job.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel.Design;
+using System.Data.SqlClient;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
@@ -78,7 +79,7 @@ public override async Task Run()
 
                         preinstalledPackagesNames.UnionWith(packagesInPath);
                     }
-
+                        
                     File.WriteAllText(_preinstalledSetPath, JsonConvert.SerializeObject(preinstalledPackagesNames));
 
                     Logger.LogInformation("Rebuilt the preinstalled package set. Found {PreinstalledPackages} package ids", preinstalledPackagesNames.Count);
@@ -139,6 +140,7 @@ protected override void ConfigureJobServices(IServiceCollection services, IConfi
             services.AddTransient<ITelemetryClient, TelemetryClientWrapper>();
 
             services.AddTransient<IPackageRevalidationStateService, PackageRevalidationStateService>();
+            services.AddTransient<IPackageRevalidationInserter, PackageRevalidationInserter>();
             services.AddTransient<IRevalidationJobStateService, RevalidationJobStateService>();
             services.AddTransient<IRevalidationStateService, RevalidationStateService>();
 

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

@@ -58,12 +58,14 @@
     <Compile Include="Services\HealthService.cs" />
     <Compile Include="Services\IGalleryService.cs" />
     <Compile Include="Services\IHealthService.cs" />
+    <Compile Include="Services\IPackageRevalidationInserter.cs" />
     <Compile Include="Services\IRevalidationQueue.cs" />
     <Compile Include="Services\IRevalidationJobStateService.cs" />
     <Compile Include="Services\IPackageRevalidationStateService.cs" />
     <Compile Include="Services\IRevalidationService.cs" />
     <Compile Include="Services\IRevalidationThrottler.cs" />
     <Compile Include="Services\ISingletonService.cs" />
+    <Compile Include="Services\PackageRevalidationInserter.cs" />
     <Compile Include="Services\RevalidationOperation.cs" />
     <Compile Include="Services\RevalidationQueue.cs" />
     <Compile Include="Services\RevalidationResult.cs" />

src/NuGet.Services.Revalidate/Services/IPackageRevalidationInserter.cs

@@ -0,0 +1,14 @@
+// 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;
+
+namespace NuGet.Services.Revalidate
+{
+    public interface IPackageRevalidationInserter
+    {
+        Task AddPackageRevalidationsAsync(IReadOnlyList<PackageRevalidation> revalidations);
+    }
+}

src/NuGet.Services.Revalidate/Services/PackageRevalidationInserter.cs

@@ -0,0 +1,95 @@
+// 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.Data;
+using System.Data.SqlClient;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using NuGet.Jobs;
+using NuGet.Jobs.Configuration;
+using NuGet.Jobs.Validation;
+using NuGet.Services.Validation;
+
+namespace NuGet.Services.Revalidate
+{
+    public class PackageRevalidationInserter : IPackageRevalidationInserter
+    {
+        private const string TableName = "[dbo].[PackageRevalidations]";
+
+        private const string PackageIdColumn = "PackageId";
+        private const string PackageNormalizedVersionColumn = "PackageNormalizedVersion";
+        private const string EnqueuedColumn = "Enqueued";
+        private const string ValidationTrackingIdColumn = "ValidationTrackingId";
+        private const string CompletedColumn = "Completed";
+
+        private readonly ISqlConnectionFactory<ValidationDbConfiguration> _connectionFactory;
+        private readonly ILogger<PackageRevalidationInserter> _logger;
+
+        public PackageRevalidationInserter(
+            ISqlConnectionFactory<ValidationDbConfiguration> connectionFactory,
+            ILogger<PackageRevalidationInserter> logger)
+        {
+            _connectionFactory = connectionFactory ?? throw new ArgumentNullException(nameof(connectionFactory));
+            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+        }
+
+        public async Task AddPackageRevalidationsAsync(IReadOnlyList<PackageRevalidation> revalidations)
+        {
+            _logger.LogDebug("Persisting package revalidations to database...");
+
+            var table = PrepareTable(revalidations);
+
+            using (var connection = await _connectionFactory.OpenAsync())
+            {
+                var bulkCopy = new SqlBulkCopy(
+                    connection,
+                    SqlBulkCopyOptions.TableLock | SqlBulkCopyOptions.FireTriggers | SqlBulkCopyOptions.UseInternalTransaction,
+                    externalTransaction: null);
+
+                foreach (DataColumn column in table.Columns)
+                {
+                    bulkCopy.ColumnMappings.Add(column.ColumnName, column.ColumnName);
+                }
+
+                bulkCopy.DestinationTableName = TableName;
+                bulkCopy.WriteToServer(table);
+            }
+
+            _logger.LogDebug("Finished persisting package revalidations to database...");
+        }
+
+        private DataTable PrepareTable(IReadOnlyList<PackageRevalidation> revalidations)
+        {
+            // Prepare the table.
+            var table = new DataTable();
+
+            table.Columns.Add(PackageIdColumn, typeof(string));
+            table.Columns.Add(PackageNormalizedVersionColumn, typeof(string));
+            table.Columns.Add(CompletedColumn, typeof(bool));
+
+            var enqueued = table.Columns.Add(EnqueuedColumn, typeof(DateTime));
+            var trackingId = table.Columns.Add(ValidationTrackingIdColumn, typeof(Guid));
+
+            enqueued.AllowDBNull = true;
+            trackingId.AllowDBNull = true;
+
+            // Populate the table.
+            foreach (var revalidation in revalidations)
+            {
+                var row = table.NewRow();
+
+                row[PackageIdColumn] = revalidation.PackageId;
+                row[PackageNormalizedVersionColumn] = revalidation.PackageNormalizedVersion;
+                row[EnqueuedColumn] = ((object)revalidation.Enqueued) ?? DBNull.Value;
+                row[ValidationTrackingIdColumn] = ((object)revalidation.ValidationTrackingId) ?? DBNull.Value;
+                row[CompletedColumn] = revalidation.Completed;
+
+                table.Rows.Add(row);
+            }
+
+            return table;
+        }
+    }
+}