imp(rbspy): change type implementation to pyroscope Backend types

Author: omarabid

examples/internal/rbspy-connect.rs

@@ -36,16 +36,7 @@ fn main() -> Result<()> {
     agent.start()?;
 
     // Profile for around 1 minute
-    std::thread::sleep(std::time::Duration::from_secs(30));
-
-    // Stop Agent
-    agent.stop()?;
-
-    std::thread::sleep(std::time::Duration::from_secs(30));
-
-    agent.start()?;
-
-    std::thread::sleep(std::time::Duration::from_secs(30));
+    std::thread::sleep(std::time::Duration::from_secs(120));
 
     agent.stop()?;
 

pyroscope_backends/pyroscope_rbspy/src/lib.rs

@@ -1,11 +1,10 @@
-use pyroscope::backend::{Backend, State};
+use pyroscope::backend::{Backend, Report, StackFrame, StackTrace, State};
 use pyroscope::error::{PyroscopeError, Result};
 
-use rbspy::{sampler::Sampler, ui::output::Outputter, StackFrame, StackTrace};
+use rbspy::sampler::Sampler;
 
-use std::collections::HashMap;
-use std::io::Write;
 use std::sync::mpsc::{channel, sync_channel, Receiver, Sender, SyncSender};
+use std::sync::{Arc, Mutex};
 
 /// Rbspy Configuration
 #[derive(Debug)]
@@ -79,9 +78,11 @@ pub struct Rbspy {
     /// Rbspy Sampler
     sampler: Option<Sampler>,
     /// StackTrace Receiver
-    stack_receiver: Option<Receiver<StackTrace>>,
+    stack_receiver: Option<Receiver<rbspy::StackTrace>>,
     /// Error Receiver
     error_receiver: Option<Receiver<std::result::Result<(), anyhow::Error>>>,
+    /// Profiling buffer
+    buffer: Arc<Mutex<Report>>,
 }
 
 impl std::fmt::Debug for Rbspy {
@@ -98,6 +99,7 @@ impl Rbspy {
             error_receiver: None,
             state: State::Uninitialized,
             config,
+            buffer: Arc::new(Mutex::new(Report::default())),
         }
     }
 }
@@ -159,8 +161,10 @@ impl Backend for Rbspy {
         let queue_size: usize = self.config.sample_rate as usize * 10 * 100;
 
         // Channel for StackTraces generated by the RubySpy Sampler
-        let (stack_sender, stack_receiver): (SyncSender<StackTrace>, Receiver<StackTrace>) =
-            sync_channel(queue_size);
+        let (stack_sender, stack_receiver): (
+            SyncSender<rbspy::StackTrace>,
+            Receiver<rbspy::StackTrace>,
+        ) = sync_channel(queue_size);
 
         // Set Error and Stack Receivers
         self.stack_receiver = Some(stack_receiver);
@@ -207,6 +211,9 @@ impl Backend for Rbspy {
             return Err(PyroscopeError::new("Rbspy: Backend is not Running"));
         }
 
+        // Get an Arc reference to the Report Buffer
+        let buffer = self.buffer.clone();
+
         // Send Errors to Log
         let errors = self
             .error_receiver
@@ -229,150 +236,61 @@ impl Backend for Rbspy {
             .ok_or_else(|| PyroscopeError::new("Rbspy: StackTrace receiver is not set"))?
             .try_iter();
 
-        // Create a new OutputFormat (collapsed). This is an object provided by rbspy.
-        // The argument should be ignored.
-        let mut outputter = RbspyFormatter::default();
-
         // Iterate over the StackTrace
         for trace in stack_trace {
-            // Write the StackTrace to the OutputFormat
-            outputter.record(&trace).map_err(|e| {
-                PyroscopeError::new(&format!("Rbspy: Error in OutputFormat: {}", e))
-            })?;
+            // convert StackTrace
+            let own_trace: StackTrace = Into::<StackTraceWrapper>::into(trace).into();
+            buffer.lock()?.record(own_trace)?;
         }
 
-        // buffer to store the output
-        let mut buffer: Vec<u8> = Vec::new();
-
-        // Create a new writer
-        let mut writer = RbspyWriter::new(&mut buffer);
+        let v8: Vec<u8> = buffer.lock()?.to_string().into_bytes();
 
-        // Push the outputter into our writer
-        outputter
-            .complete(&mut writer)
-            .map_err(|e| PyroscopeError::new(&format!("Rbspy: Error in OutputFormat: {}", e)))?;
-
-        // Flush the Writer
-        writer.flush()?;
+        buffer.lock()?.clear();
 
         // Return the writer's buffer
-        Ok(buffer)
+        Ok(v8)
     }
 }
 
-/// Rbspy Formatter
-#[derive(Default)]
-pub struct RbspyFormatter(HashMap<String, usize>);
+struct StackFrameWrapper(StackFrame);
 
-impl Outputter for RbspyFormatter {
-    fn record(&mut self, stack: &StackTrace) -> std::result::Result<(), anyhow::Error> {
-        let frame = stack
-            .iter()
-            .rev()
-            .map(|frame| format!("{}", StackFrameFormatter(frame)))
-            .collect::<Vec<String>>()
-            .join(";");
-
-        *self.0.entry(frame).or_insert(0) += 1;
-
-        Ok(())
-    }
-    fn complete(&mut self, writer: &mut dyn Write) -> std::result::Result<(), anyhow::Error> {
-        if self.0.is_empty() {
-            log::info!("Rbspy: No StackTraces reported");
-
-            return Ok(());
-        } else {
-            self.0
-                .iter()
-                .map(|(frame, count)| format!("{} {}", frame, count))
-                .collect::<Vec<String>>()
-                .iter()
-                .try_for_each(|line| writeln!(writer, "{}", line))?;
-        }
-        Ok(())
+impl From<StackFrameWrapper> for StackFrame {
+    fn from(frame: StackFrameWrapper) -> Self {
+        frame.0
     }
 }
 
-/// Custom Formatter for Rbspy StackFrames
-pub struct StackFrameFormatter<'a>(&'a StackFrame);
-
-impl<'a> std::fmt::Display for StackFrameFormatter<'a> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "{}:{} - {}",
-            self.0.relative_path, self.0.lineno, self.0.name
-        )
+impl From<rbspy::StackFrame> for StackFrameWrapper {
+    fn from(frame: rbspy::StackFrame) -> Self {
+        StackFrameWrapper(StackFrame {
+            module: None,
+            name: Some(frame.name),
+            filename: Some(frame.relative_path.clone()),
+            relative_path: Some(frame.relative_path),
+            absolute_path: frame.absolute_path,
+            line: Some(frame.lineno),
+        })
     }
 }
 
-#[cfg(test)]
-mod test_stack_frame_formatter {
-    use super::{StackFrame, StackFrameFormatter};
-
-    #[test]
-    fn test_stack_frame_formatter() {
-        let frame = StackFrame {
-            absolute_path: Some("".to_string()),
-            relative_path: "test.rb".to_string(),
-            lineno: 1,
-            name: "test".to_string(),
-        };
-        let formatter = StackFrameFormatter(&frame);
-        assert_eq!(formatter.to_string(), "test.rb:1 - test");
-    }
-}
-
-/// Rubyspy Writer
-/// This object is used to write the output of the rbspy sampler to a data buffer
-struct RbspyWriter<'a> {
-    data: Vec<u8>,
-    buffer: &'a mut Vec<u8>,
-}
-
-impl<'a> RbspyWriter<'a> {
-    /// Create a new RbspyWriter
-    pub fn new(buffer: &'a mut Vec<u8>) -> Self {
-        RbspyWriter {
-            data: Vec::new(),
-            buffer,
-        }
-    }
-}
-
-/// Implement Writer for Rbspy
-impl<'a> std::io::Write for RbspyWriter<'a> {
-    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        // push the data to the buffer
-        self.data.extend_from_slice(buf);
+struct StackTraceWrapper(StackTrace);
 
-        // return the number of bytes written
-        Ok(buf.len())
-    }
-
-    fn flush(&mut self) -> std::io::Result<()> {
-        // flush the buffer
-        self.buffer.extend_from_slice(&self.data);
-
-        Ok(())
+impl From<StackTraceWrapper> for StackTrace {
+    fn from(trace: StackTraceWrapper) -> Self {
+        trace.0
     }
 }
 
-#[cfg(test)]
-mod test_rbspy_writer {
-    use super::RbspyWriter;
-    use std::io::Write;
-
-    #[test]
-    fn test_rbspy_writer() {
-        let mut buffer: Vec<u8> = Vec::new();
-        let mut writer = RbspyWriter::new(&mut buffer);
-
-        writer.write(b"hello").unwrap();
-        writer.write(b"world").unwrap();
-        writer.flush().unwrap();
-
-        assert_eq!(buffer, b"helloworld".to_vec());
+impl From<rbspy::StackTrace> for StackTraceWrapper {
+    fn from(trace: rbspy::StackTrace) -> Self {
+        StackTraceWrapper(StackTrace {
+            pid: trace.pid.map(|pid| pid as u32),
+            thread_id: trace.thread_id.map(|id| id as u64),
+            thread_name: None,
+            frames: trace
+                .iter()
+                .map(|frame| Into::<StackFrameWrapper>::into(frame.clone()).into())
+                .collect(),
+        })
     }
 }