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
233 if plan.live_runtime_enabled() {
234 let (mutations_tx, mutations_rx) = mpsc::channel::<MutationBatch>(1024);
235 mutations_tx_guard = Some(mutations_tx.clone());
236 let bus_manager = BusManager::new();
237 let entity_cache = EntityCache::new();
238
239 if let Some(spec) = self.spec.as_ref() {
244 let snapshot_config = match self.config.snapshots.clone() {
245 Some(config) => Some(config),
246 None => match crate::snapshot::SnapshotConfig::from_env() {
247 Ok(config) => Some(config),
248 Err(e) => {
249 error!("Invalid snapshot configuration; snapshots disabled: {e:#}");
250 None
251 }
252 },
253 };
254 if let Some(snapshot_config) = snapshot_config.filter(|c| c.enabled) {
255 match crate::snapshot::SnapshotService::initialize(
256 snapshot_config,
257 spec,
258 entity_cache.clone(),
259 &self.view_index,
260 mutations_tx.clone(),
261 )
262 .await
263 {
264 Ok(service) => {
265 snapshot_runtime = Some(service.runtime());
266 snapshot_manager_handle = Some(service.spawn());
267 snapshot_service = Some(service);
268 }
269 Err(e) => {
270 error!("Failed to initialize snapshots; continuing without: {e:#}")
271 }
272 }
273 }
274 }
275
276 #[cfg(feature = "otel")]
277 let projector = Projector::new(
278 self.view_index.clone(),
279 bus_manager.clone(),
280 entity_cache.clone(),
281 mutations_rx,
282 self.metrics.clone(),
283 );
284 #[cfg(not(feature = "otel"))]
285 let projector = Projector::new(
286 self.view_index.clone(),
287 bus_manager.clone(),
288 entity_cache.clone(),
289 mutations_rx,
290 );
291 let projector = match snapshot_runtime.clone() {
292 Some(runtime) => projector.with_snapshot_runtime(runtime),
293 None => projector,
294 };
295
296 projector_handle = Some(tokio::spawn(
297 async move {
298 projector.run().await;
299 }
300 .instrument(info_span!("projector")),
301 ));
302
303 let bind_address = self
308 .config
309 .websocket
310 .as_ref()
311 .map(|ws_config| ws_config.bind_address)
312 .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
313 #[cfg(feature = "otel")]
314 let mut ws_server = WebSocketServer::new(
315 bind_address,
316 bus_manager.clone(),
317 entity_cache.clone(),
318 self.view_index.clone(),
319 self.metrics.clone(),
320 );
321 #[cfg(not(feature = "otel"))]
322 let mut ws_server = WebSocketServer::new(
323 bind_address,
324 bus_manager.clone(),
325 entity_cache.clone(),
326 self.view_index.clone(),
327 );
328
329 if let Some(max_clients) = self.websocket_max_clients {
330 ws_server = ws_server.with_max_clients(max_clients);
331 }
332 if let Some(plugin) = self.websocket_auth_plugin.clone() {
333 ws_server = ws_server.with_auth_plugin(plugin);
334 }
335 if let Some(emitter) = self.websocket_usage_emitter.clone() {
336 ws_server = ws_server.with_usage_emitter(emitter);
337 }
338 if let Some(rate_limit_config) = self.websocket_rate_limit_config {
339 ws_server = ws_server.with_rate_limit_config(rate_limit_config);
340 }
341 let (connection_acceptor, cleanup_handle) = ws_server.into_acceptor();
342 background.push(cleanup_handle);
343
344 if plan.websocket && self.config.websocket.is_some() {
345 let listener_acceptor = connection_acceptor.clone();
346 ws_handle = Some(tokio::spawn(
347 async move {
348 info!("Starting WebSocket server on {}", bind_address);
349 let listener = match TcpListener::bind(&bind_address).await {
350 Ok(listener) => listener,
351 Err(e) => {
352 error!("WebSocket server error: {}", e);
353 return;
354 }
355 };
356 if let Err(e) = listener_acceptor.serve_listener(listener).await {
357 error!("WebSocket server error: {}", e);
358 }
359 }
360 .instrument(info_span!("ws.server", %bind_address)),
361 ));
362 }
363 acceptor = Some(connection_acceptor);
364
365 if let Some(spec) = self.spec.as_ref() {
366 if let Some(parser_setup) = spec.parser_setup.clone() {
367 let program_id = spec
368 .program_ids
369 .first()
370 .cloned()
371 .unwrap_or_else(|| "unknown".to_string());
372 info!("Starting parser runtime for program: {}", program_id);
373 let health = health_monitor.clone();
374 let reconnection_config = self.config.reconnection.clone().unwrap_or_default();
375 let parser_snapshot_runtime = snapshot_runtime.clone();
376 parser_handle = Some(tokio::spawn(
377 async move {
378 let parser = async move {
379 parser_setup(mutations_tx, health, reconnection_config).await
380 };
381 let result = match parser_snapshot_runtime {
382 Some(runtime) => runtime.scope(parser).await,
383 None => parser.await,
384 };
385 if let Err(e) = result {
386 error!("Vixen parser runtime error: {}", e);
387 }
388 }
389 .instrument(info_span!("vixen.parser", %program_id)),
390 ));
391 } else {
392 info!("Spec provided but no parser_setup configured - skipping parser runtime");
393 }
394 } else {
395 info!("No spec provided - running in websocket-only mode");
396 }
397
398 let cleanup_bus = bus_manager.clone();
399 background.push(tokio::spawn(
400 async move {
401 let mut interval = tokio::time::interval(Duration::from_secs(60));
402 loop {
403 interval.tick().await;
404 let state_cleaned = cleanup_bus.cleanup_stale_state_buses().await;
405 let list_cleaned = cleanup_bus.cleanup_stale_list_buses().await;
406 if state_cleaned > 0 || list_cleaned > 0 {
407 let (state_count, list_count) = cleanup_bus.bus_counts().await;
408 info!(
409 "Bus cleanup: removed {} state, {} list buses. Current: {} state, {} list",
410 state_cleaned, list_cleaned, state_count, list_count
411 );
412 }
413 }
414 }
415 .instrument(info_span!("bus.cleanup")),
416 ));
417
418 background.push(tokio::spawn(
419 async move {
420 let mut interval = tokio::time::interval(Duration::from_secs(30));
421 loop {
422 interval.tick().await;
423 let (_state_buses, _list_buses) = bus_manager.bus_counts().await;
424 let _cache_stats = entity_cache.stats().await;
425 }
426 }
427 .instrument(info_span!("stats.reporter")),
428 ));
429 } else {
430 info!(
431 "Live runtime disabled; projection and Yellowstone resources were not initialized"
432 );
433 }
434
435 let http_shutdown = CancellationToken::new();
438 let http_health_thread = if let Some(http_health_config) = &self.config.http_health {
439 let mut http_server = HttpServer::new(http_health_config.bind_address)
440 .with_runtime_plan(plan)
441 .with_program_runtime_catalog(program_runtime_catalog)
442 .with_shutdown(http_shutdown.clone());
443 if let Some(target_id) = self.config.program_read_binding_target_id.clone() {
444 http_server = http_server.with_program_read_binding_target(target_id);
445 }
446 if let Some(target_id) = self.config.solana_gateway_target_id.clone() {
447 http_server = http_server.with_solana_gateway_target(target_id);
448 }
449 if let Some(monitor) = health_monitor.clone() {
450 http_server = http_server.with_health_monitor(monitor);
451 }
452 if let Some(runtime) = snapshot_runtime.clone() {
453 http_server = http_server.with_snapshot_runtime(runtime);
454 }
455 if let Some(plugin) = self
456 .http_auth_plugin
457 .clone()
458 .or_else(|| self.websocket_auth_plugin.clone())
459 {
460 http_server = http_server.with_auth_plugin(plugin);
461 }
462 if plan.transactions && transaction_config.enabled {
463 http_server = http_server.with_transaction_config(transaction_config.clone());
464 }
465 #[cfg(feature = "otel")]
466 {
467 http_server = http_server.with_metrics(self.metrics.clone());
468 }
469
470 let bind_addr = http_health_config.bind_address;
471 let join_handle = std::thread::Builder::new()
472 .name("health-server".into())
473 .spawn(move || {
474 let rt = tokio::runtime::Builder::new_current_thread()
475 .enable_all()
476 .build()
477 .expect("Failed to create health server runtime");
478 rt.block_on(async move {
479 let _span = info_span!("http.health", %bind_addr).entered();
480 if let Err(e) = http_server.start().await {
481 error!("HTTP health server error: {}", e);
482 }
483 });
484 })
485 .expect("Failed to spawn health server thread");
486 info!(
487 "HTTP health server running on dedicated thread at {}",
488 bind_addr
489 );
490 Some(join_handle)
491 } else {
492 None
493 };
494
495 Ok(RuntimeHandle {
496 plan,
497 health_monitor,
498 snapshot_runtime,
499 snapshot_service,
500 snapshot_manager_handle,
501 mutations_tx: mutations_tx_guard,
502 projector_handle,
503 parser_handle,
504 ws_handle,
505 background,
506 acceptor,
507 http_shutdown,
508 http_health_thread,
509 })
510 }
511}
512
513pub struct RuntimeHandle {
521 plan: crate::RuntimePlan,
522 health_monitor: Option<HealthMonitor>,
523 snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
524 snapshot_service: Option<Arc<crate::snapshot::SnapshotService>>,
525 snapshot_manager_handle: Option<JoinHandle<()>>,
526 mutations_tx: Option<mpsc::Sender<MutationBatch>>,
527 projector_handle: Option<JoinHandle<()>>,
528 parser_handle: Option<JoinHandle<()>>,
529 ws_handle: Option<JoinHandle<()>>,
530 background: Vec<JoinHandle<()>>,
531 acceptor: Option<ConnectionAcceptor>,
532 http_shutdown: CancellationToken,
533 http_health_thread: Option<std::thread::JoinHandle<()>>,
534}
535
536#[derive(Clone)]
540pub struct ConnectionServer(ConnectionAcceptor);
541
542impl ConnectionServer {
543 pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
546 self.0.serve(stream, remote_addr).await
547 }
548
549 pub fn client_count(&self) -> usize {
551 self.0.client_count()
552 }
553}
554
555const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
558
559const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
562
563const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
566
567impl RuntimeHandle {
568 pub fn plan(&self) -> crate::RuntimePlan {
570 self.plan
571 }
572
573 pub async fn is_ready(&self) -> bool {
577 let stream_ready = match self.health_monitor.as_ref() {
578 Some(monitor) => monitor.is_healthy().await,
579 None => true,
580 };
581 let snapshot_ready = self
582 .snapshot_runtime
583 .as_ref()
584 .is_none_or(crate::snapshot::SnapshotRuntime::resume_gate_ready);
585 stream_ready && snapshot_ready
586 }
587
588 pub fn client_count(&self) -> usize {
590 self.acceptor
591 .as_ref()
592 .map(ConnectionAcceptor::client_count)
593 .unwrap_or(0)
594 }
595
596 pub async fn serve_connection(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
606 match self.connection_server() {
607 Some(server) => server.serve(stream, remote_addr).await,
608 None => anyhow::bail!("this runtime has no live runtime to serve connections from"),
609 }
610 }
611
612 pub fn connection_server(&self) -> Option<ConnectionServer> {
617 self.acceptor.clone().map(ConnectionServer)
618 }
619
620 pub async fn exited(&mut self) {
623 async fn wait(handle: Option<&mut JoinHandle<()>>) {
624 match handle {
625 Some(handle) => {
626 let _ = handle.await;
627 }
628 None => std::future::pending().await,
629 }
630 }
631
632 tokio::select! {
633 _ = wait(self.ws_handle.as_mut()) => info!("WebSocket server task completed"),
634 _ = wait(self.projector_handle.as_mut()) => info!("Projector task completed"),
635 _ = wait(self.parser_handle.as_mut()) => info!("Parser runtime task completed"),
636 }
637 }
638
639 pub async fn shutdown(mut self) -> Result<()> {
655 if let Some(service) = self.snapshot_service.take() {
658 if let Some(handle) = self.snapshot_manager_handle.take() {
659 handle.abort();
660 }
661 if service.config().snapshot_on_shutdown {
662 info!("Taking final snapshot before shutdown");
663 match tokio::time::timeout(
664 SHUTDOWN_SNAPSHOT_TIMEOUT,
665 service.snapshot_now(crate::snapshot::SnapshotTrigger::Shutdown),
666 )
667 .await
668 {
669 Ok(Ok(_)) => {}
670 Ok(Err(e)) => error!("Shutdown snapshot failed: {e:#}"),
671 Err(_) => error!("Shutdown snapshot timed out"),
672 }
673 }
674 }
675 if let Some(handle) = self.snapshot_manager_handle.take() {
676 handle.abort();
677 }
678
679 if let Some(parser) = self.parser_handle.take() {
680 parser.abort();
681 let _ = parser.await;
682 }
683 if let Some(acceptor) = &self.acceptor {
684 acceptor.shutdown();
685 }
686 if let Some(ws) = self.ws_handle.take() {
687 let _ = ws.await;
688 }
689 if let Some(acceptor) = &self.acceptor {
690 if tokio::time::timeout(SESSION_DRAIN_TIMEOUT, acceptor.wait_for_sessions())
691 .await
692 .is_err()
693 {
694 warn!(
695 "Sessions did not finish within {:?} of shutdown",
696 SESSION_DRAIN_TIMEOUT
697 );
698 }
699 }
700
701 drop(self.mutations_tx.take());
702 if let Some(mut projector) = self.projector_handle.take() {
703 if tokio::time::timeout(PROJECTOR_DRAIN_TIMEOUT, &mut projector)
704 .await
705 .is_err()
706 {
707 warn!(
708 "Projector did not drain within {:?}; aborting it",
709 PROJECTOR_DRAIN_TIMEOUT
710 );
711 projector.abort();
712 let _ = projector.await;
713 }
714 }
715
716 for handle in self.background.drain(..) {
717 handle.abort();
718 }
719
720 self.http_shutdown.cancel();
721 if let Some(thread) = self.http_health_thread.take() {
722 if let Err(e) = tokio::task::spawn_blocking(move || thread.join()).await {
723 error!("Health server thread join failed: {e}");
724 }
725 }
726
727 info!("Shutting down Arete runtime");
728 Ok(())
729 }
730}