Add support for CancellationToken to TaskSeq

fsprojects · Nov 6, 2023 · cdf7938 · cdf7938
1 parent 4a01e87
commit cdf7938
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
@@ -645,6 +645,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))
+
         member inline _.Bind(computation: Async<'T>, continuation: ('T -> ResumableTSC<'U>)) =
             ResumableTSC<'U>(fun sm ->
                 let mutable awaiter =

diff --git a/src/FSharp.Control.TaskSeq/TaskSeqInternal.fs b/src/FSharp.Control.TaskSeq/TaskSeqInternal.fs
@@ -67,6 +67,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<_>) =
         checkNonNull (nameof source) source