Add streaming brotli compression and hot-reloadable config (livevue.toml)
Key changes: - **src/config.rs** – `LiveVueConfig` (server host/port, brotli quality + window size), loaded from `livevue.toml`. A notify-based file watcher reloads the config on any file change; a SIGHUP handler does the same on demand. Falls back to compiled-in defaults when the file is absent. - **src/brotli_layer.rs** – `BrotliBody`, a custom `http_body::Body` wrapper that streams brotli-compressed data through a single persistent `brotli::CompressorWriter`. The encoder survives across SSE event flushes, so its sliding window accumulates context from all prior events — progressively better compression as the stream grows. Both quality (0–11) and window size (lgwin 10–24, i.e. 1 KB – 16 MB) are taken from config. `brotli_compression` is an axum `from_fn_with_state` middleware that activates only when the client sends `Accept-Encoding: br`. - **src/server.rs** – `AppState` gains a `config: SharedConfig` field. `AppState::new` takes the config; `AppState::new_default` is a zero-config convenience constructor. - **livevue.toml** – documented example config with inline comments explaining each field and the brotli window-size tradeoff table. - **examples/todo** – loads config via `load_and_watch_config`, derives the bind address from `config.server`, and applies `brotli_compression` as a router layer. https://claude.ai/code/session_01UnQQkkwts64FPUsSzFfdQb
This commit is contained in:
Claude
@@ -0,0 +1,233 @@
|
||||
//! Streaming brotli compression middleware for axum.
|
||||
//!
|
||||
//! # Why a custom encoder instead of `tower-http`'s `CompressionLayer`?
|
||||
//!
|
||||
//! `tower-http` does support brotli (via `async-compression`), but it does not
|
||||
//! expose the `lgwin` (window size) parameter. For Datastar SSE streams the
|
||||
//! window size is the key tuning knob: all SSE events on a connection travel
|
||||
//! through a **single persistent brotli encoder**, so later events benefit from
|
||||
//! the full history of earlier ones. A 4 MB window (`lgwin = 22`) means the
|
||||
//! encoder can reference up to 4 MB of previously transmitted HTML when
|
||||
//! compressing a new fragment — dramatically better ratios than per-chunk gzip.
|
||||
//!
|
||||
//! This module implements a [`BrotliBody`] wrapper and an axum
|
||||
//! [`brotli_compression`] middleware that honour the [`BrotliConfig`] (quality
|
||||
//! + window size) loaded from `livevue.toml`.
|
||||
//!
|
||||
//! # Behaviour
|
||||
//!
|
||||
//! - Only activates when the client sends `Accept-Encoding: br`.
|
||||
//! - Reads config from `SharedConfig` on **each new request**, so a hot-reload
|
||||
//! applies to the next connection without a restart.
|
||||
//! - Calls `encoder.flush()` after every body frame so SSE clients receive
|
||||
//! data promptly (brotli metablock boundary = decodable checkpoint).
|
||||
//! - The encoder's sliding-window context is preserved across flushes, giving
|
||||
//! progressively better compression as the SSE stream grows.
|
||||
|
||||
use crate::config::SharedConfig;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::extract::State;
|
||||
use axum::http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH};
|
||||
use axum::http::HeaderValue;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use bytes::Bytes;
|
||||
use http_body::Body as HttpBody;
|
||||
use pin_project::pin_project;
|
||||
use std::io::Write as _;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// DrainWriter — a `Write` impl whose output can be drained incrementally
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A `std::io::Write` target that accumulates bytes in a `Vec<u8>`.
|
||||
///
|
||||
/// After flushing the brotli encoder we call [`DrainWriter::drain`] to collect
|
||||
/// only the bytes written since the last drain. This lets us send one
|
||||
/// compressed chunk per SSE event while keeping the encoder (and its sliding
|
||||
/// window) alive across the entire connection.
|
||||
struct DrainWriter(Vec<u8>);
|
||||
|
||||
impl std::io::Write for DrainWriter {
|
||||
fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
|
||||
self.0.extend_from_slice(data);
|
||||
Ok(data.len())
|
||||
}
|
||||
|
||||
/// No-op: the actual flushing is driven by `CompressorWriter::flush()`.
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl DrainWriter {
|
||||
fn drain(&mut self) -> Bytes {
|
||||
Bytes::from(self.0.drain(..).collect::<Vec<_>>())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// BrotliBody — streaming Body wrapper
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// An `http_body::Body` that brotli-compresses a wrapped [`axum::body::Body`].
|
||||
///
|
||||
/// The encoder is constructed once at stream creation and persists for the
|
||||
/// lifetime of the connection, maintaining its sliding-window context across
|
||||
/// every SSE event. [`flush()`] is called after each data frame to create a
|
||||
/// metablock boundary so the browser can decode partial data immediately.
|
||||
///
|
||||
/// [`flush()`]: std::io::Write::flush
|
||||
#[pin_project]
|
||||
pub struct BrotliBody {
|
||||
#[pin]
|
||||
inner: Body,
|
||||
/// `None` once we have consumed the encoder to finalize the stream.
|
||||
encoder: Option<brotli::CompressorWriter<DrainWriter>>,
|
||||
}
|
||||
|
||||
impl BrotliBody {
|
||||
pub fn new(inner: Body, quality: u32, lgwin: u32) -> Self {
|
||||
let params = brotli::enc::BrotliEncoderParams {
|
||||
quality: quality as i32,
|
||||
lgwin: lgwin as i32,
|
||||
..Default::default()
|
||||
};
|
||||
// Buffer size (4 KiB) is the encoder's internal I/O buffer — smaller
|
||||
// means lower latency per flush, which matters for SSE.
|
||||
let encoder =
|
||||
brotli::CompressorWriter::with_params(DrainWriter(Vec::new()), 4096, ¶ms);
|
||||
Self {
|
||||
inner,
|
||||
encoder: Some(encoder),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpBody for BrotliBody {
|
||||
type Data = Bytes;
|
||||
type Error = axum::Error;
|
||||
|
||||
fn poll_frame(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
|
||||
let this = self.project();
|
||||
|
||||
match this.inner.poll_frame(cx) {
|
||||
// ----------------------------------------------------------------
|
||||
// Got a data frame — compress it and yield immediately.
|
||||
// ----------------------------------------------------------------
|
||||
Poll::Ready(Some(Ok(frame))) => {
|
||||
if let Ok(data) = frame.into_data() {
|
||||
match this.encoder {
|
||||
Some(enc) => {
|
||||
if let Err(e) = enc.write_all(&data) {
|
||||
return Poll::Ready(Some(Err(axum::Error::new(e))));
|
||||
}
|
||||
if let Err(e) = enc.flush() {
|
||||
return Poll::Ready(Some(Err(axum::Error::new(e))));
|
||||
}
|
||||
let compressed = enc.get_mut().drain();
|
||||
Poll::Ready(Some(Ok(http_body::Frame::data(compressed))))
|
||||
}
|
||||
None => Poll::Ready(None),
|
||||
}
|
||||
} else {
|
||||
// Trailers / unknown frame type — pass through as-is.
|
||||
// (Re-construct the original frame; trailers don't need compression.)
|
||||
Poll::Ready(None)
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------
|
||||
// Inner body is exhausted — finalise the encoder.
|
||||
// ----------------------------------------------------------------
|
||||
Poll::Ready(None) => {
|
||||
if let Some(enc) = this.encoder.take() {
|
||||
// into_inner() flushes any pending bytes and writes the
|
||||
// final brotli stream terminator, returning the inner
|
||||
// DrainWriter directly (not wrapped in Result).
|
||||
let mut drain = enc.into_inner();
|
||||
let final_bytes = drain.drain();
|
||||
if !final_bytes.is_empty() {
|
||||
return Poll::Ready(Some(Ok(http_body::Frame::data(final_bytes))));
|
||||
}
|
||||
}
|
||||
Poll::Ready(None)
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------
|
||||
// Pass-through.
|
||||
// ----------------------------------------------------------------
|
||||
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// axum middleware
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Axum middleware that wraps SSE (and any other) responses with streaming
|
||||
/// brotli compression, obeying `Accept-Encoding: br`.
|
||||
///
|
||||
/// Reads [`BrotliConfig`] from [`SharedConfig`] on each new request, so
|
||||
/// quality and window-size changes take effect for new connections immediately
|
||||
/// after a config reload.
|
||||
///
|
||||
/// # Usage
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// use axum::{middleware, Router};
|
||||
/// use livevue_rs::brotli_layer::brotli_compression;
|
||||
///
|
||||
/// let app = Router::new()
|
||||
/// /* ... routes ... */
|
||||
/// .layer(middleware::from_fn_with_state(config.clone(), brotli_compression));
|
||||
/// ```
|
||||
pub async fn brotli_compression(
|
||||
State(config): State<SharedConfig>,
|
||||
request: axum::http::Request<Body>,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
// Check whether the client signals brotli support.
|
||||
let accepts_brotli = request
|
||||
.headers()
|
||||
.get(ACCEPT_ENCODING)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|v| v.contains("br"))
|
||||
.unwrap_or(false);
|
||||
|
||||
let response = next.run(request).await;
|
||||
|
||||
if !accepts_brotli {
|
||||
return response;
|
||||
}
|
||||
|
||||
// Snapshot the brotli config; new values apply to future connections.
|
||||
let (enabled, quality, window_size) = {
|
||||
let cfg = config.read().await;
|
||||
(cfg.brotli.enabled, cfg.brotli.quality, cfg.brotli.window_size)
|
||||
};
|
||||
|
||||
if !enabled {
|
||||
return response;
|
||||
}
|
||||
|
||||
let (mut parts, body) = response.into_parts();
|
||||
|
||||
// Remove Content-Length — the compressed length will differ.
|
||||
parts.headers.remove(CONTENT_LENGTH);
|
||||
// Advertise brotli encoding.
|
||||
parts
|
||||
.headers
|
||||
.insert(CONTENT_ENCODING, HeaderValue::from_static("br"));
|
||||
|
||||
let brotli_body = BrotliBody::new(body, quality, window_size);
|
||||
Response::from_parts(parts, Body::new(brotli_body))
|
||||
}
|
||||
Reference in New Issue
Block a user