Spaces:
Build error
Build error
File size: 4,215 Bytes
d8435ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
use std::task::{Context, Poll};
use futures_util::future::BoxFuture;
use tonic::body::BoxBody;
use tonic::codegen::http::Response;
use tonic::Code;
use tower::Service;
use tower_layer::Layer;
#[derive(Clone)]
pub struct LoggingMiddleware<T> {
inner: T,
}
#[derive(Clone)]
pub struct LoggingMiddlewareLayer;
impl LoggingMiddlewareLayer {
pub fn new() -> Self {
Self {}
}
}
impl<S> Service<tonic::codegen::http::Request<tonic::transport::Body>> for LoggingMiddleware<S>
where
S: Service<tonic::codegen::http::Request<tonic::transport::Body>, Response = Response<BoxBody>>
+ Clone,
S::Future: Send + 'static,
{
type Response = S::Response;
type Error = S::Error;
type Future = BoxFuture<'static, Result<S::Response, S::Error>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(
&mut self,
request: tonic::codegen::http::Request<tonic::transport::Body>,
) -> Self::Future {
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
let method_name = request.uri().path().to_string();
let instant = std::time::Instant::now();
let future = inner.call(request);
Box::pin(async move {
let response = future.await;
let elapsed_sec = instant.elapsed().as_secs_f32();
match response {
Err(error) => {
log::error!("gGRPC request error {}", method_name);
Err(error)
}
Ok(response_tonic) => {
let grpc_status = tonic::Status::from_header_map(response_tonic.headers());
if let Some(grpc_status) = grpc_status {
match grpc_status.code() {
Code::Ok => {
log::trace!("gRPC {} Ok {:.6}", method_name, elapsed_sec);
}
Code::Cancelled => {
// cluster mode generates a large amount of `stream error received: stream no longer needed`
log::trace!("gRPC {} {:.6}", method_name, elapsed_sec);
}
Code::DeadlineExceeded
| Code::Aborted
| Code::OutOfRange
| Code::ResourceExhausted
| Code::NotFound
| Code::InvalidArgument
| Code::AlreadyExists
| Code::FailedPrecondition
| Code::PermissionDenied
| Code::Unauthenticated => {
log::info!(
"gRPC {} failed with {} {:?} {:.6}",
method_name,
grpc_status.code(),
grpc_status.message(),
elapsed_sec,
);
}
Code::Internal
| Code::Unimplemented
| Code::Unavailable
| Code::DataLoss
| Code::Unknown => log::error!(
"gRPC {} unexpectedly failed with {} {:?} {:.6}",
method_name,
grpc_status.code(),
grpc_status.message(),
elapsed_sec,
),
};
} else {
log::trace!("gRPC {} Ok {:.6}", method_name, elapsed_sec);
}
Ok(response_tonic)
}
}
})
}
}
impl<S> Layer<S> for LoggingMiddlewareLayer {
type Service = LoggingMiddleware<S>;
fn layer(&self, service: S) -> Self::Service {
LoggingMiddleware { inner: service }
}
}
|