Skip to main content

zentinel_proxy/upstream/
mod.rs

1//! Upstream pool management module for Zentinel proxy
2//!
3//! This module handles upstream server pools, load balancing, health checking,
4//! connection pooling, and retry logic with circuit breakers.
5
6use async_trait::async_trait;
7use pingora::upstreams::peer::HttpPeer;
8use rand::seq::IndexedRandom;
9use std::collections::HashMap;
10use std::net::ToSocketAddrs;
11use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
12use std::sync::Arc;
13use std::time::Duration;
14use tokio::sync::RwLock;
15use tracing::{debug, error, info, trace, warn};
16
17use zentinel_common::{
18    errors::{ZentinelError, ZentinelResult},
19    types::{CircuitBreakerConfig, LoadBalancingAlgorithm},
20    CircuitBreaker, UpstreamId,
21};
22use zentinel_config::{UpstreamConfig, UpstreamDiscovery};
23
24use pingora_core::protocols::l4::socket::SocketAddr as PingoraSocketAddr;
25use pingora_load_balancing::discovery::ServiceDiscovery;
26use pingora_load_balancing::Backend;
27use std::collections::BTreeSet;
28
29use crate::discovery::DiscoveryConfig;
30
31// ============================================================================
32// Internal Upstream Target Type
33// ============================================================================
34
35/// Internal upstream target representation for load balancers
36///
37/// This is a simplified representation used internally by load balancers,
38/// separate from the user-facing config UpstreamTarget.
39#[derive(Debug, Clone)]
40pub struct UpstreamTarget {
41    /// Target IP address or hostname
42    pub address: String,
43    /// Target port
44    pub port: u16,
45    /// Weight for weighted load balancing
46    pub weight: u32,
47}
48
49impl UpstreamTarget {
50    /// Create a new upstream target
51    pub fn new(address: impl Into<String>, port: u16, weight: u32) -> Self {
52        Self {
53            address: address.into(),
54            port,
55            weight,
56        }
57    }
58
59    /// Create from a "host:port" string with default weight
60    pub fn from_address(addr: &str) -> Option<Self> {
61        let parts: Vec<&str> = addr.rsplitn(2, ':').collect();
62        if parts.len() == 2 {
63            let port = parts[0].parse().ok()?;
64            let address = parts[1].to_string();
65            Some(Self {
66                address,
67                port,
68                weight: 100,
69            })
70        } else {
71            None
72        }
73    }
74
75    /// Convert from config UpstreamTarget
76    pub fn from_config(config: &zentinel_config::UpstreamTarget) -> Option<Self> {
77        Self::from_address(&config.address).map(|mut t| {
78            t.weight = config.weight;
79            t
80        })
81    }
82
83    /// Get the full address string
84    pub fn full_address(&self) -> String {
85        format!("{}:{}", self.address, self.port)
86    }
87}
88
89// ============================================================================
90// Load Balancing
91// ============================================================================
92
93// Load balancing algorithm implementations
94pub mod adaptive;
95pub mod consistent_hash;
96pub mod discovery_refresh;
97pub mod drain;
98pub mod health;
99pub mod inference_health;
100pub mod least_tokens;
101pub mod locality;
102pub mod maglev;
103pub mod p2c;
104pub mod peak_ewma;
105pub mod sticky_session;
106pub mod subset;
107pub mod weighted_least_conn;
108
109// Re-export commonly used types from sub-modules
110pub use adaptive::{AdaptiveBalancer, AdaptiveConfig};
111pub use consistent_hash::{ConsistentHashBalancer, ConsistentHashConfig};
112pub use health::{ActiveHealthChecker, HealthCheckRunner};
113pub use inference_health::InferenceHealthCheck;
114pub use least_tokens::{
115    LeastTokensQueuedBalancer, LeastTokensQueuedConfig, LeastTokensQueuedTargetStats,
116};
117pub use locality::{LocalityAwareBalancer, LocalityAwareConfig};
118pub use maglev::{MaglevBalancer, MaglevConfig};
119pub use p2c::{P2cBalancer, P2cConfig};
120pub use peak_ewma::{PeakEwmaBalancer, PeakEwmaConfig};
121pub use sticky_session::{StickySessionBalancer, StickySessionRuntimeConfig};
122pub use subset::{SubsetBalancer, SubsetConfig};
123pub use weighted_least_conn::{WeightedLeastConnBalancer, WeightedLeastConnConfig};
124
125/// Request context for load balancer decisions
126#[derive(Debug, Clone)]
127pub struct RequestContext {
128    pub client_ip: Option<std::net::SocketAddr>,
129    pub headers: HashMap<String, String>,
130    pub path: String,
131    pub method: String,
132}
133
134/// Load balancer trait for different algorithms
135#[async_trait]
136pub trait LoadBalancer: Send + Sync {
137    /// Select next upstream target
138    async fn select(&self, context: Option<&RequestContext>) -> ZentinelResult<TargetSelection>;
139
140    /// The key this balancer signs session-affinity cookies with, if it signs
141    /// any.
142    ///
143    /// Used to assert that rebuilding a pool — which service discovery does
144    /// whenever the target set changes — does not rotate the key and invalidate
145    /// every cookie already issued.
146    fn session_signing_key(&self) -> Option<[u8; 32]> {
147        None
148    }
149
150    /// Report target health status
151    async fn report_health(&self, address: &str, healthy: bool);
152
153    /// Get all healthy targets
154    async fn healthy_targets(&self) -> Vec<String>;
155
156    /// Release connection (for connection tracking)
157    async fn release(&self, _selection: &TargetSelection) {
158        // Default implementation - no-op
159    }
160
161    /// Report request result (for adaptive algorithms)
162    async fn report_result(
163        &self,
164        _selection: &TargetSelection,
165        _success: bool,
166        _latency: Option<Duration>,
167    ) {
168        // Default implementation - no-op
169    }
170
171    /// Report request result by address with latency (for adaptive algorithms)
172    ///
173    /// This method allows reporting results without needing the full TargetSelection,
174    /// which is useful when the selection is not available (e.g., in logging callback).
175    /// The default implementation just calls report_health; adaptive balancers override
176    /// this to update their metrics.
177    async fn report_result_with_latency(
178        &self,
179        address: &str,
180        success: bool,
181        _latency: Option<Duration>,
182    ) {
183        // Default implementation - just report health
184        self.report_health(address, success).await;
185    }
186}
187
188/// Selected upstream target
189#[derive(Debug, Clone)]
190pub struct TargetSelection {
191    /// Target address
192    pub address: String,
193    /// Target weight
194    pub weight: u32,
195    /// Target metadata
196    pub metadata: HashMap<String, String>,
197}
198
199/// Upstream pool managing multiple backend servers
200pub struct UpstreamPool {
201    /// Pool identifier
202    id: UpstreamId,
203    /// Configured targets
204    targets: Vec<UpstreamTarget>,
205    /// Load balancer implementation
206    load_balancer: Arc<dyn LoadBalancer>,
207    /// Connection pool configuration (Pingora handles actual pooling)
208    pool_config: ConnectionPoolConfig,
209    /// HTTP version configuration
210    http_version: HttpVersionOptions,
211    /// Whether TLS is enabled for this upstream
212    tls_enabled: bool,
213    /// SNI for TLS connections
214    tls_sni: Option<String>,
215    /// TLS configuration for upstream mTLS (client certificates)
216    tls_config: Option<zentinel_config::UpstreamTlsConfig>,
217    /// Circuit breakers per target
218    circuit_breakers: Arc<RwLock<HashMap<String, CircuitBreaker>>>,
219    /// Pool statistics
220    stats: Arc<PoolStats>,
221    /// Resolved discovery source, kept so a refresh can re-resolve without
222    /// rebuilding the client (which for Consul and Kubernetes means not
223    /// re-establishing a connection every interval).
224    discovery: Option<Arc<dyn ServiceDiscovery + Send + Sync>>,
225    /// How often `discovery` is re-resolved. `None` when there is no discovery
226    /// source, or when it is one that cannot change (`static`).
227    discovery_interval: Option<Duration>,
228    /// Targets that came from the configuration rather than from discovery.
229    ///
230    /// Kept separately because a refresh replaces the discovered targets and
231    /// must not drop the configured ones alongside them.
232    static_targets: Vec<UpstreamTarget>,
233    /// The configuration this pool was built from, needed to rebuild the load
234    /// balancer when discovery changes the target set.
235    config: UpstreamConfig,
236    /// Sticky-session runtime config, carrying the HMAC key that signs affinity
237    /// cookies.
238    ///
239    /// Held on the pool rather than derived per balancer because the key is
240    /// generated randomly: rebuilding it on a discovery refresh would invalidate
241    /// every outstanding affinity cookie, resetting all sessions each time a
242    /// backend appeared or disappeared.
243    sticky_runtime: Option<StickySessionRuntimeConfig>,
244}
245
246// Note: Active health checking is handled by the PassiveHealthChecker in health.rs
247// and via load balancer health reporting. A future enhancement could add active
248// HTTP/TCP health probes here.
249
250/// Connection pool configuration for Pingora's built-in pooling
251///
252/// Note: Actual connection pooling is handled by Pingora internally.
253/// This struct holds configuration that is applied to peer options.
254#[derive(Clone)]
255pub struct ConnectionPoolConfig {
256    /// Maximum connections per target (informational - Pingora manages actual pooling)
257    pub max_connections: usize,
258    /// Maximum idle connections (informational - Pingora manages actual pooling)
259    pub max_idle: usize,
260    /// Maximum idle timeout for pooled connections
261    pub idle_timeout: Duration,
262    /// Maximum connection lifetime (None = unlimited)
263    pub max_lifetime: Option<Duration>,
264    /// Connection timeout
265    pub connection_timeout: Duration,
266    /// Read timeout
267    pub read_timeout: Duration,
268    /// Write timeout
269    pub write_timeout: Duration,
270}
271
272/// HTTP version configuration for upstream connections
273#[derive(Clone)]
274pub struct HttpVersionOptions {
275    /// Minimum HTTP version (1 or 2)
276    pub min_version: u8,
277    /// Maximum HTTP version (1 or 2)
278    pub max_version: u8,
279    /// H2 ping interval (0 to disable)
280    pub h2_ping_interval: Duration,
281    /// Maximum concurrent H2 streams per connection
282    pub max_h2_streams: usize,
283}
284
285impl ConnectionPoolConfig {
286    /// Create from upstream config
287    pub fn from_config(
288        pool_config: &zentinel_config::ConnectionPoolConfig,
289        timeouts: &zentinel_config::UpstreamTimeouts,
290    ) -> Self {
291        Self {
292            max_connections: pool_config.max_connections,
293            max_idle: pool_config.max_idle,
294            idle_timeout: Duration::from_secs(pool_config.idle_timeout_secs),
295            max_lifetime: pool_config.max_lifetime_secs.map(Duration::from_secs),
296            connection_timeout: Duration::from_secs(timeouts.connect_secs),
297            read_timeout: Duration::from_secs(timeouts.read_secs),
298            write_timeout: Duration::from_secs(timeouts.write_secs),
299        }
300    }
301}
302
303// CircuitBreaker is imported from zentinel_common
304
305/// Pool statistics
306#[derive(Default)]
307pub struct PoolStats {
308    /// Total requests
309    pub requests: AtomicU64,
310    /// Successful requests
311    pub successes: AtomicU64,
312    /// Failed requests
313    pub failures: AtomicU64,
314    /// Retried requests
315    pub retries: AtomicU64,
316    /// Circuit breaker trips
317    pub circuit_breaker_trips: AtomicU64,
318    /// Currently active requests (in-flight)
319    pub active_requests: AtomicU64,
320}
321
322/// Target information for shadow traffic
323#[derive(Debug, Clone)]
324pub struct ShadowTarget {
325    /// URL scheme (http or https)
326    pub scheme: String,
327    /// Target host
328    pub host: String,
329    /// Target port
330    pub port: u16,
331    /// SNI for TLS connections
332    pub sni: Option<String>,
333}
334
335impl ShadowTarget {
336    /// Build URL from target info and path
337    pub fn build_url(&self, path: &str) -> String {
338        let port_suffix = match (self.scheme.as_str(), self.port) {
339            ("http", 80) | ("https", 443) => String::new(),
340            _ => format!(":{}", self.port),
341        };
342        format!("{}://{}{}{}", self.scheme, self.host, port_suffix, path)
343    }
344}
345
346/// Snapshot of pool configuration for metrics/debugging
347#[derive(Debug, Clone)]
348pub struct PoolConfigSnapshot {
349    /// Maximum connections per target
350    pub max_connections: usize,
351    /// Maximum idle connections
352    pub max_idle: usize,
353    /// Idle timeout in seconds
354    pub idle_timeout_secs: u64,
355    /// Maximum connection lifetime in seconds (None = unlimited)
356    pub max_lifetime_secs: Option<u64>,
357    /// Connection timeout in seconds
358    pub connection_timeout_secs: u64,
359    /// Read timeout in seconds
360    pub read_timeout_secs: u64,
361    /// Write timeout in seconds
362    pub write_timeout_secs: u64,
363}
364
365/// Round-robin load balancer
366struct RoundRobinBalancer {
367    targets: Vec<UpstreamTarget>,
368    current: AtomicUsize,
369    health_status: Arc<RwLock<HashMap<String, bool>>>,
370}
371
372impl RoundRobinBalancer {
373    fn new(targets: Vec<UpstreamTarget>) -> Self {
374        let mut health_status = HashMap::new();
375        for target in &targets {
376            health_status.insert(target.full_address(), true);
377        }
378
379        Self {
380            targets,
381            current: AtomicUsize::new(0),
382            health_status: Arc::new(RwLock::new(health_status)),
383        }
384    }
385}
386
387#[async_trait]
388impl LoadBalancer for RoundRobinBalancer {
389    async fn select(&self, _context: Option<&RequestContext>) -> ZentinelResult<TargetSelection> {
390        trace!(
391            total_targets = self.targets.len(),
392            algorithm = "round_robin",
393            "Selecting upstream target"
394        );
395
396        let health = self.health_status.read().await;
397        let healthy_targets: Vec<_> = self
398            .targets
399            .iter()
400            .filter(|t| *health.get(&t.full_address()).unwrap_or(&true))
401            .collect();
402
403        if healthy_targets.is_empty() {
404            warn!(
405                total_targets = self.targets.len(),
406                algorithm = "round_robin",
407                "No healthy upstream targets available"
408            );
409            return Err(ZentinelError::NoHealthyUpstream);
410        }
411
412        let index = self.current.fetch_add(1, Ordering::Relaxed) % healthy_targets.len();
413        let target = healthy_targets[index];
414
415        trace!(
416            selected_target = %target.full_address(),
417            healthy_count = healthy_targets.len(),
418            index = index,
419            algorithm = "round_robin",
420            "Selected target via round robin"
421        );
422
423        Ok(TargetSelection {
424            address: target.full_address(),
425            weight: target.weight,
426            metadata: HashMap::new(),
427        })
428    }
429
430    async fn report_health(&self, address: &str, healthy: bool) {
431        trace!(
432            target = %address,
433            healthy = healthy,
434            algorithm = "round_robin",
435            "Updating target health status"
436        );
437        self.health_status
438            .write()
439            .await
440            .insert(address.to_string(), healthy);
441    }
442
443    async fn healthy_targets(&self) -> Vec<String> {
444        self.health_status
445            .read()
446            .await
447            .iter()
448            .filter_map(|(addr, &healthy)| if healthy { Some(addr.clone()) } else { None })
449            .collect()
450    }
451}
452
453/// Random load balancer - true random selection among healthy targets
454struct RandomBalancer {
455    targets: Vec<UpstreamTarget>,
456    health_status: Arc<RwLock<HashMap<String, bool>>>,
457}
458
459impl RandomBalancer {
460    fn new(targets: Vec<UpstreamTarget>) -> Self {
461        let mut health_status = HashMap::new();
462        for target in &targets {
463            health_status.insert(target.full_address(), true);
464        }
465
466        Self {
467            targets,
468            health_status: Arc::new(RwLock::new(health_status)),
469        }
470    }
471}
472
473#[async_trait]
474impl LoadBalancer for RandomBalancer {
475    async fn select(&self, _context: Option<&RequestContext>) -> ZentinelResult<TargetSelection> {
476        use rand::seq::SliceRandom;
477
478        trace!(
479            total_targets = self.targets.len(),
480            algorithm = "random",
481            "Selecting upstream target"
482        );
483
484        let health = self.health_status.read().await;
485        let healthy_targets: Vec<_> = self
486            .targets
487            .iter()
488            .filter(|t| *health.get(&t.full_address()).unwrap_or(&true))
489            .collect();
490
491        if healthy_targets.is_empty() {
492            warn!(
493                total_targets = self.targets.len(),
494                algorithm = "random",
495                "No healthy upstream targets available"
496            );
497            return Err(ZentinelError::NoHealthyUpstream);
498        }
499
500        let mut rng = rand::rng();
501        let target = healthy_targets
502            .choose(&mut rng)
503            .ok_or(ZentinelError::NoHealthyUpstream)?;
504
505        trace!(
506            selected_target = %target.full_address(),
507            healthy_count = healthy_targets.len(),
508            algorithm = "random",
509            "Selected target via random selection"
510        );
511
512        Ok(TargetSelection {
513            address: target.full_address(),
514            weight: target.weight,
515            metadata: HashMap::new(),
516        })
517    }
518
519    async fn report_health(&self, address: &str, healthy: bool) {
520        trace!(
521            target = %address,
522            healthy = healthy,
523            algorithm = "random",
524            "Updating target health status"
525        );
526        self.health_status
527            .write()
528            .await
529            .insert(address.to_string(), healthy);
530    }
531
532    async fn healthy_targets(&self) -> Vec<String> {
533        self.health_status
534            .read()
535            .await
536            .iter()
537            .filter_map(|(addr, &healthy)| if healthy { Some(addr.clone()) } else { None })
538            .collect()
539    }
540}
541
542/// Least connections load balancer
543struct LeastConnectionsBalancer {
544    targets: Vec<UpstreamTarget>,
545    connections: Arc<RwLock<HashMap<String, usize>>>,
546    health_status: Arc<RwLock<HashMap<String, bool>>>,
547}
548
549impl LeastConnectionsBalancer {
550    fn new(targets: Vec<UpstreamTarget>) -> Self {
551        let mut health_status = HashMap::new();
552        let mut connections = HashMap::new();
553
554        for target in &targets {
555            let addr = target.full_address();
556            health_status.insert(addr.clone(), true);
557            connections.insert(addr, 0);
558        }
559
560        Self {
561            targets,
562            connections: Arc::new(RwLock::new(connections)),
563            health_status: Arc::new(RwLock::new(health_status)),
564        }
565    }
566}
567
568#[async_trait]
569impl LoadBalancer for LeastConnectionsBalancer {
570    async fn select(&self, _context: Option<&RequestContext>) -> ZentinelResult<TargetSelection> {
571        trace!(
572            total_targets = self.targets.len(),
573            algorithm = "least_connections",
574            "Selecting upstream target"
575        );
576
577        let health = self.health_status.read().await;
578        let conns = self.connections.read().await;
579
580        let mut best_target = None;
581        let mut min_connections = usize::MAX;
582
583        for target in &self.targets {
584            let addr = target.full_address();
585            if !*health.get(&addr).unwrap_or(&true) {
586                trace!(
587                    target = %addr,
588                    algorithm = "least_connections",
589                    "Skipping unhealthy target"
590                );
591                continue;
592            }
593
594            let conn_count = *conns.get(&addr).unwrap_or(&0);
595            trace!(
596                target = %addr,
597                connections = conn_count,
598                "Evaluating target connection count"
599            );
600            if conn_count < min_connections {
601                min_connections = conn_count;
602                best_target = Some(target);
603            }
604        }
605
606        match best_target {
607            Some(target) => {
608                trace!(
609                    selected_target = %target.full_address(),
610                    connections = min_connections,
611                    algorithm = "least_connections",
612                    "Selected target with fewest connections"
613                );
614                Ok(TargetSelection {
615                    address: target.full_address(),
616                    weight: target.weight,
617                    metadata: HashMap::new(),
618                })
619            }
620            None => {
621                warn!(
622                    total_targets = self.targets.len(),
623                    algorithm = "least_connections",
624                    "No healthy upstream targets available"
625                );
626                Err(ZentinelError::NoHealthyUpstream)
627            }
628        }
629    }
630
631    async fn report_health(&self, address: &str, healthy: bool) {
632        trace!(
633            target = %address,
634            healthy = healthy,
635            algorithm = "least_connections",
636            "Updating target health status"
637        );
638        self.health_status
639            .write()
640            .await
641            .insert(address.to_string(), healthy);
642    }
643
644    async fn healthy_targets(&self) -> Vec<String> {
645        self.health_status
646            .read()
647            .await
648            .iter()
649            .filter_map(|(addr, &healthy)| if healthy { Some(addr.clone()) } else { None })
650            .collect()
651    }
652}
653
654/// Weighted load balancer
655struct WeightedBalancer {
656    targets: Vec<UpstreamTarget>,
657    weights: Vec<u32>,
658    current_index: AtomicUsize,
659    health_status: Arc<RwLock<HashMap<String, bool>>>,
660}
661
662#[async_trait]
663impl LoadBalancer for WeightedBalancer {
664    async fn select(&self, _context: Option<&RequestContext>) -> ZentinelResult<TargetSelection> {
665        trace!(
666            total_targets = self.targets.len(),
667            algorithm = "weighted",
668            "Selecting upstream target"
669        );
670
671        let health = self.health_status.read().await;
672        let healthy: Vec<_> = self
673            .targets
674            .iter()
675            .enumerate()
676            .filter(|(_, t)| *health.get(&t.full_address()).unwrap_or(&true))
677            .map(|(i, _)| i)
678            .collect();
679
680        if healthy.is_empty() {
681            warn!(
682                total_targets = self.targets.len(),
683                algorithm = "weighted",
684                "No healthy upstream targets available"
685            );
686            return Err(ZentinelError::NoHealthyUpstream);
687        }
688
689        // Weighted round-robin: map request counter to a weighted slot.
690        // E.g. weights [70, 30] → total 100 → slots [0..70) → target 0, [70..100) → target 1
691        let total_weight: u32 = healthy
692            .iter()
693            .map(|&i| self.weights.get(i).copied().unwrap_or(1))
694            .sum();
695
696        if total_weight == 0 {
697            return Err(ZentinelError::NoHealthyUpstream);
698        }
699
700        let slot = (self.current_index.fetch_add(1, Ordering::Relaxed) as u32) % total_weight;
701        let mut cumulative = 0u32;
702        let mut target_idx = healthy[0];
703        for &i in &healthy {
704            let w = self.weights.get(i).copied().unwrap_or(1);
705            cumulative += w;
706            if slot < cumulative {
707                target_idx = i;
708                break;
709            }
710        }
711
712        let target = &self.targets[target_idx];
713        let weight = self.weights.get(target_idx).copied().unwrap_or(1);
714
715        trace!(
716            selected_target = %target.full_address(),
717            weight = weight,
718            healthy_count = healthy.len(),
719            algorithm = "weighted",
720            "Selected target via weighted round robin"
721        );
722
723        Ok(TargetSelection {
724            address: target.full_address(),
725            weight,
726            metadata: HashMap::new(),
727        })
728    }
729
730    async fn report_health(&self, address: &str, healthy: bool) {
731        trace!(
732            target = %address,
733            healthy = healthy,
734            algorithm = "weighted",
735            "Updating target health status"
736        );
737        self.health_status
738            .write()
739            .await
740            .insert(address.to_string(), healthy);
741    }
742
743    async fn healthy_targets(&self) -> Vec<String> {
744        self.health_status
745            .read()
746            .await
747            .iter()
748            .filter_map(|(addr, &healthy)| if healthy { Some(addr.clone()) } else { None })
749            .collect()
750    }
751}
752
753/// IP hash load balancer
754struct IpHashBalancer {
755    targets: Vec<UpstreamTarget>,
756    health_status: Arc<RwLock<HashMap<String, bool>>>,
757}
758
759#[async_trait]
760impl LoadBalancer for IpHashBalancer {
761    async fn select(&self, context: Option<&RequestContext>) -> ZentinelResult<TargetSelection> {
762        trace!(
763            total_targets = self.targets.len(),
764            algorithm = "ip_hash",
765            "Selecting upstream target"
766        );
767
768        let health = self.health_status.read().await;
769        let healthy_targets: Vec<_> = self
770            .targets
771            .iter()
772            .filter(|t| *health.get(&t.full_address()).unwrap_or(&true))
773            .collect();
774
775        if healthy_targets.is_empty() {
776            warn!(
777                total_targets = self.targets.len(),
778                algorithm = "ip_hash",
779                "No healthy upstream targets available"
780            );
781            return Err(ZentinelError::NoHealthyUpstream);
782        }
783
784        // Hash the client IP to select a target
785        let (hash, client_ip_str) = if let Some(ctx) = context {
786            if let Some(ip) = &ctx.client_ip {
787                use std::hash::{Hash, Hasher};
788                let mut hasher = std::collections::hash_map::DefaultHasher::new();
789                ip.hash(&mut hasher);
790                (hasher.finish(), Some(ip.to_string()))
791            } else {
792                (0, None)
793            }
794        } else {
795            (0, None)
796        };
797
798        let idx = (hash as usize) % healthy_targets.len();
799        let target = healthy_targets[idx];
800
801        trace!(
802            selected_target = %target.full_address(),
803            client_ip = client_ip_str.as_deref().unwrap_or("unknown"),
804            hash = hash,
805            index = idx,
806            healthy_count = healthy_targets.len(),
807            algorithm = "ip_hash",
808            "Selected target via IP hash"
809        );
810
811        Ok(TargetSelection {
812            address: target.full_address(),
813            weight: target.weight,
814            metadata: HashMap::new(),
815        })
816    }
817
818    async fn report_health(&self, address: &str, healthy: bool) {
819        trace!(
820            target = %address,
821            healthy = healthy,
822            algorithm = "ip_hash",
823            "Updating target health status"
824        );
825        self.health_status
826            .write()
827            .await
828            .insert(address.to_string(), healthy);
829    }
830
831    async fn healthy_targets(&self) -> Vec<String> {
832        self.health_status
833            .read()
834            .await
835            .iter()
836            .filter_map(|(addr, &healthy)| if healthy { Some(addr.clone()) } else { None })
837            .collect()
838    }
839}
840
841impl UpstreamPool {
842    /// Create new upstream pool from configuration
843    pub async fn new(config: UpstreamConfig) -> ZentinelResult<Self> {
844        let id = UpstreamId::new(&config.id);
845
846        info!(
847            upstream_id = %config.id,
848            target_count = config.targets.len(),
849            algorithm = ?config.load_balancing,
850            "Creating upstream pool"
851        );
852
853        // Convert config targets to internal targets
854        let static_targets: Vec<UpstreamTarget> = config
855            .targets
856            .iter()
857            .filter_map(UpstreamTarget::from_config)
858            .collect();
859
860        // Resolve the discovery source, if one is configured, before the pool
861        // starts serving traffic. Discovered targets are added to the
862        // configured ones rather than replacing them, so a fixed backend can be
863        // pinned alongside a discovered set.
864        let (discovery, discovery_interval, discovered) =
865            Self::resolve_discovery(&config, &id).await;
866
867        let targets = Self::merge_targets(&static_targets, discovered);
868
869        if targets.is_empty() {
870            if discovery.is_some() {
871                // The source answered, and answered with nothing. That is a
872                // live condition (a scaled-to-zero deployment, a service with
873                // no passing instances) rather than a broken configuration, so
874                // the pool starts empty and recovers on the next refresh.
875                // Requests to it fail with "no healthy targets" until then.
876                warn!(
877                    upstream_id = %config.id,
878                    discovery = config.discovery.as_ref().map(|d| d.kind()).unwrap_or("none"),
879                    "Service discovery returned no backends; upstream starts with no targets \
880                     and will recover on the next refresh"
881                );
882            } else {
883                error!(
884                    upstream_id = %config.id,
885                    "No valid upstream targets configured"
886                );
887                return Err(ZentinelError::Config {
888                    message: "No valid upstream targets".to_string(),
889                    source: None,
890                });
891            }
892        }
893
894        for target in &targets {
895            debug!(
896                upstream_id = %config.id,
897                target = %target.full_address(),
898                weight = target.weight,
899                "Registered upstream target"
900            );
901        }
902
903        // Create load balancer
904        debug!(
905            upstream_id = %config.id,
906            algorithm = ?config.load_balancing,
907            "Creating load balancer"
908        );
909        // Built once and kept, so a discovery refresh can rebuild the balancer
910        // without rotating the key that signs affinity cookies.
911        let sticky_runtime = config
912            .sticky_session
913            .as_ref()
914            .map(StickySessionRuntimeConfig::from_config);
915
916        let load_balancer = Self::create_load_balancer(
917            &config.load_balancing,
918            &targets,
919            &config,
920            sticky_runtime.as_ref(),
921        )?;
922
923        // Create connection pool configuration (Pingora handles actual pooling)
924        debug!(
925            upstream_id = %config.id,
926            max_connections = config.connection_pool.max_connections,
927            max_idle = config.connection_pool.max_idle,
928            idle_timeout_secs = config.connection_pool.idle_timeout_secs,
929            connect_timeout_secs = config.timeouts.connect_secs,
930            read_timeout_secs = config.timeouts.read_secs,
931            write_timeout_secs = config.timeouts.write_secs,
932            "Creating connection pool configuration"
933        );
934        let pool_config =
935            ConnectionPoolConfig::from_config(&config.connection_pool, &config.timeouts);
936
937        // Create HTTP version configuration
938        let http_version = HttpVersionOptions {
939            min_version: config.http_version.min_version,
940            max_version: config.http_version.max_version,
941            h2_ping_interval: if config.http_version.h2_ping_interval_secs > 0 {
942                Duration::from_secs(config.http_version.h2_ping_interval_secs)
943            } else {
944                Duration::ZERO
945            },
946            max_h2_streams: config.http_version.max_h2_streams,
947        };
948
949        // TLS configuration
950        let tls_enabled = config.tls.is_some();
951        let tls_sni = config.tls.as_ref().and_then(|t| t.sni.clone());
952        let tls_config = config.tls.clone();
953
954        // Log mTLS configuration if present
955        if let Some(ref tls) = tls_config {
956            if tls.client_cert.is_some() {
957                info!(
958                    upstream_id = %config.id,
959                    "mTLS enabled for upstream (client certificate configured)"
960                );
961            }
962        }
963
964        if http_version.max_version >= 2 && tls_enabled {
965            info!(
966                upstream_id = %config.id,
967                "HTTP/2 enabled for upstream (via ALPN)"
968            );
969        }
970
971        // Initialize circuit breakers for each target
972
973        // Assigns default CB config if not configured, such as when the stanza is missing
974        // (and None is set for CircuitBreakerConfig)
975        let cb_config = config.circuit_breaker.unwrap_or_default();
976
977        let mut circuit_breakers = HashMap::new();
978        for target in &targets {
979            trace!(
980                upstream_id = %config.id,
981                target = %target.full_address(),
982                "Initializing circuit breaker for target, configuration {:?}",
983                cb_config
984            );
985
986            circuit_breakers.insert(target.full_address(), CircuitBreaker::new(cb_config));
987        }
988
989        let pool = Self {
990            id: id.clone(),
991            targets,
992            load_balancer,
993            pool_config,
994            http_version,
995            tls_enabled,
996            tls_sni,
997            tls_config,
998            circuit_breakers: Arc::new(RwLock::new(circuit_breakers)),
999            stats: Arc::new(PoolStats::default()),
1000            discovery,
1001            discovery_interval,
1002            static_targets,
1003            config,
1004            sticky_runtime,
1005        };
1006
1007        info!(
1008            upstream_id = %id,
1009            target_count = pool.targets.len(),
1010            "Upstream pool created successfully"
1011        );
1012
1013        Ok(pool)
1014    }
1015
1016    /// Combine configured targets with discovered ones.
1017    ///
1018    /// A configured target that the discovery source also returns appears once,
1019    /// keeping the configured entry: an operator who pinned a backend and gave
1020    /// it a weight should not have that weight doubled just because the source
1021    /// happens to list it too.
1022    fn merge_targets(
1023        static_targets: &[UpstreamTarget],
1024        discovered: Vec<UpstreamTarget>,
1025    ) -> Vec<UpstreamTarget> {
1026        let mut targets = static_targets.to_vec();
1027        let pinned: std::collections::HashSet<String> =
1028            targets.iter().map(|t| t.full_address()).collect();
1029        targets.extend(
1030            discovered
1031                .into_iter()
1032                .filter(|t| !pinned.contains(&t.full_address())),
1033        );
1034        targets
1035    }
1036
1037    /// Translate the configuration's discovery block into the runtime one.
1038    fn discovery_config(discovery: &UpstreamDiscovery) -> DiscoveryConfig {
1039        match discovery {
1040            UpstreamDiscovery::Static { backends } => DiscoveryConfig::Static {
1041                backends: backends.clone(),
1042            },
1043            UpstreamDiscovery::Dns {
1044                hostname,
1045                port,
1046                refresh_interval_secs,
1047            } => DiscoveryConfig::Dns {
1048                hostname: hostname.clone(),
1049                port: *port,
1050                refresh_interval: Duration::from_secs(*refresh_interval_secs),
1051            },
1052            UpstreamDiscovery::DnsSrv {
1053                service,
1054                refresh_interval_secs,
1055            } => DiscoveryConfig::DnsSrv {
1056                service: service.clone(),
1057                refresh_interval: Duration::from_secs(*refresh_interval_secs),
1058            },
1059            UpstreamDiscovery::Consul {
1060                address,
1061                service,
1062                datacenter,
1063                only_passing,
1064                refresh_interval_secs,
1065                tag,
1066            } => DiscoveryConfig::Consul {
1067                address: address.clone(),
1068                service: service.clone(),
1069                datacenter: datacenter.clone(),
1070                only_passing: *only_passing,
1071                refresh_interval: Duration::from_secs(*refresh_interval_secs),
1072                tag: tag.clone(),
1073            },
1074            UpstreamDiscovery::Kubernetes {
1075                namespace,
1076                service,
1077                port_name,
1078                refresh_interval_secs,
1079                kubeconfig,
1080            } => DiscoveryConfig::Kubernetes {
1081                namespace: namespace.clone(),
1082                service: service.clone(),
1083                port_name: port_name.clone(),
1084                refresh_interval: Duration::from_secs(*refresh_interval_secs),
1085                kubeconfig: kubeconfig.clone(),
1086            },
1087            UpstreamDiscovery::File {
1088                path,
1089                watch_interval_secs,
1090            } => DiscoveryConfig::File {
1091                path: path.clone(),
1092                watch_interval: Duration::from_secs(*watch_interval_secs),
1093            },
1094        }
1095    }
1096
1097    /// Convert the backends a discovery source returned into pool targets.
1098    ///
1099    /// Unix-socket backends are skipped: discovery describes networked service
1100    /// registries, and a pool target is an address/port pair.
1101    fn targets_from_backends(backends: &BTreeSet<Backend>) -> Vec<UpstreamTarget> {
1102        backends
1103            .iter()
1104            .filter_map(|backend| match &backend.addr {
1105                PingoraSocketAddr::Inet(addr) => Some(UpstreamTarget {
1106                    address: addr.ip().to_string(),
1107                    port: addr.port(),
1108                    weight: u32::try_from(backend.weight).unwrap_or(1).max(1),
1109                }),
1110                PingoraSocketAddr::Unix(_) => None,
1111            })
1112            .collect()
1113    }
1114
1115    /// Build the configured discovery source and resolve it once.
1116    ///
1117    /// A source that fails to answer is logged and treated as returning
1118    /// nothing: the pool falls back to whatever targets the configuration
1119    /// lists, and the refresh task retries on the next interval. Failing the
1120    /// whole pool here would mean one unreachable registry could stop the proxy
1121    /// from starting at all, taking every other upstream down with it.
1122    async fn resolve_discovery(
1123        config: &UpstreamConfig,
1124        id: &UpstreamId,
1125    ) -> (
1126        Option<Arc<dyn ServiceDiscovery + Send + Sync>>,
1127        Option<Duration>,
1128        Vec<UpstreamTarget>,
1129    ) {
1130        let Some(spec) = config.discovery.as_ref() else {
1131            return (None, None, Vec::new());
1132        };
1133
1134        let source = crate::discovery::build_discovery(id.as_str(), Self::discovery_config(spec));
1135
1136        // `static` never changes, so it is resolved once and never scheduled.
1137        let interval = match spec.refresh_interval_secs() {
1138            0 => None,
1139            secs => Some(Duration::from_secs(secs)),
1140        };
1141
1142        let targets = match source.discover().await {
1143            Ok((backends, _healthy)) => {
1144                let targets = Self::targets_from_backends(&backends);
1145                info!(
1146                    upstream_id = %config.id,
1147                    discovery = spec.kind(),
1148                    discovered = targets.len(),
1149                    refresh_interval_secs = interval.map(|i| i.as_secs()).unwrap_or(0),
1150                    "Resolved service discovery"
1151                );
1152                targets
1153            }
1154            Err(e) => {
1155                error!(
1156                    upstream_id = %config.id,
1157                    discovery = spec.kind(),
1158                    error = %e,
1159                    "Service discovery failed; using configured targets only"
1160                );
1161                Vec::new()
1162            }
1163        };
1164
1165        (Some(source), interval, targets)
1166    }
1167
1168    /// How often this pool's discovery source should be re-resolved, if at all.
1169    pub fn discovery_refresh_interval(&self) -> Option<Duration> {
1170        self.discovery_interval
1171    }
1172
1173    /// Re-resolve discovery and, when the target set has changed, build the
1174    /// pool that should replace this one.
1175    ///
1176    /// Returns `None` when there is nothing to do — no discovery source, the
1177    /// source failed, or it resolved to the same targets the pool already has.
1178    /// Callers install the returned pool in place of this one; the two share
1179    /// their circuit breakers and statistics, so a backend that was failing
1180    /// before the refresh is still failing after it.
1181    ///
1182    /// # Errors
1183    ///
1184    /// Never fails: a discovery source that cannot be reached leaves the pool
1185    /// serving its current targets rather than emptying it.
1186    pub async fn refreshed(&self) -> Option<UpstreamPool> {
1187        let source = self.discovery.as_ref()?;
1188
1189        let (backends, _healthy) = match source.discover().await {
1190            Ok(result) => result,
1191            Err(e) => {
1192                warn!(
1193                    upstream_id = %self.id,
1194                    error = %e,
1195                    "Service discovery refresh failed; keeping current targets"
1196                );
1197                return None;
1198            }
1199        };
1200
1201        let targets =
1202            Self::merge_targets(&self.static_targets, Self::targets_from_backends(&backends));
1203
1204        if Self::same_targets(&self.targets, &targets) {
1205            return None;
1206        }
1207
1208        let previous: Vec<String> = self.targets.iter().map(|t| t.full_address()).collect();
1209        let current: Vec<String> = targets.iter().map(|t| t.full_address()).collect();
1210        let added: Vec<&String> = current.iter().filter(|a| !previous.contains(a)).collect();
1211        let removed: Vec<&String> = previous.iter().filter(|a| !current.contains(a)).collect();
1212
1213        info!(
1214            upstream_id = %self.id,
1215            added = ?added,
1216            removed = ?removed,
1217            target_count = targets.len(),
1218            "Service discovery changed upstream targets"
1219        );
1220
1221        // Reconcile circuit breakers in place: surviving targets keep the
1222        // breaker they already had (and therefore their open/closed state),
1223        // targets that went away lose theirs so the map cannot grow without
1224        // bound as backends churn, and new targets start closed.
1225        let cb_config = self.config.circuit_breaker.unwrap_or_default();
1226        {
1227            let mut breakers = self.circuit_breakers.write().await;
1228            breakers.retain(|address, _| current.iter().any(|a| a == address));
1229            for address in &current {
1230                breakers
1231                    .entry(address.clone())
1232                    .or_insert_with(|| CircuitBreaker::new(cb_config));
1233            }
1234        }
1235
1236        let load_balancer = match Self::create_load_balancer(
1237            &self.config.load_balancing,
1238            &targets,
1239            &self.config,
1240            self.sticky_runtime.as_ref(),
1241        ) {
1242            Ok(balancer) => balancer,
1243            Err(e) => {
1244                error!(
1245                    upstream_id = %self.id,
1246                    error = %e,
1247                    "Failed to rebuild load balancer after discovery refresh; \
1248                     keeping current targets"
1249                );
1250                return None;
1251            }
1252        };
1253
1254        Some(UpstreamPool {
1255            id: self.id.clone(),
1256            targets,
1257            load_balancer,
1258            pool_config: self.pool_config.clone(),
1259            http_version: self.http_version.clone(),
1260            tls_enabled: self.tls_enabled,
1261            tls_sni: self.tls_sni.clone(),
1262            tls_config: self.tls_config.clone(),
1263            // Shared, not copied: breaker state and counters must survive the
1264            // swap or a flapping backend would look healthy on every refresh.
1265            circuit_breakers: Arc::clone(&self.circuit_breakers),
1266            stats: Arc::clone(&self.stats),
1267            discovery: self.discovery.clone(),
1268            discovery_interval: self.discovery_interval,
1269            static_targets: self.static_targets.clone(),
1270            config: self.config.clone(),
1271            sticky_runtime: self.sticky_runtime.clone(),
1272        })
1273    }
1274
1275    /// Whether two target sets are the same, ignoring order.
1276    fn same_targets(a: &[UpstreamTarget], b: &[UpstreamTarget]) -> bool {
1277        if a.len() != b.len() {
1278            return false;
1279        }
1280        let mut a: Vec<(String, u32)> = a.iter().map(|t| (t.full_address(), t.weight)).collect();
1281        let mut b: Vec<(String, u32)> = b.iter().map(|t| (t.full_address(), t.weight)).collect();
1282        a.sort();
1283        b.sort();
1284        a == b
1285    }
1286
1287    /// Create load balancer based on algorithm
1288    fn create_load_balancer(
1289        algorithm: &LoadBalancingAlgorithm,
1290        targets: &[UpstreamTarget],
1291        config: &UpstreamConfig,
1292        sticky_runtime: Option<&StickySessionRuntimeConfig>,
1293    ) -> ZentinelResult<Arc<dyn LoadBalancer>> {
1294        let balancer: Arc<dyn LoadBalancer> = match algorithm {
1295            LoadBalancingAlgorithm::RoundRobin => {
1296                Arc::new(RoundRobinBalancer::new(targets.to_vec()))
1297            }
1298            LoadBalancingAlgorithm::LeastConnections => {
1299                Arc::new(LeastConnectionsBalancer::new(targets.to_vec()))
1300            }
1301            LoadBalancingAlgorithm::Weighted => {
1302                let weights: Vec<u32> = targets.iter().map(|t| t.weight).collect();
1303                Arc::new(WeightedBalancer {
1304                    targets: targets.to_vec(),
1305                    weights,
1306                    current_index: AtomicUsize::new(0),
1307                    health_status: Arc::new(RwLock::new(HashMap::new())),
1308                })
1309            }
1310            LoadBalancingAlgorithm::IpHash => Arc::new(IpHashBalancer {
1311                targets: targets.to_vec(),
1312                health_status: Arc::new(RwLock::new(HashMap::new())),
1313            }),
1314            LoadBalancingAlgorithm::Random => Arc::new(RandomBalancer::new(targets.to_vec())),
1315            LoadBalancingAlgorithm::ConsistentHash => Arc::new(ConsistentHashBalancer::new(
1316                targets.to_vec(),
1317                ConsistentHashConfig::default(),
1318            )),
1319            LoadBalancingAlgorithm::PowerOfTwoChoices => {
1320                Arc::new(P2cBalancer::new(targets.to_vec(), P2cConfig::default()))
1321            }
1322            LoadBalancingAlgorithm::Adaptive => Arc::new(AdaptiveBalancer::new(
1323                targets.to_vec(),
1324                AdaptiveConfig::default(),
1325            )),
1326            LoadBalancingAlgorithm::LeastTokensQueued => Arc::new(LeastTokensQueuedBalancer::new(
1327                targets.to_vec(),
1328                LeastTokensQueuedConfig::default(),
1329            )),
1330            LoadBalancingAlgorithm::Maglev => Arc::new(MaglevBalancer::new(
1331                targets.to_vec(),
1332                MaglevConfig::default(),
1333            )),
1334            LoadBalancingAlgorithm::LocalityAware => Arc::new(LocalityAwareBalancer::new(
1335                targets.to_vec(),
1336                LocalityAwareConfig::default(),
1337            )),
1338            LoadBalancingAlgorithm::PeakEwma => Arc::new(PeakEwmaBalancer::new(
1339                targets.to_vec(),
1340                PeakEwmaConfig::default(),
1341            )),
1342            LoadBalancingAlgorithm::DeterministicSubset => Arc::new(SubsetBalancer::new(
1343                targets.to_vec(),
1344                SubsetConfig::default(),
1345            )),
1346            LoadBalancingAlgorithm::WeightedLeastConnections => {
1347                Arc::new(WeightedLeastConnBalancer::new(
1348                    targets.to_vec(),
1349                    WeightedLeastConnConfig::default(),
1350                ))
1351            }
1352            LoadBalancingAlgorithm::Sticky => {
1353                // Get sticky session config (required for Sticky algorithm)
1354                let sticky_config = config.sticky_session.as_ref().ok_or_else(|| {
1355                    ZentinelError::Config {
1356                        message: format!(
1357                            "Upstream '{}' uses Sticky algorithm but no sticky_session config provided",
1358                            config.id
1359                        ),
1360                        source: None,
1361                    }
1362                })?;
1363
1364                // Reuse the pool's existing key when rebuilding, so affinity
1365                // cookies issued before the rebuild still verify.
1366                let runtime_config = sticky_runtime
1367                    .cloned()
1368                    .unwrap_or_else(|| StickySessionRuntimeConfig::from_config(sticky_config));
1369
1370                // Create fallback load balancer
1371                let fallback = Self::create_load_balancer_inner(&sticky_config.fallback, targets)?;
1372
1373                info!(
1374                    upstream_id = %config.id,
1375                    cookie_name = %runtime_config.cookie_name,
1376                    cookie_ttl_secs = runtime_config.cookie_ttl_secs,
1377                    fallback_algorithm = ?sticky_config.fallback,
1378                    "Creating sticky session balancer"
1379                );
1380
1381                Arc::new(StickySessionBalancer::new(
1382                    targets.to_vec(),
1383                    runtime_config,
1384                    fallback,
1385                ))
1386            }
1387        };
1388        Ok(balancer)
1389    }
1390
1391    /// Create load balancer without sticky session support (for fallback balancers)
1392    fn create_load_balancer_inner(
1393        algorithm: &LoadBalancingAlgorithm,
1394        targets: &[UpstreamTarget],
1395    ) -> ZentinelResult<Arc<dyn LoadBalancer>> {
1396        let balancer: Arc<dyn LoadBalancer> = match algorithm {
1397            LoadBalancingAlgorithm::RoundRobin => {
1398                Arc::new(RoundRobinBalancer::new(targets.to_vec()))
1399            }
1400            LoadBalancingAlgorithm::LeastConnections => {
1401                Arc::new(LeastConnectionsBalancer::new(targets.to_vec()))
1402            }
1403            LoadBalancingAlgorithm::Weighted => {
1404                let weights: Vec<u32> = targets.iter().map(|t| t.weight).collect();
1405                Arc::new(WeightedBalancer {
1406                    targets: targets.to_vec(),
1407                    weights,
1408                    current_index: AtomicUsize::new(0),
1409                    health_status: Arc::new(RwLock::new(HashMap::new())),
1410                })
1411            }
1412            LoadBalancingAlgorithm::IpHash => Arc::new(IpHashBalancer {
1413                targets: targets.to_vec(),
1414                health_status: Arc::new(RwLock::new(HashMap::new())),
1415            }),
1416            LoadBalancingAlgorithm::Random => Arc::new(RandomBalancer::new(targets.to_vec())),
1417            LoadBalancingAlgorithm::ConsistentHash => Arc::new(ConsistentHashBalancer::new(
1418                targets.to_vec(),
1419                ConsistentHashConfig::default(),
1420            )),
1421            LoadBalancingAlgorithm::PowerOfTwoChoices => {
1422                Arc::new(P2cBalancer::new(targets.to_vec(), P2cConfig::default()))
1423            }
1424            LoadBalancingAlgorithm::Adaptive => Arc::new(AdaptiveBalancer::new(
1425                targets.to_vec(),
1426                AdaptiveConfig::default(),
1427            )),
1428            LoadBalancingAlgorithm::LeastTokensQueued => Arc::new(LeastTokensQueuedBalancer::new(
1429                targets.to_vec(),
1430                LeastTokensQueuedConfig::default(),
1431            )),
1432            LoadBalancingAlgorithm::Maglev => Arc::new(MaglevBalancer::new(
1433                targets.to_vec(),
1434                MaglevConfig::default(),
1435            )),
1436            LoadBalancingAlgorithm::LocalityAware => Arc::new(LocalityAwareBalancer::new(
1437                targets.to_vec(),
1438                LocalityAwareConfig::default(),
1439            )),
1440            LoadBalancingAlgorithm::PeakEwma => Arc::new(PeakEwmaBalancer::new(
1441                targets.to_vec(),
1442                PeakEwmaConfig::default(),
1443            )),
1444            LoadBalancingAlgorithm::DeterministicSubset => Arc::new(SubsetBalancer::new(
1445                targets.to_vec(),
1446                SubsetConfig::default(),
1447            )),
1448            LoadBalancingAlgorithm::WeightedLeastConnections => {
1449                Arc::new(WeightedLeastConnBalancer::new(
1450                    targets.to_vec(),
1451                    WeightedLeastConnConfig::default(),
1452                ))
1453            }
1454            LoadBalancingAlgorithm::Sticky => {
1455                // Sticky cannot be used as fallback (would cause infinite recursion)
1456                return Err(ZentinelError::Config {
1457                    message: "Sticky algorithm cannot be used as fallback for sticky sessions"
1458                        .to_string(),
1459                    source: None,
1460                });
1461            }
1462        };
1463        Ok(balancer)
1464    }
1465
1466    /// Select next upstream peer with selection metadata
1467    ///
1468    /// Returns the selected peer along with optional metadata from the load balancer.
1469    /// The metadata can contain sticky session information that should be passed to
1470    /// the response filter.
1471    pub async fn select_peer_with_metadata(
1472        &self,
1473        context: Option<&RequestContext>,
1474    ) -> ZentinelResult<(HttpPeer, HashMap<String, String>)> {
1475        let request_num = self.stats.requests.fetch_add(1, Ordering::Relaxed) + 1;
1476
1477        trace!(
1478            upstream_id = %self.id,
1479            request_num = request_num,
1480            target_count = self.targets.len(),
1481            "Starting peer selection with metadata"
1482        );
1483
1484        let mut attempts = 0;
1485        let max_attempts = self.targets.len() * 2;
1486
1487        while attempts < max_attempts {
1488            attempts += 1;
1489
1490            trace!(
1491                upstream_id = %self.id,
1492                attempt = attempts,
1493                max_attempts = max_attempts,
1494                "Attempting to select peer"
1495            );
1496
1497            let selection = match self.load_balancer.select(context).await {
1498                Ok(s) => s,
1499                Err(e) => {
1500                    warn!(
1501                        upstream_id = %self.id,
1502                        attempt = attempts,
1503                        error = %e,
1504                        "Load balancer selection failed"
1505                    );
1506                    continue;
1507                }
1508            };
1509
1510            trace!(
1511                upstream_id = %self.id,
1512                target = %selection.address,
1513                attempt = attempts,
1514                "Load balancer selected target"
1515            );
1516
1517            // Check circuit breaker
1518            let breakers = self.circuit_breakers.read().await;
1519            if let Some(breaker) = breakers.get(&selection.address) {
1520                if !breaker.is_closed() {
1521                    debug!(
1522                        upstream_id = %self.id,
1523                        target = %selection.address,
1524                        attempt = attempts,
1525                        "Circuit breaker is open, skipping target"
1526                    );
1527                    self.stats
1528                        .circuit_breaker_trips
1529                        .fetch_add(1, Ordering::Relaxed);
1530                    continue;
1531                }
1532            }
1533
1534            // Create peer with pooling options
1535            trace!(
1536                upstream_id = %self.id,
1537                target = %selection.address,
1538                "Creating peer for upstream (Pingora handles connection reuse)"
1539            );
1540            let peer = self.create_peer(&selection)?;
1541
1542            debug!(
1543                upstream_id = %self.id,
1544                target = %selection.address,
1545                attempt = attempts,
1546                metadata_keys = ?selection.metadata.keys().collect::<Vec<_>>(),
1547                "Selected upstream peer with metadata"
1548            );
1549
1550            self.stats.successes.fetch_add(1, Ordering::Relaxed);
1551            return Ok((peer, selection.metadata));
1552        }
1553
1554        self.stats.failures.fetch_add(1, Ordering::Relaxed);
1555        error!(
1556            upstream_id = %self.id,
1557            attempts = attempts,
1558            max_attempts = max_attempts,
1559            "Failed to select upstream after max attempts"
1560        );
1561        Err(ZentinelError::upstream(
1562            self.id.to_string(),
1563            "Failed to select upstream after max attempts",
1564        ))
1565    }
1566
1567    /// Select next upstream peer
1568    pub async fn select_peer(&self, context: Option<&RequestContext>) -> ZentinelResult<HttpPeer> {
1569        // Delegate to select_peer_with_metadata and discard metadata
1570        self.select_peer_with_metadata(context)
1571            .await
1572            .map(|(peer, _)| peer)
1573    }
1574
1575    /// Create new peer connection with connection pooling options
1576    ///
1577    /// Pingora handles actual connection pooling internally. When idle_timeout
1578    /// is set on the peer options, Pingora will keep the connection alive and
1579    /// reuse it for subsequent requests to the same upstream.
1580    fn create_peer(&self, selection: &TargetSelection) -> ZentinelResult<HttpPeer> {
1581        // Determine SNI hostname for TLS connections
1582        let sni_hostname = self.tls_sni.clone().unwrap_or_else(|| {
1583            // Extract hostname from address (strip port)
1584            selection
1585                .address
1586                .split(':')
1587                .next()
1588                .unwrap_or(&selection.address)
1589                .to_string()
1590        });
1591
1592        // Pre-resolve the address to avoid panics in Pingora's HttpPeer::new
1593        // when DNS resolution fails (e.g., when a container is killed)
1594        let resolved_address = selection
1595            .address
1596            .to_socket_addrs()
1597            .map_err(|e| {
1598                error!(
1599                    upstream = %self.id,
1600                    address = %selection.address,
1601                    error = %e,
1602                    "Failed to resolve upstream address"
1603                );
1604                ZentinelError::Upstream {
1605                    upstream: self.id.to_string(),
1606                    message: format!("DNS resolution failed for {}: {}", selection.address, e),
1607                    retryable: true,
1608                    source: None,
1609                }
1610            })?
1611            .next()
1612            .ok_or_else(|| {
1613                error!(
1614                    upstream = %self.id,
1615                    address = %selection.address,
1616                    "No addresses returned from DNS resolution"
1617                );
1618                ZentinelError::Upstream {
1619                    upstream: self.id.to_string(),
1620                    message: format!("No addresses for {}", selection.address),
1621                    retryable: true,
1622                    source: None,
1623                }
1624            })?;
1625
1626        // Use the resolved IP address to create the peer
1627        let mut peer = HttpPeer::new(resolved_address, self.tls_enabled, sni_hostname.clone());
1628
1629        // Configure connection pooling options for better performance
1630        // idle_timeout enables Pingora's connection pooling - connections are
1631        // kept alive and reused for this duration
1632        peer.options.idle_timeout = Some(self.pool_config.idle_timeout);
1633
1634        // Connection timeouts
1635        peer.options.connection_timeout = Some(self.pool_config.connection_timeout);
1636        peer.options.total_connection_timeout = Some(Duration::from_secs(10));
1637
1638        // Read/write timeouts
1639        peer.options.read_timeout = Some(self.pool_config.read_timeout);
1640        peer.options.write_timeout = Some(self.pool_config.write_timeout);
1641
1642        // Enable TCP keepalive for long-lived connections
1643        peer.options.tcp_keepalive = Some(pingora::protocols::TcpKeepalive {
1644            idle: Duration::from_secs(60),
1645            interval: Duration::from_secs(10),
1646            count: 3,
1647            // user_timeout is Linux-only
1648            #[cfg(target_os = "linux")]
1649            user_timeout: Duration::from_secs(60),
1650        });
1651
1652        // Configure HTTP version and ALPN for TLS connections
1653        if self.tls_enabled {
1654            // Set ALPN protocols based on configured HTTP version range
1655            let alpn = match (self.http_version.min_version, self.http_version.max_version) {
1656                (2, _) => {
1657                    // HTTP/2 only - use h2 ALPN
1658                    pingora::upstreams::peer::ALPN::H2
1659                }
1660                (1, 2) | (_, 2) => {
1661                    // Prefer HTTP/2 but fall back to HTTP/1.1
1662                    pingora::upstreams::peer::ALPN::H2H1
1663                }
1664                _ => {
1665                    // HTTP/1.1 only
1666                    pingora::upstreams::peer::ALPN::H1
1667                }
1668            };
1669            peer.options.alpn = alpn;
1670
1671            // Configure TLS verification options based on upstream config
1672            if let Some(ref tls_config) = self.tls_config {
1673                // Skip certificate verification if configured (DANGEROUS - testing only)
1674                if tls_config.insecure_skip_verify {
1675                    peer.options.verify_cert = false;
1676                    peer.options.verify_hostname = false;
1677                    warn!(
1678                        upstream_id = %self.id,
1679                        target = %selection.address,
1680                        "TLS certificate verification DISABLED (insecure_skip_verify=true)"
1681                    );
1682                }
1683
1684                // Set alternative CN for verification if SNI differs from actual hostname
1685                if let Some(ref sni) = tls_config.sni {
1686                    peer.options.alternative_cn = Some(sni.clone());
1687                    trace!(
1688                        upstream_id = %self.id,
1689                        target = %selection.address,
1690                        alternative_cn = %sni,
1691                        "Set alternative CN for TLS verification"
1692                    );
1693                }
1694
1695                // Configure mTLS client certificate if provided
1696                if let (Some(cert_path), Some(key_path)) =
1697                    (&tls_config.client_cert, &tls_config.client_key)
1698                {
1699                    match crate::tls::load_client_cert_key(cert_path, key_path) {
1700                        Ok(cert_key) => {
1701                            peer.client_cert_key = Some(cert_key);
1702                            info!(
1703                                upstream_id = %self.id,
1704                                target = %selection.address,
1705                                cert_path = ?cert_path,
1706                                "mTLS client certificate configured"
1707                            );
1708                        }
1709                        Err(e) => {
1710                            error!(
1711                                upstream_id = %self.id,
1712                                target = %selection.address,
1713                                error = %e,
1714                                "Failed to load mTLS client certificate"
1715                            );
1716                            return Err(ZentinelError::Tls {
1717                                message: format!("Failed to load client certificate: {}", e),
1718                                source: None,
1719                            });
1720                        }
1721                    }
1722                }
1723            }
1724
1725            trace!(
1726                upstream_id = %self.id,
1727                target = %selection.address,
1728                alpn = ?peer.options.alpn,
1729                min_version = self.http_version.min_version,
1730                max_version = self.http_version.max_version,
1731                verify_cert = peer.options.verify_cert,
1732                verify_hostname = peer.options.verify_hostname,
1733                "Configured ALPN and TLS options for HTTP version negotiation"
1734            );
1735        }
1736
1737        // Configure H2-specific settings when HTTP/2 is enabled
1738        if self.http_version.max_version >= 2 {
1739            // H2 ping interval for connection health monitoring
1740            if !self.http_version.h2_ping_interval.is_zero() {
1741                peer.options.h2_ping_interval = Some(self.http_version.h2_ping_interval);
1742                trace!(
1743                    upstream_id = %self.id,
1744                    target = %selection.address,
1745                    h2_ping_interval_secs = self.http_version.h2_ping_interval.as_secs(),
1746                    "Configured H2 ping interval"
1747                );
1748            }
1749        }
1750
1751        trace!(
1752            upstream_id = %self.id,
1753            target = %selection.address,
1754            tls = self.tls_enabled,
1755            sni = %sni_hostname,
1756            idle_timeout_secs = self.pool_config.idle_timeout.as_secs(),
1757            http_max_version = self.http_version.max_version,
1758            "Created peer with Pingora connection pooling enabled"
1759        );
1760
1761        Ok(peer)
1762    }
1763
1764    /// Report connection result for a target
1765    ///
1766    /// On failure, the circuit breaker records the failure but the load balancer
1767    /// health status is only updated when the circuit breaker transitions to Open.
1768    /// This prevents a single connection error (e.g., a stale pooled connection
1769    /// reset) from permanently removing a target from the healthy pool.
1770    pub async fn report_result(&self, target: &str, success: bool) {
1771        trace!(
1772            upstream_id = %self.id,
1773            target = %target,
1774            success = success,
1775            "Reporting connection result"
1776        );
1777
1778        if success {
1779            if let Some(breaker) = self.circuit_breakers.read().await.get(target) {
1780                breaker.record_success();
1781                trace!(
1782                    upstream_id = %self.id,
1783                    target = %target,
1784                    "Recorded success in circuit breaker"
1785                );
1786            }
1787            self.load_balancer.report_health(target, true).await;
1788        } else {
1789            let breaker_opened =
1790                if let Some(breaker) = self.circuit_breakers.read().await.get(target) {
1791                    let opened = breaker.record_failure();
1792                    debug!(
1793                        upstream_id = %self.id,
1794                        target = %target,
1795                        circuit_breaker_opened = opened,
1796                        "Recorded failure in circuit breaker"
1797                    );
1798                    opened
1799                } else {
1800                    false
1801                };
1802
1803            // Do NOT mark the target down in the load balancer when the breaker
1804            // opens. The circuit breaker is the single availability gate — the
1805            // upstream_peer selection loop checks `is_closed()`, which also runs
1806            // the timed Open->HalfOpen transition and lets a probe through after
1807            // `timeout_seconds`. Removing the target from the load balancer here
1808            // would prevent it from ever being selected again, so that probe
1809            // would never run and the target could not recover (#261).
1810
1811            self.stats.failures.fetch_add(1, Ordering::Relaxed);
1812            warn!(
1813                upstream_id = %self.id,
1814                target = %target,
1815                circuit_breaker_opened = breaker_opened,
1816                "Connection failure reported for target"
1817            );
1818        }
1819    }
1820
1821    /// Report request result with latency for adaptive load balancing
1822    ///
1823    /// This method passes latency information to the load balancer for
1824    /// adaptive weight adjustment. It updates circuit breakers and health
1825    /// status. On failure, health is only marked down when the circuit
1826    /// breaker transitions to Open, preventing stale connection resets
1827    /// from permanently removing targets.
1828    pub async fn report_result_with_latency(
1829        &self,
1830        target: &str,
1831        success: bool,
1832        latency: Option<Duration>,
1833    ) {
1834        trace!(
1835            upstream_id = %self.id,
1836            target = %target,
1837            success = success,
1838            latency_ms = latency.map(|l| l.as_millis() as u64),
1839            "Reporting result with latency for adaptive LB"
1840        );
1841
1842        // Update circuit breaker
1843        if success {
1844            if let Some(breaker) = self.circuit_breakers.read().await.get(target) {
1845                breaker.record_success();
1846            }
1847            // Always report success to the load balancer (restores health + records latency)
1848            self.load_balancer
1849                .report_result_with_latency(target, true, latency)
1850                .await;
1851        } else {
1852            // Record the failure in the circuit breaker (this may open it). We
1853            // do NOT propagate a health-down to the load balancer on open: the
1854            // selection loop's `is_closed()` check is the sole availability gate
1855            // and runs the timed half-open recovery probe. Marking the target
1856            // down here would remove it from selection and block recovery (#261).
1857            if let Some(breaker) = self.circuit_breakers.read().await.get(target) {
1858                breaker.record_failure();
1859            }
1860            self.stats.failures.fetch_add(1, Ordering::Relaxed);
1861        }
1862    }
1863
1864    /// Get pool statistics
1865    pub fn stats(&self) -> &PoolStats {
1866        &self.stats
1867    }
1868
1869    /// Get pool ID
1870    pub fn id(&self) -> &UpstreamId {
1871        &self.id
1872    }
1873
1874    /// Get target count
1875    /// Circuit-breaker state for a target, or `None` when the pool has no
1876    /// breaker for that address — which for a discovery-backed pool means the
1877    /// target is not currently in the resolved set.
1878    pub async fn circuit_breaker_state(
1879        &self,
1880        target: &str,
1881    ) -> Option<zentinel_common::types::CircuitBreakerState> {
1882        self.circuit_breakers
1883            .read()
1884            .await
1885            .get(target)
1886            .map(|breaker| breaker.state())
1887    }
1888
1889    /// The key used to sign sticky-session affinity cookies, when this upstream
1890    /// uses sticky load balancing.
1891    ///
1892    /// Exposed so tests can assert it survives a pool rebuild; rotating it would
1893    /// invalidate every cookie already issued.
1894    pub fn sticky_signing_key(&self) -> Option<[u8; 32]> {
1895        self.load_balancer.session_signing_key()
1896    }
1897
1898    /// Addresses of the pool's current targets, in selection order.
1899    pub fn target_addresses(&self) -> Vec<String> {
1900        self.targets.iter().map(|t| t.full_address()).collect()
1901    }
1902
1903    pub fn target_count(&self) -> usize {
1904        self.targets.len()
1905    }
1906
1907    /// Get pool configuration (for metrics/debugging)
1908    pub fn pool_config(&self) -> PoolConfigSnapshot {
1909        PoolConfigSnapshot {
1910            max_connections: self.pool_config.max_connections,
1911            max_idle: self.pool_config.max_idle,
1912            idle_timeout_secs: self.pool_config.idle_timeout.as_secs(),
1913            max_lifetime_secs: self.pool_config.max_lifetime.map(|d| d.as_secs()),
1914            connection_timeout_secs: self.pool_config.connection_timeout.as_secs(),
1915            read_timeout_secs: self.pool_config.read_timeout.as_secs(),
1916            write_timeout_secs: self.pool_config.write_timeout.as_secs(),
1917        }
1918    }
1919
1920    /// Check if the pool has any healthy targets.
1921    ///
1922    /// Returns true if at least one target is healthy, false if all targets are unhealthy.
1923    pub async fn has_healthy_targets(&self) -> bool {
1924        let healthy = self.load_balancer.healthy_targets().await;
1925        !healthy.is_empty()
1926    }
1927
1928    /// Select a target for shadow traffic (returns URL components)
1929    ///
1930    /// This is a simplified selection method for shadow requests that don't need
1931    /// full HttpPeer setup. Returns the target URL scheme, address, and port.
1932    pub async fn select_shadow_target(
1933        &self,
1934        context: Option<&RequestContext>,
1935    ) -> ZentinelResult<ShadowTarget> {
1936        // Use load balancer to select target
1937        let selection = self.load_balancer.select(context).await?;
1938
1939        // Check circuit breaker
1940        let breakers = self.circuit_breakers.read().await;
1941        if let Some(breaker) = breakers.get(&selection.address) {
1942            if !breaker.is_closed() {
1943                return Err(ZentinelError::upstream(
1944                    self.id.to_string(),
1945                    "Circuit breaker is open for shadow target",
1946                ));
1947            }
1948        }
1949
1950        // Parse address to get host and port
1951        let (host, port) = if selection.address.contains(':') {
1952            let parts: Vec<&str> = selection.address.rsplitn(2, ':').collect();
1953            if parts.len() == 2 {
1954                (
1955                    parts[1].to_string(),
1956                    parts[0]
1957                        .parse::<u16>()
1958                        .unwrap_or(if self.tls_enabled { 443 } else { 80 }),
1959                )
1960            } else {
1961                (
1962                    selection.address.clone(),
1963                    if self.tls_enabled { 443 } else { 80 },
1964                )
1965            }
1966        } else {
1967            (
1968                selection.address.clone(),
1969                if self.tls_enabled { 443 } else { 80 },
1970            )
1971        };
1972
1973        Ok(ShadowTarget {
1974            scheme: if self.tls_enabled { "https" } else { "http" }.to_string(),
1975            host,
1976            port,
1977            sni: self.tls_sni.clone(),
1978        })
1979    }
1980
1981    /// Check if TLS is enabled for this upstream
1982    pub fn is_tls_enabled(&self) -> bool {
1983        self.tls_enabled
1984    }
1985
1986    /// Get the number of currently active (in-flight) requests for this pool.
1987    ///
1988    /// Used by the drain tracker to determine when a pool has been fully
1989    /// drained after removal from config.
1990    pub fn active_request_count(&self) -> u64 {
1991        self.stats.active_requests.load(Ordering::Relaxed)
1992    }
1993
1994    /// Increment the active request counter. Called when a request is assigned
1995    /// to this pool.
1996    pub fn increment_active(&self) {
1997        self.stats.active_requests.fetch_add(1, Ordering::Relaxed);
1998    }
1999
2000    /// Decrement the active request counter. Called when a request completes
2001    /// (success or failure).
2002    pub fn decrement_active(&self) {
2003        let prev = self.stats.active_requests.fetch_sub(1, Ordering::Relaxed);
2004        if prev == 0 {
2005            self.stats.active_requests.fetch_add(1, Ordering::Relaxed);
2006            warn!("Attempted to decrement active request count below zero");
2007        }
2008    }
2009
2010    /// Shutdown the pool
2011    ///
2012    /// Note: Pingora manages connection pooling internally, so we just log stats.
2013    pub async fn shutdown(&self) {
2014        info!(
2015            upstream_id = %self.id,
2016            target_count = self.targets.len(),
2017            total_requests = self.stats.requests.load(Ordering::Relaxed),
2018            total_successes = self.stats.successes.load(Ordering::Relaxed),
2019            total_failures = self.stats.failures.load(Ordering::Relaxed),
2020            "Shutting down upstream pool"
2021        );
2022        // Pingora handles connection cleanup internally
2023        debug!(upstream_id = %self.id, "Upstream pool shutdown complete");
2024    }
2025}