Merge from main

Author: omarabid

Cargo.toml

@@ -64,6 +64,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/internal/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/error.rs

@@ -1,3 +1,4 @@
+
 /// Result Alias with PyroscopeError
 pub type Result<T> = std::result::Result<T, PyroscopeError>;
 
@@ -36,7 +37,9 @@ impl PyroscopeError {
 
     /// 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 {
+    where
+        E: std::error::Error + Send + Sync + 'static,
+    {
         PyroscopeError::Compat {
             msg: msg.to_string(),
             source: Box::new(source),

src/pyroscope.rs

@@ -1,16 +1,17 @@
 use std::{
     collections::HashMap,
     sync::{
-        mpsc::{channel, Receiver, Sender},
+        mpsc::{self, Sender},
         Arc, Condvar, Mutex,
     },
     thread::JoinHandle,
+    time::Duration,
 };
 
 use crate::{
     error::Result,
     session::{Session, SessionManager, SessionSignal},
-    timer::Timer,
+    timer::{Timer, TimerSignal},
 };
 
 use pyroscope_backends::pprof::Pprof;
@@ -147,6 +148,8 @@ impl PyroscopeAgentBuilder {
         }
     }
 
+    /// # use pyroscope::PyroscopeError;
+    /// # fn main() -> Result<(), PyroscopeError> {
     /// Set tags. Default is empty.
     /// # Example
     /// ```ignore
@@ -176,7 +179,7 @@ impl PyroscopeAgentBuilder {
         log::trace!(target: LOG_TAG, "Backend initialized");
 
         // Start Timer
-        let timer = Timer::default().initialize()?;
+        let timer = Timer::initialize(std::time::Duration::from_secs(10))?;
         log::trace!(target: LOG_TAG, "Timer initialized");
 
         // Start the SessionManager
@@ -201,7 +204,7 @@ impl PyroscopeAgentBuilder {
 pub struct PyroscopeAgent {
     timer: Timer,
     session_manager: SessionManager,
-    tx: Option<Sender<u64>>,
+    tx: Option<Sender<TimerSignal>>,
     handle: Option<JoinHandle<Result<()>>>,
     running: Arc<(Mutex<bool>, Condvar)>,
 
@@ -293,7 +296,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);
 
@@ -304,28 +307,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 {
+                    TimerSignal::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,
+                        )?))?
+                    }
+                    TimerSignal::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(())
@@ -346,7 +352,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(TimerSignal::NextSnapshot(0));
+            sender.send(TimerSignal::Terminate)?;
         } else {
             log::error!("PyroscopeAgent - Missing sender")
         }

src/session.rs

@@ -1,7 +1,7 @@
 use std::{
     sync::mpsc::{sync_channel, Receiver, SyncSender},
-    thread,
-    thread::JoinHandle,
+    thread::{self, JoinHandle},
+    time::Duration,
 };
 
 use crate::pyroscope::PyroscopeConfig;
@@ -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,
 }
 
@@ -159,7 +161,7 @@ impl Session {
                 ("spyName", self.config.spy_name.as_str()),
             ])
             .body(self.report)
-            .timeout(std::time::Duration::from_secs(10))
+            .timeout(Duration::from_secs(10))
             .send()?;
 
         Ok(())

src/timer/epoll.rs

@@ -1,18 +1,17 @@
-// Copyright 2021 Developers of Pyroscope.
-
-// Licensed under the Apache License, Version 2.0 <LICENSE or
-// https://www.apache.org/licenses/LICENSE-2.0>. This file may not be copied, modified, or distributed
-// except according to those terms.
-
+use super::TimerSignal;
 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};
+use std::{
+    thread::{self, JoinHandle},
+    time::Duration,
+};
 
 /// A thread that sends a notification every 10th second
 ///
@@ -22,57 +21,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<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(self) -> Result<Self> {
-        let txs = Arc::clone(&self.txs);
+    pub fn initialize(cycle: 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 = TimerSignal::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: 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 +88,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 +162,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<TimerSignal>) -> Result<()> {
         // Push Sender to a Vector of Sender(s)
         let txs = Arc::clone(&self.txs);
         txs.lock()?.push(tx);