iota_network_stack/
server.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    convert::Infallible,
7    num::NonZeroUsize,
8    task::{Context, Poll},
9    time::Duration,
10};
11
12use eyre::{Result, eyre};
13use iota_http::metrics::ListenerMetrics;
14use tokio_rustls::rustls::ServerConfig;
15use tonic::{
16    body::Body,
17    codegen::http::{HeaderValue, Request, Response},
18    server::NamedService,
19};
20use tower::{Layer, Service, ServiceBuilder};
21use tower_http::{
22    propagate_header::PropagateHeaderLayer, set_header::SetRequestHeaderLayer, trace::TraceLayer,
23};
24
25use crate::{
26    concurrency::ServiceConcurrencyLimit,
27    config::Config,
28    metrics::{
29        DefaultMetricsCallbackProvider, GRPC_ENDPOINT_PATH_HEADER, MetricsCallbackProvider,
30        MetricsHandler,
31    },
32    multiaddr::{Multiaddr, Protocol},
33};
34
35/// Server-side deadline for a request whose config sets no `request_timeout`.
36///
37/// A request holds its admission slot from the moment its headers arrive
38/// until the handler responds, including the time spent receiving the body,
39/// so without a deadline a peer that never finishes sending pins the slot
40/// indefinitely.
41pub const DEFAULT_GRPC_REQUEST_TIMEOUT: Duration = Duration::from_secs(300);
42
43pub struct ServerBuilder<M: MetricsCallbackProvider = DefaultMetricsCallbackProvider> {
44    config: Config,
45    metrics_provider: M,
46    router: tonic::service::Routes,
47    health_reporter: tonic_health::server::HealthReporter,
48    listener_metrics: Option<ListenerMetrics>,
49}
50
51impl<M: MetricsCallbackProvider> ServerBuilder<M> {
52    pub fn from_config(config: &Config, metrics_provider: M) -> Self {
53        let (health_reporter, health_service) = tonic_health::server::health_reporter();
54        let router = tonic::service::Routes::new(health_service);
55
56        Self {
57            config: config.to_owned(),
58            metrics_provider,
59            router,
60            health_reporter,
61            listener_metrics: None,
62        }
63    }
64
65    /// Measures the connections of the listener. See
66    /// [`iota_http::Config::metrics`].
67    pub fn listener_metrics(mut self, metrics: ListenerMetrics) -> Self {
68        self.listener_metrics = Some(metrics);
69        self
70    }
71
72    pub fn health_reporter(&self) -> tonic_health::server::HealthReporter {
73        self.health_reporter.clone()
74    }
75
76    /// The server-side request deadline this server enforces, unless a
77    /// request carries a shorter `grpc-timeout` header.
78    pub fn request_timeout(&self) -> Duration {
79        self.config
80            .request_timeout
81            .unwrap_or(DEFAULT_GRPC_REQUEST_TIMEOUT)
82    }
83
84    /// Add a new service to this Server.
85    pub fn add_service<S>(mut self, svc: S) -> Self
86    where
87        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
88            + NamedService
89            + Clone
90            + Send
91            + Sync
92            + 'static,
93        S::Future: Send + 'static,
94    {
95        self.router = self.router.add_service(svc);
96        self
97    }
98
99    /// Add a new service to this Server with its own concurrency limit,
100    /// enforced independently of every other service on this server.
101    ///
102    /// With `load_shed` enabled, requests over the limit are rejected
103    /// immediately with gRPC `RESOURCE_EXHAUSTED`; otherwise they wait for a
104    /// slot to free up.
105    pub fn add_service_with_concurrency_limit<S>(
106        mut self,
107        svc: S,
108        limit: NonZeroUsize,
109        load_shed: bool,
110    ) -> Self
111    where
112        S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
113            + NamedService
114            + Clone
115            + Send
116            + Sync
117            + 'static,
118        S::Future: Send + 'static,
119    {
120        self.router = self
121            .router
122            .add_service(ServiceConcurrencyLimit::new(svc, limit, load_shed));
123        self
124    }
125
126    pub async fn bind(self, addr: &Multiaddr, tls_config: Option<ServerConfig>) -> Result<Server> {
127        let request_timeout = self.request_timeout();
128        let http_config = self.config.http_config().metrics(self.listener_metrics);
129
130        let metrics_provider = self.metrics_provider;
131        let metrics = MetricsHandler::new(metrics_provider.clone());
132        let request_metrics = TraceLayer::new_for_grpc()
133            .on_request(metrics.clone())
134            .on_response(metrics.clone())
135            .on_failure(metrics);
136
137        fn add_path_to_request_header<T>(request: &Request<T>) -> Option<HeaderValue> {
138            let path = request.uri().path();
139            Some(HeaderValue::from_str(path).unwrap())
140        }
141
142        let limiting_layers = ServiceBuilder::new()
143            .option_layer(
144                self.config
145                    .load_shed
146                    .unwrap_or_default()
147                    .then_some(tower::load_shed::LoadShedLayer::new()),
148            )
149            .option_layer(
150                self.config
151                    .global_concurrency_limit
152                    .map(tower::limit::GlobalConcurrencyLimitLayer::new),
153            );
154
155        let route_layers = ServiceBuilder::new()
156            .map_request(|mut request: http::Request<_>| {
157                if let Some(connect_info) = request.extensions().get::<iota_http::ConnectInfo>() {
158                    let tonic_connect_info = tonic::transport::server::TcpConnectInfo {
159                        local_addr: Some(connect_info.local_addr),
160                        remote_addr: Some(connect_info.remote_addr),
161                    };
162                    request.extensions_mut().insert(tonic_connect_info);
163                }
164                request
165            })
166            .layer(RequestLifetimeLayer { metrics_provider })
167            .layer(SetRequestHeaderLayer::overriding(
168                GRPC_ENDPOINT_PATH_HEADER.clone(),
169                add_path_to_request_header,
170            ))
171            .layer(request_metrics)
172            .layer(PropagateHeaderLayer::new(GRPC_ENDPOINT_PATH_HEADER.clone()))
173            .layer_fn(move |service| {
174                crate::grpc_timeout::GrpcTimeout::new(service, Some(request_timeout))
175            });
176
177        let mut builder = iota_http::Builder::new().config(http_config);
178
179        let has_tls = tls_config.is_some();
180        if let Some(tls_config) = tls_config {
181            builder = builder.tls_config(tls_config);
182        }
183
184        let server_handle = builder
185            .serve(
186                addr,
187                limiting_layers.service(self.router.into_axum_router().layer(route_layers)),
188            )
189            .map_err(|e| eyre!(e))?;
190
191        let mut local_addr = update_tcp_port_in_multiaddr(addr, server_handle.local_addr().port());
192        if has_tls {
193            local_addr = local_addr.rewrite_http_to_https();
194        }
195        Ok(Server {
196            server_handle,
197            local_addr,
198            health_reporter: self.health_reporter,
199        })
200    }
201}
202
203/// TLS server name to use for the public IOTA validator interface.
204pub const IOTA_TLS_SERVER_NAME: &str = "iota";
205
206pub struct Server {
207    server_handle: iota_http::ServerHandle,
208    local_addr: Multiaddr,
209    health_reporter: tonic_health::server::HealthReporter,
210}
211
212impl Server {
213    pub async fn serve(self) -> Result<(), tonic::transport::Error> {
214        self.server_handle.wait_for_shutdown().await;
215        Ok(())
216    }
217
218    pub fn trigger_shutdown(&self) {
219        self.server_handle.trigger_shutdown();
220    }
221
222    pub fn local_addr(&self) -> &Multiaddr {
223        &self.local_addr
224    }
225
226    pub fn health_reporter(&self) -> tonic_health::server::HealthReporter {
227        self.health_reporter.clone()
228    }
229
230    pub fn handle(&self) -> &iota_http::ServerHandle {
231        &self.server_handle
232    }
233}
234
235fn update_tcp_port_in_multiaddr(addr: &Multiaddr, port: u16) -> Multiaddr {
236    addr.replace(1, |protocol| {
237        if let Protocol::Tcp(_) = protocol {
238            Some(Protocol::Tcp(port))
239        } else {
240            panic!("expected tcp protocol at index 1");
241        }
242    })
243    .expect("tcp protocol at index 1")
244}
245
246#[cfg(test)]
247mod test {
248    use std::{
249        ops::Deref,
250        sync::{Arc, Mutex},
251        time::Duration,
252    };
253
254    use fastcrypto::{ed25519::Ed25519KeyPair, traits::KeyPair};
255    use iota_metrics::test_utils::MetricsReader;
256    use tonic::Code;
257    use tonic_health::pb::{HealthCheckRequest, health_client::HealthClient};
258
259    use crate::{Multiaddr, config::Config, metrics::MetricsCallbackProvider};
260
261    #[tokio::test]
262    async fn test_metrics_layer_successful() {
263        #[derive(Clone)]
264        struct Metrics {
265            /// a flag to figure out whether the
266            /// on_request method has been called.
267            metrics_called: Arc<Mutex<bool>>,
268        }
269
270        impl MetricsCallbackProvider for Metrics {
271            fn on_request(&self, path: String) {
272                assert_eq!(path, "/grpc.health.v1.Health/Check");
273            }
274
275            fn on_response(
276                &self,
277                path: String,
278                _latency: Duration,
279                status: u16,
280                grpc_status_code: Code,
281            ) {
282                assert_eq!(path, "/grpc.health.v1.Health/Check");
283                assert_eq!(status, 200);
284                assert_eq!(grpc_status_code, Code::Ok);
285                let mut m = self.metrics_called.lock().unwrap();
286                *m = true
287            }
288        }
289
290        let metrics = Metrics {
291            metrics_called: Arc::new(Mutex::new(false)),
292        };
293
294        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
295        let config = Config::new();
296        let keypair = Ed25519KeyPair::generate(&mut rand08::thread_rng());
297
298        let server = config
299            .server_builder_with_metrics(metrics.clone())
300            .bind(
301                &address,
302                Some(iota_tls::create_rustls_server_config(
303                    keypair.copy().private(),
304                    "test".to_string(),
305                )),
306            )
307            .await
308            .unwrap();
309
310        let address = server.local_addr().to_owned();
311        let channel = config
312            .connect(
313                &address,
314                iota_tls::create_rustls_client_config(
315                    keypair.public().to_owned(),
316                    "test".to_string(),
317                    None,
318                ),
319            )
320            .await
321            .unwrap();
322        let mut client = HealthClient::new(channel);
323
324        client
325            .check(HealthCheckRequest {
326                service: "".to_owned(),
327            })
328            .await
329            .unwrap();
330
331        server.server_handle.shutdown().await;
332
333        assert!(metrics.metrics_called.lock().unwrap().deref());
334    }
335
336    #[tokio::test]
337    async fn listener_metrics_count_the_connections_of_the_server() {
338        let registry = prometheus_filtered::Registry::new();
339        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
340        let config = Config::new();
341        let server = config
342            .server_builder()
343            .listener_metrics(iota_http::metrics::ListenerMetrics::new(
344                "test",
345                &registry,
346                prometheus_filtered::MetricLevel::Info,
347            ))
348            .bind(&address, None)
349            .await
350            .unwrap();
351
352        let channel = tonic::transport::Endpoint::from_shared(format!(
353            "http://{}",
354            server.local_addr().to_socket_addr().unwrap()
355        ))
356        .unwrap()
357        .connect()
358        .await
359        .unwrap();
360        HealthClient::new(channel)
361            .check(HealthCheckRequest {
362                service: "".to_owned(),
363            })
364            .await
365            .unwrap();
366
367        let reader = MetricsReader::new(&registry).with_prefix("test");
368        assert_eq!(reader.value("inbound_connections_accepted", &[]), 1.0);
369        assert!(
370            !reader.has_family("pending_handshakes_peak"),
371            "a plain listener has no TLS metrics"
372        );
373        server.trigger_shutdown();
374    }
375
376    #[tokio::test]
377    async fn listener_metrics_time_the_tls_handshakes_of_the_server() {
378        let registry = prometheus_filtered::Registry::new();
379        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
380        let config = Config::new();
381        let keypair = Ed25519KeyPair::generate(&mut rand08::thread_rng());
382        let server = config
383            .server_builder()
384            .listener_metrics(iota_http::metrics::ListenerMetrics::new(
385                "test",
386                &registry,
387                prometheus_filtered::MetricLevel::Info,
388            ))
389            .bind(
390                &address,
391                Some(iota_tls::create_rustls_server_config(
392                    keypair.copy().private(),
393                    "test".to_string(),
394                )),
395            )
396            .await
397            .unwrap();
398
399        let channel = config
400            .connect(
401                server.local_addr(),
402                iota_tls::create_rustls_client_config(
403                    keypair.public().to_owned(),
404                    "test".to_string(),
405                    None,
406                ),
407            )
408            .await
409            .unwrap();
410        HealthClient::new(channel)
411            .check(HealthCheckRequest {
412                service: "".to_owned(),
413            })
414            .await
415            .unwrap();
416
417        let reader = MetricsReader::new(&registry).with_prefix("test");
418        assert_eq!(
419            reader
420                .histogram_totals("handshake_latency", &[("result", "completed")])
421                .count,
422            1
423        );
424        assert_eq!(reader.value("pending_handshakes_peak", &[]), 1.0);
425        server.trigger_shutdown();
426    }
427
428    #[tokio::test]
429    async fn test_metrics_layer_error() {
430        #[derive(Clone)]
431        struct Metrics {
432            /// a flag to figure out whether the
433            /// on_request method has been called.
434            metrics_called: Arc<Mutex<bool>>,
435        }
436
437        impl MetricsCallbackProvider for Metrics {
438            fn on_request(&self, path: String) {
439                assert_eq!(path, "/grpc.health.v1.Health/Check");
440            }
441
442            fn on_response(
443                &self,
444                path: String,
445                _latency: Duration,
446                status: u16,
447                grpc_status_code: Code,
448            ) {
449                assert_eq!(path, "/grpc.health.v1.Health/Check");
450                assert_eq!(status, 200);
451                // According to https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
452                // code 5 is not_found , which is what we expect to get in this case
453                assert_eq!(grpc_status_code, Code::NotFound);
454                let mut m = self.metrics_called.lock().unwrap();
455                *m = true
456            }
457        }
458
459        let metrics = Metrics {
460            metrics_called: Arc::new(Mutex::new(false)),
461        };
462
463        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
464        let config = Config::new();
465        let keypair = Ed25519KeyPair::generate(&mut rand08::thread_rng());
466
467        let server = config
468            .server_builder_with_metrics(metrics.clone())
469            .bind(
470                &address,
471                Some(iota_tls::create_rustls_server_config(
472                    keypair.copy().private(),
473                    "test".to_string(),
474                )),
475            )
476            .await
477            .unwrap();
478        let address = server.local_addr().to_owned();
479        let channel = config
480            .connect(
481                &address,
482                iota_tls::create_rustls_client_config(
483                    keypair.public().to_owned(),
484                    "test".to_string(),
485                    None,
486                ),
487            )
488            .await
489            .unwrap();
490        let mut client = HealthClient::new(channel);
491
492        // Call the healthcheck for a service that doesn't exist
493        // that should give us back an error with code 5 (not_found)
494        // https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
495        let _ = client
496            .check(HealthCheckRequest {
497                service: "non-existing-service".to_owned(),
498            })
499            .await;
500
501        server.server_handle.shutdown().await;
502
503        assert!(metrics.metrics_called.lock().unwrap().deref());
504    }
505
506    async fn test_multiaddr(address: Multiaddr) {
507        let config = Config::new();
508        let keypair = Ed25519KeyPair::generate(&mut rand08::thread_rng());
509
510        let server_handle = config
511            .server_builder()
512            .bind(
513                &address,
514                Some(iota_tls::create_rustls_server_config(
515                    keypair.copy().private(),
516                    "test".to_string(),
517                )),
518            )
519            .await
520            .unwrap();
521        let address = server_handle.local_addr().to_owned();
522        let channel = config
523            .connect(
524                &address,
525                iota_tls::create_rustls_client_config(
526                    keypair.public().to_owned(),
527                    "test".to_string(),
528                    None,
529                ),
530            )
531            .await
532            .unwrap();
533        let mut client = HealthClient::new(channel);
534
535        client
536            .check(HealthCheckRequest {
537                service: "".to_owned(),
538            })
539            .await
540            .unwrap();
541
542        server_handle.server_handle.shutdown().await;
543    }
544
545    #[tokio::test]
546    async fn dns() {
547        let address: Multiaddr = "/dns/localhost/tcp/0/http".parse().unwrap();
548        test_multiaddr(address).await;
549    }
550
551    #[tokio::test]
552    async fn ip4() {
553        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
554        test_multiaddr(address).await;
555    }
556
557    #[tokio::test]
558    async fn ip6() {
559        let address: Multiaddr = "/ip6/::1/tcp/0/http".parse().unwrap();
560        test_multiaddr(address).await;
561    }
562
563    #[test]
564    fn request_timeout_falls_back_to_the_default() {
565        let unset = Config::new();
566        assert_eq!(
567            unset.server_builder().request_timeout(),
568            super::DEFAULT_GRPC_REQUEST_TIMEOUT
569        );
570
571        let mut configured = Config::new();
572        configured.request_timeout = Some(Duration::from_secs(7));
573        assert_eq!(
574            configured.server_builder().request_timeout(),
575            Duration::from_secs(7)
576        );
577    }
578
579    /// Answers only once the whole request body has arrived, so a body that
580    /// never completes keeps the request in flight for as long as the peer
581    /// likes.
582    #[derive(Clone)]
583    struct DrainBody;
584
585    impl tower::Service<http::Request<tonic::body::Body>> for DrainBody {
586        type Response = http::Response<tonic::body::Body>;
587        type Error = std::convert::Infallible;
588        type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
589
590        fn poll_ready(
591            &mut self,
592            _cx: &mut std::task::Context<'_>,
593        ) -> std::task::Poll<Result<(), Self::Error>> {
594            std::task::Poll::Ready(Ok(()))
595        }
596
597        fn call(&mut self, request: http::Request<tonic::body::Body>) -> Self::Future {
598            use http_body_util::BodyExt as _;
599            Box::pin(async move {
600                let _ = request.into_body().collect().await;
601                Ok(tonic::Status::new(Code::Ok, "").into_http())
602            })
603        }
604    }
605
606    impl tonic::server::NamedService for DrainBody {
607        const NAME: &'static str = "test.DrainBody";
608    }
609
610    /// A request body that either is already complete or never produces its
611    /// data, mirroring a peer that sends HEADERS and then goes silent.
612    enum TestBody {
613        Empty,
614        Stalled,
615    }
616
617    impl http_body::Body for TestBody {
618        type Data = bytes::Bytes;
619        type Error = std::convert::Infallible;
620
621        fn poll_frame(
622            self: std::pin::Pin<&mut Self>,
623            _cx: &mut std::task::Context<'_>,
624        ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
625            match *self {
626                TestBody::Empty => std::task::Poll::Ready(None),
627                TestBody::Stalled => std::task::Poll::Pending,
628            }
629        }
630
631        fn is_end_stream(&self) -> bool {
632            matches!(self, TestBody::Empty)
633        }
634    }
635
636    /// A request that never finishes sending its body must be answered by the
637    /// server-side deadline, releasing the admission slot it was holding.
638    #[tokio::test]
639    async fn server_timeout_releases_a_stalled_request_and_its_admission_slot() {
640        const REQUEST_TIMEOUT: Duration = Duration::from_millis(200);
641
642        let mut config = Config::new();
643        config.request_timeout = Some(REQUEST_TIMEOUT);
644        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
645        let server = config
646            .server_builder()
647            .add_service_with_concurrency_limit(DrainBody, std::num::NonZeroUsize::MIN, false)
648            .bind(&address, None)
649            .await
650            .unwrap();
651
652        let stream = tokio::net::TcpStream::connect(server.local_addr().to_socket_addr().unwrap())
653            .await
654            .unwrap();
655        let (mut sender, connection) = hyper::client::conn::http2::handshake(
656            hyper_util::rt::TokioExecutor::new(),
657            hyper_util::rt::TokioIo::new(stream),
658        )
659        .await
660        .unwrap();
661        tokio::spawn(connection);
662
663        let request = |body| {
664            http::Request::post("/test.DrainBody/Drain")
665                .header("content-type", "application/grpc")
666                .body(body)
667                .unwrap()
668        };
669        let grpc_status = |response: &http::Response<hyper::body::Incoming>| {
670            response.headers()["grpc-status"]
671                .to_str()
672                .unwrap()
673                .to_owned()
674        };
675
676        // Occupies the only admission slot without ever sending its body.
677        let stalled = tokio::time::timeout(
678            REQUEST_TIMEOUT * 25,
679            sender.send_request(request(TestBody::Stalled)),
680        )
681        .await
682        .expect("the server must answer a stalled request once its deadline passes")
683        .unwrap();
684        assert_eq!(
685            grpc_status(&stalled),
686            (Code::DeadlineExceeded as i32).to_string()
687        );
688
689        // The slot is free again, so a complete request is served.
690        let served = sender.send_request(request(TestBody::Empty)).await.unwrap();
691        assert_eq!(grpc_status(&served), (Code::Ok as i32).to_string());
692
693        server.server_handle.shutdown().await;
694    }
695
696    /// The configured stream cap is advertised to the peer, so requests beyond
697    /// it on one connection never reach the server until an earlier one ends.
698    #[tokio::test]
699    async fn stream_cap_bounds_the_requests_one_connection_can_open() {
700        #[derive(Clone, Default)]
701        struct Started(Arc<std::sync::atomic::AtomicUsize>);
702
703        impl MetricsCallbackProvider for Started {
704            fn on_request(&self, _path: String) {}
705
706            fn on_response(
707                &self,
708                _path: String,
709                _latency: Duration,
710                _status: u16,
711                _grpc_status_code: Code,
712            ) {
713            }
714
715            fn on_start(&self, _path: &str) {
716                self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
717            }
718        }
719
720        const STREAM_CAP: u32 = 2;
721
722        let started = Started::default();
723        let mut config = Config::new();
724        config.http2_max_concurrent_streams = Some(STREAM_CAP);
725        let address: Multiaddr = "/ip4/127.0.0.1/tcp/0/http".parse().unwrap();
726        let server = config
727            .server_builder_with_metrics(started.clone())
728            .add_service(DrainBody)
729            .bind(&address, None)
730            .await
731            .unwrap();
732
733        let stream = tokio::net::TcpStream::connect(server.local_addr().to_socket_addr().unwrap())
734            .await
735            .unwrap();
736        let (sender, connection) = hyper::client::conn::http2::handshake(
737            hyper_util::rt::TokioExecutor::new(),
738            hyper_util::rt::TokioIo::new(stream),
739        )
740        .await
741        .unwrap();
742        tokio::spawn(connection);
743
744        // One more stalled request than the cap allows; the surplus waits on
745        // the client side for a stream to free up.
746        let in_flight: Vec<_> = (0..STREAM_CAP + 1)
747            .map(|_| {
748                let mut sender = sender.clone();
749                tokio::spawn(async move {
750                    sender
751                        .send_request(
752                            http::Request::post("/test.DrainBody/Drain")
753                                .body(TestBody::Stalled)
754                                .unwrap(),
755                        )
756                        .await
757                })
758            })
759            .collect();
760        tokio::time::sleep(Duration::from_millis(300)).await;
761
762        assert_eq!(
763            started.0.load(std::sync::atomic::Ordering::SeqCst),
764            STREAM_CAP as usize,
765            "the server must only see as many requests as the stream cap allows"
766        );
767
768        for request in in_flight {
769            request.abort();
770        }
771        server.server_handle.shutdown().await;
772    }
773}
774
775#[derive(Clone)]
776struct RequestLifetimeLayer<M: MetricsCallbackProvider> {
777    metrics_provider: M,
778}
779
780impl<M: MetricsCallbackProvider, S> Layer<S> for RequestLifetimeLayer<M> {
781    type Service = RequestLifetime<M, S>;
782
783    fn layer(&self, inner: S) -> Self::Service {
784        RequestLifetime {
785            inner,
786            metrics_provider: self.metrics_provider.clone(),
787            path: None,
788        }
789    }
790}
791
792#[derive(Clone)]
793struct RequestLifetime<M: MetricsCallbackProvider, S> {
794    inner: S,
795    metrics_provider: M,
796    path: Option<String>,
797}
798
799impl<M: MetricsCallbackProvider, S, RequestBody> Service<Request<RequestBody>>
800    for RequestLifetime<M, S>
801where
802    S: Service<Request<RequestBody>>,
803{
804    type Response = S::Response;
805    type Error = S::Error;
806    type Future = S::Future;
807
808    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
809        self.inner.poll_ready(cx)
810    }
811
812    fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
813        if self.path.is_none() {
814            let path = request.uri().path().to_string();
815            self.metrics_provider.on_start(&path);
816            self.path = Some(path);
817        }
818        self.inner.call(request)
819    }
820}
821
822impl<M: MetricsCallbackProvider, S> Drop for RequestLifetime<M, S> {
823    fn drop(&mut self) {
824        if let Some(path) = &self.path {
825            self.metrics_provider.on_drop(path)
826        }
827    }
828}
iota_network_stack/server.rs

iota_network_stack/
server.rs
