chore: support for pprof as transport format in ruby

Author: petethepig

.gitignore

@@ -1,3 +1,4 @@
+.idea/
 /target/
 Cargo.lock
 **/*.rs.bk

Cargo.toml

@@ -58,6 +58,7 @@ reqwest = { version = "0.11", features = ["blocking", "rustls-tls-native-roots"]
 url = "2.2.2"
 libflate = "1.2.0"
 libc = "^0.2.124"
+prost = "0.10"
 
 [dev-dependencies]
 tokio = { version = "1.18", features = ["full"] }

pyroscope_cli/Cargo.lock

@@ -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::FOLDED);
 
     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/ext/rbspy/src/lib.rs

@@ -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
 

pyroscope_ffi/ruby/lib/pyroscope.rb

@@ -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/backend/types.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.to_owned(),
+            }
+        }).collect()
+}

src/encode/folded.rs

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

src/encode/mod.rs

@@ -0,0 +1,170 @@
+use std::collections::HashMap;
+
+use prost::Message;
+
+use crate::backend::types::{EncodedReport, Report};
+use crate::encode::profiles::{Function, Label, Line, Location, Profile, Sample, ValueType};
+
+
+struct PProfBuilder {
+    profile: Profile,
+    strings: HashMap<String, i64>,
+    functions: HashMap<FunctionMirror, u64>,
+    locations: HashMap<LocationMirror, u64>,
+}
+
+#[derive(Hash, PartialEq, Eq, Clone)]
+pub struct LocationMirror {
+    pub function_id: u64,
+    pub line: i64,
+}
+
+#[derive(Hash, PartialEq, Eq, Clone)]
+pub struct FunctionMirror {
+    pub name: i64,
+    pub filename: i64,
+}
+
+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, fm: FunctionMirror) -> u64 {
+        let v = self.functions.get(&fm);
+        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: fm.name,
+            system_name: 0,
+            filename: fm.filename,
+            start_line: 0,
+        };
+        self.functions.insert(fm, id);
+        self.profile.function.push(f);
+        id
+    }
+
+    fn add_location(&mut self, lm: LocationMirror) -> u64 {
+        let v = self.locations.get(&lm);
+        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: lm.function_id,
+                line: lm.line,
+            }],
+            is_folded: false,
+        };
+        self.locations.insert(lm, id);
+        self.profile.location.push(l);
+        id
+    }
+}
+
+pub fn encode(reports: &Vec<Report>, sample_rate: u32, start_time_nanos: u64, duration_nanos: u64) -> 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: start_time_nanos as i64,
+            duration_nanos: duration_nanos as i64,
+            period_type: None,
+            period: 0,
+            comment: vec![],
+            default_sample_type: 0,
+        },
+    };
+    {
+        let count = b.add_string(&"count".to_string());
+        let samples = b.add_string(&"samples".to_string());
+        let milliseconds = b.add_string(&"milliseconds".to_string());
+        b.profile.sample_type.push(ValueType {
+            r#type: samples,
+            unit: count,
+        });
+        b.profile.period = 1_000 / sample_rate as i64;
+        b.profile.period_type = Some(ValueType {
+            r#type: 0,
+            unit: milliseconds,
+        })
+    }
+    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 = b.add_string(&sf.name.as_ref().unwrap_or(&"".to_string()));
+                let filename = b.add_string(&sf.filename.as_ref().unwrap_or(&"".to_string()));
+                let line = sf.line.unwrap_or(0) as i64;
+                let function_id = b.add_function(FunctionMirror {
+                    name: name,
+                    filename: filename,
+                });
+                let location_id = b.add_location(LocationMirror {
+                    function_id: function_id,
+                    line: line,
+                });
+                sample.location_id.push(location_id as u64);
+            }
+            let mut labels = HashMap::new();
+            for l in &stacktrace.metadata.tags {
+                let k = b.add_string(&l.key);
+                let v = b.add_string(&l.value);
+                labels.insert(k, v);
+            }
+            for l in &report.metadata.tags {
+                let k = b.add_string(&l.key);
+                let v = b.add_string(&l.value);
+                labels.insert(k, v);
+            }
+            for (k, v) in &labels {
+                sample.label.push(Label {
+                    key: *k,
+                    str: *v,
+                    num: 0,
+                    num_unit: 0,
+                })
+            }
+            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(),
+    }]
+}