RFC: Add a hook for detecting task switches.

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 switches,
and synchronize Julia's task-local environment with the library's global
or thread-local state.

Author: maleadt

src/gc.c

@@ -2799,6 +2799,8 @@ static void mark_roots(jl_gc_mark_cache_t *gc_cache, jl_gc_mark_sp_t *sp)
         gc_mark_queue_obj(gc_cache, sp, jl_an_empty_vec_any);
     if (jl_module_init_order != NULL)
         gc_mark_queue_obj(gc_cache, sp, jl_module_init_order);
+    if (jl_task_switch_hooks != NULL)
+        gc_mark_queue_obj(gc_cache, sp, jl_task_switch_hooks);
     for (size_t i = 0; i < jl_current_modules.size; i += 2) {
         if (jl_current_modules.table[i + 1] != HT_NOTFOUND) {
             gc_mark_queue_obj(gc_cache, sp, jl_current_modules.table[i]);

src/julia.h

@@ -1827,6 +1827,9 @@ JL_DLLEXPORT void JL_NORETURN jl_sig_throw(void);
 JL_DLLEXPORT void JL_NORETURN jl_rethrow_other(jl_value_t *e JL_MAYBE_UNROOTED);
 JL_DLLEXPORT void JL_NORETURN jl_no_exc_handler(jl_value_t *e);
 
+typedef void *(*jl_task_switch_hook_t)(jl_task_t *t JL_PROPAGATES_ROOT);
+JL_DLLEXPORT void jl_hook_task_switch(jl_task_switch_hook_t hook);
+
 #include "locks.h"   // requires jl_task_t definition
 
 JL_DLLEXPORT void jl_enter_handler(jl_handler_t *eh);

src/julia_internal.h

@@ -368,6 +368,7 @@ void jl_gc_track_malloced_array(jl_ptls_t ptls, jl_array_t *a) JL_NOTSAFEPOINT;
 void jl_gc_count_allocd(size_t sz) JL_NOTSAFEPOINT;
 void jl_gc_run_all_finalizers(jl_ptls_t ptls);
 void jl_release_task_stack(jl_ptls_t ptls, jl_task_t *task);
+extern jl_array_t *jl_task_switch_hooks JL_GLOBALLY_ROOTED;
 
 void gc_queue_binding(jl_binding_t *bnd) JL_NOTSAFEPOINT;
 void gc_setmark_buf(jl_ptls_t ptls, void *buf, uint8_t, size_t) JL_NOTSAFEPOINT;

src/task.c

@@ -480,6 +480,19 @@ static jl_ptls_t NOINLINE refetch_ptls(void)
     return jl_get_ptls_states();
 }
 
+jl_array_t *jl_task_switch_hooks = JL_GLOBALLY_ROOTED = NULL;
+JL_DLLEXPORT void jl_hook_task_switch(jl_task_switch_hook_t hook)
+{
+    if (jl_task_switch_hooks == NULL) {
+        jl_value_t *array_ptr_void_type = jl_apply_type2(
+            (jl_value_t *)jl_array_type, (jl_value_t *)jl_voidpointer_type, jl_box_long(1));
+        jl_task_switch_hooks = jl_alloc_array_1d(array_ptr_void_type, 0);
+    }
+    jl_array_grow_end(jl_task_switch_hooks, 1);
+    ((jl_task_switch_hook_t *)jl_array_data(
+        jl_task_switch_hooks))[jl_array_len(jl_task_switch_hooks) - 1] = hook;
+}
+
 JL_DLLEXPORT void jl_switch(void)
 {
     jl_ptls_t ptls = jl_get_ptls_states();
@@ -497,7 +510,7 @@ JL_DLLEXPORT void jl_switch(void)
     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 (t->sticky && jl_atomic_load_acquire(&t->tid) == -1) {
         // manually yielding to a task
         if (jl_atomic_compare_exchange(&t->tid, -1, ptls->tid) != -1)
@@ -507,6 +520,17 @@ JL_DLLEXPORT void jl_switch(void)
         jl_error("cannot switch to task running on another thread");
     }
 
+    if (jl_task_switch_hooks) {
+        int last_in = ptls->in_pure_callback;
+        ptls->in_pure_callback = 1;
+        for (int i = 0; i < jl_array_len(jl_task_switch_hooks); i++) {
+            jl_task_switch_hook_t hook =
+                ((jl_task_switch_hook_t *)jl_array_data(jl_task_switch_hooks))[i];
+            hook(t);
+        }
+        ptls->in_pure_callback = last_in;
+    }
+
     // Store old values on the stack and reset
     sig_atomic_t defer_signal = ptls->defer_signal;
     int8_t gc_state = jl_gc_unsafe_enter(ptls);