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 entity_resync = crate::projector::EntityResync::new();
339 let projector = projector
340 .with_journal(journal.clone())
341 .with_entity_resync(entity_resync.clone());
342
343 projector_handle = Some(tokio::spawn(async move {
349 projector.run().await;
350 }));
351
352 let bind_address = self
357 .config
358 .websocket
359 .as_ref()
360 .map(|ws_config| ws_config.bind_address)
361 .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
362 #[cfg(feature = "otel")]
363 let mut ws_server = WebSocketServer::new(
364 bind_address,
365 bus_manager.clone(),
366 entity_cache.clone(),
367 self.view_index.clone(),
368 self.metrics.clone(),
369 );
370 #[cfg(not(feature = "otel"))]
371 let mut ws_server = WebSocketServer::new(
372 bind_address,
373 bus_manager.clone(),
374 entity_cache.clone(),
375 self.view_index.clone(),
376 );
377
378 ws_server = ws_server.with_journal(journal.clone());
379 ws_server = ws_server.with_delivery_config(websocket_delivery);
380 if let Some(max_clients) = self.websocket_max_clients {
381 ws_server = ws_server.with_max_clients(max_clients);
382 }
383 if let Some(plugin) = self.websocket_auth_plugin.clone() {
384 ws_server = ws_server.with_auth_plugin(plugin);
385 }
386 if let Some(emitter) = self.websocket_usage_emitter.clone() {
387 ws_server = ws_server.with_usage_emitter(emitter);
388 }
389 if let Some(rate_limit_config) = self.websocket_rate_limit_config {
390 ws_server = ws_server.with_rate_limit_config(rate_limit_config);
391 }
392 let (connection_acceptor, cleanup_handle) = ws_server.into_acceptor();
393 background.push(cleanup_handle);
394
395 if plan.websocket && self.config.websocket.is_some() {
396 let listener_acceptor = connection_acceptor.clone();
397 ws_handle = Some(tokio::spawn(
398 async move {
399 info!("Starting WebSocket server on {}", bind_address);
400 let listener = match TcpListener::bind(&bind_address).await {
401 Ok(listener) => listener,
402 Err(e) => {
403 error!("WebSocket server error: {}", e);
404 return;
405 }
406 };
407 if let Err(e) = listener_acceptor.serve_listener(listener).await {
408 error!("WebSocket server error: {}", e);
409 }
410 }
411 .instrument(info_span!("ws.server", %bind_address)),
412 ));
413 }
414 acceptor = Some(connection_acceptor);
415
416 if let Some(spec) = self.spec.as_ref() {
417 if let Some(parser_setup) = spec.parser_setup.clone() {
418 let program_id = spec
419 .program_ids
420 .first()
421 .cloned()
422 .unwrap_or_else(|| "unknown".to_string());
423 info!("Starting parser runtime for program: {}", program_id);
424 let health = health_monitor.clone();
425 let reconnection_config = self.config.reconnection.clone().unwrap_or_default();
426 let parser_snapshot_runtime = snapshot_runtime.clone();
427 let parser_journal = journal.clone();
428 let parser_entity_resync = entity_resync.clone();
429 parser_handle = Some(tokio::spawn(async move {
435 let parser = async move {
436 parser_setup(mutations_tx, health, reconnection_config).await
437 };
438 let scoped = async move {
439 match parser_snapshot_runtime {
440 Some(runtime) => runtime.scope(parser).await,
441 None => parser.await,
442 }
443 };
444 let result = parser_entity_resync
448 .scope(parser_journal.scope(scoped))
449 .await;
450 if let Err(e) = result {
451 error!(%program_id, "Vixen parser runtime error: {}", e);
452 }
453 }));
454 } else {
455 info!("Spec provided but no parser_setup configured - skipping parser runtime");
456 }
457 } else {
458 info!("No spec provided - running in websocket-only mode");
459 }
460
461 let cleanup_bus = bus_manager.clone();
462 background.push(tokio::spawn(
463 async move {
464 let mut interval = tokio::time::interval(Duration::from_secs(60));
465 loop {
466 interval.tick().await;
467 let state_cleaned = cleanup_bus.cleanup_stale_state_buses().await;
468 let list_cleaned = cleanup_bus.cleanup_stale_list_buses().await;
469 if state_cleaned > 0 || list_cleaned > 0 {
470 let (state_count, list_count) = cleanup_bus.bus_counts().await;
471 info!(
472 "Bus cleanup: removed {} state, {} list buses. Current: {} state, {} list",
473 state_cleaned, list_cleaned, state_count, list_count
474 );
475 }
476 }
477 }
478 .instrument(info_span!("bus.cleanup")),
479 ));
480
481 background.push(tokio::spawn(
482 async move {
483 let mut interval = tokio::time::interval(Duration::from_secs(30));
484 loop {
485 interval.tick().await;
486 let (_state_buses, _list_buses) = bus_manager.bus_counts().await;
487 let _cache_stats = entity_cache.stats().await;
488 }
489 }
490 .instrument(info_span!("stats.reporter")),
491 ));
492 } else {
493 info!(
494 "Live runtime disabled; projection and Yellowstone resources were not initialized"
495 );
496 }
497
498 let http_shutdown = CancellationToken::new();
501 let http_health_thread = if let Some(http_health_config) = &self.config.http_health {
502 let mut http_server = HttpServer::new(http_health_config.bind_address)
503 .with_runtime_plan(plan)
504 .with_program_runtime_catalog(program_runtime_catalog)
505 .with_shutdown(http_shutdown.clone());
506 if let Some(target_id) = self.config.program_read_binding_target_id.clone() {
507 http_server = http_server.with_program_read_binding_target(target_id);
508 }
509 if let Some(target_id) = self.config.solana_gateway_target_id.clone() {
510 http_server = http_server.with_solana_gateway_target(target_id);
511 }
512 if let Some(monitor) = health_monitor.clone() {
513 http_server = http_server.with_health_monitor(monitor);
514 }
515 if let Some(runtime) = snapshot_runtime.clone() {
516 http_server = http_server.with_snapshot_runtime(runtime);
517 }
518 if let Some(plugin) = self
519 .http_auth_plugin
520 .clone()
521 .or_else(|| self.websocket_auth_plugin.clone())
522 {
523 http_server = http_server.with_auth_plugin(plugin);
524 }
525 if plan.transactions && transaction_config.enabled {
526 http_server = http_server.with_transaction_config(transaction_config.clone());
527 }
528 #[cfg(feature = "otel")]
529 {
530 http_server = http_server.with_metrics(self.metrics.clone());
531 }
532
533 let bind_addr = http_health_config.bind_address;
534 let join_handle = std::thread::Builder::new()
535 .name("health-server".into())
536 .spawn(move || {
537 let rt = tokio::runtime::Builder::new_current_thread()
538 .enable_all()
539 .build()
540 .expect("Failed to create health server runtime");
541 rt.block_on(async move {
542 let _span = info_span!("http.health", %bind_addr).entered();
543 if let Err(e) = http_server.start().await {
544 error!("HTTP health server error: {}", e);
545 }
546 });
547 })
548 .expect("Failed to spawn health server thread");
549 info!(
550 "HTTP health server running on dedicated thread at {}",
551 bind_addr
552 );
553 Some(join_handle)
554 } else {
555 None
556 };
557
558 Ok(RuntimeHandle {
559 plan,
560 health_monitor,
561 snapshot_runtime,
562 snapshot_service,
563 snapshot_manager_handle,
564 mutations_tx: mutations_tx_guard,
565 projector_handle,
566 parser_handle,
567 ws_handle,
568 background,
569 acceptor,
570 entity_cache: entity_cache_handle,
571 http_shutdown,
572 http_health_thread,
573 })
574 }
575}
576
577pub struct RuntimeHandle {
585 plan: crate::RuntimePlan,
586 health_monitor: Option<HealthMonitor>,
587 snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
588 snapshot_service: Option<Arc<crate::snapshot::SnapshotService>>,
589 snapshot_manager_handle: Option<JoinHandle<()>>,
590 mutations_tx: Option<mpsc::Sender<MutationBatch>>,
591 projector_handle: Option<JoinHandle<()>>,
592 parser_handle: Option<JoinHandle<()>>,
593 ws_handle: Option<JoinHandle<()>>,
594 background: Vec<JoinHandle<()>>,
595 acceptor: Option<ConnectionAcceptor>,
596 entity_cache: Option<EntityCache>,
597 http_shutdown: CancellationToken,
598 http_health_thread: Option<std::thread::JoinHandle<()>>,
599}
600
601#[derive(Clone)]
605pub struct ConnectionServer(ConnectionAcceptor);
606
607impl ConnectionServer {
608 pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
611 self.0.serve(stream, remote_addr).await
612 }
613
614 pub fn client_count(&self) -> usize {
616 self.0.client_count()
617 }
618}
619
620const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
623
624const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
627
628const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
631
632impl RuntimeHandle {
633 pub fn plan(&self) -> crate::RuntimePlan {
635 self.plan
636 }
637
638 pub async fn is_ready(&self) -> bool {
642 let stream_ready = match self.health_monitor.as_ref() {
643 Some(monitor) => monitor.is_healthy().await,
644 None => true,
645 };
646 let snapshot_ready = self
647 .snapshot_runtime
648 .as_ref()
649 .is_none_or(crate::snapshot::SnapshotRuntime::resume_gate_ready);
650 stream_ready && snapshot_ready
651 }
652
653 pub fn client_count(&self) -> usize {
655 self.acceptor
656 .as_ref()
657 .map(ConnectionAcceptor::client_count)
658 .unwrap_or(0)
659 }
660
661 pub async fn entity_cache_stats(&self) -> Option<crate::cache::CacheStats> {
665 match &self.entity_cache {
666 Some(cache) => Some(cache.stats().await),
667 None => None,
668 }
669 }
670
671 pub async fn serve_connection(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
681 match self.connection_server() {
682 Some(server) => server.serve(stream, remote_addr).await,
683 None => anyhow::bail!("this runtime has no live runtime to serve connections from"),
684 }
685 }
686
687 pub fn connection_server(&self) -> Option<ConnectionServer> {
692 self.acceptor.clone().map(ConnectionServer)
693 }
694
695 pub async fn exited(&mut self) {
698 async fn wait(handle: Option<&mut JoinHandle<()>>) {
699 match handle {
700 Some(handle) => {
701 let _ = handle.await;
702 }
703 None => std::future::pending().await,
704 }
705 }
706
707 tokio::select! {
708 _ = wait(self.ws_handle.as_mut()) => info!("WebSocket server task completed"),
709 _ = wait(self.projector_handle.as_mut()) => info!("Projector task completed"),
710 _ = wait(self.parser_handle.as_mut()) => info!("Parser runtime task completed"),
711 }
712 }
713
714 pub async fn shutdown(mut self) -> Result<()> {
730 if let Some(service) = self.snapshot_service.take() {
733 if let Some(handle) = self.snapshot_manager_handle.take() {
734 handle.abort();
735 }
736 if service.config().snapshot_on_shutdown {
737 info!("Taking final snapshot before shutdown");
738 match tokio::time::timeout(
739 SHUTDOWN_SNAPSHOT_TIMEOUT,
740 service.snapshot_now(crate::snapshot::SnapshotTrigger::Shutdown),
741 )
742 .await
743 {
744 Ok(Ok(_)) => {}
745 Ok(Err(e)) => error!("Shutdown snapshot failed: {e:#}"),
746 Err(_) => error!("Shutdown snapshot timed out"),
747 }
748 }
749 }
750 if let Some(handle) = self.snapshot_manager_handle.take() {
751 handle.abort();
752 }
753
754 if let Some(parser) = self.parser_handle.take() {
755 parser.abort();
756 let _ = parser.await;
757 }
758 if let Some(acceptor) = &self.acceptor {
759 acceptor.shutdown();
760 }
761 if let Some(ws) = self.ws_handle.take() {
762 let _ = ws.await;
763 }
764 if let Some(acceptor) = &self.acceptor {
765 if tokio::time::timeout(SESSION_DRAIN_TIMEOUT, acceptor.wait_for_sessions())
766 .await
767 .is_err()
768 {
769 warn!(
770 "Sessions did not finish within {:?} of shutdown",
771 SESSION_DRAIN_TIMEOUT
772 );
773 }
774 }
775
776 drop(self.mutations_tx.take());
777 if let Some(mut projector) = self.projector_handle.take() {
778 if tokio::time::timeout(PROJECTOR_DRAIN_TIMEOUT, &mut projector)
779 .await
780 .is_err()
781 {
782 warn!(
783 "Projector did not drain within {:?}; aborting it",
784 PROJECTOR_DRAIN_TIMEOUT
785 );
786 projector.abort();
787 let _ = projector.await;
788 }
789 }
790
791 for handle in self.background.drain(..) {
792 handle.abort();
793 }
794
795 self.http_shutdown.cancel();
796 if let Some(thread) = self.http_health_thread.take() {
797 if let Err(e) = tokio::task::spawn_blocking(move || thread.join()).await {
798 error!("Health server thread join failed: {e}");
799 }
800 }
801
802 info!("Shutting down Arete runtime");
803 Ok(())
804 }
805}