Skip to main content

appcore_api/http/
reload.rs

1//! Atomic HTTP routing-generation reload with bounded drain and rollback.
2
3use 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/// Reload phase associated with a controlled failure.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum HttpReloadPhase {
23    /// Candidate validation and health before activation.
24    Prepare,
25    /// Atomic routing activation and health confirmation.
26    Switch,
27    /// Bounded completion of requests admitted by the old generation.
28    Drain,
29    /// Restoration and drain after a failed activation.
30    Rollback,
31}
32
33/// Bounded health and drain policy for one routing reload.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct HttpReloadPolicy {
36    health_timeout: Duration,
37    drain_timeout: Duration,
38}
39
40impl HttpReloadPolicy {
41    /// Creates a policy with non-zero health and drain deadlines.
42    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    /// Returns the candidate health deadline.
60    pub fn health_timeout(self) -> Duration {
61        self.health_timeout
62    }
63
64    /// Returns the old-generation drain deadline.
65    pub fn drain_timeout(self) -> Duration {
66        self.drain_timeout
67    }
68}
69
70/// Payload-free state and counters for a reloadable HTTP host.
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub struct HttpReloadSnapshot {
73    /// Generation currently selected for new requests.
74    pub active_generation: u64,
75    /// Requests currently executing on the selected generation.
76    pub active_inflight: usize,
77    /// Whether one prepare/switch/drain transaction is active.
78    pub reload_in_progress: bool,
79    /// Reloads that switched and drained successfully.
80    pub successful_reloads: u64,
81    /// Failed reload attempts, including controlled rollbacks.
82    pub failed_reloads: u64,
83    /// Switches restored to the prior generation.
84    pub rollbacks: u64,
85}
86
87/// Controlled, redacted reload failure.
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum RuntimeHttpReloadError {
90    /// Reload requires an enabled listener.
91    ListenerDisabled,
92    /// A candidate attempted to change the listener address in-place.
93    ListenerAddressChanged,
94    /// Generation identifiers must increase monotonically.
95    StaleGeneration,
96    /// Another reload transaction already owns the coordinator.
97    ReloadInProgress,
98    /// Health or drain deadlines must be non-zero.
99    InvalidPolicy,
100    /// The candidate or active generation failed its health gate.
101    HealthGateFailed(HttpReloadPhase),
102    /// The outgoing generation did not drain before rollback.
103    DrainTimedOut,
104    /// The failed generation remained active past the rollback drain deadline.
105    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
181/// Prepared, health-gated candidate that can be consumed by one reload.
182pub 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
195/// HTTP host whose stable listener dispatches through one atomic generation.
196pub 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    /// Creates a reloadable host from an already composed initial host.
207    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    /// Validates and owns a candidate without changing live routing.
232    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    /// Returns a router that always dispatches through the active generation.
283    pub fn router(&self) -> Router {
284        dynamic_router(Arc::clone(&self.routing))
285    }
286
287    /// Runs the stable listener until cooperative shutdown is requested.
288    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    /// Runs on a listener that the composition root already bound and checked.
301    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    /// Health-checks, atomically activates, and drains one prepared generation.
319    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    /// Returns the active generation and bounded operational counters.
359    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}