xds_server/
health.rs

1//! Health service for gRPC health checking protocol.
2//!
3//! This module provides integration with `tonic-health` for the standard
4//! gRPC health checking protocol (grpc.health.v1.Health).
5//!
6//! # Example
7//!
8//! ```rust,no_run
9//! use xds_server::health::HealthService;
10//!
11//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
12//! let (health, health_svc) = HealthService::new();
13//! health.set_serving("xds").await;
14//!
15//! // `health_svc` is a ready-to-use tonic service:
16//! // Server::builder().add_service(health_svc).serve(addr).await?;
17//! # let _ = health_svc;
18//! # Ok(()) }
19//! ```
20
21use std::sync::Arc;
22use tokio::sync::Mutex;
23
24use tonic_health::server::HealthReporter;
25
26/// Health service wrapper for xDS server.
27///
28/// Provides gRPC health checking protocol support with service-level
29/// health status management.
30#[derive(Clone)]
31pub struct HealthService {
32    reporter: Arc<Mutex<HealthReporter>>,
33}
34
35impl std::fmt::Debug for HealthService {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("HealthService").finish()
38    }
39}
40
41impl HealthService {
42    /// Create a new health service and return both the wrapper and the tonic service.
43    ///
44    /// The returned tuple contains:
45    /// - `HealthService` - wrapper for managing health status
46    /// - The health server that can be added to a tonic router
47    pub fn new() -> (Self, tonic_health::pb::health_server::HealthServer<impl tonic_health::pb::health_server::Health>) {
48        let (reporter, service) = tonic_health::server::health_reporter();
49        let wrapper = Self {
50            reporter: Arc::new(Mutex::new(reporter)),
51        };
52        (wrapper, service)
53    }
54
55    /// Set a service as serving (healthy).
56    pub async fn set_serving(&self, service_name: &str) {
57        let mut reporter = self.reporter.lock().await;
58        reporter
59            .set_serving::<tonic_health::pb::health_server::HealthServer<
60                tonic_health::server::HealthService,
61            >>()
62            .await;
63        // Also set the specific service name
64        reporter
65            .set_service_status(service_name, tonic_health::ServingStatus::Serving)
66            .await;
67    }
68
69    /// Set a service as not serving (unhealthy).
70    pub async fn set_not_serving(&self, service_name: &str) {
71        let mut reporter = self.reporter.lock().await;
72        reporter
73            .set_service_status(service_name, tonic_health::ServingStatus::NotServing)
74            .await;
75    }
76
77    /// Set a service status to unknown.
78    pub async fn set_unknown(&self, service_name: &str) {
79        let mut reporter = self.reporter.lock().await;
80        reporter
81            .set_service_status(service_name, tonic_health::ServingStatus::Unknown)
82            .await;
83    }
84
85    /// Set all xDS services as serving.
86    pub async fn set_all_serving(&self) {
87        self.set_all_status(tonic_health::ServingStatus::Serving).await;
88    }
89
90    /// Set all xDS services as not serving (for graceful shutdown).
91    pub async fn set_all_not_serving(&self) {
92        self.set_all_status(tonic_health::ServingStatus::NotServing).await;
93    }
94
95    /// Set status for all known xDS services.
96    async fn set_all_status(&self, status: tonic_health::ServingStatus) {
97        let mut reporter = self.reporter.lock().await;
98
99        // Set the main server status
100        if status == tonic_health::ServingStatus::Serving {
101            reporter
102                .set_serving::<tonic_health::pb::health_server::HealthServer<
103                    tonic_health::server::HealthService,
104                >>()
105                .await;
106        }
107
108        // Set all xDS services to the requested status
109        for service in Self::xds_service_names() {
110            reporter.set_service_status(service, status).await;
111        }
112    }
113
114    /// Get the list of xDS service names for health checking.
115    ///
116    /// This returns the standard Envoy xDS service names plus a generic "xds" entry.
117    /// These names match the gRPC service names used by Envoy clients.
118    #[inline]
119    pub const fn xds_service_names() -> &'static [&'static str] {
120        &[
121            "envoy.service.discovery.v3.AggregatedDiscoveryService",
122            "envoy.service.cluster.v3.ClusterDiscoveryService",
123            "envoy.service.listener.v3.ListenerDiscoveryService",
124            "envoy.service.route.v3.RouteDiscoveryService",
125            "envoy.service.endpoint.v3.EndpointDiscoveryService",
126            "envoy.service.secret.v3.SecretDiscoveryService",
127            "xds",
128        ]
129    }
130}
131
132/// Health check configuration.
133#[derive(Debug, Clone)]
134pub struct HealthConfig {
135    /// Enable health checking.
136    pub enabled: bool,
137    /// Initial status (serving or not).
138    pub initial_serving: bool,
139}
140
141impl Default for HealthConfig {
142    fn default() -> Self {
143        Self {
144            enabled: true,
145            initial_serving: true,
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[tokio::test]
155    async fn health_service_creation() {
156        let (health, _server) = HealthService::new();
157        // Should not panic
158        health.set_all_serving().await;
159    }
160
161    #[tokio::test]
162    async fn health_service_status_transitions() {
163        let (health, _server) = HealthService::new();
164
165        health.set_serving("test-service").await;
166        health.set_not_serving("test-service").await;
167        health.set_unknown("test-service").await;
168        // Should not panic
169    }
170
171    #[test]
172    fn health_config_defaults() {
173        let config = HealthConfig::default();
174        assert!(config.enabled);
175        assert!(config.initial_serving);
176    }
177}
xds_server/health.rs

xds_server/
health.rs
