eggress_runtime/supervisor/
state.rs1use 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 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 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 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}