feat: add i/o native interposition

Author: kaahos

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

@@ -310,13 +310,21 @@ void CodeCache::saveImport(ImportId id, void** entry) {
 void CodeCache::addImport(void **entry, const char *name) {
     switch (name[0]) {
       case 'a':
-          if (strcmp(name, "aligned_alloc") == 0) {
+          if (strcmp(name, "accept") == 0) {
+              saveImport(im_accept, entry);
+          } else if (strcmp(name, "accept4") == 0) {
+              saveImport(im_accept4, entry);
+          } else if (strcmp(name, "aligned_alloc") == 0) {
               saveImport(im_aligned_alloc, entry);
           }
           break;
       case 'c':
           if (strcmp(name, "calloc") == 0) {
               saveImport(im_calloc, entry);
+          } else if (strcmp(name, "close") == 0) {
+              saveImport(im_close, entry);
+          } else if (strcmp(name, "connect") == 0) {
+              saveImport(im_connect, entry);
           }
           break;
       case 'd':
@@ -329,6 +337,13 @@ void CodeCache::addImport(void **entry, const char *name) {
               saveImport(im_free, entry);
           }
           break;
+      case 'e':
+          if (strcmp(name, "epoll_wait") == 0) {
+              saveImport(im_epoll_wait, entry);
+          } else if (strcmp(name, "epoll_pwait") == 0) {
+              saveImport(im_epoll_pwait, entry);
+          }
+          break;
       case 'm':
           if (strcmp(name, "malloc") == 0) {
               saveImport(im_malloc, entry);
@@ -343,18 +358,39 @@ void CodeCache::addImport(void **entry, const char *name) {
               saveImport(im_pthread_setspecific, entry);
           } else if (strcmp(name, "poll") == 0) {
               saveImport(im_poll, entry);
+          } else if (strcmp(name, "ppoll") == 0) {
+              saveImport(im_ppoll, entry);
+          } else if (strcmp(name, "pselect") == 0) {
+              saveImport(im_pselect, entry);
           } else if (strcmp(name, "posix_memalign") == 0) {
               saveImport(im_posix_memalign, entry);
           }
           break;
       case 'r':
           if (strcmp(name, "realloc") == 0) {
               saveImport(im_realloc, entry);
+          } else if (strcmp(name, "recv") == 0) {
+              saveImport(im_recv, entry);
+          } else if (strcmp(name, "recvfrom") == 0) {
+              saveImport(im_recvfrom, entry);
+          } else if (strcmp(name, "recvmsg") == 0) {
+              saveImport(im_recvmsg, entry);
+          } else if (strcmp(name, "read") == 0) {
+              saveImport(im_read, entry);
           }
           break;
       case 's':
           if (strcmp(name, "sigaction") == 0) {
               saveImport(im_sigaction, entry);
+          } else if (strcmp(name, "send") == 0) {
+              saveImport(im_send, entry);
+          } else if (strcmp(name, "select") == 0) {
+              saveImport(im_select, entry);
+          }
+          break;
+      case 'w':
+          if (strcmp(name, "write") == 0) {
+              saveImport(im_write, entry);
           }
           break;
     }
@@ -456,4 +492,3 @@ void CodeCache::setBuildId(const char* build_id, size_t build_id_len) {
     }
   }
 }
-

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

@@ -38,6 +38,21 @@ enum ImportId {
   im_posix_memalign,
   im_aligned_alloc,
   im_sigaction,
+  im_send,
+  im_recv,
+  im_write,
+  im_read,
+  im_close,
+  im_connect,
+  im_accept,
+  im_accept4,
+  im_recvfrom,
+  im_recvmsg,
+  im_epoll_wait,
+  im_epoll_pwait,
+  im_ppoll,
+  im_select,
+  im_pselect,
   NUM_IMPORTS
 };
 

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

@@ -1408,7 +1408,7 @@ void Recording::writeFrameTypes(Buffer *buf) {
 
 void Recording::writeThreadStates(Buffer *buf) {
   buf->putVar64(T_THREAD_STATE);
-  buf->put8(10);
+  buf->put8(11);
   buf->put8(static_cast<int>(OSThreadState::UNKNOWN));
   buf->putUtf8("UNKNOWN");
   buf->put8(static_cast<int>(OSThreadState::NEW));
@@ -1429,6 +1429,8 @@ void Recording::writeThreadStates(Buffer *buf) {
   buf->putUtf8("TERMINATED");
   buf->put8(static_cast<int>(OSThreadState::SYSCALL));
   buf->putUtf8("SYSCALL");
+  buf->put8(static_cast<int>(OSThreadState::IO_WAIT));
+  buf->putUtf8("IO_WAIT");
   flushIfNeeded(buf);
 }
 

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

@@ -376,6 +376,7 @@ static OSThreadState decodeTaskBlockObservedState(jint state) {
     case OSThreadState::OBJECT_WAIT:
     case OSThreadState::SLEEPING:
     case OSThreadState::SYSCALL:
+    case OSThreadState::IO_WAIT:
       return static_cast<OSThreadState>(state);
     default:
       return OSThreadState::UNKNOWN;

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

@@ -60,6 +60,7 @@ void Libraries::refresh() {
   _dirty.store(false, std::memory_order_release);
   updateSymbols(false);
   LibraryPatcher::patch_sigaction();
+  LibraryPatcher::install_socket_hooks();
   MallocTracer::installHooks();
   if (_remote_symbolication) {
     updateBuildIds();

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

@@ -4,6 +4,8 @@
 #include "codeCache.h"
 #include "spinLock.h"
 
+#include <atomic>
+
 #ifdef __linux__
 
 // Patch libraries' @plt entries
@@ -15,6 +17,7 @@ typedef struct _patchEntry {
   void*  _func;
 } PatchEntry;
 
+const int MAX_NATIVE_IO_HOOKS = 16;
 
 class LibraryPatcher {
 private:
@@ -28,15 +31,28 @@ class LibraryPatcher {
   static PatchEntry  _sigaction_entries[MAX_NATIVE_LIBS];
   static int         _sigaction_size;
 
+  // Separate tracking for native I/O patches. Each library can contribute one
+  // GOT slot per supported I/O hook.
+  static PatchEntry  _socket_entries[MAX_NATIVE_IO_HOOKS * MAX_NATIVE_LIBS];
+  static int         _socket_size;
+
   static void patch_library_unlocked(CodeCache* lib);
   static void patch_pthread_create();
   static void patch_pthread_setspecific();
   static void patch_sigaction_in_library(CodeCache* lib);
 public:
+  static std::atomic<bool> _socket_active;
   static void initialize();
   static void patch_libraries();
   static void unpatch_libraries();
   static void patch_sigaction();
+  static bool patch_socket_functions();
+  static void unpatch_socket_functions();
+  static inline void install_socket_hooks() {
+    if (_socket_active.load(std::memory_order_acquire)) {
+      patch_socket_functions();
+    }
+  }
 };
 
 #else
@@ -47,8 +63,11 @@ class LibraryPatcher {
   static void patch_libraries() { }
   static void unpatch_libraries() { }
   static void patch_sigaction() { }
+  static bool patch_socket_functions() { return false; }
+  static void unpatch_socket_functions() { }
+  static void install_socket_hooks() { }
 };
 
 #endif
 
-#endif // _LIBRARYPATCHER_H
+#endif // _LIBRARYPATCHER_H

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

@@ -9,6 +9,7 @@
 #include "counters.h"
 #include "profiler.h"
 #include "guards.h"
+#include "nativeSocketInterposer.h"
 
 #include <cassert>
 #include <dlfcn.h>
@@ -24,6 +25,9 @@ PatchEntry LibraryPatcher::_patched_entries[MAX_NATIVE_LIBS];
 int        LibraryPatcher::_size = 0;
 PatchEntry LibraryPatcher::_sigaction_entries[MAX_NATIVE_LIBS];
 int        LibraryPatcher::_sigaction_size = 0;
+PatchEntry LibraryPatcher::_socket_entries[MAX_NATIVE_IO_HOOKS * MAX_NATIVE_LIBS];
+int        LibraryPatcher::_socket_size = 0;
+std::atomic<bool> LibraryPatcher::_socket_active{false};
 
 void LibraryPatcher::initialize() {
   if (_profiler_name == nullptr) {
@@ -382,4 +386,114 @@ void LibraryPatcher::patch_sigaction() {
   }
 }
 
+bool LibraryPatcher::patch_socket_functions() {
+  static void* cached_originals[NativeSocketInterposer::NUM_NATIVE_IO_HOOKS] = {};
+  static bool cached = false;
+
+  const NativeSocketInterposer::NativeIoHookSpec* hooks =
+      NativeSocketInterposer::hookSpecs();
+
+  if (!cached) {
+    for (int hook_index = 0; hook_index < NativeSocketInterposer::NUM_NATIVE_IO_HOOKS;
+         hook_index++) {
+      void* original = dlsym(RTLD_NEXT, hooks[hook_index].name);
+      if (original == nullptr) {
+        original = dlsym(RTLD_DEFAULT, hooks[hook_index].name);
+      }
+      if (original == hooks[hook_index].hook) {
+        original = nullptr;
+      }
+      cached_originals[hook_index] = original;
+    }
+    cached = true;
+  }
+
+  bool any_original = false;
+  for (int hook_index = 0; hook_index < NativeSocketInterposer::NUM_NATIVE_IO_HOOKS;
+       hook_index++) {
+    any_original |= cached_originals[hook_index] != nullptr;
+  }
+  if (!any_original) {
+    return false;
+  }
+
+  const CodeCacheArray& native_libs = Libraries::instance()->native_libs();
+  int num_of_libs = native_libs.count();
+  int capped = num_of_libs <= MAX_NATIVE_LIBS ? num_of_libs : MAX_NATIVE_LIBS;
+  bool is_self[MAX_NATIVE_LIBS];
+  for (int index = 0; index < capped; index++) {
+    CodeCache* lib = native_libs.at(index);
+    is_self[index] = false;
+    if (lib == nullptr || lib->name() == nullptr) {
+      continue;
+    }
+    char path[PATH_MAX];
+    char* resolved_path = realpath(lib->name(), path);
+    is_self[index] = _profiler_name != nullptr && resolved_path != nullptr &&
+                     strcmp(resolved_path, _profiler_name) == 0;
+  }
+
+  ExclusiveLockGuard locker(&_lock);
+  if (_socket_size == 0) {
+    for (int hook_index = 0; hook_index < NativeSocketInterposer::NUM_NATIVE_IO_HOOKS;
+         hook_index++) {
+      if (cached_originals[hook_index] != nullptr) {
+        NativeSocketInterposer::setOriginalFunction(hook_index,
+                                                    cached_originals[hook_index]);
+      }
+    }
+  }
+
+  auto try_patch_slot = [](void** location, void* hook, const char* name,
+                           CodeCache* lib) {
+    if (location == nullptr) {
+      return;
+    }
+    for (int index = 0; index < _socket_size; index++) {
+      if (_socket_entries[index]._location == location) {
+        return;
+      }
+    }
+    if (_socket_size >= MAX_NATIVE_IO_HOOKS * MAX_NATIVE_LIBS) {
+      Log::warn("socket I/O patch table full, skipping %s in %s", name,
+                lib != nullptr ? lib->name() : "?");
+      return;
+    }
+    _socket_entries[_socket_size]._lib = lib;
+    _socket_entries[_socket_size]._location = location;
+    _socket_entries[_socket_size]._func =
+        reinterpret_cast<void*>(__atomic_load_n(location, __ATOMIC_ACQUIRE));
+    __atomic_store_n(location, hook, __ATOMIC_RELEASE);
+    _socket_size++;
+  };
+
+  for (int index = 0; index < capped; index++) {
+    CodeCache* lib = native_libs.at(index);
+    if (lib == nullptr || lib->name() == nullptr || is_self[index]) {
+      continue;
+    }
+    for (int hook_index = 0; hook_index < NativeSocketInterposer::NUM_NATIVE_IO_HOOKS;
+         hook_index++) {
+      if (cached_originals[hook_index] == nullptr) {
+        continue;
+      }
+      try_patch_slot(reinterpret_cast<void**>(lib->findImport(hooks[hook_index].import_id)),
+                     hooks[hook_index].hook, hooks[hook_index].name, lib);
+    }
+  }
+
+  _socket_active.store(true, std::memory_order_release);
+  return true;
+}
+
+void LibraryPatcher::unpatch_socket_functions() {
+  ExclusiveLockGuard locker(&_lock);
+  for (int index = 0; index < _socket_size; index++) {
+    __atomic_store_n(_socket_entries[index]._location, _socket_entries[index]._func,
+                     __ATOMIC_RELEASE);
+  }
+  _socket_active.store(false, std::memory_order_release);
+  _socket_size = 0;
+}
+
 #endif // __linux__