runtime: pull thread tracking into a header magic system to unify implementations

Author: nickwanninger

runtime/CMakeLists.txt

@@ -141,7 +141,6 @@ if(NOT ALASKA_CORE_ONLY) # -----------------------------------------------------
   set(SOURCES
     rt/init.cpp
     rt/halloc.cpp
-    rt/threads.cpp
     rt/compat.c
     rt/barrier.cpp
   )

runtime/include/alaska/rt/barrier.hpp

@@ -18,7 +18,9 @@
 #include <pthread.h>
 
 
-typedef struct {
+struct AlaskaThreadState {
+
+
   // Track the depth of escape of a given thread. If this number is zero,
   // the thread is in 'managed' code and will eventually poll the barrier
   uint64_t escaped;
@@ -29,28 +31,16 @@ typedef struct {
 #define ALASKA_JOIN_REASON_SAFEPOINT 1     // This thread was at a safepoint
 #define ALASKA_JOIN_REASON_ORCHESTRATOR 2  // This thread was the orchestrator
 #define ALASKA_JOIN_REASON_ABORT 3  // This thread requires the barrier abort (invalid state, for some reason)
-
-  // ...
-} alaska_thread_state_t;
-
-extern __thread alaska_thread_state_t alaska_thread_state;
-
+};
 
 namespace alaska {
   namespace barrier {
 
-    // Thread tracking lifetime
-    void add_self_thread(void);
-    void remove_self_thread(void);
-
     // Barrier operational lifetime. It is not recommended to use this interface, and
     // instead use `with_barrier` interface below:
     bool begin();
     void end();
 
-    // struct BarrierInfo {};
-    // bool with_barrier(ck::func<void(BarrierInfo &)> &&cb);
-
     // Initialization and deinitialization
     void init();
     void deinit();

runtime/include/alaska/thread_tracking.in.hpp

@@ -0,0 +1,124 @@
+/*
+ * This file is part of the Alaska Handle-Based Memory Management System
+ *
+ * Copyright (c) 2024, Nick Wanninger <ncw@u.northwestern.edu>
+ * Copyright (c) 2024, The Constellation Project
+ * All rights reserved.
+ *
+ * This is free software.  You are permitted to use, redistribute,
+ * and modify it as specified in the file "LICENSE".
+ */
+
+#include <alaska/ThreadRegistry.hpp>
+#include <alaska/rt.hpp>
+#include <alaska/rt/barrier.hpp>
+#include <dlfcn.h>
+#include <unistd.h>
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+// NOTE: this file is a little funky so far as header files go in C++ because it
+// provides implementations based on a macro begin defined *before* this file is
+// included.
+// To enable thread tracking in your alaska runtime, do the following:
+//
+//    struct my_state {};
+//    #define ALASKA_THREAD_TRACK_STATE_T struct my_state
+//    #include <alaska/thread_tracking.in.hpp>
+//
+// This will *define* several functions in the current file. Namely, it will override
+// `pthread_create` to enable tracking. Usually, this should just be included wherever
+// you implement your barrier logic - and should only be included *ONE TIME*.
+
+#ifndef ALASKA_THREAD_TRACK_STATE_T
+#warning Thread state type not defined. defaulting to int
+#define ALASKA_THREAD_TRACK_STATE_T int
+#endif
+
+// the following is the main interface to this system:
+namespace alaska::thread_tracking {
+  using StateT = ALASKA_THREAD_TRACK_STATE_T;
+  static __thread StateT my_state;
+
+  alaska::ThreadRegistry<StateT *> &threads();
+
+  // Join and leave with the current thread
+  void join();
+  void leave();
+}  // namespace alaska::thread_tracking
+
+
+
+// And this is the implementation:
+// static auto& threads(void) {
+// }
+
+namespace alaska::thread_tracking {
+  alaska::ThreadRegistry<StateT *> &threads() {
+    static alaska::ThreadRegistry<StateT *> *g_threads;
+    if (g_threads == NULL) g_threads = new alaska::ThreadRegistry<StateT *>();
+    return *g_threads;
+  }
+
+
+  void join(void) {
+#ifdef ALASKA_THREAD_TRACK_INIT
+    ALASKA_THREAD_TRACK_INIT;
+#endif
+    alaska::thread_tracking::threads().join(&alaska::thread_tracking::my_state);
+  }
+
+  void leave(void) { alaska::thread_tracking::threads().leave(); }
+}  // namespace alaska::thread_tracking
+
+
+
+
+
+// now, the icky part about wrapping around pthread_create
+struct alaska_pthread_trampoline_arg {
+  void* arg;
+  void* (*start)(void*);
+};
+
+
+static void* alaska_pthread_trampoline(void* varg) {
+  void* (*start)(void*);
+  auto* arg = (struct alaska_pthread_trampoline_arg*)varg;
+  void* thread_arg = arg->arg;
+  start = arg->start;
+  free(arg);
+
+  alaska::thread_tracking::join();
+  void* ret = start(thread_arg);
+  alaska::thread_tracking::leave();
+
+  return ret;
+}
+
+
+// Hook into thread creation by overriding the pthread_create function
+#undef pthread_create
+extern "C" int pthread_create(pthread_t* __restrict thread, const pthread_attr_t* __restrict attr,
+    void* (*start)(void*), void* __restrict arg) {
+  int rc;
+  static int (*real_create)(pthread_t* __restrict thread, const pthread_attr_t* __restrict attr,
+      void* (*start)(void*), void* __restrict arg) = NULL;
+  if (!real_create) real_create = (decltype(real_create))dlsym(RTLD_NEXT, "pthread_create");
+
+  auto* args = (struct alaska_pthread_trampoline_arg*)calloc(
+      1, sizeof(struct alaska_pthread_trampoline_arg));
+  args->arg = arg;
+  args->start = start;
+  rc = real_create(thread, attr, alaska_pthread_trampoline, args);
+  return rc;
+}
+
+static void __attribute__((constructor(102))) thread_tracking_init(void) {
+  alaska::thread_tracking::join();
+}
+
+static void __attribute__((destructor)) alaska_deinit(void) {
+  alaska::thread_tracking::leave();
+}

runtime/rt/barrier.cpp

@@ -103,35 +103,33 @@ static void patchNop(void) {
   }
 }
 
+static void setup_signal_handlers(void);
+static void clear_pending_signals(void);
 
 
 
 enum class JoinReason { Signal, Safepoint };
 
-static alaska::ThreadRegistry<alaska_thread_state_t*>* g_threads;
-static auto& threads(void) {
-  if (g_threads == NULL) g_threads = new alaska::ThreadRegistry<alaska_thread_state_t*>();
-  return *g_threads;
-}
+
+#define ALASKA_THREAD_TRACK_STATE_T AlaskaThreadState
+#define ALASKA_THREAD_TRACK_INIT setup_signal_handlers();
+#include <alaska/thread_tracking.in.hpp>
 
 
 
-__thread alaska_thread_state_t alaska_thread_state;
-static pthread_mutex_t barrier_lock = PTHREAD_MUTEX_INITIALIZER;
 
 // This is *the* barrier used in alaska_barrier to make sure threads are stopped correctly.
 static pthread_barrier_t the_barrier;
 static long barrier_last_num_threads = 0;
+static pthread_mutex_t barrier_lock = PTHREAD_MUTEX_INITIALIZER;
 
-static void setup_signal_handlers(void);
-static void clear_pending_signals(void);
 
 
 
 void alaska_remove_from_local_lock_list(void* ptr) { return; }
 static void alaska_dump_thread_states_r(void) {
-  struct alaska_thread_info* pos;
-  threads().for_each_locked([&](auto thread, alaska_thread_state_t* state) {
+  struct info* pos;
+  alaska::thread_tracking::threads().for_each_locked([&](auto thread, AlaskaThreadState* state) {
     if (state->escaped == 0) {
       printf("\e[0m. ");  // a thread will join a barrier (not out to lunch)
     } else {
@@ -142,7 +140,7 @@ static void alaska_dump_thread_states_r(void) {
 }
 
 void alaska_dump_thread_states(void) {
-  auto lk = threads().take_lock();
+  auto lk = alaska::thread_tracking::threads().take_lock();
   alaska_dump_thread_states_r();
 }
 
@@ -266,13 +264,11 @@ void alaska::barrier::get_pinned_handles(ck::set<void*>& out) {
 
 
 static void participant_join(bool leader, const ck::set<void*>& ps) {
-  printf("join from %lx with state %p%s\n", pthread_self(), &alaska_thread_state,
-      leader ? " as leader" : "");
   for (auto* p : ps) {
     record_handle(p, true);
   }
   // Wait on the barrier so everyone's state has been commited.
-  if (threads().num_threads() > 1) {
+  if (alaska::thread_tracking::threads().num_threads() > 1) {
     pthread_barrier_wait(&the_barrier);
   }
 }
@@ -282,7 +278,7 @@ static void participant_join(bool leader, const ck::set<void*>& ps) {
 
 static void participant_leave(bool leader, const ck::set<void*>& ps) {
   // wait for the the leader (and everyone else to catch up)
-  if (threads().num_threads() > 1) {
+  if (alaska::thread_tracking::threads().num_threads() > 1) {
     pthread_barrier_wait(&the_barrier);
   }
 
@@ -294,8 +290,8 @@ static void participant_leave(bool leader, const ck::set<void*>& ps) {
 
 
 void dump_thread_states(void) {
-  struct alaska_thread_info* pos;
-  threads().for_each_locked([&](auto thread, alaska_thread_state_t* state) {
+  struct info* pos;
+  alaska::thread_tracking::threads().for_each_locked([&](auto thread, AlaskaThreadState* state) {
     switch (state->join_status) {
       case ALASKA_JOIN_REASON_NOT_JOINED:
         printf("\e[41m! ");  // a thread will need to be interrupted (out to lunch)
@@ -334,27 +330,27 @@ bool alaska::barrier::begin(void) {
 
   // Take locks so nobody else tries to signal a barrier.
   pthread_mutex_lock(&barrier_lock);
-  threads().lock_thread_creation();
+  alaska::thread_tracking::threads().lock_thread_creation();
 
 
 
-  auto num_threads = threads().num_threads();
+  auto num_threads = alaska::thread_tracking::threads().num_threads();
   alaska::printf("Barrier begin from %lx:\n", pthread_self());
   alaska::printf("  num threads: %lu\n", num_threads);
 
   printf("  threads:\n");
-  threads().for_each_locked([](auto thread, alaska_thread_state_t* state) {
+  alaska::thread_tracking::threads().for_each_locked([](auto thread, AlaskaThreadState* state) {
     alaska::printf("  - %lx %p %d\n", thread, state, state->join_status);
   });
 
 
   // First, mark everyone as *not* in the barrier.
-  threads().for_each_locked([](auto thread, alaska_thread_state_t* state) {
+  alaska::thread_tracking::threads().for_each_locked([](auto thread, AlaskaThreadState* state) {
     state->join_status = ALASKA_JOIN_REASON_NOT_JOINED;
   });
 
   // Mark the orch thread (us) as joined
-  alaska_thread_state.join_status = ALASKA_JOIN_REASON_ORCHESTRATOR;
+  alaska::thread_tracking::my_state.join_status = ALASKA_JOIN_REASON_ORCHESTRATOR;
 
   // If the barrier needs resizing, do so.
   if (barrier_last_num_threads != num_threads) {
@@ -382,9 +378,8 @@ bool alaska::barrier::begin(void) {
     bool sent_signal = false;
     bool aborted = false;
 
-    threads().for_each_locked([&](auto thread, auto* state) {
+    alaska::thread_tracking::threads().for_each_locked([&](auto thread, auto* state) {
       if (state->join_status == ALASKA_JOIN_REASON_NOT_JOINED) {
-        printf("killing %lu\n", thread);
         pthread_kill(thread, SIGUSR2);
         sent_signal = true;
         signals_sent++;
@@ -422,7 +417,7 @@ void alaska::barrier::end(void) {
   participant_leave(true, locked);
 
   // Unlock all the locks we took.
-  threads().unlock_thread_creation();
+  alaska::thread_tracking::threads().unlock_thread_creation();
   pthread_mutex_unlock(&barrier_lock);
 }
 
@@ -467,26 +462,26 @@ static void alaska_barrier_signal_handler(int sig, siginfo_t* info, void* ptr) {
       printf(
           "ManagedUntracked: pc:0x%zx sig:%d invl:%d!\n", return_address, sig, invalid_state_abort);
       if (sig == SIGILL) {
-        alaska_thread_state.join_status = ALASKA_JOIN_REASON_ABORT;
+        alaska::thread_tracking::my_state.join_status = ALASKA_JOIN_REASON_ABORT;
         break;
       }
       if (not invalid_state_abort) {
         invalid_state_abort = true;
         return;
       }
-      alaska_thread_state.join_status = ALASKA_JOIN_REASON_ABORT;
+      alaska::thread_tracking::my_state.join_status = ALASKA_JOIN_REASON_ABORT;
       invalid_state_abort = false;
       break;
 
     case StackState::ManagedTracked:
       // it's possible to be at a managed poll point *and* get interrupted
       // through SIGUSR2
-      alaska_thread_state.join_status = ALASKA_JOIN_REASON_SAFEPOINT;
+      alaska::thread_tracking::my_state.join_status = ALASKA_JOIN_REASON_SAFEPOINT;
       break;
 
     case StackState::Unmanaged:
       assert(sig == SIGUSR2 && "Unmanaged code got into the barrier handler w/ the wrong signal");
-      alaska_thread_state.join_status = ALASKA_JOIN_REASON_SIGNAL;
+      alaska::thread_tracking::my_state.join_status = ALASKA_JOIN_REASON_SIGNAL;
       break;
   }
 
@@ -544,13 +539,13 @@ static void clear_pending_signals(void) {
 }
 
 
-void alaska::barrier::add_self_thread(void) {
-  setup_signal_handlers();
-  alaska_thread_state.escaped = 0;
-  threads().join(&alaska_thread_state);
-}
+// void alaska::barrier::add_self_thread(void) {
+//   setup_signal_handlers();
+//   alaska::thread_tracking::join();
+//   alaska::thread_tracking::my_state.escaped = 0;
+// }
 
-void alaska::barrier::remove_self_thread(void) { threads().leave(); }
+// void alaska::barrier::remove_self_thread(void) { alaska::thread_tracking::leave(); }
 
 /**
  * This function parses a stackmap emitted from LLVM and pushes all