substrate_prometheus_endpoint/
lib.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18mod sourced;
19
20use hyper::{http::StatusCode, Request, Response};
21use prometheus::{core::Collector, Encoder, TextEncoder};
22use std::net::SocketAddr;
23
24pub use prometheus::{
25	self,
26	core::{
27		AtomicF64 as F64, AtomicI64 as I64, AtomicU64 as U64, GenericCounter as Counter,
28		GenericCounterVec as CounterVec, GenericGauge as Gauge, GenericGaugeVec as GaugeVec,
29	},
30	exponential_buckets, Error as PrometheusError, Histogram, HistogramOpts, HistogramVec, Opts,
31	Registry,
32};
33pub use sourced::{MetricSource, SourcedCounter, SourcedGauge, SourcedMetric};
34
35type Body = http_body_util::Full<hyper::body::Bytes>;
36
37pub fn register<T: Clone + Collector + 'static>(
38	metric: T,
39	registry: &Registry,
40) -> Result<T, PrometheusError> {
41	registry.register(Box::new(metric.clone()))?;
42	Ok(metric)
43}
44
45#[derive(Debug, thiserror::Error)]
46pub enum Error {
47	/// Hyper internal error.
48	#[error(transparent)]
49	Hyper(#[from] hyper::Error),
50
51	/// Http request error.
52	#[error(transparent)]
53	Http(#[from] hyper::http::Error),
54
55	/// i/o error.
56	#[error(transparent)]
57	Io(#[from] std::io::Error),
58
59	#[error("Prometheus port {0} already in use.")]
60	PortInUse(SocketAddr),
61}
62
63async fn request_metrics(
64	req: Request<hyper::body::Incoming>,
65	registry: Registry,
66) -> Result<Response<Body>, Error> {
67	if req.uri().path() == "/metrics" {
68		let metric_families = registry.gather();
69		let mut buffer = vec![];
70		let encoder = TextEncoder::new();
71		encoder.encode(&metric_families, &mut buffer).unwrap();
72
73		Response::builder()
74			.status(StatusCode::OK)
75			.header("Content-Type", encoder.format_type())
76			.body(Body::from(buffer))
77			.map_err(Error::Http)
78	} else {
79		Response::builder()
80			.status(StatusCode::NOT_FOUND)
81			.body(Body::from("Not found."))
82			.map_err(Error::Http)
83	}
84}
85
86/// Initializes the metrics context, and starts an HTTP server
87/// to serve metrics.
88pub async fn init_prometheus(prometheus_addr: SocketAddr, registry: Registry) -> Result<(), Error> {
89	let listener = tokio::net::TcpListener::bind(&prometheus_addr).await.map_err(|e| {
90		log::error!(target: "prometheus", "Error binding to '{:#?}': {:#?}", prometheus_addr, e);
91		Error::PortInUse(prometheus_addr)
92	})?;
93
94	init_prometheus_with_listener(listener, registry).await
95}
96
97/// Init prometheus using the given listener.
98async fn init_prometheus_with_listener(
99	listener: tokio::net::TcpListener,
100	registry: Registry,
101) -> Result<(), Error> {
102	log::info!(target: "prometheus", "〽️ Prometheus exporter started at {}", listener.local_addr()?);
103
104	let server = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new());
105
106	loop {
107		let io = match listener.accept().await {
108			Ok((sock, _)) => hyper_util::rt::TokioIo::new(sock),
109			Err(e) => {
110				log::debug!(target: "prometheus", "Error accepting connection: {:?}", e);
111				continue;
112			},
113		};
114
115		let registry = registry.clone();
116
117		let conn = server
118			.serve_connection_with_upgrades(
119				io,
120				hyper::service::service_fn(move |req| request_metrics(req, registry.clone())),
121			)
122			.into_owned();
123
124		tokio::spawn(async move {
125			if let Err(err) = conn.await {
126				log::debug!(target: "prometheus", "connection error: {:?}", err);
127			}
128		});
129	}
130}
131
132#[cfg(test)]
133mod tests {
134	use super::*;
135	use http_body_util::BodyExt;
136	use hyper::Uri;
137	use hyper_util::{client::legacy::Client, rt::TokioExecutor};
138
139	const METRIC_NAME: &str = "test_test_metric_name_test_test";
140
141	#[tokio::test]
142	async fn prometheus_works() {
143		let listener =
144			tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("Creates listener");
145
146		let local_addr = listener.local_addr().expect("Returns the local addr");
147
148		let registry = Registry::default();
149		register(
150			prometheus::Counter::new(METRIC_NAME, "yeah").expect("Creates test counter"),
151			&registry,
152		)
153		.expect("Registers the test metric");
154
155		tokio::spawn(init_prometheus_with_listener(listener, registry));
156
157		let client = Client::builder(TokioExecutor::new()).build_http::<Body>();
158
159		let res = client
160			.get(Uri::try_from(&format!("http://{}/metrics", local_addr)).expect("Parses URI"))
161			.await
162			.expect("Requests metrics");
163
164		assert!(res.status().is_success());
165
166		let buf = res.into_body().collect().await.expect("Failed to read HTTP body").to_bytes();
167		let body = String::from_utf8(buf.to_vec()).expect("Converts body to String");
168
169		assert!(body.contains(&format!("{} 0", METRIC_NAME)));
170	}
171}
substrate_prometheus_endpoint/lib.rs

substrate_prometheus_endpoint/
lib.rs
