Skip to main content

arete_server/
runtime.rs

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