main
//! gRPC 请求日志：包装 tonic 的服务栈，为每次调用打印一行结构化日志。
//!
//! gRPC 的错误码放在 HTTP/2 的 trailer 里，所以响应体也必须包一层，等到 trailer
//! 到达后才能拿到真实的 `grpc-status`。这里同时兼容把状态写在响应头的情况
//! （反射、健康检查与部分异常路径）。
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Instant;
use bytes::Bytes;
use envelope::{LEGACY_REQUEST_ID_HEADER, REQUEST_ID_HEADER, new_request_id, validate_request_id};
use http::header::{HeaderName, HeaderValue};
use http::{HeaderMap, Request, Response};
use http_body::{Body, Frame, SizeHint};
use tower::{Layer, Service};
use tracing::Instrument;
/// 调用方身份头（可选，缺省表示匿名调用）。
const SERVICE_ID_HEADER: &str = "gitdata-service-id";
const GRPC_STATUS: &str = "grpc-status";
/// gRPC 请求日志层，通过 `tonic::transport::Server::layer` 装配。
#[derive(Clone, Copy, Debug, Default)]
pub struct GrpcRequestLog;
impl<S> Layer<S> for GrpcRequestLog {
type Service = GrpcRequestLogService<S>;
fn layer(&self, inner: S) -> Self::Service {
GrpcRequestLogService { inner }
}
}
#[derive(Clone, Copy, Debug)]
pub struct GrpcRequestLogService<S> {
inner: S,
}
impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for GrpcRequestLogService<S>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>>,
S::Future: Send + 'static,
ResBody: Body<Data = Bytes> + Send + 'static,
{
type Response = Response<LoggedBody<ResBody>>;
type Error = S::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(context)
}
fn call(&mut self, mut request: Request<ReqBody>) -> Self::Future {
let request_id = {
let headers = request.headers();
headers
.get(REQUEST_ID_HEADER)
.or_else(|| headers.get(LEGACY_REQUEST_ID_HEADER))
.and_then(|value| value.to_str().ok())
.and_then(validate_request_id)
.unwrap_or_else(new_request_id)
};
if let Ok(value) = HeaderValue::from_str(&request_id) {
request
.headers_mut()
.insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
}
let call = Call::new(&request, request_id.clone());
let span = tracing::info_span!(
"grpc.server.request",
"otel.kind" = "server",
"rpc.system" = "grpc",
"rpc.service" = %call.service,
"rpc.method" = %call.method,
"gitdata.service_id" = %call.caller,
request_id = %request_id,
);
let traceparent = request
.headers()
.get(crate::TRACEPARENT)
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
crate::set_parent_from_traceparent(&span, traceparent.as_deref());
let started = Instant::now();
let future = self.inner.call(request);
let context_id = request_id.clone();
Box::pin(envelope::context::with_request_id(
Some(context_id),
async move {
let response = future.await?;
let (mut parts, body) = response.into_parts();
if let Ok(value) = HeaderValue::from_str(&request_id) {
parts
.headers
.insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
}
let header_status = status_of(&parts.headers);
// 错误响应（含未认证、反射失败）不带响应体，trailer 永远不会被轮询到，
// 所以结束流的情况必须在包装时就记一笔。
let ended = body.is_end_stream();
let mut logged = LoggedBody {
inner: Box::pin(body),
call,
started,
header_status,
logged: false,
};
if ended {
logged.finish(header_status);
}
Ok(Response::from_parts(parts, logged))
}
.instrument(span),
))
}
}
/// 一次调用的静态信息，在请求进栈时抓取。
#[derive(Clone, Debug)]
struct Call {
method: String,
service: String,
request_id: String,
caller: String,
}
impl Call {
fn new<B>(request: &Request<B>, request_id: String) -> Self {
let path = request.uri().path();
let (service, method) = match path.rsplit_once('/') {
Some((service, method)) => (
service.trim_start_matches('/').to_owned(),
method.to_owned(),
),
None => (String::new(), path.to_owned()),
};
Self {
method,
service,
request_id,
caller: header(request.headers(), SERVICE_ID_HEADER),
}
}
fn emit(&self, status: Option<i64>, elapsed_ms: f64) {
let code = status.unwrap_or(-1);
let duration_ms = (elapsed_ms * 1_000.0).round() / 1_000.0;
let caller = if self.caller.is_empty() {
"anonymous"
} else {
self.caller.as_str()
};
let request_id = if self.request_id.is_empty() {
"-"
} else {
self.request_id.as_str()
};
if code == 0 {
tracing::info!(
"rpc.service" = %self.service,
"rpc.method" = %self.method,
"rpc.system" = "grpc",
"rpc.grpc.status_code" = code,
"duration_ms" = duration_ms,
"gitdata.service_id" = caller,
"request_id" = request_id,
"grpc request"
);
} else {
tracing::warn!(
"rpc.service" = %self.service,
"rpc.method" = %self.method,
"rpc.system" = "grpc",
"rpc.grpc.status_code" = code,
"duration_ms" = duration_ms,
"gitdata.service_id" = caller,
"request_id" = request_id,
"grpc request failed"
);
}
}
}
/// 透传响应体，等 trailer 出来时补上一条日志。
pub struct LoggedBody<B> {
inner: Pin<Box<B>>,
call: Call,
started: Instant,
header_status: Option<i64>,
logged: bool,
}
impl<B: Body<Data = Bytes>> LoggedBody<B> {
fn finish(&mut self, status: Option<i64>) {
if self.logged {
return;
}
self.logged = true;
let elapsed_ms = self.started.elapsed().as_secs_f64() * 1_000.0;
self.call.emit(status, elapsed_ms);
}
}
impl<B: Body<Data = Bytes>> Body for LoggedBody<B> {
type Data = Bytes;
type Error = B::Error;
fn poll_frame(
self: Pin<&mut Self>,
context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
// `Pin<Box<B>>` 让整体始终 Unpin，无需 unsafe 投影。
let this = self.get_mut();
let polled = this.inner.as_mut().poll_frame(context);
match &polled {
Poll::Ready(Some(Ok(frame))) => {
if let Some(trailers) = frame.trailers_ref() {
this.finish(status_of(trailers).or(this.header_status));
}
}
Poll::Ready(Some(Err(_))) => this.finish(this.header_status),
Poll::Ready(None) => this.finish(this.header_status),
Poll::Pending => {}
}
polled
}
fn is_end_stream(&self) -> bool {
self.inner.is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.inner.size_hint()
}
}
impl<B> fmt::Debug for LoggedBody<B> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("LoggedBody")
.field("call", &self.call)
.field("header_status", &self.header_status)
.field("logged", &self.logged)
.finish_non_exhaustive()
}
}
fn header(headers: &HeaderMap, name: &str) -> String {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_owned()
}
fn status_of(headers: &HeaderMap) -> Option<i64> {
headers
.get(GRPC_STATUS)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse().ok())
}