Implement ICancelableTask

dotnet · Feb 27, 2019 · 85c8b65 · 85c8b65
1 parent ae6703b
commit 85c8b65
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/src/Microsoft.DotNet.Helix/Sdk/DownloadFromResultsContainer.cs b/src/Microsoft.DotNet.Helix/Sdk/DownloadFromResultsContainer.cs
@@ -5,11 +5,12 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.DotNet.Helix.Sdk
 {
-    public class DownloadFromResultsContainer : BaseTask
+    public class DownloadFromResultsContainer : BaseTask, ICancelableTask
     {
         [Required]
         public ITaskItem[] WorkItems { get; set; }
@@ -30,6 +31,10 @@ public class DownloadFromResultsContainer : BaseTask
 
         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))
@@ -69,12 +74,14 @@ private async Task ExecuteCore()
             }
 
             ResultsContainer = ResultsContainer.EndsWith("/") ? ResultsContainer : ResultsContainer + "/";
-            await Task.WhenAll(WorkItems.Select(wi => DownloadFilesForWorkItem(wi, directory.FullName)));
+            await Task.WhenAll(WorkItems.Select(wi => DownloadFilesForWorkItem(wi, directory.FullName, _cancellationSource.Token)));
             return;
         }
 
-        private async Task DownloadFilesForWorkItem(ITaskItem workItem, string directoryPath)
+        private async Task DownloadFilesForWorkItem(ITaskItem workItem, string directoryPath, CancellationToken ct)
         {
+            ct.ThrowIfCancellationRequested();
+
             if (workItem.TryGetMetadata("DownloadFilesFromResults", out string files))
             {
                 string workItemName = workItem.GetMetadata("Identity");