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