如何在Rust Tonic中间件中获取gRPC状态码？

你当前的中间件是作用在HTTP底层服务层，只能拿到HTTP响应对象，而gRPC状态码并不在HTTP Body中，而是通过响应元数据（Trailer）传递的——或者当请求失败时直接以tonic::Status错误返回。下面提供两种高效的解决方案，无需解析HTTP Body：

方案一：直接适配Tonic的gRPC服务层（推荐） #

将中间件改为针对tonic::Request/tonic::Response的Service实现，直接和Tonic的gRPC API交互，能最直接、无额外开销地获取gRPC状态码。

实现代码 #

use tonic::{Request, Response, Status};
use tower::{Service, Layer};
use std::future::{Future};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Instant;
use tracing::{debug, error};

// 日志中间件结构体
pub struct GrpcLoggingMiddleware<S> {
    service_name: &'static str,
    inner: S,
}

impl<S> GrpcLoggingMiddleware<S> {
    pub fn new(inner: S, service_name: &'static str) -> Self {
        Self { inner, service_name }
    }
}

// 实现Layer trait，方便在Server中添加
impl<S> Layer<S> for GrpcLoggingMiddleware<S> {
    type Service = GrpcLoggingMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        GrpcLoggingMiddleware::new(inner, self.service_name)
    }
}

// 实现Service trait，处理Tonic的gRPC请求
impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for GrpcLoggingMiddleware<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>, Error = Status>,
    S::Future: Future<Output = Result<Response<ResBody>, Status>> + Send + 'static,
    ReqBody: Send + 'static,
    ResBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = Pin<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, request: Request<ReqBody>) -> Self::Future {
        let path = request.uri().path().to_string();
        let start = Instant::now();
        let service_name = self.service_name;
        let inner_future = self.inner.call(request);

        Box::pin(async move {
            match inner_future.await {
                Ok(response) => {
                    // 从响应元数据中读取gRPC状态码（默认0对应OK）
                    let grpc_status = response.metadata()
                        .get("grpc-status")
                        .and_then(|v| v.to_str().ok())
                        .and_then(|s| s.parse::<u32>().ok())
                        .unwrap_or(0);

                    debug!(
                        "grpc request completed | service: {}, path: {}, elapsed: {}us, grpc_status: {}",
                        service_name, path, start.elapsed().as_micros(), grpc_status
                    );
                    Ok(response)
                }
                Err(status) => {
                    error!(
                        "grpc request failed | service: {}, path: {}, elapsed: {}us, grpc_status: {}, message: {}",
                        service_name, path, start.elapsed().as_micros(), status.code() as u32, status.message()
                    );
                    Err(status)
                }
            }
        })
    }
}

使用方式 #

let logging_layer = GrpcLoggingMiddleware::new(() as _, "your_grpc_service");

Server::builder()
    .layer(logging_layer)
    .add_service(your_grpc_service_impl())
    .serve(addr).await?;

方案二：在HTTP层读取gRPC Trailer（兼容现有中间件结构） #

如果不想改动现有中间件的HTTP层结构，可以从HTTP响应的Trailer头中读取grpc-status字段，无需解析Body：

修改原中间件的poll方法 #

impl<F, Res, E> Future for ResponseFuture<F>
where
    F: Future<Output = Result<http::response::Response<Res>, E>>,
    E: std::fmt::Debug,
    Res: std::fmt::Debug,
{
    type Output = Result<http::response::Response<Res>, E>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        match this.response_future.poll(cx) {
            Poll::Ready(result) => {
                match &result {
                    Err(e) => error!("grpc response error: {:?}", e),
                    Ok(response) => {
                        // 从HTTP响应Trailer中读取gRPC状态码
                        let grpc_status = response.headers()
                            .get("grpc-status")
                            .and_then(|v| v.to_str().ok())
                            .and_then(|s| s.parse::<u32>().ok())
                            .unwrap_or(0);

                        debug!(
                            "finish grpc request | path: {}, elapsed: {}us, http_status: {}, grpc_status: {}",
                            this.path,
                            this.start.elapsed().as_micros(),
                            response.status(),
                            grpc_status
                        );
                    }
                }
                Poll::Ready(result)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

关键说明 #

1. gRPC状态码不会放在HTTP Body中，而是通过grpc-status这个Trailer字段传递，成功响应时默认值为0（对应Code::Ok）。
2. 方案一更可靠，因为直接和Tonic的gRPC服务层交互，能直接捕获tonic::Status错误，无需依赖HTTP头的解析。

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