Enable compression for gRPC and HTTP endpoint (#671)

Author: mwylde

Cargo.lock

@@ -31,7 +31,7 @@ members = [
 resolver = "2"
 
 [workspace.dependencies]
-tonic = { version = "0.11" }
+tonic = { version = "0.11", features = ["zstd"] }
 tonic-build = { version = "0.11" }
 tonic-web = { version = "0.11" }
 tonic-reflection = { version = "0.11" }

Cargo.toml

@@ -40,7 +40,7 @@ bincode = { version = "2.0.0-rc.3", features = ["serde"]}
 petgraph = {version = "0.6", features = ["serde-1"]}
 
 http = "0.2"
-tower-http = {version = "0.4", features = ["trace", "fs", "cors", "validate-request", "auth"]}
+tower-http = {version = "0.4", features = ["trace", "fs", "cors", "validate-request", "auth", "compression-zstd"]}
 axum = {version = "0.6.20", features = ["headers", "tokio", "macros"]}
 axum-extra = "0.7.4"
 h2 = "0.3.26"

crates/arroyo-api/Cargo.toml

@@ -482,6 +482,7 @@ pub async fn get_job_output(
             };
 
             let output_data: OutputData = v.into();
+
             let e = Ok(Event::default()
                 .json_data(output_data)
                 .unwrap()

crates/arroyo-api/src/jobs.rs

@@ -6,6 +6,8 @@ use serde::{Deserialize, Serialize};
 use std::net::{SocketAddr, TcpListener};
 use time::OffsetDateTime;
 use tonic::transport::Channel;
+use tower_http::compression::predicate::NotForContentType;
+use tower_http::compression::{CompressionLayer, DefaultPredicate, Predicate};
 use tracing::{error, info};
 use utoipa::OpenApi;
 
@@ -125,7 +127,14 @@ pub fn start_server(database: DatabaseSource, guard: ShutdownGuard) -> anyhow::R
     let listener = TcpListener::bind(addr)?;
     let local_addr = listener.local_addr()?;
 
-    let app = rest::create_rest_app(database, &config.controller_endpoint());
+    let app = rest::create_rest_app(database, &config.controller_endpoint()).layer(
+        CompressionLayer::new().zstd(true).compress_when(
+            DefaultPredicate::new()
+                // compression doesn't work for server-sent events
+                // (https://github.com/tokio-rs/axum/discussions/2034)
+                .and(NotForContentType::new("text/event-stream")),
+        ),
+    );
 
     info!("Starting API server on {:?}", local_addr);
     guard.into_spawn_task(wrap_start(

crates/arroyo-api/src/lib.rs

@@ -8,6 +8,8 @@ use arroyo_rpc::api_types::metrics::OperatorMetricGroup;
 use arroyo_rpc::api_types::OperatorMetricGroupCollection;
 use arroyo_rpc::grpc::rpc::controller_grpc_client::ControllerGrpcClient;
 use arroyo_rpc::grpc::rpc::JobMetricsReq;
+use tonic::codec::CompressionEncoding;
+use tonic::transport::Channel;
 use tonic::Code;
 
 /// Get a job's metrics
@@ -38,10 +40,15 @@ pub async fn get_operator_metric_groups(
     )
     .await?;
 
-    let mut controller = ControllerGrpcClient::connect(state.controller_addr)
+    let channel = Channel::builder(state.controller_addr.parse().unwrap())
+        .connect()
         .await
         .map_err(log_and_map)?;
 
+    let mut controller = ControllerGrpcClient::new(channel)
+        .accept_compressed(CompressionEncoding::Zstd)
+        .send_compressed(CompressionEncoding::Zstd);
+
     let data = match controller
         .job_metrics(JobMetricsReq { job_id: job.id })
         .await

crates/arroyo-api/src/metrics.rs

@@ -33,6 +33,7 @@ prometheus = "0.13"
 async-trait = "0.1"
 lazy_static = "1.4.0"
 chrono = "0.4"
+zstd = "0.13"
 
 arrow-schema = {workspace = true}
 

crates/arroyo-controller/Cargo.toml

@@ -36,6 +36,7 @@ use tokio::sync::mpsc::error::TrySendError;
 use tokio::sync::mpsc::Sender;
 use tokio::sync::RwLock;
 use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
+use tonic::codec::CompressionEncoding;
 use tonic::{Request, Response, Status};
 use tracing::{debug, info, warn};
 
@@ -387,7 +388,7 @@ impl ControllerGrpc for ControllerServer {
                         remove.insert(i);
                     }
                     Err(TrySendError::Full(_)) => {
-                        warn!("queue full");
+                        debug!("queue full");
                     }
                 }
             }
@@ -455,6 +456,7 @@ impl ControllerGrpc for ControllerServer {
             .ok_or_else(|| Status::not_found("No metrics for job"))?
             .clone();
 
+        // TODO: send this over in a more efficient format like protobuf
         Ok(Response::new(JobMetricsResp {
             metrics: serde_json::to_string(&metrics.get_groups().await).unwrap(),
         }))
@@ -625,7 +627,11 @@ impl ControllerServer {
             local_addr,
             arroyo_server_common::grpc_server()
                 .accept_http1(true)
-                .add_service(ControllerGrpcServer::new(self.clone()))
+                .add_service(
+                    ControllerGrpcServer::new(self.clone())
+                        .send_compressed(CompressionEncoding::Zstd)
+                        .accept_compressed(CompressionEncoding::Zstd),
+                )
                 .add_service(reflection)
                 .serve_with_incoming(TcpListenerStream::new(listener)),
         ));

crates/arroyo-controller/src/lib.rs

@@ -103,6 +103,7 @@ impl NodeServer {
             .env(JOB_ID_ENV, req.job_id.clone())
             .env("ARROYO__WORKER__TASK_SLOTS", format!("{}", slots))
             .env(RUN_ID_ENV, format!("{}", req.run_id))
+            .env("ARROYO__ADMIN__HTTP_PORT", "0")
             .kill_on_drop(true)
             .spawn()
             .map_err(|e| Status::internal(format!("Failed to start worker: {:?}", e)))?;

crates/arroyo-node/src/lib.rs

@@ -30,7 +30,7 @@ const OperatorDetail: React.FC<OperatorDetailProps> = ({ pipelineId, jobId, oper
     return (
       <Alert status="warning">
         <AlertIcon />
-        <AlertDescription>Failed to get job's metrics.</AlertDescription>
+        <AlertDescription>Failed to get job metrics.</AlertDescription>
       </Alert>
     );
   }