Add task to download files from helix results container

src/Microsoft.DotNet.Helix/JobSender/IJobDefinition.cs

@@ -19,6 +19,7 @@ public interface IJobDefinition
         IJobDefinition WithStorageAccountConnectionString(string accountConnectionString);
         IJobDefinition WithResultsContainerName(string resultsContainerName);
         IJobDefinition WithResultsStorageAccountConnectionString(string resultsAccountConnectionString);
+        IJobDefinition WithDefaultResultsContainer();
         IJobDefinition WithMaxRetryCount(int? maxRetryCount);
         Task<ISentJob> SendAsync(Action<string> log = null);
     }

src/Microsoft.DotNet.Helix/JobSender/ISentJob.cs

@@ -5,5 +5,7 @@ namespace Microsoft.DotNet.Helix.Client
     public interface ISentJob
     {
         string CorrelationId { get; }
+        string ResultsContainerUri { get; }
+        string ResultsContainerReadSAS { get; }
     }
 }

src/Microsoft.DotNet.Helix/JobSender/JobDefinition.cs

@@ -21,6 +21,7 @@ internal class JobDefinition : IJobDefinitionWithSource,
     {
         private readonly Dictionary<string, string> _properties;
         private readonly List<WorkItemDefinition> _workItems;
+        private bool _withDefaultResultsContainer;
 
         public JobDefinition(IJob jobApi)
         {
@@ -133,6 +134,12 @@ public IJobDefinition WithResultsStorageAccountConnectionString(string resultsAc
             return this;
         }
 
+        public IJobDefinition WithDefaultResultsContainer()
+        {
+            _withDefaultResultsContainer = true;
+            return this;
+        }
+
         public async Task<ISentJob> SendAsync(Action<string> log = null)
         {
             IBlobHelper storage;
@@ -148,12 +155,16 @@ public async Task<ISentJob> SendAsync(Action<string> log = null)
             IBlobContainer storageContainer = await storage.GetContainerAsync(TargetContainerName);
             var jobList = new List<JobListEntry>();
 
-            IBlobHelper resultsStorage = null;
             IBlobContainer resultsStorageContainer = null;
             if (!string.IsNullOrEmpty(ResultsStorageAccountConnectionString))
             {
-                resultsStorage = new ConnectionStringBlobHelper(ResultsStorageAccountConnectionString);
-                resultsStorageContainer = await resultsStorage.GetContainerAsync(TargetContainerName);
+
+                IBlobHelper resultsStorage = new ConnectionStringBlobHelper(ResultsStorageAccountConnectionString);
+                resultsStorageContainer = await resultsStorage.GetContainerAsync(TargetResultsContainerName);
+            }
+            else if (_withDefaultResultsContainer)
+            {
+                resultsStorageContainer = await storage.GetContainerAsync(TargetResultsContainerName);
             }
 
             List<string> correlationPayloadUris =
@@ -195,7 +206,7 @@ public async Task<ISentJob> SendAsync(Action<string> log = null)
                 ex => log?.Invoke($"Starting job failed with {ex}\nRetrying..."));
 
 
-            return new SentJob(JobApi, newJob);
+            return new SentJob(JobApi, newJob, resultsStorageContainer?.Uri, string.IsNullOrEmpty(Creator) ? resultsStorageContainer?.ReadSas : string.Empty);
         }
 
         public IJobDefinitionWithTargetQueue WithBuild(string buildNumber)

src/Microsoft.DotNet.Helix/JobSender/SentJob.cs

@@ -9,13 +9,17 @@ namespace Microsoft.DotNet.Helix.Client
 {
     internal class SentJob : ISentJob
     {
-        public SentJob(IJob jobApi, JobCreationResult newJob)
+        public SentJob(IJob jobApi, JobCreationResult newJob, string resultsContainerUri, string resultsContainerReadSAS)
         {
             JobApi = jobApi;
             CorrelationId = newJob.Name;
+            ResultsContainerUri = resultsContainerUri;
+            ResultsContainerReadSAS = resultsContainerReadSAS;
         }
 
         public IJob JobApi { get; }
         public string CorrelationId { get; }
+        public string ResultsContainerUri { get; }
+        public string ResultsContainerReadSAS { get; }
     }
 }

src/Microsoft.DotNet.Helix/Sdk/DownloadFromResultsContainer.cs

@@ -0,0 +1,116 @@
+using Microsoft.Build.Framework;
+using Microsoft.WindowsAzure.Storage;
+using Microsoft.WindowsAzure.Storage.Auth;
+using Microsoft.WindowsAzure.Storage.Blob;
+using System;
+using System.IO;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.DotNet.Helix.Sdk
+{
+    public class DownloadFromResultsContainer : BaseTask, ICancelableTask
+    {
+        [Required]
+        public ITaskItem[] WorkItems { get; set; }
+
+        [Required]
+        public string ResultsContainer { get; set; }
+
+        [Required]
+        public string OutputDirectory { get; set; }
+
+        [Required]
+        public string JobId { get; set; }
+
+        [Required]
+        public ITaskItem[] MetadataToWrite { get; set; }
+
+        public string ResultsContainerReadSAS { get; set; }
+
+        private const string MetadataFile = "metadata.txt";
+
+        private readonly CancellationTokenSource _cancellationSource = new CancellationTokenSource();
+
+        public void Cancel() => _cancellationSource.Cancel();
+
+        public override bool Execute()
+        {
+            if (string.IsNullOrEmpty(ResultsContainer))
+            {
+                LogRequiredParameterError(nameof(ResultsContainer));
+            }
+
+            if (string.IsNullOrEmpty(OutputDirectory))
+            {
+                LogRequiredParameterError(nameof(OutputDirectory));
+            }
+
+            if (string.IsNullOrEmpty(JobId))
+            {
+                LogRequiredParameterError(nameof(JobId));
+            }
+
+            if (!Log.HasLoggedErrors)
+            {
+                Log.LogMessage(MessageImportance.High, $"Downloading result files for job {JobId}");
+                ExecuteCore().GetAwaiter().GetResult();
+            }
+
+            return !Log.HasLoggedErrors;
+        }
+
+        private async Task ExecuteCore()
+        {
+            DirectoryInfo directory = Directory.CreateDirectory(Path.Combine(OutputDirectory, JobId));
+            using (FileStream stream = File.Open(Path.Combine(directory.FullName, MetadataFile), FileMode.Create, FileAccess.Write))
+            using (var writer = new StreamWriter(stream))
+            {
+                foreach (ITaskItem metadata in MetadataToWrite)
+                {
+                    await writer.WriteLineAsync(metadata.GetMetadata("Identity"));
+                }
+            }
+
+            ResultsContainer = ResultsContainer.EndsWith("/") ? ResultsContainer : ResultsContainer + "/";
+            await Task.WhenAll(WorkItems.Select(wi => DownloadFilesForWorkItem(wi, directory.FullName, _cancellationSource.Token)));
+            return;
+        }
+
+        private async Task DownloadFilesForWorkItem(ITaskItem workItem, string directoryPath, CancellationToken ct)
+        {
+            ct.ThrowIfCancellationRequested();
+
+            if (workItem.TryGetMetadata("DownloadFilesFromResults", out string files))
+            {
+                string workItemName = workItem.GetMetadata("Identity");
+                string[] filesToDownload = files.Split(';');
+
+                DirectoryInfo destinationDir = Directory.CreateDirectory(Path.Combine(directoryPath, workItemName));
+                foreach (var file in filesToDownload)
+                {
+                    try
+                    {
+                        string destinationFile = Path.Combine(destinationDir.FullName, file);
+                        Log.LogMessage(MessageImportance.Normal, $"Downloading {file} => {destinationFile}...");
+
+                        var uri = new Uri($"{ResultsContainer}{workItemName}/{file}");
+                        CloudBlob blob = string.IsNullOrEmpty(ResultsContainerReadSAS) ? new CloudBlob(uri) : new CloudBlob(uri, new StorageCredentials(ResultsContainerReadSAS));
+                        await blob.DownloadToFileAsync(destinationFile, FileMode.Create);
+                    }
+                    catch (StorageException e)
+                    {
+                        Log.LogWarning($"Failed to download {workItemName}/{file} blob from results container: {e.Message}");
+                    }
+                }
+            };
+            return;
+        }
+
+        private void LogRequiredParameterError(string parameter)
+        {
+            Log.LogError($"Required parameter {parameter} string was null or empty");
+        }
+    }
+}

src/Microsoft.DotNet.Helix/Sdk/SendHelixJob.cs

@@ -68,6 +68,18 @@ public class SendHelixJob : HelixTask
         [Output]
         public string JobCorrelationId { get; set; }
 
+        /// <summary>
+        ///   When the task finishes, the results container uri should be available in case we want to download files.
+        /// </summary>
+        [Output]
+        public string ResultsContainerUri { get; set; }
+
+        /// <summary>
+        ///   If the job is internal, we need to give the DownloadFromResultsContainer task the Write SAS to download files.
+        /// </summary>
+        [Output]
+        public string ResultsContainerReadSAS { get; set; }
+
         /// <summary>
         ///   A collection of commands that will run for each work item before any work item commands.
         ///   Use ';' to separate commands and escape a ';' with ';;'
@@ -128,6 +140,11 @@ public class SendHelixJob : HelixTask
         /// </summary>
         public int MaxRetryCount { get; set; }
 
+        /// <summary>
+        /// Currently, if we're using the download results feature, we need to return the results container back to know where to download results from. Currently, the Helix API is the one that provides this container and we have no way to get it back. If this property is set to true, we will create the container before sending the job and tell the API to use this container.
+        /// </summary>
+        public bool UsingDownloadResultsFeature { get; set; }
+
         private CommandPayload _commandPayload;
 
         protected override async Task ExecuteCore()
@@ -199,6 +216,11 @@ protected override async Task ExecuteCore()
                     }
                 }
 
+                if (UsingDownloadResultsFeature)
+                {
+                    def = def.WithDefaultResultsContainer();
+                }
+
                 // don't send the job if we have errors
                 if (Log.HasLoggedErrors)
                 {
@@ -209,6 +231,8 @@ protected override async Task ExecuteCore()
 
                 ISentJob job = await def.SendAsync(msg => Log.LogMessage(msg));
                 JobCorrelationId = job.CorrelationId;
+                ResultsContainerUri = job.ResultsContainerUri;
+                ResultsContainerReadSAS = job.ResultsContainerReadSAS;
             }
 
             string mcUri = await GetMissionControlResultUri();

src/Microsoft.DotNet.Helix/Sdk/tools/Microsoft.DotNet.Helix.Sdk.MonoQueue.targets

@@ -44,6 +44,10 @@
     <PropertyGroup Condition="$(IsPosixShell)">
       <HelixPreCommands>set -x;$(HelixPreCommands)</HelixPreCommands>
     </PropertyGroup>
+    <PropertyGroup>
+      <_usingDownloadResultsFeature>false</_usingDownloadResultsFeature>
+      <_usingDownloadResultsFeature Condition="'@(HelixWorkItem -> HasMetadata('DownloadFilesFromResults'))' != ''">true</_usingDownloadResultsFeature>
+    </PropertyGroup>
     <SendHelixJob Source="$(HelixSource)"
                   Type="$(HelixType)"
                   Build="$(HelixBuild)"
@@ -57,12 +61,18 @@
                   PostCommands="$(HelixPostCommands)"
                   CorrelationPayloads="@(HelixCorrelationPayload)"
                   WorkItems="@(HelixWorkItem)"
-                  HelixProperties="@(HelixProperties)">
+                  HelixProperties="@(HelixProperties)"
+                  UsingDownloadResultsFeature="$(_usingDownloadResultsFeature)">
       <Output TaskParameter="JobCorrelationId" PropertyName="HelixJobId"/>
+      <Output TaskParameter="ResultsContainerUri" PropertyName="HelixResultsContainer"/>
+      <Output TaskParameter="ResultsContainerReadSAS" PropertyName="HelixResultsContainerReadSAS"/>
     </SendHelixJob>
     <ItemGroup>
       <SentJob Include="$(HelixJobId)">
         <WorkItemCount>@(HelixWorkItem->Count())</WorkItemCount>
+        <HelixTargetQueue>$(HelixTargetQueue)</HelixTargetQueue>
+        <ResultsContainerUri>$(HelixResultsContainer)</ResultsContainerUri>
+        <ResultsContainerReadSAS>$(HelixResultsContainerReadSAS)</ResultsContainerReadSAS>
       </SentJob>
     </ItemGroup>
     <Message Text="Sent Helix Job $(HelixJobId)" Importance="High" />

src/Microsoft.DotNet.Helix/Sdk/tools/Microsoft.DotNet.Helix.Sdk.props

@@ -15,4 +15,5 @@
   <UsingTask TaskName="StartAzurePipelinesTestRun" AssemblyFile="$(MicrosoftDotNetHelixSdkTasksAssembly)"/>
   <UsingTask TaskName="StopAzurePipelinesTestRun" AssemblyFile="$(MicrosoftDotNetHelixSdkTasksAssembly)"/>
   <UsingTask TaskName="CheckAzurePipelinesTestRun" AssemblyFile="$(MicrosoftDotNetHelixSdkTasksAssembly)"/>
+  <UsingTask TaskName="DownloadFromResultsContainer" AssemblyFile="$(MicrosoftDotNetHelixSdkTasksAssembly)"/>
 </Project>

src/Microsoft.DotNet.Helix/Sdk/tools/download-results/DownloadFromResultsContainer.props

@@ -0,0 +1,5 @@
+<Project>
+  <PropertyGroup>
+    <_HelixMultiQueueTargets>$(_HelixMultiQueueTargets);$(MSBuildThisFileDirectory)DownloadFromResultsContainer.targets</_HelixMultiQueueTargets>
+  </PropertyGroup>
+</Project>

src/Microsoft.DotNet.Helix/Sdk/tools/download-results/DownloadFromResultsContainer.targets

@@ -0,0 +1,33 @@
+<Project>
+  <Target Name="DownloadFromResultsContainer"
+          Condition="$(WaitForWorkItemCompletion)"
+          AfterTargets="CoreTest"
+          Inputs="unused"
+          Outputs="%(SentJob.Identity)">
+    <ItemGroup>
+      <_workItemsWithDownloadMetadata Include="@(HelixWorkItem)" Condition="'%(HelixWorkItem.DownloadFilesFromResults)' != ''" />
+
+      <HelixDownloadResultsMetadata Include="HelixQueue=%(SentJob.HelixTargetQueue)" />
+      <HelixDownloadResultsMetadata Condition="'$(HelixConfiguration)' != ''" Include="Configuration=$(HelixConfiguration)" />
+      <HelixDownloadResultsMetadata Condition="'$(HelixArchitecture)' != ''" Include="Architecture=$(HelixArchitecture)" />
+    </ItemGroup>
+
+    <PropertyGroup>
+      <HelixResultsDestinationDir Condition="'$(HelixResultsDestinationDir)' == '' AND '$(BUILD_SOURCESDIRECTORY)' != ''">$([MSBuild]::NormalizePath('$(BUILD_SOURCESDIRECTORY)', 'artifacts', helixresults'))</HelixResultsDestinationDir>
+
+      <_shouldDownloadResults>false</_shouldDownloadResults>
+      <_shouldDownloadResults Condition="'@(_workItemsWithDownloadMetadata)' != '' AND '$(HelixResultsDestinationDir)' != ''">true</_shouldDownloadResults>
+    </PropertyGroup>
+
+    <Warning Text="DownloadFromResultsContainer will be skipped for job %(SentJob.Identity) becaue results container uri is empty" Condition="'%(SentJob.ResultsContainerUri)' == '' AND $(_shouldDownloadResults)" />
+
+    <DownloadFromResultsContainer
+        Condition="$(_shouldDownloadResults) AND '%(SentJob.ResultsContainerUri)' != ''"
+        WorkItems="@(_workItemsWithDownloadMetadata)"
+        ResultsContainer="%(SentJob.ResultsContainerUri)"
+        OutputDirectory="$(HelixResultsDestinationDir)"
+        MetadataToWrite="@(HelixDownloadResultsMetadata)"
+        JobId="%(SentJob.Identity)"
+        ResultsContainerReadSAS="%(SentJob.ResultsContainerReadSAS)" />
+  </Target>
+</Project>