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/task.c

@@ -480,6 +480,13 @@ static jl_ptls_t NOINLINE refetch_ptls(void)
     return jl_get_ptls_states();
 }
 
+typedef void *(*jl_task_switch_hook_t)(jl_task_t *t JL_PROPAGATES_ROOT);
+jl_task_switch_hook_t task_switch_hook = NULL;
+JL_DLLEXPORT void jl_hook_task_switch(jl_task_switch_hook_t hook)
+{
+    task_switch_hook = hook;
+}
+
 JL_DLLEXPORT void jl_switch(void)
 {
     jl_ptls_t ptls = jl_get_ptls_states();
@@ -497,7 +504,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 +514,13 @@ JL_DLLEXPORT void jl_switch(void)
         jl_error("cannot switch to task running on another thread");
     }
 
+    if (task_switch_hook) {
+        int last_in = ptls->in_pure_callback;
+        ptls->in_pure_callback = 1;
+        task_switch_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);