example: Use hyper to serve metrics

adamchalmers · adamchalmers · commit 5021289bf15e · 2022-09-19T17:35:46.000-05:00
Signed-off-by: Adam Chalmers &lt;adam.s.chalmers@gmail.com&gt;
diff --git a/Cargo.toml b/Cargo.toml
@@ -34,10 +34,17 @@ quickcheck = "1"
 rand = "0.8.4"
 tide = "0.16"
 actix-web = "4"
+tokio = { version = "1", features = ["rt-multi-thread", "net", "macros", "signal"] }
+hyper = { version = "0.14.16", features = ["server", "http1", "tcp"] }
+lazy_static = "1.4"
 
 [build-dependencies]
 prost-build = { version = "0.9.0", optional = true }
 
+[[example]]
+name = "hyper"
+crate-type = ["bin"]
+
 [[bench]]
 name = "family"
 harness = false
diff --git a/examples/hyper.rs b/examples/hyper.rs
@@ -0,0 +1,89 @@
+use hyper::{
+    service::{make_service_fn, service_fn},
+    Body, Request, Response, Server,
+};
+use prometheus_client::{encoding::text::encode, metrics::counter::Counter, registry::Registry};
+use std::{
+    future::Future,
+    io,
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    pin::Pin,
+    sync::Arc,
+    time::Duration,
+};
+use tokio::signal::unix::{signal, SignalKind};
+
+lazy_static::lazy_static! {
+    pub static ref EXAMPLE_COUNTER: Counter = Counter::default();
+}
+
+#[tokio::main]
+async fn main() {
+    let registry = register_metrics();
+
+    // Spawn a server to serve the OpenMetrics endpoint.
+    let metrics_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8001);
+    tokio::task::spawn(start_server(metrics_addr, registry));
+
+    loop {
+        EXAMPLE_COUNTER.inc();
+        tokio::time::sleep(Duration::from_secs(1)).await;
+    }
+}
+
+pub fn register_metrics() -> Registry {
+    let mut registry = <Registry>::with_prefix("tokio_hyper_example");
+
+    registry.register(
+        "example_counter",
+        "An example counter that will get incremented automatically",
+        Box::new(EXAMPLE_COUNTER.clone()),
+    );
+
+    registry
+}
+
+/// Start a HTTP server to report metrics.
+pub async fn start_server(metrics_addr: SocketAddr, registry: Registry) {
+    let mut shutdown_stream = signal(SignalKind::terminate()).unwrap();
+
+    eprintln!("Starting metrics server on {metrics_addr}");
+
+    let registry = Arc::new(registry);
+    Server::bind(&metrics_addr)
+        .serve(make_service_fn(move |_conn| {
+            let registry = registry.clone();
+            async move {
+                let handler = make_handler(registry);
+                Ok::<_, io::Error>(service_fn(handler))
+            }
+        }))
+        .with_graceful_shutdown(async move {
+            shutdown_stream.recv().await;
+        })
+        .await
+        .unwrap();
+}
+
+/// This function returns a HTTP handler (i.e. another function)
+pub fn make_handler(
+    registry: Arc<Registry>,
+) -> impl Fn(Request<Body>) -> Pin<Box<dyn Future<Output = io::Result<Response<Body>>> + Send>> {
+    // This closure accepts a request and responds with the OpenMetrics encoding of our metrics.
+    move |_req: Request<Body>| {
+        let reg = registry.clone();
+        Box::pin(async move {
+            let mut buf = Vec::new();
+            encode(&mut buf, &reg.clone()).map(|_| {
+                let body = std::str::from_utf8(buf.as_slice()).unwrap().to_string();
+                Response::builder()
+                    .header(
+                        hyper::header::CONTENT_TYPE,
+                        "application/openmetrics-text; version=1.0.0; charset=utf-8",
+                    )
+                    .body(Body::from(body))
+                    .unwrap()
+            })
+        })
+    }
+}
