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