use std::time::Duration;
use crate::core::{CoreError, CoreResult};
pub struct ProfilingHandle {
inner: ProfilingHandleInner,
shutdown_timeout: Duration,
}
enum ProfilingHandleInner {
Disabled,
Pyroscope(Option<pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>>),
}
impl ProfilingHandle {
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,
}
}
pub fn is_running(&self) -> bool {
matches!(self.inner, ProfilingHandleInner::Pyroscope(Some(_)))
}
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()
}
}