From de29125f20c38f4d2147f841dc8bd4b4e6038fe0 Mon Sep 17 00:00:00 2001
From: Travis Jensen <travisje@microsoft.com>
Date: Tue, 9 Jan 2024 17:38:06 -0700
Subject: [PATCH] Feature/forceparallelactions (#57)

* Force ParallelActions by yielding back the task

* Fixing flaky test.
---
 .../test/TreeWalkerUnitTests.cs                  | 16 ++++++++++++----
 Forge.TreeWalker/src/TreeWalkerSession.cs        |  3 +++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/Forge.TreeWalker.UnitTests/test/TreeWalkerUnitTests.cs b/Forge.TreeWalker.UnitTests/test/TreeWalkerUnitTests.cs
index d8f1d0e..1c520ec 100644
--- a/Forge.TreeWalker.UnitTests/test/TreeWalkerUnitTests.cs
+++ b/Forge.TreeWalker.UnitTests/test/TreeWalkerUnitTests.cs
@@ -468,17 +468,25 @@ public void TestTreeWalkerSession_WalkTree_ActionWithDelay_CancelWalkTree()
             // This gives us time to start WalkTree before calling CancelWalkTree.
             this.TestInitialize(jsonSchema: ForgeSchemaHelper.ActionDelay_ContinuationOnTimeout_RetryPolicy_TimeoutInAction);
 
-            // Test - WalkTree then CancelWalkTree while WalkTree is running and expect the Status to be Cancelled.
+            // Test - WalkTree then CancelWalkTree while WalkTree is running and expect the Status to be either:
+            //   1. CancelledBeforeExecution if TaskCanceledException is thrown from WalkTree. This happens when the ForgeAction gets canceled and honored by Forge.
+            //   2. Cancelled if OperationCanceledException is thrown from WalkTree. This happens when some other Task wins the race to get cancelled first, such as nodeTimeoutTask.
             Task<string> task = this.session.WalkTree("Root");
             Thread.Sleep(25);
             this.session.CancelWalkTree();
-            Assert.ThrowsException<OperationCanceledException>(() => 
+
+            try
             {
                 string temp = task.GetAwaiter().GetResult();
-            }, "Expected WalkTree to throw exception after calling CancelWalkTree.");
+                Assert.Fail("Expected WalkTree to throw exception after calling CancelWalkTree.");
+            }
+            catch { }
 
             string actualStatus = this.session.Status;
-            Assert.AreEqual("Cancelled", actualStatus, "Expected WalkTree to be cancelled after calling CancelWalkTree.");
+            if (actualStatus != "Cancelled" && actualStatus != "CancelledBeforeExecution")
+            {
+                Assert.Fail($"Expected WalkTree to be Cancelled or CancelledBeforeExecution after calling CancelWalkTree, but it was {actualStatus}.");
+            }
         }
 
         [TestMethod]
diff --git a/Forge.TreeWalker/src/TreeWalkerSession.cs b/Forge.TreeWalker/src/TreeWalkerSession.cs
index 30325ad..6ffd413 100644
--- a/Forge.TreeWalker/src/TreeWalkerSession.cs
+++ b/Forge.TreeWalker/src/TreeWalkerSession.cs
@@ -714,6 +714,9 @@ internal async Task ExecuteActionWithRetry(
             ActionDefinition actionDefinition,
             CancellationToken token)
         {
+            // Ensure ForgeActions run in parallel by yielding back the task.
+            await Task.Yield();
+
             // Initialize values. Default infinite timeout. Default RetryPolicyType.None.
             int retryCount = 0;
             Exception innerException = null;