如何在Rust Hyper 0.14中计算HTTP请求与响应的大小

由于Hyper 0.14的HTTP请求/响应Body是流式结构，无法直接获取总大小，核心解决思路是包装原始Body流，在数据流经时累计字节数，同时保持Body的Stream特性，再通过中间件集成到网关的请求处理流程中。

实现步骤 #

1. 定义Body包装器用于字节统计 #

首先实现一个通用的CountedBody结构体，包裹原始Body流，每次读取数据时更新累计字节数：

use hyper::{Body, Error};
use bytes::Bytes;
use futures_core::Stream;
use std::pin::Pin;
use std::task::{Context, Poll};

#[derive(Debug)]
struct CountedBody<S> {
    inner: S,
    bytes_count: usize,
}

impl<S> CountedBody<S> {
    // 创建新的统计Body
    fn new(inner: S) -> Self {
        CountedBody {
            inner,
            bytes_count: 0,
        }
    }

    // 获取累计的字节数
    fn get_count(&self) -> usize {
        self.bytes_count
    }
}

// 实现Stream trait，转发流数据并统计字节
impl<S> Stream for CountedBody<S>
where
    S: Stream<Item = Result<Bytes, Error>> + Unpin,
{
    type Item = Result<Bytes, Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(cx) {
            Poll::Ready(Some(Ok(bytes))) => {
                self.bytes_count += bytes.len();
                Poll::Ready(Some(Ok(bytes)))
            }
            // 转发错误或流结束信号
            other => other,
        }
    }
}

2. 实现统计中间件 #

基于Hyper的Service trait实现中间件，将请求/响应的Body替换为CountedBody，并在流处理完成后获取统计结果：

use hyper::{service::Service, Request, Response};
use std::future::Future;
use std::task::{Context, Poll};
use tokio;

struct CountingMiddleware<T> {
    inner: T,
}

impl<T> CountingMiddleware<T> {
    fn new(inner: T) -> Self {
        CountingMiddleware { inner }
    }
}

impl<T, ReqBody, ResBody> Service<Request<ReqBody>> for CountingMiddleware<T>
where
    T: Service<Request<ReqBody>, Response = Response<ResBody>, Error = Error>,
    T::Future: 'static + Send,
    ReqBody: Stream<Item = Result<Bytes, Error>> + Unpin + 'static,
    ResBody: Stream<Item = Result<Bytes, Error>> + Unpin + 'static,
{
    type Response = Response<CountedBody<ResBody>>;
    type Error = Error;
    type Future = Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        // 包装请求Body以统计请求大小
        let mut counted_req_body = CountedBody::new(req.into_body());
        let req = Request::from_parts(req.into_parts().0, counted_req_body);

        let inner_fut = self.inner.call(req);

        Box::pin(async move {
            let res = inner_fut.await?;
            // 包装响应Body以统计响应大小
            let mut counted_res_body = CountedBody::new(res.into_body());
            let mut res = Response::from_parts(res.into_parts().0, counted_res_body);

            // 异步统计请求大小：等待请求Body流消费完成
            let req_body = req.body_mut();
            tokio::spawn(async move {
                while let Some(_) = req_body.next().await {}
                println!("Request size: {} bytes", req_body.get_count());
            });

            // 异步统计响应大小：等待响应Body流消费完成
            let res_body = res.body_mut();
            tokio::spawn(async move {
                while let Some(_) = res_body.next().await {}
                println!("Response size: {} bytes", res_body.get_count());
            });

            Ok(res)
        })
    }
}

3. 集成到网关服务 #

将中间件包装到你的核心服务上，启动Hyper服务器：

use hyper::{Server, service::make_service_fn, service::service_fn};
use std::convert::Infallible;

async fn handle_request(req: Request<Body>) -> Result<Response<Body>, Infallible> {
    // 你的核心请求处理逻辑
    Ok(Response::new(Body::from("Hello from gateway")))
}

#[tokio::main]
async fn main() {
    let addr = ([127, 0, 0, 1], 3000).into();

    let make_svc = make_service_fn(|_conn| async {
        // 用统计中间件包装核心服务
        let service = CountingMiddleware::new(service_fn(handle_request));
        Ok::<_, Infallible>(service)
    });

    let server = Server::bind(&addr).serve(make_svc);

    if let Err(e) = server.await {
        eprintln!("server error: {}", e);
    }
}

关键注意事项 #

• 分块编码处理：如果请求/响应使用Transfer-Encoding: chunked，无法通过Content-Length头获取大小，必须通过流统计才能得到准确值。
• 计数时机：只有当整个Body流被完全消费后，get_count()返回的才是完整的字节数。如果需要在处理早期获取大小，可优先检查Content-Length头，不存在时再回退到流统计。
• 线程安全扩展：如果需要在多个组件间共享计数，可将bytes_count替换为Arc<AtomicUsize>，避免线程安全问题。

内容的提问来源于stack exchange，提问作者Developer