rpc/
stack.rs

1//! Generated-service-independent client and server assembly.
2
3use crate::{BearerToken, RpcMetrics, RpcMetricsLayer};
4use futures::FutureExt;
5use http::{Request, Response};
6use http_body::{Body, Frame};
7use pin_project_lite::pin_project;
8use rust_zero_core::{
9    AdaptiveShedder, AuthFailure, CircuitBreaker, CircuitBreakerConfig, CircuitBreakerPermit,
10    CircuitOutcome, LoadShedderConfig, LogContext, LogField, LogLevel, Logger, ShedPermit,
11    TraceContext, TraceFlags,
12};
13use std::{
14    future::Future,
15    panic::AssertUnwindSafe,
16    pin::Pin,
17    sync::Arc,
18    task::{Context, Poll},
19    time::{Duration, Instant},
20};
21use tonic::{body::BoxBody, Code, Status};
22use tower::{Layer, Service};
23
24type Validator<T> = dyn Fn(&str) -> Option<T> + Send + Sync;
25
26#[derive(Clone)]
27struct RpcSlowLogConfig {
28    logger: Logger,
29    threshold: Duration,
30}
31
32/// Creates the common server layer once and applies it to every generated Tonic service.
33///
34/// The layer installs authentication, W3C trace extraction, panic recovery, adaptive admission
35/// control, and complete unary/streaming metrics. Add it to [`tonic::transport::Server`] before
36/// adding generated services; those services need no per-service interceptors.
37pub struct RpcServerStackBuilder<T = ()> {
38    metrics: RpcMetrics,
39    auth: Option<Arc<Validator<T>>>,
40    shedder: Option<AdaptiveShedder>,
41    slow_log: Option<RpcSlowLogConfig>,
42}
43
44impl RpcServerStackBuilder<()> {
45    pub fn new(metrics: RpcMetrics) -> Self {
46        Self {
47            metrics,
48            auth: None,
49            shedder: None,
50            slow_log: None,
51        }
52    }
53}
54
55impl<T> RpcServerStackBuilder<T>
56where
57    T: Clone + Send + Sync + 'static,
58{
59    pub fn with_bearer_auth<U>(
60        self,
61        validator: impl Fn(&str) -> Option<U> + Send + Sync + 'static,
62    ) -> RpcServerStackBuilder<U>
63    where
64        U: Clone + Send + Sync + 'static,
65    {
66        RpcServerStackBuilder {
67            metrics: self.metrics,
68            auth: Some(Arc::new(validator)),
69            shedder: self.shedder,
70            slow_log: self.slow_log,
71        }
72    }
73
74    pub fn with_load_shedder(mut self, config: LoadShedderConfig) -> Self {
75        self.shedder = Some(AdaptiveShedder::new(config));
76        self
77    }
78
79    /// Logs every completed gRPC call and classifies calls at or above `threshold` as slow.
80    /// Streaming calls are finalized from their trailers so the record contains the final status.
81    pub fn with_slow_call_logging(mut self, logger: Logger, threshold: Duration) -> Self {
82        assert!(
83            !threshold.is_zero(),
84            "gRPC slow-call threshold must be positive"
85        );
86        self.slow_log = Some(RpcSlowLogConfig { logger, threshold });
87        self
88    }
89
90    pub fn build(self) -> RpcServerStack<T> {
91        RpcServerStack {
92            metrics: RpcMetricsLayer::new(self.metrics),
93            auth: self.auth,
94            shedder: self.shedder,
95            slow_log: self.slow_log,
96        }
97    }
98}
99
100pub struct RpcServerStack<T> {
101    metrics: RpcMetricsLayer,
102    auth: Option<Arc<Validator<T>>>,
103    shedder: Option<AdaptiveShedder>,
104    slow_log: Option<RpcSlowLogConfig>,
105}
106
107impl<T> Clone for RpcServerStack<T> {
108    fn clone(&self) -> Self {
109        Self {
110            metrics: self.metrics.clone(),
111            auth: self.auth.clone(),
112            shedder: self.shedder.clone(),
113            slow_log: self.slow_log.clone(),
114        }
115    }
116}
117
118impl<S, T> Layer<S> for RpcServerStack<T>
119where
120    T: Clone + Send + Sync + 'static,
121{
122    type Service = RpcServerStackService<crate::metrics::RpcMetricsService<S>, T>;
123
124    fn layer(&self, inner: S) -> Self::Service {
125        RpcServerStackService {
126            inner: self.metrics.layer(inner),
127            auth: self.auth.clone(),
128            shedder: self.shedder.clone(),
129            slow_log: self.slow_log.clone(),
130        }
131    }
132}
133
134#[derive(Clone)]
135pub struct RpcServerStackService<S, T> {
136    inner: S,
137    auth: Option<Arc<Validator<T>>>,
138    shedder: Option<AdaptiveShedder>,
139    slow_log: Option<RpcSlowLogConfig>,
140}
141
142impl<S, T, RequestBody, ResponseBody> Service<Request<RequestBody>> for RpcServerStackService<S, T>
143where
144    S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
145    S::Future: Send + 'static,
146    S::Error: Send + 'static,
147    RequestBody: Send + 'static,
148    ResponseBody: Body<Data = bytes::Bytes> + Send + 'static,
149    ResponseBody::Error: Into<tonic::codegen::StdError>,
150    T: Clone + Send + Sync + 'static,
151{
152    type Response = Response<BoxBody>;
153    type Error = S::Error;
154    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
155
156    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
157        self.inner.poll_ready(context)
158    }
159
160    fn call(&mut self, mut request: Request<RequestBody>) -> Self::Future {
161        let trace = request
162            .headers()
163            .get("traceparent")
164            .and_then(|value| value.to_str().ok())
165            .and_then(|value| TraceContext::parse(value).ok())
166            .map(|parent| parent.child())
167            .unwrap_or_else(|| TraceContext::root(TraceFlags::SAMPLED));
168        let mut slow_call = self.slow_log.as_ref().map(|config| {
169            RpcSlowCall::new(
170                config.clone(),
171                request.uri().path().to_owned(),
172                request
173                    .headers()
174                    .get("grpc-timeout")
175                    .and_then(|value| value.to_str().ok())
176                    .and_then(parse_grpc_timeout),
177                trace.clone(),
178            )
179        });
180        request.extensions_mut().insert(trace);
181
182        if let Some(auth) = &self.auth {
183            let identity = request
184                .headers()
185                .get("authorization")
186                .and_then(|value| value.to_str().ok())
187                .and_then(bearer_token)
188                .and_then(|token| auth(token));
189            let Some(identity) = identity else {
190                if let Some(call) = slow_call.take() {
191                    call.finish(Code::Unauthenticated);
192                }
193                return Box::pin(async {
194                    Ok(Status::unauthenticated(format!(
195                        "{}: {}",
196                        AuthFailure::InvalidCredentials.code(),
197                        AuthFailure::InvalidCredentials.message()
198                    ))
199                    .into_http())
200                });
201            };
202            request.extensions_mut().insert(identity);
203        }
204
205        let permit = match &self.shedder {
206            Some(shedder) => match shedder.try_acquire() {
207                Some(permit) => Some(permit),
208                None => {
209                    if let Some(call) = slow_call.take() {
210                        call.finish(Code::ResourceExhausted);
211                    }
212                    return Box::pin(async {
213                        Ok(Status::resource_exhausted("gRPC server is overloaded").into_http())
214                    });
215                }
216            },
217            None => None,
218        };
219
220        let future = self.inner.call(request);
221        Box::pin(async move {
222            let result = AssertUnwindSafe(future).catch_unwind().await;
223            match result {
224                Ok(Ok(response)) => {
225                    let header_code = grpc_status(response.headers());
226                    let (parts, body) = response.into_parts();
227                    let mut body = RpcServerLogBody {
228                        inner: body,
229                        slow_call,
230                        _shed_permit: permit,
231                    };
232                    if let Some(code) = header_code {
233                        body.finish(code);
234                    } else if !parts.status.is_success() {
235                        body.finish_transport_error();
236                    }
237                    Ok(Response::from_parts(parts, tonic::body::boxed(body)))
238                }
239                Ok(Err(error)) => {
240                    if let Some(call) = slow_call.take() {
241                        call.finish_transport_error();
242                    }
243                    Err(error)
244                }
245                Err(_) => {
246                    if let Some(call) = slow_call.take() {
247                        call.finish(Code::Internal);
248                    }
249                    Ok(Status::internal("gRPC handler panicked").into_http())
250                }
251            }
252        })
253    }
254}
255
256struct RpcSlowCall {
257    config: RpcSlowLogConfig,
258    method: String,
259    deadline: Option<Duration>,
260    trace: TraceContext,
261    started: Instant,
262}
263
264impl RpcSlowCall {
265    fn new(
266        config: RpcSlowLogConfig,
267        method: String,
268        deadline: Option<Duration>,
269        trace: TraceContext,
270    ) -> Self {
271        Self {
272            config,
273            method,
274            deadline,
275            trace,
276            started: Instant::now(),
277        }
278    }
279
280    fn finish(self, code: Code) {
281        self.emit(code_name(code), code == Code::DeadlineExceeded);
282    }
283
284    fn finish_transport_error(self) {
285        self.emit("transport_error", false);
286    }
287
288    fn emit(self, status: &'static str, status_deadline_exceeded: bool) {
289        let elapsed = self.started.elapsed();
290        let slow = elapsed >= self.config.threshold;
291        let deadline_exceeded =
292            status_deadline_exceeded || self.deadline.is_some_and(|deadline| elapsed >= deadline);
293        let mut fields = vec![
294            LogField::new("transport", "grpc"),
295            LogField::new("rpc_role", "server"),
296            LogField::new("method", self.method),
297            LogField::new("status", status),
298            LogField::new("elapsed_ms", duration_millis(elapsed)),
299            LogField::new("slow_threshold_ms", duration_millis(self.config.threshold)),
300            LogField::new("slow", slow),
301            LogField::new("deadline_exceeded", deadline_exceeded),
302        ];
303        if let Some(deadline) = self.deadline {
304            fields.push(LogField::new("deadline_ms", duration_millis(deadline)));
305        }
306        let context = LogContext::new().with_trace(self.trace);
307        let _ = self.config.logger.log_with_context(
308            if slow { LogLevel::Slow } else { LogLevel::Info },
309            "grpc request completed",
310            Some(&context),
311            fields,
312        );
313    }
314}
315
316pin_project! {
317    struct RpcServerLogBody<B> {
318        #[pin]
319        inner: B,
320        slow_call: Option<RpcSlowCall>,
321        _shed_permit: Option<ShedPermit>,
322    }
323
324    impl<B> PinnedDrop for RpcServerLogBody<B> {
325        fn drop(this: Pin<&mut Self>) {
326            let this = this.project();
327            if let Some(call) = this.slow_call.take() {
328                call.finish(Code::Cancelled);
329            }
330        }
331    }
332}
333
334impl<B> RpcServerLogBody<B> {
335    fn finish(&mut self, code: Code) {
336        if let Some(call) = self.slow_call.take() {
337            call.finish(code);
338        }
339    }
340
341    fn finish_transport_error(&mut self) {
342        if let Some(call) = self.slow_call.take() {
343            call.finish_transport_error();
344        }
345    }
346}
347
348impl<B> Body for RpcServerLogBody<B>
349where
350    B: Body,
351{
352    type Data = B::Data;
353    type Error = B::Error;
354
355    fn poll_frame(
356        self: Pin<&mut Self>,
357        context: &mut Context<'_>,
358    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
359        let mut this = self.project();
360        match this.inner.as_mut().poll_frame(context) {
361            Poll::Ready(Some(Ok(frame))) => {
362                if let Some(code) = frame.trailers_ref().and_then(grpc_status) {
363                    if let Some(call) = this.slow_call.take() {
364                        call.finish(code);
365                    }
366                }
367                Poll::Ready(Some(Ok(frame)))
368            }
369            Poll::Ready(Some(Err(error))) => {
370                if let Some(call) = this.slow_call.take() {
371                    call.finish_transport_error();
372                }
373                Poll::Ready(Some(Err(error)))
374            }
375            Poll::Ready(None) => {
376                if let Some(call) = this.slow_call.take() {
377                    call.finish(Code::Ok);
378                }
379                Poll::Ready(None)
380            }
381            Poll::Pending => Poll::Pending,
382        }
383    }
384
385    fn is_end_stream(&self) -> bool {
386        self.inner.is_end_stream()
387    }
388
389    fn size_hint(&self) -> http_body::SizeHint {
390        self.inner.size_hint()
391    }
392}
393
394fn bearer_token(value: &str) -> Option<&str> {
395    let (scheme, token) = value.split_once(char::is_whitespace)?;
396    let token = token.trim();
397    (scheme.eq_ignore_ascii_case("bearer")
398        && !token.is_empty()
399        && !token.contains(char::is_whitespace))
400    .then_some(token)
401}
402
403/// Builds one transport service that can be passed directly to any generated Tonic client.
404///
405/// The stack installs bearer credentials, W3C trace propagation, bounded metrics, a default
406/// deadline, and protocol-aware circuit breaking. Final streaming statuses are observed from
407/// trailers before circuit health is recorded.
408pub struct RpcClientStackBuilder {
409    metrics: RpcMetrics,
410    token: Option<BearerToken>,
411    default_timeout: Option<Duration>,
412    breaker: Option<Arc<CircuitBreaker>>,
413}
414
415impl RpcClientStackBuilder {
416    pub fn new(metrics: RpcMetrics) -> Self {
417        Self {
418            metrics,
419            token: None,
420            default_timeout: None,
421            breaker: None,
422        }
423    }
424
425    pub fn with_bearer_token(mut self, token: BearerToken) -> Self {
426        self.token = Some(token);
427        self
428    }
429
430    pub fn with_default_timeout(mut self, timeout: Duration) -> Self {
431        assert!(!timeout.is_zero(), "RPC timeout must be greater than zero");
432        self.default_timeout = Some(timeout);
433        self
434    }
435
436    pub fn with_circuit_breaker(mut self, config: CircuitBreakerConfig) -> Self {
437        self.breaker = Some(Arc::new(CircuitBreaker::new(config)));
438        self
439    }
440
441    pub fn build(self) -> RpcClientStack {
442        RpcClientStack {
443            metrics: RpcMetricsLayer::new(self.metrics),
444            token: self.token,
445            default_timeout: self.default_timeout,
446            breaker: self.breaker,
447        }
448    }
449}
450
451#[derive(Clone)]
452pub struct RpcClientStack {
453    metrics: RpcMetricsLayer,
454    token: Option<BearerToken>,
455    default_timeout: Option<Duration>,
456    breaker: Option<Arc<CircuitBreaker>>,
457}
458
459impl<S> Layer<S> for RpcClientStack {
460    type Service = RpcClientStackService<crate::metrics::RpcMetricsService<S>>;
461
462    fn layer(&self, inner: S) -> Self::Service {
463        RpcClientStackService {
464            inner: self.metrics.layer(inner),
465            token: self.token.clone(),
466            default_timeout: self.default_timeout,
467            breaker: self.breaker.clone(),
468        }
469    }
470}
471
472#[derive(Clone)]
473pub struct RpcClientStackService<S> {
474    inner: S,
475    token: Option<BearerToken>,
476    default_timeout: Option<Duration>,
477    breaker: Option<Arc<CircuitBreaker>>,
478}
479
480impl<S, RequestBody, ResponseBody> Service<Request<RequestBody>> for RpcClientStackService<S>
481where
482    S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
483    S::Future: Send + 'static,
484    S::Error: Send + 'static,
485    RequestBody: Send + 'static,
486    ResponseBody: Body<Data = bytes::Bytes> + Send + 'static,
487    ResponseBody::Error: Into<tonic::codegen::StdError>,
488{
489    type Response = Response<BoxBody>;
490    type Error = S::Error;
491    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
492
493    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
494        self.inner.poll_ready(context)
495    }
496
497    fn call(&mut self, mut request: Request<RequestBody>) -> Self::Future {
498        if let Some(token) = &self.token {
499            request.headers_mut().insert(
500                "authorization",
501                http::HeaderValue::from_bytes(token.authorization().as_encoded_bytes())
502                    .expect("ASCII gRPC metadata is a valid HTTP header"),
503            );
504        }
505
506        let parent = request
507            .extensions()
508            .get::<TraceContext>()
509            .cloned()
510            .or_else(|| {
511                request
512                    .headers()
513                    .get("traceparent")
514                    .and_then(|value| value.to_str().ok())
515                    .and_then(|value| TraceContext::parse(value).ok())
516            });
517        let trace = parent
518            .as_ref()
519            .map(TraceContext::child)
520            .unwrap_or_else(|| TraceContext::root(TraceFlags::SAMPLED));
521        request.headers_mut().insert(
522            "traceparent",
523            trace
524                .traceparent()
525                .parse()
526                .expect("generated traceparent is a valid HTTP header"),
527        );
528        request.extensions_mut().insert(trace);
529
530        if !request.headers().contains_key("grpc-timeout") {
531            if let Some(timeout) = self.default_timeout {
532                request.headers_mut().insert(
533                    "grpc-timeout",
534                    grpc_timeout(timeout)
535                        .parse()
536                        .expect("formatted gRPC timeout is a valid HTTP header"),
537                );
538            }
539        }
540
541        let permit = match &self.breaker {
542            Some(breaker) => match breaker.acquire() {
543                Some(permit) => Some(permit),
544                None => {
545                    return Box::pin(async {
546                        Ok(Status::unavailable("gRPC dependency circuit is open").into_http())
547                    });
548                }
549            },
550            None => None,
551        };
552
553        let future = self.inner.call(request);
554        Box::pin(async move {
555            match future.await {
556                Ok(response) => {
557                    let header_code = grpc_status(response.headers());
558                    let (parts, body) = response.into_parts();
559                    let mut wrapped = RpcCircuitBody {
560                        inner: body,
561                        permit,
562                    };
563                    if let Some(code) = header_code {
564                        wrapped.finish(status_outcome(code));
565                    } else if !parts.status.is_success() {
566                        wrapped.finish(CircuitOutcome::Failure);
567                    }
568                    Ok(Response::from_parts(parts, tonic::body::boxed(wrapped)))
569                }
570                Err(error) => {
571                    if let Some(permit) = permit {
572                        permit.finish(false);
573                    }
574                    Err(error)
575                }
576            }
577        })
578    }
579}
580
581pin_project! {
582    struct RpcCircuitBody<B> {
583        #[pin]
584        inner: B,
585        permit: Option<CircuitBreakerPermit>,
586    }
587}
588
589impl<B> RpcCircuitBody<B> {
590    fn finish(&mut self, outcome: CircuitOutcome) {
591        if let Some(permit) = self.permit.take() {
592            permit.finish_with_outcome(outcome);
593        }
594    }
595}
596
597impl<B> Body for RpcCircuitBody<B>
598where
599    B: Body,
600{
601    type Data = B::Data;
602    type Error = B::Error;
603
604    fn poll_frame(
605        self: Pin<&mut Self>,
606        context: &mut Context<'_>,
607    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
608        let mut this = self.project();
609        match this.inner.as_mut().poll_frame(context) {
610            Poll::Ready(Some(Ok(frame))) => {
611                if let Some(code) = frame.trailers_ref().and_then(grpc_status) {
612                    if let Some(permit) = this.permit.take() {
613                        permit.finish_with_outcome(status_outcome(code));
614                    }
615                }
616                Poll::Ready(Some(Ok(frame)))
617            }
618            Poll::Ready(Some(Err(error))) => {
619                if let Some(permit) = this.permit.take() {
620                    permit.finish(false);
621                }
622                Poll::Ready(Some(Err(error)))
623            }
624            Poll::Ready(None) => {
625                if let Some(permit) = this.permit.take() {
626                    permit.finish(true);
627                }
628                Poll::Ready(None)
629            }
630            Poll::Pending => Poll::Pending,
631        }
632    }
633
634    fn is_end_stream(&self) -> bool {
635        self.inner.is_end_stream()
636    }
637
638    fn size_hint(&self) -> http_body::SizeHint {
639        self.inner.size_hint()
640    }
641}
642
643fn grpc_status(headers: &http::HeaderMap) -> Option<Code> {
644    headers
645        .get("grpc-status")
646        .and_then(|value| value.to_str().ok())
647        .and_then(|value| value.parse::<i32>().ok())
648        .map(Code::from_i32)
649}
650
651fn parse_grpc_timeout(value: &str) -> Option<Duration> {
652    let (amount, unit) = value.split_at(value.len().checked_sub(1)?);
653    if amount.is_empty() || amount.len() > 8 {
654        return None;
655    }
656    let amount = amount.parse::<u64>().ok()?;
657    match unit {
658        "n" => Some(Duration::from_nanos(amount)),
659        "u" => Some(Duration::from_micros(amount)),
660        "m" => Some(Duration::from_millis(amount)),
661        "S" => Some(Duration::from_secs(amount)),
662        "M" => Some(Duration::from_secs(amount.saturating_mul(60))),
663        "H" => Some(Duration::from_secs(amount.saturating_mul(3_600))),
664        _ => None,
665    }
666}
667
668fn duration_millis(duration: Duration) -> u64 {
669    duration.as_millis().min(u128::from(u64::MAX)) as u64
670}
671
672fn code_name(code: Code) -> &'static str {
673    match code {
674        Code::Ok => "ok",
675        Code::Cancelled => "cancelled",
676        Code::Unknown => "unknown",
677        Code::InvalidArgument => "invalid_argument",
678        Code::DeadlineExceeded => "deadline_exceeded",
679        Code::NotFound => "not_found",
680        Code::AlreadyExists => "already_exists",
681        Code::PermissionDenied => "permission_denied",
682        Code::ResourceExhausted => "resource_exhausted",
683        Code::FailedPrecondition => "failed_precondition",
684        Code::Aborted => "aborted",
685        Code::OutOfRange => "out_of_range",
686        Code::Unimplemented => "unimplemented",
687        Code::Internal => "internal",
688        Code::Unavailable => "unavailable",
689        Code::DataLoss => "data_loss",
690        Code::Unauthenticated => "unauthenticated",
691    }
692}
693
694fn status_outcome(code: Code) -> CircuitOutcome {
695    if code == Code::Cancelled {
696        return CircuitOutcome::Cancellation;
697    }
698    if matches!(
699        code,
700        Code::DeadlineExceeded
701            | Code::Internal
702            | Code::Unavailable
703            | Code::DataLoss
704            | Code::Unimplemented
705            | Code::ResourceExhausted
706    ) {
707        CircuitOutcome::Failure
708    } else {
709        CircuitOutcome::Success
710    }
711}
712
713fn grpc_timeout(timeout: Duration) -> String {
714    let nanos = timeout.as_nanos().max(1);
715    for (unit, divisor) in [
716        ('n', 1_u128),
717        ('u', 1_000),
718        ('m', 1_000_000),
719        ('S', 1_000_000_000),
720        ('M', 60_000_000_000),
721        ('H', 3_600_000_000_000),
722    ] {
723        let value = nanos.saturating_add(divisor - 1) / divisor;
724        if value <= 99_999_999 {
725            return format!("{value}{unit}");
726        }
727    }
728    "99999999H".to_owned()
729}
730
731#[cfg(test)]
732mod client_tests {
733    use super::*;
734    use crate::RpcMetricMode;
735    use std::{
736        convert::Infallible,
737        io::{self, Write},
738        sync::{
739            atomic::{AtomicUsize, Ordering},
740            Mutex,
741        },
742    };
743    use tower::ServiceExt;
744
745    #[derive(Clone)]
746    struct FailingTransport {
747        calls: Arc<AtomicUsize>,
748    }
749
750    impl Service<Request<()>> for FailingTransport {
751        type Response = Response<TrailerBody>;
752        type Error = Infallible;
753        type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
754
755        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
756            Poll::Ready(Ok(()))
757        }
758
759        fn call(&mut self, request: Request<()>) -> Self::Future {
760            self.calls.fetch_add(1, Ordering::Relaxed);
761            assert_eq!(
762                request
763                    .headers()
764                    .get("authorization")
765                    .and_then(|value| value.to_str().ok()),
766                Some("Bearer secret")
767            );
768            assert_eq!(
769                request
770                    .headers()
771                    .get("grpc-timeout")
772                    .and_then(|value| value.to_str().ok()),
773                Some("250000u")
774            );
775            assert!(request.headers().contains_key("traceparent"));
776            std::future::ready(Ok(Response::new(TrailerBody(true))))
777        }
778    }
779
780    struct TrailerBody(bool);
781
782    impl Body for TrailerBody {
783        type Data = bytes::Bytes;
784        type Error = Status;
785
786        fn poll_frame(
787            mut self: Pin<&mut Self>,
788            _: &mut Context<'_>,
789        ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
790            if !self.0 {
791                return Poll::Ready(None);
792            }
793            self.0 = false;
794            let mut trailers = http::HeaderMap::new();
795            trailers.insert("grpc-status", "14".parse().unwrap());
796            Poll::Ready(Some(Ok(Frame::trailers(trailers))))
797        }
798    }
799
800    #[derive(Clone, Default)]
801    struct SharedWriter(Arc<Mutex<Vec<u8>>>);
802
803    impl Write for SharedWriter {
804        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
805            self.0.lock().unwrap().extend_from_slice(bytes);
806            Ok(bytes.len())
807        }
808
809        fn flush(&mut self) -> io::Result<()> {
810            Ok(())
811        }
812    }
813
814    #[derive(Clone)]
815    struct SlowServerTransport;
816
817    impl Service<Request<()>> for SlowServerTransport {
818        type Response = Response<TrailerBody>;
819        type Error = Infallible;
820        type Future =
821            Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
822
823        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
824            Poll::Ready(Ok(()))
825        }
826
827        fn call(&mut self, _: Request<()>) -> Self::Future {
828            Box::pin(async {
829                tokio::time::sleep(Duration::from_millis(5)).await;
830                Ok(Response::new(TrailerBody(true)))
831            })
832        }
833    }
834
835    #[tokio::test]
836    async fn client_stack_composes_headers_metrics_deadline_and_trailer_breaking() {
837        let registry = rust_zero_core::Metrics::new();
838        let metrics = RpcMetrics::new(
839            &registry,
840            "stacked",
841            RpcMetricMode::Client,
842            ["/echo.Echo/Call"],
843        )
844        .unwrap();
845        let calls = Arc::new(AtomicUsize::new(0));
846        let mut service = RpcClientStackBuilder::new(metrics)
847            .with_bearer_token(BearerToken::new("secret").unwrap())
848            .with_default_timeout(Duration::from_millis(250))
849            .with_circuit_breaker(CircuitBreakerConfig::new(1, Duration::from_secs(60)))
850            .build()
851            .layer(FailingTransport {
852                calls: Arc::clone(&calls),
853            });
854
855        let response = service
856            .ready()
857            .await
858            .unwrap()
859            .call(Request::builder().uri("/echo.Echo/Call").body(()).unwrap())
860            .await
861            .unwrap();
862        let mut body = Box::pin(response.into_body());
863        std::future::poll_fn(|context| body.as_mut().poll_frame(context)).await;
864
865        let rejected = service
866            .ready()
867            .await
868            .unwrap()
869            .call(Request::builder().uri("/echo.Echo/Call").body(()).unwrap())
870            .await
871            .unwrap();
872        assert_eq!(rejected.headers().get("grpc-status").unwrap(), "14");
873        assert_eq!(calls.load(Ordering::Relaxed), 1);
874        assert!(registry.render().contains(
875            "stacked_rpc_client_requests_total{method=\"/echo.Echo/Call\",code=\"14\"} 1"
876        ));
877    }
878
879    #[tokio::test]
880    async fn server_stack_logs_final_status_deadline_and_slow_classification() {
881        let registry = rust_zero_core::Metrics::new();
882        let metrics = RpcMetrics::new(
883            &registry,
884            "server_log",
885            RpcMetricMode::Server,
886            ["/echo.Echo/Call"],
887        )
888        .unwrap();
889        let output = SharedWriter::default();
890        let logger =
891            Logger::to_writer(rust_zero_core::LogConfig::console("rpc"), output.clone()).unwrap();
892        let mut service = RpcServerStackBuilder::new(metrics)
893            .with_slow_call_logging(logger, Duration::from_millis(1))
894            .build()
895            .layer(SlowServerTransport);
896
897        let response = service
898            .ready()
899            .await
900            .unwrap()
901            .call(
902                Request::builder()
903                    .uri("/echo.Echo/Call")
904                    .header("grpc-timeout", "1m")
905                    .body(())
906                    .unwrap(),
907            )
908            .await
909            .unwrap();
910        let mut body = Box::pin(response.into_body());
911        std::future::poll_fn(|context| body.as_mut().poll_frame(context)).await;
912
913        let bytes = output.0.lock().unwrap().clone();
914        let record: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
915        assert_eq!(record["level"], "slow");
916        assert_eq!(record["transport"], "grpc");
917        assert_eq!(record["rpc_role"], "server");
918        assert_eq!(record["method"], "/echo.Echo/Call");
919        assert_eq!(record["status"], "unavailable");
920        assert_eq!(record["slow"], true);
921        assert_eq!(record["slow_threshold_ms"], 1);
922        assert_eq!(record["deadline_ms"], 1);
923        assert_eq!(record["deadline_exceeded"], true);
924        assert!(record["trace_id"].as_str().is_some());
925    }
926
927    #[tokio::test]
928    async fn server_shedder_holds_permit_until_stream_completion() {
929        let registry = rust_zero_core::Metrics::new();
930        let metrics = RpcMetrics::new(
931            &registry,
932            "server_shed",
933            RpcMetricMode::Server,
934            ["/echo.Echo/Call"],
935        )
936        .unwrap();
937        let mut service = RpcServerStackBuilder::new(metrics)
938            .with_load_shedder(LoadShedderConfig::production(1))
939            .build()
940            .layer(SlowServerTransport);
941
942        let first = service
943            .ready()
944            .await
945            .unwrap()
946            .call(Request::builder().uri("/echo.Echo/Call").body(()).unwrap())
947            .await
948            .unwrap();
949        let rejected = service
950            .ready()
951            .await
952            .unwrap()
953            .call(Request::builder().uri("/echo.Echo/Call").body(()).unwrap())
954            .await
955            .unwrap();
956        assert_eq!(rejected.headers().get("grpc-status").unwrap(), "8");
957
958        let mut body = Box::pin(first.into_body());
959        std::future::poll_fn(|context| body.as_mut().poll_frame(context)).await;
960        drop(body);
961
962        let admitted = service
963            .ready()
964            .await
965            .unwrap()
966            .call(Request::builder().uri("/echo.Echo/Call").body(()).unwrap())
967            .await
968            .unwrap();
969        assert_ne!(
970            admitted
971                .headers()
972                .get("grpc-status")
973                .map(|value| value.as_bytes()),
974            Some(b"8".as_slice())
975        );
976    }
977
978    #[test]
979    fn grpc_deadlines_use_the_smallest_exact_unit() {
980        assert_eq!(grpc_timeout(Duration::from_nanos(7)), "7n");
981        assert_eq!(grpc_timeout(Duration::from_millis(250)), "250000u");
982        assert_eq!(grpc_timeout(Duration::from_secs(100)), "100000m");
983    }
984}
rpc/stack.rs

rpc/
stack.rs
