Merge pull request #339 from fsprojects/repo-assist/fix-issue-167-withCancellation-128d379345c09895

[Repo Assist] adds TaskSeq.withCancellation

Author: dsyme

release-notes.txt

@@ -10,6 +10,7 @@ Release notes:
     - adds TaskSeq.unfold and TaskSeq.unfoldAsync, #289
     - adds TaskSeq.chunkBySize (closes #258) and TaskSeq.windowed, #289
     - fixes: CancellationToken passed to GetAsyncEnumerator is now honored in MoveNextAsync, #179
+    - adds TaskSeq.withCancellation, #167
 
 0.6.0
     - fixes: async { for item in taskSeq do ... } no longer wraps exceptions in AggregateException, #129

src/FSharp.Control.TaskSeq.Test/FSharp.Control.TaskSeq.Test.fsproj

@@ -71,6 +71,7 @@
     <Compile Include="TaskSeq.Let.Tests.fs" />
     <Compile Include="TaskSeq.Using.Tests.fs" />
     <Compile Include="TaskSeq.CancellationToken.Tests.fs" />
+    <Compile Include="TaskSeq.WithCancellation.Tests.fs" />
   </ItemGroup>
 
   <ItemGroup>

src/FSharp.Control.TaskSeq.Test/TaskSeq.WithCancellation.Tests.fs

@@ -0,0 +1,172 @@
+module TaskSeq.Tests.``WithCancellation``
+
+open System
+open System.Collections.Generic
+open System.Threading
+open System.Threading.Tasks
+
+open Xunit
+open FsUnit.Xunit
+
+open FSharp.Control
+
+/// A simple IAsyncEnumerable whose GetAsyncEnumerator records the token it was called with.
+type TokenCapturingSeq<'T>(items: 'T list) =
+    let mutable capturedToken = CancellationToken.None
+
+    member _.CapturedToken = capturedToken
+
+    interface IAsyncEnumerable<'T> with
+        member _.GetAsyncEnumerator(ct) =
+            capturedToken <- ct
+
+            let source = taskSeq {
+                for x in items do
+                    yield x
+            }
+
+            source.GetAsyncEnumerator(ct)
+
+module ``Null check`` =
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: null source throws ArgumentNullException`` () =
+        assertNullArg
+        <| fun () -> TaskSeq.withCancellation CancellationToken.None null
+
+module ``Token threading`` =
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: passes supplied token to GetAsyncEnumerator`` () = task {
+        let source = TokenCapturingSeq([ 1; 2; 3 ])
+        use cts = new CancellationTokenSource()
+
+        let wrapped = TaskSeq.withCancellation cts.Token (source :> IAsyncEnumerable<_>)
+        let! _ = TaskSeq.toArrayAsync wrapped
+        source.CapturedToken |> should equal cts.Token
+    }
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: overrides any token passed to GetAsyncEnumerator`` () = task {
+        let source = TokenCapturingSeq([ 1; 2; 3 ])
+        use cts = new CancellationTokenSource()
+
+        let wrapped = TaskSeq.withCancellation cts.Token (source :> IAsyncEnumerable<_>)
+
+        // Consume with a different token; withCancellation should win
+        use outerCts = new CancellationTokenSource()
+        let enum = wrapped.GetAsyncEnumerator(outerCts.Token)
+
+        while! enum.MoveNextAsync() do
+            ()
+
+        source.CapturedToken |> should equal cts.Token
+    }
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: CancellationToken.None passes through correctly`` () = task {
+        let source = TokenCapturingSeq([ 10; 20 ])
+
+        let wrapped = TaskSeq.withCancellation CancellationToken.None (source :> IAsyncEnumerable<_>)
+        let! _ = TaskSeq.toArrayAsync wrapped
+        source.CapturedToken |> should equal CancellationToken.None
+    }
+
+module ``Cancellation behaviour`` =
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: pre-cancelled token causes OperationCanceledException on iteration`` () = task {
+        use cts = new CancellationTokenSource()
+        cts.Cancel()
+
+        let source = taskSeq {
+            while true do
+                yield 1
+        }
+
+        let wrapped = TaskSeq.withCancellation cts.Token source
+
+        fun () -> TaskSeq.iter ignore wrapped |> Task.ignore
+        |> should throwAsync typeof<OperationCanceledException>
+    }
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: token cancelled mid-iteration raises OperationCanceledException`` () = task {
+        use cts = new CancellationTokenSource()
+
+        let source = taskSeq {
+            for i in 1..100 do
+                yield i
+        }
+
+        let wrapped = TaskSeq.withCancellation cts.Token source
+
+        fun () ->
+            task {
+                let mutable count = 0
+                use enum = wrapped.GetAsyncEnumerator(CancellationToken.None)
+
+                while! enum.MoveNextAsync() do
+                    count <- count + 1
+
+                    if count = 3 then
+                        cts.Cancel()
+            }
+            |> Task.ignore
+        |> should throwAsync typeof<OperationCanceledException>
+    }
+
+module ``Sequence contents`` =
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: empty source produces empty sequence`` () =
+        TaskSeq.empty<int>
+        |> TaskSeq.withCancellation CancellationToken.None
+        |> verifyEmpty
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: finite source produces all items`` () = task {
+        let! result =
+            taskSeq {
+                for i in 1..10 do
+                    yield i
+            }
+            |> TaskSeq.withCancellation CancellationToken.None
+            |> TaskSeq.toArrayAsync
+
+        result |> should equal [| 1..10 |]
+    }
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: can be used with TaskSeq combinators`` () = task {
+        use cts = new CancellationTokenSource()
+
+        let! result =
+            taskSeq {
+                for i in 1..5 do
+                    yield i
+            }
+            |> TaskSeq.withCancellation cts.Token
+            |> TaskSeq.map (fun x -> x * 2)
+            |> TaskSeq.toArrayAsync
+
+        result |> should equal [| 2; 4; 6; 8; 10 |]
+    }
+
+    [<Fact>]
+    let ``TaskSeq-withCancellation: can be piped like .WithCancellation usage pattern`` () = task {
+        use cts = new CancellationTokenSource()
+        let mutable collected = ResizeArray()
+
+        let source = taskSeq {
+            for i in 1..5 do
+                yield i
+        }
+
+        do!
+            source
+            |> TaskSeq.withCancellation cts.Token
+            |> TaskSeq.iterAsync (fun x -> task { collected.Add(x) })
+
+        collected |> Seq.toArray |> should equal [| 1..5 |]
+    }

src/FSharp.Control.TaskSeq/TaskSeq.fs

@@ -258,6 +258,13 @@ type TaskSeq private () =
                 yield c
         }
 
+    static member withCancellation (cancellationToken: CancellationToken) (source: TaskSeq<'T>) =
+        Internal.checkNonNull (nameof source) source
+
+        { new IAsyncEnumerable<'T> with
+            member _.GetAsyncEnumerator(_ct) = source.GetAsyncEnumerator(cancellationToken)
+        }
+
     //
     // Utility functions
     //

src/FSharp.Control.TaskSeq/TaskSeq.fsi

@@ -1,6 +1,7 @@
 namespace FSharp.Control
 
 open System.Collections.Generic
+open System.Threading
 open System.Threading.Tasks
 
 [<AutoOpen>]
@@ -602,6 +603,23 @@ type TaskSeq =
     /// <exception cref="T:ArgumentNullException">Thrown when the input sequence is null.</exception>
     static member ofAsyncArray: source: Async<'T> array -> TaskSeq<'T>
 
+    /// <summary>
+    /// Returns a task sequence that, when iterated, passes the given <paramref name="cancellationToken" /> to the
+    /// underlying <see cref="IAsyncEnumerable&lt;'T&gt;" />. This is the equivalent of calling
+    /// <c>.WithCancellation(cancellationToken)</c> on an <see cref="IAsyncEnumerable&lt;'T&gt;" />.
+    /// </summary>
+    /// <remarks>
+    /// The <paramref name="cancellationToken" /> supplied to this function overrides any token that would otherwise
+    /// be passed to the enumerator. This is useful when consuming sequences from libraries such as Entity Framework,
+    /// which accept a <see cref="CancellationToken" /> through <c>GetAsyncEnumerator</c>.
+    /// </remarks>
+    ///
+    /// <param name="cancellationToken">The cancellation token to pass to <c>GetAsyncEnumerator</c>.</param>
+    /// <param name="source">The input task sequence.</param>
+    /// <returns>A task sequence that uses the given <paramref name="cancellationToken" /> when iterated.</returns>
+    /// <exception cref="T:ArgumentNullException">Thrown when the input task sequence is null.</exception>
+    static member withCancellation: cancellationToken: CancellationToken -> source: TaskSeq<'T> -> TaskSeq<'T>
+
     /// <summary>
     /// Views each item in the input task sequence as <see cref="obj" />, boxing value types.
     /// </summary>