Skip to main content

eggress_runtime/supervisor/
state.rs

1//! Shared runtime state and the canonical reload transaction.
2//!
3//! [`RuntimeState`] is the single owner of the compiled snapshot, routing
4//! service, metrics handle, readiness flag, connection accounting, UDP
5//! registry, health manager, and (feature-gated) reverse/admin state.
6//! [`RuntimeState::apply_compiled_config`] is the one canonical reload
7//! transaction used by file-backed reload, SIGHUP handling, and embed
8//! string/file/compiled reload entry points.
9
10use std::sync::atomic::{AtomicBool, AtomicU64};
11use std::sync::{Arc, Mutex};
12use std::time::Instant;
13
14use arc_swap::ArcSwap;
15use tokio_util::sync::CancellationToken;
16use tokio_util::task::TaskTracker;
17
18use eggress_routing::health::HealthManager;
19use eggress_routing::upstream::UpstreamRuntime;
20use eggress_routing::SharedRoutingService;
21
22use crate::snapshot::{compile_runtime_snapshot, CompiledRuntimeSnapshot};
23
24#[cfg(feature = "operations")]
25use super::operations::RuntimeAdminState;
26use super::reload::{classify_reload_config, ReloadResult};
27
28pub struct RuntimeState {
29    pub snapshot: Arc<ArcSwap<CompiledRuntimeSnapshot>>,
30    pub routing: Arc<SharedRoutingService>,
31    pub metrics: Arc<dyn eggress_server::SessionMetrics>,
32    pub runtime_metrics: Arc<dyn eggress_metrics::RuntimeMetrics>,
33    pub readiness: Arc<AtomicBool>,
34    pub start_time: Instant,
35    pub active_connections: Arc<AtomicU64>,
36    pub connection_counter: Arc<AtomicU64>,
37    pub admin_local_addr: Arc<Mutex<Option<std::net::SocketAddr>>>,
38    pub listener_addrs: Arc<Mutex<Vec<Option<std::net::SocketAddr>>>>,
39    #[cfg(feature = "operations")]
40    pub(crate) admin_snapshot: Arc<ArcSwap<RuntimeAdminState>>,
41    pub health: Arc<Mutex<Option<HealthManager>>>,
42    pub health_cancel: CancellationToken,
43    pub health_runtime: Mutex<Option<tokio::runtime::Handle>>,
44    pub udp_registry: Arc<eggress_udp::registry::UdpAssociationRegistry>,
45    pub udp_metrics: Arc<eggress_udp::metrics::UdpMetrics>,
46    #[cfg(feature = "extended")]
47    pub shadowsocks_metrics: Arc<eggress_protocol_shadowsocks::ShadowsocksMetrics>,
48    pub udp_tasks: TaskTracker,
49    pub transparent_accepted_total: Arc<AtomicU64>,
50    pub transparent_original_dst_failed_total: Arc<AtomicU64>,
51    #[cfg(feature = "reverse")]
52    pub reverse_registry: Arc<eggress_admin::ReverseRegistry>,
53    #[cfg(feature = "reverse")]
54    pub reverse_metrics: Arc<eggress_protocol_reverse::metrics::ReverseMetrics>,
55}
56
57impl RuntimeState {
58    pub fn generation(&self) -> u64 {
59        self.snapshot.load().generation
60    }
61
62    /// Canonical reload transaction shared by file-backed supervisor reload,
63    /// SIGHUP handling, and embed string/file reload entry points.
64    ///
65    /// Applies a newly compiled [`eggress_config::compile::RuntimeConfig`] to
66    /// the running state with all side effects centralized:
67    /// classification, snapshot compilation, snapshot publication, routing
68    /// swap, admin publication, health restart, H2 pool invalidation, and
69    /// metrics recording. Failure preserves the prior generation.
70    ///
71    /// Callers differ only in how `new_config` is obtained (file load vs.
72    /// string parse). Supervisors with stored `rt_config` must update that
73    /// bookkeeping on `Applied`; the snapshot itself remains authoritative
74    /// for the next classification.
75    pub fn apply_compiled_config(
76        &self,
77        new_config: &eggress_config::compile::RuntimeConfig,
78    ) -> ReloadResult {
79        let prev_snapshot = self.snapshot.load();
80        if let Err(reason) = classify_reload_config(
81            &prev_snapshot.listeners,
82            &prev_snapshot.timeouts,
83            prev_snapshot.admin.as_ref(),
84            new_config,
85        ) {
86            self.runtime_metrics.record_reload(false);
87            return ReloadResult::Rejected { reason };
88        }
89
90        let prev_ref: Option<&CompiledRuntimeSnapshot> = Some(&prev_snapshot);
91        let new_snapshot = match compile_runtime_snapshot(new_config, prev_ref) {
92            Ok(snapshot) => snapshot,
93            Err(error) => {
94                self.runtime_metrics.record_reload(false);
95                return ReloadResult::Failed {
96                    error: format!("snapshot build: {error}"),
97                };
98            }
99        };
100
101        let upstream_count = new_snapshot.upstreams.len();
102        let generation = new_snapshot.generation;
103
104        // Snapshot must be published before the router swap. Readers that
105        // observe the new generation via `snapshot.load()` pull the router
106        // from that same snapshot Arc, so any reader seeing the new
107        // generation also sees the router that belongs to it.
108        let new_snapshot = Arc::new(new_snapshot);
109        self.snapshot.store(new_snapshot.clone());
110        self.routing.swap_arc(new_snapshot.router.clone());
111        #[cfg(feature = "operations")]
112        self.publish_admin_snapshot(new_snapshot.clone());
113
114        self.restart_health_probes();
115        eggress_protocol_http::H2_POOL_REGISTRY.clear();
116
117        self.runtime_metrics.set_config_generation(generation);
118        self.runtime_metrics.record_reload(true);
119
120        ReloadResult::Applied {
121            generation,
122            upstreams: upstream_count,
123        }
124    }
125
126    #[cfg(feature = "operations")]
127    pub(crate) fn publish_admin_snapshot(&self, snapshot: Arc<CompiledRuntimeSnapshot>) {
128        let listener_addrs = self.admin_snapshot.load().listener_addrs.clone();
129        self.admin_snapshot.store(Arc::new(RuntimeAdminState {
130            snapshot,
131            listener_addrs,
132        }));
133    }
134
135    #[cfg(feature = "operations")]
136    pub(crate) fn publish_admin_listener_addrs(
137        &self,
138        snapshot: Arc<CompiledRuntimeSnapshot>,
139        listener_addrs: Vec<Option<std::net::SocketAddr>>,
140    ) {
141        self.admin_snapshot.store(Arc::new(RuntimeAdminState {
142            snapshot,
143            listener_addrs,
144        }));
145    }
146
147    /// Restart health probes for the upstreams in the current snapshot.
148    pub fn restart_health_probes(&self) {
149        let mut guard = self.health.lock().unwrap_or_else(|error| {
150            tracing::warn!("health manager state was poisoned; resetting it: {error}");
151            let mut guard = error.into_inner();
152            *guard = None;
153            self.health.clear_poison();
154            guard
155        });
156        if let Some(ref mut health) = *guard {
157            health.stop_all();
158        }
159        let upstreams: Vec<Arc<UpstreamRuntime>> =
160            self.snapshot.load().upstreams.values().cloned().collect();
161        if !upstreams.is_empty() {
162            let mut health = HealthManager::new(self.health_cancel.clone());
163            if let Some(handle) = self
164                .health_runtime
165                .lock()
166                .unwrap_or_else(|error| error.into_inner())
167                .clone()
168            {
169                health.start_probes_on(&handle, &upstreams);
170            }
171            *guard = Some(health);
172        } else {
173            *guard = None;
174        }
175    }
176}