Skip to main content

arete_server/
runtime.rs

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::SolanaGatewayUsageObserver;
17use crate::Spec;
18use crate::WebSocketAuthPlugin;
19use crate::WebSocketUsageEmitter;
20use anyhow::Result;
21use std::net::SocketAddr;
22use std::sync::Arc;
23use std::time::Duration;
24use tokio::net::{TcpListener, TcpStream};
25use tokio::sync::mpsc;
26use tokio::task::JoinHandle;
27use tokio_util::sync::CancellationToken;
28use tracing::{error, info, info_span, warn, Instrument};
29
30#[cfg(feature = "otel")]
31use crate::metrics::Metrics;
32
33/// Wait for shutdown signal (SIGINT on all platforms, SIGTERM on Unix)
34async fn shutdown_signal() {
35    let ctrl_c = async {
36        tokio::signal::ctrl_c()
37            .await
38            .expect("Failed to install Ctrl+C handler");
39    };
40
41    #[cfg(unix)]
42    let terminate = async {
43        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
44            .expect("Failed to install SIGTERM handler")
45            .recv()
46            .await;
47    };
48
49    #[cfg(not(unix))]
50    let terminate = std::future::pending::<()>();
51
52    tokio::select! {
53        _ = ctrl_c => {
54            info!("Received SIGINT (Ctrl+C), initiating shutdown");
55        }
56        _ = terminate => {
57            info!("Received SIGTERM, initiating graceful shutdown");
58        }
59    }
60}
61
62pub struct Runtime {
63    config: ServerConfig,
64    view_index: Arc<ViewIndex>,
65    spec: Option<Spec>,
66    program_runtime_catalog: ProgramRuntimeCatalog,
67    materialized_views: Option<MaterializedViewRegistry>,
68    websocket_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
69    http_auth_plugin: Option<Arc<dyn WebSocketAuthPlugin>>,
70    websocket_usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
71    solana_gateway_usage_observer: Option<Arc<dyn SolanaGatewayUsageObserver>>,
72    websocket_max_clients: Option<usize>,
73    websocket_rate_limit_config: Option<RateLimitConfig>,
74    #[cfg(feature = "otel")]
75    metrics: Option<Arc<Metrics>>,
76}
77
78impl Runtime {
79    #[cfg(feature = "otel")]
80    pub fn new(config: ServerConfig, view_index: ViewIndex, metrics: Option<Arc<Metrics>>) -> Self {
81        Self {
82            config,
83            view_index: Arc::new(view_index),
84            spec: None,
85            program_runtime_catalog: ProgramRuntimeCatalog::default(),
86            materialized_views: None,
87            websocket_auth_plugin: None,
88            http_auth_plugin: None,
89            websocket_usage_emitter: None,
90            solana_gateway_usage_observer: None,
91            websocket_max_clients: None,
92            websocket_rate_limit_config: None,
93            metrics,
94        }
95    }
96
97    #[cfg(not(feature = "otel"))]
98    pub fn new(config: ServerConfig, view_index: ViewIndex) -> Self {
99        Self {
100            config,
101            view_index: Arc::new(view_index),
102            spec: None,
103            program_runtime_catalog: ProgramRuntimeCatalog::default(),
104            materialized_views: None,
105            websocket_auth_plugin: None,
106            http_auth_plugin: None,
107            websocket_usage_emitter: None,
108            solana_gateway_usage_observer: None,
109            websocket_max_clients: None,
110            websocket_rate_limit_config: None,
111        }
112    }
113
114    pub fn with_spec(mut self, spec: Spec) -> Result<Self> {
115        self.program_runtime_catalog =
116            ProgramRuntimeCatalog::try_new(spec.program_runtime_definitions.clone())?;
117        self.spec = Some(spec);
118        Ok(self)
119    }
120
121    pub fn with_materialized_views(mut self, registry: MaterializedViewRegistry) -> Self {
122        self.materialized_views = Some(registry);
123        self
124    }
125
126    pub fn with_websocket_auth_plugin(
127        mut self,
128        websocket_auth_plugin: Arc<dyn WebSocketAuthPlugin>,
129    ) -> Self {
130        self.websocket_auth_plugin = Some(websocket_auth_plugin);
131        self
132    }
133
134    pub fn with_http_auth_plugin(mut self, http_auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
135        self.http_auth_plugin = Some(http_auth_plugin);
136        self
137    }
138
139    pub fn with_websocket_usage_emitter(
140        mut self,
141        websocket_usage_emitter: Arc<dyn WebSocketUsageEmitter>,
142    ) -> Self {
143        self.websocket_usage_emitter = Some(websocket_usage_emitter);
144        self
145    }
146
147    pub fn with_solana_gateway_usage_observer(
148        mut self,
149        observer: Arc<dyn SolanaGatewayUsageObserver>,
150    ) -> Self {
151        self.solana_gateway_usage_observer = Some(observer);
152        self
153    }
154
155    pub fn with_websocket_max_clients(mut self, websocket_max_clients: usize) -> Self {
156        self.websocket_max_clients = Some(websocket_max_clients);
157        self
158    }
159
160    /// Configure rate limiting for WebSocket connections.
161    ///
162    /// This sets global rate limits such as maximum connections per IP,
163    /// timeouts, and rate windows. Per-subject limits are controlled
164    /// via AuthContext.Limits from the authentication token.
165    pub fn with_websocket_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
166        self.websocket_rate_limit_config = Some(config);
167        self
168    }
169
170    /// Return the immutable capability plan selected by the builder.
171    pub fn plan(&self) -> crate::RuntimePlan {
172        self.config.runtime_plan
173    }
174
175    /// Start everything and block until a shutdown signal, then stop cleanly.
176    ///
177    /// This is [`spawn`](Self::spawn) followed by waiting for SIGINT/SIGTERM
178    /// (or for a core task to exit) and [`RuntimeHandle::shutdown`]. Callers
179    /// that embed the server in a larger process should use `spawn` directly
180    /// and decide for themselves when to stop.
181    pub async fn run(self) -> Result<()> {
182        let mut handle = self.spawn().await?;
183        info!("Arete runtime is running. Press Ctrl+C to stop.");
184
185        tokio::select! {
186            _ = handle.exited() => {}
187            _ = shutdown_signal() => {}
188        }
189
190        handle.shutdown().await
191    }
192
193    /// Start the runtime's tasks and return a handle that owns them.
194    ///
195    /// Everything `run` starts is started here: projector, parser, snapshot
196    /// manager, bus cleanup, stats, and - only when configured - the WebSocket
197    /// listener and the HTTP health server. Nothing here installs signal
198    /// handlers or blocks; the handle is how the caller waits, serves
199    /// connections it accepted itself, and shuts the runtime down.
200    ///
201    /// A handle owns its runtime's channel, buses, cache and snapshot state
202    /// outright; nothing is shared through process globals, so a test or an
203    /// application can start and stop runtimes independently.
204    pub async fn spawn(self) -> Result<RuntimeHandle> {
205        info!("Starting Arete runtime");
206
207        let plan = self.config.runtime_plan;
208        let transaction_config = if plan.transactions {
209            match self.config.transactions.clone() {
210                Some(config) => config,
211                None => TransactionConfig::from_env()?,
212            }
213        } else {
214            TransactionConfig::default()
215        };
216        if plan.transactions && !transaction_config.enabled {
217            anyhow::bail!(
218                "the runtime plan enables transactions but transaction configuration is disabled"
219            );
220        }
221        let program_runtime_catalog = self.program_runtime_catalog.clone();
222
223        let health_monitor = if plan.health {
224            self.config
225                .health
226                .as_ref()
227                .map(|health_config| HealthMonitor::new(health_config.clone()))
228        } else {
229            None
230        };
231        let mut background = Vec::new();
232        if let Some(monitor) = &health_monitor {
233            background.push(monitor.start().await);
234            info!("Health monitoring enabled");
235        }
236
237        let mut projector_handle = None;
238        let mut ws_handle = None;
239        let mut parser_handle = None;
240        let mut mutations_tx_guard = None;
241        let mut snapshot_service: Option<Arc<crate::snapshot::SnapshotService>> = None;
242        let mut snapshot_manager_handle = None;
243        let mut snapshot_runtime = None;
244        let mut acceptor = None;
245        let mut entity_cache_handle = None;
246
247        if plan.live_runtime_enabled() {
248            let (mutations_tx, mutations_rx) = mpsc::channel::<MutationBatch>(1024);
249            mutations_tx_guard = Some(mutations_tx.clone());
250            let websocket_delivery = match self.config.websocket_delivery.clone() {
251                Some(config) => {
252                    config.validate()?;
253                    config
254                }
255                None => WebSocketDeliveryConfig::from_env()?,
256            };
257            info!(
258                list_bus_capacity = websocket_delivery.list_bus_capacity,
259                collection_coalesce_ms = ?websocket_delivery.collection_coalesce_ms,
260                "WebSocket delivery configured"
261            );
262            let bus_manager = BusManager::with_capacity(websocket_delivery.list_bus_capacity);
263            let entity_cache = EntityCache::new();
264            entity_cache_handle = Some(entity_cache.clone());
265
266            // Retained event tape for replayable append subscriptions. The
267            // builder wins over the process env so one host can enable replay
268            // for a single deployment and size it independently. A bad
269            // configuration disables replay rather than failing startup, the
270            // same posture snapshots take below.
271            let journal_config = match self.config.journal.clone() {
272                Some(config) => config,
273                None => match crate::journal::JournalConfig::from_env() {
274                    Ok(config) => config,
275                    Err(e) => {
276                        error!("Invalid journal configuration; event replay disabled: {e:#}");
277                        crate::journal::JournalConfig::default()
278                    }
279                },
280            };
281            let journal = Arc::new(crate::journal::EventJournal::new(journal_config));
282            if journal.is_enabled() {
283                info!(
284                    max_bytes_per_view = journal.config().max_bytes_per_view,
285                    max_records_per_view = journal.config().max_records_per_view,
286                    max_age_secs = journal.config().max_age.as_secs(),
287                    "Event replay enabled for append views"
288                );
289            }
290
291            // Restore state from the latest snapshot (when enabled) before the
292            // WebSocket server spawns, so the first client's snapshot-on-subscribe
293            // is already warm. The VM portion is stashed for the generated
294            // runtime to hydrate before it connects to Yellowstone.
295            if let Some(spec) = self.spec.as_ref() {
296                let snapshot_config = match self.config.snapshots.clone() {
297                    Some(config) => Some(config),
298                    None => match crate::snapshot::SnapshotConfig::from_env() {
299                        Ok(config) => Some(config),
300                        Err(e) => {
301                            error!("Invalid snapshot configuration; snapshots disabled: {e:#}");
302                            None
303                        }
304                    },
305                };
306                if let Some(snapshot_config) = snapshot_config.filter(|c| c.enabled) {
307                    match crate::snapshot::SnapshotService::initialize(
308                        snapshot_config,
309                        spec,
310                        entity_cache.clone(),
311                        &self.view_index,
312                        journal.clone(),
313                        mutations_tx.clone(),
314                    )
315                    .await
316                    {
317                        Ok(service) => {
318                            snapshot_runtime = Some(service.runtime());
319                            snapshot_manager_handle = Some(service.spawn());
320                            snapshot_service = Some(service);
321                        }
322                        Err(e) => {
323                            error!("Failed to initialize snapshots; continuing without: {e:#}")
324                        }
325                    }
326                }
327            }
328
329            #[cfg(feature = "otel")]
330            let projector = Projector::new(
331                self.view_index.clone(),
332                bus_manager.clone(),
333                entity_cache.clone(),
334                mutations_rx,
335                self.metrics.clone(),
336            );
337            #[cfg(not(feature = "otel"))]
338            let projector = Projector::new(
339                self.view_index.clone(),
340                bus_manager.clone(),
341                entity_cache.clone(),
342                mutations_rx,
343            );
344            let projector = match snapshot_runtime.clone() {
345                Some(runtime) => projector.with_snapshot_runtime(runtime),
346                None => projector,
347            };
348            // The parser's VM links itself to this when it registers, so the
349            // projector can ask it for entities the bounded cache dropped.
350            let entity_resync = crate::projector::EntityResync::new();
351            let projector = projector
352                .with_journal(journal.clone())
353                .with_entity_resync(entity_resync.clone());
354
355            // The projector runs for the lifetime of the server. Giving the
356            // task a span would make that span the parent of every batch
357            // whose producer does not carry an explicit context, creating an
358            // unbounded trace. `Projector::run` instead enters the bounded
359            // batch span before it processes and logs each batch.
360            projector_handle = Some(tokio::spawn(async move {
361                projector.run().await;
362            }));
363
364            // The connection-serving half of the WebSocket server exists
365            // whenever there is a live runtime, so a caller that owns its own
366            // listener can hand in the connections it accepts. The listener
367            // is bound only when a WebSocket address is configured.
368            let bind_address = self
369                .config
370                .websocket
371                .as_ref()
372                .map(|ws_config| ws_config.bind_address)
373                .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
374            #[cfg(feature = "otel")]
375            let mut ws_server = WebSocketServer::new(
376                bind_address,
377                bus_manager.clone(),
378                entity_cache.clone(),
379                self.view_index.clone(),
380                self.metrics.clone(),
381            );
382            #[cfg(not(feature = "otel"))]
383            let mut ws_server = WebSocketServer::new(
384                bind_address,
385                bus_manager.clone(),
386                entity_cache.clone(),
387                self.view_index.clone(),
388            );
389
390            ws_server = ws_server.with_journal(journal.clone());
391            ws_server = ws_server.with_delivery_config(websocket_delivery);
392            if let Some(max_clients) = self.websocket_max_clients {
393                ws_server = ws_server.with_max_clients(max_clients);
394            }
395            if let Some(plugin) = self.websocket_auth_plugin.clone() {
396                ws_server = ws_server.with_auth_plugin(plugin);
397            }
398            if let Some(emitter) = self.websocket_usage_emitter.clone() {
399                ws_server = ws_server.with_usage_emitter(emitter);
400            }
401            if let Some(rate_limit_config) = self.websocket_rate_limit_config {
402                ws_server = ws_server.with_rate_limit_config(rate_limit_config);
403            }
404            let (connection_acceptor, cleanup_handle) = ws_server.into_acceptor();
405            background.push(cleanup_handle);
406
407            if plan.websocket && self.config.websocket.is_some() {
408                let listener_acceptor = connection_acceptor.clone();
409                ws_handle = Some(tokio::spawn(
410                    async move {
411                        info!("Starting WebSocket server on {}", bind_address);
412                        let listener = match TcpListener::bind(&bind_address).await {
413                            Ok(listener) => listener,
414                            Err(e) => {
415                                error!("WebSocket server error: {}", e);
416                                return;
417                            }
418                        };
419                        if let Err(e) = listener_acceptor.serve_listener(listener).await {
420                            error!("WebSocket server error: {}", e);
421                        }
422                    }
423                    .instrument(info_span!("ws.server", %bind_address)),
424                ));
425            }
426            acceptor = Some(connection_acceptor);
427
428            if let Some(spec) = self.spec.as_ref() {
429                if let Some(parser_setup) = spec.parser_setup.clone() {
430                    let program_id = spec
431                        .program_ids
432                        .first()
433                        .cloned()
434                        .unwrap_or_else(|| "unknown".to_string());
435                    info!("Starting parser runtime for program: {}", program_id);
436                    let health = health_monitor.clone();
437                    let reconnection_config = self.config.reconnection.clone().unwrap_or_default();
438                    let parser_snapshot_runtime = snapshot_runtime.clone();
439                    let parser_journal = journal.clone();
440                    let parser_entity_resync = entity_resync.clone();
441                    // The parser runs for the lifetime of the server, like
442                    // the projector. A span on the task would be the current
443                    // span of every update it processes, and OpenTelemetry
444                    // keeps an open span's events until it closes: for a task
445                    // that runs as long as the server, an unbounded buffer.
446                    parser_handle = Some(tokio::spawn(async move {
447                        let parser = async move {
448                            parser_setup(mutations_tx, health, reconnection_config).await
449                        };
450                        let scoped = async move {
451                            match parser_snapshot_runtime {
452                                Some(runtime) => runtime.scope(parser).await,
453                                None => parser.await,
454                            }
455                        };
456                        // The tape is in scope even with snapshots off, so
457                        // a runtime that abandons its checkpoint can still
458                        // mark the hole it just created.
459                        let result = parser_entity_resync
460                            .scope(parser_journal.scope(scoped))
461                            .await;
462                        if let Err(e) = result {
463                            error!(%program_id, "Vixen parser runtime error: {}", e);
464                        }
465                    }));
466                } else {
467                    info!("Spec provided but no parser_setup configured - skipping parser runtime");
468                }
469            } else {
470                info!("No spec provided - running in websocket-only mode");
471            }
472
473            let cleanup_bus = bus_manager.clone();
474            background.push(tokio::spawn(
475                async move {
476                    let mut interval = tokio::time::interval(Duration::from_secs(60));
477                    loop {
478                        interval.tick().await;
479                        let state_cleaned = cleanup_bus.cleanup_stale_state_buses().await;
480                        let list_cleaned = cleanup_bus.cleanup_stale_list_buses().await;
481                        if state_cleaned > 0 || list_cleaned > 0 {
482                            let (state_count, list_count) = cleanup_bus.bus_counts().await;
483                            info!(
484                                "Bus cleanup: removed {} state, {} list buses. Current: {} state, {} list",
485                                state_cleaned, list_cleaned, state_count, list_count
486                            );
487                        }
488                    }
489                }
490                .instrument(info_span!("bus.cleanup")),
491            ));
492
493            background.push(tokio::spawn(
494                async move {
495                    let mut interval = tokio::time::interval(Duration::from_secs(30));
496                    loop {
497                        interval.tick().await;
498                        let (_state_buses, _list_buses) = bus_manager.bus_counts().await;
499                        let _cache_stats = entity_cache.stats().await;
500                    }
501                }
502                .instrument(info_span!("stats.reporter")),
503            ));
504        } else {
505            info!(
506                "Live runtime disabled; projection and Yellowstone resources were not initialized"
507            );
508        }
509
510        // Run the HTTP server on a dedicated OS thread with its own single-threaded
511        // tokio runtime so liveness remains responsive under projection load.
512        let http_shutdown = CancellationToken::new();
513        let http_health_thread = if let Some(http_health_config) = &self.config.http_health {
514            let mut http_server = HttpServer::new(http_health_config.bind_address)
515                .with_runtime_plan(plan)
516                .with_program_runtime_catalog(program_runtime_catalog)
517                .with_shutdown(http_shutdown.clone());
518            if let Some(target_id) = self.config.program_read_binding_target_id.clone() {
519                http_server = http_server.with_program_read_binding_target(target_id);
520            }
521            if let Some(target_id) = self.config.solana_gateway_target_id.clone() {
522                http_server = http_server.with_solana_gateway_target(target_id);
523            }
524            if let Some(monitor) = health_monitor.clone() {
525                http_server = http_server.with_health_monitor(monitor);
526            }
527            if let Some(runtime) = snapshot_runtime.clone() {
528                http_server = http_server.with_snapshot_runtime(runtime);
529            }
530            if let Some(plugin) = self
531                .http_auth_plugin
532                .clone()
533                .or_else(|| self.websocket_auth_plugin.clone())
534            {
535                http_server = http_server.with_auth_plugin(plugin);
536            }
537            if plan.transactions && transaction_config.enabled {
538                http_server = http_server.with_transaction_config(transaction_config.clone());
539            }
540            if let Some(observer) = self.solana_gateway_usage_observer.clone() {
541                http_server = http_server.with_solana_gateway_usage_observer(observer);
542            }
543            #[cfg(feature = "otel")]
544            {
545                http_server = http_server.with_metrics(self.metrics.clone());
546            }
547
548            let bind_addr = http_health_config.bind_address;
549            let join_handle = std::thread::Builder::new()
550                .name("health-server".into())
551                .spawn(move || {
552                    let rt = tokio::runtime::Builder::new_current_thread()
553                        .enable_all()
554                        .build()
555                        .expect("Failed to create health server runtime");
556                    rt.block_on(async move {
557                        let _span = info_span!("http.health", %bind_addr).entered();
558                        if let Err(e) = http_server.start().await {
559                            error!("HTTP health server error: {}", e);
560                        }
561                    });
562                })
563                .expect("Failed to spawn health server thread");
564            info!(
565                "HTTP health server running on dedicated thread at {}",
566                bind_addr
567            );
568            Some(join_handle)
569        } else {
570            None
571        };
572
573        Ok(RuntimeHandle {
574            plan,
575            health_monitor,
576            snapshot_runtime,
577            snapshot_service,
578            snapshot_manager_handle,
579            mutations_tx: mutations_tx_guard,
580            projector_handle,
581            parser_handle,
582            ws_handle,
583            background,
584            acceptor,
585            entity_cache: entity_cache_handle,
586            websocket_usage_emitter: self.websocket_usage_emitter,
587            http_shutdown,
588            http_health_thread,
589        })
590    }
591}
592
593/// A running [`Runtime`], owned by whoever called [`Runtime::spawn`].
594///
595/// Dropping the handle does **not** stop the runtime; the tasks it owns keep
596/// running on the tokio runtime. Call [`shutdown`](Self::shutdown) to stop
597/// them and release the memory they hold. This is deliberate: the same
598/// semantics as dropping a `JoinHandle`, and what lets `run` hand the handle
599/// across a `select!`.
600pub struct RuntimeHandle {
601    plan: crate::RuntimePlan,
602    health_monitor: Option<HealthMonitor>,
603    snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
604    snapshot_service: Option<Arc<crate::snapshot::SnapshotService>>,
605    snapshot_manager_handle: Option<JoinHandle<()>>,
606    mutations_tx: Option<mpsc::Sender<MutationBatch>>,
607    projector_handle: Option<JoinHandle<()>>,
608    parser_handle: Option<JoinHandle<()>>,
609    ws_handle: Option<JoinHandle<()>>,
610    background: Vec<JoinHandle<()>>,
611    acceptor: Option<ConnectionAcceptor>,
612    entity_cache: Option<EntityCache>,
613    websocket_usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
614    http_shutdown: CancellationToken,
615    http_health_thread: Option<std::thread::JoinHandle<()>>,
616}
617
618/// Serves caller-accepted connections against one runtime. Obtained from
619/// [`RuntimeHandle::connection_server`]; cheap to clone into the task that
620/// owns each connection.
621#[derive(Clone)]
622pub struct ConnectionServer(ConnectionAcceptor);
623
624impl ConnectionServer {
625    /// Serve one accepted connection until the peer disconnects or the
626    /// runtime shuts down. See [`RuntimeHandle::serve_connection`].
627    pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
628        self.0.serve(stream, remote_addr).await
629    }
630
631    /// Number of WebSocket clients currently connected to the runtime.
632    pub fn client_count(&self) -> usize {
633        self.0.client_count()
634    }
635}
636
637/// How long `shutdown` waits for listener-spawned sessions to finish their
638/// cleanup after being told to stop.
639const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
640
641/// How long `shutdown` lets the projector drain queued batches after the
642/// producers have stopped.
643const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
644
645/// How long `shutdown` waits for accepted usage events to reach the API or
646/// durable spool storage.
647const USAGE_EMITTER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
648
649/// Bound on the final snapshot, chosen to fit inside the platform's
650/// termination grace period.
651const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
652
653impl RuntimeHandle {
654    /// The capability plan this runtime was started with.
655    pub fn plan(&self) -> crate::RuntimePlan {
656        self.plan
657    }
658
659    /// Whether this runtime should take traffic: the stream is healthy and,
660    /// after a snapshot restore, the parser has caught back up to the slot
661    /// tip. The same test the HTTP `/ready` endpoint applies.
662    pub async fn is_ready(&self) -> bool {
663        let stream_ready = match self.health_monitor.as_ref() {
664            Some(monitor) => monitor.is_healthy().await,
665            None => true,
666        };
667        let snapshot_ready = self
668            .snapshot_runtime
669            .as_ref()
670            .is_none_or(crate::snapshot::SnapshotRuntime::resume_gate_ready);
671        stream_ready && snapshot_ready
672    }
673
674    /// Number of WebSocket clients currently connected to this runtime.
675    pub fn client_count(&self) -> usize {
676        self.acceptor
677            .as_ref()
678            .map(ConnectionAcceptor::client_count)
679            .unwrap_or(0)
680    }
681
682    /// What this runtime's entity cache holds: the entities kept per view
683    /// for snapshot-on-subscribe. For an embedder accounting for the
684    /// runtime's memory. `None` when the runtime has no live runtime.
685    pub async fn entity_cache_stats(&self) -> Option<crate::cache::CacheStats> {
686        match &self.entity_cache {
687            Some(cache) => Some(cache.stats().await),
688            None => None,
689        }
690    }
691
692    /// Serve a TCP connection the caller accepted, as this runtime's
693    /// WebSocket server would have: handshake, authentication, then the
694    /// subscription session until the peer disconnects.
695    ///
696    /// Resolves when the session ends, or when the runtime shuts down. Fails
697    /// if the runtime has no live runtime (no buses to subscribe to). Callers
698    /// that serve from an accept loop should take a
699    /// [`connection_server`](Self::connection_server), which is cheap to clone
700    /// into each connection's task.
701    pub async fn serve_connection(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
702        match self.connection_server() {
703            Some(server) => server.serve(stream, remote_addr).await,
704            None => anyhow::bail!("this runtime has no live runtime to serve connections from"),
705        }
706    }
707
708    /// A clonable handle that serves connections against this runtime.
709    ///
710    /// `None` when the runtime has no live runtime. Sessions served through a
711    /// clone are ended by [`shutdown`](Self::shutdown) like any other.
712    pub fn connection_server(&self) -> Option<ConnectionServer> {
713        self.acceptor.clone().map(ConnectionServer)
714    }
715
716    /// Resolves when a core task - projector, parser or WebSocket listener -
717    /// exits on its own. `run` treats that as a reason to shut down.
718    pub async fn exited(&mut self) {
719        async fn wait(handle: Option<&mut JoinHandle<()>>) {
720            match handle {
721                Some(handle) => {
722                    let _ = handle.await;
723                }
724                None => std::future::pending().await,
725            }
726        }
727
728        tokio::select! {
729            _ = wait(self.ws_handle.as_mut()) => info!("WebSocket server task completed"),
730            _ = wait(self.projector_handle.as_mut()) => info!("Projector task completed"),
731            _ = wait(self.parser_handle.as_mut()) => info!("Parser runtime task completed"),
732        }
733    }
734
735    /// Stop the runtime: take the final snapshot if configured, stop
736    /// producing, let the projector drain, close sessions, then stop
737    /// everything else.
738    ///
739    /// Order matters. The final snapshot is taken *first*, while the parser
740    /// is still running: capture takes the barrier exclusively, so it waits
741    /// for every in-flight update to reach the projector and records one
742    /// consistent cut. Aborting the parser before that could cut an update
743    /// between its VM write and its batch, and the snapshot would keep the
744    /// write without the projection. Only then is the parser aborted; anything
745    /// it produces after the snapshot is simply not restored. The sender is
746    /// dropped so the projector exits once its queue is empty, bounded by a
747    /// timeout after which it is aborted and awaited. Sessions are ended
748    /// through their normal cleanup, the HTTP health server is signalled and
749    /// its thread joined, and the remaining tasks are aborted.
750    pub async fn shutdown(mut self) -> Result<()> {
751        // Final snapshot before anything stops, so planned stops restart
752        // near-lossless. Bounded to fit inside a termination grace period.
753        if let Some(service) = self.snapshot_service.take() {
754            if let Some(handle) = self.snapshot_manager_handle.take() {
755                handle.abort();
756            }
757            if service.config().snapshot_on_shutdown {
758                info!("Taking final snapshot before shutdown");
759                match tokio::time::timeout(
760                    SHUTDOWN_SNAPSHOT_TIMEOUT,
761                    service.snapshot_now(crate::snapshot::SnapshotTrigger::Shutdown),
762                )
763                .await
764                {
765                    Ok(Ok(_)) => {}
766                    Ok(Err(e)) => error!("Shutdown snapshot failed: {e:#}"),
767                    Err(_) => error!("Shutdown snapshot timed out"),
768                }
769            }
770        }
771        if let Some(handle) = self.snapshot_manager_handle.take() {
772            handle.abort();
773        }
774
775        if let Some(parser) = self.parser_handle.take() {
776            parser.abort();
777            let _ = parser.await;
778        }
779        if let Some(acceptor) = &self.acceptor {
780            acceptor.shutdown();
781        }
782        if let Some(ws) = self.ws_handle.take() {
783            let _ = ws.await;
784        }
785        let sessions_drained = if let Some(acceptor) = &self.acceptor {
786            if tokio::time::timeout(SESSION_DRAIN_TIMEOUT, acceptor.wait_for_sessions())
787                .await
788                .is_err()
789            {
790                warn!(
791                    "Sessions did not finish within {:?} of shutdown",
792                    SESSION_DRAIN_TIMEOUT
793                );
794                false
795            } else {
796                true
797            }
798        } else {
799            true
800        };
801
802        let usage_events_drained = if !sessions_drained {
803            false
804        } else if let Some(acceptor) = &self.acceptor {
805            if tokio::time::timeout(
806                USAGE_EMITTER_SHUTDOWN_TIMEOUT,
807                acceptor.wait_for_usage_events(),
808            )
809            .await
810            .is_err()
811            {
812                warn!(
813                    "WebSocket usage event tasks did not finish within {:?}",
814                    USAGE_EMITTER_SHUTDOWN_TIMEOUT
815                );
816                false
817            } else {
818                true
819            }
820        } else {
821            true
822        };
823
824        if let Some(emitter) = self.websocket_usage_emitter.take() {
825            let operation = if usage_events_drained {
826                emitter.shutdown()
827            } else {
828                // A timed-out session can still enqueue its final usage event.
829                // Keep the receiver alive for that late producer, but put an
830                // ordered persistence barrier behind everything already in its
831                // queue so shutdown cannot strand the accepted backlog.
832                emitter.flush()
833            };
834            if tokio::time::timeout(USAGE_EMITTER_SHUTDOWN_TIMEOUT, operation)
835                .await
836                .is_err()
837            {
838                warn!(
839                    operation = if usage_events_drained {
840                        "shutdown"
841                    } else {
842                        "flush"
843                    },
844                    "WebSocket usage emitter operation did not finish within {:?}",
845                    USAGE_EMITTER_SHUTDOWN_TIMEOUT
846                );
847            }
848        }
849
850        drop(self.mutations_tx.take());
851        if let Some(mut projector) = self.projector_handle.take() {
852            if tokio::time::timeout(PROJECTOR_DRAIN_TIMEOUT, &mut projector)
853                .await
854                .is_err()
855            {
856                warn!(
857                    "Projector did not drain within {:?}; aborting it",
858                    PROJECTOR_DRAIN_TIMEOUT
859                );
860                projector.abort();
861                let _ = projector.await;
862            }
863        }
864
865        for handle in self.background.drain(..) {
866            handle.abort();
867        }
868
869        self.http_shutdown.cancel();
870        if let Some(thread) = self.http_health_thread.take() {
871            if let Err(e) = tokio::task::spawn_blocking(move || thread.join()).await {
872                error!("Health server thread join failed: {e}");
873            }
874        }
875
876        info!("Shutting down Arete runtime");
877        Ok(())
878    }
879}