Merge pull request #978 from tgnottingham/update-flamegraph-crox-support

Update flamegraph and crox support

Author: Mark-Simulacrum

site/src/request_handlers/self_profile.rs

@@ -25,14 +25,33 @@ pub async fn handle_self_profile_processed_download(
         body.benchmark,
         body.run_name
     );
+
     let start = Instant::now();
-    let pieces = match get_pieces(body, ctxt).await {
+
+    let (url, is_tarball) = match handle_self_profile_raw(body, ctxt).await {
+        Ok(v) => (v.url, v.is_tarball),
+        Err(e) => {
+            let mut resp = Response::new(e.into());
+            *resp.status_mut() = StatusCode::BAD_REQUEST;
+            return resp;
+        }
+    };
+
+    if is_tarball {
+        let mut resp =
+            Response::new("Processing legacy format self-profile data is not supported".into());
+        *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+        return resp;
+    }
+
+    let data = match get_self_profile_raw_data(&url).await {
         Ok(v) => v,
         Err(e) => return e,
     };
-    log::trace!("got pieces {:?} in {:?}", pieces, start.elapsed());
 
-    let output = match crate::self_profile::generate(&title, pieces, params) {
+    log::trace!("got data in {:?}", start.elapsed());
+
+    let output = match crate::self_profile::generate(&title, data, params) {
         Ok(c) => c,
         Err(e) => {
             log::error!("Failed to generate json {:?}", e);
@@ -166,22 +185,10 @@ fn get_self_profile_data(
     Ok(profile)
 }
 
-async fn get_pieces(
-    body: crate::api::self_profile_raw::Request,
-    ctxt: &SiteCtxt,
-) -> Result<crate::self_profile::Pieces, Response> {
-    let res = handle_self_profile_raw(body, ctxt).await;
-    let url = match res {
-        Ok(v) => v.url,
-        Err(e) => {
-            let mut resp = Response::new(e.into());
-            *resp.status_mut() = StatusCode::BAD_REQUEST;
-            return Err(resp);
-        }
-    };
+async fn get_self_profile_raw_data(url: &str) -> Result<Vec<u8>, Response> {
     log::trace!("downloading {}", url);
 
-    let resp = match reqwest::get(&url).await {
+    let resp = match reqwest::get(url).await {
         Ok(r) => r,
         Err(e) => {
             let mut resp = Response::new(format!("{:?}", e).into());
@@ -197,7 +204,7 @@ async fn get_pieces(
         return Err(resp);
     }
 
-    let tarball = match resp.bytes().await {
+    let compressed = match resp.bytes().await {
         Ok(b) => b,
         Err(e) => {
             let mut resp =
@@ -206,18 +213,19 @@ async fn get_pieces(
             return Err(resp);
         }
     };
-    let tarball = tar::Archive::new(std::io::BufReader::new(snap::read::FrameDecoder::new(
-        tarball.reader(),
-    )));
-    let pieces = match crate::self_profile::Pieces::from_tarball(tarball) {
+
+    let mut data = Vec::new();
+
+    match snap::read::FrameDecoder::new(compressed.reader()).read_to_end(&mut data) {
         Ok(v) => v,
         Err(e) => {
-            let mut resp = Response::new(format!("could not extract from tarball: {:?}", e).into());
+            let mut resp = Response::new(format!("could not decode: {:?}", e).into());
             *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
             return Err(resp);
         }
     };
-    Ok(pieces)
+
+    Ok(data)
 }
 
 pub async fn handle_self_profile_raw_download(

site/src/self_profile.rs

@@ -3,8 +3,6 @@
 
 use anyhow::Context;
 use std::collections::HashMap;
-use std::fmt;
-use std::io::Read;
 
 pub mod crox;
 pub mod flamegraph;
@@ -17,7 +15,7 @@ pub struct Output {
 
 pub fn generate(
     title: &str,
-    pieces: Pieces,
+    self_profile_data: Vec<u8>,
     mut params: HashMap<String, String>,
 ) -> anyhow::Result<Output> {
     let removed = params.remove("type");
@@ -27,7 +25,7 @@ pub fn generate(
                 .context("crox opts")?;
             Ok(Output {
                 filename: "chrome_profiler.json",
-                data: crox::generate(pieces, opt).context("crox")?,
+                data: crox::generate(self_profile_data, opt).context("crox")?,
                 is_download: true,
             })
         }
@@ -36,56 +34,10 @@ pub fn generate(
                 .context("flame opts")?;
             Ok(Output {
                 filename: "flamegraph.svg",
-                data: flamegraph::generate(title, pieces, opt).context("flame")?,
+                data: flamegraph::generate(title, self_profile_data, opt).context("flame")?,
                 is_download: false,
             })
         }
         _ => anyhow::bail!("Unknown type, specify type={crox,flamegraph}"),
     }
 }
-
-pub struct Pieces {
-    pub string_data: Vec<u8>,
-    pub string_index: Vec<u8>,
-    pub events: Vec<u8>,
-}
-
-impl fmt::Debug for Pieces {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("Pieces")
-            .field("string_data", &self.string_data.len())
-            .field("string_index", &self.string_index.len())
-            .field("events", &self.events.len())
-            .finish()
-    }
-}
-
-impl Pieces {
-    pub fn from_tarball<R: std::io::Read>(mut tarball: tar::Archive<R>) -> anyhow::Result<Pieces> {
-        let mut pieces = Pieces {
-            string_data: Vec::new(),
-            string_index: Vec::new(),
-            events: Vec::new(),
-        };
-
-        for entry in tarball.entries().context("entries")? {
-            let mut entry = entry.context("tarball entry")?;
-            let path = entry.path_bytes();
-            if *path == *b"self-profile.string_index" {
-                entry
-                    .read_to_end(&mut pieces.string_index)
-                    .context("reading string index")?;
-            } else if *path == *b"self-profile.string_data" {
-                entry
-                    .read_to_end(&mut pieces.string_data)
-                    .context("reading string data")?;
-            } else if *path == *b"self-profile.events" {
-                entry
-                    .read_to_end(&mut pieces.events)
-                    .context("reading events")?;
-            }
-        }
-
-        Ok(pieces)
-    }
-}

site/src/self_profile/crox.rs

@@ -125,14 +125,13 @@ fn get_args(full_event: &analyzeme::Event) -> Option<HashMap<String, String>> {
 }
 
 /// Returns JSON blob fit, `chrome_profiler.json`.
-pub fn generate(pieces: super::Pieces, opt: Opt) -> anyhow::Result<Vec<u8>> {
+pub fn generate(self_profile_data: Vec<u8>, opt: Opt) -> anyhow::Result<Vec<u8>> {
     let mut serializer = serde_json::Serializer::new(Vec::new());
 
     let mut seq = serializer.serialize_seq(None)?;
 
-    let data =
-        ProfilingData::from_buffers(pieces.string_data, pieces.string_index, pieces.events, None)
-            .map_err(|e| anyhow::format_err!("{:?}", e))?;
+    let data = ProfilingData::from_paged_buffer(self_profile_data)
+        .map_err(|e| anyhow::format_err!("{:?}", e))?;
 
     let thread_to_collapsed_thread =
         generate_thread_to_collapsed_thread_mapping(opt.collapse_threads, &data);

site/src/self_profile/flamegraph.rs

@@ -5,10 +5,9 @@ use inferno::flamegraph::{from_lines, Options as FlamegraphOptions};
 #[derive(serde::Deserialize, Debug)]
 pub struct Opt {}
 
-pub fn generate(title: &str, pieces: super::Pieces, _: Opt) -> anyhow::Result<Vec<u8>> {
-    let profiling_data =
-        ProfilingData::from_buffers(pieces.string_data, pieces.string_index, pieces.events, None)
-            .map_err(|e| anyhow::format_err!("{:?}", e))?;
+pub fn generate(title: &str, self_profile_data: Vec<u8>, _: Opt) -> anyhow::Result<Vec<u8>> {
+    let profiling_data = ProfilingData::from_paged_buffer(self_profile_data)
+        .map_err(|e| anyhow::format_err!("{:?}", e))?;
 
     let recorded_stacks = collapse_stacks(&profiling_data)
         .iter()