async_profiler_agent/reporter/
local.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! A reporter that reports into a directory.
5
6use async_trait::async_trait;
7use chrono::SecondsFormat;
8use std::path::{Path, PathBuf};
9use std::time::SystemTime;
10use thiserror::Error;
11
12use crate::metadata::ReportMetadata;
13
14use super::Reporter;
15
16#[derive(Error, Debug)]
17enum LocalReporterError {
18    #[error("{0}")]
19    IoError(#[from] std::io::Error),
20}
21
22/// A reporter that reports into a directory.
23///
24/// The files are reported with the filename `yyyy-mm-ddTHH-MM-SSZ.jfr`
25///
26/// Unlike other reporters (e.g. S3), the local reporter will
27/// flush any pending JFR data on drop when the profiler task is cancelled
28/// (for example, during Tokio runtime shutdown). Other reporters require
29/// an explicit call to [`RunningProfiler::stop`] to ensure the last sample
30/// is uploaded.
31///
32/// It does not currently use the metadata, so if you are using
33/// [LocalReporter] alone, rather than inside a [MultiReporter], you
34/// can just use [AgentMetadata::NoMetadata] as metadata.
35///
36/// [AgentMetadata::NoMetadata]: crate::metadata::AgentMetadata::NoMetadata
37/// [MultiReporter]: crate::reporter::multi::MultiReporter
38/// [`RunningProfiler::stop`]: crate::profiler::RunningProfiler::stop
39///
40/// ### Example
41///
42/// ```
43/// # use async_profiler_agent::metadata::AgentMetadata;
44/// # use async_profiler_agent::profiler::{ProfilerBuilder, SpawnError};
45/// # #[tokio::main]
46/// # async fn main() -> Result<(), SpawnError> {
47/// let profiler = ProfilerBuilder::default()
48///    .with_local_reporter("/tmp/profiles")
49///    .build();
50/// # if false { // don't spawn the profiler in doctests
51/// let profiler = profiler.spawn_controllable()?;
52/// // ... your program goes here
53/// profiler.stop().await; // make sure the last profile is flushed
54/// # }
55/// # Ok(())
56/// # }
57/// ```
58#[derive(Debug)]
59pub struct LocalReporter {
60    directory: PathBuf,
61}
62
63impl LocalReporter {
64    /// Instantiate a new LocalReporter writing into the provided directory.
65    pub fn new(directory: impl Into<PathBuf>) -> Self {
66        LocalReporter {
67            directory: directory.into(),
68        }
69    }
70
71    fn jfr_file_name() -> String {
72        let time: chrono::DateTime<chrono::Utc> = SystemTime::now().into();
73        let time = time
74            .to_rfc3339_opts(SecondsFormat::Secs, true)
75            .replace(":", "-");
76        format!("{time}.jfr")
77    }
78
79    /// Writes the jfr file to disk.
80    async fn report_profiling_data(
81        &self,
82        jfr: Vec<u8>,
83        _metadata_obj: &ReportMetadata<'_>,
84    ) -> Result<(), std::io::Error> {
85        let file_name = Self::jfr_file_name();
86        tracing::debug!("reporting {file_name}");
87        tokio::fs::write(self.directory.join(file_name), jfr).await?;
88        Ok(())
89    }
90}
91
92#[async_trait]
93impl Reporter for LocalReporter {
94    async fn report(
95        &self,
96        jfr: Vec<u8>,
97        metadata: &ReportMetadata,
98    ) -> Result<(), Box<dyn std::error::Error + Send>> {
99        self.report_profiling_data(jfr, metadata)
100            .await
101            .map_err(|e| Box::new(LocalReporterError::IoError(e)) as _)
102    }
103
104    fn report_blocking(
105        &self,
106        jfr_path: &Path,
107        _metadata: &ReportMetadata,
108    ) -> Result<(), Box<dyn std::error::Error + Send>> {
109        let file_name = Self::jfr_file_name();
110        tracing::debug!("reporting {file_name} (blocking)");
111        std::fs::copy(jfr_path, self.directory.join(file_name))
112            .map(|_| ())
113            .map_err(|e| Box::new(LocalReporterError::IoError(e)) as _)
114    }
115}
116
117#[cfg(test)]
118mod test {
119    use std::path::Path;
120
121    use crate::{
122        metadata::DUMMY_METADATA,
123        reporter::{Reporter, local::LocalReporter},
124    };
125
126    #[tokio::test]
127    async fn test_local_reporter() {
128        let dir = tempfile::tempdir().unwrap();
129        let reporter = LocalReporter::new(dir.path());
130        reporter
131            .report(b"JFR".into(), &DUMMY_METADATA)
132            .await
133            .unwrap();
134        let jfr_file = std::fs::read_dir(dir.path())
135            .unwrap()
136            .flat_map(|f| f.ok())
137            .filter(|f| {
138                Path::new(&f.file_name())
139                    .extension()
140                    .is_some_and(|e| e == "jfr")
141            })
142            .next()
143            .unwrap();
144        assert_eq!(tokio::fs::read(jfr_file.path()).await.unwrap(), b"JFR");
145    }
146
147    #[test]
148    fn test_local_reporter_report_blocking() {
149        let dir = tempfile::tempdir().unwrap();
150        let src = dir.path().join("input.jfr");
151        std::fs::write(&src, b"JFR-DROP").unwrap();
152        let out_dir = tempfile::tempdir().unwrap();
153        let reporter = LocalReporter::new(out_dir.path());
154        reporter.report_blocking(&src, &DUMMY_METADATA).unwrap();
155        let jfr_file = std::fs::read_dir(out_dir.path())
156            .unwrap()
157            .flat_map(|f| f.ok())
158            .filter(|f| {
159                Path::new(&f.file_name())
160                    .extension()
161                    .is_some_and(|e| e == "jfr")
162            })
163            .next()
164            .unwrap();
165        assert_eq!(std::fs::read(jfr_file.path()).unwrap(), b"JFR-DROP");
166    }
167}
async_profiler_agent/reporter/local.rs

async_profiler_agent/reporter/
local.rs
