Fix compression for deferred responses (#2986)

We replace tower-http's `CompressionLayer` with a custom stream transformation. This is necessary because tower-http uses async-compression, which buffers data until the end of the stream to then write it, ensuring a better compression. This is incompatible with the multipart protocol for `@defer`, which requires chunks to be sent as soon as possible. So we need to compress them independently.

This extracts parts of the codec module of async-compression, which so far is not public, and makes a streaming wrapper above it that flushes the compressed data on every response in the stream.

This is expected to be temporary, as we have in flight PRs for async-compression:
- Nullus157/async-compression#155
- Nullus157/async-compression#178

With Nullus157/async-compression#150 we might be able to at least remove the vendored code

Author: Geal

.changesets/fix_geal_fix_defer_compression.md

@@ -0,0 +1,7 @@
+### Fix compression for deferred responses ([Issue #1572](https://github.com/apollographql/router/issues/1572))
+
+We replace tower-http's `CompressionLayer` with a custom stream transformation. This is necessary because tower-http uses async-compression, which buffers data until the end of the stream to then write it, ensuring a better compression. This is incompatible with the multipart protocol for `@defer`, which requires chunks to be sent as soon as possible. So we need to compress them independently.
+
+This extracts parts of the codec module of async-compression, which so far is not public, and makes a streaming wrapper above it that flushes the compressed data on every response in the stream.
+
+By [@Geal](https://github.com/Geal) in https://github.com/apollographql/router/pull/2986

Cargo.lock

@@ -286,6 +286,7 @@ dependencies = [
  "axum",
  "backtrace",
  "base64 0.20.0",
+ "brotli",
  "buildstructor 0.5.2",
  "bytes",
  "ci_info",
@@ -401,6 +402,8 @@ dependencies = [
  "wiremock",
  "wsl",
  "yaml-rust",
+ "zstd",
+ "zstd-safe",
 ]
 
 [[package]]
@@ -7073,3 +7076,33 @@ dependencies = [
  "quote",
  "syn 2.0.13",
 ]
+
+[[package]]
+name = "zstd"
+version = "0.12.3+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76eea132fb024e0e13fd9c2f5d5d595d8a967aa72382ac2f9d39fcc95afd0806"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "6.0.5+zstd.1.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d56d9e60b4b1758206c238a10165fbcae3ca37b01744e394c463463f6529d23b"
+dependencies = [
+ "libc",
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.8+zstd.1.5.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5556e6ee25d32df2586c098bbfa278803692a20d0ab9565e049480d52707ec8c"
+dependencies = [
+ "cc",
+ "libc",
+ "pkg-config",
+]

apollo-router/Cargo.toml

@@ -203,6 +203,10 @@ yaml-rust = "0.4.5"
 wsl = "0.1.0"
 tokio-rustls = "0.23.4"
 http-serde = "1.1.2"
+memchr = "2.5.0"
+brotli = "3.3.4"
+zstd = "0.12.3"
+zstd-safe = "6.0.5"
 
 [target.'cfg(macos)'.dependencies]
 uname = "0.1.1"

apollo-router/src/axum_factory/axum_http_server_factory.rs

@@ -19,6 +19,9 @@ use futures::channel::oneshot;
 use futures::future::join;
 use futures::future::join_all;
 use futures::prelude::*;
+use http::header::ACCEPT_ENCODING;
+use http::header::CONTENT_ENCODING;
+use http::HeaderValue;
 use http::Request;
 use http_body::combinators::UnsyncBoxBody;
 use hyper::Body;
@@ -32,10 +35,6 @@ use tokio_rustls::TlsAcceptor;
 use tower::service_fn;
 use tower::BoxError;
 use tower::ServiceExt;
-use tower_http::compression::predicate::NotForContentType;
-use tower_http::compression::CompressionLayer;
-use tower_http::compression::DefaultPredicate;
-use tower_http::compression::Predicate;
 use tower_http::trace::TraceLayer;
 
 use super::listeners::ensure_endpoints_consistency;
@@ -45,6 +44,7 @@ use super::listeners::ListenersAndRouters;
 use super::utils::decompress_request_body;
 use super::utils::PropagatingMakeSpan;
 use super::ListenAddrAndRouter;
+use crate::axum_factory::compression::Compressor;
 use crate::axum_factory::listeners::get_extra_listeners;
 use crate::axum_factory::listeners::serve_router_on_listen_addr;
 use crate::configuration::Configuration;
@@ -329,12 +329,7 @@ where
         ))
         .layer(TraceLayer::new_for_http().make_span_with(PropagatingMakeSpan { entitlement }))
         .layer(Extension(service_factory))
-        .layer(cors)
-        // Compress the response body, except for multipart responses such as with `@defer`.
-        // This is a work-around for https://github.com/apollographql/router/issues/1572
-        .layer(CompressionLayer::new().compress_when(
-            DefaultPredicate::new().and(NotForContentType::const_new("multipart/")),
-        ));
+        .layer(cors);
 
     let route = endpoints_on_main_listener
         .into_iter()
@@ -434,6 +429,11 @@ async fn handle_graphql(
 
     let request: router::Request = http_request.into();
     let context = request.context.clone();
+    let accept_encoding = request
+        .router_request
+        .headers()
+        .get(ACCEPT_ENCODING)
+        .cloned();
 
     let res = service.oneshot(request).await;
     let dur = context.busy_time().await;
@@ -467,7 +467,24 @@ async fn handle_graphql(
         }
         Ok(response) => {
             tracing::info!(counter.apollo_router_session_count_active = -1,);
-            response.response.into_response()
+            let (mut parts, body) = response.response.into_parts();
+
+            let opt_compressor = accept_encoding
+                .as_ref()
+                .and_then(|value| value.to_str().ok())
+                .and_then(|v| Compressor::new(v.split(',').map(|s| s.trim())));
+            let body = match opt_compressor {
+                None => body,
+                Some(compressor) => {
+                    parts.headers.insert(
+                        CONTENT_ENCODING,
+                        HeaderValue::from_static(compressor.content_encoding()),
+                    );
+                    Body::wrap_stream(compressor.process(body))
+                }
+            };
+
+            http::Response::from_parts(parts, body).into_response()
         }
     }
 }

apollo-router/src/axum_factory/compression/codec/brotli/encoder.rs

@@ -0,0 +1,112 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+use std::fmt;
+use std::io::Error;
+use std::io::ErrorKind;
+use std::io::Result;
+
+use brotli::enc::backward_references::BrotliEncoderParams;
+use brotli::enc::encode::BrotliEncoderCompressStream;
+use brotli::enc::encode::BrotliEncoderCreateInstance;
+use brotli::enc::encode::BrotliEncoderHasMoreOutput;
+use brotli::enc::encode::BrotliEncoderIsFinished;
+use brotli::enc::encode::BrotliEncoderOperation;
+use brotli::enc::encode::BrotliEncoderStateStruct;
+use brotli::enc::StandardAlloc;
+
+use crate::axum_factory::compression::codec::Encode;
+use crate::axum_factory::compression::util::PartialBuffer;
+
+pub(crate) struct BrotliEncoder {
+    state: BrotliEncoderStateStruct<StandardAlloc>,
+}
+
+impl BrotliEncoder {
+    pub(crate) fn new(params: BrotliEncoderParams) -> Self {
+        let mut state = BrotliEncoderCreateInstance(StandardAlloc::default());
+        state.params = params;
+        Self { state }
+    }
+
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+        op: BrotliEncoderOperation,
+    ) -> Result<()> {
+        let in_buf = input.unwritten();
+        let out_buf = output.unwritten_mut();
+
+        let mut input_len = 0;
+        let mut output_len = 0;
+
+        if BrotliEncoderCompressStream(
+            &mut self.state,
+            op,
+            &mut in_buf.len(),
+            in_buf,
+            &mut input_len,
+            &mut out_buf.len(),
+            out_buf,
+            &mut output_len,
+            &mut None,
+            &mut |_, _, _, _| (),
+        ) <= 0
+        {
+            return Err(Error::new(ErrorKind::Other, "brotli error"));
+        }
+
+        input.advance(input_len);
+        output.advance(output_len);
+
+        Ok(())
+    }
+}
+
+impl Encode for BrotliEncoder {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<()> {
+        self.encode(
+            input,
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
+        )
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_FLUSH,
+        )?;
+
+        Ok(BrotliEncoderHasMoreOutput(&self.state) == 0)
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_FINISH,
+        )?;
+
+        Ok(BrotliEncoderIsFinished(&self.state) == 1)
+    }
+}
+
+impl fmt::Debug for BrotliEncoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("BrotliEncoder")
+            .field("compress", &"<no debug>")
+            .finish()
+    }
+}

apollo-router/src/axum_factory/compression/codec/brotli/mod.rs

@@ -0,0 +1,5 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+mod encoder;
+
+pub(crate) use self::encoder::BrotliEncoder;

apollo-router/src/axum_factory/compression/codec/deflate/encoder.rs

@@ -0,0 +1,46 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+use std::io::Result;
+
+use flate2::Compression;
+
+use crate::axum_factory::compression::codec::Encode;
+use crate::axum_factory::compression::codec::FlateEncoder;
+use crate::axum_factory::compression::util::PartialBuffer;
+
+#[derive(Debug)]
+pub(crate) struct DeflateEncoder {
+    inner: FlateEncoder,
+}
+
+impl DeflateEncoder {
+    pub(crate) fn new(level: Compression) -> Self {
+        Self {
+            inner: FlateEncoder::new(level, false),
+        }
+    }
+}
+
+impl Encode for DeflateEncoder {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<()> {
+        self.inner.encode(input, output)
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.inner.flush(output)
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.inner.finish(output)
+    }
+}

apollo-router/src/axum_factory/compression/codec/deflate/mod.rs

@@ -0,0 +1,5 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+mod encoder;
+
+pub(crate) use self::encoder::DeflateEncoder;