Reduce throttling by adding distributed locks

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Azure.Sdk.Tools.PipelineWitness.csproj

@@ -10,6 +10,7 @@
     <PackageReference Include="Azure.Storage.Blobs" Version="12.13.0" />
     <PackageReference Include="Azure.Storage.Queues" Version="12.11.0" />
     <PackageReference Include="Microsoft.ApplicationInsights.AspNetCore" Version="2.21.0" />
+    <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.29.0" />
     <PackageReference Include="Microsoft.Extensions.Azure" Version="1.3.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration.AzureAppConfiguration" Version="5.1.0" />
     <PackageReference Include="Microsoft.Extensions.Http" Version="6.0.0" />

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/BlobUploadProcessor.cs

@@ -13,11 +13,13 @@ namespace Azure.Sdk.Tools.PipelineWitness
     using System.Text.RegularExpressions;
     using System.Threading.Tasks;
 
+    using Azure.Sdk.Tools.PipelineWitness.Configuration;
     using Azure.Sdk.Tools.PipelineWitness.Services;
     using Azure.Sdk.Tools.PipelineWitness.Services.FailureAnalysis;
     using Azure.Storage.Blobs;
     using Azure.Storage.Blobs.Models;
     using Azure.Storage.Queues;
+
     using Microsoft.Extensions.Logging;
     using Microsoft.Extensions.Options;
     using Microsoft.TeamFoundation.Build.WebApi;

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Configuration/ISecretClientProvider.cs

@@ -0,0 +1,11 @@
+using System;
+
+using Azure.Security.KeyVault.Secrets;
+
+namespace Azure.Sdk.Tools.PipelineWitness.Configuration
+{
+    public interface ISecretClientProvider
+    {
+        SecretClient GetSecretClient(Uri vaultUri);
+    }
+}

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/PipelineWitnessSettings.cs → tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Configuration/PipelineWitnessSettings.cs

@@ -2,7 +2,7 @@
 using System.Collections.Generic;
 using System.Runtime.Serialization;
 
-namespace Azure.Sdk.Tools.PipelineWitness
+namespace Azure.Sdk.Tools.PipelineWitness.Configuration
 {
     public class PipelineWitnessSettings
     {
@@ -12,12 +12,17 @@ public class PipelineWitnessSettings
         public string KeyVaultUri { get; set; }
 
         /// <summary>
-        /// Gets or sets uri of the storage account use for queue processing
+        /// Gets or sets uri of the cosmos account to use
+        /// </summary>
+        public string CosmosAccountUri { get; set; }
+
+        /// <summary>
+        /// Gets or sets uri of the storage account to use for queue processing
         /// </summary>
         public string QueueStorageAccountUri { get; set; }
 
         /// <summary>
-        /// Gets or sets uri of the blob storage account use for blob export
+        /// Gets or sets uri of the blob storage account to use for blob export
         /// </summary>
         public string BlobStorageAccountUri { get; set; }
 
@@ -80,5 +85,25 @@ public class PipelineWitnessSettings
         /// Gets or sets the definition id of the pipeline owners extraction build 
         /// </summary>
         public int PipelineOwnersDefinitionId { get; set; }
+
+        /// <summary>
+        /// Gets or sets the database to use 
+        /// </summary>
+        public string CosmosDatabase { get; set; }
+
+        /// <summary>
+        /// Gets or sets the container to use for async locks
+        /// </summary>
+        public string CosmosAsyncLockContainer { get; set; }
+
+        /// <summary>
+        /// Gets or sets the authorization key for the Cosmos account
+        /// </summary>
+        public string CosmosAuthorizationKey { get; set; }
+
+        /// <summary>
+        /// Gets or sets the access token to use for Azure DevOps clients
+        /// </summary>
+        public string DevopsAccessToken { get; set; }
     }
 }

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Configuration/PostConfigureKeyVaultSettings.cs

@@ -0,0 +1,60 @@
+using System;
+using System.Linq;
+using System.Text.RegularExpressions;
+
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+
+namespace Azure.Sdk.Tools.PipelineWitness.Configuration
+{
+    public class PostConfigureKeyVaultSettings<T> : IPostConfigureOptions<T> where T : class
+    {
+        private static readonly Regex secretRegex = new Regex(@"(?<vault>https://.*?\.vault\.azure\.net)/secrets/(?<secret>.*)", RegexOptions.Compiled, TimeSpan.FromSeconds(5));
+        private readonly ILogger logger;
+        private readonly ISecretClientProvider secretClientProvider;
+
+        public PostConfigureKeyVaultSettings(ILogger<PostConfigureKeyVaultSettings<T>> logger, ISecretClientProvider secretClientProvider)
+        {
+            this.logger = logger;
+            this.secretClientProvider = secretClientProvider;
+        }
+
+        public void PostConfigure(string name, T options)
+        {
+            var stringProperties = typeof(T)
+                .GetProperties()
+                .Where(x => x.PropertyType == typeof(string));
+
+            foreach (var property in stringProperties)
+            {
+                var value = (string)property.GetValue(options);
+
+                if (value != null)
+                {
+                    var match = secretRegex.Match(value);
+
+                    if (match.Success)
+                    {
+                        var vaultUrl = match.Groups["vault"].Value;
+                        var secretName = match.Groups["secret"].Value;
+
+                        try
+                        {
+                            var secretClient = this.secretClientProvider.GetSecretClient(new Uri(vaultUrl));
+                            this.logger.LogInformation("Replacing setting property {PropertyName} with value from secret {SecretUrl}", property.Name, value);
+
+                            var response = secretClient.GetSecret(secretName);
+                            var secret = response.Value;
+
+                            property.SetValue(options, secret.Value);
+                        }
+                        catch (Exception exception)
+                        {
+                            this.logger.LogError(exception, "Unable to read secret {SecretName} from vault {VaultUrl}", secretName, vaultUrl);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Configuration/SecretClientProvider.cs

@@ -0,0 +1,22 @@
+using System;
+
+using Azure.Core;
+using Azure.Security.KeyVault.Secrets;
+
+namespace Azure.Sdk.Tools.PipelineWitness.Configuration
+{
+    public class SecretClientProvider : ISecretClientProvider
+    {
+        private readonly TokenCredential tokenCredential;
+
+        public SecretClientProvider(TokenCredential tokenCredential)
+        {
+            this.tokenCredential = tokenCredential;
+        }
+
+        public SecretClient GetSecretClient(Uri vaultUri)
+        {
+            return new SecretClient(vaultUri, this.tokenCredential);
+        }
+    }
+}

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/CosmosDatabaseInitializer.cs

@@ -0,0 +1,43 @@
+using System.Threading;
+using System.Threading.Tasks;
+
+using Azure.Sdk.Tools.PipelineWitness.Configuration;
+
+using Microsoft.Azure.Cosmos;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Options;
+
+namespace Azure.Sdk.Tools.PipelineWitness
+{
+    internal class CosmosDatabaseInitializer : IHostedService
+    {
+        private readonly CosmosClient cosmosClient;
+        private readonly IOptions<PipelineWitnessSettings> options;
+
+        public CosmosDatabaseInitializer(CosmosClient cosmosClient, IOptions<PipelineWitnessSettings> options)
+        {
+            this.cosmosClient = cosmosClient;
+            this.options = options;
+        }
+
+        public async Task StartAsync(CancellationToken cancellationToken)
+        {
+            var settings = this.options.Value;
+
+            Database database = await this.cosmosClient.CreateDatabaseIfNotExistsAsync(settings.CosmosDatabase, cancellationToken: cancellationToken);
+
+            await database.CreateContainerIfNotExistsAsync(
+                new ContainerProperties 
+                { 
+                    Id = settings.CosmosAsyncLockContainer,
+                    PartitionKeyPath = "/id",                    
+                }, 
+                cancellationToken: cancellationToken);
+        }
+
+        public Task StopAsync(CancellationToken cancellationToken)
+        {
+            return Task.CompletedTask;
+        }
+    }
+}

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Services/AzurePipelinesBuildDefinitionWorker.cs

@@ -6,6 +6,8 @@
 using Microsoft.VisualStudio.Services.WebApi;
 using System.Diagnostics;
 using Microsoft.Extensions.Options;
+using Azure.Sdk.Tools.PipelineWitness.Services.WorkTokens;
+using Azure.Sdk.Tools.PipelineWitness.Configuration;
 
 namespace Azure.Sdk.Tools.PipelineWitness.Services
 {
@@ -14,33 +16,51 @@ public class AzurePipelinesBuildDefinitionWorker : BackgroundService
         private readonly ILogger<AzurePipelinesBuildDefinitionWorker> logger;
         private readonly BlobUploadProcessor runProcessor;
         private readonly IOptions<PipelineWitnessSettings> options;
+        private IAsyncLockProvider asyncLockProvider;
 
         public AzurePipelinesBuildDefinitionWorker(
             ILogger<AzurePipelinesBuildDefinitionWorker> logger,
             BlobUploadProcessor runProcessor,
+            IAsyncLockProvider asyncLockProvider,
             IOptions<PipelineWitnessSettings> options)
         {
             this.logger = logger;
             this.runProcessor = runProcessor;
             this.options = options;
+            this.asyncLockProvider = asyncLockProvider;
         }
 
         protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         {
+            var processEvery = TimeSpan.FromMinutes(60);
+
             while (true)
             {
                 var stopWatch = Stopwatch.StartNew();
                 var settings = this.options.Value;
 
-                foreach (var project in settings.Projects)
+                try
+                {
+                    await using var asyncLock = await this.asyncLockProvider.GetLockAsync("UpdateBuildDefinitions", processEvery, stoppingToken);
+
+                    // if there's no asyncLock, this process has alread completed in the last hour
+                    if (asyncLock != null)
+                    {
+                        foreach (var project in settings.Projects)
+                        {
+                            await this.runProcessor.UploadBuildDefinitionBlobsAsync(settings.Account, project);
+                        }
+                    }
+                }
+                catch(Exception ex)
                 {
-                    await this.runProcessor.UploadBuildDefinitionBlobsAsync(settings.Account, project);
+                    this.logger.LogError(ex, "Error processing build definitions");
                 }
 
                 var duration = settings.BuildDefinitionLoopPeriod - stopWatch.Elapsed;
                 if (duration > TimeSpan.Zero)
                 {
-                    await Task.Delay(duration);
+                    await Task.Delay(duration, stoppingToken);
                 }
             }
         }

tools/pipeline-witness/Azure.Sdk.Tools.PipelineWitness/Services/BuildCompleteQueueWorker.cs

@@ -2,12 +2,16 @@
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+
+using Azure.Sdk.Tools.PipelineWitness.Configuration;
 using Azure.Storage.Queues;
 using Azure.Storage.Queues.Models;
+
 using Microsoft.ApplicationInsights;
 using Microsoft.ApplicationInsights.DataContracts;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
+
 using Newtonsoft.Json.Linq;
 
 namespace Azure.Sdk.Tools.PipelineWitness.Services