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