pyroscope 1.0.9 → 1.1.0

data/ext/rbspy/src/{backend.rs → rbspy_backend.rs}

@@ -1,4 +1,4 @@
-use pyroscope::{
+use crate::{
     backend::{
         Backend, BackendConfig, Report, ReportBatch, ReportData, StackBuffer, StackFrame,
         StackTrace, ThreadTag, ThreadTagsSet,
@@ -192,8 +192,8 @@ impl From<(rbspy::StackTrace, &BackendConfig)> for StackTraceWrapper {
     }
 }
 
-pub fn self_thread_id() -> pyroscope::ThreadId {
+pub fn self_thread_id() -> crate::ThreadId {
     // for rbspy we use pthread_t as thread id
     // https://github.com/ruby/ruby/blob/54a74c42033e42869e69e7dc9e67efa1faf225be/include/ruby/thread_native.h#L41
-    pyroscope::ThreadId::pthread_self()
+    crate::ThreadId::pthread_self()
 }

data/ext/rbspy/src/session.rs

@@ -0,0 +1,275 @@
+use std::{
+    io::Write,
+    sync::mpsc::{sync_channel, Receiver, SyncSender},
+    thread::{self, JoinHandle},
+    time::Duration,
+};
+
+use crate::encode::gen::push::{PushRequest, RawProfileSeries, RawSample};
+use crate::encode::gen::types::LabelPair;
+use crate::{
+    backend::{Report, ReportBatch, ReportData},
+    encode::pprof,
+    pyroscope::PyroscopeConfig,
+    utils::get_time_range,
+    Result,
+};
+use libflate::gzip::Encoder;
+use prost::Message;
+use reqwest::Url;
+use uuid::Uuid;
+
+const LOG_TAG: &str = "Pyroscope::Session";
+const LABEL_SCOPE_NAME: &str = "otel.scope.name";
+const LABEL_SCOPE_VERSION: &str = "otel.scope.version";
+const LABEL_PROCESS_RUNTIME_NAME: &str = "process.runtime.name";
+const LABEL_PROCESS_RUNTIME_VERSION: &str = "process.runtime.version";
+const LABEL_SERVICE_NAME: &str = "service_name";
+const LABEL_PROFILE_NAME: &str = "__name__";
+const SCOPE_NAME: &str = "com.grafana.pyroscope/ruby";
+
+/// Session Signal
+///
+/// This enum is used to send data to the session thread. It can also kill the session thread.
+pub enum SessionSignal {
+    /// Send session data to the session thread.
+    Session(Box<Session>),
+    /// Kill the session thread.
+    Kill,
+}
+
+/// Manage sessions and send data to the server.
+#[derive(Debug)]
+pub struct SessionManager {
+    /// The SessionManager thread.
+    pub handle: Option<JoinHandle<Result<()>>>,
+    /// Channel to send data to the SessionManager thread.
+    pub tx: SyncSender<SessionSignal>,
+}
+
+impl SessionManager {
+    /// Create a new SessionManager
+    pub fn new() -> Result<Self> {
+        log::info!(target: LOG_TAG, "Creating SessionManager");
+
+        // Create a channel for sending and receiving sessions
+        let (tx, rx): (SyncSender<SessionSignal>, Receiver<SessionSignal>) = sync_channel(10);
+
+        // Create a thread for the SessionManager
+        let handle = Some(thread::spawn(move || {
+            log::trace!(target: LOG_TAG, "Started");
+            let client = reqwest::blocking::Client::new();
+            while let Ok(signal) = rx.recv() {
+                match signal {
+                    SessionSignal::Session(session) => {
+                        // Send the session
+                        // Matching is done here (instead of ?) to avoid breaking
+                        // the SessionManager thread if the server is not available.
+                        match (*session).push(&client) {
+                            Ok(_) => log::trace!("SessionManager - Session sent"),
+                            Err(e) => log::error!("SessionManager - Failed to send session: {e}"),
+                        }
+                    }
+                    SessionSignal::Kill => {
+                        // Kill the session manager
+                        log::trace!(target: LOG_TAG, "Kill signal received");
+                        return Ok(());
+                    }
+                }
+            }
+            Ok(())
+        }));
+
+        Ok(SessionManager { handle, tx })
+    }
+
+    /// Push a new session into the SessionManager
+    pub fn push(&self, session: SessionSignal) -> Result<()> {
+        // Push the session into the SessionManager
+        self.tx.send(session)?;
+
+        log::trace!(target: LOG_TAG, "SessionSignal pushed");
+
+        Ok(())
+    }
+}
+
+pub struct Session {
+    pub config: PyroscopeConfig,
+    pub batch: ReportBatch,
+    // unix time todo remove comment, use types
+    pub from: u64,
+    // unix time todo remove comment, use types
+    pub until: u64,
+}
+
+impl Session {
+    /// Create a new Session
+    /// # Example
+    /// ```ignore
+    /// let config = PyroscopeConfig::new("https://localhost:8080", "my-app", 100, "pyspy", "0.8.16");
+    /// let report = vec![1, 2, 3];
+    /// let until = 154065120;
+    /// let session = Session::new(until, config, report)?;
+    /// ```
+    pub fn new(until: u64, config: PyroscopeConfig, batch: ReportBatch) -> Result<Self> {
+        log::info!(target: LOG_TAG, "Creating Session");
+
+        // get_time_range should be used with "from". We balance this by reducing
+        // 10s from the returned range.
+        let time_range = get_time_range(until)?;
+
+        Ok(Self {
+            config,
+            batch,
+            from: time_range.from - 10,
+            until: time_range.until - 10,
+        })
+    }
+
+    fn push(self, client: &reqwest::blocking::Client) -> Result<()> {
+        log::info!(target: LOG_TAG, "Sending Session: {} - {}", self.from, self.until);
+
+        let ReportBatch { profile_type, data } = self.batch;
+
+        let raw_profile = match data {
+            ReportData::RawPprof(pprof_bytes) => {
+                if self.config.func.is_some() {
+                    log::warn!(target: LOG_TAG, "report transform function is not supported with raw pprof backends (e.g. jemalloc)");
+                }
+                pprof_bytes
+            }
+            ReportData::Reports(reports) => {
+                let transformed: Vec<Report>;
+                let encode_input = match self.config.func {
+                    None => &reports,
+                    Some(f) => {
+                        transformed = reports.iter().map(|r| f(r.to_owned())).collect();
+                        &transformed
+                    }
+                };
+                pprof::encode(
+                    encode_input,
+                    self.config.sample_rate,
+                    self.from * 1_000_000_000,
+                    (self.until - self.from) * 1_000_000_000,
+                )
+                .encode_to_vec()
+            }
+        };
+
+        let labels = labels_for_profile(&self.config, profile_type);
+        let req = PushRequest {
+            series: vec![RawProfileSeries {
+                labels,
+                samples: vec![RawSample {
+                    raw_profile,
+                    id: Uuid::new_v4().to_string(),
+                }],
+            }],
+        };
+
+        let req = Self::gzip(&req.encode_to_vec())?;
+
+        let mut url = Url::parse(&self.config.url)?;
+        url.path_segments_mut()
+            .unwrap()
+            .push("push.v1.PusherService")
+            .push("Push");
+
+        let mut req_builder = client
+            .post(url.as_str())
+            .header(
+                "User-Agent",
+                format!(
+                    "pyroscope-rs/{}/{} reqwest",
+                    self.config.spy_name, self.config.spy_version
+                ),
+            )
+            .header("Content-Type", "application/proto")
+            .header("Content-Encoding", "gzip");
+
+        if let Some(basic_auth) = &self.config.basic_auth {
+            req_builder = req_builder.basic_auth(
+                basic_auth.username.clone(),
+                Some(basic_auth.password.clone()),
+            );
+        }
+        if let Some(tenant_id) = &self.config.tenant_id {
+            req_builder = req_builder.header("X-Scope-OrgID", tenant_id);
+        }
+        for (k, v) in &self.config.http_headers {
+            req_builder = req_builder.header(k, v);
+        }
+
+        let mut response = req_builder
+            .body(req)
+            .timeout(Duration::from_secs(10))
+            .send()?;
+
+        let status = response.status();
+
+        if status.is_success() {
+            let mut sink = std::io::sink();
+            _ = response.copy_to(&mut sink);
+        } else {
+            let resp = response.text();
+            let resp = match &resp {
+                Ok(t) => t,
+                Err(_) => "",
+            };
+            log::error!(target: LOG_TAG, "Sending Session failed {} {}", status.as_u16(), resp);
+        }
+        Ok(())
+    }
+
+    fn gzip(report: &[u8]) -> Result<Vec<u8>> {
+        let mut encoder = Encoder::new(Vec::new())?;
+        encoder.write_all(report)?;
+        let compressed_data = encoder.finish().into_result()?;
+        Ok(compressed_data)
+    }
+}
+
+fn labels_for_profile(config: &PyroscopeConfig, profile_type: String) -> Vec<LabelPair> {
+    let mut labels: Vec<LabelPair> = Vec::with_capacity(6 + config.tags.len());
+    labels.push(LabelPair {
+        name: LABEL_PROFILE_NAME.to_string(),
+        value: profile_type,
+    });
+    for (k, v) in &config.tags {
+        if k == LABEL_PROFILE_NAME {
+            continue;
+        }
+        labels.push(LabelPair {
+            name: k.clone(),
+            value: v.clone(),
+        })
+    }
+    push_label_if_absent(&mut labels, LABEL_SERVICE_NAME, &config.application_name);
+    push_label_if_absent(&mut labels, LABEL_SCOPE_NAME, SCOPE_NAME);
+    push_label_if_absent(&mut labels, LABEL_SCOPE_VERSION, &config.spy_version);
+    push_label_if_absent(
+        &mut labels,
+        LABEL_PROCESS_RUNTIME_NAME,
+        &config.runtime_name,
+    );
+    push_label_if_absent(
+        &mut labels,
+        LABEL_PROCESS_RUNTIME_VERSION,
+        &config.runtime_version,
+    );
+
+    labels
+}
+
+fn push_label_if_absent(labels: &mut Vec<LabelPair>, name: &str, value: &str) {
+    if labels.iter().any(|label| label.name == name) {
+        return;
+    }
+
+    labels.push(LabelPair {
+        name: name.to_string(),
+        value: value.to_string(),
+    });
+}

data/ext/rbspy/src/timer/epoll.rs

@@ -0,0 +1,268 @@
+use super::TimerSignal;
+use crate::{
+    utils::{check_err, get_time_range},
+    PyroscopeError, Result,
+};
+
+use std::sync::{
+    mpsc::{channel, Sender},
+    Arc, Mutex,
+};
+use std::{
+    thread::{self, JoinHandle},
+    time::Duration,
+};
+
+const LOG_TAG: &str = "Pyroscope::Timer";
+
+/// A thread that sends a notification every 10th second
+///
+/// Timer will send an event to attached listeners (mpsc::Sender) every 10th
+/// second (...10, ...20, ...)
+///
+/// The Timer thread will run continously until all Senders are dropped.
+/// The Timer thread will be joined when all Senders are dropped.
+
+#[derive(Debug)]
+pub struct Timer {
+    /// A vector to store listeners
+    txs: Arc<Mutex<Vec<Sender<TimerSignal>>>>,
+
+    /// Thread handle
+    pub handle: Option<JoinHandle<Result<()>>>,
+}
+
+impl Timer {
+    /// Initialize Timer and run a thread to send events to attached listeners
+    pub fn initialize(cycle: Duration) -> Result<Self> {
+        log::info!(target: LOG_TAG, "Initializing Timer");
+
+        let txs = Arc::new(Mutex::new(Vec::new()));
+
+        // Add a dummy tx so the below thread does not terminate early
+        let (tx, _rx) = channel();
+        txs.lock()?.push(tx);
+
+        let timer_fd = Timer::set_timerfd(cycle)?;
+        let epoll_fd = Timer::create_epollfd(timer_fd)?;
+
+        let handle = Some({
+            let txs = txs.clone();
+            thread::spawn(move || {
+                loop {
+                    // Exit thread if there are no listeners
+                    if txs.lock()?.is_empty() {
+                        // Close file descriptors
+                        unsafe { libc::close(timer_fd) };
+                        unsafe { libc::close(epoll_fd) };
+
+                        log::info!(target: LOG_TAG, "Timer thread terminated");
+
+                        return Ok::<_, PyroscopeError>(());
+                    }
+
+                    // Fire @ 10th sec
+                    let res = Timer::epoll_wait(timer_fd, epoll_fd);
+                    if matches!(&res, Err(PyroscopeError::Io(err)) if err.kind() == std::io::ErrorKind::Interrupted)
+                    {
+                        continue;
+                    }
+                    res?;
+
+                    // Get the current time range
+                    let from = TimerSignal::NextSnapshot(get_time_range(0)?.from);
+
+                    log::trace!(target: LOG_TAG, "Timer fired @ {from}");
+
+                    // Iterate through Senders
+                    txs.lock()?.iter().for_each(|tx| {
+                        // Send event to attached Sender
+                        match tx.send(from) {
+                            Ok(_) => {
+                                log::trace!(target: LOG_TAG, "Sent event to listener @ {:?}", &tx)
+                            }
+                            Err(_e) => {} // There could be a less confusing message, or this
+                                          // refactored to avoid a first sender
+                                          //log::warn!(
+                                          //target: LOG_TAG,
+                                          //"Failed to send event to listener @ {:?} - {}",
+                                          //&tx,
+                                          //e
+                                          //),
+                        }
+                    });
+                }
+            })
+        });
+
+        Ok(Self { handle, txs })
+    }
+
+    /// create and set a timer file descriptor
+    fn set_timerfd(cycle: Duration) -> Result<libc::c_int> {
+        // Set the timer to use the system time.
+        let clockid: libc::clockid_t = libc::CLOCK_REALTIME;
+        // Non-blocking file descriptor
+        let clock_flags: libc::c_int = libc::TFD_NONBLOCK;
+
+        // Create timer fd
+        let tfd = timerfd_create(clockid, clock_flags)?;
+
+        // Get the next event time
+        let first_fire = get_time_range(0)?.until;
+
+        // new_value sets the Timer
+        let mut new_value = libc::itimerspec {
+            it_interval: libc::timespec {
+                tv_sec: cycle.as_secs() as libc::c_long,
+                tv_nsec: cycle.subsec_nanos() as libc::c_long,
+            },
+            it_value: libc::timespec {
+                tv_sec: first_fire as libc::c_long,
+                tv_nsec: 0,
+            },
+        };
+
+        // Empty itimerspec object
+        let mut old_value = libc::itimerspec {
+            it_interval: libc::timespec {
+                tv_sec: 0,
+                tv_nsec: 0,
+            },
+            it_value: libc::timespec {
+                tv_sec: 0,
+                tv_nsec: 0,
+            },
+        };
+
+        let set_flags = libc::TFD_TIMER_ABSTIME;
+
+        // Set the timer
+        timerfd_settime(tfd, set_flags, &mut new_value, &mut old_value)?;
+
+        // Return file descriptor
+        Ok(tfd)
+    }
+
+    /// Create a new epoll file descriptor and add the timer to its interests
+    fn create_epollfd(timer_fd: libc::c_int) -> Result<libc::c_int> {
+        // create a new epoll fd
+        let epoll_fd = epoll_create1(0)?;
+
+        // event to pull
+        let mut event = libc::epoll_event {
+            events: libc::EPOLLIN as u32,
+            u64: 1,
+        };
+
+        let epoll_flags = libc::EPOLL_CTL_ADD;
+
+        // add event to the epoll
+        epoll_ctl(epoll_fd, epoll_flags, timer_fd, &mut event)?;
+
+        // return epoll fd
+        Ok(epoll_fd)
+    }
+
+    /// Wait for an event on the epoll file descriptor
+    fn epoll_wait(timer_fd: libc::c_int, epoll_fd: libc::c_int) -> Result<()> {
+        // vector to store events
+        let mut events = Vec::with_capacity(1);
+
+        // wait for the timer to fire an event. This is function will block.
+        unsafe {
+            epoll_wait(epoll_fd, events.as_mut_ptr(), 1, -1)?;
+        }
+
+        // read the value from the timerfd. This is required to re-arm the timer.
+        let mut buffer: u64 = 0;
+        let bufptr: *mut _ = &mut buffer;
+        unsafe {
+            read(timer_fd, bufptr as *mut libc::c_void, 8)?;
+        }
+
+        Ok(())
+    }
+
+    /// Attach an mpsc::Sender to Timer
+    ///
+    /// Timer will dispatch an event with the timestamp of the current instant,
+    /// every 10th second to all attached senders
+    pub fn attach_listener(&mut self, tx: Sender<TimerSignal>) -> Result<()> {
+        // Push Sender to a Vector of Sender(s)
+        let txs = Arc::clone(&self.txs);
+        txs.lock()?.push(tx);
+
+        Ok(())
+    }
+
+    /// Clear the listeners (txs) from Timer. This will shutdown the Timer thread
+    pub fn drop_listeners(&mut self) -> Result<()> {
+        let txs = Arc::clone(&self.txs);
+        txs.lock()?.clear();
+
+        Ok(())
+    }
+}
+
+// Wrapper for libc functions.
+//
+// Error wrapper for some libc functions used by the library. This only does
+// Error (-1 return) wrapping. Alternatively, the nix crate could be used
+// instead of expanding this wrappers (if more functions and types are used
+// from libc)
+
+// libc::timerfd wrapper
+pub fn timerfd_create(clockid: libc::clockid_t, clock_flags: libc::c_int) -> Result<i32> {
+    check_err(unsafe { libc::timerfd_create(clockid, clock_flags) })
+}
+
+/// libc::timerfd_settime wrapper
+pub fn timerfd_settime(
+    timer_fd: i32,
+    set_flags: libc::c_int,
+    new_value: &mut libc::itimerspec,
+    old_value: &mut libc::itimerspec,
+) -> Result<()> {
+    check_err(unsafe { libc::timerfd_settime(timer_fd, set_flags, new_value, old_value) })?;
+    Ok(())
+}
+
+/// libc::epoll_create1 wrapper
+pub fn epoll_create1(epoll_flags: libc::c_int) -> Result<i32> {
+    check_err(unsafe { libc::epoll_create1(epoll_flags) })
+}
+
+/// libc::epoll_ctl wrapper
+pub fn epoll_ctl(
+    epoll_fd: i32,
+    epoll_flags: libc::c_int,
+    timer_fd: i32,
+    event: &mut libc::epoll_event,
+) -> Result<()> {
+    check_err(unsafe { libc::epoll_ctl(epoll_fd, epoll_flags, timer_fd, event) })?;
+    Ok(())
+}
+
+/// libc::epoll_wait wrapper
+///
+/// # Safety
+/// This function is a wrapper for libc::epoll_wait.
+pub unsafe fn epoll_wait(
+    epoll_fd: i32,
+    events: *mut libc::epoll_event,
+    maxevents: libc::c_int,
+    timeout: libc::c_int,
+) -> Result<()> {
+    check_err(libc::epoll_wait(epoll_fd, events, maxevents, timeout))?;
+    Ok(())
+}
+
+/// libc::read wrapper
+///
+/// # Safety
+/// This function is a wrapper for libc::read.
+pub unsafe fn read(timer_fd: i32, bufptr: *mut libc::c_void, count: libc::size_t) -> Result<()> {
+    check_err(libc::read(timer_fd, bufptr, count))?;
+    Ok(())
+}