Add support for CancellationToken to TaskSeq

fsprojects · Dec 17, 2022 · 3a9a47c · 3a9a47c
1 parent fdcc270
commit 3a9a47c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs b/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs
@@ -657,6 +657,12 @@ module HighPriority =
                     sm.Data.current <- ValueNone
                     false)
 
+        // Binding to a cancellation token. This allows `do! someCancellationToken`
+        member inline _.Bind(myToken: CancellationToken, continuation: (unit -> ResumableTSC<'T>)) : ResumableTSC<'T> =
+            ResumableTSC<'T>(fun sm ->
+                sm.Data.cancellationToken <- myToken
+                (continuation ()).Invoke(&sm))
+
 [<AutoOpen>]
 module TaskSeqBuilder =
     /// Builds an asynchronous task sequence based on IAsyncEnumerable<'T> using computation expression syntax.

diff --git a/src/FSharp.Control.TaskSeq/TaskSeqInternal.fs b/src/FSharp.Control.TaskSeq/TaskSeqInternal.fs
@@ -60,6 +60,11 @@ module internal TaskSeqInternal =
         KeyNotFoundException("The predicate function or index did not satisfy any item in the async sequence.")
         |> raise
 
+    let inline withCancellationToken (cancellationToken: CancellationToken) (source: taskSeq<'T>) = taskSeq {
+        do! cancellationToken
+        yield! source
+    }
+
     let isEmpty (source: taskSeq<_>) = task {
         use e = source.GetAsyncEnumerator(CancellationToken())
         let! step = e.MoveNextAsync()