AsyncSeq.merge

Author: eulerfx

src/FSharpx.Async/Async.fs

@@ -5,6 +5,7 @@
 namespace FSharpx.Control
 open System
 open System.Threading
+open System.Threading.Tasks
 
 // ----------------------------------------------------------------------------
 
@@ -92,4 +93,25 @@ module AsyncExtensions =
       static member bindChoices (f:'a -> Async<Choice<'b, 'e2>>) (a:Async<Choice<'a, 'e1>>) : Async<Choice<'b, Choice<'e1, 'e2>>> = 
         a |> Async.bind (function
           | Choice1Of2 a' -> f a' |> Async.map (function Choice1Of2 b -> Choice1Of2 b | Choice2Of2 e2 -> Choice2Of2 (Choice2Of2 e2))
-          | Choice2Of2 e1 -> Choice2Of2 (Choice1Of2 e1) |> async.Return)
+          | Choice2Of2 e1 -> Choice2Of2 (Choice1Of2 e1) |> async.Return)
+
+      /// Creates a computation which produces a tuple consiting of the value produces by the first
+      /// argument computation to complete and a handle to the other computation. The second computation
+      /// to complete is memoized.
+      static member internal chooseBoth (a:Async<'a>) (b:Async<'a>) : Async<'a * Async<'a>> =
+        Async.FromContinuations <| fun (ok,err,cnc) ->
+          let state = ref 0            
+          let tcs = TaskCompletionSource<'a>()            
+          let inline ok a =
+            if (Interlocked.CompareExchange(state, 1, 0) = 0) then
+              ok (a, tcs.Task |> Async.AwaitTask)
+            else
+              tcs.SetResult a
+          let inline err (ex:exn) =
+            if (Interlocked.CompareExchange(state, 1, 0) = 0) then err ex
+            else tcs.SetException ex
+          let inline cnc ex =
+            if (Interlocked.CompareExchange(state, 1, 0) = 0) then cnc ex
+            else tcs.SetCanceled()
+          Async.StartWithContinuations(a, ok, err, cnc)
+          Async.StartWithContinuations(b, ok, err, cnc)

src/FSharpx.Async/AsyncSeq.fs

@@ -605,7 +605,23 @@ module AsyncSeq =
             return! loop tl            
       }
       return! loop s
-    }          
+    }
+
+  /// Merges two async sequences into an async sequence non-deterministically.
+  let rec merge (a:AsyncSeq<'a>) (b:AsyncSeq<'a>) : AsyncSeq<'a> = async {
+    let! one,other = Async.chooseBoth a b
+    match one with
+    | Nil -> return! other
+    | Cons(hd,tl) ->
+      return Cons(hd, merge tl other) }
+
+  /// Merges all specified async sequences into an async sequence non-deterministically.
+  let rec mergeAll (ss:AsyncSeq<'a> list) : AsyncSeq<'a> =
+    match ss with
+    | [] -> empty
+    | [s] -> s
+    | [a;b] -> merge a b
+    | hd::tl -> merge hd (mergeAll tl)
 
 
 

tests/FSharpx.Async.Tests/AsyncSeqTests.fs

@@ -60,7 +60,6 @@ let ``AsyncSeq.interleave``() =
   let s1 = AsyncSeq.ofSeq ["a";"b";"c"]
   let s2 = AsyncSeq.ofSeq [1;2;3]
   let merged = AsyncSeq.interleave s1 s2 |> AsyncSeq.toList |> Async.RunSynchronously
-  printfn "%A" merged
   Assert.True([Choice1Of2 "a" ; Choice2Of2 1 ; Choice1Of2 "b" ; Choice2Of2 2 ; Choice1Of2 "c" ; Choice2Of2 3] = merged)
 
 
@@ -180,6 +179,27 @@ let ``AsyncSeq.filterAsync``() =
   Assert.True(EQ expected actual)
 
 
+[<Test>]
+let ``AsyncSeq.merge``() =
+  let ls1 = [1;2;3;4;5]
+  let ls2 = [6;7;8;9;10]
+  let actual = AsyncSeq.merge (AsyncSeq.ofSeq ls1) (AsyncSeq.ofSeq ls2) |> AsyncSeq.toList |> Async.RunSynchronously |> Set.ofList
+  let expected = ls1 @ ls2 |> Set.ofList
+  Assert.True((expected = actual))
+
+
+[<Test>]
+let ``AsyncSeq.merge should be fair``() =  
+  let s1 = asyncSeq {
+    do! Async.Sleep 10
+    yield 1
+  }
+  let s2 = asyncSeq {
+    yield 2
+  }
+  let actual = AsyncSeq.merge s1 s2
+  let expected = [2;1] |> AsyncSeq.ofSeq
+  Assert.True(EQ expected actual)