use brotli to compress self profile response

rust-lang · Aug 23, 2024 · 2ab84d3 · 2ab84d3
1 parent 191132b
commit 2ab84d3
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git a/site/src/request_handlers/self_profile.rs b/site/src/request_handlers/self_profile.rs
@@ -3,6 +3,7 @@ use std::io::Read;
 use std::sync::Arc;
 use std::time::Instant;
 
+use brotli::enc::BrotliEncoderParams;
 use bytes::Buf;
 use database::selector;
 use database::ArtifactId;
@@ -14,11 +15,13 @@ use crate::api::self_profile::ArtifactSizeDelta;
 use crate::api::{self_profile, self_profile_processed, self_profile_raw, ServerResult};
 use crate::load::SiteCtxt;
 use crate::self_profile::{get_or_download_self_profile, get_self_profile_raw_data};
+use crate::server::maybe_compressed_response;
 use crate::server::{Response, ResponseHeaders};
 
 pub async fn handle_self_profile_processed_download(
     body: self_profile_processed::Request,
     ctxt: &SiteCtxt,
+    allow_compression: bool,
 ) -> http::Response<hyper::Body> {
     log::info!("handle_self_profile_processed_download({:?})", body);
     let mut params = body.params.clone();
@@ -148,7 +151,18 @@ pub async fn handle_self_profile_processed_download(
         hyper::header::HeaderValue::from_static("https://profiler.firefox.com"),
     );
 
-    builder.body(hyper::Body::from(output.data)).unwrap()
+    if output.filename.ends_with("json") && allow_compression {
+        maybe_compressed_response(
+            builder,
+            output.data,
+            &Some(BrotliEncoderParams {
+                quality: 4,
+                ..Default::default()
+            }),
+        )
+    } else {
+        builder.body(hyper::Body::from(output.data)).unwrap()
+    }
 }
 
 // Add query data entries to `profile` for any queries in `base_profile` which are not present in

diff --git a/site/src/server.rs b/site/src/server.rs
@@ -342,7 +342,10 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
     let allow_compression = req
         .headers()
         .get(hyper::header::ACCEPT_ENCODING)
-        .map_or(false, |e| e.to_str().unwrap().contains("br"));
+        .and_then(|e| e.to_str().ok())
+        .map_or(false, |s| {
+            s.split(',').any(|part| part.trim().starts_with("br"))
+        });
 
     let compression = if allow_compression {
         // In tests on /perf/graphs and /perf/get, quality = 2 reduces size by 20-40% compared to 0,
@@ -441,7 +444,12 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
         "/perf/processed-self-profile" => {
             let ctxt: Arc<SiteCtxt> = server.ctxt.read().as_ref().unwrap().clone();
             let req = check!(parse_query_string(req.uri()));
-            return Ok(request_handlers::handle_self_profile_processed_download(req, &ctxt).await);
+            return Ok(request_handlers::handle_self_profile_processed_download(
+                req,
+                &ctxt,
+                allow_compression,
+            )
+            .await);
         }
         _ if req.method() == http::Method::GET => return Ok(not_found()),
         _ => {}
@@ -730,7 +738,7 @@ where
     }
 }
 
-fn maybe_compressed_response(
+pub fn maybe_compressed_response(
     response: http::response::Builder,
     body: Vec<u8>,
     compression: &Option<BrotliEncoderParams>,