use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use http::Request;
use tonic::{body::Body, metadata::MetadataMap};
use tower::Service;
use crate::context::{RequestContext, layer::ContextService, run_with_context};
type GrpcRequest = Request<Body>;
impl<S> Service<GrpcRequest> for ContextService<S>
where
S: Service<GrpcRequest> + Clone + Send + 'static,
S::Future: Send,
{
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, req: GrpcRequest) -> Self::Future {
let inner = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, inner);
let metadata = MetadataMap::from_headers(req.headers().clone());
let ctx = RequestContext::from_grpc(&metadata);
Box::pin(async move { run_with_context(ctx, inner.call(req)).await })
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::{get_context, layer::ContextLayer};
use tonic::body::Body;
use tower::{Layer, ServiceExt};
#[tokio::test]
async fn test_grpc_context_layer_with_request_id() {
let layer = ContextLayer;
let service = layer.layer(tower::service_fn(|_req: GrpcRequest| async {
let ctx = get_context().expect("get_context should work within context layer");
let request_id = ctx.request_id().to_string();
assert_eq!(request_id, "grpc-test-123");
Ok::<_, std::convert::Infallible>(http::Response::new(Body::empty()))
}));
let req = Request::builder()
.uri("/test.Service/Method")
.method("POST")
.header("x-request-id", "grpc-test-123")
.body(Body::empty())
.unwrap();
let response: http::Response<Body> = service.oneshot(req).await.unwrap();
assert_eq!(response.status(), 200);
}
#[tokio::test]
async fn test_grpc_context_layer_auto_generate_request_id() {
let layer = ContextLayer;
let service = layer.layer(tower::service_fn(|_req: GrpcRequest| async {
let ctx = get_context().expect("get_context should work within context layer");
let request_id = ctx.request_id().to_string();
assert!(!request_id.is_empty());
Ok::<_, std::convert::Infallible>(http::Response::new(Body::empty()))
}));
let req = Request::builder()
.uri("/test.Service/Method")
.method("POST")
.body(Body::empty())
.unwrap();
let response: http::Response<Body> = service.oneshot(req).await.unwrap();
assert_eq!(response.status(), 200);
}
}