From cda24d35a626ec814095a8c010d943f63f096835 Mon Sep 17 00:00:00 2001
From: Scott Beca <scottbeca@mightygamesgroup.com>
Date: Mon, 6 Nov 2023 21:35:53 +1100
Subject: [PATCH] Implement some more extension functions

---
 api_list.include.md                           |  12 ++
 readme.md                                     |  12 ++
 src/Consume/Consume.cs                        |  14 ++
 .../PolyfillExtensions_CancellationToken.cs   | 137 ++++++++++++++++++
 ...yfillExtensions_CancellationTokenSource.cs |  41 +++++-
 src/Polyfill/PolyfillExtensions_Process.cs    | 117 +++++++++++++++
 src/Tests/BuildApiTest.cs                     |   5 +-
 ...lyfillExtensionsTests_CancellationToken.cs |  69 +++++++++
 ...ExtensionsTests_CancellationTokenSource.cs |  69 +++++++++
 src/Tests/PolyfillExtensionsTests_Process.cs  |  27 ++++
 10 files changed, 499 insertions(+), 4 deletions(-)
 create mode 100644 src/Polyfill/PolyfillExtensions_CancellationToken.cs
 create mode 100644 src/Polyfill/PolyfillExtensions_Process.cs
 create mode 100644 src/Tests/PolyfillExtensionsTests_CancellationToken.cs
 create mode 100644 src/Tests/PolyfillExtensionsTests_CancellationTokenSource.cs
 create mode 100644 src/Tests/PolyfillExtensionsTests_Process.cs

diff --git a/api_list.include.md b/api_list.include.md
index 02f7218c..c44d606e 100644
--- a/api_list.include.md
+++ b/api_list.include.md
@@ -59,6 +59,11 @@
  * `Boolean TryFormat(Span<Char>, Int32&, ReadOnlySpan<Char>, IFormatProvider)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.decimal.tryformat)
 
 
+### Process
+
+ * `Task WaitForExitAsync(CancellationToken)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.diagnostics.process.waitforexitasync)
+
+
 ### Double
 
  * `Boolean TryFormat(Span<Char>, Int32&, ReadOnlySpan<Char>, IFormatProvider)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.double.tryformat)
@@ -208,6 +213,13 @@
  * `Boolean Equals(ReadOnlySpan<Char>)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.text.stringbuilder.equals#system-text-stringbuilder-equals(system-readonlyspan((system-char))))
 
 
+### CancellationToken
+
+ * `CancellationTokenRegistration Register(Action<Object,CancellationToken>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.register#system-threading-cancellationtoken-register(system-action((system-object-system-threading-cancellationtoken))-system-object))
+ * `CancellationTokenRegistration UnsafeRegister(Action<Object>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object))-system-object))
+ * `CancellationTokenRegistration UnsafeRegister(Action<Object,CancellationToken>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object-system-threading-cancellationtoken))-system-object))
+
+
 ### CancellationTokenSource
 
  * `Task CancelAsync()` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtokensource.cancelasync)
diff --git a/readme.md b/readme.md
index f9827642..cb91a334 100644
--- a/readme.md
+++ b/readme.md
@@ -414,6 +414,11 @@ The class `PolyfillExtensions` includes the following extension methods:
  * `Boolean TryFormat(Span<Char>, Int32&, ReadOnlySpan<Char>, IFormatProvider)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.decimal.tryformat)
 
 
+### Process
+
+ * `Task WaitForExitAsync(CancellationToken)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.diagnostics.process.waitforexitasync)
+
+
 ### Double
 
  * `Boolean TryFormat(Span<Char>, Int32&, ReadOnlySpan<Char>, IFormatProvider)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.double.tryformat)
@@ -563,6 +568,13 @@ The class `PolyfillExtensions` includes the following extension methods:
  * `Boolean Equals(ReadOnlySpan<Char>)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.text.stringbuilder.equals#system-text-stringbuilder-equals(system-readonlyspan((system-char))))
 
 
+### CancellationToken
+
+ * `CancellationTokenRegistration Register(Action<Object,CancellationToken>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.register#system-threading-cancellationtoken-register(system-action((system-object-system-threading-cancellationtoken))-system-object))
+ * `CancellationTokenRegistration UnsafeRegister(Action<Object>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object))-system-object))
+ * `CancellationTokenRegistration UnsafeRegister(Action<Object,CancellationToken>, Object)` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object-system-threading-cancellationtoken))-system-object))
+
+
 ### CancellationTokenSource
 
  * `Task CancelAsync()` [reference](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtokensource.cancelasync)
diff --git a/src/Consume/Consume.cs b/src/Consume/Consume.cs
index d4162431..97c4923e 100644
--- a/src/Consume/Consume.cs
+++ b/src/Consume/Consume.cs
@@ -138,6 +138,20 @@ async Task CancellationTokenSource()
         await source.CancelAsync();
     }
 
+    void CancellationTokenUnsafeRegister()
+    {
+        var source = new CancellationTokenSource();
+        var token = source.Token;
+        token.UnsafeRegister((state) => {}, null);
+        token.UnsafeRegister((state, token) => {}, null);
+    }
+
+    async Task ProcessWaitForExitAsync()
+    {
+        var process = new Process();
+        await process.WaitForExitAsync();
+    }
+
     async Task StreamReaderReadToEndAsync()
     {
         var reader = new StreamReader(new MemoryStream());
diff --git a/src/Polyfill/PolyfillExtensions_CancellationToken.cs b/src/Polyfill/PolyfillExtensions_CancellationToken.cs
new file mode 100644
index 00000000..bff06e5a
--- /dev/null
+++ b/src/Polyfill/PolyfillExtensions_CancellationToken.cs
@@ -0,0 +1,137 @@
+// <auto-generated />
+
+#nullable enable
+
+#pragma warning disable
+
+using System;
+using System.Threading;
+using Link = System.ComponentModel.DescriptionAttribute;
+
+static partial class PolyfillExtensions
+{
+
+#if !NETCOREAPP3_0_OR_GREATER
+
+    /// <summary>
+    /// Registers a delegate that will be called when this
+    /// <see cref="CancellationToken">CancellationToken</see> is canceled.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// If this token is already in the canceled state, the delegate will be run immediately and synchronously.
+    /// Any exception the delegate generates will be propagated out of this method call.
+    /// </para>
+    /// <para>
+    /// <see cref="ExecutionContext">ExecutionContext</see> is not captured nor flowed
+    /// to the callback's invocation.
+    /// </para>
+    /// </remarks>
+    /// <param name="callback">The delegate to be executed when the <see cref="CancellationToken">CancellationToken</see> is canceled.</param>
+    /// <param name="state">The state to pass to the <paramref name="callback"/> when the delegate is invoked.  This may be null.</param>
+    /// <returns>The <see cref="CancellationTokenRegistration"/> instance that can
+    /// be used to unregister the callback.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="callback"/> is null.</exception>
+    [Link("https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object))-system-object)")]
+    public static CancellationTokenRegistration UnsafeRegister(this CancellationToken target, Action<object?> callback, object? state)
+    {
+        if (callback is null)
+        {
+            throw new ArgumentNullException(nameof(callback));
+        }
+
+        // The main difference between UnsafeRegister and Register appears to be that UnsafeRegister callbacks don't capture and use the execution context.
+        // So to emulate that here, let's suppress the execution context if needed before calling Register.
+        // This idea was taken from UniTask and how they implemented their RegisterWithoutCaptureExecutionContext extension methods:
+        // https://github.com/Cysharp/UniTask/blob/master/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs
+
+        var restoreFlow = false;
+        if (!ExecutionContext.IsFlowSuppressed())
+        {
+            ExecutionContext.SuppressFlow();
+            restoreFlow = true;
+        }
+
+        try
+        {
+            return target.Register(callback, state, false);
+        }
+        finally
+        {
+            if (restoreFlow)
+            {
+                ExecutionContext.RestoreFlow();
+            }
+        }
+    }
+
+#endif
+
+#if !NET6_0_OR_GREATER
+
+    /// <summary>Registers a delegate that will be called when this <see cref="CancellationToken">CancellationToken</see> is canceled.</summary>
+    /// <remarks>
+    /// If this token is already in the canceled state, the delegate will be run immediately and synchronously. Any exception the delegate
+    /// generates will be propagated out of this method call. The current <see cref="ExecutionContext">ExecutionContext</see>, if one exists,
+    /// will be captured along with the delegate and will be used when executing it. The current <see cref="SynchronizationContext"/> is not captured.
+    /// </remarks>
+    /// <param name="callback">The delegate to be executed when the <see cref="CancellationToken">CancellationToken</see> is canceled.</param>
+    /// <param name="state">The state to pass to the <paramref name="callback"/> when the delegate is invoked.  This may be null.</param>
+    /// <returns>The <see cref="CancellationTokenRegistration"/> instance that can be used to unregister the callback.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="callback"/> is null.</exception>
+    [Link("https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.register#system-threading-cancellationtoken-register(system-action((system-object-system-threading-cancellationtoken))-system-object)")]
+    public static CancellationTokenRegistration Register(this CancellationToken target, Action<object?, CancellationToken> callback, object? state)
+    {
+        if (callback is null)
+        {
+            throw new ArgumentNullException(nameof(callback));
+        }
+
+        return target.Register((data) => callback(data, target), state, useSynchronizationContext: false);
+    }
+
+    /// <summary>Registers a delegate that will be called when this <see cref="CancellationToken">CancellationToken</see> is canceled.</summary>
+    /// <remarks>
+    /// If this token is already in the canceled state, the delegate will be run immediately and synchronously. Any exception the delegate
+    /// generates will be propagated out of this method call. <see cref="ExecutionContext"/> is not captured nor flowed to the callback's invocation.
+    /// </remarks>
+    /// <param name="callback">The delegate to be executed when the <see cref="CancellationToken">CancellationToken</see> is canceled.</param>
+    /// <param name="state">The state to pass to the <paramref name="callback"/> when the delegate is invoked.  This may be null.</param>
+    /// <returns>The <see cref="CancellationTokenRegistration"/> instance that can be used to unregister the callback.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="callback"/> is null.</exception>
+    [Link("https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken.unsaferegister#system-threading-cancellationtoken-unsaferegister(system-action((system-object-system-threading-cancellationtoken))-system-object)")]
+    public static CancellationTokenRegistration UnsafeRegister(this CancellationToken target, Action<object?, CancellationToken> callback, object? state) 
+    {
+        if (callback is null)
+        {
+            throw new ArgumentNullException(nameof(callback));
+        }
+
+        // The main difference between UnsafeRegister and Register appears to be that UnsafeRegister callbacks don't capture and use the execution context.
+        // So to emulate that here, let's suppress the execution context if needed before calling Register.
+        // This idea was taken from UniTask and how they implemented their RegisterWithoutCaptureExecutionContext extension methods:
+        // https://github.com/Cysharp/UniTask/blob/master/src/UniTask/Assets/Plugins/UniTask/Runtime/CancellationTokenExtensions.cs
+
+        var restoreFlow = false;
+        if (!ExecutionContext.IsFlowSuppressed())
+        {
+            ExecutionContext.SuppressFlow();
+            restoreFlow = true;
+        }
+
+        try
+        {
+            Action<object> internalCallback = (data) => callback(data, target);
+            return target.Register(internalCallback, state, false);
+        }
+        finally
+        {
+            if (restoreFlow)
+            {
+                ExecutionContext.RestoreFlow();
+            }
+        }
+    }
+
+#endif
+}
diff --git a/src/Polyfill/PolyfillExtensions_CancellationTokenSource.cs b/src/Polyfill/PolyfillExtensions_CancellationTokenSource.cs
index ca9283f5..121c0728 100644
--- a/src/Polyfill/PolyfillExtensions_CancellationTokenSource.cs
+++ b/src/Polyfill/PolyfillExtensions_CancellationTokenSource.cs
@@ -10,11 +10,48 @@ static partial class PolyfillExtensions
 {
 #if !NET8_0_OR_GREATER
 
+    /// <summary>Communicates a request for cancellation asynchronously.</summary>
+    /// <remarks>
+    /// <para>
+    /// The associated <see cref="CancellationToken" /> will be notified of the cancellation
+    /// and will synchronously transition to a state where <see cref="CancellationToken.IsCancellationRequested"/> returns true.
+    /// Any callbacks or cancelable operations registered with the <see cref="CancellationToken"/>  will be executed asynchronously,
+    /// with the returned <see cref="Task"/> representing their eventual completion.
+    /// </para>
+    /// <para>
+    /// Callbacks registered with the token should not throw exceptions.
+    /// However, any such exceptions that are thrown will be aggregated into an <see cref="AggregateException"/>,
+    /// such that one callback throwing an exception will not prevent other registered callbacks from being executed.
+    /// </para>
+    /// <para>
+    /// The <see cref="ExecutionContext"/> that was captured when each callback was registered
+    /// will be reestablished when the callback is invoked.
+    /// </para>
+    /// </remarks>
+    /// <exception cref="ObjectDisposedException">This <see cref="CancellationTokenSource"/> has been disposed.</exception>
     [Link("https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtokensource.cancelasync")]
     public static Task CancelAsync(this CancellationTokenSource target)
     {
-         target.Cancel();
-         return Task.CompletedTask;
+        if (target.IsCancellationRequested)
+        {
+            // If cancellation has already been requested then we can just return immediately
+            return Task.CompletedTask;
+        }
+        else
+        {
+            // Run sync Cancel call in Task to avoid possible deadlock.
+            // As an example, the CancellationTokenSource_CancelAsync_CallbacksInvokedAsynchronously test
+            // will hit a deadlock if we try to just call Cancel directly without it being run in a task
+            Task task = Task.Run(() => target.Cancel());
+
+            while (!target.IsCancellationRequested)
+            {
+                // Don't return until we know that the cancellation request has started, to match the
+                // state that the real implemenation for CancelAsync would be in after being called
+            }
+
+            return task;
+        }
     }
 
 #endif
diff --git a/src/Polyfill/PolyfillExtensions_Process.cs b/src/Polyfill/PolyfillExtensions_Process.cs
new file mode 100644
index 00000000..11031172
--- /dev/null
+++ b/src/Polyfill/PolyfillExtensions_Process.cs
@@ -0,0 +1,117 @@
+// <auto-generated />
+
+#pragma warning disable
+
+#if !NET5_0_OR_GREATER
+
+using System;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Link = System.ComponentModel.DescriptionAttribute;
+
+static partial class PolyfillExtensions
+{
+    /// <summary>
+    /// Instructs the Process component to wait for the associated process to exit, or
+    /// for the <paramref name="cancellationToken"/> to be canceled.
+    /// </summary>
+    /// <remarks>
+    /// Calling this method will set <see cref="EnableRaisingEvents"/> to <see langword="true" />.
+    /// </remarks>
+    /// <returns>
+    /// A task that will complete when the process has exited, cancellation has been requested,
+    /// or an error occurs.
+    /// </returns>
+    [Link("https://learn.microsoft.com/en-us/dotnet/api/system.diagnostics.process.waitforexitasync")]
+    public static async Task WaitForExitAsync(this Process target, CancellationToken cancellationToken = default)
+    {
+        // Because the process has already started by the time this method is called,
+        // we're in a race against the process to set up our exit handlers before the process
+        // exits. As a result, there are several different flows that must be handled:
+        //
+        // CASE 1: WE ENABLE EVENTS
+        // This is the "happy path". In this case we enable events.
+        //
+        // CASE 1.1: PROCESS EXITS OR IS CANCELED AFTER REGISTERING HANDLER
+        // This case continues the "happy path". The process exits or waiting is canceled after
+        // registering the handler and no special cases are needed.
+        //
+        // CASE 1.2: PROCESS EXITS BEFORE REGISTERING HANDLER
+        // It's possible that the process can exit after we enable events but before we register
+        // the handler. In that case we must check for exit after registering the handler.
+        //
+        //
+        // CASE 2: PROCESS EXITS BEFORE ENABLING EVENTS
+        // The process may exit before we attempt to enable events. In that case EnableRaisingEvents
+        // will throw an exception like this:
+        //     System.InvalidOperationException : Cannot process request because the process (42) has exited.
+        // In this case we catch the InvalidOperationException. If the process has exited, our work
+        // is done and we return. If for any reason (now or in the future) enabling events fails
+        // and the process has not exited, bubble the exception up to the user.
+        //
+        //
+        // CASE 3: USER ALREADY ENABLED EVENTS
+        // In this case the user has already enabled raising events. Re-enabling events is a no-op
+        // as the value hasn't changed. However, no-op also means that if the process has already
+        // exited, EnableRaisingEvents won't throw an exception.
+        //
+        // CASE 3.1: PROCESS EXITS OR IS CANCELED AFTER REGISTERING HANDLER
+        // (See CASE 1.1)
+        //
+        // CASE 3.2: PROCESS EXITS BEFORE REGISTERING HANDLER
+        // (See CASE 1.2)
+
+        if (!target.HasExited)
+        {
+            // Early out for cancellation before doing more expensive work
+            cancellationToken.ThrowIfCancellationRequested();
+        }
+
+        try
+        {
+            // CASE 1: We enable events
+            // CASE 2: Process exits before enabling events (and throws an exception)
+            // CASE 3: User already enabled events (no-op)
+            target.EnableRaisingEvents = true;
+        }
+        catch (InvalidOperationException)
+        {
+            // CASE 2: If the process has exited, our work is done, otherwise bubble the
+            // exception up to the user
+            if (target.HasExited)
+            {
+                return;
+            }
+
+            throw;
+        }
+
+        var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        EventHandler handler = (_, _) => tcs.TrySetResult(null);
+        target.Exited += handler;
+
+        try
+        {
+            if (target.HasExited)
+            {
+                // CASE 1.2 & CASE 3.2: Handle race where the process exits before registering the handler
+            }
+            else
+            {
+                // CASE 1.1 & CASE 3.1: Process exits or is canceled here
+                using (cancellationToken.UnsafeRegister(static (s, cancellationToken) => ((TaskCompletionSource<object>)s!).TrySetCanceled(cancellationToken), tcs))
+                {
+                    await tcs.Task.ConfigureAwait(false);
+                }
+            }
+        }
+        finally
+        {
+            target.Exited -= handler;
+        }
+    }
+}
+
+#endif
diff --git a/src/Tests/BuildApiTest.cs b/src/Tests/BuildApiTest.cs
index c0c92d22..43b36f85 100644
--- a/src/Tests/BuildApiTest.cs
+++ b/src/Tests/BuildApiTest.cs
@@ -4,6 +4,7 @@ class BuildApiTest
 {
     static string[] namespacesToClean =
     {
+        "System.Diagnostics.",
         "System.Collections.Generic.",
         "System.Threading.Tasks.",
         "System.Threading.",
@@ -17,8 +18,8 @@ class BuildApiTest
     public void Run()
     {
         var solutionDirectory = SolutionDirectoryFinder.Find();
-        var path = Path.Combine(solutionDirectory, @"Consume\bin\Debug\netstandard2.0\Consume.dll");
-        var md = Path.Combine(solutionDirectory, @"..\api_list.include.md");
+        var path = Path.Combine(solutionDirectory, "Consume", "bin", "Debug", "netstandard2.0", "Consume.dll");
+        var md = Path.Combine(solutionDirectory, "..", "api_list.include.md");
         File.Delete(md);
         using var module = Mono.Cecil.ModuleDefinition.ReadModule(path);
         var extensions = module.GetTypes().Single(_ => _.Name == nameof(PolyfillExtensions));
diff --git a/src/Tests/PolyfillExtensionsTests_CancellationToken.cs b/src/Tests/PolyfillExtensionsTests_CancellationToken.cs
new file mode 100644
index 00000000..d0fc66de
--- /dev/null
+++ b/src/Tests/PolyfillExtensionsTests_CancellationToken.cs
@@ -0,0 +1,69 @@
+partial class PolyfillExtensionsTests
+{
+    [Test]
+    public void CancellationToken_Register_Exceptions()
+    {
+        CancellationToken token = default;
+
+#nullable disable
+        Assert.Throws<ArgumentNullException>(() => token.Register((Action<object, CancellationToken>)null, null));
+
+        Assert.Throws<ArgumentNullException>(() => token.UnsafeRegister((Action<object>)null, null));
+        Assert.Throws<ArgumentNullException>(() => token.UnsafeRegister((Action<object, CancellationToken>)null, null));
+#nullable enable
+    }
+
+    [Test]
+    [TestCase(false)]
+    [TestCase(true)]
+    public static void CancellationToken_Register_ExecutionContextFlowsIfExpected(bool callbackWithToken)
+    {
+        var cts = new CancellationTokenSource();
+
+        const int Iters = 5;
+        int invoked = 0;
+
+        AsyncLocal<int> al = new AsyncLocal<int>();
+        for (int i = 1; i <= Iters; i++)
+        {
+            bool flowExecutionContext = i % 2 == 0;
+
+            al.Value = i;
+            Action<object?> callback = s =>
+            {
+                invoked++;
+                Assert.AreEqual(flowExecutionContext ? (int)s! : 0, al.Value);
+            };
+
+            CancellationToken ct = cts.Token;
+            if (flowExecutionContext && callbackWithToken)
+            {
+                ct.Register((s, t) =>
+                {
+                    Assert.AreEqual(ct, t);
+                    callback(s);
+                }, i);
+            }
+            else if (flowExecutionContext)
+            {
+                ct.Register(callback, i);
+            }
+            else if (callbackWithToken)
+            {
+                ct.UnsafeRegister((s, t) =>
+                {
+                    Assert.AreEqual(ct, t);
+                    callback(s);
+                }, i);
+            }
+            else
+            {
+                ct.UnsafeRegister(callback, i);
+            }
+        }
+        al.Value = 0;
+
+        cts.Cancel();
+        Assert.AreEqual(Iters, invoked);
+    }
+}
diff --git a/src/Tests/PolyfillExtensionsTests_CancellationTokenSource.cs b/src/Tests/PolyfillExtensionsTests_CancellationTokenSource.cs
new file mode 100644
index 00000000..ffc04e73
--- /dev/null
+++ b/src/Tests/PolyfillExtensionsTests_CancellationTokenSource.cs
@@ -0,0 +1,69 @@
+partial class PolyfillExtensionsTests
+{
+    private static bool IsCompletedSuccessfully(Task task)
+    {
+#if NETFRAMEWORK || NETSTANDARD
+        return task.Status == TaskStatus.RanToCompletion;
+#else
+        return task.IsCompletedSuccessfully;
+#endif
+    }
+
+    [Test]
+    [Ignore("This test is taken directly from the .NET repo but we can't match the real CancelAsync logic exactly, so differ slightly and can't pass this test")]
+    public static void CancellationTokenSource_CancelAsync_NoRegistrations_CallbackCompletesImmediately()
+    {
+        var cts = new CancellationTokenSource();
+        Assert.True(IsCompletedSuccessfully(cts.CancelAsync()));
+        Assert.True(cts.IsCancellationRequested);
+
+        cts = new CancellationTokenSource();
+        cts.Token.Register(() => { }).Dispose();
+        Assert.True(IsCompletedSuccessfully(cts.CancelAsync()));
+        Assert.True(cts.IsCancellationRequested);
+    }
+
+    [Test]
+    public static async Task CancellationTokenSource_CancelAsync_CallbacksInvokedAsynchronously()
+    {
+        var cts = new CancellationTokenSource();
+
+        var mres = new ManualResetEventSlim();
+        cts.Token.Register(mres.Wait);
+
+        Task t = cts.CancelAsync();
+        Assert.False(t.IsCompleted);
+        Assert.True(cts.IsCancellationRequested);
+
+        Assert.True(IsCompletedSuccessfully(cts.CancelAsync())); // secondary call completes immediately
+
+        mres.Set();
+        await t;
+    }
+
+    [Test]
+    public static void CancellationTokenSource_CancelAsync_AllCallbacksInvoked()
+    {
+        const int Iters = 1000;
+
+        int sum = 0;
+        int callingThreadId = Environment.CurrentManagedThreadId;
+
+        var cts = new CancellationTokenSource();
+        for (int i = 1; i <= Iters; i++)
+        {
+            cts.Token.Register(s =>
+            {
+                sum += (int)s!;
+            }, i);
+        }
+
+        Task t = cts.CancelAsync();
+        Assert.True(cts.IsCancellationRequested);
+
+        ((IAsyncResult)t).AsyncWaitHandle.WaitOne(); // synchronously block without inlining to ensure this thread isn't reused
+        t.Wait(); // propagate any exceptions
+
+        Assert.AreEqual(Iters * (Iters + 1) / 2, sum);
+    }
+}
diff --git a/src/Tests/PolyfillExtensionsTests_Process.cs b/src/Tests/PolyfillExtensionsTests_Process.cs
new file mode 100644
index 00000000..9410fd6f
--- /dev/null
+++ b/src/Tests/PolyfillExtensionsTests_Process.cs
@@ -0,0 +1,27 @@
+partial class PolyfillExtensionsTests
+{
+    [Test, RequiresThread]
+    [TestCase(0)]  // poll
+    [TestCase(10)] // real timeout
+    public void Process_CurrentProcess_WaitAsyncNeverCompletes(int milliseconds)
+    {
+        using (var cts = new CancellationTokenSource(milliseconds))
+        {
+            CancellationToken token = cts.Token;
+            Process process = Process.GetCurrentProcess();
+
+            OperationCanceledException? ex = Assert.CatchAsync(async () => await process.WaitForExitAsync(token)) as OperationCanceledException;
+
+            Assert.IsNotNull(ex);
+            Assert.AreEqual(token, ex!.CancellationToken);
+            Assert.False(process.HasExited);
+        }
+    }
+
+    [Test, RequiresThread]
+    public void Process_WaitForExitAsync_NotDirected_ThrowsInvalidOperationException()
+    {
+        var process = new Process();
+        Assert.ThrowsAsync<InvalidOperationException>(async () => await process.WaitForExitAsync());
+    }
+}