From 7771275cc6b216ce4ef471002408a6e15538d1bd Mon Sep 17 00:00:00 2001
From: Arthur Petukhovsky <petuhovskiy@yandex.ru>
Date: Fri, 17 May 2024 16:30:27 +0000
Subject: [PATCH] Add metrics

---
 safekeeper/src/metrics.rs          | 28 ++++++++++++++++++++++++++--
 safekeeper/src/timeline_manager.rs | 18 ++++++++++++------
 safekeeper/src/wal_backup.rs       |  4 +++-
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/safekeeper/src/metrics.rs b/safekeeper/src/metrics.rs
index 044386d0a9..1e965393e3 100644
--- a/safekeeper/src/metrics.rs
+++ b/safekeeper/src/metrics.rs
@@ -11,8 +11,9 @@ use futures::Future;
 use metrics::{
     core::{AtomicU64, Collector, Desc, GenericCounter, GenericGaugeVec, Opts},
     proto::MetricFamily,
-    register_int_counter, register_int_counter_pair_vec, register_int_counter_vec, Gauge,
-    IntCounter, IntCounterPairVec, IntCounterVec, IntGaugeVec,
+    register_int_counter, register_int_counter_pair, register_int_counter_pair_vec,
+    register_int_counter_vec, Gauge, IntCounter, IntCounterPair, IntCounterPairVec, IntCounterVec,
+    IntGaugeVec,
 };
 use once_cell::sync::Lazy;
 
@@ -162,6 +163,29 @@ pub static PARTIAL_BACKUP_UPLOADED_BYTES: Lazy<IntCounter> = Lazy::new(|| {
     )
     .expect("Failed to register safekeeper_partial_backup_uploaded_bytes_total counter")
 });
+pub static MANAGER_ITERATIONS_TOTAL: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "safekeeper_manager_iterations_total",
+        "Number of iterations of the timeline manager task"
+    )
+    .expect("Failed to register safekeeper_manager_iterations_total counter")
+});
+pub static MANAGER_ACTIVE_CHANGES: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "safekeeper_manager_active_changes_total",
+        "Number of timeline active status changes in the timeline manager task"
+    )
+    .expect("Failed to register safekeeper_manager_active_changes_total counter")
+});
+pub static WAL_BACKUP_TASKS: Lazy<IntCounterPair> = Lazy::new(|| {
+    register_int_counter_pair!(
+        "safekeeper_wal_backup_tasks_started_total",
+        "Number of active WAL backup tasks",
+        "safekeeper_wal_backup_tasks_finished_total",
+        "Number of finished WAL backup tasks",
+    )
+    .expect("Failed to register safekeeper_wal_backup_tasks_finished_total counter")
+});
 
 pub const LABEL_UNKNOWN: &str = "unknown";
 
diff --git a/safekeeper/src/timeline_manager.rs b/safekeeper/src/timeline_manager.rs
index fc86f655b8..8077d1d30c 100644
--- a/safekeeper/src/timeline_manager.rs
+++ b/safekeeper/src/timeline_manager.rs
@@ -4,6 +4,7 @@ use tracing::{info, instrument, warn};
 use utils::lsn::Lsn;
 
 use crate::{
+    metrics::{MANAGER_ACTIVE_CHANGES, MANAGER_ITERATIONS_TOTAL},
     timeline::{PeerInfo, ReadGuardSharedState, Timeline},
     timelines_set::TimelinesSet,
     wal_backup::{self, WalBackupTaskHandle},
@@ -73,6 +74,8 @@ pub async fn main_task(
     let mut backup_task: Option<WalBackupTaskHandle> = None;
 
     let last_state = 'outer: loop {
+        MANAGER_ITERATIONS_TOTAL.inc();
+
         let state_snapshot = StateSnapshot::new(tli.read_shared_state().await, heartbeat_timeout);
         let num_computes = *num_computes_rx.borrow();
 
@@ -97,10 +100,12 @@ pub async fn main_task(
         if tli_broker_active.set(is_active) {
             // write log if state has changed
             info!(
-                "timeline {} active={} now, remote_consistent_lsn={}, commit_lsn={}",
-                ttid, is_active, state_snapshot.remote_consistent_lsn, state_snapshot.commit_lsn,
+                "timeline active={} now, remote_consistent_lsn={}, commit_lsn={}",
+                is_active, state_snapshot.remote_consistent_lsn, state_snapshot.commit_lsn,
             );
 
+            MANAGER_ACTIVE_CHANGES.inc();
+
             if !is_active {
                 // TODO: maybe use tokio::spawn?
                 if let Err(e) = tli.maybe_persist_control_file().await {
@@ -115,16 +120,17 @@ pub async fn main_task(
         tli.broker_active
             .store(is_active, std::sync::atomic::Ordering::SeqCst);
 
-        // sleep to make the loop less busy
-        tokio::time::sleep(REFRESH_INTERVAL).await;
-
         // wait until something changes
         tokio::select! {
             _ = cancellation_rx.changed() => {
                 // timeline was deleted
                 break 'outer state_snapshot;
             }
-            _ = state_version_rx.changed() => {
+            _ = async {
+                // don't wake up on every state change, but at most every REFRESH_INTERVAL
+                tokio::time::sleep(REFRESH_INTERVAL).await;
+                let _ = state_version_rx.changed().await;
+            } => {
                 // state was updated
             }
             _ = num_computes_rx.changed() => {
diff --git a/safekeeper/src/wal_backup.rs b/safekeeper/src/wal_backup.rs
index 8a5f1b7df7..dead6c44a3 100644
--- a/safekeeper/src/wal_backup.rs
+++ b/safekeeper/src/wal_backup.rs
@@ -29,7 +29,7 @@ use tracing::*;
 
 use utils::{id::TenantTimelineId, lsn::Lsn};
 
-use crate::metrics::{BACKED_UP_SEGMENTS, BACKUP_ERRORS};
+use crate::metrics::{BACKED_UP_SEGMENTS, BACKUP_ERRORS, WAL_BACKUP_TASKS};
 use crate::timeline::{PeerInfo, Timeline};
 use crate::timeline_manager::StateSnapshot;
 use crate::{GlobalTimelines, SafeKeeperConf, WAL_BACKUP_RUNTIME};
@@ -209,6 +209,8 @@ async fn backup_task_main(
     parallel_jobs: usize,
     mut shutdown_rx: Receiver<()>,
 ) {
+    let _guard = WAL_BACKUP_TASKS.guard();
+
     info!("started");
     let res = GlobalTimelines::get(ttid);
     if let Err(e) = res {