allow configurable accumulation_cycle

Author: drahnr

Cargo.toml

@@ -23,6 +23,7 @@ libc = "^0.2.66"
 [dev-dependencies]
 tokio = { version = "1.13", features = ["full"] }
 pretty_env_logger = "0.4.0"
+assert_matches = "1"
 
 [profile.dev]
 opt-level=0

examples/internals-timer.rs

@@ -1,18 +1,17 @@
 extern crate pyroscope;
 
-use std::sync::mpsc::channel;
-use std::sync::mpsc::{Receiver, Sender};
+use std::sync::mpsc;
 
 use pyroscope::timer::Timer;
 
 fn main() {
     // Initialize the Timer
-    let mut timer = Timer::default().initialize().unwrap();
+    let mut timer = Timer::initialize(std::time::Duration::from_secs(10)).unwrap();
 
     // Create a streaming channel
-    let (tx, rx): (Sender<u64>, Receiver<u64>) = channel();
+    let (tx, rx) = mpsc::channel();
 
-    let (tx2, rx2): (Sender<u64>, Receiver<u64>) = channel();
+    let (tx2, rx2) = mpsc::channel();
 
     // Attach tx to Timer
     timer.attach_listener(tx).unwrap();

src/pyroscope.rs

@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     sync::{
-        mpsc::{channel, Receiver, Sender},
+        mpsc::{self, Sender},
         Arc, Condvar, Mutex,
     },
     thread::JoinHandle,
@@ -16,6 +16,28 @@ use crate::{
 
 const LOG_TAG: &str = "Pyroscope::Agent";
 
+
+/// A signal sent from the agent to the timer.
+///
+/// Either schedules another wake-up, or asks
+/// the timer thread to terminate.
+#[derive(Debug, Clone, Copy)]
+pub enum AgentSignal {
+    // Thread termination was requested.
+    Terminate,
+    // When to take the next snapshot using the `Backend`.
+    NextSnapshot(u64),
+}
+
+impl std::fmt::Display for AgentSignal {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Self::Terminate => write!(f, "Terminate"),
+            Self::NextSnapshot(when) => write!(f, "NextSnapshot({})", when),
+        }
+    }
+}
+
 /// Pyroscope Agent Configuration. This is the configuration that is passed to the agent.
 /// # Example
 /// ```
@@ -32,6 +54,8 @@ pub struct PyroscopeConfig {
     pub tags: HashMap<String, String>,
     /// Sample rate used in Hz
     pub sample_rate: i32,
+    /// How long to accumulate data for before sending it upstream
+    pub accumulation_cycle: std::time::Duration,
     // TODO
     // log_level
     // auth_token
@@ -52,6 +76,7 @@ impl PyroscopeConfig {
             application_name: application_name.as_ref().to_owned(),
             tags: HashMap::new(),
             sample_rate: 100i32,
+            accumulation_cycle: std::time::Duration::from_secs(10),
         }
     }
 
@@ -69,6 +94,26 @@ impl PyroscopeConfig {
         }
     }
 
+    /// Override the accumulation cycle.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # fn main() -> Result<(), PyroscopeError> {
+    /// use std::time::Duration;
+    /// use pyroscope::pyroscope::PyroscopeConfig;
+    /// let config = PyroscopeConfig::new("http://localhost:8080", "my-app")
+    ///    .accumulation_cycle(Duration::from_millis(4587))?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn accumulation_cycle(self, accumulation_cycle: std::time::Duration) -> Self {
+        Self {
+            accumulation_cycle,
+            ..self
+        }
+    }
+
     /// Set the tags
     /// # Example
     /// ```ignore
@@ -134,7 +179,9 @@ impl PyroscopeAgentBuilder {
     /// ?;
     /// ```
     pub fn backend<T>(self, backend: T) -> Self
-        where T: 'static + Backend {
+    where
+        T: 'static + Backend,
+    {
         Self {
             backend: Arc::new(Mutex::new(backend)),
             ..self
@@ -145,7 +192,7 @@ impl PyroscopeAgentBuilder {
     /// # Example
     /// ```ignore
     /// let builder = PyroscopeAgentBuilder::new("http://localhost:8080", "my-app")
-    /// .sample_rate(99)
+    /// .sample_rate(113)
     /// .build()
     /// ?;
     /// ```
@@ -156,6 +203,28 @@ impl PyroscopeAgentBuilder {
         }
     }
 
+    /// Set the accumulation cycle. Default value is 10 seconds.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # fn main() -> Result<(), PyroscopeError> {
+    /// use std::time::Duration;
+    ///
+    /// let builder = PyroscopeAgentBuilder::new("http://localhost:8080", "my-app")
+    /// .accumulation_cycle(Duration::from_secs(3))
+    /// .build()
+    /// ?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn accumulation_cycle(self, acc_cycle: impl Into<std::time::Duration>) -> Self {
+        Self {
+            config: self.config.accumulation_cycle(acc_cycle.into()),
+            ..self
+        }
+    }
+
     /// Set tags. Default is empty.
     /// # Example
     /// ```ignore
@@ -179,7 +248,7 @@ impl PyroscopeAgentBuilder {
         log::trace!(target: LOG_TAG, "Backend initialized");
 
         // Start Timer
-        let timer = Timer::default().initialize()?;
+        let timer = Timer::initialize(self.config.accumulation_cycle.clone())?;
         log::trace!(target: LOG_TAG, "Timer initialized");
 
         // Start the SessionManager
@@ -204,7 +273,7 @@ impl PyroscopeAgentBuilder {
 pub struct PyroscopeAgent {
     timer: Timer,
     session_manager: SessionManager,
-    tx: Option<Sender<u64>>,
+    tx: Option<Sender<AgentSignal>>,
     handle: Option<JoinHandle<Result<()>>>,
     running: Arc<(Mutex<bool>, Condvar)>,
 
@@ -287,7 +356,7 @@ impl PyroscopeAgent {
         *running = true;
         drop(running);
 
-        let (tx, rx): (Sender<u64>, Receiver<u64>) = channel();
+        let (tx, rx) = mpsc::channel();
         self.timer.attach_listener(tx.clone())?;
         self.tx = Some(tx);
 
@@ -298,28 +367,31 @@ impl PyroscopeAgent {
         self.handle = Some(std::thread::spawn(move || {
             log::trace!(target: LOG_TAG, "Main Thread started");
 
-            while let Ok(until) = rx.recv() {
-                log::trace!(target: LOG_TAG, "Sending session {}", until);
-
-                // Generate report from backend
-                let report = backend.lock()?.report()?;
-
-                // Send new Session to SessionManager
-                stx.send(SessionSignal::Session(Session::new(
-                    until,
-                    config.clone(),
-                    report,
-                )?))?;
-
-                if until == 0 {
-                    log::trace!(target: LOG_TAG, "Session Killed");
-
-                    let (lock, cvar) = &*pair;
-                    let mut running = lock.lock()?;
-                    *running = false;
-                    cvar.notify_one();
-
-                    return Ok(());
+            while let Ok(signal) = rx.recv() {
+                match signal {
+                    AgentSignal::NextSnapshot(until) => {
+                        log::trace!(target: LOG_TAG, "Sending session {}", until);
+
+                        // Generate report from backend
+                        let report = backend.lock()?.report()?;
+
+                        // Send new Session to SessionManager
+                        stx.send(SessionSignal::Session(Session::new(
+                            until,
+                            config.clone(),
+                            report,
+                        )?))?
+                    }
+                    AgentSignal::Terminate => {
+                        log::trace!(target: LOG_TAG, "Session Killed");
+
+                        let (lock, cvar) = &*pair;
+                        let mut running = lock.lock()?;
+                        *running = false;
+                        cvar.notify_one();
+
+                        return Ok(());
+                    }
                 }
             }
             Ok(())
@@ -345,7 +417,9 @@ impl PyroscopeAgent {
         log::debug!(target: LOG_TAG, "Stopping");
         // get tx and send termination signal
         if let Some(sender) = self.tx.take() {
-            sender.send(0)?;
+            // best effort
+            let _ = sender.send(AgentSignal::NextSnapshot(0));
+            sender.send(AgentSignal::Terminate)?;
         } else {
             log::error!("PyroscopeAgent - Missing sender")
         }

src/session.rs

@@ -84,7 +84,9 @@ impl SessionManager {
 pub struct Session {
     pub config: PyroscopeConfig,
     pub report: Vec<u8>,
+    // unix time
     pub from: u64,
+    // unix time
     pub until: u64,
 }
 

src/timer/epoll.rs

@@ -4,12 +4,14 @@
 // https://www.apache.org/licenses/LICENSE-2.0>. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use crate::pyroscope::AgentSignal;
 use crate::utils::check_err;
 use crate::utils::get_time_range;
+use crate::PyroscopeError;
 use crate::Result;
 
 use std::sync::{
-    mpsc::{channel, Receiver, Sender},
+    mpsc::{channel, Sender},
     Arc, Mutex,
 };
 use std::{thread, thread::JoinHandle};
@@ -22,57 +24,59 @@ use std::{thread, thread::JoinHandle};
 /// The Timer thread will run continously until all Senders are dropped.
 /// The Timer thread will be joined when all Senders are dropped.
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct Timer {
     /// A vector to store listeners (mpsc::Sender)
-    txs: Arc<Mutex<Vec<Sender<u64>>>>,
+    txs: Arc<Mutex<Vec<Sender<AgentSignal>>>>,
 
     /// Thread handle
     pub handle: Option<JoinHandle<Result<()>>>,
 }
 
 impl Timer {
     /// Initialize Timer and run a thread to send events to attached listeners
-    pub fn initialize(self) -> Result<Self> {
-        let txs = Arc::clone(&self.txs);
+    pub fn initialize(cycle: std::time::Duration) -> Result<Self> {
+        let txs = Arc::new(Mutex::new(Vec::new()));
 
-        // Add Default tx
-        let (tx, _rx): (Sender<u64>, Receiver<u64>) = channel();
+        // 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()?;
+        let timer_fd = Timer::set_timerfd(cycle)?;
         let epoll_fd = Timer::create_epollfd(timer_fd)?;
 
-        let handle = Some(thread::spawn(move || {
-            loop {
-                // Exit thread if there are no listeners
-                if txs.lock()?.len() == 0 {
-                    // Close file descriptors
-                    unsafe { libc::close(timer_fd) };
-                    unsafe { libc::close(epoll_fd) };
-
-                    return Ok(());
+        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) };
+                        return Ok::<_, PyroscopeError>(());
+                    }
+
+                    // Fire @ 10th sec
+                    Timer::epoll_wait(timer_fd, epoll_fd)?;
+
+                    // Get the current time range
+                    let from = AgentSignal::NextSnapshot(get_time_range(0)?.from);
+
+                    // Iterate through Senders
+                    txs.lock()?.iter().for_each(|tx| {
+                        // Send event to attached Sender
+                        if tx.send(from).is_ok() {}
+                    });
                 }
+            })
+        });
 
-                // Fire @ 10th sec
-                Timer::epoll_wait(timer_fd, epoll_fd)?;
-
-                // Get the current time range
-                let from = get_time_range(0)?.from;
-
-                // Iterate through Senders
-                txs.lock()?.iter().for_each(|tx| {
-                    // Send event to attached Sender
-                    if tx.send(from).is_ok() {}
-                });
-            }
-        }));
-
-        Ok(Self { handle, ..self })
+        Ok(Self { handle, txs })
     }
 
     /// create and set a timer file descriptor
-    fn set_timerfd() -> Result<libc::c_int> {
+    fn set_timerfd(cycle: std::time::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
@@ -87,8 +91,8 @@ impl Timer {
         // new_value sets the Timer
         let mut new_value = libc::itimerspec {
             it_interval: libc::timespec {
-                tv_sec: 10,
-                tv_nsec: 0,
+                tv_sec: cycle.as_secs() as i64,
+                tv_nsec: cycle.subsec_nanos() as i64,
             },
             it_value: libc::timespec {
                 tv_sec: first_fire as i64,
@@ -161,7 +165,7 @@ impl 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<u64>) -> Result<()> {
+    pub fn attach_listener(&mut self, tx: Sender<AgentSignal>) -> Result<()> {
         // Push Sender to a Vector of Sender(s)
         let txs = Arc::clone(&self.txs);
         txs.lock()?.push(tx);