rs-zero 0.2.7

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use std::time::Duration;

use crate::core::{CoreError, CoreResult};

/// Running profiling handle.
pub struct ProfilingHandle {
    inner: ProfilingHandleInner,
    shutdown_timeout: Duration,
}

enum ProfilingHandleInner {
    Disabled,
    Pyroscope(Option<pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>>),
}

impl ProfilingHandle {
    /// Creates a disabled handle that can be shut down safely.
    pub fn disabled() -> Self {
        Self {
            inner: ProfilingHandleInner::Disabled,
            shutdown_timeout: Duration::ZERO,
        }
    }

    pub(crate) fn pyroscope(
        agent: pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>,
        shutdown_timeout: Duration,
    ) -> Self {
        Self {
            inner: ProfilingHandleInner::Pyroscope(Some(agent)),
            shutdown_timeout,
        }
    }

    /// Returns whether profiling is currently running.
    pub fn is_running(&self) -> bool {
        matches!(self.inner, ProfilingHandleInner::Pyroscope(Some(_)))
    }

    /// Stops the profiler and releases agent resources.
    pub fn shutdown(mut self) -> CoreResult<()> {
        let ProfilingHandleInner::Pyroscope(agent) = &mut self.inner else {
            return Ok(());
        };
        let Some(agent) = agent.take() else {
            return Ok(());
        };
        let started_at = std::time::Instant::now();
        let ready = agent
            .stop()
            .map_err(|error| CoreError::Profiling(error.to_string()))?;
        ready.shutdown();
        if !self.shutdown_timeout.is_zero() && started_at.elapsed() > self.shutdown_timeout {
            return Err(CoreError::Profiling(format!(
                "pyroscope shutdown exceeded {:?}",
                self.shutdown_timeout
            )));
        }
        Ok(())
    }
}

impl std::fmt::Debug for ProfilingHandle {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ProfilingHandle")
            .field("running", &self.is_running())
            .field("shutdown_timeout", &self.shutdown_timeout)
            .finish()
    }
}