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
33async 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 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 pub fn plan(&self) -> crate::RuntimePlan {
194 self.config.runtime_plan
195 }
196
197 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 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 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 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 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 projector_handle = Some(tokio::spawn(async move {
383 projector.run().await;
384 }));
385
386 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 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 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 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
621pub 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#[derive(Clone)]
650pub struct ConnectionServer(ConnectionAcceptor);
651
652impl ConnectionServer {
653 pub async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
656 self.0.serve(stream, remote_addr).await
657 }
658
659 pub fn client_count(&self) -> usize {
661 self.0.client_count()
662 }
663}
664
665const SESSION_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
668
669const PROJECTOR_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
672
673const USAGE_EMITTER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
676
677const SHUTDOWN_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(20);
680
681impl RuntimeHandle {
682 pub fn plan(&self) -> crate::RuntimePlan {
684 self.plan
685 }
686
687 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 pub fn client_count(&self) -> usize {
704 self.acceptor
705 .as_ref()
706 .map(ConnectionAcceptor::client_count)
707 .unwrap_or(0)
708 }
709
710 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 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 pub fn connection_server(&self) -> Option<ConnectionServer> {
741 self.acceptor.clone().map(ConnectionServer)
742 }
743
744 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 pub async fn shutdown(mut self) -> Result<()> {
779 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 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}