Add environment variable for tracking child processes

docs/src/configuration.md

@@ -179,3 +179,10 @@ Whenever to use a more intrusive, faster unwinding algorithm; enabled by default
 
 Setting it to `0` will on average significantly slow down unwinding. This option
 is provided only for debugging purposes.
+
+### `MEMORY_PROFILER_TRACK_CHILD_PROCESSES`
+
+*Default: `0`*
+
+If set to `1`, bytehound will also track the memory allocations of child processes spawned by the
+profiled process.

integration-tests/src/tests.rs

@@ -1278,3 +1278,32 @@ fn test_cross_thread_alloc_non_culled() {
 
     assert!( iter.next().is_none() );
 }
+
+#[test]
+fn test_track_spawned_children() {
+    let cwd = workdir();
+
+    compile_with_flags( "spawn-child.c", &["-o", "spawn-child"] );
+
+    run_on_target(
+        &cwd,
+        "./spawn-child",
+        EMPTY_ARGS,
+        &[
+            ("LD_PRELOAD", preload_path().into_os_string()),
+            ("MEMORY_PROFILER_LOG", "debug".into()),
+            ("MEMORY_PROFILER_TRACK_CHILD_PROCESSES", "1".into())
+        ]
+    ).assert_success();
+
+    eprintln!("cwd: {}", cwd.display());
+
+    for entry in std::fs::read_dir(cwd).unwrap() {
+        let entry = entry.unwrap();
+        if entry.file_name().to_str().unwrap().starts_with("memory-profiling_ls") {
+            return;
+        }
+    }
+
+    panic!("Memory profiling file for spawned process not found");
+}

integration-tests/test-programs/spawn-child.c

@@ -0,0 +1,22 @@
+#include <stdlib.h>
+#include <unistd.h>
+#include <sys/wait.h>
+
+int main() {
+    usleep( 100000 );
+    malloc( 10001 );
+
+    pid_t pid = fork();
+    if( pid == 0 ) {
+        // Child
+        if (execl("/usr/bin/ls", "/usr/bin/ls", NULL) == -1) {
+            return 1;
+        }
+        return 0;
+    }
+
+    waitpid(pid, NULL, 0);
+    malloc( 10003 );
+
+    return 0;
+}

preload/src/global.rs

@@ -7,7 +7,7 @@ use std::thread;
 use crate::arc_lite::ArcLite;
 use crate::event::{InternalAllocationId, InternalEvent, send_event};
 use crate::spin_lock::{SpinLock, SpinLockGuard};
-use crate::syscall;
+use crate::{opt, syscall};
 use crate::unwind::{ThreadUnwindState, prepare_to_start_unwinding};
 use crate::timestamp::Timestamp;
 use crate::allocation_tracker::AllocationTracker;
@@ -553,7 +553,9 @@ fn initialize_stage_2() {
     crate::init::initialize_atexit_hook();
     crate::init::initialize_signal_handlers();
 
-    std::env::remove_var( "LD_PRELOAD" );
+    if !opt::get().track_child_processes {
+        std::env::remove_var( "LD_PRELOAD" );
+    }
 
     info!( "Stage 2 initialization finished" );
 }

preload/src/opt.rs

@@ -22,7 +22,8 @@ pub struct Opts {
     pub backtrace_cache_size_level_2: usize,
     pub cull_temporary_allocations: bool,
     pub temporary_allocation_lifetime_threshold: u64,
-    pub temporary_allocation_pending_threshold: Option< usize >
+    pub temporary_allocation_pending_threshold: Option< usize >,
+    pub track_child_processes: bool
 }
 
 static mut OPTS: Opts = Opts {
@@ -48,6 +49,7 @@ static mut OPTS: Opts = Opts {
     cull_temporary_allocations: false,
     temporary_allocation_lifetime_threshold: 10000,
     temporary_allocation_pending_threshold: None,
+    track_child_processes: false
 };
 
 trait ParseVar: Sized {
@@ -143,7 +145,9 @@ pub unsafe fn initialize() {
         "MEMORY_PROFILER_TEMPORARY_ALLOCATION_LIFETIME_THRESHOLD"
             => &mut opts.temporary_allocation_lifetime_threshold,
         "MEMORY_PROFILER_TEMPORARY_ALLOCATION_PENDING_THRESHOLD"
-            => &mut opts.temporary_allocation_pending_threshold
+            => &mut opts.temporary_allocation_pending_threshold,
+        "MEMORY_PROFILER_TRACK_CHILD_PROCESSES"
+            => &mut opts.track_child_processes
     }
 
     opts.is_initialized = true;