1use crate::bus::BusManager;
2use crate::cache::EntityCache;
3use crate::config::ServerConfig;
4use crate::config::TransactionConfig;
5use crate::health::HealthMonitor;
6use crate::http_server::HttpServer;
7use crate::materialized_view::MaterializedViewRegistry;
8use crate::mutation_batch::MutationBatch;
9use crate::program_runtime::ProgramRuntimeCatalog;
10use crate::projector::Projector;
11use crate::view::ViewIndex;
12use crate::websocket::client_manager::RateLimitConfig;
13use crate::websocket::server::ConnectionAcceptor;
14use crate::websocket::WebSocketServer;
15use crate::Spec;
16use crate::WebSocketAuthPlugin;
17use crate::WebSocketUsageEmitter;
18use anyhow::Result;
19use std::net::SocketAddr;
20use std::sync::Arc;
21use std::time::Duration;
22use tokio::net::{TcpListener, TcpStream};
23use tokio::sync::mpsc;
24use tokio::task::JoinHandle;
25use tokio_util::sync::CancellationToken;
26use tracing::{error, info, info_span, warn, Instrument};
27
28#[cfg(feature = "otel")]
29use crate::metrics::Metrics;
30
31async fn shutdown_signal() {
33 let ctrl_c = async {
34 tokio::signal::ctrl_c()
35 .await
36 .expect("Failed to install Ctrl+C handler");
37 };
38
39 #[cfg(unix)]
40 let terminate = async {
41 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
42 .expect("Failed to install SIGTERM handler")
43 .recv()
44 .await;
45 };
46
47 #[cfg(not(unix))]
48 let terminate = std::future::pending::<()>();
49
50 tokio::select! {
51 _ = ctrl_c => {
52 info!("Received SIGINT (Ctrl+C), initiating shutdown");
53 }
54 _ = terminate => {
55 info!("Received SIGTERM, initiating graceful shutdown");
56 }
57 }
58}
59
60pub struct Runtime {
61 config: ServerConfig,
62 view_index: Arc<ViewIndex>,
63 spec: Option<Spec>,
64 program_runtime_catalog: ProgramRuntimeCatalog,
65 materialized_views: Option<MaterializedViewRegistry>,
66 websocket_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
67 http_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
68 websocket_usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
69 websocket_max_clients: Option<usize>,
70 websocket_rate_limit_config: Option<RateLimitConfig>,
71 #[cfg(feature = "otel")]
72 metrics: Option<Arc<Metrics>>,
73}
74
75impl Runtime {
76 #[cfg(feature = "otel")]
77 pub fn new(config: ServerConfig, view_index: ViewIndex, metrics: Option<Arc<Metrics>>) -> Self {
78 Self {
79 config,
80 view_index: Arc::new(view_index),
81 spec: None,
82 program_runtime_catalog: ProgramRuntimeCatalog::default(),
83 materialized_views: None,
84 websocket_auth_plugin: None,
85 http_auth_plugin: None,
86 websocket_usage_emitter: None,
87 websocket_max_clients: None,
88 websocket_rate_limit_config: None,
89 metrics,
90 }
91 }
92
93 #[cfg(not(feature = "otel"))]
94 pub fn new(config: ServerConfig, view_index: ViewIndex) -> Self {
95 Self {
96 config,
97 view_index: Arc::new(view_index),
98 spec: None,
99 program_runtime_catalog: ProgramRuntimeCatalog::default(),
100 materialized_views: None,
101 websocket_auth_plugin: None,
102 http_auth_plugin: None,
103 websocket_usage_emitter: None,
104 websocket_max_clients: None,
105 websocket_rate_limit_config: None,
106 }
107 }
108
109 pub fn with_spec(mut self, spec: Spec) -> Result<Self> {
110 self.program_runtime_catalog =
111 ProgramRuntimeCatalog::try_new(spec.program_runtime_definitions.clone())?;
112 self.spec = Some(spec);
113 Ok(self)
114 }
115
116 pub fn with_materialized_views(mut self, registry: MaterializedViewRegistry) -> Self {
117 self.materialized_views = Some(registry);
118 self
119 }
120
121 pub fn with_websocket_auth_plugin(
122 mut self,
123 websocket_auth_plugin: Arc<dyn WebSocketAuthPlugin>,
124 ) -> Self {
125 self.websocket_auth_plugin = Some(websocket_auth_plugin);
126 self
127 }
128
129 pub fn with_http_auth_plugin(mut self, http_auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
130 self.http_auth_plugin = Some(http_auth_plugin);
131 self
132 }
133
134 pub fn with_websocket_usage_emitter(
135 mut self,
136 websocket_usage_emitter: Arc<dyn WebSocketUsageEmitter>,
137 ) -> Self {
138 self.websocket_usage_emitter = Some(websocket_usage_emitter);
139 self
140 }
141
142 pub fn with_websocket_max_clients(mut self, websocket_max_clients: usize) -> Self {
143 self.websocket_max_clients = Some(websocket_max_clients);
144 self
145 }
146
147 pub fn with_websocket_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
153 self.websocket_rate_limit_config = Some(config);
154 self
155 }
156
157 pub fn plan(&self) -> crate::RuntimePlan {
159 self.config.runtime_plan
160 }
161
162 pub async fn run(self) -> Result<()> {
169 let mut handle = self.spawn().await?;
170 info!("Arete runtime is running. Press Ctrl+C to stop.");
171
172 tokio::select! {
173 _ = handle.exited() => {}
174 _ = shutdown_signal() => {}
175 }
176
177 handle.shutdown().await
178 }
179
180 pub async fn spawn(self) -> Result<RuntimeHandle> {
192 info!("Starting Arete runtime");
193
194 let plan = self.config.runtime_plan;
195 let transaction_config = if plan.transactions {
196 match self.config.transactions.clone() {
197 Some(config) => config,
198 None => TransactionConfig::from_env()?,
199 }
200 } else {
201 TransactionConfig::default()
202 };
203 if plan.transactions && !transaction_config.enabled {
204 anyhow::bail!(
205 "the runtime plan enables transactions but transaction configuration is disabled"
206 );
207 }
208 let program_runtime_catalog = self.program_runtime_catalog.clone();
209
210 let health_monitor = if plan.health {
211 self.config
212 .health
213 .as_ref()
214 .map(|health_config| HealthMonitor::new(health_config.clone()))
215 } else {
216 None
217 };
218 let mut background = Vec::new();
219 if let Some(monitor) = &health_monitor {
220 background.push(monitor.start().await);
221 info!("Health monitoring enabled");
222 }
223
224 let mut projector_handle = None;
225 let mut ws_handle = None;
226 let mut parser_handle = None;
227 let mut mutations_tx_guard = None;
228 let mut snapshot_service: Option<Arc<crate::snapshot::SnapshotService>> = None;
229 let mut snapshot_manager_handle = None;
230 let mut snapshot_runtime = None;
231 let mut acceptor = None;
232 let mut entity_cache_handle = None;
233
234 if plan.live_runtime_enabled() {
235 let (mutations_tx, mutations_rx) = mpsc::channel::<MutationBatch>(1024);
236 mutations_tx_guard = Some(mutations_tx.clone());
237 let bus_manager = BusManager::new();
238 let entity_cache = EntityCache::new();
239 entity_cache_handle = Some(entity_cache.clone());
240
241 let journal_config = match self.config.journal.clone() {
247 Some(config) => config,
248 None => match crate::journal::JournalConfig::from_env() {
249 Ok(config) => config,
250 Err(e) => {
251 error!("Invalid journal configuration; event replay disabled: {e:#}");
252 crate::journal::JournalConfig::default()
253 }
254 },
255 };
256 let journal = Arc::new(crate::journal::EventJournal::new(journal_config));
257 if journal.is_enabled() {
258 info!(
259 max_bytes_per_view = journal.config().max_bytes_per_view,
260 max_records_per_view = journal.config().max_records_per_view,
261 max_age_secs = journal.config().max_age.as_secs(),
262 "Event replay enabled for append views"
263 );
264 }
265
266 if let Some(spec) = self.spec.as_ref() {
271 let snapshot_config = match self.config.snapshots.clone() {
272 Some(config) => Some(config),
273 None => match crate::snapshot::SnapshotConfig::from_env() {
274 Ok(config) => Some(config),
275 Err(e) => {
276 error!("Invalid snapshot configuration; snapshots disabled: {e:#}");
277 None
278 }
279 },
280 };
281 if let Some(snapshot_config) = snapshot_config.filter(|c| c.enabled) {
282 match crate::snapshot::SnapshotService::initialize(
283 snapshot_config,
284 spec,
285 entity_cache.clone(),
286 &self.view_index,
287 journal.clone(),
288 mutations_tx.clone(),
289 )
290 .await
291 {
292 Ok(service) => {
293 snapshot_runtime = Some(service.runtime());
294 snapshot_manager_handle = Some(service.spawn());
295 snapshot_service = Some(service);
296 }
297 Err(e) => {
298 error!("Failed to initialize snapshots; continuing without: {e:#}")
299 }
300 }
301 }
302 }
303
304 #[cfg(feature = "otel")]
305 let projector = Projector::new(
306 self.view_index.clone(),
307 bus_manager.clone(),
308 entity_cache.clone(),
309 mutations_rx,
310 self.metrics.clone(),
311 );
312 #[cfg(not(feature = "otel"))]
313 let projector = Projector::new(
314 self.view_index.clone(),
315 bus_manager.clone(),
316 entity_cache.clone(),
317 mutations_rx,
318 );
319 let projector = match snapshot_runtime.clone() {
320 Some(runtime) => projector.with_snapshot_runtime(runtime),
321 None => projector,
322 };
323 let projector = projector.with_journal(journal.clone());
324
325 projector_handle = Some(tokio::spawn(async move {
331 projector.run().await;
332 }));
333
334 let bind_address = self
339 .config
340 .websocket
341 .as_ref()
342 .map(|ws_config| ws_config.bind_address)
343 .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
344 #[cfg(feature = "otel")]
345 let mut ws_server = WebSocketServer::new(
346 bind_address,
347 bus_manager.clone(),
348 entity_cache.clone(),
349 self.view_index.clone(),
350 self.metrics.clone(),
351 );
352 #[cfg(not(feature = "otel"))]
353 let mut ws_server = WebSocketServer::new(
354 bind_address,
355 bus_manager.clone(),
356 entity_cache.clone(),
357 self.view_index.clone(),
358 );
359
360 ws_server = ws_server.with_journal(journal.clone());
361 if let Some(max_clients) = self.websocket_max_clients {
362 ws_server = ws_server.with_max_clients(max_clients);
363 }
364 if let Some(plugin) = self.websocket_auth_plugin.clone() {
365 ws_server = ws_server.with_auth_plugin(plugin);
366 }
367 if let Some(emitter) = self.websocket_usage_emitter.clone() {
368 ws_server = ws_server.with_usage_emitter(emitter);
369 }
370 if let Some(rate_limit_config) = self.websocket_rate_limit_config {
371 ws_server = ws_server.with_rate_limit_config(rate_limit_config);
372 }
373 let (connection_acceptor, cleanup_handle) = ws_server.into_acceptor();
374 background.push(cleanup_handle);
375
376 if plan.websocket && self.config.websocket.is_some() {
377 let listener_acceptor = connection_acceptor.clone();
378 ws_handle = Some(tokio::spawn(
379 async move {
380 info!("Starting WebSocket server on {}", bind_address);
381 let listener = match TcpListener::bind(&bind_address).await {
382 Ok(listener) => listener,
383 Err(e) => {
384 error!("WebSocket server error: {}", e);
385 return;
386 }
387 };
388 if let Err(e) = listener_acceptor.serve_listener(listener).await {
389 error!("WebSocket server error: {}", e);
390 }
391 }
392 .instrument(info_span!("ws.server", %bind_address)),
393 ));
394 }
395 acceptor = Some(connection_acceptor);
396
397 if let Some(spec) = self.spec.as_ref() {
398 if let Some(parser_setup) = spec.parser_setup.clone() {
399 let program_id = spec
400 .program_ids
401 .first()
402 .cloned()
403 .unwrap_or_else(|| "unknown".to_string());
404 info!("Starting parser runtime for program: {}", program_id);
405 let health = health_monitor.clone();
406 let reconnection_config = self.config.reconnection.clone().unwrap_or_default();
407 let parser_snapshot_runtime = snapshot_runtime.clone();
408 let parser_journal = journal.clone();
409 parser_handle = Some(tokio::spawn(
410 async move {
411 let parser = async move {
412 parser_setup(mutations_tx, health, reconnection_config).await
413 };
414 let scoped = async move {
415 match parser_snapshot_runtime {
416 Some(runtime) => runtime.scope(parser).await,
417 None => parser.await,
418 }
419 };
420 let result = parser_journal.scope(scoped).await;
424 if let Err(e) = result {
425 error!("Vixen parser runtime error: {}", e);
426 }
427 }
428 .instrument(info_span!("vixen.parser", %program_id)),
429 ));
430 } else {
431 info!("Spec provided but no parser_setup configured - skipping parser runtime");
432 }
433 } else {
434 info!("No spec provided - running in websocket-only mode");
435 }
436
437 let cleanup_bus = bus_manager.clone();
438 background.push(tokio::spawn(
439 async move {
440 let mut interval = tokio::time::interval(Duration::from_secs(60));
441 loop {
442 interval.tick().await;
443 let state_cleaned = cleanup_bus.cleanup_stale_state_buses().await;
444 let list_cleaned = cleanup_bus.cleanup_stale_list_buses().await;
445 if state_cleaned > 0 || list_cleaned > 0 {
446 let (state_count, list_count) = cleanup_bus.bus_counts().await;
447 info!(
448 "Bus cleanup: removed {} state, {} list buses. Current: {} state, {} list",
449 state_cleaned, list_cleaned, state_count, list_count
450 );
451 }
452 }
453 }
454 .instrument(info_span!("bus.cleanup")),
455 ));
456
457 background.push(tokio::spawn(
458 async move {
459 let mut interval = tokio::time::interval(Duration::from_secs(30));
460 loop {
461 interval.tick().await;
462 let (_state_buses, _list_buses) = bus_manager.bus_counts().await;
463 let _cache_stats = entity_cache.stats().await;
464 }
465 }
466 .instrument(info_span!("stats.reporter")),
467 ));
468 } else {
469 info!(
470 "Live runtime disabled; projection and Yellowstone resources were not initialized"
471 );
472 }
473
474 let http_shutdown = CancellationToken::new();
477 let http_health_thread = if let Some(http_health_config) = &self.config.http_health {
478 let mut http_server = HttpServer::new(http_health_config.bind_address)
479 .with_runtime_plan(plan)
480 .with_program_runtime_catalog(program_runtime_catalog)
481 .with_shutdown(http_shutdown.clone());
482 if let Some(target_id) = self.config.program_read_binding_target_id.clone() {
483 http_server = http_server.with_program_read_binding_target(target_id);
484 }
485 if let Some(target_id) = self.config.solana_gateway_target_id.clone() {
486 http_server = http_server.with_solana_gateway_target(target_id);
487 }
488 if let Some(monitor) = health_monitor.clone() {
489 http_server = http_server.with_health_monitor(monitor);
490 }
491 if let Some(runtime) = snapshot_runtime.clone() {
492 http_server = http_server.with_snapshot_runtime(runtime);
493 }
494 if let Some(plugin) = self
495 .http_auth_plugin
496 .clone()
497 .or_else(|| self.websocket_auth_plugin.clone())
498 {
499 http_server = http_server.with_auth_plugin(plugin);
500 }
501 if plan.transactions && transaction_config.enabled {
502 http_server = http_server.with_transaction_config(transaction_config.clone());
503 }
504 #[cfg(feature = "otel")]
505 {
506 http_server = http_server.with_metrics(self.metrics.clone());
507 }
508
509 let bind_addr = http_health_config.bind_address;
510 let join_handle = std::thread::Builder::new()
511 .name("health-server".into())
512 .spawn(move || {
513 let rt = tokio::runtime::Builder::new_current_thread()
514 .enable_all()
515 .build()
516 .expect("Failed to create health server runtime");
517 rt.block_on(async move {
518 let _span = info_span!("http.health", %bind_addr).entered();
519 if let Err(e) = http_server.start().await {
520 error!("HTTP health server error: {}", e);
521 }
522 });
523 })
524 .expect("Failed to spawn health server thread");
525 info!(
526 "HTTP health server running on dedicated thread at {}",
527 bind_addr
528 );
529 Some(join_handle)
530 } else {
531 None
532 };
533
534 Ok(RuntimeHandle {
535 plan,
536 health_monitor,
537 snapshot_runtime,
538 snapshot_service,
539 snapshot_manager_handle,
540 mutations_tx: mutations_tx_guard,
541 projector_handle,
542 parser_handle,
543 ws_handle,
544 background,
545 acceptor,
546 entity_cache: entity_cache_handle,
547 http_shutdown,
548 http_health_thread,
549 })
550 }
551}
552
553pub struct RuntimeHandle {
561 plan: crate::RuntimePlan,
562 health_monitor: Option<HealthMonitor>,
563 snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
564 snapshot_service: Option<Arc<crate::snapshot::SnapshotService>>,
565 snapshot_manager_handle: Option<JoinHandle<()>>,
566 mutations_tx: Option<mpsc::Sender<MutationBatch>>,
567 projector_handle: Option<JoinHandle<()>>,
568 parser_handle: Option<JoinHandle<()>>,
569 ws_handle: Option<JoinHandle<()>>,
570 background: Vec<JoinHandle<()>>,
571 acceptor: Option<ConnectionAcceptor>,
572 entity_cache: Option<EntityCache>,
573 http_shutdown: CancellationToken,
574 http_health_thread: Option<std::thread::JoinHandle<()>>,
575}
576
577#[derive(Clone)]
581pub struct ConnectionServer(ConnectionAcceptor);
582
583impl ConnectionServer {
584 pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
587 self.0.serve(stream, remote_addr).await
588 }
589
590 pub fn client_count(&self) -> usize {
592 self.0.client_count()
593 }
594}
595
596const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
599
600const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
603
604const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
607
608impl RuntimeHandle {
609 pub fn plan(&self) -> crate::RuntimePlan {
611 self.plan
612 }
613
614 pub async fn is_ready(&self) -> bool {
618 let stream_ready = match self.health_monitor.as_ref() {
619 Some(monitor) => monitor.is_healthy().await,
620 None => true,
621 };
622 let snapshot_ready = self
623 .snapshot_runtime
624 .as_ref()
625 .is_none_or(crate::snapshot::SnapshotRuntime::resume_gate_ready);
626 stream_ready && snapshot_ready
627 }
628
629 pub fn client_count(&self) -> usize {
631 self.acceptor
632 .as_ref()
633 .map(ConnectionAcceptor::client_count)
634 .unwrap_or(0)
635 }
636
637 pub async fn entity_cache_stats(&self) -> Option<crate::cache::CacheStats> {
641 match &self.entity_cache {
642 Some(cache) => Some(cache.stats().await),
643 None => None,
644 }
645 }
646
647 pub async fn serve_connection(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
657 match self.connection_server() {
658 Some(server) => server.serve(stream, remote_addr).await,
659 None => anyhow::bail!("this runtime has no live runtime to serve connections from"),
660 }
661 }
662
663 pub fn connection_server(&self) -> Option<ConnectionServer> {
668 self.acceptor.clone().map(ConnectionServer)
669 }
670
671 pub async fn exited(&mut self) {
674 async fn wait(handle: Option<&mut JoinHandle<()>>) {
675 match handle {
676 Some(handle) => {
677 let _ = handle.await;
678 }
679 None => std::future::pending().await,
680 }
681 }
682
683 tokio::select! {
684 _ = wait(self.ws_handle.as_mut()) => info!("WebSocket server task completed"),
685 _ = wait(self.projector_handle.as_mut()) => info!("Projector task completed"),
686 _ = wait(self.parser_handle.as_mut()) => info!("Parser runtime task completed"),
687 }
688 }
689
690 pub async fn shutdown(mut self) -> Result<()> {
706 if let Some(service) = self.snapshot_service.take() {
709 if let Some(handle) = self.snapshot_manager_handle.take() {
710 handle.abort();
711 }
712 if service.config().snapshot_on_shutdown {
713 info!("Taking final snapshot before shutdown");
714 match tokio::time::timeout(
715 SHUTDOWN_SNAPSHOT_TIMEOUT,
716 service.snapshot_now(crate::snapshot::SnapshotTrigger::Shutdown),
717 )
718 .await
719 {
720 Ok(Ok(_)) => {}
721 Ok(Err(e)) => error!("Shutdown snapshot failed: {e:#}"),
722 Err(_) => error!("Shutdown snapshot timed out"),
723 }
724 }
725 }
726 if let Some(handle) = self.snapshot_manager_handle.take() {
727 handle.abort();
728 }
729
730 if let Some(parser) = self.parser_handle.take() {
731 parser.abort();
732 let _ = parser.await;
733 }
734 if let Some(acceptor) = &self.acceptor {
735 acceptor.shutdown();
736 }
737 if let Some(ws) = self.ws_handle.take() {
738 let _ = ws.await;
739 }
740 if let Some(acceptor) = &self.acceptor {
741 if tokio::time::timeout(SESSION_DRAIN_TIMEOUT, acceptor.wait_for_sessions())
742 .await
743 .is_err()
744 {
745 warn!(
746 "Sessions did not finish within {:?} of shutdown",
747 SESSION_DRAIN_TIMEOUT
748 );
749 }
750 }
751
752 drop(self.mutations_tx.take());
753 if let Some(mut projector) = self.projector_handle.take() {
754 if tokio::time::timeout(PROJECTOR_DRAIN_TIMEOUT, &mut projector)
755 .await
756 .is_err()
757 {
758 warn!(
759 "Projector did not drain within {:?}; aborting it",
760 PROJECTOR_DRAIN_TIMEOUT
761 );
762 projector.abort();
763 let _ = projector.await;
764 }
765 }
766
767 for handle in self.background.drain(..) {
768 handle.abort();
769 }
770
771 self.http_shutdown.cancel();
772 if let Some(thread) = self.http_health_thread.take() {
773 if let Err(e) = tokio::task::spawn_blocking(move || thread.join()).await {
774 error!("Health server thread join failed: {e}");
775 }
776 }
777
778 info!("Shutting down Arete runtime");
779 Ok(())
780 }
781}