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::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
32/// Wait for shutdown signal (SIGINT on all platforms, SIGTERM on Unix)
33async 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    /// Configure rate limiting for WebSocket connections.
149    ///
150    /// This sets global rate limits such as maximum connections per IP,
151    /// timeouts, and rate windows. Per-subject limits are controlled
152    /// via AuthContext.Limits from the authentication token.
153    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    /// Return the immutable capability plan selected by the builder.
159    pub fn plan(&self) -> crate::RuntimePlan {
160        self.config.runtime_plan
161    }
162
163    /// Start everything and block until a shutdown signal, then stop cleanly.
164    ///
165    /// This is [`spawn`](Self::spawn) followed by waiting for SIGINT/SIGTERM
166    /// (or for a core task to exit) and [`RuntimeHandle::shutdown`]. Callers
167    /// that embed the server in a larger process should use `spawn` directly
168    /// and decide for themselves when to stop.
169    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    /// Start the runtime's tasks and return a handle that owns them.
182    ///
183    /// Everything `run` starts is started here: projector, parser, snapshot
184    /// manager, bus cleanup, stats, and - only when configured - the WebSocket
185    /// listener and the HTTP health server. Nothing here installs signal
186    /// handlers or blocks; the handle is how the caller waits, serves
187    /// connections it accepted itself, and shuts the runtime down.
188    ///
189    /// A handle owns its runtime's channel, buses, cache and snapshot state
190    /// outright; nothing is shared through process globals, so a test or an
191    /// application can start and stop runtimes independently.
192    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            // Retained event tape for replayable append subscriptions. The
255            // builder wins over the process env so one host can enable replay
256            // for a single deployment and size it independently. A bad
257            // configuration disables replay rather than failing startup, the
258            // same posture snapshots take below.
259            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            // Restore state from the latest snapshot (when enabled) before the
280            // WebSocket server spawns, so the first client's snapshot-on-subscribe
281            // is already warm. The VM portion is stashed for the generated
282            // runtime to hydrate before it connects to Yellowstone.
283            if let Some(spec) = self.spec.as_ref() {
284                let snapshot_config = match self.config.snapshots.clone() {
285                    Some(config) => Some(config),
286                    None => match crate::snapshot::SnapshotConfig::from_env() {
287                        Ok(config) => Some(config),
288                        Err(e) => {
289                            error!("Invalid snapshot configuration; snapshots disabled: {e:#}");
290                            None
291                        }
292                    },
293                };
294                if let Some(snapshot_config) = snapshot_config.filter(|c| c.enabled) {
295                    match crate::snapshot::SnapshotService::initialize(
296                        snapshot_config,
297                        spec,
298                        entity_cache.clone(),
299                        &self.view_index,
300                        journal.clone(),
301                        mutations_tx.clone(),
302                    )
303                    .await
304                    {
305                        Ok(service) => {
306                            snapshot_runtime = Some(service.runtime());
307                            snapshot_manager_handle = Some(service.spawn());
308                            snapshot_service = Some(service);
309                        }
310                        Err(e) => {
311                            error!("Failed to initialize snapshots; continuing without: {e:#}")
312                        }
313                    }
314                }
315            }
316
317            #[cfg(feature = "otel")]
318            let projector = Projector::new(
319                self.view_index.clone(),
320                bus_manager.clone(),
321                entity_cache.clone(),
322                mutations_rx,
323                self.metrics.clone(),
324            );
325            #[cfg(not(feature = "otel"))]
326            let projector = Projector::new(
327                self.view_index.clone(),
328                bus_manager.clone(),
329                entity_cache.clone(),
330                mutations_rx,
331            );
332            let projector = match snapshot_runtime.clone() {
333                Some(runtime) => projector.with_snapshot_runtime(runtime),
334                None => projector,
335            };
336            let projector = projector.with_journal(journal.clone());
337
338            // The projector runs for the lifetime of the server. Giving the
339            // task a span would make that span the parent of every batch
340            // whose producer does not carry an explicit context, creating an
341            // unbounded trace. `Projector::run` instead enters the bounded
342            // batch span before it processes and logs each batch.
343            projector_handle = Some(tokio::spawn(async move {
344                projector.run().await;
345            }));
346
347            // The connection-serving half of the WebSocket server exists
348            // whenever there is a live runtime, so a caller that owns its own
349            // listener can hand in the connections it accepts. The listener
350            // is bound only when a WebSocket address is configured.
351            let bind_address = self
352                .config
353                .websocket
354                .as_ref()
355                .map(|ws_config| ws_config.bind_address)
356                .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));
357            #[cfg(feature = "otel")]
358            let mut ws_server = WebSocketServer::new(
359                bind_address,
360                bus_manager.clone(),
361                entity_cache.clone(),
362                self.view_index.clone(),
363                self.metrics.clone(),
364            );
365            #[cfg(not(feature = "otel"))]
366            let mut ws_server = WebSocketServer::new(
367                bind_address,
368                bus_manager.clone(),
369                entity_cache.clone(),
370                self.view_index.clone(),
371            );
372
373            ws_server = ws_server.with_journal(journal.clone());
374            ws_server = ws_server.with_delivery_config(websocket_delivery);
375            if let Some(max_clients) = self.websocket_max_clients {
376                ws_server = ws_server.with_max_clients(max_clients);
377            }
378            if let Some(plugin) = self.websocket_auth_plugin.clone() {
379                ws_server = ws_server.with_auth_plugin(plugin);
380            }
381            if let Some(emitter) = self.websocket_usage_emitter.clone() {
382                ws_server = ws_server.with_usage_emitter(emitter);
383            }
384            if let Some(rate_limit_config) = self.websocket_rate_limit_config {
385                ws_server = ws_server.with_rate_limit_config(rate_limit_config);
386            }
387            let (connection_acceptor, cleanup_handle) = ws_server.into_acceptor();
388            background.push(cleanup_handle);
389
390            if plan.websocket && self.config.websocket.is_some() {
391                let listener_acceptor = connection_acceptor.clone();
392                ws_handle = Some(tokio::spawn(
393                    async move {
394                        info!("Starting WebSocket server on {}", bind_address);
395                        let listener = match TcpListener::bind(&bind_address).await {
396                            Ok(listener) => listener,
397                            Err(e) => {
398                                error!("WebSocket server error: {}", e);
399                                return;
400                            }
401                        };
402                        if let Err(e) = listener_acceptor.serve_listener(listener).await {
403                            error!("WebSocket server error: {}", e);
404                        }
405                    }
406                    .instrument(info_span!("ws.server", %bind_address)),
407                ));
408            }
409            acceptor = Some(connection_acceptor);
410
411            if let Some(spec) = self.spec.as_ref() {
412                if let Some(parser_setup) = spec.parser_setup.clone() {
413                    let program_id = spec
414                        .program_ids
415                        .first()
416                        .cloned()
417                        .unwrap_or_else(|| "unknown".to_string());
418                    info!("Starting parser runtime for program: {}", program_id);
419                    let health = health_monitor.clone();
420                    let reconnection_config = self.config.reconnection.clone().unwrap_or_default();
421                    let parser_snapshot_runtime = snapshot_runtime.clone();
422                    let parser_journal = journal.clone();
423                    // The parser runs for the lifetime of the server, like
424                    // the projector. A span on the task would be the current
425                    // span of every update it processes, and OpenTelemetry
426                    // keeps an open span's events until it closes: for a task
427                    // that runs as long as the server, an unbounded buffer.
428                    parser_handle = Some(tokio::spawn(async move {
429                        let parser = async move {
430                            parser_setup(mutations_tx, health, reconnection_config).await
431                        };
432                        let scoped = async move {
433                            match parser_snapshot_runtime {
434                                Some(runtime) => runtime.scope(parser).await,
435                                None => parser.await,
436                            }
437                        };
438                        // The tape is in scope even with snapshots off, so
439                        // a runtime that abandons its checkpoint can still
440                        // mark the hole it just created.
441                        let result = parser_journal.scope(scoped).await;
442                        if let Err(e) = result {
443                            error!(%program_id, "Vixen parser runtime error: {}", e);
444                        }
445                    }));
446                } else {
447                    info!("Spec provided but no parser_setup configured - skipping parser runtime");
448                }
449            } else {
450                info!("No spec provided - running in websocket-only mode");
451            }
452
453            let cleanup_bus = bus_manager.clone();
454            background.push(tokio::spawn(
455                async move {
456                    let mut interval = tokio::time::interval(Duration::from_secs(60));
457                    loop {
458                        interval.tick().await;
459                        let state_cleaned = cleanup_bus.cleanup_stale_state_buses().await;
460                        let list_cleaned = cleanup_bus.cleanup_stale_list_buses().await;
461                        if state_cleaned > 0 || list_cleaned > 0 {
462                            let (state_count, list_count) = cleanup_bus.bus_counts().await;
463                            info!(
464                                "Bus cleanup: removed {} state, {} list buses. Current: {} state, {} list",
465                                state_cleaned, list_cleaned, state_count, list_count
466                            );
467                        }
468                    }
469                }
470                .instrument(info_span!("bus.cleanup")),
471            ));
472
473            background.push(tokio::spawn(
474                async move {
475                    let mut interval = tokio::time::interval(Duration::from_secs(30));
476                    loop {
477                        interval.tick().await;
478                        let (_state_buses, _list_buses) = bus_manager.bus_counts().await;
479                        let _cache_stats = entity_cache.stats().await;
480                    }
481                }
482                .instrument(info_span!("stats.reporter")),
483            ));
484        } else {
485            info!(
486                "Live runtime disabled; projection and Yellowstone resources were not initialized"
487            );
488        }
489
490        // Run the HTTP server on a dedicated OS thread with its own single-threaded
491        // tokio runtime so liveness remains responsive under projection load.
492        let http_shutdown = CancellationToken::new();
493        let http_health_thread = if let Some(http_health_config) = &self.config.http_health {
494            let mut http_server = HttpServer::new(http_health_config.bind_address)
495                .with_runtime_plan(plan)
496                .with_program_runtime_catalog(program_runtime_catalog)
497                .with_shutdown(http_shutdown.clone());
498            if let Some(target_id) = self.config.program_read_binding_target_id.clone() {
499                http_server = http_server.with_program_read_binding_target(target_id);
500            }
501            if let Some(target_id) = self.config.solana_gateway_target_id.clone() {
502                http_server = http_server.with_solana_gateway_target(target_id);
503            }
504            if let Some(monitor) = health_monitor.clone() {
505                http_server = http_server.with_health_monitor(monitor);
506            }
507            if let Some(runtime) = snapshot_runtime.clone() {
508                http_server = http_server.with_snapshot_runtime(runtime);
509            }
510            if let Some(plugin) = self
511                .http_auth_plugin
512                .clone()
513                .or_else(|| self.websocket_auth_plugin.clone())
514            {
515                http_server = http_server.with_auth_plugin(plugin);
516            }
517            if plan.transactions && transaction_config.enabled {
518                http_server = http_server.with_transaction_config(transaction_config.clone());
519            }
520            #[cfg(feature = "otel")]
521            {
522                http_server = http_server.with_metrics(self.metrics.clone());
523            }
524
525            let bind_addr = http_health_config.bind_address;
526            let join_handle = std::thread::Builder::new()
527                .name("health-server".into())
528                .spawn(move || {
529                    let rt = tokio::runtime::Builder::new_current_thread()
530                        .enable_all()
531                        .build()
532                        .expect("Failed to create health server runtime");
533                    rt.block_on(async move {
534                        let _span = info_span!("http.health", %bind_addr).entered();
535                        if let Err(e) = http_server.start().await {
536                            error!("HTTP health server error: {}", e);
537                        }
538                    });
539                })
540                .expect("Failed to spawn health server thread");
541            info!(
542                "HTTP health server running on dedicated thread at {}",
543                bind_addr
544            );
545            Some(join_handle)
546        } else {
547            None
548        };
549
550        Ok(RuntimeHandle {
551            plan,
552            health_monitor,
553            snapshot_runtime,
554            snapshot_service,
555            snapshot_manager_handle,
556            mutations_tx: mutations_tx_guard,
557            projector_handle,
558            parser_handle,
559            ws_handle,
560            background,
561            acceptor,
562            entity_cache: entity_cache_handle,
563            http_shutdown,
564            http_health_thread,
565        })
566    }
567}
568
569/// A running [`Runtime`], owned by whoever called [`Runtime::spawn`].
570///
571/// Dropping the handle does **not** stop the runtime; the tasks it owns keep
572/// running on the tokio runtime. Call [`shutdown`](Self::shutdown) to stop
573/// them and release the memory they hold. This is deliberate: the same
574/// semantics as dropping a `JoinHandle`, and what lets `run` hand the handle
575/// across a `select!`.
576pub struct RuntimeHandle {
577    plan: crate::RuntimePlan,
578    health_monitor: Option<HealthMonitor>,
579    snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
580    snapshot_service: Option<Arc<crate::snapshot::SnapshotService>>,
581    snapshot_manager_handle: Option<JoinHandle<()>>,
582    mutations_tx: Option<mpsc::Sender<MutationBatch>>,
583    projector_handle: Option<JoinHandle<()>>,
584    parser_handle: Option<JoinHandle<()>>,
585    ws_handle: Option<JoinHandle<()>>,
586    background: Vec<JoinHandle<()>>,
587    acceptor: Option<ConnectionAcceptor>,
588    entity_cache: Option<EntityCache>,
589    http_shutdown: CancellationToken,
590    http_health_thread: Option<std::thread::JoinHandle<()>>,
591}
592
593/// Serves caller-accepted connections against one runtime. Obtained from
594/// [`RuntimeHandle::connection_server`]; cheap to clone into the task that
595/// owns each connection.
596#[derive(Clone)]
597pub struct ConnectionServer(ConnectionAcceptor);
598
599impl ConnectionServer {
600    /// Serve one accepted connection until the peer disconnects or the
601    /// runtime shuts down. See [`RuntimeHandle::serve_connection`].
602    pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
603        self.0.serve(stream, remote_addr).await
604    }
605
606    /// Number of WebSocket clients currently connected to the runtime.
607    pub fn client_count(&self) -> usize {
608        self.0.client_count()
609    }
610}
611
612/// How long `shutdown` waits for listener-spawned sessions to finish their
613/// cleanup after being told to stop.
614const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
615
616/// How long `shutdown` lets the projector drain queued batches after the
617/// producers have stopped.
618const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
619
620/// Bound on the final snapshot, chosen to fit inside the platform's
621/// termination grace period.
622const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
623
624impl RuntimeHandle {
625    /// The capability plan this runtime was started with.
626    pub fn plan(&self) -> crate::RuntimePlan {
627        self.plan
628    }
629
630    /// Whether this runtime should take traffic: the stream is healthy and,
631    /// after a snapshot restore, the parser has caught back up to the slot
632    /// tip. The same test the HTTP `/ready` endpoint applies.
633    pub async fn is_ready(&self) -> bool {
634        let stream_ready = match self.health_monitor.as_ref() {
635            Some(monitor) => monitor.is_healthy().await,
636            None => true,
637        };
638        let snapshot_ready = self
639            .snapshot_runtime
640            .as_ref()
641            .is_none_or(crate::snapshot::SnapshotRuntime::resume_gate_ready);
642        stream_ready && snapshot_ready
643    }
644
645    /// Number of WebSocket clients currently connected to this runtime.
646    pub fn client_count(&self) -> usize {
647        self.acceptor
648            .as_ref()
649            .map(ConnectionAcceptor::client_count)
650            .unwrap_or(0)
651    }
652
653    /// What this runtime's entity cache holds: the entities kept per view
654    /// for snapshot-on-subscribe. For an embedder accounting for the
655    /// runtime's memory. `None` when the runtime has no live runtime.
656    pub async fn entity_cache_stats(&self) -> Option<crate::cache::CacheStats> {
657        match &self.entity_cache {
658            Some(cache) => Some(cache.stats().await),
659            None => None,
660        }
661    }
662
663    /// Serve a TCP connection the caller accepted, as this runtime's
664    /// WebSocket server would have: handshake, authentication, then the
665    /// subscription session until the peer disconnects.
666    ///
667    /// Resolves when the session ends, or when the runtime shuts down. Fails
668    /// if the runtime has no live runtime (no buses to subscribe to). Callers
669    /// that serve from an accept loop should take a
670    /// [`connection_server`](Self::connection_server), which is cheap to clone
671    /// into each connection's task.
672    pub async fn serve_connection(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
673        match self.connection_server() {
674            Some(server) => server.serve(stream, remote_addr).await,
675            None => anyhow::bail!("this runtime has no live runtime to serve connections from"),
676        }
677    }
678
679    /// A clonable handle that serves connections against this runtime.
680    ///
681    /// `None` when the runtime has no live runtime. Sessions served through a
682    /// clone are ended by [`shutdown`](Self::shutdown) like any other.
683    pub fn connection_server(&self) -> Option<ConnectionServer> {
684        self.acceptor.clone().map(ConnectionServer)
685    }
686
687    /// Resolves when a core task - projector, parser or WebSocket listener -
688    /// exits on its own. `run` treats that as a reason to shut down.
689    pub async fn exited(&mut self) {
690        async fn wait(handle: Option<&mut JoinHandle<()>>) {
691            match handle {
692                Some(handle) => {
693                    let _ = handle.await;
694                }
695                None => std::future::pending().await,
696            }
697        }
698
699        tokio::select! {
700            _ = wait(self.ws_handle.as_mut()) => info!("WebSocket server task completed"),
701            _ = wait(self.projector_handle.as_mut()) => info!("Projector task completed"),
702            _ = wait(self.parser_handle.as_mut()) => info!("Parser runtime task completed"),
703        }
704    }
705
706    /// Stop the runtime: take the final snapshot if configured, stop
707    /// producing, let the projector drain, close sessions, then stop
708    /// everything else.
709    ///
710    /// Order matters. The final snapshot is taken *first*, while the parser
711    /// is still running: capture takes the barrier exclusively, so it waits
712    /// for every in-flight update to reach the projector and records one
713    /// consistent cut. Aborting the parser before that could cut an update
714    /// between its VM write and its batch, and the snapshot would keep the
715    /// write without the projection. Only then is the parser aborted; anything
716    /// it produces after the snapshot is simply not restored. The sender is
717    /// dropped so the projector exits once its queue is empty, bounded by a
718    /// timeout after which it is aborted and awaited. Sessions are ended
719    /// through their normal cleanup, the HTTP health server is signalled and
720    /// its thread joined, and the remaining tasks are aborted.
721    pub async fn shutdown(mut self) -> Result<()> {
722        // Final snapshot before anything stops, so planned stops restart
723        // near-lossless. Bounded to fit inside a termination grace period.
724        if let Some(service) = self.snapshot_service.take() {
725            if let Some(handle) = self.snapshot_manager_handle.take() {
726                handle.abort();
727            }
728            if service.config().snapshot_on_shutdown {
729                info!("Taking final snapshot before shutdown");
730                match tokio::time::timeout(
731                    SHUTDOWN_SNAPSHOT_TIMEOUT,
732                    service.snapshot_now(crate::snapshot::SnapshotTrigger::Shutdown),
733                )
734                .await
735                {
736                    Ok(Ok(_)) => {}
737                    Ok(Err(e)) => error!("Shutdown snapshot failed: {e:#}"),
738                    Err(_) => error!("Shutdown snapshot timed out"),
739                }
740            }
741        }
742        if let Some(handle) = self.snapshot_manager_handle.take() {
743            handle.abort();
744        }
745
746        if let Some(parser) = self.parser_handle.take() {
747            parser.abort();
748            let _ = parser.await;
749        }
750        if let Some(acceptor) = &self.acceptor {
751            acceptor.shutdown();
752        }
753        if let Some(ws) = self.ws_handle.take() {
754            let _ = ws.await;
755        }
756        if let Some(acceptor) = &self.acceptor {
757            if tokio::time::timeout(SESSION_DRAIN_TIMEOUT, acceptor.wait_for_sessions())
758                .await
759                .is_err()
760            {
761                warn!(
762                    "Sessions did not finish within {:?} of shutdown",
763                    SESSION_DRAIN_TIMEOUT
764                );
765            }
766        }
767
768        drop(self.mutations_tx.take());
769        if let Some(mut projector) = self.projector_handle.take() {
770            if tokio::time::timeout(PROJECTOR_DRAIN_TIMEOUT, &mut projector)
771                .await
772                .is_err()
773            {
774                warn!(
775                    "Projector did not drain within {:?}; aborting it",
776                    PROJECTOR_DRAIN_TIMEOUT
777                );
778                projector.abort();
779                let _ = projector.await;
780            }
781        }
782
783        for handle in self.background.drain(..) {
784            handle.abort();
785        }
786
787        self.http_shutdown.cancel();
788        if let Some(thread) = self.http_health_thread.take() {
789            if let Err(e) = tokio::task::spawn_blocking(move || thread.join()).await {
790                error!("Health server thread join failed: {e}");
791            }
792        }
793
794        info!("Shutting down Arete runtime");
795        Ok(())
796    }
797}