From cdf7938ede3879e66ce5ed59f872869fbc03d052 Mon Sep 17 00:00:00 2001 From: Abel Braaksma Date: Sat, 17 Dec 2022 03:40:38 +0100 Subject: [PATCH] Add support for `CancellationToken` to TaskSeq --- src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs | 6 ++++++ src/FSharp.Control.TaskSeq/TaskSeqInternal.fs | 5 +++++ 2 files changed, 11 insertions(+) diff --git a/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs b/src/FSharp.Control.TaskSeq/TaskSeqBuilder.fs index 0a6905e3..18d0f7ae 100644 --- 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 index 7ea572a5..99b3c378 100644 --- 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