From bb9bdf74eccef1550a2a78febe4152eb02335e1d Mon Sep 17 00:00:00 2001
From: WenyXu <wenymedia@gmail.com>
Date: Mon, 13 May 2024 02:44:34 +0000
Subject: [PATCH] feat: export metric endpoint

---
 Cargo.lock              |  3 +++
 tests-chaos/Cargo.toml  |  3 +++
 tests-chaos/src/main.rs | 48 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index c9290a11fe..8662b40437 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10065,11 +10065,14 @@ checksum = "3369f5ac52d5eb6ab48c6b4ffdc8efbcad6b89c765749064ba298f2c68a16a76"
 name = "tests-chaos"
 version = "0.7.2"
 dependencies = [
+ "axum",
+ "axum-macros",
  "common-error",
  "common-macro",
  "common-telemetry",
  "common-time",
  "nix 0.26.4",
+ "prometheus",
  "rand",
  "rand_chacha",
  "reqwest",
diff --git a/tests-chaos/Cargo.toml b/tests-chaos/Cargo.toml
index 762f571eb1..ebdc6f30e7 100644
--- a/tests-chaos/Cargo.toml
+++ b/tests-chaos/Cargo.toml
@@ -5,11 +5,14 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+axum-macros = "0.3.8"
+axum.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
 common-telemetry.workspace = true
 common-time = { workspace = true }
 nix = { version = "0.26", features = ["process"] }
+prometheus.workspace = true
 rand = { workspace = true }
 rand_chacha = "0.3.1"
 reqwest.workspace = true
diff --git a/tests-chaos/src/main.rs b/tests-chaos/src/main.rs
index a7ce75d471..7aed8f941d 100644
--- a/tests-chaos/src/main.rs
+++ b/tests-chaos/src/main.rs
@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
+use std::net::SocketAddr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
+use axum::extract::{Query, State};
+use axum::Router;
 use bare::process::{Pid, ProcessManager};
 use common_telemetry::{info, warn};
 use nix::sys::signal::Signal;
@@ -39,14 +42,57 @@ mod bare;
 mod error;
 mod utils;
 
+use axum::routing::get;
+use prometheus::{Encoder, TextEncoder};
+
 use crate::error::Result;
 use crate::utils::{generate_create_table_expr, get_conf_path, path_to_stdio, render_config_file};
 const DEFAULT_LOG_LEVEL: &str = "--log-level=debug,hyper=warn,tower=warn,datafusion=warn,reqwest=warn,sqlparser=warn,h2=info,opendal=info";
 
+#[derive(Copy, Clone)]
+pub struct MetricsHandler;
+
+impl MetricsHandler {
+    pub fn render(&self) -> String {
+        let mut buffer = Vec::new();
+        let encoder = TextEncoder::new();
+        // Gather the metrics.
+        let metric_families = prometheus::gather();
+        // Encode them to send.
+        match encoder.encode(&metric_families, &mut buffer) {
+            Ok(_) => match String::from_utf8(buffer) {
+                Ok(s) => s,
+                Err(e) => e.to_string(),
+            },
+            Err(e) => e.to_string(),
+        }
+    }
+}
+
+#[axum_macros::debug_handler]
+pub async fn metrics(
+    State(state): State<MetricsHandler>,
+    Query(_params): Query<HashMap<String, String>>,
+) -> String {
+    state.render()
+}
+
 #[tokio::main]
 async fn main() {
     common_telemetry::init_default_ut_logging();
 
+    tokio::spawn(async move {
+        let app = Router::new()
+            .route("/metric", get(metrics))
+            .with_state(MetricsHandler);
+        // run our app with hyper, listening globally on port 3000
+        let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
+        axum::Server::bind(&addr)
+            .serve(app.into_make_service())
+            .await
+            .unwrap(); 
+    });
+
     let state = Arc::new(TestState {
         killed: AtomicBool::new(false),
     });