kube-yaseg/kuyaseg-controller/src/reporting/mod.rs

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

use anyhow::Context;
use confique::Config;
use opentelemetry::TraceId;
use tokio::sync::broadcast::{
    self,
    Sender,
};
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

use crate::reporting::k8s::K8sEvent;
use crate::reporting::prometheus::ProLabel;
use crate::reporting::tracing_evt::{
    TracingK8sEventLayer,
    TracingK8sEventLayerConfig,
};

pub mod k8s;
pub mod prometheus;
pub mod tracing_evt;

pub fn span_trace_id() -> Option<TraceId> {
    use opentelemetry::trace::TraceContextExt;
    use tracing_opentelemetry::OpenTelemetrySpanExt;

    match tracing::Span::current().context().span().span_context().trace_id() {
        TraceId::INVALID => None,
        other => Some(other),
    }
}

#[derive(Config)]
pub struct ReportingConfig {
    #[config(env = "KUYASEG_REPORT_K8S")]
    #[config(default = true)]
    pub k8s_events: bool,
    #[config(env = "KUYASEG_REPORT_PROMETHEUS")]
    #[config(default = false)]
    pub prometheus: bool,
}

/// You should await on handle somewhere
pub struct ReportingInit {
    pub tracing_layer: TracingK8sEventLayer,
    pub handle: JoinSet<()>,
    pub pro_registry: Option<prometheus_client::registry::Registry>,
}

/// For opentelemetry TraceId to be present, tracing_opentelemetry::layer must also be set
pub fn init(k8s_client: kube::Client, ct: CancellationToken) -> anyhow::Result<ReportingInit> {
    let config = ReportingConfig::builder()
        .env()
        .load()
        .context("Load Reporting Configuration")?;

    let mut reporting_jobs = JoinSet::new();

    let k8s_tx = config
        .k8s_events
        .then(|| init_k8s_events(k8s_client, &mut reporting_jobs, ct.clone()));
    let (pro_tx, pro_registry) = config
        .prometheus
        .then(|| init_pro_metrics(&mut reporting_jobs, ct.clone()))
        .unzip();

    Ok(ReportingInit {
        tracing_layer: TracingK8sEventLayer::new(TracingK8sEventLayerConfig {
            k8s_events: k8s_tx,
            prometheus: pro_tx,
        }),
        handle: reporting_jobs,
        pro_registry,
    })
}

fn init_k8s_events(client: kube::Client, joinset: &mut JoinSet<()>, ct: CancellationToken) -> Sender<Arc<K8sEvent>> {
    use kube::runtime::events::{
        Recorder,
        Reporter,
    };

    use crate::reporting::k8s::K8sEventRecorder;

    const WARN_THRESHOLD: usize = 10;
    const MAX_CAPACITY: usize = 100;

    let (tx, mut rx) = broadcast::channel::<Arc<K8sEvent>>(512);
    let reporter = Reporter {
        controller: env!("CARGO_PKG_NAME").into(),
        instance: std::env::var("CONTROLLER_POD_NAME").ok(),
    };
    let recorder = Recorder::new(client, reporter);
    let recorder = K8sEventRecorder::new(recorder);

    let rtx = tx.clone(); // for dumb recovery
    joinset.spawn(async move {
        let mut queue: VecDeque<Arc<K8sEvent>> = VecDeque::new();
        let mut warned = false;
        let mut drain_tick = tokio::time::interval(Duration::from_secs(6)); // 10 event/min
        let mut rec_limit = 4;

        while rec_limit > 0 {
            tokio::select! {
                biased;
                _ = ct.cancelled() => break,
                _ = drain_tick.tick() => {
                    if let Some(ev) = queue.pop_front() {
                        recorder.record(&ev).await;
                    }
                },
                msg = rx.recv() => {
                    match msg {
                        Ok(ev) => {
                            if queue.is_empty() {
                                recorder.record(&ev).await;
                                warned = false;
                            } else {
                                if queue.len() >= MAX_CAPACITY {
                                    queue.pop_front();
                                }
                                queue.push_back(ev);
                                if queue.len() > WARN_THRESHOLD && !warned {
                                    tracing::warn!(k13t.no_report=true, debug=?queue.len(), "K8s event queue exceeded threshold");
                                    warned = true;
                                } else if queue.len() <= WARN_THRESHOLD {
                                    warned = false;
                                }
                            }
                            rec_limit = 4;
                        }
                        Err(e) => {
                            tracing::error!(k13t.no_report=true, debug=?e, "Failed to receive k8s event from tracing");
                            rx = rtx.subscribe(); // dumb recover strategy
                            tokio::time::sleep(Duration::from_secs(3)).await;
                            rec_limit -= 1;
                        }
                    }
                }
            }
        }
        tracing::error!(
            k13t.no_report = true,
            "K8s Event reporting from tracing stopped"
        );
    });

    tx
}

fn init_pro_metrics(
    joinset: &mut JoinSet<()>,
    ct: CancellationToken,
) -> (Sender<Arc<ProLabel>>, prometheus_client::registry::Registry) {
    use prometheus_client::registry::Registry;

    use crate::reporting::prometheus::ProMetrics;

    let (tx, mut rx) = broadcast::channel::<Arc<ProLabel>>(512);
    let metrics = ProMetrics::new();

    let rtx = tx.clone();
    let promt = metrics.clone();
    let consumer_ct = ct.clone();
    joinset.spawn(async move {
        let mut rec_limit = 4;
        while rec_limit > 0 {
            tokio::select! {
                biased;
                _ = consumer_ct.cancelled() => break,
                msg = rx.recv() => {
                    match msg {
                        Ok(met) => {
                            promt.measure(&met);
                            rec_limit = 4;
                        }
                        Err(e) => {
                            tracing::error!(k13t.no_report=true, debug=?e, "Failed to receive prometheus measurement from tracing");
                            rx = rtx.subscribe(); // dumb recover strategy;
                            tokio::time::sleep(Duration::from_secs(3)).await;
                            rec_limit -= 1;
                        }
                    }
                }
            }
        }
        tracing::error!(
            k13t.no_report=true,
            "Prometheus measurement from tracing stopped"
        );
    });

    let registry = {
        let mut registry = Registry::default();
        metrics.register(&mut registry);
        registry
    };

    (tx, registry)
}