Async enumerable cancellation

microsoft · Jun 16, 2024 · df90762 · df90762
1 parent f90d315
commit df90762
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 5 deletions.
diff --git a/src/NodeApi/Interop/JSCollectionExtensions.cs b/src/NodeApi/Interop/JSCollectionExtensions.cs
@@ -199,12 +199,17 @@ internal sealed class JSAsyncIterableEnumerator<T> : IAsyncEnumerator<T>, IDispo
     private readonly JSValue.To<T> _fromJS;
     private readonly JSReference _iteratorReference;
     private JSReference? _currentReference;
+    private CancellationToken _cancellation;
 
-    internal JSAsyncIterableEnumerator(JSValue iterable, JSValue.To<T> fromJS)
+    internal JSAsyncIterableEnumerator(
+        JSValue iterable,
+        JSValue.To<T> fromJS,
+        CancellationToken cancellation)
     {
         _fromJS = fromJS;
         _iteratorReference = new JSReference(iterable.CallMethod(JSSymbol.AsyncIterator));
         _currentReference = null;
+        _cancellation = cancellation;
     }
 
     public async ValueTask<bool> MoveNextAsync()
@@ -214,7 +219,7 @@ public async ValueTask<bool> MoveNextAsync()
             _currentReference?.Dispose();
 
             JSPromise nextPromise = (JSPromise)iterator.CallMethod("next");
-            JSValue nextResult = await nextPromise.AsTask();
+            JSValue nextResult = await nextPromise.AsTask(_cancellation);
             JSValue done = nextResult["done"];
 
             if (done.IsBoolean() && (bool)done)
@@ -334,10 +339,10 @@ internal JSAsyncIterableEnumerable(JSValue iterable, JSValue.To<T> fromJS)
     bool IEquatable<JSValue>.Equals(JSValue other)
         => _iterableReference.Run((iterable) => iterable.Equals(other));
 
-    public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken)
+    public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellation)
     {
         return _iterableReference.Run(
-            (iterable) => new JSAsyncIterableEnumerator<T>(iterable, _fromJS));
+            (iterable) => new JSAsyncIterableEnumerator<T>(iterable, _fromJS, cancellation));
     }
 
     public void Dispose()

diff --git a/src/NodeApi/JSPromiseExtensions.cs b/src/NodeApi/JSPromiseExtensions.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.JavaScript.NodeApi;
@@ -27,12 +28,41 @@ public static Task<JSValue> AsTask(this JSPromise promise)
         return completion.Task;
     }
 
-    public static async Task<T> AsTask<T>(this JSPromise promise, JSValue.To<T> fromJS)
+    public static Task<JSValue> AsTask(this JSPromise promise, CancellationToken cancellation)
+    {
+        TaskCompletionSource<JSValue> completion = new();
+        cancellation.Register(() => completion.TrySetCanceled(cancellation));
+        promise.Then(
+            (JSValue value) =>
+            {
+                completion.TrySetResult(value);
+                return default;
+            },
+            (JSError error) =>
+            {
+                completion.TrySetException(new JSException(error));
+                return default;
+            });
+        return completion.Task;
+    }
+
+    public static async Task<T> AsTask<T>(
+        this JSPromise promise,
+        JSValue.To<T> fromJS)
     {
         Task<JSValue> jsTask = promise.AsTask();
         return fromJS(await jsTask);
     }
 
+    public static async Task<T> AsTask<T>(
+        this JSPromise promise,
+        JSValue.To<T> fromJS,
+        CancellationToken cancellation)
+    {
+        Task<JSValue> jsTask = promise.AsTask(cancellation);
+        return fromJS(await jsTask);
+    }
+
     public static JSPromise AsPromise(this Task task)
     {
         if (task.Status == TaskStatus.RanToCompletion)