implement streaming compression

Geal · Geal · commit 3a3412c2ff7e · 2023-04-24T16:47:56.000+02:00
this makes an assumption that async-commpression cannot make: every
chunk in the body stream has to be compressed and flushed directly, we
should not wait for more data to come and get better compression.
This is due to the multipart protocol for defer: we know each chunk
represent eitehr the primary or a deferred response, and should be sent
as soon as possible
diff --git a/apollo-router/src/axum_factory/axum_http_server_factory.rs b/apollo-router/src/axum_factory/axum_http_server_factory.rs
@@ -15,10 +15,14 @@ use axum::middleware::Next;
 use axum::response::*;
 use axum::routing::get;
 use axum::Router;
+use flate2::Compression;
 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,16 +36,17 @@ use tokio_rustls::TlsAcceptor;
 use tower::service_fn;
 use tower::BoxError;
 use tower::ServiceExt;
-use tower_http::compression::CompressionLayer;
 use tower_http::trace::TraceLayer;
 
+use super::compression::codec::DeflateEncoder;
 use super::listeners::ensure_endpoints_consistency;
 use super::listeners::ensure_listenaddrs_consistency;
 use super::listeners::extra_endpoints;
 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;
@@ -326,10 +331,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());
+        .layer(cors);
 
     let route = endpoints_on_main_listener
         .into_iter()
@@ -429,6 +431,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;
@@ -462,7 +469,25 @@ 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();
+
+            println!("will compress response, accept-encoding == {accept_encoding:?}");
+            let first: Option<&str> = accept_encoding
+                .as_ref()
+                .and_then(|value| value.to_str().ok())
+                .and_then(|v| v.split(',').map(|s| s.trim()).next());
+            println!("first: {first:?}");
+            let body = if first.is_none() {
+                body
+            } else {
+                let compressor = Compressor::Deflate(DeflateEncoder::new(Compression::fast()));
+                parts
+                    .headers
+                    .insert(CONTENT_ENCODING, HeaderValue::from_static("deflate"));
+                Body::wrap_stream(compressor.process(body))
+            };
+
+            http::Response::from_parts(parts, body).into_response()
         }
     }
 }
diff --git a/apollo-router/src/axum_factory/compression/mod.rs b/apollo-router/src/axum_factory/compression/mod.rs
@@ -1,7 +1,135 @@
-// 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
-//#[macro_use]
-//mod macros;
+use bytes::{Bytes, BytesMut};
+use futures::{Stream, StreamExt};
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tower::BoxError;
+
+use self::{
+    codec::{DeflateEncoder, Encode, GzipEncoder},
+    util::PartialBuffer,
+};
+
 pub(crate) mod codec;
 pub(crate) mod unshared;
 pub(crate) mod util;
+
+pub(crate) enum Compressor {
+    //Identity,
+    Deflate(DeflateEncoder),
+    Gzip(GzipEncoder),
+    //Brotli(BrotliEncoder),
+    //Zstd,
+    //others?
+}
+
+//FIXME: we should call finish at the end
+impl Compressor {
+    pub(crate) fn process(
+        mut self,
+        mut stream: hyper::Body,
+    ) -> impl Stream<Item = Result<Bytes, BoxError>>
+where {
+        let (tx, rx) = mpsc::channel(10);
+
+        tokio::task::spawn(async move {
+            while let Some(data) = stream.next().await {
+                match data {
+                    Err(e) => {
+                        tx.send(Err(e.into())).await;
+                    }
+                    Ok(data) => {
+                        let mut buf = BytesMut::zeroed(1024);
+                        let mut written = 0usize;
+
+                        let mut partial_input = PartialBuffer::new(&*data);
+                        loop {
+                            let mut partial_output = PartialBuffer::new(&mut buf);
+                            partial_output.advance(written);
+
+                            match self.encode(&mut partial_input, &mut partial_output) {
+                                Err(e) => panic!("{e:?}"),
+                                Ok(()) => {}
+                            }
+
+                            let read = partial_input.written().len();
+                            written += partial_output.written().len();
+                            println!("encode: read from input: {read}, written = {written}");
+
+                            if !partial_input.unwritten().is_empty() {
+                                // there was not enough space in the output buffer to compress everything,
+                                // so we resize and add more data
+                                if partial_output.unwritten().is_empty() {
+                                    let _ = partial_output.into_inner();
+                                    buf.reserve(written);
+                                }
+                            } else {
+                                // FIXME: what happens if we try to flush in a full buffer
+                                match self.flush(&mut partial_output) {
+                                    Err(e) => panic!("{e:?}"),
+                                    Ok(_) => {
+                                        let flushed = partial_output.written().len() - written;
+                                        println!("flush with buffer of size {flushed}");
+                                        let _ = partial_output.into_inner();
+                                        buf.resize(flushed, 0);
+                                        tx.send(Ok(buf.freeze())).await;
+                                        break;
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            let buf = BytesMut::zeroed(64);
+            let mut partial_output = PartialBuffer::new(buf);
+
+            match self.finish(&mut partial_output) {
+                Err(e) => panic!("{e:?}"),
+                Ok(b) => {
+                    let len = partial_output.written().len();
+                    println!("finish with buffer of size {}", len);
+
+                    let mut buf = partial_output.into_inner();
+                    buf.resize(len, 0);
+                    tx.send(Ok(buf.freeze())).await;
+                }
+            }
+            //tx.send(partial_output.into_inner().freeze());
+        });
+        ReceiverStream::new(rx)
+    }
+}
+
+impl Encode for Compressor {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> std::io::Result<()> {
+        match self {
+            Compressor::Deflate(e) => e.encode(input, output),
+            Compressor::Gzip(e) => e.encode(input, output),
+        }
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> std::io::Result<bool> {
+        match self {
+            Compressor::Deflate(e) => e.flush(output),
+            Compressor::Gzip(e) => e.flush(output),
+        }
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> std::io::Result<bool> {
+        match self {
+            Compressor::Deflate(e) => e.finish(output),
+            Compressor::Gzip(e) => e.finish(output),
+        }
+    }
+}
