RAI: Add heartbeat capability

Author: kpamnany

src/gc-stock.c

@@ -3395,6 +3395,9 @@ static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection)
     return recollect;
 }
 
+extern int jl_heartbeat_pause(void);
+extern int jl_heartbeat_resume(void);
+
 JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
 {
     JL_PROBE_GC_BEGIN(collection);
@@ -3437,6 +3440,7 @@ JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
     // existence of the thread in the jl_n_threads count.
     //
     // TODO: concurrently queue objects
+    jl_heartbeat_pause();
     jl_fence();
     gc_n_threads = jl_atomic_load_acquire(&jl_n_threads);
     gc_all_tls_states = jl_atomic_load_relaxed(&jl_all_tls_states);
@@ -3468,6 +3472,7 @@ JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
 
     gc_n_threads = 0;
     gc_all_tls_states = NULL;
+    jl_heartbeat_resume();
     jl_safepoint_end_gc();
     jl_gc_state_set(ptls, old_state, JL_GC_STATE_WAITING);
     JL_PROBE_GC_END();

src/gc-stock.h

@@ -499,6 +499,9 @@ extern uv_sem_t gc_sweep_assists_needed;
 extern _Atomic(int) gc_n_threads_marking;
 extern _Atomic(int) gc_n_threads_sweeping_pools;
 extern _Atomic(int) n_threads_running;
+extern _Atomic(int) gc_n_threads_sweeping_stacks;
+extern _Atomic(int) gc_ptls_sweep_idx;
+extern _Atomic(int) gc_stack_free_idx;
 extern uv_barrier_t thread_init_done;
 void gc_mark_queue_all_roots(jl_ptls_t ptls, jl_gc_markqueue_t *mq);
 void gc_mark_finlist_(jl_gc_markqueue_t *mq, jl_value_t *fl_parent, jl_value_t **fl_begin, jl_value_t **fl_end) JL_NOTSAFEPOINT;

src/init.c

@@ -863,6 +863,8 @@ JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel)
     _finish_julia_init(rel, ptls, ct);
 }
 
+void jl_init_heartbeat(void);
+
 static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_task_t *ct)
 {
     JL_TIMING(JULIA_INIT, JULIA_INIT);
@@ -916,6 +918,11 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
     jl_start_gc_threads();
     uv_barrier_wait(&thread_init_done);
 
+    if (jl_base_module != NULL) {
+        // requires code in Base
+        jl_init_heartbeat();
+    }
+
     jl_gc_enable(1);
 
     if (jl_options.image_file && (!jl_generating_output() || jl_options.incremental) && jl_module_init_order) {

src/julia_internal.h

@@ -224,6 +224,7 @@ extern volatile size_t profile_bt_size_max;
 extern volatile size_t profile_bt_size_cur;
 extern volatile int profile_running;
 extern volatile int profile_all_tasks;
+extern int heartbeat_tid; // Mostly used to ensure we skip this thread in the CPU profiler. XXX: not implemented on Windows
 // Ensures that we can safely read the `live_tasks`field of every TLS when profiling.
 // We want to avoid the case that a GC gets interleaved with `jl_profile_task` and shrinks
 // the `live_tasks` array while we are reading it or frees tasks that are being profiled.
@@ -239,6 +240,7 @@ extern uv_mutex_t bt_data_prof_lock;
 #define PROFILE_STATE_THREAD_NOT_SLEEPING (1)
 #define PROFILE_STATE_THREAD_SLEEPING (2)
 #define PROFILE_STATE_WALL_TIME_PROFILING (3)
+extern _Atomic(int) n_threads_running;
 void jl_profile_task(void);
 
 // number of cycles since power-on

src/options.h

@@ -144,6 +144,9 @@
 #define MACHINE_EXCLUSIVE_NAME          "JULIA_EXCLUSIVE"
 #define DEFAULT_MACHINE_EXCLUSIVE       0
 
+// heartbeats
+#define JL_HEARTBEAT_THREAD
+
 // sanitizer defaults ---------------------------------------------------------
 
 // Automatically enable MEMDEBUG and KEEP_BODIES for the sanitizers

src/signals-mach.c

@@ -826,6 +826,10 @@ void *mach_profile_listener(void *arg)
             for (int idx = nthreads; idx-- > 0; ) {
                 // Stop the threads in random order.
                 int i = randperm[idx];
+                // skip heartbeat thread
+                if (i == heartbeat_tid) {
+                    continue;
+                }
                 jl_profile_thread_mach(i);
             }
         }

src/signals-unix.c

@@ -872,6 +872,10 @@ static void do_profile(void *ctx)
     for (int idx = nthreads; idx-- > 0; ) {
         // Stop the threads in the random order.
         int tid = randperm[idx];
+        // skip heartbeat thread
+        if (tid == heartbeat_tid) {
+            return;
+        }
         // do backtrace for profiler
         if (!profile_running)
             return;

src/stackwalk.c

@@ -1402,9 +1402,22 @@ JL_DLLEXPORT void jl_print_backtrace(void) JL_NOTSAFEPOINT
     jlbacktrace();
 }
 
-// Print backtraces for all live tasks, for all threads, to jl_safe_printf stderr
+extern int jl_inside_heartbeat_thread(void);
+extern int jl_heartbeat_pause(void);
+extern int jl_heartbeat_resume(void);
+
+// Print backtraces for all live tasks, for all threads, to jl_safe_printf
+// stderr. This can take a _long_ time!
 JL_DLLEXPORT void jl_print_task_backtraces(int show_done) JL_NOTSAFEPOINT
 {
+    // disable heartbeats to prevent heartbeat loss while running this,
+    // unless this is called from the heartbeat thread itself; in that
+    // situation, the thread is busy running this and it will not be
+    // updating the missed heartbeats counter
+    if (!jl_inside_heartbeat_thread()) {
+        jl_heartbeat_pause();
+    }
+
     size_t nthreads = jl_atomic_load_acquire(&jl_n_threads);
     jl_ptls_t *allstates = jl_atomic_load_relaxed(&jl_all_tls_states);
     int ctid = jl_threadid() + 1;
@@ -1463,6 +1476,10 @@ JL_DLLEXPORT void jl_print_task_backtraces(int show_done) JL_NOTSAFEPOINT
         jl_safe_printf("thread (%d) ==== End thread %d\n", ctid, ptls2->tid + 1);
     }
     jl_safe_printf("thread (%d) ++++ Done\n", ctid);
+
+    if (!jl_inside_heartbeat_thread()) {
+        jl_heartbeat_resume();
+    }
 }
 
 #ifdef __cplusplus