Refactor, document, make more robust and remove a lock

Author: Liozou

src/C/context.jl

@@ -30,18 +30,59 @@ function _atpyexit()
 end
 
 
-const MAIN_THREAD_TASK_LOCK = ReentrantLock()
-const MAIN_THREAD_CHANNEL_INPUT = Channel(1)
-const MAIN_THREAD_CHANNEL_OUTPUT = Channel(1)
-
-# Execute f() on the main thread.
-function on_main_thread(f)
-    @lock MAIN_THREAD_TASK_LOCK begin
-        put!(MAIN_THREAD_CHANNEL_INPUT, f)
-        take!(MAIN_THREAD_CHANNEL_OUTPUT)
+function setup_onfixedthread()
+    channel_input = Channel(1)
+    channel_output = Channel(1)
+    islaunched = Ref(false) # use Ref to avoid closure boxing of variable
+    function launch_worker(tid)
+        islaunched[] && error("Cannot launch more than once: call setup_onfixedthread again if need be.")
+        islaunched[] = true
+        worker_task = Task() do
+            while true
+                f = take!(channel_input)
+                put!(channel_output, f())
+            end
+        end
+        # code adapted from set_task_tid! in StableTasks.jl, itself taken from Dagger.jl
+        worker_task.sticky = true
+        for _ in 1:100
+            # try to fix the task id to tid, retrying up to 100 times
+            ret = ccall(:jl_set_task_tid, Cint, (Any, Cint), worker_task, tid-1)
+            if ret == 1
+                break # success
+            elseif ret == 0
+                yield()
+            else
+                error("Unexpected retcode from jl_set_task_tid: $ret")
+            end
+        end
+        if Threads.threadid(worker_task) != tid
+            error("Failed setting the thread ID to $tid.")
+        end
+        schedule(worker_task)
+    end
+    function onfixedthread(f)
+        put!(channel_input, f)
+        take!(channel_output)
     end
+    launch_worker, onfixedthread
 end
 
+# launch_on_main_thread is used in init_context(), after which on_main_thread becomes usable
+const launch_on_main_thread, on_main_thread = setup_onfixedthread()
+
+"""
+    on_main_thread(f)
+
+Execute `f()` on the main thread.
+
+!!! warning
+    The value returned by `on_main_thread(f)` cannot be type-inferred by the compiler:
+    if necessary, use explicit type annotations such as `on_main_thread(f)::T`, where `T` is
+    the expected return type.
+"""
+on_main_thread
+
 
 function init_context()
 
@@ -254,14 +295,7 @@ function init_context()
         "Only Python 3.9+ is supported, this is Python $(CTX.version) at $(CTX.exe_path===missing ? "unknown location" : CTX.exe_path).",
     )
 
-    main_thread_task = Task() do
-        while true
-            f = take!(MAIN_THREAD_CHANNEL_INPUT)
-            put!(MAIN_THREAD_CHANNEL_OUTPUT, f())
-        end
-    end
-    set_task_tid!(main_thread_task, Threads.threadid())
-    schedule(main_thread_task)
+    launch_on_main_thread(Threads.threadid()) # makes on_main_thread usable
 
     @debug "Initialized PythonCall.jl" CTX.is_embedded CTX.is_initialized CTX.exe_path CTX.lib_path CTX.lib_ptr CTX.pyprogname CTX.pyhome CTX.version
 
@@ -283,26 +317,3 @@ const PYTHONCALL_PKGID = Base.PkgId(PYTHONCALL_UUID, "PythonCall")
 
 const PYCALL_UUID = Base.UUID("438e738f-606a-5dbb-bf0a-cddfbfd45ab0")
 const PYCALL_PKGID = Base.PkgId(PYCALL_UUID, "PyCall")
-
-
-# taken from StableTasks.jl, itself taken from Dagger.jl
-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