1use std::net::SocketAddr;
2use std::sync::atomic::AtomicUsize;
3use std::sync::{Arc, RwLock};
4use std::time::Instant;
5
6use alopex_cluster::{
7 ClusterManager, ClusterMode, ClusterStatusSnapshot, MembershipSource, NodeRole, NodeState,
8 TableLifecycleEffect,
9};
10use alopex_core::kv::any::AnyKV;
11use alopex_core::kv::async_adapter::{AsyncKVStoreAdapter, AsyncKVTransactionAdapter};
12use alopex_core::kv::AsyncKVStore;
13use alopex_core::types::TxnMode;
14use alopex_sql::catalog::{Catalog, CatalogError, PersistentCatalog};
15use alopex_sql::storage::async_storage::AsyncTxnBridge;
16use alopex_sql::storage::erased::ErasedAsyncSqlTransaction;
17use tokio::sync::{broadcast, Semaphore};
18use tracing::info;
19
20use crate::audit::AuditLogger;
21use crate::auth::AuthMiddleware;
22use crate::config::ServerConfig;
23use crate::error::{Result, ServerError};
24use crate::metrics::Metrics;
25use crate::ops::backup::BackupCoordinator;
26use crate::ops::memory::MemoryControlPolicy;
27use crate::ops::recovery::{RecoveryCoordinator, RecoveryInfo};
28use crate::ops::restore::RestoreCoordinator;
29use crate::ops::state::{LifecycleStateManager, Mode};
30use crate::session::{CatalogRollbackEffect, SessionConfig, SessionManager, TransactionFactory};
31use crate::tls;
32
33pub struct Server {
34 pub state: Arc<ServerState>,
35}
36
37pub struct ServerState {
38 pub config: ServerConfig,
39 pub cluster_manager: Arc<RwLock<ClusterManager>>,
40 pub store: Arc<AnyKV>,
41 pub catalog: Arc<RwLock<dyn Catalog + Send + Sync>>,
42 pub async_store: Arc<AsyncKVStoreAdapter<AnyKV>>,
43 pub session_manager: Arc<SessionManager>,
44 pub metrics: Metrics,
45 pub audit: AuditLogger,
46 pub auth: AuthMiddleware,
47 pub start_time: Instant,
48 pub lifecycle_state: Arc<LifecycleStateManager>,
49 pub recovery_info: RecoveryInfo,
50 pub backup_coordinator: BackupCoordinator,
51 pub restore_coordinator: RestoreCoordinator,
52 pub admission_permits: Arc<Semaphore>,
53 pub admission_waiters: AtomicUsize,
54}
55
56#[derive(Clone, Debug, PartialEq, Eq)]
57pub struct ClusterStartupDiagnostics {
58 pub metadata_schema_version: u32,
59 pub mode: ClusterMode,
60 pub node_id: String,
61 pub cluster_id: Option<String>,
62 pub advertised_endpoint: Option<String>,
63 pub role: NodeRole,
64 pub lifecycle_state: NodeState,
65 pub membership_source: MembershipSource,
66 pub degraded: bool,
67 pub http_bind: SocketAddr,
68 pub grpc_bind: SocketAddr,
69 pub admin_bind: SocketAddr,
70}
71
72impl Server {
73 pub fn new(config: ServerConfig) -> Result<Self> {
74 config.validate()?;
75 let cluster_manager = Arc::new(RwLock::new(
76 ClusterManager::new(config.cluster_manager_config()?)
77 .map_err(|err| ServerError::InvalidConfig(err.to_string()))?,
78 ));
79 let (store, recovery_info) = RecoveryCoordinator::open_store(&config.data_dir)?;
80 let lifecycle_state = Arc::new(LifecycleStateManager::new(Mode::Normal));
81 RecoveryCoordinator::apply_initial_mode(&lifecycle_state, &recovery_info);
82
83 let store = Arc::new(store);
84 let catalog = load_catalog(store.clone())?;
85 let async_store = Arc::new(AsyncKVStoreAdapter::from_arc(
86 store.clone(),
87 TxnMode::ReadWrite,
88 ));
89 let metrics = Metrics::new()?;
90 let audit = AuditLogger::new(config.audit_log_output.clone())?;
91 let auth = AuthMiddleware::new(config.auth_mode.clone());
92
93 let txn_factory = build_txn_factory(async_store.clone(), catalog.clone(), metrics.clone());
94 let session_manager = Arc::new(SessionManager::new(
95 SessionConfig {
96 ttl: config.session_ttl,
97 },
98 txn_factory,
99 ));
100 let data_dir = config.data_dir.clone();
101 let checkpoint = {
102 let store = store.clone();
103 Arc::new(move || match store.as_ref() {
104 AnyKV::Lsm(kv) => {
105 kv.checkpoint()?;
106 Ok(())
107 }
108 _ => Err(ServerError::BadRequest(
109 "checkpoint unsupported for current storage engine".to_string(),
110 )),
111 })
112 };
113 let backup_coordinator =
114 BackupCoordinator::new(data_dir.clone(), lifecycle_state.clone(), checkpoint);
115 let restore_coordinator = RestoreCoordinator::new(data_dir, lifecycle_state.clone());
116 let admission_permits = Arc::new(Semaphore::new(config.max_concurrency));
117
118 Ok(Self {
119 state: Arc::new(ServerState {
120 config,
121 cluster_manager,
122 store,
123 catalog,
124 async_store,
125 session_manager,
126 metrics,
127 audit,
128 auth,
129 start_time: Instant::now(),
130 lifecycle_state,
131 recovery_info,
132 backup_coordinator,
133 restore_coordinator,
134 admission_permits,
135 admission_waiters: AtomicUsize::new(0),
136 }),
137 })
138 }
139
140 pub async fn run(self) -> Result<()> {
141 if self.state.config.tracing_enabled {
142 init_tracing();
143 }
144 emit_cluster_startup_diagnostics(&self.state)?;
145 info!(
146 max_concurrency = self.state.config.max_concurrency,
147 max_queue_len = self.state.config.max_queue_len,
148 query_timeout_ms = self.state.config.query_timeout.as_millis(),
149 "Admission control configuration applied"
150 );
151
152 let (shutdown_tx, _) = broadcast::channel(2);
153 let http_state = self.state.clone();
154 let admin_state = self.state.clone();
155 let grpc_state = self.state.clone();
156 let cleanup_state = self.state.clone();
157 let http_shutdown = shutdown_tx.subscribe();
158 let admin_shutdown = shutdown_tx.subscribe();
159 let grpc_shutdown = shutdown_tx.subscribe();
160 let cleanup_shutdown = shutdown_tx.subscribe();
161
162 let tls_config = if let Some(tls) = &self.state.config.tls {
163 let config = tls::build_rustls_config(tls)?;
164 Some(axum_server::tls_rustls::RustlsConfig::from_config(config))
165 } else {
166 None
167 };
168
169 let http_task = tokio::spawn(run_http(http_state, tls_config.clone(), http_shutdown));
170 let admin_task = tokio::spawn(run_admin(admin_state, tls_config, admin_shutdown));
171 let grpc_task = tokio::spawn(run_grpc(grpc_state, grpc_shutdown));
172 let cleanup_task = tokio::spawn(run_cleanup(cleanup_state, cleanup_shutdown));
173
174 wait_for_shutdown(shutdown_tx.clone()).await;
175 let _ = shutdown_tx.send(());
176
177 http_task
178 .await
179 .map_err(|err| ServerError::Internal(err.to_string()))??;
180 admin_task
181 .await
182 .map_err(|err| ServerError::Internal(err.to_string()))??;
183 grpc_task
184 .await
185 .map_err(|err| ServerError::Internal(err.to_string()))??;
186 cleanup_task
187 .await
188 .map_err(|err| ServerError::Internal(err.to_string()))??;
189
190 self.state.audit.flush()?;
191 Ok(())
192 }
193}
194
195impl ServerState {
196 pub fn cluster_status_snapshot(&self) -> Result<ClusterStatusSnapshot> {
197 let snapshot = self
198 .cluster_manager
199 .read()
200 .map_err(|err| ServerError::Internal(format!("cluster manager lock poisoned: {err}")))?
201 .status_snapshot();
202 Ok(snapshot)
203 }
204
205 pub fn cluster_join(&self) -> Result<ClusterStatusSnapshot> {
206 self.cluster_membership_transition("join")
207 }
208
209 pub fn cluster_leave(&self) -> Result<ClusterStatusSnapshot> {
210 self.cluster_membership_transition("leave")
211 }
212
213 fn cluster_membership_transition(&self, action: &str) -> Result<ClusterStatusSnapshot> {
214 let mut manager = self.cluster_manager.write().map_err(|err| {
215 ServerError::Internal(format!("cluster manager lock poisoned: {err}"))
216 })?;
217 if manager.status_snapshot().mode == ClusterMode::SingleNode {
218 return Err(ServerError::BadRequest(format!(
219 "cluster {action} requires cluster_aware mode"
220 )));
221 }
222 let snapshot = match action {
223 "join" => manager.join(),
224 "leave" => manager.leave(),
225 _ => unreachable!("cluster membership action is fixed by caller"),
226 }
227 .map_err(|err| ServerError::BadRequest(err.to_string()))?;
228 Ok(snapshot)
229 }
230
231 pub fn cluster_startup_diagnostics(&self) -> Result<ClusterStartupDiagnostics> {
232 ClusterStartupDiagnostics::from_state(self)
233 }
234
235 pub fn apply_table_lifecycle_effects(&self, effects: Vec<TableLifecycleEffect>) -> Result<()> {
236 if effects.is_empty() {
237 return Ok(());
238 }
239 let mut manager = self.cluster_manager.write().map_err(|err| {
240 ServerError::Internal(format!("cluster manager lock poisoned: {err}"))
241 })?;
242 for effect in effects {
243 manager
244 .apply_table_lifecycle_effect(effect)
245 .map_err(|err| ServerError::Internal(err.to_string()))?;
246 }
247 Ok(())
248 }
249
250 pub fn apply_catalog_rollback_effects(
251 &self,
252 effects: Vec<CatalogRollbackEffect>,
253 ) -> Result<()> {
254 if effects.is_empty() {
255 return Ok(());
256 }
257 let mut catalog = self
258 .catalog
259 .write()
260 .map_err(|_| ServerError::Internal("catalog lock poisoned".into()))?;
261 for effect in effects.into_iter().rev() {
262 match effect {
263 CatalogRollbackEffect::DropTable { table_name } => {
264 if catalog.table_exists(&table_name) {
265 catalog
266 .drop_table(&table_name)
267 .map_err(|err| ServerError::Internal(err.to_string()))?;
268 }
269 }
270 CatalogRollbackEffect::CreateTable { table } => {
271 if !catalog.table_exists(&table.name) {
272 catalog
273 .create_table(*table)
274 .map_err(|err| ServerError::Internal(err.to_string()))?;
275 }
276 }
277 }
278 }
279 Ok(())
280 }
281
282 pub async fn begin_sql_txn(
283 &self,
284 ) -> Result<AsyncTxnBridge<'static, AsyncKVTransactionAdapter>> {
285 let txn = self.async_store.begin_async().await?;
286 let mut bridge =
287 AsyncTxnBridge::with_catalog(txn, TxnMode::ReadWrite, self.catalog.clone());
288 let policy = MemoryControlPolicy::from_env_with_metrics(self.metrics.clone()).sql_policy();
289 bridge.set_memory_policy(policy);
290 Ok(bridge)
291 }
292}
293
294impl ClusterStartupDiagnostics {
295 fn from_state(state: &ServerState) -> Result<Self> {
296 let snapshot = state.cluster_status_snapshot()?;
297 let identity = &snapshot.identity;
298
299 Ok(Self {
300 metadata_schema_version: snapshot.schema_version,
301 mode: snapshot.mode,
302 node_id: identity.node_id.as_str().to_string(),
303 cluster_id: identity
304 .cluster_id
305 .as_ref()
306 .map(|cluster_id| cluster_id.as_str().to_string()),
307 advertised_endpoint: identity
308 .advertised_endpoint
309 .as_ref()
310 .map(|endpoint| endpoint.as_str().to_string()),
311 role: identity.role,
312 lifecycle_state: identity.lifecycle_state,
313 membership_source: snapshot.membership.source,
314 degraded: snapshot.degraded,
315 http_bind: state.config.http_bind,
316 grpc_bind: state.config.grpc_bind,
317 admin_bind: state.config.admin_bind,
318 })
319 }
320}
321
322fn emit_cluster_startup_diagnostics(state: &ServerState) -> Result<()> {
323 let diagnostics = state.cluster_startup_diagnostics()?;
324 info!(
325 metadata_schema_version = diagnostics.metadata_schema_version,
326 cluster_mode = ?diagnostics.mode,
327 node_id = %diagnostics.node_id,
328 cluster_id = ?diagnostics.cluster_id,
329 advertised_endpoint = ?diagnostics.advertised_endpoint,
330 role = ?diagnostics.role,
331 lifecycle_state = ?diagnostics.lifecycle_state,
332 membership_source = ?diagnostics.membership_source,
333 degraded = diagnostics.degraded,
334 http_bind = %diagnostics.http_bind,
335 grpc_bind = %diagnostics.grpc_bind,
336 admin_bind = %diagnostics.admin_bind,
337 "Cluster startup configuration applied"
338 );
339 Ok(())
340}
341
342fn build_txn_factory(
343 store: Arc<AsyncKVStoreAdapter<AnyKV>>,
344 catalog: Arc<RwLock<dyn Catalog + Send + Sync>>,
345 metrics: Metrics,
346) -> TransactionFactory {
347 Arc::new(move || {
348 let store = store.clone();
349 let catalog = catalog.clone();
350 let metrics = metrics.clone();
351 Box::pin(async move {
352 let txn = store.begin_async().await?;
353 let mut bridge: AsyncTxnBridge<'static, AsyncKVTransactionAdapter> =
354 AsyncTxnBridge::with_catalog(txn, TxnMode::ReadWrite, catalog);
355 let policy = MemoryControlPolicy::from_env_with_metrics(metrics).sql_policy();
356 bridge.set_memory_policy(policy);
357 Ok(Box::new(bridge) as Box<dyn ErasedAsyncSqlTransaction>)
358 })
359 })
360}
361
362fn load_catalog(store: Arc<AnyKV>) -> Result<Arc<RwLock<dyn Catalog + Send + Sync>>> {
363 let catalog = match PersistentCatalog::load(store.clone()) {
364 Ok(catalog) => catalog,
365 Err(CatalogError::Kv(alopex_core::Error::NotFound)) => PersistentCatalog::new(store),
366 Err(err) => return Err(ServerError::Catalog(err)),
367 };
368 let catalog: Arc<RwLock<dyn Catalog + Send + Sync>> = Arc::new(RwLock::new(catalog));
369 Ok(catalog)
370}
371
372async fn run_http(
373 state: Arc<ServerState>,
374 tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
375 mut shutdown: broadcast::Receiver<()>,
376) -> Result<()> {
377 let app = crate::http::router(state.clone());
378 let addr = state.config.http_bind;
379
380 if let Some(tls) = tls_config {
381 let handle = axum_server::Handle::new();
382 let shutdown_handle = handle.clone();
383 tokio::spawn(async move {
384 let _ = shutdown.recv().await;
385 shutdown_handle.graceful_shutdown(Some(std::time::Duration::from_secs(10)));
386 });
387 axum_server::bind_rustls(addr, tls)
388 .handle(handle)
389 .serve(app.into_make_service())
390 .await
391 .map_err(|err| ServerError::Internal(err.to_string()))?;
392 } else {
393 let shutdown_signal = async move {
394 let _ = shutdown.recv().await;
395 };
396 let listener = tokio::net::TcpListener::bind(addr)
397 .await
398 .map_err(ServerError::Io)?;
399 axum::serve(listener, app.into_make_service())
400 .with_graceful_shutdown(shutdown_signal)
401 .await
402 .map_err(|err| ServerError::Internal(err.to_string()))?;
403 }
404 Ok(())
405}
406
407async fn run_admin(
408 state: Arc<ServerState>,
409 tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
410 mut shutdown: broadcast::Receiver<()>,
411) -> Result<()> {
412 let app = crate::http::admin_router(state.clone());
413 let addr = state.config.admin_bind;
414
415 if let Some(tls) = tls_config {
416 let handle = axum_server::Handle::new();
417 let shutdown_handle = handle.clone();
418 tokio::spawn(async move {
419 let _ = shutdown.recv().await;
420 shutdown_handle.graceful_shutdown(Some(std::time::Duration::from_secs(10)));
421 });
422 axum_server::bind_rustls(addr, tls)
423 .handle(handle)
424 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
425 .await
426 .map_err(|err| ServerError::Internal(err.to_string()))?;
427 } else {
428 let shutdown_signal = async move {
429 let _ = shutdown.recv().await;
430 };
431 let listener = tokio::net::TcpListener::bind(addr)
432 .await
433 .map_err(ServerError::Io)?;
434 axum::serve(
435 listener,
436 app.into_make_service_with_connect_info::<SocketAddr>(),
437 )
438 .with_graceful_shutdown(shutdown_signal)
439 .await
440 .map_err(|err| ServerError::Internal(err.to_string()))?;
441 }
442 Ok(())
443}
444
445async fn run_grpc(state: Arc<ServerState>, shutdown: broadcast::Receiver<()>) -> Result<()> {
446 let addr = state.config.grpc_bind;
447 crate::grpc::serve(state, addr, shutdown).await
448}
449
450async fn run_cleanup(state: Arc<ServerState>, mut shutdown: broadcast::Receiver<()>) -> Result<()> {
451 let mut interval = tokio::time::interval(state.config.session_ttl);
452 loop {
453 tokio::select! {
454 _ = interval.tick() => {
455 state.session_manager.cleanup_expired();
456 }
457 _ = shutdown.recv() => break,
458 }
459 }
460 Ok(())
461}
462
463async fn wait_for_shutdown(signal: broadcast::Sender<()>) {
464 #[cfg(unix)]
465 let mut term = match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
466 Ok(signal) => signal,
467 Err(_) => {
468 let _ = tokio::signal::ctrl_c().await;
469 let _ = signal.send(());
470 return;
471 }
472 };
473
474 #[cfg(unix)]
475 tokio::select! {
476 _ = tokio::signal::ctrl_c() => {}
477 _ = term.recv() => {}
478 }
479
480 #[cfg(not(unix))]
481 {
482 let _ = tokio::signal::ctrl_c().await;
483 }
484
485 let _ = signal.send(());
486}
487
488fn init_tracing() {
489 let _ = tracing_subscriber::fmt()
490 .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
491 .try_init();
492}
493
494#[cfg(test)]
495mod tests {
496 use super::*;
497 use crate::config::ClusterServerConfig;
498
499 #[test]
500 fn server_state_exposes_default_single_node_cluster_manager() {
501 let temp = tempfile::tempdir().unwrap();
502 let server = Server::new(ServerConfig {
503 data_dir: temp.path().to_path_buf(),
504 ..ServerConfig::default()
505 })
506 .unwrap();
507
508 let snapshot = server.state.cluster_status_snapshot().unwrap();
509 assert_eq!(snapshot.mode, ClusterMode::SingleNode);
510 assert_eq!(snapshot.identity.node_id.as_str(), "local");
511 assert_eq!(snapshot.identity.lifecycle_state, NodeState::Unconfigured);
512 assert_eq!(snapshot.membership.source, MembershipSource::LocalDefault);
513 assert!(!snapshot.degraded);
514
515 let diagnostics = server.state.cluster_startup_diagnostics().unwrap();
516 assert_eq!(diagnostics.mode, ClusterMode::SingleNode);
517 assert_eq!(diagnostics.node_id, "local");
518 assert_eq!(diagnostics.cluster_id, None);
519 assert_eq!(
520 diagnostics.membership_source,
521 MembershipSource::LocalDefault
522 );
523 assert_eq!(
524 diagnostics.metadata_schema_version,
525 alopex_cluster::CLUSTER_METADATA_SCHEMA_VERSION
526 );
527 assert!(!diagnostics.degraded);
528 }
529
530 #[test]
531 fn server_state_exposes_cluster_aware_configured_identity() {
532 let temp = tempfile::tempdir().unwrap();
533 let server = Server::new(ServerConfig {
534 data_dir: temp.path().to_path_buf(),
535 cluster: ClusterServerConfig {
536 mode: ClusterMode::ClusterAware,
537 node_id: Some("node-a".to_string()),
538 cluster_id: Some("cluster-a".to_string()),
539 advertised_endpoint: Some("127.0.0.1:7001".to_string()),
540 role: NodeRole::Worker,
541 lifecycle_state: NodeState::Active,
542 membership_source_available: false,
543 ..ClusterServerConfig::default()
544 },
545 ..ServerConfig::default()
546 })
547 .unwrap();
548
549 let snapshot = server.state.cluster_status_snapshot().unwrap();
550 assert_eq!(snapshot.mode, ClusterMode::ClusterAware);
551 assert_eq!(snapshot.identity.node_id.as_str(), "node-a");
552 assert_eq!(snapshot.identity.role, NodeRole::Worker);
553 assert_eq!(snapshot.identity.lifecycle_state, NodeState::Active);
554 assert_eq!(snapshot.membership.source, MembershipSource::Chirps);
555 assert_eq!(snapshot.membership.members.len(), 1);
556 assert!(snapshot.degraded);
557
558 let diagnostics = server.state.cluster_startup_diagnostics().unwrap();
559 assert_eq!(diagnostics.mode, ClusterMode::ClusterAware);
560 assert_eq!(diagnostics.node_id, "node-a");
561 assert_eq!(diagnostics.cluster_id.as_deref(), Some("cluster-a"));
562 assert_eq!(
563 diagnostics.advertised_endpoint.as_deref(),
564 Some("127.0.0.1:7001")
565 );
566 assert_eq!(diagnostics.role, NodeRole::Worker);
567 assert_eq!(diagnostics.membership_source, MembershipSource::Chirps);
568 assert!(diagnostics.degraded);
569 }
570
571 #[test]
572 fn server_new_rejects_invalid_cluster_identity_before_store_open() {
573 let result = Server::new(ServerConfig {
574 cluster: ClusterServerConfig {
575 mode: ClusterMode::ClusterAware,
576 node_id: Some("node-a".to_string()),
577 cluster_id: Some("cluster-a".to_string()),
578 advertised_endpoint: Some("127.0.0.1:7001".to_string()),
579 lifecycle_state: NodeState::Unconfigured,
580 ..ClusterServerConfig::default()
581 },
582 ..ServerConfig::default()
583 });
584
585 let err = match result {
586 Ok(_) => panic!("invalid cluster identity should fail"),
587 Err(err) => err,
588 };
589 assert!(err.to_string().contains("unconfigured lifecycle state"));
590 }
591}