tests tools: fix gauges/counters in prometheus metrics (#67)

* Fix gauges/counters Move values to gauges/counters as it needs. * add dump interval for prometheus to bgtask --------- Co-authored-by: Kirill Fomichev <[email protected]>
rpcpool · Jun 1, 2023 · e140dfc · e140dfc
1 parent 8722ae1
commit e140dfc
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/tests/bgtask_creator/Cargo.toml b/tests/bgtask_creator/Cargo.toml
@@ -17,4 +17,5 @@ sea-orm = { version = "0.10.6", features = ["macros", "runtime-tokio-rustls", "s
 sea-query = { version = "0.28.1", features = ["postgres-array"] }
 solana-sdk = "1.14.10"
 sqlx = { version = "0.6.2", features = ["macros", "runtime-tokio-rustls", "postgres", "uuid", "offline", "json"] }
-tokio = { version = "1.23.0", features = ["macros", "rt-multi-thread", "fs"] }
+tokio = { version = "1.23.0", features = ["macros", "rt-multi-thread"] }
+txn_forwarder = { path = "../txn_forwarder" }
diff --git a/tests/bgtask_creator/src/main.rs b/tests/bgtask_creator/src/main.rs
@@ -1,5 +1,4 @@
 use {
-    anyhow::Context,
     clap::{value_parser, Arg, ArgAction, Command},
     digital_asset_types::dao::{
         asset, asset_authority, asset_creators, asset_data, asset_grouping,
@@ -14,19 +13,18 @@ use {
         metrics::setup_metrics,
         tasks::{BgTask, DownloadMetadata, DownloadMetadataTask, IntoTaskData, TaskManager},
     },
-    prometheus::{IntGaugeVec, Opts, Registry, TextEncoder},
+    prometheus::{IntGaugeVec, Opts, Registry},
     sea_orm::{
         entity::*, query::*, DbBackend, DeleteResult, EntityTrait, JsonValue, SqlxPostgresConnector,
     },
     solana_sdk::pubkey::Pubkey,
     sqlx::types::chrono::Utc,
     std::{collections::HashMap, path::PathBuf, str::FromStr, sync::Arc, time},
-    tokio::fs,
+    tokio::time::Duration,
+    txn_forwarder::save_metrics,
 };
 
 lazy_static::lazy_static! {
-    pub static ref REGISTRY: Registry = Registry::new();
-
     pub static ref BGTASK_SHOW: IntGaugeVec = IntGaugeVec::new(
         Opts::new("bgtask_show", "Number of assets in tasks"),
         &["type", "kind"]
@@ -50,9 +48,6 @@ async fn main() -> anyhow::Result<()> {
     init_logger();
     info!("Starting bgtask creator");
 
-    REGISTRY.register(Box::new(BGTASK_SHOW.clone())).unwrap();
-    REGISTRY.register(Box::new(BGTASK_CREATE.clone())).unwrap();
-
     let matches = Command::new("bgtaskcreator")
         .arg(
             Arg::new("config")
@@ -112,6 +107,13 @@ async fn main() -> anyhow::Result<()> {
                 .required(false)
                 .action(ArgAction::Set),
         )
+        .arg(
+            Arg::new("prom_save_interval")
+                .long("prom-save-interval")
+                .help("Prometheus metrics file update interval")
+                .required(false)
+                .action(ArgAction::Set),
+        )
         .subcommand(
             Command::new("show").about("Show tasks").arg(
                 Arg::new("print")
@@ -132,6 +134,20 @@ async fn main() -> anyhow::Result<()> {
         .subcommand(Command::new("delete").about("Delete ALL pending background tasks"))
         .get_matches();
 
+    let registry = Registry::new();
+    registry.register(Box::new(BGTASK_SHOW.clone()))?;
+    registry.register(Box::new(BGTASK_CREATE.clone()))?;
+    let metrics_jh = save_metrics(
+        registry,
+        matches.get_one::<String>("prom").cloned(),
+        Duration::from_millis(
+            matches
+                .get_one::<u64>("prom_save_interval")
+                .cloned()
+                .unwrap_or(1_000),
+        ),
+    );
+
     let config_path = matches.get_one::<PathBuf>("config");
     if let Some(config_path) = config_path {
         info!("Loading config from: {}", config_path.display());
@@ -389,14 +405,7 @@ WHERE
         }
     }
 
-    if let Some(prom) = matches.get_one::<String>("prom") {
-        let metrics = TextEncoder::new()
-            .encode_to_string(&REGISTRY.gather())
-            .context("could not encode custom metrics")?;
-        fs::write(prom, metrics).await?;
-    }
-
-    Ok(())
+    metrics_jh.await
 }
 
 fn find_by_type<'a>(

diff --git a/tests/tree-status/src/main.rs b/tests/tree-status/src/main.rs
@@ -8,7 +8,7 @@ use {
         stream::{self, StreamExt},
     },
     log::{debug, error, info},
-    prometheus::{IntCounter, IntGaugeVec, Opts, Registry},
+    prometheus::{IntGauge, IntGaugeVec, Opts, Registry},
     sea_orm::{
         sea_query::{Expr, Value},
         ColumnTrait, ConnectionTrait, DatabaseConnection, DbBackend, DbErr, EntityTrait,
@@ -67,11 +67,11 @@ lazy_static::lazy_static! {
         &["tree"]
     ).unwrap();
 
-    pub static ref TREE_STATUS_LEAVES_COMPLETED: IntCounter = IntCounter::new(
+    pub static ref TREE_STATUS_LEAVES_COMPLETED: IntGauge = IntGauge::new(
         "tree_status_leaves_completed", "Number of complete trees"
     ).unwrap();
 
-    pub static ref TREE_STATUS_LEAVES_INCOMPLETE: IntCounter = IntCounter::new(
+    pub static ref TREE_STATUS_LEAVES_INCOMPLETE: IntGauge = IntGauge::new(
         "tree_status_leaves_incomplete", "Number of incomplete trees"
     ).unwrap();
 
@@ -246,7 +246,7 @@ async fn main() -> anyhow::Result<()> {
     if let Some(group) = args.prom_group.clone() {
         labels.insert("group".to_owned(), group);
     }
-    let registry = Registry::new_custom(None, Some(labels)).unwrap();
+    let registry = Registry::new_custom(None, Some(labels))?;
     registry.register(Box::new(TREE_STATUS_MAX_SEQ.clone()))?;
     registry.register(Box::new(TREE_STATUS_MISSING_SEQ.clone()))?;
     registry.register(Box::new(TREE_STATUS_LEAVES_COMPLETED.clone()))?;

diff --git a/tests/txn_forwarder/src/main.rs b/tests/txn_forwarder/src/main.rs
@@ -9,7 +9,7 @@ use {
     log::info,
     plerkle_messenger::{MessengerConfig, ACCOUNT_STREAM, TRANSACTION_STREAM},
     plerkle_serialization::serializer::seralize_encoded_transaction_with_status,
-    prometheus::{IntGaugeVec, Opts, Registry},
+    prometheus::{IntCounterVec, Opts, Registry},
     solana_client::{
         nonblocking::rpc_client::RpcClient, rpc_config::RpcTransactionConfig,
         rpc_request::RpcRequest,
@@ -34,7 +34,7 @@ use {
 };
 
 lazy_static::lazy_static! {
-    pub static ref TXN_FORWARDER_SENT: IntGaugeVec = IntGaugeVec::new(
+    pub static ref TXN_FORWARDER_SENT: IntCounterVec = IntCounterVec::new(
         Opts::new("txn_forwarder_sent", "Number of sent transactions"),
         &["tree"]
     ).unwrap();