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