Retry jl_set_task_tid on failure

src/Dagger.jl

@@ -25,6 +25,7 @@ include("utils/dagdebug.jl")
 
 # Distributed data
 include("utils/locked-object.jl")
+include("utils/tasks.jl")
 include("options.jl")
 include("processor.jl")
 include("scopes.jl")

src/processor.jl

@@ -166,12 +166,7 @@ function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...); @n
         TimespanLogging.prof_task_put!(tls.sch_handle.thunk_id.id)
         @invokelatest f(args...; kwargs...)
     end
-    task.sticky = true
-    ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), task, proc.tid-1)
-    if ret == 0
-        error("jl_set_task_tid == 0")
-    end
-    @assert Threads.threadid(task) == proc.tid
+    set_task_tid!(task, proc.tid)
     schedule(task)
     try
         fetch(task)

src/sch/Sch.jl

@@ -1286,8 +1286,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
             lock(istate.queue) do _
                 tid = task_tid_for_processor(to_proc)
                 if tid !== nothing
-                    t.sticky = true
-                    ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), t, tid-1)
+                    Dagger.set_task_tid!(t, tid)
                 else
                     t.sticky = false
                 end
@@ -1299,8 +1298,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
     end
     tid = task_tid_for_processor(to_proc)
     if tid !== nothing
-        proc_run_task.sticky = true
-        ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), proc_run_task, tid-1)
+        Dagger.set_task_tid!(proc_run_task, tid)
     else
         proc_run_task.sticky = false
     end

src/utils/tasks.jl

@@ -0,0 +1,20 @@
+function set_task_tid!(task::Task, tid::Integer)
+    task.sticky = true
+    ctr = 0
+    while true
+        ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), task, tid-1)
+        if ret == 1
+            break
+        elseif ret == 0
+            yield()
+        else
+            error("Unexpected retcode from jl_set_task_tid: $ret")
+        end
+        ctr += 1
+        if ctr > 10
+            @warn "Setting task TID to $tid failed, giving up!"
+            return
+        end
+    end
+    @assert Threads.threadid(task) == tid "jl_set_task_tid failed!"
+end