diff --git a/Cargo.lock b/Cargo.lock index 28ef665..c8c95c2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -223,6 +223,7 @@ dependencies = [ "chrono", "clap", "hashbrown", + "http-body", "itertools 0.14.0", "rand", "rand_chacha", diff --git a/Cargo.toml b/Cargo.toml index f402615..390adf7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ axum-server = { version = "0.7.2", features = ["tls-rustls"] } chrono = "0.4.40" clap = { version = "4.5.37", features = ["derive"] } hashbrown = "0.15.2" +http-body = "1.0.1" itertools = "0.14.0" rand = "0.9.1" rand_chacha = "0.9.0" diff --git a/src/main.rs b/src/main.rs index 5829fc0..3bcd53f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,21 +1,32 @@ mod generator; -use axum::{extract::Path, response::Html, routing::get, Router}; +use axum::{ + body::{Body, Bytes}, + extract::Path, + response::{Html, IntoResponse, Response}, + routing::get, + Router, +}; use axum_server::{bind_rustls, tls_rustls::RustlsConfig}; use clap::Parser; use hashbrown::HashMap; +use http_body::Frame; use itertools::Itertools; use rand::{prelude::*, Rng as _}; use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng}; use std::{ borrow::Cow, + convert::Infallible, path::PathBuf, + pin::Pin, sync::{ atomic::{AtomicU64, Ordering}, Arc, }, + task::{Context, Poll}, time::Duration, }; +use tokio::time::{interval, Interval}; pub type Rng = ChaCha8Rng; @@ -43,6 +54,42 @@ fn create_rng(seed_bytes: impl IntoIterator) -> Rng { const COUNT_FILE: &str = "count.txt"; +const SLOW_CHUNK_SIZE: usize = 50; +const SLOW_DURATION: Duration = Duration::from_millis(250); + +struct SlowBody { + bytes: Bytes, + interval: Interval, +} + +impl http_body::Body for SlowBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + match self.interval.poll_tick(cx) { + Poll::Ready(_) => { + if self.bytes.is_empty() { + Poll::Ready(None) + } else { + let split_len = self.bytes.len().min(SLOW_CHUNK_SIZE); + Poll::Ready(Some(Ok(Frame::data(self.bytes.split_to(split_len))))) + } + } + Poll::Pending => Poll::Pending, + } + } +} + +impl IntoResponse for SlowBody { + fn into_response(self) -> Response { + Html(Body::new(self)).into_response() + } +} + #[tokio::main] async fn main() { let args = Args::parse(); @@ -105,7 +152,7 @@ async fn main() { content }); - Html(format!( + let html = format!( " @@ -119,7 +166,12 @@ async fn main() { " - )) + ); + + SlowBody { + bytes: html.into(), + interval: interval(SLOW_DURATION), + } }), ) };