Add mremap and shmat/shmdt instrumentation

include/common_symbol_errors.hpp

@@ -14,10 +14,14 @@ using namespace std::string_view_literals;
 
 namespace ddprof {
 
-inline constexpr std::array<std::string_view, 6> k_common_frame_names = {
-    "[truncated]"sv,      "[unknown mapping]"sv,
-    "[unwind failure]"sv, "[incomplete]"sv,
-    "[lost]"sv,           "[maximum pids]"sv};
+inline constexpr std::array<std::string_view, 7> k_common_frame_names = {
+    "[truncated]"sv,
+    "[unknown mapping]"sv,
+    "[unwind failure]"sv,
+    "[incomplete]"sv,
+    "[lost]"sv,
+    "[maximum pids]"sv,
+    "[live-alloc cleared]"sv};
 
 enum SymbolErrors : std::uint8_t {
   truncated_stack = 0,
@@ -26,6 +30,7 @@ enum SymbolErrors : std::uint8_t {
   incomplete_stack,
   lost_event,
   max_pids,
+  live_alloc_cleared,
 };
 
 } // namespace ddprof

include/ddprof_stats.hpp

@@ -46,7 +46,11 @@ namespace ddprof {
   X(PPROF_SIZE, "pprof.size", STAT_GAUGE)                                      \
   X(PROFILE_DURATION, "profile.duration_ms", STAT_GAUGE)                       \
   X(AGGREGATION_AVG_TIME, "aggregation.avg_time_ns", STAT_GAUGE)               \
-  X(BACKPOPULATE_COUNT, "backpopulate.count", STAT_GAUGE)
+  X(BACKPOPULATE_COUNT, "backpopulate.count", STAT_GAUGE)                      \
+  X(LIVE_ALLOC_SNAPSHOT_BYTES, "live_alloc.snapshot.bytes", STAT_GAUGE)        \
+  X(LIVE_ALLOC_CLEARED_STACKS, "live_alloc.snapshot.cleared_stacks",           \
+    STAT_GAUGE)                                                                \
+  X(LIVE_ALLOC_DROPPED_PIDS, "live_alloc.snapshot.dropped_pids", STAT_GAUGE)
 
 // Expand the enum/index for the individual stats
 enum DDPROF_STATS : uint8_t { STATS_TABLE(X_ENUM) STATS_LEN };

include/live_allocation.hpp

@@ -11,6 +11,8 @@
 #include "unwind_output_hash.hpp"
 
 #include <cstddef>
+#include <deque>
+#include <string>
 #include <sys/types.h>
 #include <unordered_map>
 
@@ -55,6 +57,23 @@ class LiveAllocation {
   // NOLINTNEXTLINE(misc-non-private-member-variables-in-classes)
   WatcherVector _watcher_vector;
 
+  // Owns the string storage backing string_views inside UnwindOutputs that
+  // were restored from a snapshot. Live entries created from incoming
+  // allocation events continue to use views into Process / base-frame
+  // tables; this storage is only used by snapshot-restore.
+  // NOLINTNEXTLINE(misc-non-private-member-variables-in-classes)
+  std::deque<std::string> _restored_strings;
+
+  // Returns a string_view backed by _restored_strings. Empty input maps
+  // to an empty view without allocation.
+  std::string_view intern_restored_string(std::string_view sv) {
+    if (sv.empty()) {
+      return {};
+    }
+    _restored_strings.emplace_back(sv);
+    return _restored_strings.back();
+  }
+
   void register_library_state(int watcher_pos, pid_t pid,
                               uint32_t address_conflict_count,
                               uint32_t tracked_address_count,

include/live_allocation_snapshot.hpp

@@ -0,0 +1,129 @@
+// 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.
+
+// Live-allocation snapshot: portable serialisation of LiveAllocation across
+// worker resets.
+//
+// The aggregator state in LiveAllocation refers to per-worker handles
+// (libdatadog dictionary pointers, indices into SymbolHdr tables, string_views
+// into Process/base-frame caches). None of these survive a worker fork, so
+// before a worker restart we capture a fully self-owned snapshot via
+// `capture_snapshot`, serialise it through a memfd held by the parent, and
+// re-intern everything into the freshly-built tables of the new worker via
+// `restore_snapshot`.
+
+#pragma once
+
+#include "ddprof_defs.hpp"
+#include "live_allocation.hpp"
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <sys/types.h>
+#include <vector>
+
+namespace ddprof {
+
+struct SymbolHdr;
+
+namespace live_alloc_snapshot {
+
+// Fully self-owned mirror of one FunLoc.
+struct FunLocPortable {
+  ProcessAddress_t ip{};
+  ElfAddress_t elf_addr{};
+  uint32_t lineno{};
+  std::string fn_name;
+  std::string fn_system_name;
+  std::string fn_file;
+  ElfAddress_t map_low{};
+  ElfAddress_t map_high{};
+  Offset_t map_offset{};
+  std::string map_filename;
+  std::string map_build_id;
+};
+
+// Fully self-owned mirror of one UnwindOutput.
+struct UnwindOutputPortable {
+  std::vector<FunLocPortable> locs;
+  std::string container_id;
+  std::string exe_name;
+  std::string thread_name;
+  int pid{};
+  int tid{};
+};
+
+struct AddressEntry {
+  uintptr_t addr{};
+  int64_t value{};
+  uint32_t stack_idx{}; // index into Snapshot::stacks
+};
+
+struct PidEntry {
+  int watcher_pos{};
+  pid_t pid{};
+  uint32_t address_conflict_count{};
+  uint32_t tracked_address_count{};
+  std::vector<AddressEntry> addresses;
+};
+
+// Special stack index meaning "stack was dropped during budget enforcement".
+// On restore this is materialised as a synthetic single-frame stack pointing
+// at the [live-alloc cleared] common frame, ensuring per-PID heap totals are
+// preserved even when detail is shed.
+inline constexpr uint32_t k_cleared_stack_idx =
+    std::numeric_limits<uint32_t>::max();
+
+struct Snapshot {
+  std::vector<UnwindOutputPortable> stacks;
+  std::vector<PidEntry> pids;
+  // How many distinct allocation addresses had their stack remapped to the
+  // synthetic cleared stack because of the size budget.
+  uint32_t cleared_addresses{};
+  // How many (watcher_pos, pid) entries we dropped entirely because even
+  // minimal accounting did not fit in the budget.
+  uint32_t dropped_pids{};
+};
+
+// Default and hard-ceiling sizes for the serialised snapshot.
+inline constexpr std::size_t k_default_max_snapshot_bytes = 4UL * 1024 * 1024;
+inline constexpr std::size_t k_hard_max_snapshot_bytes = 20UL * 1024 * 1024;
+
+// Read a Snapshot out of an in-memory LiveAllocation, by resolving symbol /
+// mapping IDs through `symbol_hdr`. The result has no references to the
+// originating worker.
+//
+// If `max_bytes` is non-zero and the projected serialised size exceeds it,
+// the snapshot is degraded in a value-preserving way: low-value stacks are
+// remapped to the synthetic cleared stack first, and finally whole pids are
+// dropped from the lowest aggregate value upwards. The `cleared_addresses` /
+// `dropped_pids` counters in the result report how aggressive the degradation
+// had to be.
+Snapshot capture_snapshot(const LiveAllocation &live_alloc,
+                          const SymbolHdr &symbol_hdr,
+                          std::size_t max_bytes = k_default_max_snapshot_bytes);
+
+// Serialise / deserialise a Snapshot to/from a binary blob.
+void serialize(const Snapshot &snapshot, std::vector<uint8_t> &out);
+bool deserialize(const uint8_t *data, std::size_t size, Snapshot &out);
+
+// Write the binary blob into `fd` (memfd). Truncates fd to the blob size on
+// success. On error, returns false and leaves fd in an unspecified state.
+bool write_to_fd(int fd, const Snapshot &snapshot);
+
+// Read a snapshot blob out of `fd`. Returns false if fd is empty / unreadable
+// / malformed.
+bool read_from_fd(int fd, Snapshot &out);
+
+// Re-intern a snapshot into `symbol_hdr` and populate the empty `live_alloc`
+// state. Restored UnwindOutputs hold string_views into
+// LiveAllocation::_restored_strings.
+void restore_snapshot(const Snapshot &snapshot, LiveAllocation &live_alloc,
+                      SymbolHdr &symbol_hdr);
+
+} // namespace live_alloc_snapshot
+
+} // namespace ddprof

include/persistent_worker_state.hpp

@@ -15,6 +15,13 @@ struct PersistentWorkerState {
   // Why not volatile ? Although several threads can update the number of
   // cycles, by design Only a single thread reads and writes to this variable.
   uint32_t profile_seq;
+  // memfd holding the most recent serialized live-allocation snapshot.
+  // The fd is opened by the parent in main_loop and inherited by every
+  // worker child. A child that is restarting writes its serialized
+  // LiveAllocation state to this fd just before exiting; the next child
+  // reads it during worker_library_init so live-heap tracking survives
+  // worker resets. -1 if snapshotting is disabled.
+  int live_alloc_snapshot_fd;
 };
 
 } // namespace ddprof

src/ddprof_worker.cc

@@ -6,12 +6,15 @@
 #include "ddprof_worker.hpp"
 
 #include "ddprof_context.hpp"
+#include "ddprof_context_lib.hpp"
 #include "ddprof_perf_event.hpp"
 #include "ddprof_stats.hpp"
 #include "dso_hdr.hpp"
 #include "exporter/ddprof_exporter.hpp"
+#include "live_allocation_snapshot.hpp"
 #include "logger.hpp"
 #include "perf.hpp"
+#include "persistent_worker_state.hpp"
 #include "pevent_lib.hpp"
 #include "pprof/ddprof_pprof.hpp"
 #include "procutils.hpp"
@@ -561,6 +564,27 @@ DDRes worker_library_init(DDProfContext &ctx,
     ctx.worker_ctx.exp[1] = nullptr;
     ctx.worker_ctx.pprof[0] = nullptr;
     ctx.worker_ctx.pprof[1] = nullptr;
+
+    // If the previous worker handed us a live-allocation snapshot, replay
+    // it before the poll loop starts draining new events. Restored entries
+    // use string storage owned by LiveAllocation itself.
+    if (persistent_worker_state &&
+        persistent_worker_state->live_alloc_snapshot_fd >= 0 &&
+        context_allocation_profiling_watcher_idx(ctx) != -1) {
+      live_alloc_snapshot::Snapshot snap;
+      if (live_alloc_snapshot::read_from_fd(
+              persistent_worker_state->live_alloc_snapshot_fd, snap)) {
+        if (!snap.stacks.empty() || !snap.pids.empty()) {
+          live_alloc_snapshot::restore_snapshot(snap,
+                                                ctx.worker_ctx.live_allocation,
+                                                ctx.worker_ctx.us->symbol_hdr);
+          LG_NTC("[live-alloc] Snapshot restored: stacks=%zu pids=%zu "
+                 "cleared=%u dropped_pids=%u",
+                 snap.stacks.size(), snap.pids.size(), snap.cleared_addresses,
+                 snap.dropped_pids);
+        }
+      }
+    }
   }
   CatchExcept2DDRes();
   return {};

src/lib/symbol_overrides.cc

@@ -22,7 +22,9 @@
 #include <ctime>
 #include <dlfcn.h>
 #include <malloc.h>
+#include <sys/ipc.h>
 #include <sys/mman.h>
+#include <sys/shm.h>
 
 #if defined(__GNUC__) && !defined(__clang__)
 #  define NOEXCEPT noexcept
@@ -35,6 +37,14 @@ extern "C" {
 // Declaration of reallocarray is only available starting from glibc 2.28
 DDPROF_WEAK void *reallocarray(void *ptr, size_t nmemb, size_t nmenb) NOEXCEPT;
 DDPROF_WEAK void *pvalloc(size_t size) NOEXCEPT;
+DDPROF_WEAK void *mremap(void *old_address, size_t old_size, size_t new_size,
+                         int flags, ...) NOEXCEPT;
+DDPROF_WEAK void *shmat(int shmid, const void *shmaddr, int shmflg) NOEXCEPT;
+DDPROF_WEAK int shmdt(const void *shmaddr) NOEXCEPT;
+// brk()/sbrk() are not instrumented:
+// - They grow/shrink the program break, not individual allocations
+// - Would double-count if malloc uses them internally
+// - Tracking would require maintaining state of what's actually allocated
 DDPROF_WEAK int __libc_allocate_rtsig(int high) NOEXCEPT;
 // NOLINTEND
 
@@ -857,6 +867,57 @@ struct Munmap_Hook : HookBase {
   }
 };
 
+struct MremapHook : HookBase {
+  static constexpr auto name = "mremap";
+  using FuncType = decltype(&::mremap);
+  static inline FuncType ref{};
+
+  static void *hook(void *old_address, size_t old_size, size_t new_size,
+                    int flags) noexcept {
+    AllocTrackerHelperMmap helper;
+    if (likely(old_address) && helper) {
+      ddprof::AllocationTracker::track_deallocation_s(
+          reinterpret_cast<uintptr_t>(old_address), *helper.tl_state(), true);
+    }
+    auto *new_ptr = ref(old_address, old_size, new_size, flags);
+    if (likely(new_size && new_ptr != MAP_FAILED)) {
+      helper.track(new_ptr, new_size);
+    }
+    return new_ptr;
+  }
+};
+
+struct ShmatHook : HookBase {
+  static constexpr auto name = "shmat";
+  using FuncType = decltype(&::shmat);
+  static inline FuncType ref{};
+
+  static void *hook(int shmid, const void *shmaddr, int shmflg) noexcept {
+    void *ptr = ref(shmid, shmaddr, shmflg);
+    if (ptr != reinterpret_cast<void *>(-1)) {
+      AllocTrackerHelperMmap helper;
+      // Query the segment size via shmctl
+      struct shmid_ds buf;
+      if (shmctl(shmid, IPC_STAT, &buf) == 0) {
+        helper.track(ptr, buf.shm_segsz);
+      }
+    }
+    return ptr;
+  }
+};
+
+struct ShmdtHook : HookBase {
+  static constexpr auto name = "shmdt";
+  using FuncType = decltype(&::shmdt);
+  static inline FuncType ref{};
+
+  static int hook(const void *shmaddr) noexcept {
+    DeallocTrackerHelperMmap helper;
+    helper.track(const_cast<void *>(shmaddr));
+    return ref(shmaddr);
+  }
+};
+
 template <typename T> void register_hook() {
   g_symbol_overrides->register_override(T::name,
                                         reinterpret_cast<uintptr_t>(&T::hook),
@@ -922,6 +983,9 @@ void register_hooks() {
   register_hook<MunmapHook>();
   register_hook<Mmap_Hook>();
   register_hook<Munmap_Hook>();
+  register_hook<MremapHook>();
+  register_hook<ShmatHook>();
+  register_hook<ShmdtHook>();
 
   register_hook<MallocxHook>();
   register_hook<RallocxHook>();