Add multidomain support to Build Artifacts

src/Agent.Plugins/Artifact/FileContainerProvider.cs

@@ -11,6 +11,7 @@
 using Microsoft.VisualStudio.Services.Agent.Util;
 using Microsoft.VisualStudio.Services.BlobStore.Common;
 using Microsoft.VisualStudio.Services.BlobStore.WebApi;
+using Microsoft.VisualStudio.Services.BlobStore.WebApi.Contracts;
 using Microsoft.VisualStudio.Services.Content.Common;
 using Microsoft.VisualStudio.Services.Content.Common.Tracing;
 using Microsoft.VisualStudio.Services.FileContainer;
@@ -149,27 +150,44 @@ private async Task DownloadFileContainerAsync(IEnumerable<FileContainerItem> ite
             // Only initialize these clients if we know we need to download from Blobstore
             // If a client cannot connect to Blobstore, we shouldn't stop them from downloading from FCS
             var downloadFromBlob = !AgentKnobs.DisableBuildArtifactsToBlob.GetValue(context).AsBoolean();
-            DedupStoreClient dedupClient = null;
+            Dictionary<IDomainId,DedupStoreClient> dedupClientTable = new Dictionary<IDomainId, DedupStoreClient>();
             BlobStoreClientTelemetryTfs clientTelemetry = null;
             if (downloadFromBlob && fileItems.Any(x => x.BlobMetadata != null))
             {
+                // this is not the most efficient but good enough for now:
+                var domains = fileItems.Select(x => GetDomainIdAndDedupIdFromArtifactHash(x.BlobMetadata.ArtifactHash).domainId).Distinct();
+                DedupStoreClient dedupClient = null;
                 try
                 {
-                    (dedupClient, clientTelemetry) = await DedupManifestArtifactClientFactory.Instance.CreateDedupClientAsync(
-                        false,
-                        (str) => this.tracer.Info(str),
-                        this.connection,
-                        DedupManifestArtifactClientFactory.Instance.GetDedupStoreClientMaxParallelism(context),
+                    BlobstoreClientSettings clientSettings = await BlobstoreClientSettings.GetClientSettingsAsync(
+                        connection,
                         Microsoft.VisualStudio.Services.BlobStore.WebApi.Contracts.Client.BuildArtifact,
+                        tracer,
                         cancellationToken);
+
+                    foreach(var domainId in domains)
+                    {
+                        (dedupClient, clientTelemetry) = DedupManifestArtifactClientFactory.Instance.CreateDedupClient(
+                            this.connection,
+                            domainId,
+                            DedupManifestArtifactClientFactory.Instance.GetDedupStoreClientMaxParallelism(context),
+                            clientSettings.GetRedirectTimeout(),
+                            false,
+                            (str) => this.tracer.Info(str),
+                            cancellationToken);
+
+                        dedupClientTable.Add(domainId, dedupClient);
+                    }
                 }
                 catch (SocketException e)
                 {
                     ExceptionsUtil.HandleSocketException(e, connection.Uri.ToString(), context.Warning);
+                    // Fall back to streaming through TFS if we cannot reach blobstore for any reason
+                    downloadFromBlob = false;
                 }
                 catch
                 {
-                    var blobStoreHost = dedupClient.Client.BaseAddress.Host;
+                    var blobStoreHost = dedupClient?.Client.BaseAddress.Host;
                     var allowListLink = BlobStoreWarningInfoProvider.GetAllowListLinkForCurrentPlatform();
                     var warningMessage = StringUtil.Loc("BlobStoreDownloadWarning", blobStoreHost, allowListLink);
 
@@ -191,7 +209,8 @@ await AsyncHttpRetryHelper.InvokeVoidAsync(
                             tracer.Info($"Downloading: {targetPath}");
                             if (item.BlobMetadata != null && downloadFromBlob)
                             {
-                                await this.DownloadFileFromBlobAsync(context, containerIdAndRoot, targetPath, projectId, item, dedupClient, clientTelemetry, cancellationToken);
+                                var client = dedupClientTable[GetDomainIdAndDedupIdFromArtifactHash(item.BlobMetadata.ArtifactHash).domainId];
+                                await this.DownloadFileFromBlobAsync(context, containerIdAndRoot, targetPath, projectId, item, client, clientTelemetry, cancellationToken);
                             }
                             else
                             {
@@ -336,6 +355,22 @@ private async Task<Stream> DownloadFileAsync(
             return responseStream;
         }
 
+        private static (IDomainId domainId, DedupIdentifier dedupId) GetDomainIdAndDedupIdFromArtifactHash(string artifactHash)
+        {
+            string[] parts = artifactHash.Split(',');
+            if(parts.Length == 1)
+            {
+                // legacy format is always in the default domain:
+                return (WellKnownDomainIds.DefaultDomainId, DedupIdentifier.Deserialize(parts[0]));
+            }
+            else if(parts.Length==2)
+            {
+                // Multidomain format is in the form of <domainId>,<dedupId>
+                return (DomainIdFactory.Create(parts[0]), DedupIdentifier.Deserialize(parts[1]));
+            }
+            throw new ArgumentException($"Invalid artifact hash: {artifactHash}", nameof(artifactHash));
+        }
+
         private async Task DownloadFileFromBlobAsync(
             AgentTaskPluginExecutionContext context,
             (long, string) containerIdAndRoot,
@@ -346,7 +381,7 @@ private async Task DownloadFileFromBlobAsync(
             BlobStoreClientTelemetryTfs clientTelemetry,
             CancellationToken cancellationToken)
         {
-            var dedupIdentifier = DedupIdentifier.Deserialize(item.BlobMetadata.ArtifactHash);
+            (var domainId, var dedupIdentifier) = GetDomainIdAndDedupIdFromArtifactHash(item.BlobMetadata.ArtifactHash);
 
             var downloadRecord = clientTelemetry.CreateRecord<BuildArtifactActionRecord>((level, uri, type) =>
                 new BuildArtifactActionRecord(level, uri, type, nameof(DownloadFileContainerAsync), context));

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
 {
@@ -44,14 +45,15 @@ internal async Task UploadAsync(
             // Get the client settings, if any.
             var tracer = DedupManifestArtifactClientFactory.CreateArtifactsTracer(verbose: false, (str) => context.Output(str));
             VssConnection connection = context.VssConnection;
-            var clientSettings = await DedupManifestArtifactClientFactory.GetClientSettingsAsync(
+            var clientSettings = await BlobstoreClientSettings.GetClientSettingsAsync(
                 connection,
                 Microsoft.VisualStudio.Services.BlobStore.WebApi.Contracts.Client.PipelineArtifact,
                 tracer,
                 cancellationToken);
 
-            // Get the default domain to use:
-            IDomainId domainId = DedupManifestArtifactClientFactory.GetDefaultDomainId(clientSettings, tracer);
+            // Check if the pipeline has an override domain set, if not, use the default domain from the client settings.
+            string overrideDomain = AgentKnobs.SendPipelineArtifactsToBlobstoreDomain.GetValue(context).AsString();
+            IDomainId domainId = String.IsNullOrWhiteSpace(overrideDomain) ? clientSettings.GetDefaultDomainId() : DomainIdFactory.Create(overrideDomain);
 
             var (dedupManifestClient, clientTelemetry) = DedupManifestArtifactClientFactory.Instance
                 .CreateDedupManifestClient(

src/Agent.Sdk/Knob/AgentKnobs.cs

@@ -374,6 +374,18 @@ public class AgentKnobs
             new RuntimeKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB"),
             new EnvironmentKnobSource("DISABLE_BUILD_ARTIFACTS_TO_BLOB"),
             new BuiltInDefaultKnobSource("false"));
+        public static readonly Knob SendBuildArtifactsToBlobstoreDomain = new Knob(
+            nameof(SendBuildArtifactsToBlobstoreDomain),
+            "When set, defines the domain to use to send Build artifacts to.",
+            new RuntimeKnobSource("SEND_BUILD_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new EnvironmentKnobSource("SEND_BUILD_ARTIFACT_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new BuiltInDefaultKnobSource(string.Empty));
+        public static readonly Knob SendPipelineArtifactsToBlobstoreDomain = new Knob(
+            nameof(SendPipelineArtifactsToBlobstoreDomain),
+            "When set, defines the domain to use to send Pipeline artifacts to.",
+            new RuntimeKnobSource("SEND_PIPELINE_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new EnvironmentKnobSource("SEND_PIPELINE_ARTIFACT_ARTIFACTS_TO_BLOBSTORE_DOMAIN"),
+            new BuiltInDefaultKnobSource(string.Empty));
 
         public static readonly Knob EnableIncompatibleBuildArtifactsPathResolution = new Knob(
             nameof(EnableIncompatibleBuildArtifactsPathResolution),

src/Agent.Worker/Build/FileContainerServer.cs

@@ -3,8 +3,13 @@
 
 using Agent.Sdk.Knob;
 using Agent.Sdk.Util;
+using BuildXL.Cache.ContentStore.Hashing;
+using Microsoft.TeamFoundation.DistributedTask.WebApi;
 using Microsoft.VisualStudio.Services.Agent.Blob;
 using Microsoft.VisualStudio.Services.Agent.Util;
+using Microsoft.VisualStudio.Services.BlobStore.WebApi;
+using Microsoft.VisualStudio.Services.BlobStore.Common;
+using Microsoft.VisualStudio.Services.BlobStore.WebApi.Contracts;
 using Microsoft.VisualStudio.Services.FileContainer.Client;
 using System;
 using System.Collections.Concurrent;
@@ -18,8 +23,6 @@
 using System.Net.Http;
 using System.Net;
 using System.Net.Sockets;
-using Microsoft.TeamFoundation.DistributedTask.WebApi;
-using Microsoft.VisualStudio.Services.BlobStore.WebApi;
 
 
 namespace Microsoft.VisualStudio.Services.Agent.Worker.Build
@@ -326,6 +329,16 @@ private async Task<UploadResult> UploadAsync(IAsyncCommandContext context, int u
 
             return new UploadResult(failedFiles, uploadedSize);
         }
+        public static string CreateDomainHash(IDomainId domainId, DedupIdentifier dedupId)
+        {
+            if (domainId != WellKnownDomainIds.DefaultDomainId)
+            {
+                // Only use the new format domainId,dedupId if we aren't going to the default domain as this is a breaking change:
+                return $"{domainId.Serialize()},{dedupId.ValueString}";
+            }
+            // We are still uploading to the default domain so use the don't use the new format:
+            return dedupId.ValueString;
+        }
 
         private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, IReadOnlyList<string> files, int concurrentUploads, CancellationToken token)
         {
@@ -342,20 +355,33 @@ private async Task<UploadResult> BlobUploadAsync(IAsyncCommandContext context, I
             BlobStoreClientTelemetryTfs clientTelemetry = null;
             try
             {
+
                 var verbose = String.Equals(context.GetVariableValueOrDefault("system.debug"), "true", StringComparison.InvariantCultureIgnoreCase);
-                int maxParallelism = context.GetHostContext().GetService<IConfigurationStore>().GetSettings().MaxDedupParallelism;
-                (dedupClient, clientTelemetry) = await DedupManifestArtifactClientFactory.Instance
-                    .CreateDedupClientAsync(
+                Action<string> tracer = (str) => context.Output(str);
+
+                var clientSettings = await BlobstoreClientSettings.GetClientSettingsAsync(
+                        _connection,
+                        Microsoft.VisualStudio.Services.BlobStore.WebApi.Contracts.Client.BuildArtifact,
+                        DedupManifestArtifactClientFactory.CreateArtifactsTracer(verbose, tracer),
+                        token);
+
+                // Check if the pipeline has an override domain set, if not, use the default domain from the client settings.
+                string overrideDomain = AgentKnobs.SendBuildArtifactsToBlobstoreDomain.GetValue(context).AsString();
+                IDomainId domainId = String.IsNullOrWhiteSpace(overrideDomain) ? clientSettings.GetDefaultDomainId() : DomainIdFactory.Create(overrideDomain);
+
+                (dedupClient, clientTelemetry) = DedupManifestArtifactClientFactory.Instance
+                    .CreateDedupClient(
+                        _connection,
+                        domainId,
+                        context.GetHostContext().GetService<IConfigurationStore>().GetSettings().MaxDedupParallelism,
+                        clientSettings.GetRedirectTimeout(),
                         verbose,
-                        (str) => context.Output(str),
-                        this._connection,
-                        maxParallelism,
-                        BlobStore.WebApi.Contracts.Client.BuildArtifact,
+                        tracer,
                         token);
 
                 // Upload to blobstore
                 var results = await BlobStoreUtils.UploadBatchToBlobstore(verbose, files, (level, uri, type) =>
-                    new BuildArtifactActionRecord(level, uri, type, nameof(BlobUploadAsync), context), (str) => context.Output(str), dedupClient, clientTelemetry, token, enableReporting: true);
+                    new BuildArtifactActionRecord(level, uri, type, nameof(BlobUploadAsync), context), tracer, dedupClient, clientTelemetry, token, enableReporting: true);
 
                 // Associate with TFS
                 context.Output(StringUtil.Loc("AssociateFiles"));
@@ -373,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.
@@ -419,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();
 
@@ -443,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;

src/Microsoft.VisualStudio.Services.Agent/Blob/BlobStoreUtils.cs

@@ -120,7 +120,7 @@ private static async Task<List<BlobFileInfo>> GenerateHashes(IReadOnlyList<strin
                     var itemPath = filePaths[i];
                     try
                     {
-                        var dedupNode = await ChunkerHelper.CreateFromFileAsync(FileSystem.Instance, itemPath, cancellationToken, false);
+                        var dedupNode = await ChunkerHelper.CreateFromFileAsync(FileSystem.Instance, itemPath, false, ChunkerHelper.DefaultChunkHashType, cancellationToken);
                         nodes[i] = new BlobFileInfo
                         {
                             Path = itemPath,