fix: memory profile data cannot be compressed

deepflowio · Sep 17, 2024 · 27749ee · 27749ee
1 parent 451006e
commit 27749ee
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 6 deletions.
diff --git a/agent/src/ebpf/mod.rs b/agent/src/ebpf/mod.rs
@@ -489,7 +489,7 @@ extern "C" {
     pub fn set_feature_regex(idx: c_int, pattern: *const c_char) -> c_int;
     /*
      * @brief Add regex-matched process list for feature.
-     * 
+     *
      * @param feature Refers to a specific feature module, value: FEATURE_*
      * @param pids Address of the process list
      * @param num Number of elements in the process list

diff --git a/agent/src/ebpf_dispatcher.rs b/agent/src/ebpf_dispatcher.rs
@@ -534,7 +534,7 @@ impl EbpfCollector {
                 let Some(m_ctx) = (ctx as *mut memory_profile::MemoryContext).as_mut() else {
                     return;
                 };
-                m_ctx.update(data);
+                m_ctx.update(data, PROFILE_STACK_COMPRESSION);
                 m_ctx.report(
                     Duration::from_nanos(ts_nanos),
                     EBPF_PROFILE_SENDER.as_mut().unwrap(),

diff --git a/agent/src/ebpf_dispatcher/memory_profile.rs b/agent/src/ebpf_dispatcher/memory_profile.rs
@@ -27,6 +27,7 @@ use std::{
 
 use log::{debug, trace, warn};
 use procfs::process::Process;
+use zstd::bulk::compress;
 
 use public::{proto::metric, queue::DebugSender};
 
@@ -62,6 +63,7 @@ impl ProcessAllocInfo {
         stacks: &mut HashMap<u32, metric::Profile>,
         data: &ebpf::stack_profile_data,
         update_trace_str: bool,
+        stack_compression: bool,
     ) {
         let entry = stacks.entry(data.u_stack_id).or_insert_with(|| {
             let mut profile = metric::Profile::default();
@@ -77,6 +79,16 @@ impl ProcessAllocInfo {
             profile.data =
                 slice::from_raw_parts(data.stack_data as *mut u8, data.stack_data_len as usize)
                     .to_vec();
+            if stack_compression {
+                match compress(&profile.data, 0) {
+                    Ok(compressed_data) => {
+                        profile.data_compressed = true;
+                        profile.data = compressed_data;
+                    }
+                    Err(e) => debug!("failed to compress ebpf memory profile: {:?}", e),
+                }
+            }
+
             let container_id =
                 CStr::from_ptr(data.container_id.as_ptr() as *const libc::c_char).to_string_lossy();
             if let Some(policy_getter) = POLICY_GETTER.as_ref() {
@@ -94,7 +106,7 @@ impl ProcessAllocInfo {
         }
     }
 
-    unsafe fn update(&mut self, data: &ebpf::stack_profile_data) {
+    unsafe fn update(&mut self, data: &ebpf::stack_profile_data, stack_compression: bool) {
         // Memory allocations and frees may come in any order in a multi-core envrionment.
         // Make best effort to handle such situation.
         if data.count == 0 {
@@ -127,7 +139,7 @@ impl ProcessAllocInfo {
             }
         } else {
             // allocs
-            Self::update_stack(&mut self.alloc, data, false);
+            Self::update_stack(&mut self.alloc, data, false, stack_compression);
             // for languages without free (i.e. JAVA), not recording in_use info, only allocs
             if data.mem_addr == 0 {
                 return;
@@ -148,7 +160,7 @@ impl ProcessAllocInfo {
                         stack_id: data.u_stack_id,
                         size: data.count,
                     });
-                    Self::update_stack(&mut self.in_use, data, true);
+                    Self::update_stack(&mut self.in_use, data, true, stack_compression);
                 }
             } else {
                 self.allocated_addrs.insert(
@@ -161,7 +173,7 @@ impl ProcessAllocInfo {
                         frees: 0,
                     },
                 );
-                Self::update_stack(&mut self.in_use, data, true);
+                Self::update_stack(&mut self.in_use, data, true, stack_compression);
             }
         }
     }