RFC: Add task hooks for create and switch events

Certain libraries are configured using global or thread-local state
instead of passing handles to every function. CUDA, for example, has a
`cudaSetDevice` function that binds a device to the current thread for
all future API calls. This is at odds with Julia's task-based
concurrency, which presents an execution environment that's local to the
current task (e.g., in the case of CUDA, using a different device).

This PR adds a hook mechanism that can be used to detect task creation
and task switches, allowing synchronization of Julia's task-local
environment with a library's global or thread-local state.

Author: maleadt

base/task.jl

@@ -299,6 +299,31 @@ function task_local_storage(body::Function, key, val)
     end
 end
 
+# N.B. `task` must either be the current task, or a known-stopped task
+function attach_task_hook!(task::Task, hook::Ptr{Cvoid})
+    if task.hooks === nothing
+        hooks = Ptr{Cvoid}[]
+        task.hooks = hooks
+    else
+        hooks = task.hooks
+    end
+    if findfirst(==(hook), hooks) === nothing
+        push!(hooks, hook)
+    end
+    return
+end
+attach_task_hook!(hook::Ptr{Cvoid}) = attach_task_hook!(current_task(), hook)
+
+function detach_task_hook!(task::Task, hook::Ptr{Cvoid})
+    if task.hooks === nothing
+        return
+    end
+    hooks = task.hooks::Vector{Ptr{Cvoid}}
+    deleteat!(hooks, findall(==(hook), hooks))
+    return
+end
+detach_task_hook!(hook::Ptr{Cvoid}) = detach_task_hook!(current_task(), hook)
+
 # just wait for a task to be done, no error propagation
 function _wait(t::Task)
     if !istaskdone(t)

src/jltypes.c

@@ -2769,7 +2769,7 @@ void jl_init_types(void) JL_GC_DISABLED
                         NULL,
                         jl_any_type,
                         jl_emptysvec,
-                        jl_perm_symsvec(16,
+                        jl_perm_symsvec(17,
                                         "next",
                                         "queue",
                                         "storage",
@@ -2782,11 +2782,12 @@ void jl_init_types(void) JL_GC_DISABLED
                                         "rngState2",
                                         "rngState3",
                                         "rngState4",
+                                        "hooks",
                                         "_state",
                                         "sticky",
                                         "_isexception",
                                         "priority"),
-                        jl_svec(16,
+                        jl_svec(17,
                                 jl_any_type,
                                 jl_any_type,
                                 jl_any_type,
@@ -2799,6 +2800,7 @@ void jl_init_types(void) JL_GC_DISABLED
                                 jl_uint64_type,
                                 jl_uint64_type,
                                 jl_uint64_type,
+                                jl_any_type,
                                 jl_uint8_type,
                                 jl_bool_type,
                                 jl_bool_type,

src/julia.h

@@ -1925,6 +1925,7 @@ typedef struct _jl_task_t {
     // 4 byte padding on 32-bit systems
     // uint32_t padding0;
     uint64_t rngState[JL_RNG_SIZE];
+    jl_value_t *hooks;
     _Atomic(uint8_t) _state;
     uint8_t sticky; // record whether this Task can be migrated to a new thread
     _Atomic(uint8_t) _isexception; // set if `result` is an exception to throw or that we exited with
@@ -1979,6 +1980,9 @@ JL_DLLEXPORT void JL_NORETURN jl_no_exc_handler(jl_value_t *e, jl_task_t *ct);
 JL_DLLEXPORT JL_CONST_FUNC jl_gcframe_t **(jl_get_pgcstack)(void) JL_GLOBALLY_ROOTED JL_NOTSAFEPOINT;
 #define jl_current_task (container_of(jl_get_pgcstack(), jl_task_t, gcstack))
 
+typedef void *(*jl_task_switch_hook_t)(uint8_t is_switch,
+                                       jl_value_t *t JL_PROPAGATES_ROOT);
+
 #include "julia_locks.h"   // requires jl_task_t definition
 
 JL_DLLEXPORT void jl_enter_handler(jl_handler_t *eh);

src/task.c

@@ -290,9 +290,35 @@ JL_NO_ASAN static void restore_stack2(jl_task_t *t, jl_ptls_t ptls, jl_task_t *l
 /* Rooted by the base module */
 static _Atomic(jl_function_t*) task_done_hook_func JL_GLOBALLY_ROOTED = NULL;
 
+void run_task_hooks(uint8_t code, jl_value_t *t)
+{
+    jl_task_t *ct = jl_current_task;
+    jl_array_t *hooks = (jl_array_t *)ct->hooks;
+    if ((jl_value_t *)hooks != jl_nothing) {
+        jl_ptls_t ptls = ct->ptls;
+        int last_in = ptls->in_pure_callback;
+        ptls->in_pure_callback = 1;
+        for (int i = 0; i < jl_array_len(hooks); i++) {
+            jl_task_switch_hook_t hook =
+                ((jl_task_switch_hook_t *)jl_array_data(hooks))[i];
+            JL_TRY {
+                hook(code, t);
+            }
+            JL_CATCH {
+                jl_printf((JL_STREAM*)STDERR_FILENO, "WARNING: task hook threw an error:\n");
+                jl_static_show((JL_STREAM*)STDERR_FILENO, jl_current_exception());
+                jl_printf((JL_STREAM*)STDERR_FILENO, "\n");
+                jlbacktrace(); // written to STDERR_FILENO
+            }
+        }
+        ptls->in_pure_callback = last_in;
+    }
+}
+
 void JL_NORETURN jl_finish_task(jl_task_t *t)
 {
     jl_task_t *ct = jl_current_task;
+    run_task_hooks(4, jl_nothing);
     JL_PROBE_RT_FINISH_TASK(ct);
     JL_SIGATOMIC_BEGIN();
     if (jl_atomic_load_relaxed(&t->_isexception))
@@ -635,10 +661,12 @@ JL_DLLEXPORT void jl_switch(void) JL_NOTSAFEPOINT_LEAVE JL_NOTSAFEPOINT_ENTER
     if (ptls->in_finalizer)
         jl_error("task switch not allowed from inside gc finalizer");
     if (ptls->in_pure_callback)
-        jl_error("task switch not allowed from inside staged nor pure functions");
+        jl_error("task switch not allowed from inside staged nor pure functions or callbacks");
     if (!jl_set_task_tid(t, jl_atomic_load_relaxed(&ct->tid))) // manually yielding to a task
         jl_error("cannot switch to task running on another thread");
 
+    run_task_hooks(2, jl_nothing);
+
     JL_PROBE_RT_PAUSE_TASK(ct);
 
     // Store old values on the stack and reset
@@ -688,6 +716,7 @@ JL_DLLEXPORT void jl_switch(void) JL_NOTSAFEPOINT_LEAVE JL_NOTSAFEPOINT_ENTER
         jl_sigint_safepoint(ptls);
 
     JL_PROBE_RT_RUN_TASK(ct);
+    run_task_hooks(3, jl_nothing);
     jl_gc_unsafe_leave(ptls, gc_state);
 }
 
@@ -1078,6 +1107,7 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion
     t->next = jl_nothing;
     t->queue = jl_nothing;
     t->tls = jl_nothing;
+    t->hooks = jl_nothing;
     jl_atomic_store_relaxed(&t->_state, JL_TASK_STATE_RUNNABLE);
     t->start = start;
     t->result = jl_nothing;
@@ -1119,6 +1149,11 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion
 #ifdef _COMPILER_ASAN_ENABLED_
     t->ctx.asan_fake_stack = NULL;
 #endif
+
+    JL_GC_PUSH1(&t);
+    run_task_hooks(0, (jl_value_t *)t);
+    JL_GC_POP();
+
     return t;
 }
 
@@ -1234,6 +1269,9 @@ CFI_NORETURN
 
     ct->started = 1;
     JL_PROBE_RT_START_TASK(ct);
+
+    run_task_hooks(1, jl_nothing);
+
     if (jl_atomic_load_relaxed(&ct->_isexception)) {
         record_backtrace(ptls, 0);
         jl_push_excstack(&ct->excstack, ct->result,
@@ -1670,6 +1708,7 @@ jl_task_t *jl_init_root_task(jl_ptls_t ptls, void *stack_lo, void *stack_hi)
     ct->next = jl_nothing;
     ct->queue = jl_nothing;
     ct->tls = jl_nothing;
+    ct->hooks = jl_nothing;
     jl_atomic_store_relaxed(&ct->_state, JL_TASK_STATE_RUNNABLE);
     ct->start = NULL;
     ct->result = jl_nothing;

test/staged.jl

@@ -179,7 +179,7 @@ let gf_err, tsk = @async nothing # create a Task for yield to try to run
         yield()
         gf_err_ref[] += 1000
     end
-    Expected = ErrorException("task switch not allowed from inside staged nor pure functions")
+    Expected = ErrorException("task switch not allowed from inside staged nor pure functions or callbacks")
     @test_throws Expected gf_err()
     @test_throws Expected gf_err()
     @test gf_err_ref[] == 4