Middleware for compressing SSE streams

[stevenengler]

Currently tower-http's CompressionLayer middleware does not support server-sent event streams due to buffering within the encoder, which causes the SSE events to be delayed. For this reason tower does not apply compression to "text/event-stream" responses.

Since the lack of compression is a bit painful when using SSE with htmx (streaming html without compression can use way too much bandwidth), I wrote a small proof-of-concept middleware for applying gzip compression in axum that can be used as a workaround. I'm hoping this helps someone else who wants to apply compression to their SSE streams, since I haven't seen any other examples of it.

use axum::body::{Body, BodyDataStream, Bytes};
use axum::extract::Request;
use axum::http::{header};
use axum::middleware::Next;
use axum::response::Response;
use flate2::write::GzEncoder;
use flate2::Compression;
use futures_util::stream::Stream;

use std::io::Write;
use std::pin::{pin, Pin};
use std::task::Context;
use std::task::Poll;

pub async fn compress_sse(request: Request, next: Next) -> Response {
    let response = next.run(request).await;

    // TODO: check if response already has a `CONTENT_ENCODING` header
    // TODO: check if request allows gzip compression with the `ACCEPT_ENCODING` header

    if let Some(content_type) = response.headers().get(header::CONTENT_TYPE) {
        if content_type.as_bytes() != b"text/event-stream" {
            return response;
        }
    } else {
        return response;
    }

    let (mut parts, body) = response.into_parts();

    let body = body.into_data_stream();
    let body = Body::from_stream(CompressedStream::new(body));

    parts.headers.insert(
        header::CONTENT_ENCODING,
        header::HeaderValue::from_static("gzip"),
    );

    Response::from_parts(parts, body)
}

struct CompressedStream {
    inner: BodyDataStream,
    compression: GzEncoder<Vec<u8>>,
}

impl CompressedStream {
    pub fn new(body: BodyDataStream) -> Self {
        Self {
            inner: body,
            compression: GzEncoder::new(Vec::new(), Compression::default()),
        }
    }
}

impl Stream for CompressedStream {
    type Item = Result<Bytes, axum::Error>;

    #[inline]
    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match pin!(&mut self.inner).as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(x))) => {
                self.compression.write_all(&x).unwrap();
                self.compression.flush().unwrap();

                let mut buf = Vec::new();
                std::mem::swap(&mut buf, self.compression.get_mut());

                Poll::Ready(Some(Ok(buf.into())))
            }
            x => x,
        }
    }
}

and use it with .layer(axum::middleware::from_fn(compress_sse)).