Skip to main content

reth_node_metrics/
server.rs

1use crate::{
2    chain::ChainSpecInfo,
3    hooks::{Hook, Hooks},
4    process::register_process_metrics,
5    recorder::install_prometheus_recorder,
6    storage::StorageSettingsInfo,
7    version::VersionInfo,
8};
9use bytes::Bytes;
10use eyre::WrapErr;
11use http::{header::CONTENT_TYPE, HeaderValue, Request, Response, StatusCode};
12use http_body_util::{BodyExt, Full};
13use metrics::describe_gauge;
14use metrics_process::Collector;
15use reqwest::Client;
16use reth_metrics::metrics::Unit;
17use reth_tasks::TaskExecutor;
18use std::{convert::Infallible, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
19use tokio::sync::Semaphore;
20use tower::Service;
21use tower_http::compression::Compression;
22
23/// Bound rendering and compression work across all metrics connections.
24const MAX_CONCURRENT_METRICS_REQUESTS: usize = 4;
25
26/// Configuration for the [`MetricServer`]
27#[derive(Debug)]
28pub struct MetricServerConfig {
29    listen_addr: SocketAddr,
30    version_info: VersionInfo,
31    chain_spec_info: ChainSpecInfo,
32    storage_settings_info: Option<StorageSettingsInfo>,
33    task_executor: TaskExecutor,
34    hooks: Hooks,
35    push_gateway_url: Option<String>,
36    push_gateway_interval: Duration,
37    pprof_dump_dir: PathBuf,
38}
39
40impl MetricServerConfig {
41    /// Create a new [`MetricServerConfig`] with the given configuration
42    pub const fn new(
43        listen_addr: SocketAddr,
44        version_info: VersionInfo,
45        chain_spec_info: ChainSpecInfo,
46        task_executor: TaskExecutor,
47        hooks: Hooks,
48        pprof_dump_dir: PathBuf,
49    ) -> Self {
50        Self {
51            listen_addr,
52            hooks,
53            task_executor,
54            version_info,
55            chain_spec_info,
56            storage_settings_info: None,
57            push_gateway_url: None,
58            push_gateway_interval: Duration::from_secs(5),
59            pprof_dump_dir,
60        }
61    }
62
63    /// Set the storage settings information to expose over prometheus.
64    pub fn with_storage_settings_info(mut self, info: StorageSettingsInfo) -> Self {
65        self.storage_settings_info = Some(info);
66        self
67    }
68
69    /// Set the gateway URL and interval for pushing metrics
70    pub fn with_push_gateway(mut self, url: Option<String>, interval: Duration) -> Self {
71        self.push_gateway_url = url;
72        self.push_gateway_interval = interval;
73        self
74    }
75}
76
77/// [`MetricServer`] responsible for serving the metrics endpoint
78#[derive(Debug)]
79pub struct MetricServer {
80    config: MetricServerConfig,
81}
82
83impl MetricServer {
84    /// Create a new [`MetricServer`] with the given configuration
85    pub const fn new(config: MetricServerConfig) -> Self {
86        Self { config }
87    }
88
89    /// Spawns the metrics server
90    pub async fn serve(&self) -> eyre::Result<()> {
91        let MetricServerConfig {
92            listen_addr,
93            hooks,
94            task_executor,
95            version_info,
96            chain_spec_info,
97            storage_settings_info,
98            push_gateway_url,
99            push_gateway_interval,
100            pprof_dump_dir,
101        } = &self.config;
102
103        let hooks_for_endpoint = hooks.clone();
104        let executor_for_hooks = task_executor.clone();
105        self.start_endpoint(
106            *listen_addr,
107            Arc::new(move || {
108                hooks_for_endpoint.refresh_background(&executor_for_hooks);
109                hooks_for_endpoint.iter().for_each(|hook| hook());
110            }),
111            task_executor.clone(),
112            pprof_dump_dir.clone(),
113        )
114        .await
115        .wrap_err_with(|| format!("Could not start Prometheus endpoint at {listen_addr}"))?;
116
117        // Start push-gateway task if configured
118        if let Some(url) = push_gateway_url {
119            self.start_push_gateway_task(
120                url.clone(),
121                *push_gateway_interval,
122                hooks.clone(),
123                task_executor.clone(),
124            )?;
125        }
126
127        // Describe metrics after recorder installation
128        describe_db_metrics();
129        describe_static_file_metrics();
130        describe_rocksdb_metrics();
131        Collector::default().describe();
132        describe_memory_stats();
133        describe_io_stats();
134
135        version_info.register_version_metrics();
136        chain_spec_info.register_chain_spec_metrics();
137        if let Some(storage_settings_info) = storage_settings_info {
138            storage_settings_info.register_storage_settings_metrics();
139        }
140        register_process_metrics();
141
142        Ok(())
143    }
144
145    async fn start_endpoint<F: Hook>(
146        &self,
147        listen_addr: SocketAddr,
148        hook: Arc<F>,
149        task_executor: TaskExecutor,
150        pprof_dump_dir: PathBuf,
151    ) -> eyre::Result<()> {
152        let listener = tokio::net::TcpListener::bind(listen_addr)
153            .await
154            .wrap_err("Could not bind to address")?;
155
156        tracing::info!(target: "reth::cli", "Starting metrics endpoint at {}", listener.local_addr().unwrap());
157
158        let executor = task_executor.clone();
159        let requests = Arc::new(Semaphore::new(MAX_CONCURRENT_METRICS_REQUESTS));
160        task_executor.spawn_with_graceful_shutdown_signal(async move |mut signal| loop {
161            let io = tokio::select! {
162                _ = &mut signal => break,
163                io = listener.accept() => {
164                    match io {
165                        Ok((stream, _remote_addr)) => stream,
166                        Err(err) => {
167                            tracing::error!(%err, "failed to accept connection");
168                            continue;
169                        }
170                    }
171                }
172            };
173
174            let handle = install_prometheus_recorder();
175            let hook = hook.clone();
176            let pprof_dump_dir = pprof_dump_dir.clone();
177            let executor = executor.clone();
178            let compression_executor = executor.clone();
179            let service = tower::service_fn(move |req: Request<_>| {
180                let hook = hook.clone();
181                let pprof_dump_dir = pprof_dump_dir.clone();
182                let executor = executor.clone();
183                async move {
184                    let response =
185                        handle_request(req.uri().path(), hook, executor, handle, &pprof_dump_dir)
186                            .await;
187                    Ok::<_, Infallible>(response)
188                }
189            });
190
191            let service = Compression::new(service).gzip(true).no_br().no_deflate().no_zstd();
192            let requests = requests.clone();
193            let service = tower::service_fn(move |req| {
194                let mut service = service.clone();
195                let executor = compression_executor.clone();
196                let permit = requests.clone().try_acquire_owned();
197                async move {
198                    let Ok(permit) = permit else {
199                        return Ok::<_, Infallible>(metrics_error_response(
200                            StatusCode::SERVICE_UNAVAILABLE,
201                            "Metrics collection is busy. Try again later.",
202                        ))
203                    };
204                    let response = service.call(req).await?;
205                    // CompressionBody performs CPU work when polled. Fully collect it on a
206                    // blocking worker so Hyper only polls the completed bytes on Tokio workers.
207                    let span = tracing::Span::current();
208                    let result = executor
209                        .spawn_blocking(move || {
210                            let _entered = span.enter();
211                            // Keep the concurrency slot even if the client disconnects.
212                            let _permit = permit;
213                            tokio::runtime::Handle::current().block_on(async move {
214                                let (parts, body) = response.into_parts();
215                                body.collect().await.map(|body| {
216                                    Response::from_parts(parts, Full::new(body.to_bytes()))
217                                })
218                            })
219                        })
220                        .await;
221                    let response = match result {
222                        Ok(Ok(response)) => response,
223                        error => {
224                            tracing::error!(target: "reth::cli", ?error, "Failed to encode metrics response");
225                            metrics_error_response(
226                                StatusCode::INTERNAL_SERVER_ERROR,
227                                "Failed to encode metrics response.",
228                            )
229                        }
230                    };
231                    Ok(response)
232                }
233            });
234
235            let mut shutdown = signal.clone().ignore_guard();
236            tokio::task::spawn(async move {
237                let _ = jsonrpsee_server::serve_with_graceful_shutdown(io, service, &mut shutdown)
238                    .await
239                    .inspect_err(|error| tracing::debug!(%error, "failed to serve request"));
240            });
241        });
242
243        Ok(())
244    }
245
246    /// Starts a background task to push metrics to a metrics gateway
247    fn start_push_gateway_task(
248        &self,
249        url: String,
250        interval: Duration,
251        hooks: Hooks,
252        task_executor: TaskExecutor,
253    ) -> eyre::Result<()> {
254        let client = Client::builder()
255            .build()
256            .wrap_err("Could not create HTTP client to push metrics to gateway")?;
257        let executor = task_executor.clone();
258        task_executor.spawn_with_graceful_shutdown_signal(async move |mut signal| {
259            tracing::info!(url = %url, interval = ?interval, "Starting task to push metrics to gateway");
260            let handle = install_prometheus_recorder();
261            loop {
262                tokio::select! {
263                    _ = &mut signal => {
264                        tracing::info!("Shutting down task to push metrics to gateway");
265                        break;
266                    }
267                    _ = tokio::time::sleep(interval) => {
268                        hooks.refresh_background(&executor);
269                        let hooks = hooks.clone();
270                        let metrics_handle = handle.handle().clone();
271                        let metrics = match executor.spawn_blocking(move || {
272                            hooks.iter().for_each(|hook| hook());
273                            metrics_handle.render()
274                        }).await {
275                            Ok(metrics) => metrics,
276                            Err(err) => {
277                                tracing::warn!(%err, "Failed to collect metrics for gateway");
278                                continue;
279                            }
280                        };
281                        match client.put(&url).header("Content-Type", "text/plain").body(metrics).send().await {
282                            Ok(response) => {
283                                if !response.status().is_success() {
284                                    tracing::warn!(
285                                        status = %response.status(),
286                                        "Failed to push metrics to gateway"
287                                    );
288                                }
289                            }
290                            Err(err) => {
291                                tracing::warn!(%err, "Failed to push metrics to gateway");
292                            }
293                        }
294                    }
295                }
296            }
297        });
298        Ok(())
299    }
300}
301
302fn metrics_error_response(status: StatusCode, message: &'static str) -> Response<Full<Bytes>> {
303    let mut response = Response::new(Full::new(Bytes::from_static(message.as_bytes())));
304    *response.status_mut() = status;
305    response
306}
307
308fn describe_db_metrics() {
309    describe_gauge!("db.table_size", Unit::Bytes, "The size of a database table (in bytes)");
310    describe_gauge!("db.table_pages", "The number of database pages for a table");
311    describe_gauge!("db.table_entries", "The number of entries for a table");
312    describe_gauge!("db.freelist", "The number of pages on the freelist");
313    describe_gauge!("db.page_size", Unit::Bytes, "The size of a database page (in bytes)");
314    describe_gauge!(
315        "db.timed_out_not_aborted_transactions",
316        "Number of timed out transactions that were not aborted by the user yet"
317    );
318}
319
320fn describe_static_file_metrics() {
321    describe_gauge!("static_files.segment_size", Unit::Bytes, "The size of a static file segment");
322    describe_gauge!("static_files.segment_files", "The number of files for a static file segment");
323    describe_gauge!(
324        "static_files.segment_entries",
325        "The number of entries for a static file segment"
326    );
327}
328
329fn describe_rocksdb_metrics() {
330    describe_gauge!(
331        "rocksdb.table_size",
332        Unit::Bytes,
333        "The estimated size of a RocksDB table (SST + memtable)"
334    );
335    describe_gauge!("rocksdb.table_entries", "The estimated number of keys in a RocksDB table");
336    describe_gauge!(
337        "rocksdb.pending_compaction_bytes",
338        Unit::Bytes,
339        "Bytes pending compaction for a RocksDB table"
340    );
341    describe_gauge!("rocksdb.sst_size", Unit::Bytes, "The size of SST files for a RocksDB table");
342    describe_gauge!(
343        "rocksdb.memtable_size",
344        Unit::Bytes,
345        "The size of memtables for a RocksDB table"
346    );
347    describe_gauge!(
348        "rocksdb.wal_size",
349        Unit::Bytes,
350        "The total size of WAL (Write-Ahead Log) files. Important: this is not included in table_size or sst_size metrics"
351    );
352}
353
354#[cfg(all(feature = "jemalloc", unix))]
355fn describe_memory_stats() {
356    describe_gauge!(
357        "jemalloc.active",
358        Unit::Bytes,
359        "Total number of bytes in active pages allocated by the application"
360    );
361    describe_gauge!(
362        "jemalloc.allocated",
363        Unit::Bytes,
364        "Total number of bytes allocated by the application"
365    );
366    describe_gauge!(
367        "jemalloc.mapped",
368        Unit::Bytes,
369        "Total number of bytes in active extents mapped by the allocator"
370    );
371    describe_gauge!(
372        "jemalloc.metadata",
373        Unit::Bytes,
374        "Total number of bytes dedicated to jemalloc metadata"
375    );
376    describe_gauge!(
377        "jemalloc.resident",
378        Unit::Bytes,
379        "Total number of bytes in physically resident data pages mapped by the allocator"
380    );
381    describe_gauge!(
382        "jemalloc.retained",
383        Unit::Bytes,
384        "Total number of bytes in virtual memory mappings that were retained rather than \
385        being returned to the operating system via e.g. munmap(2)"
386    );
387}
388
389#[cfg(not(all(feature = "jemalloc", unix)))]
390const fn describe_memory_stats() {}
391
392#[cfg(target_os = "linux")]
393fn describe_io_stats() {
394    use metrics::describe_counter;
395
396    describe_counter!("io.rchar", "Characters read");
397    describe_counter!("io.wchar", "Characters written");
398    describe_counter!("io.syscr", "Read syscalls");
399    describe_counter!("io.syscw", "Write syscalls");
400    describe_counter!("io.read_bytes", Unit::Bytes, "Bytes read");
401    describe_counter!("io.write_bytes", Unit::Bytes, "Bytes written");
402    describe_counter!("io.cancelled_write_bytes", Unit::Bytes, "Cancelled write bytes");
403}
404
405#[cfg(not(target_os = "linux"))]
406const fn describe_io_stats() {}
407
408async fn handle_request<F: Hook>(
409    path: &str,
410    hook: Arc<F>,
411    executor: TaskExecutor,
412    handle: &crate::recorder::PrometheusRecorder,
413    pprof_dump_dir: &PathBuf,
414) -> Response<Full<Bytes>> {
415    match path {
416        "/debug/pprof/heap" => handle_pprof_heap(pprof_dump_dir),
417        "/debug/tokio/dump" => handle_tokio_dump().await,
418        _ => {
419            let metrics_handle = handle.handle().clone();
420            let metrics = match executor
421                .spawn_blocking(move || {
422                    hook();
423                    metrics_handle.render()
424                })
425                .await
426            {
427                Ok(metrics) => metrics,
428                Err(err) => {
429                    let mut response = Response::new(Full::new(Bytes::from(format!(
430                        "Failed to collect metrics: {err}"
431                    ))));
432                    *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
433                    return response;
434                }
435            };
436            let mut response = Response::new(Full::new(Bytes::from(metrics)));
437            response.headers_mut().insert(CONTENT_TYPE, HeaderValue::from_static("text/plain"));
438            response
439        }
440    }
441}
442
443#[cfg(all(feature = "jemalloc-prof", unix))]
444fn handle_pprof_heap(pprof_dump_dir: &PathBuf) -> Response<Full<Bytes>> {
445    use http::header::CONTENT_ENCODING;
446
447    match jemalloc_pprof::PROF_CTL.as_ref() {
448        Some(prof_ctl) => match prof_ctl.try_lock() {
449            Ok(_) => match jemalloc_pprof_dump(pprof_dump_dir) {
450                Ok(pprof) => {
451                    let mut response = Response::new(Full::new(Bytes::from(pprof)));
452                    response
453                        .headers_mut()
454                        .insert(CONTENT_TYPE, HeaderValue::from_static("application/octet-stream"));
455                    response
456                        .headers_mut()
457                        .insert(CONTENT_ENCODING, HeaderValue::from_static("gzip"));
458                    response
459                }
460                Err(err) => {
461                    let mut response = Response::new(Full::new(Bytes::from(format!(
462                        "Failed to dump pprof: {err}"
463                    ))));
464                    *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
465                    response
466                }
467            },
468            Err(_) => {
469                let mut response = Response::new(Full::new(Bytes::from_static(
470                    b"Profile dump already in progress. Try again later.",
471                )));
472                *response.status_mut() = StatusCode::SERVICE_UNAVAILABLE;
473                response
474            }
475        },
476        None => {
477            let mut response = Response::new(Full::new(Bytes::from_static(
478                b"jemalloc profiling not enabled. \
479                 Set MALLOC_CONF=prof:true or rebuild with jemalloc-prof feature.",
480            )));
481            *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
482            response
483        }
484    }
485}
486
487/// Equivalent to [`jemalloc_pprof::JemallocProfCtl::dump`], but accepts a directory that the
488/// temporary pprof file will be written to. The file is deleted when the function exits.
489#[cfg(all(feature = "jemalloc-prof", unix))]
490fn jemalloc_pprof_dump(pprof_dump_dir: &PathBuf) -> eyre::Result<Vec<u8>> {
491    use std::{ffi::CString, io::BufReader};
492
493    use mappings::MAPPINGS;
494    use pprof_util::parse_jeheap;
495    use tempfile::NamedTempFile;
496
497    reth_fs_util::create_dir_all(pprof_dump_dir)?;
498    let f = NamedTempFile::new_in(pprof_dump_dir)?;
499    let path = CString::new(f.path().as_os_str().as_encoded_bytes()).unwrap();
500
501    // SAFETY: "prof.dump" is documented as being writable and taking a C string as input:
502    // http://jemalloc.net/jemalloc.3.html#prof.dump
503    unsafe { tikv_jemalloc_ctl::raw::write(b"prof.dump\0", path.as_ptr()) }?;
504
505    let dump_reader = BufReader::new(f);
506    let profile =
507        parse_jeheap(dump_reader, MAPPINGS.as_deref()).map_err(|err| eyre::eyre!(Box::new(err)))?;
508    let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
509
510    Ok(pprof)
511}
512
513#[cfg(not(all(feature = "jemalloc-prof", unix)))]
514fn handle_pprof_heap(_pprof_dump_dir: &PathBuf) -> Response<Full<Bytes>> {
515    let mut response = Response::new(Full::new(Bytes::from_static(
516        b"jemalloc pprof support not compiled. Rebuild with the jemalloc-prof feature.",
517    )));
518    *response.status_mut() = StatusCode::NOT_IMPLEMENTED;
519    response
520}
521
522#[cfg(tokio_unstable)]
523async fn handle_tokio_dump() -> Response<Full<Bytes>> {
524    let handle = tokio::runtime::Handle::current();
525    let dump = handle.dump().await;
526
527    let mut output = String::new();
528    for (i, task) in dump.tasks().iter().enumerate() {
529        let trace = task.trace();
530        output.push_str(&format!("task {i}:\n{trace}\n\n"));
531    }
532
533    let mut response = Response::new(Full::new(Bytes::from(output)));
534    response.headers_mut().insert(CONTENT_TYPE, HeaderValue::from_static("text/plain"));
535    response
536}
537
538#[cfg(not(tokio_unstable))]
539async fn handle_tokio_dump() -> Response<Full<Bytes>> {
540    let mut response = Response::new(Full::new(Bytes::from_static(
541        b"tokio task dump not available. Rebuild with RUSTFLAGS=\"--cfg tokio_unstable\" and tokio's `taskdump` feature.",
542    )));
543    *response.status_mut() = StatusCode::NOT_IMPLEMENTED;
544    response
545}
546
547#[cfg(test)]
548mod tests {
549    use super::*;
550    use flate2::read::GzDecoder;
551    use http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, VARY};
552    use reqwest::Client;
553    use reth_tasks::Runtime;
554    use socket2::{Domain, Socket, Type};
555    use std::{
556        io::Read,
557        net::{SocketAddr, TcpListener},
558        sync::{
559            atomic::{AtomicUsize, Ordering},
560            mpsc, Mutex,
561        },
562    };
563
564    fn get_random_available_addr() -> SocketAddr {
565        let addr = &"127.0.0.1:0".parse::<SocketAddr>().unwrap().into();
566        let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
567        socket.set_reuse_address(true).unwrap();
568        socket.bind(addr).unwrap();
569        socket.listen(1).unwrap();
570        let listener = TcpListener::from(socket);
571        listener.local_addr().unwrap()
572    }
573
574    #[tokio::test(flavor = "multi_thread")]
575    async fn test_metrics_endpoint() {
576        // Install the recorder before serve() so gauge registrations are captured,
577        // mirroring how start_prometheus_endpoint() works in the real node launch.
578        install_prometheus_recorder();
579
580        let chain_spec_info = ChainSpecInfo { name: "test".to_string() };
581        let storage_settings_info = StorageSettingsInfo {
582            storage_v2: true,
583            pruning_mode: "archive",
584            prune_config: r#"{"block_interval":5}"#.to_string(),
585        };
586        let version_info = VersionInfo {
587            version: "test",
588            build_timestamp: "test",
589            cargo_features: "test",
590            git_sha: "test",
591            target_triple: "test",
592            build_profile: "test",
593        };
594
595        let runtime = Runtime::test();
596
597        let hooks = Hooks::builder().build();
598
599        let listen_addr = get_random_available_addr();
600        let config = MetricServerConfig::new(
601            listen_addr,
602            version_info,
603            chain_spec_info,
604            runtime.clone(),
605            hooks,
606            std::env::temp_dir(),
607        )
608        .with_storage_settings_info(storage_settings_info);
609
610        let server = MetricServer::new(config);
611        server.serve().await.unwrap();
612
613        // Send request to the metrics endpoint
614        let url = format!("http://{listen_addr}");
615        let client = Client::builder()
616            .no_gzip()
617            .no_brotli()
618            .no_deflate()
619            .no_zstd()
620            .timeout(Duration::from_secs(10))
621            .build()
622            .unwrap();
623        let response = client.get(&url).send().await.unwrap();
624        assert!(response.status().is_success());
625        assert_eq!(response.headers().get(CONTENT_ENCODING), None);
626
627        // Check the response body
628        let body = response.text().await.unwrap();
629        assert!(body.contains("reth_process_cpu_seconds_total"));
630        assert!(body.contains("reth_process_start_time_seconds"));
631        assert!(body.contains("process_cli_args"), "expected process_cli_args metric in output");
632        assert!(body.contains("reth_storage_settings"), "expected storage settings metric");
633        assert!(body.contains("storage_v2=\"true\""), "expected storage v2 label");
634        assert!(body.contains("pruning_mode=\"archive\""), "expected pruning mode label");
635        assert!(body.contains("prune_config="), "expected prune config label");
636
637        let storage_metric =
638            body.lines().find(|line| line.starts_with("reth_storage_settings{")).unwrap();
639        for (accept_encoding, gzip) in [
640            ("gzip", true),
641            ("br, gzip;q=0.5", true),
642            ("gzip;q=0", false),
643            ("identity", false),
644            ("br", false),
645        ] {
646            let response =
647                client.get(&url).header(ACCEPT_ENCODING, accept_encoding).send().await.unwrap();
648            assert_eq!(response.status(), StatusCode::OK);
649            assert_eq!(response.headers()[CONTENT_TYPE], "text/plain");
650            assert_eq!(response.headers()[VARY], "accept-encoding");
651            assert_eq!(
652                response.headers().get(CONTENT_ENCODING),
653                gzip.then_some(&HeaderValue::from_static("gzip")),
654                "Accept-Encoding: {accept_encoding}"
655            );
656            let bytes = response.bytes().await.unwrap();
657            let decoded = if gzip {
658                let mut decoded = String::new();
659                GzDecoder::new(bytes.as_ref()).read_to_string(&mut decoded).unwrap();
660                decoded
661            } else {
662                String::from_utf8(bytes.to_vec()).unwrap()
663            };
664            assert_eq!(
665                decoded.lines().find(|line| line.starts_with("reth_storage_settings{")),
666                Some(storage_metric)
667            );
668        }
669
670        // Fill all request slots with blocked collectors on separate connections. Excess
671        // scrapes must fail promptly without starting more blocking work.
672        let busy_addr = get_random_available_addr();
673        let started = Arc::new(Semaphore::new(0));
674        let (release, wait_for_release) = mpsc::channel::<()>();
675        let wait_for_release = Mutex::new(wait_for_release);
676        let hook = {
677            let started = started.clone();
678            Arc::new(move || {
679                started.add_permits(1);
680                assert_eq!(
681                    wait_for_release.lock().unwrap().recv_timeout(Duration::from_secs(10)),
682                    Err(mpsc::RecvTimeoutError::Disconnected)
683                );
684            })
685        };
686        server
687            .start_endpoint(busy_addr, hook, runtime.clone(), std::env::temp_dir())
688            .await
689            .unwrap();
690        let busy_url = format!("http://{busy_addr}");
691        let mut scrapes = tokio::task::JoinSet::new();
692        for _ in 0..MAX_CONCURRENT_METRICS_REQUESTS {
693            let client = client.clone();
694            let url = busy_url.clone();
695            scrapes.spawn(async move {
696                let response =
697                    client.get(url).header(ACCEPT_ENCODING, "gzip").send().await.unwrap();
698                assert_eq!(response.status(), StatusCode::OK);
699                assert_eq!(response.headers()[CONTENT_ENCODING], "gzip");
700                response.bytes().await.unwrap();
701            });
702        }
703        tokio::time::timeout(
704            Duration::from_secs(10),
705            started.acquire_many(MAX_CONCURRENT_METRICS_REQUESTS as u32),
706        )
707        .await
708        .unwrap()
709        .unwrap()
710        .forget();
711        let response = tokio::time::timeout(
712            Duration::from_secs(10),
713            client.get(&busy_url).header(ACCEPT_ENCODING, "gzip").send(),
714        )
715        .await
716        .unwrap()
717        .unwrap();
718        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
719        assert_eq!(started.available_permits(), 0);
720        // Closing the channel releases every collector, including during panic unwinding.
721        drop(release);
722        while let Some(result) =
723            tokio::time::timeout(Duration::from_secs(10), scrapes.join_next()).await.unwrap()
724        {
725            result.unwrap();
726        }
727        assert_eq!(client.get(&busy_url).send().await.unwrap().status(), StatusCode::OK);
728
729        // Make sure the runtime is dropped after the test runs.
730        drop(runtime);
731    }
732
733    #[tokio::test(flavor = "multi_thread")]
734    async fn test_background_hooks_do_not_block_collection() {
735        install_prometheus_recorder();
736
737        let collections = Arc::new(AtomicUsize::new(0));
738        let (release, wait_for_release) = mpsc::channel();
739        let wait_for_release = Mutex::new(Some(wait_for_release));
740        let completed = Arc::new(tokio::sync::Notify::new());
741        let hooks = Hooks::builder()
742            .with_background_interval(Duration::from_secs(60))
743            .with_background_hook({
744                let collections = collections.clone();
745                let completed = completed.clone();
746                move || {
747                    // Dropping the sender also unblocks the hook if a scrape assertion fails.
748                    if let Some(wait_for_release) = wait_for_release.lock().unwrap().take() &&
749                        wait_for_release.recv().is_err()
750                    {
751                        return
752                    }
753                    let collections = collections.fetch_add(1, Ordering::Relaxed) + 1;
754                    metrics::gauge!("test_background_hook_collections").set(collections as f64);
755                    completed.notify_one();
756                }
757            })
758            .build();
759
760        let runtime = Runtime::test();
761        let listen_addr = get_random_available_addr();
762        let config = MetricServerConfig::new(
763            listen_addr,
764            VersionInfo {
765                version: "test",
766                build_timestamp: "test",
767                cargo_features: "test",
768                git_sha: "test",
769                target_triple: "test",
770                build_profile: "test",
771            },
772            ChainSpecInfo { name: "test".to_string() },
773            runtime.clone(),
774            hooks,
775            std::env::temp_dir(),
776        );
777
778        MetricServer::new(config).serve().await.unwrap();
779
780        // the first scrape kicks off the background hook but must not wait for it
781        let url = format!("http://{listen_addr}");
782        let started = std::time::Instant::now();
783        let response =
784            Client::new().get(&url).timeout(Duration::from_millis(500)).send().await.unwrap();
785        assert!(response.status().is_success());
786        assert!(started.elapsed() < Duration::from_secs(1), "scrape waited for the hook");
787        assert_eq!(collections.load(Ordering::Relaxed), 0);
788
789        // once it completes, its values are rendered by the following scrape
790        release.send(()).unwrap();
791        tokio::time::timeout(Duration::from_secs(10), completed.notified()).await.unwrap();
792        let body = Client::new().get(&url).send().await.unwrap().text().await.unwrap();
793        assert!(
794            body.contains("test_background_hook_collections"),
795            "background hook metric missing: {body}"
796        );
797
798        // and scrapes within the interval do not collect again
799        assert_eq!(collections.load(Ordering::Relaxed), 1);
800
801        drop(runtime);
802    }
803}