turbomcp_server/
lifecycle.rs

1//! Server lifecycle management and graceful shutdown
2
3use std::sync::Arc;
4use tokio::sync::{RwLock, broadcast};
5use tokio::time::Instant;
6
7/// Server lifecycle manager
8#[derive(Debug)]
9pub struct ServerLifecycle {
10    /// Current server state
11    state: Arc<RwLock<ServerState>>,
12    /// Shutdown signal broadcaster
13    shutdown_tx: broadcast::Sender<()>,
14    /// Health status
15    health: Arc<RwLock<HealthStatus>>,
16}
17
18/// Server states
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ServerState {
21    /// Server is starting up
22    Starting,
23    /// Server is running normally
24    Running,
25    /// Server is shutting down
26    ShuttingDown,
27    /// Server has stopped
28    Stopped,
29}
30
31/// Health status information
32#[derive(Debug, Clone)]
33pub struct HealthStatus {
34    /// Overall health
35    pub healthy: bool,
36    /// Health check timestamp
37    pub timestamp: Instant,
38    /// Health details
39    pub details: Vec<HealthCheck>,
40}
41
42/// Individual health check
43#[derive(Debug, Clone)]
44pub struct HealthCheck {
45    /// Check name
46    pub name: String,
47    /// Check status
48    pub healthy: bool,
49    /// Check message
50    pub message: Option<String>,
51    /// Check timestamp
52    pub timestamp: Instant,
53}
54
55/// Shutdown signal
56pub type ShutdownSignal = broadcast::Receiver<()>;
57
58impl ServerLifecycle {
59    /// Create a new lifecycle manager
60    #[must_use]
61    pub fn new() -> Self {
62        let (shutdown_tx, _) = broadcast::channel(16);
63
64        Self {
65            state: Arc::new(RwLock::new(ServerState::Starting)),
66            shutdown_tx,
67            health: Arc::new(RwLock::new(HealthStatus {
68                healthy: true,
69                timestamp: Instant::now(),
70                details: Vec::new(),
71            })),
72        }
73    }
74
75    /// Get current server state
76    pub async fn state(&self) -> ServerState {
77        *self.state.read().await
78    }
79
80    /// Set server state
81    pub async fn set_state(&self, state: ServerState) {
82        *self.state.write().await = state;
83    }
84
85    /// Start the server
86    pub async fn start(&self) {
87        self.set_state(ServerState::Running).await;
88        tracing::info!("Server started");
89    }
90
91    /// Initiate graceful shutdown
92    pub async fn shutdown(&self) {
93        self.set_state(ServerState::ShuttingDown).await;
94        let _ = self.shutdown_tx.send(());
95        tracing::info!("Server shutdown initiated");
96    }
97
98    /// Subscribe to shutdown signals
99    #[must_use]
100    pub fn shutdown_signal(&self) -> ShutdownSignal {
101        self.shutdown_tx.subscribe()
102    }
103
104    /// Get health status
105    pub async fn health(&self) -> HealthStatus {
106        self.health.read().await.clone()
107    }
108
109    /// Update health status
110    pub async fn update_health(&self, healthy: bool, details: Vec<HealthCheck>) {
111        let mut health = self.health.write().await;
112        health.healthy = healthy;
113        health.timestamp = Instant::now();
114        health.details = details;
115    }
116
117    /// Add health check
118    pub async fn add_health_check(&self, check: HealthCheck) {
119        let mut health = self.health.write().await;
120        health.details.push(check);
121        health.healthy = health.details.iter().all(|c| c.healthy);
122        health.timestamp = Instant::now();
123    }
124}
125
126impl Default for ServerLifecycle {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl HealthStatus {
133    /// Create a healthy status
134    #[must_use]
135    pub fn healthy() -> Self {
136        Self {
137            healthy: true,
138            timestamp: Instant::now(),
139            details: Vec::new(),
140        }
141    }
142
143    /// Create an unhealthy status
144    #[must_use]
145    pub fn unhealthy() -> Self {
146        Self {
147            healthy: false,
148            timestamp: Instant::now(),
149            details: Vec::new(),
150        }
151    }
152}
153
154impl HealthCheck {
155    /// Create a healthy check
156    pub fn healthy(name: impl Into<String>) -> Self {
157        Self {
158            name: name.into(),
159            healthy: true,
160            message: None,
161            timestamp: Instant::now(),
162        }
163    }
164
165    /// Create an unhealthy check
166    pub fn unhealthy(name: impl Into<String>, message: impl Into<String>) -> Self {
167        Self {
168            name: name.into(),
169            healthy: false,
170            message: Some(message.into()),
171            timestamp: Instant::now(),
172        }
173    }
174}
175
176// Comprehensive tests in separate file (tokio/axum pattern)
177#[cfg(test)]
178mod tests;