Add USDT output for native allocation profiling

ddprof-lib/build.gradle.kts

@@ -17,6 +17,7 @@
 import com.datadoghq.native.model.Platform
 import com.datadoghq.native.util.PlatformUtils
 import org.gradle.api.publish.maven.tasks.AbstractPublishToMaven
+import org.gradle.api.tasks.Exec
 import org.gradle.api.tasks.VerificationTask
 
 plugins {
@@ -173,6 +174,27 @@ afterEvaluate {
       archiveVersion.set(componentVersion)
     }
 
+    if (name == "release") {
+      tasks.register("verifyNativeAllocationUsdtRelease", Exec::class) {
+        group = "verification"
+        description = "Verify the packaged Linux release USDT allocation probe"
+        dependsOn(assembleJarTask)
+
+        val verifier = rootProject.layout.projectDirectory.file("utils/verify-native-allocation-usdt.sh")
+        val releaseJar = assembleJarTask.flatMap { it.archiveFile }
+        inputs.file(verifier)
+        inputs.file(releaseJar)
+
+        onlyIf {
+          PlatformUtils.currentPlatform == Platform.LINUX &&
+            !PlatformUtils.isNativeSkipped(project)
+        }
+        doFirst {
+          commandLine("bash", verifier.asFile.absolutePath, releaseJar.get().asFile.absolutePath)
+        }
+      }
+    }
+
     // Create consumable configuration for inter-project dependencies
     // This allows other projects to depend on specific build configurations
     configurations.create(name) {

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

@@ -425,6 +425,20 @@ Error Arguments::parse(const char *args) {
         msg = "nativemem must be >= 0";
       }
 
+      CASE("nativememoutput")
+      _nativemem_output_set = true;
+      if (value == NULL || *value == '\0') {
+        msg = "nativememoutput requires jfr, usdt, or both";
+      } else if (strcasecmp(value, "jfr") == 0) {
+        _nativemem_output = NATIVE_MEMORY_JFR;
+      } else if (strcasecmp(value, "usdt") == 0) {
+        _nativemem_output = NATIVE_MEMORY_USDT;
+      } else if (strcasecmp(value, "both") == 0) {
+        _nativemem_output = NATIVE_MEMORY_BOTH;
+      } else {
+        msg = "Invalid nativememoutput value";
+      }
+
       CASE("natsock")
       if (value != NULL) {
         _nativesocket_interval = parseUnits(value, NANOS);

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

@@ -112,6 +112,19 @@ enum Ring {
   RING_ANY = RING_KERNEL | RING_USER,
 };
 
+enum NativeMemoryOutput {
+  NATIVE_MEMORY_JFR,
+  NATIVE_MEMORY_USDT,
+  NATIVE_MEMORY_BOTH,
+  NATIVE_MEMORY_NOOP
+};
+
+#if defined(__linux__)
+const NativeMemoryOutput DEFAULT_NATIVE_MEMORY_OUTPUT = NATIVE_MEMORY_BOTH;
+#else
+const NativeMemoryOutput DEFAULT_NATIVE_MEMORY_OUTPUT = NATIVE_MEMORY_JFR;
+#endif
+
 enum Style {
   STYLE_SIMPLE = 1,
   STYLE_DOTTED = 2,
@@ -254,6 +267,8 @@ class Arguments {
   u8 _reference_chains_tuned_mask;
   // Explicit opt-in for the legacy whole-graph JVMTI FollowReferences walk.
   long _nativemem;
+  NativeMemoryOutput _nativemem_output;
+  bool _nativemem_output_set;
   int  _jstackdepth;
   int _safe_mode;
   StackWalkFeatures _features;
@@ -305,6 +320,8 @@ class Arguments {
         _reference_chains_first_pass_budget(DEFAULT_REFERENCE_CHAINS_FIRST_PASS_BUDGET),
         _reference_chains_tuned_mask(0),
         _nativemem(-1),
+        _nativemem_output(DEFAULT_NATIVE_MEMORY_OUTPUT),
+        _nativemem_output_set(false),
         _jstackdepth(DEFAULT_JSTACKDEPTH),
         _safe_mode(0),
         _features{1, 1, 1, 1, 1, 1},

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

@@ -10,10 +10,12 @@
 #include <stdlib.h>
 #include <string.h>
 #include "codeCache.h"
+#include "common.h"
 #include "counters.h"
 #include "guards.h"
 #include "libraries.h"
 #include "mallocTracer.h"
+#include "nativeAllocationUsdt.h"
 #include "os.h"
 #include "pidController.h"
 #include "profiler.h"
@@ -87,6 +89,7 @@ u64 MallocTracer::_configured_interval;
 volatile u64 MallocTracer::_sample_count;
 volatile u64 MallocTracer::_last_config_update_ts;
 volatile bool MallocTracer::_running = false;
+NativeMemoryOutput MallocTracer::_output = NATIVE_MEMORY_JFR;
 PidController MallocTracer::_pid(MallocTracer::TARGET_SAMPLES_PER_WINDOW,
     31, 511, 3, MallocTracer::CONFIG_UPDATE_CHECK_PERIOD_SECS, 15);
 
@@ -102,6 +105,20 @@ static pthread_t _current_thread;
 static volatile bool _nested_malloc = false;
 static volatile bool _nested_posix_memalign = false;
 
+#ifdef UNIT_TEST
+static MallocTracer::AcquireCurrentTestHook _test_acquire_current;
+static MallocTracer::CriticalSectionTestHook _test_critical_section;
+static MallocTracer::ShouldSampleTestHook _test_should_sample;
+static MallocTracer::JfrSinkTestHook _test_jfr_sink;
+static MallocTracer::UsdtActiveTestHook _test_usdt_active;
+static MallocTracer::UsdtEmitTestHook _test_usdt_emit;
+static MallocTracer::StageTestHook _test_stage;
+
+static inline void testStage(const char* stage) {
+    if (_test_stage != NULL) _test_stage(stage);
+}
+#endif
+
 // Test if calloc() implementation calls malloc()
 static void* nested_malloc_hook(size_t size) {
     if (pthread_self() == _current_thread) {
@@ -313,15 +330,33 @@ void MallocTracer::maybeRecord(void* ret, size_t size) {
     if (running() && ret && size) {
         // Even we are not in a signal handler, we cannot malloc or
         // we may get into indefinite loop
+#ifdef UNIT_TEST
+        ProfiledThread* prof_thread = _test_acquire_current == NULL
+            ? ProfiledThread::acquireCurrent() : _test_acquire_current();
+        testStage("acquire");
+#else
         ProfiledThread* prof_thread = ProfiledThread::acquireCurrent();
+#endif
         if (prof_thread == nullptr) {
             Counters::increment(SAMPLES_DROPPED_THREAD_LOCAL);
             return;
         }
         CriticalSection cs(prof_thread);
+#ifdef UNIT_TEST
+        bool entered = _test_critical_section == NULL
+            ? cs.entered() : _test_critical_section(cs.entered());
+        testStage("critical-section");
+        if (entered) {
+            bool sampled = _test_should_sample == NULL
+                ? shouldSample(size) : _test_should_sample(size);
+            testStage("sampling");
+            if (sampled) recordMalloc(ret, size, prof_thread->tid());
+        }
+#else
         if (cs.entered() && shouldSample(size)) {
             recordMalloc(ret, size, prof_thread->tid());
         }
+#endif
     }
 }
 
@@ -347,18 +382,54 @@ void MallocTracer::recordMalloc(void* address, size_t size, int tid) {
     SAMPLER_PERF_PROBE(SP_NATIVEMEM);
 
     u64 current_interval = __atomic_load_n(&_interval, __ATOMIC_ACQUIRE);
-    MallocEvent event;
-    event._start_time = TSC::ticks();
-    event._address = (uintptr_t)address;
-    event._size = size;
-    // _interval == 0 means sample every allocation; weight is 1.0.
-    if (size == 0 || current_interval <= 1) {
-        event._weight = 1.0f;
-    } else {
-        event._weight = (float)(1.0 / (1.0 - exp(-(double)size / (double)current_interval)));
+    if (_output == NATIVE_MEMORY_JFR || _output == NATIVE_MEMORY_BOTH) {
+        MallocEvent event;
+        event._start_time = TSC::ticks();
+        event._address = (uintptr_t)address;
+        event._size = size;
+        // _interval == 0 means sample every allocation; weight is 1.0.
+        if (size == 0 || current_interval <= 1) {
+            event._weight = 1.0f;
+        } else {
+            event._weight = (float)(1.0 / (1.0 - exp(-(double)size / (double)current_interval)));
+        }
+
+#ifdef UNIT_TEST
+        testStage("jfr");
+        if (_test_jfr_sink != NULL) {
+            _test_jfr_sink(address, size, tid, &event);
+        } else {
+            Profiler::instance()->recordSample(NULL, size, tid, BCI_NATIVE_MALLOC, 0, &event);
+        }
+#else
+        Profiler::instance()->recordSample(NULL, size, tid, BCI_NATIVE_MALLOC, 0, &event);
+#endif
     }
 
-    Profiler::instance()->recordSample(NULL, size, tid, BCI_NATIVE_MALLOC, 0, &event);
+    bool usdt_active = false;
+    if ((_output == NATIVE_MEMORY_USDT || _output == NATIVE_MEMORY_BOTH) &&
+        size <= (size_t)INT64_MAX) {
+#ifdef UNIT_TEST
+        usdt_active = _test_usdt_active == NULL
+            ? isNativeAllocationUsdtEnabled() : _test_usdt_active();
+#else
+        usdt_active = isNativeAllocationUsdtEnabled();
+#endif
+    }
+    if (usdt_active) {
+        uint64_t weighted_bytes = nativeAllocationUsdtWeight(size, current_interval);
+#ifdef UNIT_TEST
+        testStage("usdt");
+        if (_test_usdt_emit != NULL) {
+            _test_usdt_emit(address, (uint64_t)size, weighted_bytes);
+        } else {
+            emitNativeAllocationUsdt(address, (uint64_t)size, weighted_bytes);
+        }
+#else
+        emitNativeAllocationUsdt(address, (uint64_t)size,
+                                 weighted_bytes);
+#endif
+    }
 
     u64 current_samples = __atomic_add_fetch(&_sample_count, 1, __ATOMIC_RELAXED);
     if ((current_samples % TARGET_SAMPLES_PER_WINDOW) == 0) {
@@ -375,9 +446,18 @@ void MallocTracer::recordMalloc(void* address, size_t size, int tid) {
             }
         }
     }
+#ifdef UNIT_TEST
+    testStage("accounting");
+#endif
 }
 
 Error MallocTracer::start(Arguments& args) {
+    _output = args._nativemem_output;
+    if (_output == NATIVE_MEMORY_NOOP) {
+        __atomic_store_n(&_running, false, __ATOMIC_RELEASE);
+        return Error::OK;
+    }
+
     _configured_interval = args._nativemem > 0 ? args._nativemem : 0;
     __atomic_store_n(&_interval, _configured_interval, __ATOMIC_RELEASE);
     __atomic_store_n(&_bytes_until_sample,
@@ -403,8 +483,59 @@ Error MallocTracer::start(Arguments& args) {
     return Error::OK;
 }
 
+Error MallocTracer::check(Arguments& args) {
+    if (args._nativemem_output_set && args._nativemem < 0) {
+        return Error("nativememoutput requires nativemem");
+    }
+#if !defined(__linux__)
+    if (args._nativemem_output == NATIVE_MEMORY_USDT) {
+        LOG_WARN("nativememoutput=usdt is supported only on Linux; "
+                 "native memory profiling will be disabled");
+        args._nativemem_output = NATIVE_MEMORY_NOOP;
+    } else if (args._nativemem_output == NATIVE_MEMORY_BOTH) {
+        LOG_WARN("USDT output is supported only on Linux; continuing with JFR only");
+        args._nativemem_output = NATIVE_MEMORY_JFR;
+    }
+#endif
+    return Error::OK;
+}
+
 void MallocTracer::stop() {
     // Ideally, we should reset original malloc entries, but it's not currently safe
     // in the view of library unloading. Consider using dl_iterate_phdr.
     __atomic_store_n(&_running, false, __ATOMIC_RELEASE);
 }
+
+#ifdef UNIT_TEST
+void MallocTracer::setTestHooks(AcquireCurrentTestHook acquire,
+                                CriticalSectionTestHook critical_section,
+                                ShouldSampleTestHook should_sample,
+                                JfrSinkTestHook jfr_sink,
+                                UsdtActiveTestHook usdt_active,
+                                UsdtEmitTestHook usdt_emit,
+                                StageTestHook stage) {
+    _test_acquire_current = acquire;
+    _test_critical_section = critical_section;
+    _test_should_sample = should_sample;
+    _test_jfr_sink = jfr_sink;
+    _test_usdt_active = usdt_active;
+    _test_usdt_emit = usdt_emit;
+    _test_stage = stage;
+}
+
+void MallocTracer::resetTestHooks() {
+    setTestHooks(NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+    configureForTest(NATIVE_MEMORY_JFR, 0, false);
+}
+
+void MallocTracer::configureForTest(NativeMemoryOutput output, u64 interval,
+                                    bool running) {
+    _output = output;
+    _configured_interval = interval;
+    __atomic_store_n(&_interval, interval, __ATOMIC_RELEASE);
+    __atomic_store_n(&_bytes_until_sample, 0, __ATOMIC_RELEASE);
+    __atomic_store_n(&_sample_count, (u64)0, __ATOMIC_RELEASE);
+    __atomic_store_n(&_last_config_update_ts, OS::nanotime(), __ATOMIC_RELEASE);
+    __atomic_store_n(&_running, running, __ATOMIC_RELEASE);
+}
+#endif

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

@@ -13,6 +13,8 @@
 #include "mutex.h"
 #include "pidController.h"
 
+class ProfiledThread;
+
 // Manages GOT-patching for malloc interception across all loaded native libraries.
 class MallocHooker {
   private:
@@ -41,6 +43,7 @@ class MallocTracer : public Engine {
     static const int TARGET_SAMPLES_PER_WINDOW = 100;
 
     static volatile bool _running;
+    static NativeMemoryOutput _output;
     static PidController _pid;
 
     static u64 nextPoissonInterval();
@@ -53,6 +56,7 @@ class MallocTracer : public Engine {
     }
 
     Error start(Arguments& args);
+    Error check(Arguments& args) override;
     void stop();
 
     static inline bool running() {
@@ -65,9 +69,30 @@ class MallocTracer : public Engine {
 
     static void recordMalloc(void* address, size_t size, int tid);
 
-    // Entry point for the malloc hooks (mallocTracer.cpp): gates on the
-    // sampling decision before touching thread-local state.
+    // Entry point for malloc hooks. Thread-local and critical-section guards
+    // must be acquired before the shared sampling state is touched.
     static void maybeRecord(void* ret, size_t size);
+
+#ifdef UNIT_TEST
+    using AcquireCurrentTestHook = ProfiledThread* (*)();
+    using CriticalSectionTestHook = bool (*)(bool);
+    using ShouldSampleTestHook = bool (*)(size_t);
+    using JfrSinkTestHook = void (*)(void*, size_t, int, MallocEvent*);
+    using UsdtActiveTestHook = bool (*)();
+    using UsdtEmitTestHook = void (*)(void*, uint64_t, uint64_t);
+    using StageTestHook = void (*)(const char*);
+
+    static void setTestHooks(AcquireCurrentTestHook acquire,
+                             CriticalSectionTestHook critical_section,
+                             ShouldSampleTestHook should_sample,
+                             JfrSinkTestHook jfr_sink,
+                             UsdtActiveTestHook usdt_active,
+                             UsdtEmitTestHook usdt_emit,
+                             StageTestHook stage);
+    static void resetTestHooks();
+    static void configureForTest(NativeMemoryOutput output, u64 interval,
+                                 bool running);
+#endif
 };
 
 #endif // _MALLOCTRACER_H