From 90e9f3e476d1fab1048d78301535d7e94f1d39eb Mon Sep 17 00:00:00 2001
From: Erik Mavrinac <erikmav@microsoft.com>
Date: Wed, 10 Jan 2024 09:28:09 -0800
Subject: [PATCH] Break dependency from CoW to System.Threading.Tasks.Dataflow
 to avoid dependency load problems in a few projects (#525)

---
 src/CopyOnWrite/Copy.cs | 52 ++++++++++++++---------------------------
 1 file changed, 17 insertions(+), 35 deletions(-)

diff --git a/src/CopyOnWrite/Copy.cs b/src/CopyOnWrite/Copy.cs
index e3061f6..3e96956 100644
--- a/src/CopyOnWrite/Copy.cs
+++ b/src/CopyOnWrite/Copy.cs
@@ -11,7 +11,8 @@
 using System.IO;
 using System.Runtime.InteropServices;
 using System.Threading;
-using System.Threading.Tasks.Dataflow;
+using System.Threading.Tasks;
+using Task = Microsoft.Build.Utilities.Task;
 
 #nullable enable annotations
 
@@ -554,19 +555,18 @@ private void MakeFileWriteable(FileState file, bool logActivity)
 
             // Lockless flags updated from each thread - each needs to be a processor word for atomicity.
             var successFlags = new IntPtr[DestinationFiles.Length];
-            var actionBlockOptions = new ExecutionDataflowBlockOptions
-            {
-                MaxDegreeOfParallelism = parallelism,
-                CancellationToken = _cancellationTokenSource.Token,
-                EnsureOrdered = parallelism == 1,
-            };
-            var partitionCopyActionBlock = new ActionBlock<List<int>>(
-                async (List<int> partition) =>
-                {
-                    // Break from synchronous thread context of caller to get onto thread pool thread.
-                    await System.Threading.Tasks.Task.Yield();
 
-                    for (int partitionIndex = 0; partitionIndex < partition.Count && !_cancellationTokenSource.IsCancellationRequested; partitionIndex++)
+            Parallel.ForEach(partitionsByDestination.Values,
+                new ParallelOptions
+                {
+                    CancellationToken = _cancellationTokenSource.Token,
+                    MaxDegreeOfParallelism = parallelism,
+                },
+                (List<int> partition) =>
+                {
+                    for (int partitionIndex = 0;
+                         partitionIndex < partition.Count && !_cancellationTokenSource.IsCancellationRequested;
+                         partitionIndex++)
                     {
                         int fileIndex = partition[partitionIndex];
                         ITaskItem sourceItem = SourceFiles[fileIndex];
@@ -583,9 +583,9 @@ private void MakeFileWriteable(FileState file, bool logActivity)
                         if (!copyComplete)
                         {
                             if (DoCopyIfNecessary(
-                                new FileState(sourceItem.ItemSpec),
-                                new FileState(destItem.ItemSpec),
-                                copyFile))
+                                    new FileState(sourceItem.ItemSpec),
+                                    new FileState(destItem.ItemSpec),
+                                    copyFile))
                             {
                                 copyComplete = true;
                             }
@@ -602,25 +602,7 @@ private void MakeFileWriteable(FileState file, bool logActivity)
                             successFlags[fileIndex] = (IntPtr)1;
                         }
                     }
-                },
-                actionBlockOptions);
-
-            foreach (List<int> partition in partitionsByDestination.Values)
-            {
-                bool partitionAccepted = partitionCopyActionBlock.Post(partition);
-                if (_cancellationTokenSource.IsCancellationRequested)
-                {
-                    break;
-                }
-                else if (!partitionAccepted)
-                {
-                    // Retail assert...
-                    LogError("Failed posting a file copy to an ActionBlock. Should not happen with block at max int capacity.");
-                }
-            }
-
-            partitionCopyActionBlock.Complete();
-            partitionCopyActionBlock.Completion.GetAwaiter().GetResult();
+                });
 
             // Assemble an in-order list of destination items that succeeded.
             destinationFilesSuccessfullyCopied = new List<ITaskItem>(DestinationFiles.Length);