Add configuration to save profiles to disk (#1837)

DataDog · Dec 20, 2022 · 319b2af · 319b2af
1 parent 828d579
commit 319b2af
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 13 deletions.
diff --git a/profiling/src/config.rs b/profiling/src/config.rs
@@ -126,6 +126,7 @@ pub(crate) enum ConfigId {
     ProfilingEndpointCollectionEnabled,
     ProfilingExperimentalCpuTimeEnabled,
     ProfilingLogLevel,
+    ProfilingOutputPprof,
 
     // todo: do these need to be kept in sync with the tracer?
     AgentHost,
@@ -147,6 +148,11 @@ impl ConfigId {
             ProfilingExperimentalCpuTimeEnabled => b"DD_PROFILING_EXPERIMENTAL_CPU_TIME_ENABLED\0",
             ProfilingLogLevel => b"DD_PROFILING_LOG_LEVEL\0",
 
+            /* Note: this is meant only for debugging and testing. Please don't
+             * advertise this in the docs.
+             */
+            ProfilingOutputPprof => b"DD_PROFILING_OUTPUT_PPROF\0",
+
             AgentHost => b"DD_AGENT_HOST\0",
             Env => b"DD_ENV\0",
             Service => b"DD_SERVICE\0",
@@ -182,6 +188,13 @@ pub(crate) unsafe fn profiling_experimental_cpu_time_enabled() -> bool {
     get_bool(ProfilingExperimentalCpuTimeEnabled, true)
 }
 
+/// # Safety
+/// This function must only be called after config has been initialized in
+/// rinit, and before it is uninitialized in mshutdown.
+pub(crate) unsafe fn profiling_output_pprof() -> Option<Cow<'static, str>> {
+    get_str(ProfilingOutputPprof)
+}
+
 unsafe fn get_bool(id: ConfigId, default: bool) -> bool {
     get_value(id).try_into().unwrap_or(default)
 }
@@ -364,6 +377,16 @@ pub(crate) fn minit(module_number: libc::c_int) {
                     ini_change: Some(zai_config_system_ini_change),
                     parser: Some(parse_level_filter),
                 },
+                zai_config_entry {
+                    id: transmute(ProfilingOutputPprof),
+                    name: ProfilingOutputPprof.env_var_name(),
+                    type_: ZAI_CONFIG_TYPE_STRING,
+                    default_encoded_value: ZaiStringView::new(),
+                    aliases: std::ptr::null_mut(),
+                    aliases_count: 0,
+                    ini_change: Some(zai_config_system_ini_change),
+                    parser: Some(parse_utf8_string),
+                },
                 zai_config_entry {
                     id: transmute(AgentHost),
                     name: AgentHost.env_var_name(),
@@ -472,6 +495,10 @@ mod tests {
                 "datadog.profiling.experimental_cpu_time_enabled",
             ),
             (b"DD_PROFILING_LOG_LEVEL\0", "datadog.profiling.log_level"),
+            (
+                b"DD_PROFILING_OUTPUT_PPROF\0",
+                "datadog.profiling.output_pprof",
+            ),
         ];
 
         for (env_name, expected_ini_name) in cases {

diff --git a/profiling/src/profiling.rs b/profiling/src/profiling.rs
@@ -560,6 +560,12 @@ impl Uploader {
     }
 
     pub fn run(&self) {
+        /* Safety: Called from Profiling::new, which is after config is
+         * initialized, and before it's destroyed in mshutdown.
+         */
+        let pprof_filename = unsafe { crate::config::profiling_output_pprof() };
+        let mut i = 0;
+
         loop {
             /* Since profiling uploads are going over the Internet and not just
              * the local network, it would be ideal if they were the lowest
@@ -581,19 +587,25 @@ impl Uploader {
                 },
 
                 recv(self.upload_receiver) -> message => match message {
-                    Ok(upload_message) => match Self::upload(upload_message) {
-                        Ok(status) => {
-                            if status >= 400 {
-                                warn!(
-                                    "Unexpected HTTP status when sending profile (HTTP {}).",
-                                    status
-                                )
-                            } else {
-                                info!("Successfully uploaded profile (HTTP {}).", status)
-                            }
-                        }
-                        Err(err) => {
-                            warn!("Failed to upload profile: {}", err)
+                    Ok(upload_message) => {
+                        match pprof_filename.as_ref() {
+                            Some(filename) => {
+                                let r = upload_message.profile.serialize(None, None).unwrap();
+                                i += 1;
+                                std::fs::write(format!("{filename}.{i}"), r.buffer).expect("write to succeed")
+                            },
+                            None => match Self::upload(upload_message) {
+                                Ok(status) => {
+                                    if status >= 400 {
+                                        warn!("Unexpected HTTP status when sending profile (HTTP {status}).")
+                                    } else {
+                                        info!("Successfully uploaded profile (HTTP {status}).")
+                                    }
+                                }
+                                Err(err) => {
+                                    warn!("Failed to upload profile: {err}")
+                                }
+                            },
                         }
                     },
                     _ => {