From c70816d5062f8f66663a2a9cff7ee9ecd1cc2a3c Mon Sep 17 00:00:00 2001
From: jbe2277 <jbe2277@users.noreply.github.com>
Date: Thu, 16 Jan 2025 21:40:12 +0100
Subject: [PATCH] WAF: Add TaskHelper.Run overloads with CancellationToken
 support

---
 .../Foundation/TaskHelperTest.cs              | 19 ++++++++++++++
 .../System.Waf.Core/Foundation/TaskHelper.cs  | 25 +++++++++++++++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/System.Waf/System.Waf/System.Waf.Core.Test/Foundation/TaskHelperTest.cs b/src/System.Waf/System.Waf/System.Waf.Core.Test/Foundation/TaskHelperTest.cs
index 45942193..00e38657 100644
--- a/src/System.Waf/System.Waf/System.Waf.Core.Test/Foundation/TaskHelperTest.cs
+++ b/src/System.Waf/System.Waf/System.Waf.Core.Test/Foundation/TaskHelperTest.cs
@@ -1,7 +1,9 @@
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Waf.Foundation;
+using System.Waf.UnitTesting;
 
 namespace Test.Waf.Foundation
 {
@@ -16,6 +18,23 @@ public async Task RunTest()
             Assert.AreEqual(42, result);
         }
 
+        [TestMethod]
+        public void RunWithCancellationTest()
+        {
+            using var context = UnitTestSynchronizationContext.Create();
+
+            var task = TaskHelper.Run(() => 42, TaskScheduler.FromCurrentSynchronizationContext());
+            var result = task.GetResult(context);
+            Assert.AreEqual(42, result);
+
+            bool called = false;
+            var cts = new CancellationTokenSource();
+            var task2 = TaskHelper.Run(() => called = true, TaskScheduler.FromCurrentSynchronizationContext(), cts.Token);
+            cts.Cancel();
+            AssertHelper.ExpectedException<TaskCanceledException>(() => task2.Wait(context));
+            Assert.IsFalse(called);
+        }
+
         [TestMethod]
         public async Task NoWaitTest()
         {
diff --git a/src/System.Waf/System.Waf/System.Waf.Core/Foundation/TaskHelper.cs b/src/System.Waf/System.Waf/System.Waf.Core/Foundation/TaskHelper.cs
index 67ac76ad..8e120c2c 100644
--- a/src/System.Waf/System.Waf/System.Waf.Core/Foundation/TaskHelper.cs
+++ b/src/System.Waf/System.Waf/System.Waf.Core/Foundation/TaskHelper.cs
@@ -12,7 +12,7 @@ public static class TaskHelper
         /// <returns>A task that represents the work queued to execute in the ThreadPool.</returns>
         public static Task Run(Action action, TaskScheduler scheduler)
         {
-            return Task.Factory.StartNew(action, CancellationToken.None, TaskCreationOptions.DenyChildAttach, scheduler);
+            return Run(action, scheduler, CancellationToken.None);
         }
 
         /// <summary>Queues the specified work to run on the thread pool and returns a proxy for the task returned by function.</summary>
@@ -22,7 +22,28 @@ public static Task Run(Action action, TaskScheduler scheduler)
         /// <returns>A task that represents the work queued to execute in the ThreadPool.</returns>
         public static Task<T> Run<T>(Func<T> action, TaskScheduler scheduler)
         {
-            return Task<T>.Factory.StartNew(action, CancellationToken.None, TaskCreationOptions.DenyChildAttach, scheduler);
+            return Run(action, scheduler, CancellationToken.None);
+        }
+
+        /// <summary>Queues the specified work to run on the thread pool and returns a Task object that represents that work.</summary>
+        /// <param name="action">The work to execute asynchronously</param>
+        /// <param name="scheduler">The TaskScheduler that is used to schedule the created Task.</param>
+        /// <param name="cancellation">The cancellation token that will be assigned to the new task.</param>
+        /// <returns>A task that represents the work queued to execute in the ThreadPool.</returns>
+        public static Task Run(Action action, TaskScheduler scheduler, CancellationToken cancellation)
+        {
+            return Task.Factory.StartNew(action, cancellation, TaskCreationOptions.DenyChildAttach, scheduler);
+        }
+
+        /// <summary>Queues the specified work to run on the thread pool and returns a proxy for the task returned by function.</summary>
+        /// <typeparam name="T">The type of the result produced by this task.</typeparam>
+        /// <param name="action">The work to execute asynchronously</param>
+        /// <param name="scheduler">The TaskScheduler that is used to schedule the created Task.</param>
+        /// <param name="cancellation">The cancellation token that will be assigned to the new task.</param>
+        /// <returns>A task that represents the work queued to execute in the ThreadPool.</returns>
+        public static Task<T> Run<T>(Func<T> action, TaskScheduler scheduler, CancellationToken cancellation)
+        {
+            return Task<T>.Factory.StartNew(action, cancellation, TaskCreationOptions.DenyChildAttach, scheduler);
         }
 
         /// <summary>