cli-batteries 0.5.0

Batteries included opinionated command line interfaces
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

#![cfg(feature = "prometheus")]
use crate::{default_from_clap, shutdown::await_shutdown};
use clap::Parser;
use eyre::{bail, ensure, Result as EyreResult, WrapErr as _};
use hyper::{
    body::HttpBody,
    header::CONTENT_TYPE,
    service::{make_service_fn, service_fn},
    Body, Method, Request, Response, Server,
};
use once_cell::sync::Lazy;
use prometheus::{
    opts, register_counter, register_gauge, register_histogram, Counter, Encoder as _, Gauge,
    Histogram,
};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use tracing::{error, info, instrument, trace};
use url::{Host, Url};

#[cfg(feature = "otlp")]
use crate::trace_from_headers;

// TODO: Spans, traceId and SpanKind trace_span!("request", "otel.kind" =
// %SpanKind::Server, "http.url" = ..),

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Parser)]
#[group(skip)]
pub struct Options {
    /// Prometheus scrape endpoint
    // See <https://github.com/prometheus/prometheus/wiki/Default-port-allocations>
    #[clap(long, env, default_value = "http://127.0.0.1:9998/metrics")]
    pub prometheus: Url,
}

default_from_clap!(Options);

static REQ_COUNTER: Lazy<Counter> = Lazy::new(|| {
    register_counter!(opts!(
        "prometheus_requests_total",
        "Number of Prometheus scrape requests made."
    ))
    .unwrap()
});
static REQ_BODY_GAUGE: Lazy<Gauge> = Lazy::new(|| {
    register_gauge!(opts!(
        "prometheus_response_size_bytes",
        "The Prometheus response sizes in bytes."
    ))
    .unwrap()
});
static REQ_HISTOGRAM: Lazy<Histogram> = Lazy::new(|| {
    register_histogram!(
        "prometheus_request_duration_seconds",
        "The Prometheus request latencies in seconds."
    )
    .unwrap()
});

#[allow(clippy::unnecessary_wraps)]
#[allow(clippy::needless_pass_by_value)]
#[allow(clippy::unused_async)]
async fn serve_req(_req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
    // let encoder = prometheus::ProtobufEncoder;
    let encoder = prometheus::TextEncoder;
    let metric_families = prometheus::gather();
    let mut buffer = vec![];
    let result = encoder.encode(&metric_families, &mut buffer);
    if let Err(e) = result {
        error!("Internal server error: {}", e);
        let response = Response::builder()
            .status(500)
            .body(Body::from(e.to_string()))
            .unwrap(); // TODO
        return Ok(response);
    }

    let response = Response::builder()
        .status(200)
        .header(CONTENT_TYPE, encoder.format_type())
        .body(Body::from(buffer))
        .unwrap(); // TODO

    Ok(response)
}

#[allow(clippy::unused_async)] // We are implementing an interface
#[instrument(level="debug", name="prometheus_request", skip(req), fields(http.uri = %req.uri(), http.method = %req.method()))]
async fn route(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
    #[cfg(feature = "otlp")]
    trace_from_headers(req.headers());

    trace!("Receiving request at path {}", req.uri());
    REQ_COUNTER.inc();
    let timer = REQ_HISTOGRAM.start_timer();

    let response = match (req.method(), req.uri().path()) {
        (&Method::GET, "/metrics") => serve_req(req).await?,
        _ => Response::builder()
            .status(404)
            .body(Body::from("404"))
            .unwrap(),
    };

    #[allow(clippy::cast_precision_loss)]
    REQ_BODY_GAUGE.set(response.body().size_hint().lower() as f64);
    timer.observe_duration();
    Ok(response)
}

pub async fn main(options: Options) -> EyreResult<()> {
    ensure!(
        options.prometheus.scheme() == "http",
        "Only http:// is supported in {}",
        options.prometheus
    );
    ensure!(
        options.prometheus.path() == "/metrics",
        "Only /metrics is supported in {}",
        options.prometheus
    );
    let ip: IpAddr = match options.prometheus.host() {
        Some(Host::Ipv4(ip)) => ip.into(),
        Some(Host::Ipv6(ip)) => ip.into(),
        Some(_) => bail!("Cannot bind {}", options.prometheus),
        None => Ipv4Addr::LOCALHOST.into(),
    };
    let port = options.prometheus.port().unwrap_or(9998);
    let addr = SocketAddr::new(ip, port);

    let server = Server::try_bind(&addr)
        .wrap_err("Could not bind Prometheus server port")?
        .serve(make_service_fn(|_| async {
            Ok::<_, hyper::Error>(service_fn(route))
        }))
        .with_graceful_shutdown(await_shutdown());
    info!(url = %options.prometheus, "Metrics server listening");

    server.await?;
    Ok(())
}