From b2c121d911fe424a43529b0907fff3bb764e356e Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Sat, 31 Aug 2024 14:01:55 -0700
Subject: [PATCH] add support for compressed http output

The client needs to request this via Accept-Encoding.
We support either gzip or deflate; we want this to be
as cheap as possible to compress.

Note that this doesn't enable requesting compressed output in
kcli, because compression increases the latency around collecting
metrics and we want that (typically locally-run) command to have
lower latency for queue-summary and top.
---
 Cargo.lock                                    | 19 +++++++++++++++++++
 crates/kcli/Cargo.toml                        |  2 +-
 crates/kumo-server-common/Cargo.toml          |  2 +-
 .../kumo-server-common/src/http_server/mod.rs |  7 +++++++
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index dcc1c0cd..3e3fee76 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -239,6 +239,19 @@ dependencies = [
  "pin-project-lite",
 ]
 
+[[package]]
+name = "async-compression"
+version = "0.4.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c90a406b4495d129f00461241616194cb8a032c8d1c53c657f0961d5f8e0498"
+dependencies = [
+ "flate2",
+ "futures-core",
+ "memchr",
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "async-executor"
 version = "1.13.0"
@@ -5043,6 +5056,7 @@ version = "0.11.27"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dd67538700a17451e7cba03ac727fb961abb7607553461627b97de0b89cf4a62"
 dependencies = [
+ "async-compression",
  "base64 0.21.7",
  "bytes",
  "encoding_rs",
@@ -5069,6 +5083,7 @@ dependencies = [
  "system-configuration",
  "tokio",
  "tokio-rustls 0.24.1",
+ "tokio-util",
  "tower-service",
  "url",
  "wasm-bindgen",
@@ -6621,12 +6636,16 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5"
 dependencies = [
+ "async-compression",
  "bitflags 2.6.0",
  "bytes",
+ "futures-core",
  "http 1.1.0",
  "http-body 1.0.1",
  "http-body-util",
  "pin-project-lite",
+ "tokio",
+ "tokio-util",
  "tower-layer",
  "tower-service",
  "tracing",
diff --git a/crates/kcli/Cargo.toml b/crates/kcli/Cargo.toml
index 06afd41c..05f34660 100644
--- a/crates/kcli/Cargo.toml
+++ b/crates/kcli/Cargo.toml
@@ -18,7 +18,7 @@ lexicmp = "0.1"
 message = {path="../message", default-features=false}
 ordermap = {version="0.5", features=["serde"]}
 kumo-api-types = {path="../kumo-api-types", default-features=false}
-reqwest = {workspace=true, default-features=false, features=["json", "rustls-tls"]}
+reqwest = {workspace=true, default-features=false, features=["json", "rustls-tls", "deflate", "gzip"]}
 serde = {version="1.0", features=["derive"]}
 serde_json = "1.0"
 tabout = "0.3"
diff --git a/crates/kumo-server-common/Cargo.toml b/crates/kumo-server-common/Cargo.toml
index 79652c07..096ead34 100644
--- a/crates/kumo-server-common/Cargo.toml
+++ b/crates/kumo-server-common/Cargo.toml
@@ -58,7 +58,7 @@ throttle = {path="../throttle"}
 tokio = {workspace=true, features=["full", "tracing"]}
 tokio-metrics = "0.3.1"
 tokio-metrics-collector = "0.2.1"
-tower-http = {version="0.5", features=["trace"]}
+tower-http = {version="0.5", features=["trace", "compression-deflate", "compression-gzip"]}
 tracing = "0.1"
 tracing-appender = "0.2"
 tracing-subscriber = {version="0.3", features=["env-filter", "std", "fmt", "json"]}
diff --git a/crates/kumo-server-common/src/http_server/mod.rs b/crates/kumo-server-common/src/http_server/mod.rs
index aac9e8da..bbd0ac77 100644
--- a/crates/kumo-server-common/src/http_server/mod.rs
+++ b/crates/kumo-server-common/src/http_server/mod.rs
@@ -15,6 +15,7 @@ use prometheus::Histogram;
 use serde::Deserialize;
 use std::net::{IpAddr, SocketAddr, TcpListener};
 use std::sync::Arc;
+use tower_http::compression::CompressionLayer;
 use tower_http::trace::TraceLayer;
 use utoipa::openapi::security::{Http, HttpAuthScheme, SecurityScheme};
 use utoipa::OpenApi;
@@ -149,6 +150,11 @@ impl HttpListenerParams {
     pub async fn start(self, router_and_docs: RouterAndDocs) -> anyhow::Result<()> {
         let api_docs = router_and_docs.make_docs();
 
+        let compression_layer: CompressionLayer = CompressionLayer::new()
+            .deflate(true)
+            .gzip(true)
+            .quality(tower_http::CompressionLevel::Fastest);
+
         let app = router_and_docs
             .router
             .layer(DefaultBodyLimit::max(
@@ -169,6 +175,7 @@ impl HttpListenerParams {
                 },
                 auth_middleware,
             ))
+            .layer(compression_layer)
             .layer(TraceLayer::new_for_http());
         let socket = TcpListener::bind(&self.listen)
             .with_context(|| format!("listen on {}", self.listen))?;