Switch tracking unit for dbsync from total query time to total connection time (#35)

* feat: updated metrics collector to use connection time instead query time

* chore: adjusted docker image to build with essential packages

* chore: adjusted query to get only dmtr users

* fix: removed dmtr header

Author: paulobressan

operator/Cargo.lock

@@ -26,6 +26,8 @@ bech32 = "0.9.1"
 sha3 = "0.10.8"
 lazy_static = "1.4.0"
 deadpool-postgres = "0.12.1"
+chrono = "0.4.38"
+reqwest = { version = "0.12.4", features = ["json"] }
 
 [[bin]]
 name = "controller"

operator/Cargo.toml

@@ -2,6 +2,9 @@ FROM rust:1.74-slim-buster as build
 
 WORKDIR /app
 
+RUN apt update 
+RUN apt install -y build-essential pkg-config libssl-dev cmake
+
 COPY ./Cargo.lock ./Cargo.lock
 COPY ./Cargo.toml ./Cargo.toml
 COPY ./src ./src
@@ -12,4 +15,4 @@ FROM rust:1.74-slim-buster
 
 COPY --from=build /app/target/release/controller .
 
-CMD ["./controller"]
+CMD ["./controller"]

operator/Dockerfile

@@ -17,6 +17,7 @@ pub struct Config {
     pub dcu_per_second: HashMap<String, f64>,
 
     pub metrics_delay: Duration,
+    pub prometheus_url: String,
     pub statement_timeout: u64,
 }
 
@@ -64,6 +65,8 @@ impl Config {
                 .expect("METRICS_DELAY must be a number"),
         );
 
+        let prometheus_url = env::var("PROMETHEUS_URL").expect("PROMETHEUS_URL must be set");
+
         let statement_timeout = env::var("STATEMENT_TIMEOUT")
             .unwrap_or("120000".to_string())
             .parse::<u64>()
@@ -75,6 +78,7 @@ impl Config {
             db_max_connections,
             dcu_per_second,
             metrics_delay,
+            prometheus_url,
             statement_timeout,
         }
     }
@@ -93,6 +97,7 @@ mod tests {
         );
         env::set_var("DCU_PER_SECOND", "preview=5,preprod=5,mainnet=5");
         env::set_var("METRICS_DELAY", "100");
+        env::set_var("PROMETHEUS_URL", "localhost");
         env::set_var("STATEMENT_TIMEOUT", "100");
 
         let config = Config::from_env();

operator/src/config.rs

@@ -31,6 +31,9 @@ pub enum Error {
 
     #[error("Config Error: {0}")]
     ConfigError(String),
+
+    #[error("Http Request error: {0}")]
+    HttpError(String),
 }
 
 impl Error {
@@ -75,6 +78,11 @@ impl From<bech32::Error> for Error {
         Error::Bech32Error(value)
     }
 }
+impl From<reqwest::Error> for Error {
+    fn from(value: reqwest::Error) -> Self {
+        Error::HttpError(value.to_string())
+    }
+}
 
 #[derive(Clone)]
 pub struct State {

operator/src/lib.rs

@@ -1,14 +1,11 @@
-use futures::future;
+use chrono::Utc;
 use kube::{api::ListParams, Api, Client, Resource, ResourceExt};
 use prometheus::{opts, IntCounterVec, Registry};
-use std::{collections::HashMap, sync::Arc};
-use tracing::{error, info, instrument};
+use serde::{Deserialize, Deserializer};
+use std::sync::Arc;
+use tracing::{error, info, instrument, warn};
 
-use crate::{
-    get_config,
-    postgres::{Postgres, UserStatements},
-    DbSyncPort, Error, State,
-};
+use crate::{get_config, Config, DbSyncPort, Error, State};
 
 #[derive(Clone)]
 pub struct Metrics {
@@ -135,13 +132,14 @@ pub async fn run_metrics_collector(state: Arc<State>) {
             .await
             .expect("failed to create kube client");
 
-        let config = get_config();
-
-        let mut metrics_state: HashMap<String, HashMap<String, UserStatements>> = HashMap::new();
-
         let crds_api = Api::<DbSyncPort>::all(client.clone());
 
+        let config = get_config();
+        let mut last_execution = Utc::now();
+
         loop {
+            tokio::time::sleep(config.metrics_delay).await;
+
             let crds_result = crds_api.list(&ListParams::default()).await;
             if let Err(error) = crds_result {
                 error!(error = error.to_string(), "error to get k8s resources");
@@ -150,97 +148,115 @@ pub async fn run_metrics_collector(state: Arc<State>) {
             }
             let crds = crds_result.unwrap();
 
-            for crd in crds.items.iter().filter(|i| i.status.is_some()) {
-                let status = crd.status.as_ref().unwrap();
+            let end = Utc::now();
+            let interval = (end - last_execution).num_seconds();
 
-                let pg_connections_result = state.get_pg_by_network(&crd.spec.network);
-                if let Err(error) = pg_connections_result {
-                    error!(error = error.to_string());
-                    state.metrics.metrics_failure(&error);
+            last_execution = end;
+
+            let query = format!(
+                "sum by (user) (avg_over_time(pgbouncer_pools_client_active_connections{{user=~\"dmtr_.*\"}}[{interval}s] @ {})) > 0",
+                end.timestamp_millis() / 1000
+            );
+            dbg!(&query);
+
+            let response = collect_prometheus_metrics(config, query).await;
+            if let Err(err) = response {
+                error!(error = err.to_string(), "error to make prometheus request");
+                state.metrics.metrics_failure(&err);
+                continue;
+            }
+            let response = response.unwrap();
+
+            for result in response.data.result {
+                let crd = crds
+                    .iter()
+                    .filter(|c| c.status.is_some())
+                    .find(|c| c.status.as_ref().unwrap().username.eq(&result.metric.user));
+
+                if crd.is_none() {
+                    warn!(user = result.metric.user, "username doesnt have a crd");
                     continue;
                 }
 
-                let user_statements_result =
-                    get_user_statements(&status.username, pg_connections_result.unwrap()).await;
-                if let Err(error) = user_statements_result {
-                    error!(error = error.to_string(), "error get user statements");
+                let crd = crd.unwrap();
+
+                let dcu_per_second = config.dcu_per_second.get(&crd.spec.network);
+                if dcu_per_second.is_none() {
+                    let error = Error::ConfigError(format!(
+                        "dcu_per_second not configured to {} network",
+                        &crd.spec.network
+                    ));
+                    error!(error = error.to_string());
                     state.metrics.metrics_failure(&error);
                     continue;
                 }
 
-                let user_statements = user_statements_result.unwrap();
-
-                let latest_user_statement = metrics_state
-                    .entry(crd.spec.network.clone())
-                    .or_default()
-                    .get(&user_statements.usename);
-
-                if let Some(latest_user_statement) = latest_user_statement {
-                    let total_exec_time =
-                        user_statements.total_exec_time - latest_user_statement.total_exec_time;
-
-                    if total_exec_time == 0.0 {
-                        continue;
-                    }
-
-                    let dcu_per_second = config.dcu_per_second.get(&crd.spec.network);
-                    if dcu_per_second.is_none() {
-                        let error = Error::ConfigError(format!(
-                            "dcu_per_second not configured to {} network",
-                            &crd.spec.network
-                        ));
-                        error!(error = error.to_string());
-                        state.metrics.metrics_failure(&error);
-                        continue;
-                    }
-
-                    let dcu_per_second = dcu_per_second.unwrap();
-
-                    let dcu = (total_exec_time / 1000.) * dcu_per_second;
-                    state.metrics.count_dcu_consumed(
-                        &crd.namespace().unwrap(),
-                        &crd.spec.network,
-                        dcu,
-                    );
-                }
+                let dcu_per_second = dcu_per_second.unwrap();
+                let total_exec_time = result.value * (interval as f64);
 
-                metrics_state
-                    .entry(crd.spec.network.clone())
-                    .and_modify(|statements| {
-                        statements.insert(user_statements.usename.clone(), user_statements);
-                    });
+                let dcu = total_exec_time * dcu_per_second;
+                state
+                    .metrics
+                    .count_dcu_consumed(&crd.namespace().unwrap(), &crd.spec.network, dcu);
             }
-
-            tokio::time::sleep(config.metrics_delay).await;
         }
     });
 }
 
-async fn get_user_statements(
-    username: &str,
-    pg_connections: &[Postgres],
-) -> Result<UserStatements, Error> {
-    let tasks = future::join_all(
-        pg_connections
-            .iter()
-            .map(|pg| pg.find_metrics_by_user(username)),
-    )
-    .await;
-
-    let mut user_statements_all_host = UserStatements {
-        usename: username.into(),
-        total_exec_time: 0.,
-    };
-
-    for user_statements_by_host_result in tasks.into_iter() {
-        let user_statements_by_host = user_statements_by_host_result?;
-        if user_statements_by_host.is_none() {
-            continue;
-        }
-
-        let user_statements_by_host = user_statements_by_host.unwrap();
-        user_statements_all_host.total_exec_time += user_statements_by_host.total_exec_time;
+async fn collect_prometheus_metrics(
+    config: &Config,
+    query: String,
+) -> Result<PrometheusResponse, Error> {
+    let client = reqwest::Client::builder().build().unwrap();
+
+    let response = client
+        .get(format!("{}/query?query={query}", config.prometheus_url))
+        .send()
+        .await?;
+
+    let status = response.status();
+    if status.is_client_error() || status.is_server_error() {
+        error!(status = status.to_string(), "request status code fail");
+        return Err(Error::HttpError(format!(
+            "Prometheus request error. Status: {} Query: {}",
+            status, query
+        )));
     }
 
-    Ok(user_statements_all_host)
+    Ok(response.json().await.unwrap())
+}
+
+#[derive(Debug, Deserialize)]
+struct PrometheusDataResultMetric {
+    user: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct PrometheusDataResult {
+    metric: PrometheusDataResultMetric,
+    #[serde(deserialize_with = "deserialize_value")]
+    value: f64,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct PrometheusData {
+    result: Vec<PrometheusDataResult>,
+}
+
+#[derive(Debug, Deserialize)]
+struct PrometheusResponse {
+    data: PrometheusData,
+}
+
+fn deserialize_value<'de, D>(deserializer: D) -> Result<f64, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let value: Vec<serde_json::Value> = Deserialize::deserialize(deserializer)?;
+    Ok(value.into_iter().as_slice()[1]
+        .as_str()
+        .unwrap()
+        .parse::<f64>()
+        .unwrap())
 }

operator/src/metrics.rs

@@ -1,7 +1,7 @@
 use std::str::FromStr;
 
 use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
-use tokio_postgres::{NoTls, Row};
+use tokio_postgres::NoTls;
 
 use crate::{get_config, Error};
 
@@ -111,42 +111,4 @@ impl Postgres {
 
         Ok(result.is_some())
     }
-
-    pub async fn find_metrics_by_user(
-        &self,
-        username: &str,
-    ) -> Result<Option<UserStatements>, Error> {
-        let query_metrics = "select
-            usename,
-            sum(total_exec_time) as total_exec_time
-        from
-            pg_stat_statements
-        inner join
-            pg_catalog.pg_user on pg_catalog.pg_user.usesysid = userid
-        where 
-            pg_catalog.pg_user.usename = $1
-        group by
-            usename;";
-
-        let client = self.pool.get().await?;
-
-        let stmt = client.prepare(query_metrics).await?;
-        let result = client.query_opt(&stmt, &[&username]).await?;
-
-        Ok(result.as_ref().map(|row| row.into()))
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct UserStatements {
-    pub usename: String,
-    pub total_exec_time: f64,
-}
-impl From<&Row> for UserStatements {
-    fn from(row: &Row) -> Self {
-        Self {
-            usename: row.get("usename"),
-            total_exec_time: row.get("total_exec_time"),
-        }
-    }
 }

operator/src/postgres.rs

@@ -171,6 +171,8 @@ spec:
               value: "preview=5"
             - name: METRICS_DELAY
               value: "30"
+            - name: PROMETHEUS_URL
+              value: localhost 
             - name: RUST_LOG
               value: info,kube=debug,controller=debug
 ---