mz_orchestratord/

metrics.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;

use axum::{Extension, Router, body::Body, routing::get};
use http::{HeaderMap, Method, Request, Response, StatusCode};
use k8s_controller::PrometheusMetrics;
use prometheus::{Encoder, TextEncoder};
use tower_http::{classify::ServerErrorsFailureClass, trace::TraceLayer};
use tracing::{Level, Span};

use mz_ore::metric;
use mz_ore::metrics::{MetricsRegistry, UIntGauge};

#[derive(Debug)]
pub struct Metrics {
    pub is_leader: UIntGauge,
    pub environmentd_needs_update: UIntGauge,
    /// The reconciliation metrics of every controller, as defined by
    /// `k8s_controller`: counts and durations of reconciliation passes, and
    /// of the steps within them.
    pub reconcile: Arc<PrometheusMetrics>,
}

impl Metrics {
    pub fn register_into(registry: &MetricsRegistry) -> Self {
        Self {
            is_leader: registry.register(
                metric! {
                    name: "orchestratord_is_leader",
                    help: "Whether this operator replica holds the controller leadership lease, and is therefore the replica reconciling. Summed across the replicas this should be 1. A sustained 0 means no replica can take the lease, for instance because the service account lacks permission on leases, and the operator is reconciling nothing.",
                }),
            environmentd_needs_update: registry.register(
                metric! {
                    name: "environmentd_needs_update",
                    help: "Count of organizations in this cluster which are running outdated pod templates. Only the operator replica holding the leadership lease reconciles, so the others report zero.",
                }),
            reconcile: {
                let reconcile = PrometheusMetrics::new("orchestratord")
                    .expect("valid reconciliation metric definitions");
                registry.register_collector(reconcile.clone());
                Arc::new(reconcile)
            },
        }
    }

    /// Records that this replica has taken the leadership lease.
    pub fn leadership_acquired(&self) {
        self.is_leader.set(1);
    }

    /// Records that this replica no longer holds the leadership lease, and
    /// resets the metrics that only mean anything while it is reconciling.
    ///
    /// Those are derived from what reconciliation observed, and the process
    /// outlives its own leadership: it keeps serving the conversion webhook
    /// after losing the lease. Without this, a former leader would go on
    /// publishing its last observation forever, so summing across the
    /// replicas would count the same organizations once per past leader.
    pub fn leadership_lost(&self) {
        self.is_leader.set(0);
        self.environmentd_needs_update.set(0);
    }
}

pub fn router(registry: MetricsRegistry) -> Router {
    add_tracing_layer(
        Router::new()
            .route("/metrics", get(metrics))
            .layer(Extension(registry)),
    )
}

#[allow(clippy::unused_async)]
async fn metrics(Extension(registry): Extension<MetricsRegistry>) -> (StatusCode, Vec<u8>) {
    let mut buf = vec![];
    let encoder = TextEncoder::new();
    let metric_families = registry.gather();
    encoder.encode(&metric_families, &mut buf).unwrap();
    (StatusCode::OK, buf)
}

///   Adds a tracing layer that reports an `INFO` level span per
///   request and reports a `WARN` event when a handler returns a
///   server error to the given Axum Router
///
///   This accepts a router instead of returning a layer itself
///   to avoid dealing with defining generics over a bunch of closures
///   (see <https://users.rust-lang.org/t/how-to-encapsulate-a-builder-that-depends-on-a-closure/71139/6>)
///
///   And this also can't be returned as a Router::new()::layer(TraceLayer)...
///   because the TraceLayer needs to be added to a Router after
///   all routes are defined, as it won't trace any routes defined
///   on the router after it's attached.
fn add_tracing_layer<S>(router: Router<S>) -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    router.layer(
        TraceLayer::new_for_http()
            .make_span_with(|request: &Request<Body>| {
                // This ugly macro is needed, unfortunately (and
                // copied from tower-http), because
                // `tracing::span!` required the level argument to
                // be static. Meaning we can't just pass
                // `self.level`.
                macro_rules! make_span {
                        ($level:expr) => {
                            tracing::span!(
                                $level,
                                "HTTP request",
                                "request.uri" = %request.uri(),
                                "request.version" = ?request.version(),
                                "request.method" = %request.method(),
                                "request.headers" = tracing::field::Empty,
                                "response.status" = tracing::field::Empty,
                                "response.status_code" = tracing::field::Empty,
                                "response.headers" = tracing::field::Empty,
                            )
                        }
                    }
                let span = if ["/api/health", "/metrics"].contains(&request.uri().path())
                    || request.method() == Method::OPTIONS
                {
                    make_span!(Level::DEBUG)
                } else {
                    make_span!(Level::INFO)
                };

                if let Ok(s) = serde_json::to_string(&display_headers(request.headers().clone())) {
                    span.record("request.headers", s);
                }

                span
            })
            .on_response(|response: &Response<Body>, _latency, span: &Span| {
                span.record(
                    "response.status",
                    &tracing::field::display(response.status()),
                );
                span.record("response.status_code", response.status().as_u16());
                if let Ok(s) = serde_json::to_string(&display_headers(response.headers().clone())) {
                    span.record("response.headers", s);
                }

                // Emit an event at the same level as the span. For the same reason as noted in the comment
                // above we can't use `tracing::event!(dynamic_level, ...)` since the level argument
                // needs to be static
                if span
                    .metadata()
                    .and_then(|m| Some(m.level()))
                    .unwrap_or(&Level::DEBUG)
                    == &Level::DEBUG
                {
                    tracing::debug!("HTTP response generated");
                } else {
                    tracing::info!("HTTP response generated");
                }
            })
            .on_failure(
                |error: ServerErrorsFailureClass, _latency: Duration, _span: &Span| {
                    tracing::warn!(error = ?error, "HTTP request handling error");
                },
            ),
    )
}

fn display_headers(mut headers: HeaderMap) -> BTreeMap<String, String> {
    // Don't log Authorization headers
    _ = headers.remove(http::header::AUTHORIZATION);

    headers
        .into_iter()
        .filter_map(|(k, v)| {
            k.map(|k| {
                (
                    k.to_string(),
                    String::from_utf8_lossy(v.as_bytes()).to_string(),
                )
            })
        })
        .collect()
}