Wait for input/output handlers to finish when closing a StreamStore

Otherwise it's possible that the scheduler will close the DTask before all
outputs have been sent, which would cause the downstream tasks to hang. This is
how it could happen:
1. A streaming task starts.
2. The output handler task calls `take!(::ProcessRingBuffer)` on an output
   buffer, finds it empty, and `yield()`'s.
3. The task executes, pushes its output to the output buffers, reaches `max_evals`
   and finishes.
4. The scheduler finishes the corresponding DTask.
5. The `take!(::ProcessRingBuffer)` call resumes. The buffer isn't empty anymore
   but it calls `task_may_cancel(; must_force=true)` before continuing and
   throws an exception since the scheduler has finished the DTask. The result is
   that the last output is never sent, and the exeption is swallowed by the
   output handler started by `initialize_output_stream!()`.
6. Downstream tasks don't get that last result so they never reach `max_evals`
   and spin forever.

Fixed by storing the handler tasks in the `StreamStore` and closing them in
`close(::StreamStore)`.

Also increased the timeout of the 'Single task running forever' task because it
will sometimes timeout before the default 10s is up.

Author: JamesWrigley

src/stream.jl

@@ -12,13 +12,18 @@ mutable struct StreamStore{T,B}
     open::Bool
     migrating::Bool
     lock::Threads.Condition
+
+    input_handlers::Dict{UInt, Task}
+    output_handlers::Dict{UInt, Task}
+
     StreamStore{T,B}(uid::UInt, input_buffer_amount::Integer, output_buffer_amount::Integer) where {T,B} =
         new{T,B}(uid, zeros(Int, 0),
                  Dict{UInt,Any}(), Dict{UInt,Any}(),
                  Dict{UInt,B}(), Dict{UInt,B}(),
                  input_buffer_amount, output_buffer_amount,
                  Dict{UInt,Any}(), Dict{UInt,Any}(),
-                 true, false, Threads.Condition())
+                 true, false, Threads.Condition(),
+                 Dict{UInt, Task}(), Dict{UInt, Task}())
 end
 
 function tid_to_uid(thunk_id)
@@ -164,6 +169,19 @@ function Base.close(store::StreamStore)
         end
         notify(store.lock)
     end
+
+    # We have to close the input fetchers for the input handlers to finish
+    for fetcher in values(store.input_fetchers)
+        close(fetcher.chan)
+    end
+
+    # Wait for the handlers to finish
+    for handler in values(store.input_handlers)
+        wait(handler)
+    end
+    for handler in values(store.output_handlers)
+        wait(handler)
+    end
 end
 
 # FIXME: Just pass Stream directly, rather than its uid
@@ -229,7 +247,7 @@ function initialize_input_stream!(our_store::StreamStore{OT,OB}, input_stream::S
     end
     thunk_id = STREAM_THUNK_ID[]
     tls = get_tls()
-    Sch.errormonitor_tracked("streaming input: $input_uid -> $our_uid", Threads.@spawn begin
+    t = Sch.errormonitor_tracked("streaming input: $input_uid -> $our_uid", Threads.@spawn begin
         set_tls!(tls)
         STREAM_THUNK_ID[] = thunk_id
         try
@@ -247,9 +265,14 @@ function initialize_input_stream!(our_store::StreamStore{OT,OB}, input_stream::S
             @dagdebug STREAM_THUNK_ID[] :stream "input stream closed"
         end
     end)
+
+    our_store.input_handlers[input_uid] = t
+
     return StreamingValue(buffer)
 end
+
 initialize_input_stream!(our_store::StreamStore, arg) = arg
+
 function initialize_output_stream!(our_store::StreamStore{T,B}, output_uid::UInt) where {T,B}
     @assert islocked(our_store.lock)
     @dagdebug STREAM_THUNK_ID[] :stream "initializing output stream $output_uid"
@@ -260,7 +283,7 @@ function initialize_output_stream!(our_store::StreamStore{T,B}, output_uid::UInt
     output_fetcher = our_store.output_fetchers[output_uid]
     thunk_id = STREAM_THUNK_ID[]
     tls = get_tls()
-    Sch.errormonitor_tracked("streaming output: $our_uid -> $output_uid", Threads.@spawn begin
+    t = Sch.errormonitor_tracked("streaming output: $our_uid -> $output_uid", Threads.@spawn begin
         set_tls!(tls)
         STREAM_THUNK_ID[] = thunk_id
         try
@@ -282,6 +305,8 @@ function initialize_output_stream!(our_store::StreamStore{T,B}, output_uid::UInt
             @dagdebug thunk_id :stream "output stream closed"
         end
     end)
+
+    our_store.output_handlers[output_uid] = t
 end
 
 Base.put!(stream::Stream, @nospecialize(value)) = put!(stream.store, value)

test/streaming.jl

@@ -37,12 +37,14 @@ function catch_interrupt(f)
         rethrow(err)
     end
 end
+
 function merge_testset!(inner::Test.DefaultTestSet)
     outer = Test.get_testset()
     append!(outer.results, inner.results)
     outer.n_passed += inner.n_passed
 end
-function test_finishes(f, message::String; ignore_timeout=false, max_evals=10)
+
+function test_finishes(f, message::String; timeout=10, ignore_timeout=false, max_evals=10)
     t = @eval Threads.@spawn begin
         tset = nothing
         try
@@ -61,7 +63,7 @@ function test_finishes(f, message::String; ignore_timeout=false, max_evals=10)
         end
         return tset
     end
-    timed_out = timedwait(()->istaskdone(t), 10) == :timed_out
+    timed_out = timedwait(()->istaskdone(t), timeout) == :timed_out
     if timed_out
         if !ignore_timeout
             @warn "Testing task timed out: $message"
@@ -89,7 +91,7 @@ for idx in 1:5
     end
 
     @testset "Single Task Control Flow ($scope_str)" begin
-        @test !test_finishes("Single task running forever"; max_evals=1_000_000, ignore_timeout=true) do
+        @test !test_finishes("Single task running forever"; timeout=15, max_evals=1_000_000, ignore_timeout=true) do
             local x
             Dagger.spawn_streaming() do
                 x = Dagger.@spawn scope=rand(scopes) () -> begin
@@ -98,6 +100,7 @@ for idx in 1:5
                     return y
                 end
             end
+
             @test_throws_unwrap InterruptException fetch(x)
         end