Add knob to override the storage domain used.

b-barthel · Jan 26, 2024 · 5f550fd · 5f550fd
1 parent e907322
commit 5f550fd
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 13 deletions.
diff --git a/src/Agent.Plugins/Artifact/PipelineArtifactServer.cs b/src/Agent.Plugins/Artifact/PipelineArtifactServer.cs
@@ -18,6 +18,7 @@
 using Microsoft.VisualStudio.Services.WebApi;
 using Microsoft.VisualStudio.Services.Agent.Util;
 using Microsoft.VisualStudio.Services.BlobStore.Common;
+using Agent.Sdk.Knob;
 
 namespace Agent.Plugins
 {
@@ -50,8 +51,9 @@ internal async Task UploadAsync(
                 tracer,
                 cancellationToken);
 
-            // Get the default domain to use:
-            IDomainId domainId = clientSettings.GetDefaultDomainId();
+            // Check if the pipeline has an override domain set, if not, use the default domain from the client settings.
+            string overrideDomain = AgentKnobs.SendArtifactsToBlobstoreDomain.GetValue(context).AsString();
+            IDomainId domainId = String.IsNullOrWhiteSpace(overrideDomain) ? clientSettings.GetDefaultDomainId() : DomainIdFactory.Create(overrideDomain);
 
             var (dedupManifestClient, clientTelemetry) = DedupManifestArtifactClientFactory.Instance
                 .CreateDedupManifestClient(

diff --git a/src/Agent.Sdk/Knob/AgentKnobs.cs b/src/Agent.Sdk/Knob/AgentKnobs.cs
@@ -374,12 +374,12 @@ public class AgentKnobs
             new RuntimeKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB"),
             new EnvironmentKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB"),
             new BuiltInDefaultKnobSource("false"));
-        public static readonly Knob DisableBuildArtifactsToBlobAlternateDomain = new Knob(
-            nameof(DisableBuildArtifactsToBlobAlternateDomain),
-            "The agent may upload build artifacts to alternate domains in Blobstore. If enabled, all artifacts will go to the legacy domain.",
-            new RuntimeKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB_ALTERNATE_DOMAIN"),
-            new EnvironmentKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB_ALTERNATE_DOMAIN"),
-            new BuiltInDefaultKnobSource("false"));
+        public static readonly Knob SendArtifactsToBlobstoreDomain = new Knob(
+            nameof(SendArtifactsToBlobstoreDomain),
+            "When set, defines the domain to use to send (Build|Pipeline) artifacts to.",
+            new RuntimeKnobSource("SEND_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new EnvironmentKnobSource("SEND_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new BuiltInDefaultKnobSource(string.Empty));
 
         public static readonly Knob EnableIncompatibleBuildArtifactsPathResolution = new Knob(
             nameof(EnableIncompatibleBuildArtifactsPathResolution),

diff --git a/src/Agent.Worker/Build/FileContainerServer.cs b/src/Agent.Worker/Build/FileContainerServer.cs
@@ -355,7 +355,6 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
             BlobStoreClientTelemetryTfs clientTelemetry = null;
             try
             {
-                bool forceDefaultDomain = AgentKnobs.DisableBuildArtifactsToBlobAlternateDomain.GetValue(context).AsBoolean();
 
                 var verbose = String.Equals(context.GetVariableValueOrDefault("system.debug"), "true", StringComparison.InvariantCultureIgnoreCase);
                 Action<string> tracer = (str) => context.Output(str);
@@ -366,7 +365,9 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
                         DedupManifestArtifactClientFactory.CreateArtifactsTracer(verbose, tracer),
                         token);
 
-                IDomainId domainId = forceDefaultDomain ? WellKnownDomainIds.DefaultDomainId : clientSettings.GetDefaultDomainId();
+                // Check if the pipeline has an override domain set, if not, use the default domain from the client settings.
+                string overrideDomain = AgentKnobs.SendArtifactsToBlobstoreDomain.GetValue(context).AsString();
+                IDomainId domainId = String.IsNullOrWhiteSpace(overrideDomain) ? clientSettings.GetDefaultDomainId() : DomainIdFactory.Create(overrideDomain);
 
                 (dedupClient, clientTelemetry) = DedupManifestArtifactClientFactory.Instance
                     .CreateDedupClient(
@@ -377,6 +378,7 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
                         verbose,
                         tracer,
                         token);
+
                 // Upload to blobstore
                 var results = await BlobStoreUtils.UploadBatchToBlobstore(verbose, files, (level, uri, type) =>
                     new BuildArtifactActionRecord(level, uri, type, nameof(BlobUploadAsync), context), tracer, dedupClient, clientTelemetry, token, enableReporting: true);
@@ -397,7 +399,7 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
                 var parallelAssociateTasks = new List<Task<UploadResult>>();
                 for (int uploader = 0; uploader < concurrentUploads; uploader++)
                 {
-                    parallelAssociateTasks.Add(AssociateAsync(context, queue, token));
+                    parallelAssociateTasks.Add(AssociateAsync(context, domainId, queue, token));
                 }
 
                 // Wait for parallel associate tasks to finish.
@@ -443,7 +445,7 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
             return uploadResult;
         }
 
-        private async Task<UploadResult> AssociateAsync(IAsyncCommandContext context, ConcurrentQueue<BlobFileInfo> associateQueue, CancellationToken token)
+        private async Task<UploadResult> AssociateAsync(IAsyncCommandContext context, IDomainId domainId, ConcurrentQueue<BlobFileInfo> associateQueue, CancellationToken token)
         {
             var uploadResult = new UploadResult();
 
@@ -467,7 +469,7 @@ private async Task<UploadResult> AssociateAsync(IAsyncCommandContext context, Co
                     {
                         var length = (long)file.Node.TransitiveContentBytes;
                         response = await retryHelper.Retry(async () => await _fileContainerHttpClient.CreateItemForArtifactUpload(_containerId, itemPath, _projectId,
-                            file.DedupId.ValueString, length, token),
+                            CreateDomainHash(domainId, file.DedupId), length, token),
                                                     (retryCounter) => (int)Math.Pow(retryCounter, 2) * 5,
                                                     (exception) => true);
                         uploadResult.TotalFileSizeUploaded += length;