dwctl 8.38.2

The Doubleword Control Layer - A self-hostable observability and analytics platform for LLM applications
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! Database pool metrics sampling.
//!
//! Provides a background task that periodically samples database pool state
//! and records metrics for observability.

use std::time::Duration;

use metrics::gauge;
use sqlx::PgPool;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info};

/// Configuration for pool metrics sampling
#[derive(Debug, Clone)]
pub struct PoolMetricsConfig {
    /// How often to sample pool metrics
    pub sample_interval: Duration,
}

impl Default for PoolMetricsConfig {
    fn default() -> Self {
        Self {
            sample_interval: Duration::from_secs(5),
        }
    }
}

/// A named pool for metrics labeling
pub struct LabeledPool {
    pub name: &'static str,
    pub pool: PgPool,
}

/// Start the pool metrics sampler background task.
///
/// This task periodically samples the pool state and records:
/// - `db_pool_connections_total` - Total connections in the pool
/// - `db_pool_connections_idle` - Idle connections available
/// - `db_pool_connections_in_use` - Connections currently in use
/// - `db_pool_connections_max` - Maximum configured connections
///
/// All metrics are labeled with `pool` to distinguish between different pools.
pub async fn run_pool_metrics_sampler(
    pools: Vec<LabeledPool>,
    config: PoolMetricsConfig,
    shutdown: CancellationToken,
) -> anyhow::Result<()> {
    info!(
        "Starting pool metrics sampler for {} pools with {:?} interval",
        pools.len(),
        config.sample_interval
    );

    // Record max connections once at startup (it doesn't change)
    for labeled in &pools {
        let max = labeled.pool.options().get_max_connections();
        gauge!("dwctl_db_pool_connections_max", "pool" => labeled.name.to_string()).set(max as f64);
    }

    let mut interval = tokio::time::interval(config.sample_interval);
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    loop {
        tokio::select! {
            _ = shutdown.cancelled() => {
                info!("Pool metrics sampler shutting down");
                break;
            }
            _ = interval.tick() => {
                for labeled in &pools {
                    let size = labeled.pool.size();
                    let idle = labeled.pool.num_idle();
                    let in_use = size as usize - idle;

                    gauge!("dwctl_db_pool_connections_total", "pool" => labeled.name.to_string())
                        .set(size as f64);
                    gauge!("dwctl_db_pool_connections_idle", "pool" => labeled.name.to_string())
                        .set(idle as f64);
                    gauge!("dwctl_db_pool_connections_in_use", "pool" => labeled.name.to_string())
                        .set(in_use as f64);

                    debug!(
                        pool = labeled.name,
                        size,
                        idle,
                        in_use,
                        "Sampled pool metrics"
                    );
                }
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[sqlx::test]
    async fn test_pool_metrics_sampler_runs_and_shuts_down(pool: PgPool) {
        let shutdown = CancellationToken::new();
        let shutdown_clone = shutdown.clone();

        let pools = vec![LabeledPool {
            name: "test",
            pool: pool.clone(),
        }];

        let config = PoolMetricsConfig {
            sample_interval: Duration::from_millis(10),
        };

        // Spawn the sampler
        let handle = tokio::spawn(async move { run_pool_metrics_sampler(pools, config, shutdown_clone).await });

        // Let it run for a few samples
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Verify task is still running
        assert!(!handle.is_finished(), "Sampler should still be running");

        // Signal shutdown
        shutdown.cancel();

        // Should complete gracefully
        let result = handle.await.unwrap();
        assert!(result.is_ok(), "Sampler should exit cleanly");
    }

    #[test]
    fn test_pool_metrics_config_default() {
        let config = PoolMetricsConfig::default();
        assert_eq!(config.sample_interval, Duration::from_secs(5));
    }
}