Make Dagger.spawn safely nestable

Author: jpsamaroo

Diff for: src/sch/eager.jl

@@ -29,18 +29,44 @@ function init_eager()
     end
 end
 
-"Sets the scheduler's cached pressure indicator for the specified worker."
-set_pressure!(h::SchedulerHandle, pid::Int, proctype::Type, pressure::Int) =
-    exec!(_set_pressure!, h, pid, proctype, pressure)
-function _set_pressure!(ctx, state, task, tid, (pid, proctype, pressure))
-    state.worker_pressure[pid][proctype] = pressure
-    ACTIVE_TASKS[state.uid][proctype] = pressure # HACK-ish
+"Adjusts the scheduler's cached pressure indicator for the specified worker by
+the specified amount."
+function adjust_pressure!(h::SchedulerHandle, proctype::Type, pressure)
+    uid = Dagger.get_tls().sch_uid
+    lock(ACTIVE_TASKS_LOCK) do
+        ACTIVE_TASKS[uid][proctype][] += pressure
+    end
+    exec!(_adjust_pressure!, h, myid(), proctype, pressure)
+end
+function _adjust_pressure!(ctx, state, task, tid, (pid, proctype, pressure))
+    state.worker_pressure[pid][proctype] += pressure
     nothing
 end
 
+"Allows a thunk to safely wait on another thunk, by temporarily reducing its
+effective pressure to 0."
+function thunk_yield(f)
+    if Dagger.in_thunk()
+        h = sch_handle()
+        tls = Dagger.get_tls()
+        proctype = typeof(tls.processor)
+        util = tls.utilization
+        adjust_pressure!(h, proctype, -util)
+        try
+            f()
+        finally
+            adjust_pressure!(h, proctype, util)
+        end
+    else
+        f()
+    end
+end
+
 function eager_thunk()
     h = sch_handle()
-    set_pressure!(h, 1, Dagger.ThreadProc, 0) # HACK: Don't apply pressure from this thunk
+    util = Dagger.get_tls().utilization
+    # Don't apply pressure from this thunk
+    adjust_pressure!(h, Dagger.ThreadProc, -util)
     while isopen(EAGER_THUNK_CHAN)
         try
             future, uid, f, args, opts = take!(EAGER_THUNK_CHAN)

Diff for: src/thunk.jl

@@ -90,9 +90,13 @@ end
 ThunkFuture(x::Integer) = ThunkFuture(Future(x))
 ThunkFuture() = ThunkFuture(Future())
 Base.isready(t::ThunkFuture) = isready(t.future)
-Base.wait(t::ThunkFuture) = wait(t.future)
+Base.wait(t::ThunkFuture) = Dagger.Sch.thunk_yield() do
+    wait(t.future)
+end
 function Base.fetch(t::ThunkFuture; proc=OSProc())
-    error, value = move(proc, fetch(t.future))
+    error, value = Dagger.Sch.thunk_yield() do
+        move(proc, fetch(t.future))
+    end
     if error
         throw(value)
     end

Diff for: test/thunk.jl

@@ -1,6 +1,13 @@
 import Dagger: @par, @spawn, spawn
 
-@everywhere checkwid() = myid()==1
+@everywhere begin
+    checkwid() = myid()==1
+    function dynamic_fib(n)
+        n <= 1 && return n
+        t = Dagger.spawn(dynamic_fib, n-1)
+        return (fetch(t)::Int) + dynamic_fib(n-2)
+    end
+end
 
 @testset "@par" begin
     @testset "per-call" begin
@@ -138,5 +145,8 @@ end
         @test fetch(Distributed.@spawnat 2 !(Dagger.Sch.EAGER_INIT[]))
         @test a isa Dagger.EagerThunk
         @test fetch(a) == 3
+
+        # Mild stress-test
+        @test fetch(dynamic_fib(10)) == 55
     end
 end