zentinel_proxy/upstream/
drain.rs

1//! Backend drain lifecycle tracking.
2//!
3//! Monitors old upstream pools after a config reload, tracking active
4//! connections and emitting structured events when backends are fully drained.
5//!
6//! Lifecycle states:
7//! - `Active`: backend is receiving new connections
8//! - `Draining`: backend removed from config, existing connections finishing
9//! - `Drained`: all connections completed, safe to terminate backend
10//!
11//! This replaces the previous blind 60-second sleep with active monitoring.
12
13use std::collections::HashMap;
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16use tracing::{debug, info, warn};
17
18use super::UpstreamPool;
19
20/// Backend drain lifecycle state.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum BackendState {
23    /// Backend is active and receiving new connections.
24    Active,
25    /// Backend has been removed from config. No new connections are being
26    /// sent, but existing connections are still in flight.
27    Draining,
28    /// All connections have completed. Safe to shut down the backend.
29    Drained,
30}
31
32impl std::fmt::Display for BackendState {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        match self {
35            BackendState::Active => write!(f, "active"),
36            BackendState::Draining => write!(f, "draining"),
37            BackendState::Drained => write!(f, "drained"),
38        }
39    }
40}
41
42/// Tracks the drain lifecycle of old upstream pools after a config reload.
43///
44/// When pools are replaced during a reload, the old pools are handed to
45/// the `DrainTracker` which monitors their active request count and
46/// emits structured log events when they transition through the lifecycle.
47pub struct DrainTracker {
48    /// Maximum time to wait for drain before force-shutting down.
49    max_drain_time: Duration,
50    /// Poll interval for checking connection counts.
51    poll_interval: Duration,
52}
53
54impl DrainTracker {
55    pub fn new(max_drain_time: Duration, poll_interval: Duration) -> Self {
56        Self {
57            max_drain_time,
58            poll_interval,
59        }
60    }
61
62    /// Monitor old pools and emit drain lifecycle events.
63    ///
64    /// This runs as a spawned task. It polls each pool's active request count
65    /// and emits structured events as they transition from Draining to Drained.
66    pub async fn track_pools(&self, pools: HashMap<String, Arc<UpstreamPool>>) {
67        if pools.is_empty() {
68            return;
69        }
70
71        let pool_count = pools.len();
72        info!(
73            pool_count = pool_count,
74            "Starting drain tracking for removed upstream pools"
75        );
76
77        // Emit Draining events
78        for (name, pool) in &pools {
79            let active = pool.active_request_count();
80            info!(
81                upstream_id = %name,
82                active_requests = active,
83                state = %BackendState::Draining,
84                "Backend entering drain state"
85            );
86        }
87
88        let start = Instant::now();
89        let mut pending: HashMap<String, Arc<UpstreamPool>> = pools;
90
91        while !pending.is_empty() && start.elapsed() < self.max_drain_time {
92            tokio::time::sleep(self.poll_interval).await;
93
94            let mut newly_drained = Vec::new();
95
96            for (name, pool) in &pending {
97                let active = pool.active_request_count();
98
99                if active == 0 {
100                    let drain_duration = start.elapsed();
101                    info!(
102                        upstream_id = %name,
103                        drain_duration_ms = drain_duration.as_millis(),
104                        drain_duration_secs = drain_duration.as_secs_f64(),
105                        state = %BackendState::Drained,
106                        "Backend fully drained, safe to terminate"
107                    );
108                    newly_drained.push(name.clone());
109                } else {
110                    debug!(
111                        upstream_id = %name,
112                        active_requests = active,
113                        elapsed_ms = start.elapsed().as_millis(),
114                        state = %BackendState::Draining,
115                        "Backend still draining"
116                    );
117                }
118            }
119
120            for name in newly_drained {
121                if let Some(pool) = pending.remove(&name) {
122                    pool.shutdown().await;
123                }
124            }
125        }
126
127        // Force shutdown any remaining pools that didn't drain in time
128        for (name, pool) in &pending {
129            let active = pool.active_request_count();
130            warn!(
131                upstream_id = %name,
132                active_requests = active,
133                max_drain_time_secs = self.max_drain_time.as_secs(),
134                state = "drain_timeout",
135                "Backend drain timeout exceeded, force shutting down"
136            );
137            pool.shutdown().await;
138        }
139
140        if pool_count > 0 {
141            info!(
142                pool_count = pool_count,
143                total_duration_ms = start.elapsed().as_millis(),
144                "Drain tracking complete for all removed pools"
145            );
146        }
147    }
148}
149
150impl Default for DrainTracker {
151    fn default() -> Self {
152        Self {
153            max_drain_time: Duration::from_secs(60),
154            poll_interval: Duration::from_secs(1),
155        }
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn backend_state_display() {
165        assert_eq!(BackendState::Active.to_string(), "active");
166        assert_eq!(BackendState::Draining.to_string(), "draining");
167        assert_eq!(BackendState::Drained.to_string(), "drained");
168    }
169}
zentinel_proxy/upstream/drain.rs

zentinel_proxy/upstream/
drain.rs
