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