pyroscope 1.0.9 → 1.1.0

data/ext/rbspy/src/encode/pprof.rs

@@ -0,0 +1,159 @@
+use std::collections::HashMap;
+
+use crate::backend::types::Report;
+use crate::encode::gen::google::{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 let Some(v) = v {
+            return *v;
+        }
+        assert_ne!(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 let Some(v) = v {
+            return *v;
+        }
+        assert_ne!(self.functions.len(), self.profile.function.len() + 1);
+        let id: u64 = self.functions.len() as u64 + 1;
+        let f = Function {
+            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 let Some(v) = v {
+            return *v;
+        }
+        assert_ne!(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,
+) -> Profile {
+    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,
+        },
+    };
+    b.add_string(&"".to_string());
+    {
+        let cpu = b.add_string(&"cpu".to_string());
+        let nanoseconds = b.add_string(&"nanoseconds".to_string());
+        b.profile.sample_type.push(ValueType {
+            r#type: cpu,
+            unit: nanoseconds,
+        });
+        b.profile.period = 1_000_000_000 / sample_rate as i64;
+        b.profile.period_type = Some(ValueType {
+            r#type: cpu,
+            unit: nanoseconds,
+        });
+    }
+    for report in reports {
+        for (stacktrace, value) in &report.data {
+            let mut sample = Sample {
+                location_id: vec![],
+                value: vec![*value as i64 * b.profile.period],
+                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, filename });
+                let location_id = b.add_location(LocationMirror { function_id, line });
+                sample.location_id.push(location_id);
+            }
+            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);
+        }
+    }
+    b.profile
+}

data/ext/rbspy/src/error.rs

@@ -0,0 +1,67 @@
+pub type Result<T> = std::result::Result<T, PyroscopeError>;
+
+/// Error type of Pyroscope
+#[non_exhaustive]
+#[derive(thiserror::Error, Debug)]
+pub enum PyroscopeError {
+    #[error("Other: {}", &.0)]
+    AdHoc(String),
+
+    #[error("{msg}: {source:?}")]
+    Compat {
+        msg: String,
+        #[source]
+        source: Box<dyn std::error::Error + Send + Sync + 'static>,
+    },
+
+    #[error("BackendImpl error")]
+    BackendImpl,
+
+    #[error(transparent)]
+    Reqwest(#[from] reqwest::Error),
+
+    #[error(transparent)]
+    ParseError(#[from] url::ParseError),
+
+    #[error(transparent)]
+    TimeSource(#[from] std::time::SystemTimeError),
+
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+
+    #[error(transparent)]
+    Json(#[from] serde_json::Error),
+}
+
+impl PyroscopeError {
+    /// Create a new instance of PyroscopeError
+    pub fn new(msg: &str) -> Self {
+        PyroscopeError::AdHoc(msg.to_string())
+    }
+
+    /// Create a new instance of PyroscopeError with source
+    pub fn new_with_source<E>(msg: &str, source: E) -> Self
+    where
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        PyroscopeError::Compat {
+            msg: msg.to_string(),
+            source: Box::new(source),
+        }
+    }
+}
+
+impl<T> From<std::sync::PoisonError<T>> for PyroscopeError {
+    fn from(_err: std::sync::PoisonError<T>) -> Self {
+        PyroscopeError::AdHoc("Poison Error".to_owned())
+    }
+}
+
+impl<T: 'static + Send + Sync> From<std::sync::mpsc::SendError<T>> for PyroscopeError {
+    fn from(err: std::sync::mpsc::SendError<T>) -> Self {
+        PyroscopeError::Compat {
+            msg: String::from("SendError Error"),
+            source: Box::new(err),
+        }
+    }
+}

data/ext/rbspy/src/ffikit.rs

@@ -0,0 +1,76 @@
+use crate::backend::Tag;
+use crate::error::{PyroscopeError, Result};
+use crate::pyroscope::{PyroscopeAgentBuilder, PyroscopeAgentRunning};
+use crate::{PyroscopeAgent, ThreadId};
+use lazy_static::lazy_static;
+use std::sync::{
+    mpsc::{self, Receiver, Sender},
+    Mutex,
+};
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum Signal {
+    Kill,
+    AddThreadTag(ThreadId, Tag),
+    RemoveThreadTag(ThreadId, Tag),
+}
+
+const TAG: &str = "pyroscope::ffikit";
+
+lazy_static! {
+    static ref SENDER: Mutex<Option<Sender<Signal>>> = Mutex::new(None);
+}
+pub fn run(agent: PyroscopeAgentBuilder) -> Result<()> {
+    let mut sender_holder = SENDER.lock()?;
+    if (*sender_holder).is_some() {
+        return Err(PyroscopeError::new("FFI channel already initialized"));
+    }
+
+    let agent = agent.build()?;
+
+    let agent = agent.start()?;
+
+    let (sender, receiver): (Sender<Signal>, Receiver<Signal>) = mpsc::channel();
+
+    *sender_holder = Some(sender);
+
+    std::thread::spawn(move || {
+        while let Ok(signal) = receiver.recv() {
+            match signal {
+                Signal::Kill => {
+                    if let Err(err) = stop(agent) {
+                        log::error!(target: TAG, "failed to stop agent {err}");
+                    }
+                    break;
+                }
+                Signal::AddThreadTag(thread_id, tag) => {
+                    if let Err(err) = agent.add_thread_tag(thread_id, tag) {
+                        log::error!(target: TAG, "failed to add tag {err}");
+                    }
+                }
+                Signal::RemoveThreadTag(thread_id, tag) => {
+                    if let Err(err) = agent.remove_thread_tag(thread_id, tag) {
+                        log::error!(target: TAG, "failed to remove tag {err}");
+                    }
+                }
+            }
+        }
+    });
+
+    Ok(())
+}
+
+pub fn send(signal: Signal) -> Result<()> {
+    if let Some(sender) = &*SENDER.lock()? {
+        sender.send(signal)?;
+    } else {
+        return Err(PyroscopeError::new("FFI channel not initialized"));
+    }
+    Ok(())
+}
+
+fn stop(agent: PyroscopeAgent<PyroscopeAgentRunning>) -> Result<()> {
+    agent.stop()?;
+    *SENDER.lock()? = None;
+    Ok(())
+}

data/ext/rbspy/src/lib.rs

@@ -1,4 +1,17 @@
-mod backend;
+pub use crate::pyroscope::PyroscopeAgent;
+pub use error::{PyroscopeError, Result};
+
+pub mod backend;
+pub mod encode;
+pub mod error;
+pub mod ffikit;
+pub mod pyroscope;
+pub mod session;
+pub mod timer;
+
+mod rbspy_backend;
+mod utils;
+pub use utils::ThreadId;
 
 use rbspy::sampler::Sampler;
 use remoteprocess::Pid;
@@ -6,9 +19,9 @@ use std::env;
 use std::ffi::CStr;
 use std::os::raw::c_char;
 
-use crate::backend::Rbspy;
-use pyroscope::backend::{BackendConfig, BackendImpl, Report, StackFrame, Tag};
-use pyroscope::pyroscope::PyroscopeAgentBuilder;
+use crate::backend::{BackendConfig, BackendImpl, Report, StackFrame, Tag};
+use crate::pyroscope::PyroscopeAgentBuilder;
+use crate::rbspy_backend::Rbspy;
 
 const LOG_TAG: &str = "Pyroscope::rbspy::ffi";
 const RBSPY_NAME: &str = "rbspy";
@@ -110,6 +123,8 @@ pub unsafe extern "C" fn initialize_agent(
     oncpu: bool,
     report_pid: bool,
     report_thread_id: bool,
+    runtime_name: *const c_char,
+    runtime_version: *const c_char,
     tags: *const c_char,
     tenant_id: *const c_char,
     http_headers_json: *const c_char,
@@ -134,6 +149,16 @@ pub unsafe extern "C" fn initialize_agent(
         .unwrap()
         .to_string();
 
+    let runtime_name = unsafe { CStr::from_ptr(runtime_name) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let runtime_version = unsafe { CStr::from_ptr(runtime_version) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
     let tags_string = unsafe { CStr::from_ptr(tags) }
         .to_str()
         .unwrap()
@@ -172,7 +197,8 @@ pub unsafe extern "C" fn initialize_agent(
         rbspy,
     )
     .func(transform_report_with_current_dir)
-    .tags(tags);
+    .tags(tags)
+    .runtime(runtime_name, runtime_version);
 
     if !basic_auth_user.is_empty() && !basic_auth_password.is_empty() {
         agent_builder = agent_builder.basic_auth(basic_auth_user, basic_auth_password);
@@ -182,28 +208,28 @@ pub unsafe extern "C" fn initialize_agent(
         agent_builder = agent_builder.tenant_id(tenant_id);
     }
 
-    let http_headers = pyroscope::pyroscope::parse_http_headers_json(http_headers_json);
+    let http_headers = crate::pyroscope::parse_http_headers_json(http_headers_json);
     match http_headers {
         Ok(http_headers) => {
             agent_builder = agent_builder.http_headers(http_headers);
         }
         Err(e) => match e {
-            pyroscope::PyroscopeError::Json(e) => {
+            crate::PyroscopeError::Json(e) => {
                 log::error!(target: LOG_TAG, "parse_http_headers_json error {}", e);
             }
-            pyroscope::PyroscopeError::AdHoc(e) => {
+            crate::PyroscopeError::AdHoc(e) => {
                 log::error!(target: LOG_TAG, "parse_http_headers_json {}", e);
             }
             _ => {}
         },
     }
 
-    pyroscope::ffikit::run(agent_builder).is_ok()
+    crate::ffikit::run(agent_builder).is_ok()
 }
 
 #[no_mangle]
 pub extern "C" fn drop_agent() -> bool {
-    pyroscope::ffikit::send(pyroscope::ffikit::Signal::Kill).is_ok()
+    crate::ffikit::send(crate::ffikit::Signal::Kill).is_ok()
 }
 
 #[no_mangle]
@@ -216,8 +242,8 @@ pub unsafe extern "C" fn add_thread_tag(key: *const c_char, value: *const c_char
         .unwrap()
         .to_owned();
 
-    pyroscope::ffikit::send(pyroscope::ffikit::Signal::AddThreadTag(
-        backend::self_thread_id(),
+    crate::ffikit::send(crate::ffikit::Signal::AddThreadTag(
+        rbspy_backend::self_thread_id(),
         Tag { key, value },
     ))
     .is_ok()
@@ -233,8 +259,8 @@ pub unsafe extern "C" fn remove_thread_tag(key: *const c_char, value: *const c_c
         .unwrap()
         .to_owned();
 
-    pyroscope::ffikit::send(pyroscope::ffikit::Signal::RemoveThreadTag(
-        backend::self_thread_id(),
+    crate::ffikit::send(crate::ffikit::Signal::RemoveThreadTag(
+        rbspy_backend::self_thread_id(),
         Tag { key, value },
     ))
     .is_ok()
@@ -260,7 +286,7 @@ fn string_to_tags(tags: &str) -> Vec<(&str, &str)> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use pyroscope::backend::{Report, StackFrame, StackTrace};
+    use crate::backend::{Report, StackFrame, StackTrace};
     use std::collections::HashMap;
 
     fn report_with_filename(filename: &str) -> Report {