pprof encoder

Author: korniltsev

pyroscope_ffi/ruby/ext/rbspy/src/lib.rs

@@ -10,6 +10,7 @@ use pyroscope_rbspy::{rbspy_backend, RbspyConfig};
 
 use pyroscope::{pyroscope::Compression, PyroscopeAgent};
 use pyroscope::backend::{Report, StackFrame, Tag};
+use pyroscope::pyroscope::ReportEncoding;
 
 pub fn transform_report(report: Report) -> Report {
     let cwd = env::current_dir().unwrap();
@@ -107,7 +108,8 @@ pub extern "C" fn initialize_logging(logging_level: u32) -> bool {
 pub extern "C" fn initialize_agent(
     application_name: *const c_char, server_address: *const c_char, auth_token: *const c_char,
     sample_rate: u32, detect_subprocesses: bool, oncpu: bool, report_pid: bool,
-    report_thread_id: bool, tags: *const c_char, compression: *const c_char
+    report_thread_id: bool, tags: *const c_char, compression: *const c_char,
+    report_encoding: *const c_char
 ) -> bool {
     // Initialize FFIKit
     let recv = ffikit::initialize_ffi().unwrap();
@@ -137,7 +139,14 @@ pub extern "C" fn initialize_agent(
         .unwrap()
         .to_string();
 
+    let report_encoding = unsafe { CStr::from_ptr(report_encoding) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
     let compression = Compression::from_str(&compression_string);
+    let report_encoding = ReportEncoding::from_str(&report_encoding)
+        .unwrap_or(ReportEncoding::PPROF);
 
     let pid = std::process::id();
 
@@ -156,7 +165,8 @@ pub extern "C" fn initialize_agent(
     let mut agent_builder = PyroscopeAgent::builder(server_address, application_name)
         .backend(rbspy)
         .func(transform_report)
-        .tags(tags);
+        .tags(tags)
+        .report_encoding(report_encoding);
 
     if auth_token != "" {
         agent_builder = agent_builder.auth_token(auth_token);

pyroscope_ffi/ruby/lib/pyroscope.rb

@@ -8,7 +8,7 @@ module Rust
     extend FFI::Library
     ffi_lib File.expand_path(File.dirname(__FILE__)) + "/rbspy/rbspy.#{RbConfig::CONFIG["DLEXT"]}"
     attach_function :initialize_logging, [:int], :bool
-    attach_function :initialize_agent, [:string, :string, :string, :int, :bool, :bool, :bool, :bool, :string, :string], :bool
+    attach_function :initialize_agent, [:string, :string, :string, :int, :bool, :bool, :bool, :bool, :string, :string, :string], :bool
     attach_function :add_thread_tag, [:uint64, :string, :string], :bool
     attach_function :remove_thread_tag, [:uint64, :string, :string], :bool
     attach_function :add_global_tag, [:string, :string], :bool
@@ -22,7 +22,7 @@ module Utils
     attach_function :thread_id, [], :uint64
   end
 
-  Config = Struct.new(:application_name, :app_name, :server_address, :auth_token, :log_level, :sample_rate, :detect_subprocesses, :oncpu, :report_pid, :report_thread_id, :tags, :compression) do
+  Config = Struct.new(:application_name, :app_name, :server_address, :auth_token, :log_level, :sample_rate, :detect_subprocesses, :oncpu, :report_pid, :report_thread_id, :tags, :compression, :report_encoding) do
     def initialize(*)
       super
       # defaults:
@@ -37,6 +37,7 @@ def initialize(*)
       self.log_level = 'error'
       self.tags = {}
       self.compression = 'gzip'
+      self.report_encoding = 'pprof'
     end
   end
 
@@ -79,7 +80,8 @@ def configure
         @config.report_pid || false,
         @config.report_thread_id || false,
         tags_to_string(@config.tags || {}),
-        @config.compression || ""
+        @config.compression || "",
+        @config.report_encoding || "pprof"
       )
     end
 

src/backend/types.rs

@@ -130,6 +130,15 @@ pub struct Report {
     pub metadata: Metadata,
 }
 
+#[derive(Debug)]
+pub struct EncodedReport {
+    pub format: String,
+    pub content_type: String,
+    pub content_encoding: String,
+    pub data: Vec<u8>,
+    pub metadata: Metadata,
+}
+
 /// Custom implementation of the Hash trait for Report.
 /// Only the metadata is hashed.
 impl Hash for Report {

src/encode/folded.rs

@@ -0,0 +1,17 @@
+
+use crate::backend;
+
+use backend::types::{Report, EncodedReport};
+
+pub fn encode(reports: Vec<Report>) -> Vec<EncodedReport> {
+    reports.into_iter()
+        .map(|r| {
+            EncodedReport {
+                format: "folded".to_string(),
+                content_type: "binary/octet-stream".to_string(),
+                content_encoding: "".to_string(),
+                data: r.to_string().into_bytes(),
+                metadata: r.metadata,
+            }
+        }).collect()
+}

src/encode/mod.rs

@@ -0,0 +1,5 @@
+pub mod folded;
+pub mod pprof;
+pub mod profiles;
+
+

src/encode/pprof.rs

@@ -0,0 +1,129 @@
+use std::collections::HashMap;
+
+use prost::Message;
+
+use crate::backend::types::{EncodedReport, Report};
+use crate::encode::profiles::{Function, Line, Location, Profile, Sample, ValueType};
+
+
+struct PProfBuilder {
+    profile: Profile,
+    strings: HashMap<String, i64>,
+    functions: HashMap<i64, u64>,
+    locations: HashMap<u64, u64>,
+}
+
+impl PProfBuilder {
+    fn add_string(&mut self, s: &String) -> i64 {
+        let v = self.strings.get(s);
+        if v.is_some() {
+            return *v.unwrap();
+        }
+        assert!(self.strings.len() != self.profile.string_table.len() + 1);
+        let id: i64 = self.strings.len() as i64;
+        self.strings.insert(s.to_owned(), id);
+        self.profile.string_table.push(s.to_owned());
+        id
+    }
+
+    fn add_function(&mut self, name: i64) -> u64 {
+        let v = self.functions.get(&name);
+        if v.is_some() {
+            return *v.unwrap();
+        }
+        assert!(self.functions.len() != self.profile.function.len() + 1);
+        let id: u64 = self.functions.len() as u64 + 1;
+        let f = Function {
+            id: id,
+            name: name,
+            system_name: 0,
+            filename: 0,
+            start_line: 0,
+        };
+        self.functions.insert(name, id);
+        self.profile.function.push(f);
+        id
+    }
+
+    fn add_location(&mut self, function_id: u64) -> u64 {
+        let v = self.locations.get(&function_id);
+        if v.is_some() {
+            return *v.unwrap();
+        }
+        assert!(self.locations.len() != self.profile.location.len() + 1);
+        let id: u64 = self.locations.len() as u64 + 1;
+        let l = Location {
+            id,
+            mapping_id: 0,
+            address: 0,
+            line: vec![Line {
+                function_id: function_id,
+                line: 0,
+            }],
+            is_folded: false,
+        };
+        self.locations.insert(function_id, id);
+        self.profile.location.push(l);
+        id
+    }
+}
+
+pub fn encode(reports: Vec<Report>) -> Vec<EncodedReport> {
+    let mut b = PProfBuilder {
+        strings: HashMap::new(),
+        functions: HashMap::new(),
+        locations: HashMap::new(),
+        profile: Profile {
+            sample_type: vec![],
+            sample: vec![],
+            mapping: vec![],
+            location: vec![],
+            function: vec![],
+            string_table: vec![],
+            drop_frames: 0,
+            keep_frames: 0,
+            time_nanos: 0,
+            duration_nanos: 0,
+            period_type: None,
+            period: 0,
+            comment: vec![],
+            default_sample_type: 0,
+        },
+    };
+    {
+        let cpu = b.add_string(&"cpu".to_string());
+        let samples = b.add_string(&"samples".to_string());
+        b.profile.sample_type.push(ValueType {
+            r#type: samples,
+            unit: cpu,
+        });
+    }
+    for report in &reports {
+        for (stacktrace, value) in &report.data {
+            let mut sample = Sample {
+                location_id: vec![],
+                value: vec![*value as i64],
+                label: vec![],
+            };
+            for sf in &stacktrace.frames {
+                let name = format!("{}:{} - {}",
+                                   sf.filename.as_ref().unwrap_or(&"".to_string()),
+                                   sf.line.unwrap_or(0),
+                                   sf.name.as_ref().unwrap_or(&"".to_string()));
+                let name = b.add_string(&name);
+                let function_id = b.add_function(name);
+                let location_id = b.add_location(function_id);
+                sample.location_id.push(location_id as u64);
+            }
+            b.profile.sample.push(sample);
+        }
+    }
+
+    vec![EncodedReport {
+        format: "pprof".to_string(),
+        content_type: "binary/octet-stream".to_string(),
+        content_encoding: "".to_string(),
+        data: b.profile.encode_to_vec(),
+        metadata: Default::default(),
+    }]
+}

src/pprof/pprof.proto renamed to src/encode/profiles.proto

@@ -39,7 +39,7 @@
 syntax = "proto3";
 
 //package perftools.profiles;
-package pprof;
+package profiles;
 
 option java_package = "com.google.perftools.profiles";
 option java_outer_classname = "ProfileProto";

src/pprof/pprof.rs renamed to src/encode/profiles.rs

@@ -39,6 +39,8 @@
 //! agent_ready.shutdown();
 //! ```
 
+extern crate core;
+
 // Re-exports structs
 pub use crate::pyroscope::PyroscopeAgent;
 pub use error::{PyroscopeError, Result};
@@ -52,4 +54,4 @@ pub mod timer;
 
 // Private modules
 mod utils;
-mod pprof;
+mod encode;