ruvector_metrics/
health.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::time::Instant;
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
6#[serde(rename_all = "lowercase")]
7pub enum HealthStatus {
8    Healthy,
9    Degraded,
10    Unhealthy,
11}
12
13#[derive(Debug, Serialize, Deserialize)]
14pub struct HealthResponse {
15    pub status: HealthStatus,
16    pub version: String,
17    pub uptime_seconds: u64,
18}
19
20#[derive(Debug, Serialize, Deserialize)]
21pub struct ReadinessResponse {
22    pub status: HealthStatus,
23    pub collections_count: usize,
24    pub total_vectors: usize,
25    pub details: HashMap<String, CollectionHealth>,
26}
27
28#[derive(Debug, Serialize, Deserialize)]
29pub struct CollectionHealth {
30    pub status: HealthStatus,
31    pub vectors_count: usize,
32    pub last_updated: Option<String>,
33}
34
35#[derive(Debug)]
36pub struct CollectionStats {
37    pub name: String,
38    pub vectors_count: usize,
39    pub last_updated: Option<chrono::DateTime<chrono::Utc>>,
40}
41
42pub struct HealthChecker {
43    start_time: Instant,
44    version: String,
45}
46
47impl HealthChecker {
48    /// Create a new health checker
49    pub fn new() -> Self {
50        Self {
51            start_time: Instant::now(),
52            version: env!("CARGO_PKG_VERSION").to_string(),
53        }
54    }
55
56    /// Create a health checker with custom version
57    pub fn with_version(version: String) -> Self {
58        Self {
59            start_time: Instant::now(),
60            version,
61        }
62    }
63
64    /// Get basic health status
65    pub fn health(&self) -> HealthResponse {
66        HealthResponse {
67            status: HealthStatus::Healthy,
68            version: self.version.clone(),
69            uptime_seconds: self.start_time.elapsed().as_secs(),
70        }
71    }
72
73    /// Get detailed readiness status
74    pub fn readiness(&self, collections: &[CollectionStats]) -> ReadinessResponse {
75        let total_vectors: usize = collections.iter().map(|c| c.vectors_count).sum();
76
77        let mut details = HashMap::new();
78        for collection in collections {
79            let status = if collection.vectors_count > 0 {
80                HealthStatus::Healthy
81            } else {
82                HealthStatus::Degraded
83            };
84
85            details.insert(
86                collection.name.clone(),
87                CollectionHealth {
88                    status,
89                    vectors_count: collection.vectors_count,
90                    last_updated: collection.last_updated.map(|dt| dt.to_rfc3339()),
91                },
92            );
93        }
94
95        let overall_status = if collections.is_empty() {
96            HealthStatus::Degraded
97        } else if details.values().all(|c| c.status == HealthStatus::Healthy) {
98            HealthStatus::Healthy
99        } else if details.values().any(|c| c.status == HealthStatus::Healthy) {
100            HealthStatus::Degraded
101        } else {
102            HealthStatus::Unhealthy
103        };
104
105        ReadinessResponse {
106            status: overall_status,
107            collections_count: collections.len(),
108            total_vectors,
109            details,
110        }
111    }
112}
113
114impl Default for HealthChecker {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_health_checker_new() {
126        let checker = HealthChecker::new();
127        let health = checker.health();
128
129        assert_eq!(health.status, HealthStatus::Healthy);
130        assert_eq!(health.version, env!("CARGO_PKG_VERSION"));
131        // Uptime is always >= 0 for u64, so just check it exists
132        let _ = health.uptime_seconds;
133    }
134
135    #[test]
136    fn test_readiness_empty_collections() {
137        let checker = HealthChecker::new();
138        let readiness = checker.readiness(&[]);
139
140        assert_eq!(readiness.status, HealthStatus::Degraded);
141        assert_eq!(readiness.collections_count, 0);
142        assert_eq!(readiness.total_vectors, 0);
143    }
144
145    #[test]
146    fn test_readiness_with_collections() {
147        let checker = HealthChecker::new();
148        let collections = vec![
149            CollectionStats {
150                name: "test1".to_string(),
151                vectors_count: 100,
152                last_updated: Some(chrono::Utc::now()),
153            },
154            CollectionStats {
155                name: "test2".to_string(),
156                vectors_count: 200,
157                last_updated: None,
158            },
159        ];
160
161        let readiness = checker.readiness(&collections);
162
163        assert_eq!(readiness.status, HealthStatus::Healthy);
164        assert_eq!(readiness.collections_count, 2);
165        assert_eq!(readiness.total_vectors, 300);
166        assert_eq!(readiness.details.len(), 2);
167    }
168
169    #[test]
170    fn test_readiness_with_empty_collection() {
171        let checker = HealthChecker::new();
172        let collections = vec![CollectionStats {
173            name: "empty".to_string(),
174            vectors_count: 0,
175            last_updated: None,
176        }];
177
178        let readiness = checker.readiness(&collections);
179
180        // Collection exists but is empty (degraded), so overall is Unhealthy
181        // since no collections are in healthy state
182        assert_eq!(readiness.status, HealthStatus::Unhealthy);
183        assert_eq!(readiness.collections_count, 1);
184        assert_eq!(readiness.total_vectors, 0);
185    }
186
187    #[test]
188    fn test_collection_health_status() {
189        let checker = HealthChecker::new();
190        let collections = vec![
191            CollectionStats {
192                name: "healthy".to_string(),
193                vectors_count: 100,
194                last_updated: Some(chrono::Utc::now()),
195            },
196            CollectionStats {
197                name: "degraded".to_string(),
198                vectors_count: 0,
199                last_updated: None,
200            },
201        ];
202
203        let readiness = checker.readiness(&collections);
204
205        assert_eq!(
206            readiness.details.get("healthy").unwrap().status,
207            HealthStatus::Healthy
208        );
209        assert_eq!(
210            readiness.details.get("degraded").unwrap().status,
211            HealthStatus::Degraded
212        );
213    }
214}