ruvector_metrics/
health.rs1use 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 pub fn new() -> Self {
50 Self {
51 start_time: Instant::now(),
52 version: env!("CARGO_PKG_VERSION").to_string(),
53 }
54 }
55
56 pub fn with_version(version: String) -> Self {
58 Self {
59 start_time: Instant::now(),
60 version,
61 }
62 }
63
64 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 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 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 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}