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