diff --git a/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs b/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs index e8f1f729..7cfbcf78 100644 --- 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)) + [] 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 index 89ce51e3..6904a60e 100644 --- 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()