Instrument native threads

ddprof-lib/src/main/cpp/codeCache.h

@@ -226,14 +226,18 @@ class CodeCacheArray {
 
   CodeCache *operator[](int index) { return _libs[index]; }
 
-  int count() { return __atomic_load_n(&_count, __ATOMIC_ACQUIRE); }
+  int count() const { return __atomic_load_n(&_count, __ATOMIC_ACQUIRE); }
 
   void add(CodeCache *lib) {
     int index = __atomic_load_n(&_count, __ATOMIC_ACQUIRE);
     _libs[index] = lib;
     __atomic_store_n(&_count, index + 1, __ATOMIC_RELEASE);
   }
 
+  CodeCache* at(int index) const {
+    return _libs[index];
+  }
+
   long long memoryUsage() {
     int count = __atomic_load_n(&_count, __ATOMIC_ACQUIRE);
     long long totalUsage = 0;

ddprof-lib/src/main/cpp/ctimer_linux.cpp

@@ -23,6 +23,7 @@
 #include "profiler.h"
 #include "vmStructs.h"
 #include <assert.h>
+#include <dlfcn.h>
 #include <stdlib.h>
 #include <sys/syscall.h>
 #include <time.h>
@@ -32,11 +33,84 @@
 #define SIGEV_THREAD_ID 4
 #endif
 
-static inline clockid_t thread_cpu_clock(unsigned int tid) {
-  return ((~tid) << 3) | 6; // CPUCLOCK_SCHED | CPUCLOCK_PERTHREAD_MASK
+typedef void* (*func_start_routine)(void*);
+
+// Patch libraries' @plt entries
+typedef struct _patchEntry {
+    // library's @plt location
+    void** _location;
+    // original function
+    void*  _func;
+} PatchEntry;
+
+static PatchEntry* patched_entries = nullptr;
+static volatile int num_of_entries = 0;
+
+typedef struct _startRoutineArg {
+    func_start_routine _func;
+    void*              _arg;
+} StartRoutineArg;
+
+static void* start_routine_wrapper(void* args) {
+    StartRoutineArg* data = (StartRoutineArg*)args;
+    ProfiledThread::initCurrentThread();
+    int tid = ProfiledThread::currentTid();
+    Profiler::registerThread(tid);
+    void* result = data->_func(data->_arg);
+    Profiler::unregisterThread(tid);
+    ProfiledThread::release();
+    free(args);
+    return result;
 }
 
-static void **_pthread_entry = NULL;
+static int pthread_create_hook(pthread_t* thread,
+                          const pthread_attr_t* attr,
+                          func_start_routine start_routine,
+                          void* arg) {
+  StartRoutineArg* data = (StartRoutineArg*)malloc(sizeof(StartRoutineArg));
+  data->_func = start_routine;
+  data->_arg = arg;
+  return pthread_create(thread, attr, start_routine_wrapper, (void*)data);
+}
+
+static Error patch_libraries_for_hotspot_or_zing() {
+   Dl_info info;
+   void* caller_address = __builtin_return_address(0); // Get return address of caller
+
+   if (!dladdr(caller_address, &info)) {
+      return Error("Cannot resolve current library name");
+   }
+   TEST_LOG("Profiler library name: %s", info.dli_fname );
+
+  const CodeCacheArray& native_libs = Libraries::instance()->native_libs();
+  int count = 0;
+  int num_of_libs = native_libs.count();
+  size_t size = num_of_libs * sizeof(PatchEntry);
+  patched_entries = (PatchEntry*)malloc(size);
+  memset((void*)patched_entries, 0, size);
+  TEST_LOG("Patching libraries");
+
+  for (int index = 0; index < num_of_libs; index++) {
+     CodeCache* lib = native_libs.at(index);
+     // Don't patch self
+     if (strcmp(lib->name(), info.dli_fname) == 0) {
+        continue;
+     }
+
+     void** pthread_create_location = (void**)lib->findImport(im_pthread_create);
+     if (pthread_create_location != nullptr) {
+       TEST_LOG("Patching %s", lib->name());
+
+       patched_entries[count]._location = pthread_create_location;
+       patched_entries[count]._func = (void*)__atomic_load_n(pthread_create_location, __ATOMIC_RELAXED);
+        __atomic_store_n(pthread_create_location, (void*)pthread_create_hook, __ATOMIC_RELAXED);
+        count++;
+     }
+  }
+  // Publish everything, including patched entries
+  __atomic_store_n(&num_of_entries, count, __ATOMIC_SEQ_CST);
+  return Error::OK;
+}
 
 // Intercept thread creation/termination by patching libjvm's GOT entry for
 // pthread_setspecific(). HotSpot puts VMThread into TLS on thread start, and
@@ -62,21 +136,49 @@ static int pthread_setspecific_hook(pthread_key_t key, const void *value) {
   }
 }
 
-static void **lookupThreadEntry() {
-  // Depending on Zing version, pthread_setspecific is called either from
-  // libazsys.so or from libjvm.so
-  if (VM::isZing()) {
-    CodeCache *libazsys = Libraries::instance()->findLibraryByName("libazsys");
-    if (libazsys != NULL) {
-      void **entry = libazsys->findImport(im_pthread_setspecific);
-      if (entry != NULL) {
-        return entry;
-      }
-    }
+static Error patch_libraries_for_J9_or_musl() {
+   CodeCache *lib = Libraries::instance()->findJvmLibrary("libj9thr");
+   void** func_location = lib->findImport(im_pthread_setspecific);
+   if (func_location != nullptr) {
+       patched_entries = (PatchEntry*)malloc(sizeof(PatchEntry));
+       patched_entries[0]._location = func_location;
+       patched_entries[0]._func = (void*)__atomic_load_n(func_location, __ATOMIC_RELAXED);
+       __atomic_store_n(func_location, (void*)pthread_setspecific_hook, __ATOMIC_RELAXED);
+
+      // Publish everything, including patched entries
+      __atomic_store_n(&num_of_entries, 1, __ATOMIC_SEQ_CST);
   }
+  return Error::OK;
+}
 
-  CodeCache *lib = Libraries::instance()->findJvmLibrary("libj9thr");
-  return lib != NULL ? lib->findImport(im_pthread_setspecific) : NULL;
+static Error patch_libraries() {
+  if ((VM::isHotspot() || VM::isZing()) && !OS::isMusl()) {
+    return patch_libraries_for_hotspot_or_zing();
+  } else {
+    return patch_libraries_for_J9_or_musl();
+  }
+}
+
+static void unpatch_libraries() {
+  int count = __atomic_load_n(&num_of_entries, __ATOMIC_RELAXED);
+  PatchEntry* tmp = patched_entries;
+  patched_entries = nullptr;
+  __atomic_store_n(&num_of_entries, 0, __ATOMIC_SEQ_CST);
+
+  for (int index = 0; index < count; index++) {
+     if (tmp[index]._location != nullptr) {
+       __atomic_store_n(tmp[index]._location, tmp[index]._func, __ATOMIC_RELAXED);
+     }
+  }
+  __atomic_thread_fence(__ATOMIC_SEQ_CST);
+  if (tmp != nullptr) {
+    free((void*)tmp);
+  }
+}
+
+
+static inline clockid_t thread_cpu_clock(unsigned int tid) {
+  return ((~tid) << 3) | 6; // CPUCLOCK_SCHED | CPUCLOCK_PERTHREAD_MASK
 }
 
 long CTimer::_interval;
@@ -135,11 +237,6 @@ void CTimer::unregisterThread(int tid) {
 }
 
 Error CTimer::check(Arguments &args) {
-  if (_pthread_entry == NULL &&
-      (_pthread_entry = lookupThreadEntry()) == NULL) {
-    return Error("Could not set pthread hook");
-  }
-
   timer_t timer;
   if (timer_create(CLOCK_THREAD_CPUTIME_ID, NULL, &timer) < 0) {
     return Error("Failed to create CPU timer");
@@ -153,10 +250,7 @@ Error CTimer::start(Arguments &args) {
   if (args._interval < 0) {
     return Error("interval must be positive");
   }
-  if (_pthread_entry == NULL &&
-      (_pthread_entry = lookupThreadEntry()) == NULL) {
-    return Error("Could not set pthread hook");
-  }
+
   _interval = args.cpuSamplerInterval();
   _cstack = args._cstack;
   _signal = SIGPROF;
@@ -170,9 +264,10 @@ Error CTimer::start(Arguments &args) {
 
   OS::installSignalHandler(_signal, signalHandler);
 
-  // Enable pthread hook before traversing currently running threads
-  __atomic_store_n(_pthread_entry, (void *)pthread_setspecific_hook,
-                   __ATOMIC_RELEASE);
+  Error err = patch_libraries();
+  if (err) {
+    return err;
+  }
 
   // Register all existing threads
   Error result = Error::OK;
@@ -190,8 +285,7 @@ Error CTimer::start(Arguments &args) {
 }
 
 void CTimer::stop() {
-  __atomic_store_n(_pthread_entry, (void *)pthread_setspecific,
-                   __ATOMIC_RELEASE);
+  unpatch_libraries();
   for (int i = 0; i < _max_timers; i++) {
     unregisterThread(i);
   }

ddprof-lib/src/main/cpp/libraries.h

@@ -24,6 +24,10 @@ class Libraries {
     return &instance;
   }
 
+  const CodeCacheArray& native_libs() const {
+    return _native_libs;
+  }
+
   // Delete copy constructor and assignment operator to prevent copies
   Libraries(const Libraries&) = delete;
   Libraries& operator=(const Libraries&) = delete;

ddprof-lib/src/main/cpp/utils.h

@@ -16,12 +16,12 @@ inline bool is_aligned(const T* ptr, size_t alignment) noexcept {
     auto iptr = reinterpret_cast<uintptr_t>(ptr);
 
     // Check if the integer value is a multiple of the alignment
-    return (iptr & ~(alignment - 1) == 0);
+    return ((iptr & (~(alignment - 1))) == 0);
 }
 
 inline size_t align_down(size_t size, size_t alignment) noexcept {
     assert(is_power_of_2(alignment));
-    return size & ~(alignment - 1);
+    return size & (~(alignment - 1));
 }
 
 inline size_t align_up(size_t size, size_t alignment) noexcept {

ddprof-test/build.gradle

@@ -1,14 +1,62 @@
 plugins {
   id 'java'
+  id 'java-library'
   id 'application'
 }
 
 repositories {
   mavenCentral()
 }
 
+// 1. Define paths and properties
+def nativeSrcDir = file('src/test/cpp')
+def jniHeadersDir = layout.buildDirectory.dir("generated/jni-headers").get().asFile
+def outputLibDir = layout.buildDirectory.dir("libs/native").get().asFile
+// Define the name of your JNI library (e.g., "ddproftest" becomes libddproftest.so/ddproftest.dll/libddproftest.dylib)
+def libraryName = "ddproftest"
+
+// Determine OS-specific file extensions and library names
+def osName = org.gradle.internal.os.OperatingSystem.current()
+def libFileExtension = (os().isMacOsX() ? "dylib" : "so")
+def libraryFileName = "lib${libraryName}.${libFileExtension}"
+
+// 2. Generate JNI headers using javac
+tasks.named('compileJava') {
+  // Tell javac to generate the JNI headers into the specified directory
+  options.compilerArgs += ['-h', jniHeadersDir]
+}
+
+// 3. Define a task to compile the native code
+tasks.register('buildNativeJniLibrary', Exec) {
+  description 'Compiles the JNI C/C++ sources into a shared library'
+  group 'build'
+
+  // Ensure Java compilation (and thus header generation) happens first
+  dependsOn tasks.named('compileJava')
+
+  // Clean up previous build artifacts
+  doFirst {
+    outputLibDir.mkdirs()
+  }
+
+  // Assume GCC/Clang on Linux/macOS
+  commandLine 'gcc'
+  args "-I${System.getenv('JAVA_HOME')}/include" // Standard JNI includes
+  if (os().isMacOsX()) {
+    args "-I${System.getenv('JAVA_HOME')}/include/darwin" // macOS-specific includes
+    args "-dynamiclib" // Build a dynamic library on macOS
+  } else if (os().isLinux()) {
+    args "-I${System.getenv('JAVA_HOME')}/include/linux" // Linux-specific includes
+    args "-fPIC"
+    args "-shared" // Build a shared library on Linux
+  }
+  args nativeSrcDir.listFiles()*.getAbsolutePath() // Source files
+  args "-o", "${outputLibDir.absolutePath}/${libraryFileName}" // Output file path
+}
+
 apply from: rootProject.file('common.gradle')
 
+
 def addCommonTestDependencies(Configuration configuration) {
   configuration.dependencies.add(project.dependencies.create('org.junit.jupiter:junit-jupiter-api:5.9.2'))
   configuration.dependencies.add(project.dependencies.create('org.junit.jupiter:junit-jupiter-engine:5.9.2'))
@@ -217,6 +265,8 @@ task unwindingReport {
 }
 
 tasks.withType(Test).configureEach {
+  dependsOn tasks.named('buildNativeJniLibrary')
+
   // this is a shared configuration for all test tasks
   onlyIf {
     !project.hasProperty('skip-tests')
@@ -229,7 +279,7 @@ tasks.withType(Test).configureEach {
 
   jvmArgs "-Dddprof_test.keep_jfrs=${keepRecordings}", '-Djdk.attach.allowAttachSelf', '-Djol.tryWithSudo=true',
     "-Dddprof_test.config=${config}", "-Dddprof_test.ci=${project.hasProperty('CI')}", "-Dddprof.disable_unsafe=true", '-XX:ErrorFile=build/hs_err_pid%p.log', '-XX:+ResizeTLAB',
-    '-Xmx512m', '-XX:OnError=/tmp/do_stuff.sh'
+    '-Xmx512m', '-XX:OnError=/tmp/do_stuff.sh', "-Djava.library.path=${outputLibDir.absolutePath}"
 
   def javaHome = System.getenv("JAVA_TEST_HOME")
   if (javaHome == null) {
@@ -274,4 +324,4 @@ gradle.projectsEvaluated {
       testTask.dependsOn gtestTask
     }
   }
-}
+}