1use super::RuntimeHttpHost;
4use arc_swap::ArcSwap;
5use axum::body::Body;
6use axum::extract::{Request, State};
7use axum::http::{Response, StatusCode};
8use axum::routing::any;
9use axum::Router;
10use std::fmt;
11use std::io;
12use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15use tower::ServiceExt;
16
17const ROUTING_RETRY_DELAY: Duration = Duration::from_millis(1);
18const MAX_RELOAD_TIMEOUT: Duration = Duration::from_secs(60);
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum HttpReloadPhase {
23 Prepare,
25 Switch,
27 Drain,
29 Rollback,
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct HttpReloadPolicy {
36 health_timeout: Duration,
37 drain_timeout: Duration,
38}
39
40impl HttpReloadPolicy {
41 pub fn new(
43 health_timeout: Duration,
44 drain_timeout: Duration,
45 ) -> Result<Self, RuntimeHttpReloadError> {
46 if health_timeout.is_zero()
47 || drain_timeout.is_zero()
48 || health_timeout > MAX_RELOAD_TIMEOUT
49 || drain_timeout > MAX_RELOAD_TIMEOUT
50 {
51 return Err(RuntimeHttpReloadError::InvalidPolicy);
52 }
53 Ok(Self {
54 health_timeout,
55 drain_timeout,
56 })
57 }
58
59 pub fn health_timeout(self) -> Duration {
61 self.health_timeout
62 }
63
64 pub fn drain_timeout(self) -> Duration {
66 self.drain_timeout
67 }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub struct HttpReloadSnapshot {
73 pub active_generation: u64,
75 pub active_inflight: usize,
77 pub reload_in_progress: bool,
79 pub successful_reloads: u64,
81 pub failed_reloads: u64,
83 pub rollbacks: u64,
85}
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum RuntimeHttpReloadError {
90 ListenerDisabled,
92 ListenerAddressChanged,
94 StaleGeneration,
96 ReloadInProgress,
98 InvalidPolicy,
100 HealthGateFailed(HttpReloadPhase),
102 DrainTimedOut,
104 RollbackDrainTimedOut,
106}
107
108impl fmt::Display for RuntimeHttpReloadError {
109 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
110 match self {
111 Self::ListenerDisabled => formatter.write_str("HTTP reload listener is disabled"),
112 Self::ListenerAddressChanged => {
113 formatter.write_str("HTTP reload requires a prepared listener generation")
114 }
115 Self::StaleGeneration => formatter.write_str("HTTP routing generation must increase"),
116 Self::ReloadInProgress => formatter.write_str("HTTP reload is already in progress"),
117 Self::InvalidPolicy => formatter.write_str("HTTP reload policy is invalid"),
118 Self::HealthGateFailed(phase) => {
119 write!(formatter, "HTTP reload health gate failed during {phase:?}")
120 }
121 Self::DrainTimedOut => formatter.write_str("HTTP routing generation drain timed out"),
122 Self::RollbackDrainTimedOut => {
123 formatter.write_str("HTTP rollback generation drain timed out")
124 }
125 }
126 }
127}
128
129impl std::error::Error for RuntimeHttpReloadError {}
130
131struct RoutingGeneration {
132 id: u64,
133 router: Router,
134 accepting: AtomicBool,
135 inflight: AtomicUsize,
136}
137
138impl RoutingGeneration {
139 fn new(id: u64, router: Router) -> Self {
140 Self {
141 id,
142 router,
143 accepting: AtomicBool::new(true),
144 inflight: AtomicUsize::new(0),
145 }
146 }
147
148 fn try_admit(self: &Arc<Self>) -> Option<RoutingPermit> {
149 if !self.accepting.load(Ordering::Acquire) {
150 return None;
151 }
152 self.inflight
153 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
154 current.checked_add(1)
155 })
156 .ok()?;
157 if self.accepting.load(Ordering::Acquire) {
158 return Some(RoutingPermit {
159 generation: Arc::clone(self),
160 });
161 }
162 self.inflight.fetch_sub(1, Ordering::AcqRel);
163 None
164 }
165}
166
167struct RoutingPermit {
168 generation: Arc<RoutingGeneration>,
169}
170
171impl Drop for RoutingPermit {
172 fn drop(&mut self) {
173 self.generation.inflight.fetch_sub(1, Ordering::AcqRel);
174 }
175}
176
177struct RoutingTable {
178 active: ArcSwap<RoutingGeneration>,
179}
180
181pub struct PreparedRuntimeHttpGeneration {
183 generation: Arc<RoutingGeneration>,
184}
185
186impl fmt::Debug for PreparedRuntimeHttpGeneration {
187 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
188 formatter
189 .debug_struct("PreparedRuntimeHttpGeneration")
190 .field("generation", &self.generation.id)
191 .finish_non_exhaustive()
192 }
193}
194
195pub struct ReloadableRuntimeHttpHost {
197 config: super::HttpApiConfig,
198 routing: Arc<RoutingTable>,
199 reload_in_progress: AtomicBool,
200 successful_reloads: AtomicU64,
201 failed_reloads: AtomicU64,
202 rollbacks: AtomicU64,
203}
204
205impl ReloadableRuntimeHttpHost {
206 pub fn new(
208 initial_generation: u64,
209 host: RuntimeHttpHost,
210 ) -> Result<Self, RuntimeHttpReloadError> {
211 if !host.config().enabled {
212 return Err(RuntimeHttpReloadError::ListenerDisabled);
213 }
214 if initial_generation == 0 {
215 return Err(RuntimeHttpReloadError::StaleGeneration);
216 }
217 let config = host.config().clone();
218 let initial = Arc::new(RoutingGeneration::new(initial_generation, host.router()));
219 Ok(Self {
220 config,
221 routing: Arc::new(RoutingTable {
222 active: ArcSwap::from(initial),
223 }),
224 reload_in_progress: AtomicBool::new(false),
225 successful_reloads: AtomicU64::new(0),
226 failed_reloads: AtomicU64::new(0),
227 rollbacks: AtomicU64::new(0),
228 })
229 }
230
231 pub fn prepare(
233 &self,
234 generation: u64,
235 host: RuntimeHttpHost,
236 ) -> Result<PreparedRuntimeHttpGeneration, RuntimeHttpReloadError> {
237 let active = self.routing.active.load();
238 if generation <= active.id {
239 return Err(RuntimeHttpReloadError::StaleGeneration);
240 }
241 if !host.config().enabled {
242 return Err(RuntimeHttpReloadError::ListenerDisabled);
243 }
244 if host.config().host != self.config.host || host.config().port != self.config.port {
245 return Err(RuntimeHttpReloadError::ListenerAddressChanged);
246 }
247 Ok(PreparedRuntimeHttpGeneration {
248 generation: Arc::new(RoutingGeneration::new(generation, host.router())),
249 })
250 }
251
252 #[cfg(test)]
253 pub(super) fn new_for_test(
254 initial_generation: u64,
255 config: super::HttpApiConfig,
256 router: Router,
257 ) -> Self {
258 let initial = Arc::new(RoutingGeneration::new(initial_generation, router));
259 Self {
260 config,
261 routing: Arc::new(RoutingTable {
262 active: ArcSwap::from(initial),
263 }),
264 reload_in_progress: AtomicBool::new(false),
265 successful_reloads: AtomicU64::new(0),
266 failed_reloads: AtomicU64::new(0),
267 rollbacks: AtomicU64::new(0),
268 }
269 }
270
271 #[cfg(test)]
272 pub(super) fn prepare_router_for_test(
273 &self,
274 generation: u64,
275 router: Router,
276 ) -> PreparedRuntimeHttpGeneration {
277 PreparedRuntimeHttpGeneration {
278 generation: Arc::new(RoutingGeneration::new(generation, router)),
279 }
280 }
281
282 pub fn router(&self) -> Router {
284 dynamic_router(Arc::clone(&self.routing))
285 }
286
287 pub fn run_until_shutdown(&self, shutdown: Arc<AtomicBool>) -> io::Result<()> {
289 let address = format!("{}:{}", self.config.host, self.config.port);
290 let runtime = tokio::runtime::Builder::new_current_thread()
291 .enable_all()
292 .build()
293 .map_err(io::Error::other)?;
294 runtime.block_on(async move {
295 let listener = tokio::net::TcpListener::bind(address).await?;
296 serve_listener(listener, self.router(), shutdown).await
297 })
298 }
299
300 pub fn run_on_listener_until_shutdown(
302 &self,
303 listener: std::net::TcpListener,
304 shutdown: Arc<AtomicBool>,
305 ) -> io::Result<()> {
306 listener.set_nonblocking(true)?;
307 let router = self.router();
308 let runtime = tokio::runtime::Builder::new_current_thread()
309 .enable_all()
310 .build()
311 .map_err(io::Error::other)?;
312 runtime.block_on(async move {
313 let listener = tokio::net::TcpListener::from_std(listener)?;
314 serve_listener(listener, router, shutdown).await
315 })
316 }
317
318 pub async fn reload(
320 &self,
321 prepared: PreparedRuntimeHttpGeneration,
322 policy: HttpReloadPolicy,
323 ) -> Result<(), RuntimeHttpReloadError> {
324 let _guard = match ReloadGuard::acquire(&self.reload_in_progress) {
325 Ok(guard) => guard,
326 Err(error) => return self.fail(error),
327 };
328 let previous = self.routing.active.load_full();
329 if prepared.generation.id <= previous.id {
330 return self.fail(RuntimeHttpReloadError::StaleGeneration);
331 }
332 if !probe_health(&prepared.generation.router, policy.health_timeout).await {
333 return self.fail(RuntimeHttpReloadError::HealthGateFailed(
334 HttpReloadPhase::Prepare,
335 ));
336 }
337
338 previous.accepting.store(false, Ordering::Release);
339 self.routing.active.store(Arc::clone(&prepared.generation));
340 if !probe_health(&prepared.generation.router, policy.health_timeout).await {
341 return self
342 .rollback(previous, prepared.generation, policy)
343 .await
344 .and(Err(RuntimeHttpReloadError::HealthGateFailed(
345 HttpReloadPhase::Switch,
346 )));
347 }
348 if !wait_for_drain(&previous, policy.drain_timeout).await {
349 return self
350 .rollback(previous, prepared.generation, policy)
351 .await
352 .and(Err(RuntimeHttpReloadError::DrainTimedOut));
353 }
354 increment(&self.successful_reloads);
355 Ok(())
356 }
357
358 pub fn snapshot(&self) -> HttpReloadSnapshot {
360 let active = self.routing.active.load();
361 HttpReloadSnapshot {
362 active_generation: active.id,
363 active_inflight: active.inflight.load(Ordering::Acquire),
364 reload_in_progress: self.reload_in_progress.load(Ordering::Acquire),
365 successful_reloads: self.successful_reloads.load(Ordering::Relaxed),
366 failed_reloads: self.failed_reloads.load(Ordering::Relaxed),
367 rollbacks: self.rollbacks.load(Ordering::Relaxed),
368 }
369 }
370
371 async fn rollback(
372 &self,
373 previous: Arc<RoutingGeneration>,
374 failed: Arc<RoutingGeneration>,
375 policy: HttpReloadPolicy,
376 ) -> Result<(), RuntimeHttpReloadError> {
377 failed.accepting.store(false, Ordering::Release);
378 previous.accepting.store(true, Ordering::Release);
379 self.routing.active.store(previous);
380 increment(&self.rollbacks);
381 increment(&self.failed_reloads);
382 if wait_for_drain(&failed, policy.drain_timeout).await {
383 Ok(())
384 } else {
385 Err(RuntimeHttpReloadError::RollbackDrainTimedOut)
386 }
387 }
388
389 fn fail<T>(&self, error: RuntimeHttpReloadError) -> Result<T, RuntimeHttpReloadError> {
390 increment(&self.failed_reloads);
391 Err(error)
392 }
393}
394
395struct ReloadGuard<'a> {
396 flag: &'a AtomicBool,
397}
398
399impl<'a> ReloadGuard<'a> {
400 fn acquire(flag: &'a AtomicBool) -> Result<Self, RuntimeHttpReloadError> {
401 flag.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
402 .map_err(|_| RuntimeHttpReloadError::ReloadInProgress)?;
403 Ok(Self { flag })
404 }
405}
406
407impl Drop for ReloadGuard<'_> {
408 fn drop(&mut self) {
409 self.flag.store(false, Ordering::Release);
410 }
411}
412
413fn dynamic_router(routing: Arc<RoutingTable>) -> Router {
414 Router::new()
415 .fallback(any(dispatch_active_generation))
416 .with_state(routing)
417}
418
419async fn dispatch_active_generation(
420 State(routing): State<Arc<RoutingTable>>,
421 request: Request,
422) -> Response<Body> {
423 loop {
424 let generation = routing.active.load_full();
425 if let Some(_permit) = generation.try_admit() {
426 return generation
427 .router
428 .clone()
429 .oneshot(request)
430 .await
431 .unwrap_or_else(|never| match never {});
432 }
433 tokio::time::sleep(ROUTING_RETRY_DELAY).await;
434 }
435}
436
437async fn probe_health(router: &Router, timeout: Duration) -> bool {
438 let request = Request::get("/v1/health").body(Body::empty());
439 let Ok(request) = request else {
440 return false;
441 };
442 tokio::time::timeout(timeout, router.clone().oneshot(request))
443 .await
444 .ok()
445 .and_then(Result::ok)
446 .is_some_and(|response| response.status() == StatusCode::OK)
447}
448
449async fn wait_for_drain(generation: &RoutingGeneration, timeout: Duration) -> bool {
450 let deadline = Instant::now() + timeout;
451 loop {
452 if generation.inflight.load(Ordering::Acquire) == 0 {
453 return true;
454 }
455 if Instant::now() >= deadline {
456 return false;
457 }
458 tokio::time::sleep(ROUTING_RETRY_DELAY).await;
459 }
460}
461
462async fn serve_listener(
463 listener: tokio::net::TcpListener,
464 router: Router,
465 shutdown: Arc<AtomicBool>,
466) -> io::Result<()> {
467 axum::serve(listener, router)
468 .with_graceful_shutdown(super::wait_for_shutdown(shutdown))
469 .await
470}
471
472fn increment(counter: &AtomicU64) {
473 let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
474 Some(current.saturating_add(1))
475 });
476}