Make allocation profiling work after exec

include/constants.hpp

@@ -48,6 +48,10 @@ DDPROF_CONSTREXPR const char *k_profiler_lib_env_variable =
 DDPROF_CONSTREXPR const char *k_startup_wait_ms_env_variable =
     "DD_PROFILING_NATIVE_STARTUP_WAIT_MS";
 
+// Env variable to disable allocation profiling of of exec'd processes
+DDPROF_CONSTREXPR const char *k_allocation_profiling_follow_execs =
+    "DD_PROFILING_NATIVE_ALLOCATION_PROFILING_FOLLOW_EXECS";
+
 DDPROF_CONSTREXPR const char *k_libdd_profiling_name = "libdd_profiling.so";
 
 DDPROF_CONSTREXPR const char *k_libdd_profiling_embedded_name =

include/daemonize.hpp

@@ -5,6 +5,8 @@
 
 #pragma once
 
+#include "unique_fd.hpp"
+
 #include <sys/types.h>
 
 namespace ddprof {
@@ -16,6 +18,8 @@ struct DaemonizeResult {
       temp_pid; // -1 on failure, 0 for initial process, > 0 for daemon process
   pid_t parent_pid; // pid of process initiating daemonize
   pid_t daemon_pid; // pid of daemon process
+  UniqueFd
+      pipe_fd; // pipe to communicate from daemon process to initial process
 };
 
 // Daemonization function

include/ddprof_cli.hpp

@@ -67,7 +67,8 @@ struct DDProfCLI {
   bool show_samples{false};
   bool fault_info{true};
   bool help_extended{false};
-  int socket{-1};
+  std::string socket_path;
+  int pipefd_to_library{-1};
   bool continue_exec{false};
   bool timeline{false};
 

include/ddprof_cmdline.hpp

@@ -5,26 +5,15 @@
 
 #pragma once
 
-#include <cstddef>
-#include <cstdint>
+#include <span>
+#include <string_view>
 
 namespace ddprof {
 
-/**************************** Cmdline Helpers *********************************/
-// Helper functions for processing commandline arguments.
-//
-// Note that `arg_yesno(,1)` is not the same as `!arg(,0)` or vice-versa.  This
-// is mostly because a parameter whose default value is true needs to check
-// very specifically for disablement, but the failover is to retain enable
-//
-// That said, it might be better to be more correct and only accept input of
-// the specified form, returning error otherwise.
+/// Returns index to element that matches str (case insensitive), otherwise -1
+int arg_which(std::string_view str, std::span<const std::string_view> str_set);
 
-/// Returns index to element that compars to str, otherwise -1
-int arg_which(const char *str, char const *const *set, int sz_set);
-
-bool arg_inset(const char *str, char const *const *set, int sz_set);
-
-bool arg_yesno(const char *str, int mode);
+bool arg_yes(std::string_view str);
+bool arg_no(std::string_view str);
 
 } // namespace ddprof

include/ddprof_context.hpp

@@ -28,7 +28,8 @@ struct DDProfContext {
     pid_t pid{0}; // ! only use for perf attach (can be -1 in global mode)
     uint32_t worker_period{}; // exports between worker refreshes
     int dd_profiling_fd{-1};  // opened file descriptor to our internal lib
-    ddprof::UniqueFd sockfd;
+    std::string socket_path;
+    UniqueFd pipefd_to_library;
     bool show_samples{false};
     bool timeline{false};
     cpu_set_t cpu_affinity{};
@@ -39,6 +40,7 @@ struct DDProfContext {
     std::chrono::milliseconds loaded_libs_check_interval{0};
   } params;
 
+  ddprof::UniqueFd socket_fd;
   PerfClockSource perf_clock_source{PerfClockSource::kNoClock};
   std::vector<PerfWatcher> watchers;
   ExporterInput exp_input;

include/defer.hpp

@@ -22,6 +22,12 @@ template <class F> ddprof::scope_exit<F> make_defer(F &&f) {
   return ddprof::scope_exit<F>{std::forward<F>(f)};
 }
 
+// Allows to execute a deferred operation early (before scope end)
+template <typename F> void exec_defer(ddprof::scope_exit<F> &&scope_object) {
+  const ddprof::scope_exit<F> local{std::move(scope_object)};
+  (void)local;
+}
+
 #define DEFER_(LINE) zz_defer##LINE
 #define DEFER(LINE) DEFER_(LINE)
 #define defer                                                                  \

include/ipc.hpp

@@ -11,8 +11,10 @@
 
 #include <chrono>
 #include <functional>
+#include <latch>
 #include <span>
 #include <system_error>
+#include <thread>
 
 namespace ddprof {
 
@@ -28,6 +30,8 @@ class UnixSocket {
   using socket_t = int;
 
   explicit UnixSocket(socket_t handle) noexcept : _handle(handle) {}
+  explicit UnixSocket(UniqueFd &&handle) noexcept
+      : _handle(std::move(handle)) {}
 
   void close(std::error_code &ec) noexcept;
 
@@ -65,6 +69,7 @@ struct RequestMessage {
   enum { kProfilerInfo = 0x1 };
   // request is bit mask of request flags
   uint32_t request = 0;
+  pid_t pid = -1;
 };
 
 struct RingBufferInfo {
@@ -90,27 +95,36 @@ struct ReplyMessage {
   uint32_t stack_sample_size = 0;
 };
 
-class Client {
-public:
-  explicit Client(UnixSocket &&socket,
-                  std::chrono::microseconds timeout = kDefaultSocketTimeout);
+DDRes send(const UnixSocket &socket, const RequestMessage &msg);
+DDRes send(const UnixSocket &socket, const ReplyMessage &msg);
+DDRes receive(const UnixSocket &socket, RequestMessage &msg);
+DDRes receive(const UnixSocket &socket, ReplyMessage &msg);
 
-  ReplyMessage get_profiler_info();
+UniqueFd create_server_socket(std::string_view path) noexcept;
+UniqueFd create_client_socket(std::string_view path) noexcept;
+DDRes get_profiler_info(UniqueFd &&socket, std::chrono::microseconds timeout,
+                        ReplyMessage *reply) noexcept;
 
-private:
-  UnixSocket _socket;
-};
+bool is_socket_abstract(std::string_view path) noexcept;
 
-class Server {
+class WorkerServer {
 public:
-  explicit Server(UnixSocket &&socket,
-                  std::chrono::microseconds timeout = kDefaultSocketTimeout);
-
-  using ReplyFunc = std::function<ReplyMessage(const RequestMessage &)>;
-  void waitForRequest(const ReplyFunc &func);
+  WorkerServer(const WorkerServer &) = delete;
+  WorkerServer &operator=(const WorkerServer &) = delete;
+  ~WorkerServer();
 
 private:
-  UnixSocket _socket;
+  friend WorkerServer start_worker_server(int socket, const ReplyMessage &msg);
+
+  WorkerServer(int socket, const ReplyMessage &msg);
+  void event_loop();
+
+  int _socket;
+  std::latch _latch;
+  ReplyMessage _msg;
+  std::jthread _loop_thread;
 };
 
+WorkerServer start_worker_server(int socket, const ReplyMessage &msg);
+
 } // namespace ddprof

include/prng.hpp

@@ -0,0 +1,80 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
+// Datadog, Inc.
+
+#pragma once
+
+#include <cstdint>
+#include <random>
+#include <string>
+
+namespace ddprof {
+// The "xoshiro256** 1.0" generator.
+// C++ port by Arthur O'Dwyer (2021).
+// https://quuxplusone.github.io/blog/2021/11/23/xoshiro/ Based on the C version
+// by David Blackman and Sebastiano Vigna (2018),
+// https://prng.di.unimi.it/xoshiro256starstar.c
+
+// NOLINTBEGIN(readability-magic-numbers)
+class xoshiro256ss {
+  uint64_t s[4]{};
+
+  static constexpr uint64_t splitmix64(uint64_t &x) {
+    uint64_t z = (x += 0x9e3779b97f4a7c15UL);
+    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9UL;
+    z = (z ^ (z >> 27)) * 0x94d049bb133111ebUL;
+    return z ^ (z >> 31);
+  }
+
+  static constexpr uint64_t rotl(uint64_t x, int k) {
+    return (x << k) | (x >> (64 - k));
+  }
+
+public:
+  constexpr explicit xoshiro256ss() : xoshiro256ss(0) {}
+
+  constexpr explicit xoshiro256ss(uint64_t seed) {
+    s[0] = splitmix64(seed);
+    s[1] = splitmix64(seed);
+    s[2] = splitmix64(seed);
+    s[3] = splitmix64(seed);
+  }
+
+  using result_type = uint64_t;
+  static constexpr uint64_t min() { return 0; }
+  static constexpr uint64_t max() { return static_cast<uint64_t>(-1); }
+
+  constexpr uint64_t operator()() {
+    const uint64_t result = rotl(s[1] * 5, 7) * 9;
+    const uint64_t t = s[1] << 17;
+    s[2] ^= s[0];
+    s[3] ^= s[1];
+    s[1] ^= s[2];
+    s[0] ^= s[3];
+    s[2] ^= t;
+    s[3] = rotl(s[3], 45);
+    return result;
+  }
+};
+// NOLINTEND(readability-magic-numbers)
+
+inline constexpr char charset[] = "0123456789"
+                                  "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+                                  "abcdefghijklmnopqrstuvwxyz";
+
+template <typename TRandomGenerator>
+std::string generate_random_string(TRandomGenerator &engine,
+                                   std::size_t length) {
+  // NOLINTNEXTLINE(misc-const-correctness)
+  std::uniform_int_distribution<std::size_t> distribution(0,
+                                                          sizeof(charset) - 2);
+  std::string result;
+  result.reserve(length);
+  for (std::size_t i = 0; i < length; ++i) {
+    result.push_back(charset[distribution(engine)]);
+  }
+  return result;
+}
+
+} // namespace ddprof

src/daemonize.cc

@@ -5,7 +5,9 @@
 
 #include "daemonize.hpp"
 
-#include <csignal> //
+#include "unique_fd.hpp"
+
+#include <csignal>
 #include <cstdlib>
 #include <fcntl.h>
 #include <sys/types.h>
@@ -23,23 +25,26 @@ DaemonizeResult daemonize_error() {
 
 DaemonizeResult daemonize() {
   int pipefd[2];
-  if (pipe2(pipefd, O_CLOEXEC) == -1) {
+  if (pipe2(pipefd, 0) == -1) {
     return daemonize_error();
   }
 
+  UniqueFd readfd{pipefd[0]};
+  UniqueFd writefd{pipefd[1]};
+
   const pid_t parent_pid = getpid();
-  pid_t temp_pid = fork(); // "middle" (temporary) PID
+  pid_t temp_pid = fork();
 
   if (temp_pid == -1) {
     return daemonize_error();
   }
 
-  if (temp_pid == 0) { // If I'm the temp PID enter branch
-    close(pipefd[0]);
+  if (temp_pid == 0) { // Intermediate (temporary) process
+    readfd.reset();
 
     temp_pid = getpid();
     pid_t child_pid = fork();
-    if (child_pid != 0) { // If I'm the temp PID again, enter branch
+    if (child_pid != 0) { // Intermediate (temporary) process again
 
       struct sigaction sa;
       if (sigemptyset(&sa.sa_mask) == -1) {
@@ -60,28 +65,29 @@ DaemonizeResult daemonize() {
               child_pid};
     }
 
+    // Daemon process
     child_pid = getpid();
-    if (write(pipefd[1], &child_pid, sizeof(child_pid)) != sizeof(child_pid)) {
+    if (write(writefd.get(), &child_pid, sizeof(child_pid)) !=
+        sizeof(child_pid)) {
       exit(1);
     }
-    close(pipefd[1]);
-    // If I'm the child PID, then leave and attach profiler
-    return {DaemonizeResult::DaemonProcess, temp_pid, parent_pid, child_pid};
+    return {DaemonizeResult::DaemonProcess, temp_pid, parent_pid, child_pid,
+            std::move(writefd)};
   }
 
-  close(pipefd[1]);
+  // Initial process
+  writefd.reset();
 
   pid_t grandchild_pid;
-  if (read(pipefd[0], &grandchild_pid, sizeof(grandchild_pid)) !=
+  if (read(readfd.get(), &grandchild_pid, sizeof(grandchild_pid)) !=
       sizeof(grandchild_pid)) {
     return daemonize_error();
   }
 
-  // If I'm the target PID, then now it's time to wait until my
-  // child, the middle PID, returns.
+  // Wait until intermediate process terminates
   waitpid(temp_pid, nullptr, 0);
-  return {DaemonizeResult::InitialProcess, temp_pid, parent_pid,
-          grandchild_pid};
+  return {DaemonizeResult::InitialProcess, temp_pid, parent_pid, grandchild_pid,
+          std::move(readfd)};
 }
 
-} // namespace ddprof
+} // namespace ddprof

src/ddprof_cli.cc

@@ -297,10 +297,14 @@ int DDProfCLI::parse(int argc, const char *argv[]) {
       app.add_flag("--help_extended", help_extended, "Show extended options")
           ->group(""));
   extended_options.push_back(
-      app.add_option("--socket", socket,
-                     "Profiler's IPC socket, as a file descriptor")
+      app.add_option("--socket", socket_path, "Profiler socket path")
           ->envname("DD_PROFILING_NATIVE_SOCKET")
           ->group(""));
+  extended_options.push_back(
+      app.add_option("--pipefd", pipefd_to_library,
+                     "Pipe file descriptor to communicate with library that "
+                     "spawned the profiler")
+          ->group(""));
   extended_options.push_back(
       app.add_option("--stack_sample_size", default_stack_sample_size,
                      "Sample size for the user's stack."
@@ -417,7 +421,13 @@ void DDProfCLI::print() const {
     PRINT_NFO("  - global: %s", global ? "true" : "false");
   }
   if (!command_line.empty()) {
-    PRINT_NFO("  - command line(wrapper mode): %s", command_line[0].c_str());
+    std::string command_line_str = "[" + command_line[0];
+    std::for_each(std::next(command_line.begin()), command_line.end(),
+                  [&command_line_str](const std::string &el) {
+                    command_line_str += ", " + el;
+                  });
+    command_line_str += "]";
+    PRINT_NFO("  - command line(wrapper mode): %s", command_line_str.c_str());
   }
   PRINT_NFO("  - upload_period: %lds",
             std::chrono::seconds{upload_period}.count());