feat: support zstd response encoding based on Accept-Encoding

crates/hotblocks/Cargo.toml

@@ -11,6 +11,7 @@ bytes = { workspace = true }
 chrono = { workspace = true, features = ["std"] }
 clap =  { workspace = true, features = ["derive"] }
 flate2 = { workspace = true }
+zstd = "0.13"
 futures = { workspace = true }
 ouroboros = { workspace = true }
 prometheus-client = { workspace = true }

crates/hotblocks/src/api.rs

@@ -160,9 +160,11 @@ impl IntoResponse for ResponseWithMetadata {
 async fn stream(
     Extension(app): Extension<AppRef>,
     Path(dataset_id): Path<DatasetId>,
+    headers: HeaderMap,
     Json(query): Json<Query>,
 ) -> impl IntoResponse {
-    let response = stream_internal(app, dataset_id, query, false).await;
+    let encoding = encoding_from_headers(&headers);
+    let response = stream_internal(app, dataset_id, query, false, encoding).await;
     ResponseWithMetadata::new()
         .with_dataset_id(dataset_id)
         .with_endpoint("/stream")
@@ -172,20 +174,34 @@ async fn stream(
 async fn finalized_stream(
     Extension(app): Extension<AppRef>,
     Path(dataset_id): Path<DatasetId>,
+    headers: HeaderMap,
     Json(query): Json<Query>,
 ) -> impl IntoResponse {
-    let response = stream_internal(app, dataset_id, query, true).await;
+    let encoding = encoding_from_headers(&headers);
+    let response = stream_internal(app, dataset_id, query, true, encoding).await;
     ResponseWithMetadata::new()
         .with_dataset_id(dataset_id)
         .with_endpoint("/finalized_stream")
         .with_response(|| response)
 }
 
+use axum::http::HeaderMap;
+use crate::encoding::ContentEncoding;
+
+fn encoding_from_headers(headers: &HeaderMap) -> ContentEncoding {
+    let accept = headers
+        .get("accept-encoding")
+        .and_then(|v| v.to_str().ok());
+    ContentEncoding::from_accept_encoding(accept)
+        .unwrap_or(ContentEncoding::Gzip)
+}
+
 async fn stream_internal(
     app: AppRef,
     dataset_id: DatasetId,
     query: Query,
     finalized: bool,
+    encoding: ContentEncoding,
 ) -> Response {
     let dataset = get_dataset!(app, dataset_id);
 
@@ -194,17 +210,18 @@ async fn stream_internal(
     }
 
     let query_result = if finalized {
-        app.query_service.query_finalized(&dataset, query).await
+        app.query_service.query_finalized(&dataset, query, encoding).await
     } else {
-        app.query_service.query(&dataset, query).await
+        app.query_service.query(&dataset, query, encoding).await
     };
 
     match query_result {
         Ok(stream) => {
             let mut res = Response::builder()
                 .status(200)
                 .header("content-type", "text/plain")
-                .header("content-encoding", "gzip");
+                .header("content-encoding", encoding.as_str())
+                .header("vary", "Accept-Encoding");
 
             if let Some(finalized_head) = stream.finalized_head() {
                 if finalized {

crates/hotblocks/src/encoding.rs

@@ -0,0 +1,166 @@
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ContentEncoding {
+    Gzip,
+    Zstd,
+}
+
+impl ContentEncoding {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            ContentEncoding::Gzip => "gzip",
+            ContentEncoding::Zstd => "zstd",
+        }
+    }
+
+    /// Parse Accept-Encoding header respecting quality values.
+    /// See <https://datatracker.ietf.org/doc/html/rfc7231#section-5.3.4>
+    /// At equal q, prefers zstd over gzip.
+    /// Returns None if neither gzip nor zstd is accepted.
+    pub fn from_accept_encoding(header: Option<&str>) -> Option<Self> {
+        let header = header?;
+
+        let mut best: Option<Self> = None;
+        let mut best_q: f32 = -1.0;
+
+        for part in header.split(',') {
+            let part = part.trim();
+            let (name, q) = if let Some((name, params)) = part.split_once(';') {
+                let q = params
+                    .split(';')
+                    .find_map(|p| p.trim().strip_prefix("q="))
+                    .and_then(|v| v.trim().parse::<f32>().ok())
+                    .unwrap_or(1.0);
+                (name.trim(), q)
+            } else {
+                (part, 1.0)
+            };
+
+            if q <= 0.0 {
+                continue;
+            }
+
+            let encoding = match name {
+                "zstd" => Some(ContentEncoding::Zstd),
+                "gzip" => Some(ContentEncoding::Gzip),
+                _ => None,
+            };
+
+            if let Some(enc) = encoding {
+                if q > best_q || (q == best_q && enc == ContentEncoding::Zstd) {
+                    best_q = q;
+                    best = Some(enc);
+                }
+            }
+        }
+
+        best
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn none_header() {
+        assert_eq!(ContentEncoding::from_accept_encoding(None), None);
+    }
+
+    #[test]
+    fn empty_header() {
+        assert_eq!(ContentEncoding::from_accept_encoding(Some("")), None);
+    }
+
+    #[test]
+    fn only_zstd() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("zstd")),
+            Some(ContentEncoding::Zstd)
+        );
+    }
+
+    #[test]
+    fn only_gzip() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip")),
+            Some(ContentEncoding::Gzip)
+        );
+    }
+
+    #[test]
+    fn zstd_and_gzip_equal_q_prefers_zstd() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip, zstd")),
+            Some(ContentEncoding::Zstd)
+        );
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("zstd, gzip")),
+            Some(ContentEncoding::Zstd)
+        );
+    }
+
+    #[test]
+    fn gzip_higher_q() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip;q=1.0, zstd;q=0.5")),
+            Some(ContentEncoding::Gzip)
+        );
+    }
+
+    #[test]
+    fn zstd_higher_q() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip;q=0.5, zstd;q=1.0")),
+            Some(ContentEncoding::Zstd)
+        );
+    }
+
+    #[test]
+    fn unsupported_encodings_only() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("deflate, br")),
+            None
+        );
+    }
+
+    #[test]
+    fn mixed_with_unsupported() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("br, deflate, gzip;q=0.8")),
+            Some(ContentEncoding::Gzip)
+        );
+    }
+
+    #[test]
+    fn q_zero_excluded() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip;q=0, zstd")),
+            Some(ContentEncoding::Zstd)
+        );
+    }
+
+    #[test]
+    fn q_zero_both() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("gzip;q=0, zstd;q=0")),
+            None
+        );
+    }
+
+    #[test]
+    fn whitespace_handling() {
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("  gzip ; q=0.5 , zstd ; q=0.9 ")),
+            Some(ContentEncoding::Zstd)
+        );
+    }
+
+    #[test]
+    fn wildcard_ignored() {
+        // We don't handle *, just pick from what we support
+        assert_eq!(
+            ContentEncoding::from_accept_encoding(Some("*, gzip;q=0.5")),
+            Some(ContentEncoding::Gzip)
+        );
+    }
+}

crates/hotblocks/src/main.rs

@@ -3,6 +3,7 @@ mod cli;
 mod data_service;
 mod dataset_config;
 mod dataset_controller;
+mod encoding;
 mod errors;
 mod metrics;
 mod query;

crates/hotblocks/src/query/mod.rs

@@ -6,4 +6,5 @@ mod static_snapshot;
 
 pub use executor::QueryExecutorCollector;
 pub use response::*;
+
 pub use service::*;

crates/hotblocks/src/query/response.rs

@@ -81,6 +81,7 @@ impl QueryResponse {
         query: Query,
         only_finalized: bool,
         time_limit: Option<Duration>,
+        encoding: crate::encoding::ContentEncoding,
     ) -> anyhow::Result<Self> {
         let Some(slot) = executor.get_slot() else {
             bail!(Busy)
@@ -90,7 +91,7 @@ impl QueryResponse {
         let mut runner = slot
             .run(move |slot| -> anyhow::Result<_> {
                 let mut runner =
-                    RunningQuery::new(db, dataset_id, &query, only_finalized).map(Box::new)?;
+                    RunningQuery::new(db, dataset_id, &query, only_finalized, encoding).map(Box::new)?;
                 next_run(&mut runner, slot)?;
                 Ok(runner)
             })