Merge remote-tracking branch 'origin/0.5.4' into 0.5.4

Author: omarabid

Cargo.toml

@@ -56,6 +56,7 @@ log = "0.4"
 names = "0.13.0"
 reqwest = { version = "0.11", features = ["blocking", "rustls-tls-native-roots"]}
 url = "2.2.2"
+libflate = "1.2.0"
 libc = "^0.2.124"
 
 [dev-dependencies]

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

@@ -1,12 +1,15 @@
-use ffikit::Signal;
-use pyroscope::backend::{Report, StackFrame, Tag};
-use pyroscope::PyroscopeAgent;
-use pyroscope_rbspy::{rbspy_backend, RbspyConfig};
 use std::collections::hash_map::DefaultHasher;
 use std::env;
 use std::ffi::CStr;
 use std::hash::Hasher;
 use std::os::raw::c_char;
+use std::str::FromStr;
+
+use ffikit::Signal;
+use pyroscope_rbspy::{rbspy_backend, RbspyConfig};
+
+use pyroscope::{pyroscope::Compression, PyroscopeAgent};
+use pyroscope::backend::{Report, StackFrame, Tag};
 
 pub fn transform_report(report: Report) -> Report {
     let cwd = env::current_dir().unwrap();
@@ -104,7 +107,7 @@ 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,
+    report_thread_id: bool, tags: *const c_char, compression: *const c_char
 ) -> bool {
     // Initialize FFIKit
     let recv = ffikit::initialize_ffi().unwrap();
@@ -129,6 +132,13 @@ pub extern "C" fn initialize_agent(
         .unwrap()
         .to_string();
 
+    let compression_string = unsafe { CStr::from_ptr(compression) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let compression = Compression::from_str(&compression_string);
+
     let pid = std::process::id();
 
     let rbspy_config = RbspyConfig::new(pid.try_into().unwrap())
@@ -152,6 +162,10 @@ pub extern "C" fn initialize_agent(
         agent_builder = agent_builder.auth_token(auth_token);
     }
 
+    if let Ok(compression) = compression {
+        agent_builder = agent_builder.compression(compression);
+    }
+
     let agent = agent_builder.build().unwrap();
 
     let agent_running = agent.start().unwrap();

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], :bool
+    attach_function :initialize_agent, [:string, :string, :string, :int, :bool, :bool, :bool, :bool, :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,8 +22,10 @@ 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) 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) do
     def initialize(*)
+      super
+      # defaults:
       self.application_name = ''
       self.server_address = 'http://localhost:4040'
       self.auth_token = ''
@@ -34,7 +36,7 @@ def initialize(*)
       self.report_thread_id = false
       self.log_level = 'error'
       self.tags = {}
-      super
+      self.compression = 'gzip'
     end
   end
 
@@ -63,10 +65,11 @@ def configure
 
       # Initialize Logging
       Rust.initialize_logging(@log_level)
-      
+
 
       # initialize Pyroscope Agent
       Rust.initialize_agent(
+        # these are defaults in case user-provided values are nil:
         @config.app_name || @config.application_name || "",
         @config.server_address || "",
         @config.auth_token || "",
@@ -75,7 +78,8 @@ def configure
         @config.oncpu || false,
         @config.report_pid || false,
         @config.report_thread_id || false,
-        tags_to_string(@config.tags || {})
+        tags_to_string(@config.tags || {}),
+        @config.compression || ""
       )
     end
 

pyroscope_ffi/ruby/lib/pyroscope/version.rb

@@ -1,3 +1,3 @@
 module Pyroscope
-  VERSION = '0.3.2'.freeze
+  VERSION = '0.4.0'.freeze
 end

src/pyroscope.rs

@@ -1,6 +1,7 @@
 use std::{
     collections::HashMap,
     marker::PhantomData,
+    str::FromStr,
     sync::{
         mpsc::{self, Sender},
         Arc, Condvar, Mutex,
@@ -18,6 +19,7 @@ use crate::{
 };
 
 use crate::backend::BackendImpl;
+use crate::pyroscope::Compression::GZIP;
 
 const LOG_TAG: &str = "Pyroscope::Agent";
 
@@ -44,6 +46,8 @@ pub struct PyroscopeConfig {
     pub auth_token: Option<String>,
     /// Function to apply
     pub func: Option<fn(Report) -> Report>,
+    /// Pyroscope http request body compression
+    pub compression: Option<Compression>,
 }
 
 impl Default for PyroscopeConfig {
@@ -59,6 +63,7 @@ impl Default for PyroscopeConfig {
             spy_name: "undefined".to_string(),
             auth_token: None,
             func: None,
+            compression: None,
         }
     }
 }
@@ -80,6 +85,7 @@ impl PyroscopeConfig {
             spy_name: String::from("undefined"), // Spy Name should be set by the backend
             auth_token: None,             // No authentication token
             func: None,                   // No function
+            compression: None,
         }
     }
 
@@ -150,6 +156,22 @@ impl PyroscopeConfig {
             ..self
         }
     }
+
+
+    /// Set the http request body compression.
+    ///
+    /// # Example
+    /// ```ignore
+    /// use pyroscope::pyroscope::PyroscopeConfig;
+    /// let config = PyroscopeConfig::new("http://localhost:8080", "my-app")
+    ///     .compression(GZIP);
+    /// ```
+    pub fn compression(self, compression: Compression) -> Self {
+        Self {
+            compression: Some(compression),
+            ..self
+        }
+    }
 }
 
 /// PyroscopeAgent Builder
@@ -288,6 +310,22 @@ impl PyroscopeAgentBuilder {
         }
     }
 
+    /// Set the http request body compression.
+    ///
+    /// # Example
+    /// ```ignore
+    /// use pyroscope::pyroscope::PyroscopeConfig;
+    /// let agent = PyroscopeAgentBuilder::new("http://localhost:8080", "my-app")
+    ///     .compression(GZIP)
+    ///     .build();
+    /// ```
+    pub fn compression(self, compression: Compression) -> Self {
+        Self {
+            config: self.config.compression(compression),
+            ..self
+        }
+    }
+
     /// Initialize the backend, timer and return a PyroscopeAgent with Ready
     /// state. While you can call this method, you should call it through the
     /// `PyroscopeAgent.build()` method.
@@ -344,6 +382,21 @@ impl PyroscopeAgentBuilder {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum Compression {
+    GZIP
+}
+
+impl FromStr for Compression {
+    type Err = ();
+    fn from_str(input: &str) -> std::result::Result<Compression, Self::Err> {
+        match input {
+            "gzip" => Ok(GZIP),
+            _ => Err(()),
+        }
+    }
+}
+
 /// This trait is used to encode the state of the agent.
 pub trait PyroscopeAgentState {}
 

src/session.rs

@@ -2,13 +2,15 @@ use std::{
     sync::mpsc::{sync_channel, Receiver, SyncSender},
     thread::{self, JoinHandle},
     time::Duration,
+    io::Write,
 };
 
 use reqwest::Url;
+use libflate::gzip::Encoder;
 
 use crate::{
     backend::Report,
-    pyroscope::PyroscopeConfig,
+    pyroscope::{PyroscopeConfig, Compression},
     utils::{get_time_range, merge_tags_with_app_name},
     Result,
 };
@@ -192,6 +194,15 @@ impl Session {
         if let Some(auth_token) = self.config.auth_token.clone() {
             req_builder = req_builder.bearer_auth(auth_token);
         }
+        let body = match &self.config.compression {
+            None => report_u8,
+            Some(Compression::GZIP) => {
+                req_builder = req_builder.header("Content-Encoding", "gzip");
+                let mut encoder = Encoder::new(Vec::new()).unwrap();
+                encoder.write_all(&report_u8).unwrap();
+                encoder.finish().into_result().unwrap()
+            }
+        };
 
         // Send the request
         req_builder
@@ -203,7 +214,7 @@ impl Session {
                 ("sampleRate", &format!("{}", self.config.sample_rate)),
                 ("spyName", self.config.spy_name.as_str()),
             ])
-            .body(report_u8)
+            .body(body)
             .timeout(Duration::from_secs(10))
             .send()?;
         Ok(())