fix: implement Envoy-compliant connection draining with mixed protocol support

Eeshu-Yadav · Eeshu-Yadav · commit 99c2e88770fe · 2025-09-24T02:30:17.000+05:30
- Add ListenerProtocolConfig enum to handle HTTP, TCP, and mixed protocol listeners
- Respect HttpConnectionManager.drain_timeout field from configuration
- Support listeners with both http_connection_manager and tcp_proxy filters
- Remove ambiguous 'Gradual' strategy, align with Envoy's draining behavior
- Add initiate_listener_drain_from_filter_analysis() for proper integration
Signed-off-by: Eeshu-Yadav &lt;eeshuyadav123@gmail.com&gt;
diff --git a/orion-lib/src/listeners/drain_signaling.rs b/orion-lib/src/listeners/drain_signaling.rs
@@ -15,8 +15,9 @@
 //
 //
 
+
 use crate::{Error, Result};
-use orion_configuration::config::listener::DrainType as ConfigDrainType;
+use orion_configuration::config::listener::{DrainType as ConfigDrainType, FilterChain, MainFilter};
 use pingora_timeout::fast_timeout::fast_timeout;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -25,6 +26,17 @@ use tokio::sync::RwLock;
 use tokio::time::sleep;
 use tracing::{debug, info, warn};
 
+#[derive(Debug, Clone)]
+pub enum ListenerProtocolConfig {
+    Http { drain_timeout: Option<Duration> },
+    Tcp,
+    Mixed { 
+        http_drain_timeout: Option<Duration>,
+        has_tcp: bool,
+        has_http: bool,
+    },
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum DrainScenario {
     HealthCheckFail,
@@ -46,8 +58,13 @@ impl DrainScenario {
 pub enum DrainStrategy {
     Tcp { global_timeout: Duration },
     Http { global_timeout: Duration, drain_timeout: Duration },
+    Mixed { 
+        global_timeout: Duration, 
+        http_drain_timeout: Duration,
+        tcp_connections: bool,
+        http_connections: bool,
+    },
     Immediate,
-    Gradual,
 }
 
 #[derive(Debug, Clone)]
@@ -102,9 +119,10 @@ impl ListenerDrainContext {
 
     pub fn is_timeout_exceeded(&self) -> bool {
         let global_timeout = match &self.strategy {
-            DrainStrategy::Tcp { global_timeout } | DrainStrategy::Http { global_timeout, .. } => *global_timeout,
+            DrainStrategy::Tcp { global_timeout } 
+            | DrainStrategy::Http { global_timeout, .. }
+            | DrainStrategy::Mixed { global_timeout, .. } => *global_timeout,
             DrainStrategy::Immediate => Duration::from_secs(0),
-            DrainStrategy::Gradual => Duration::from_secs(600),
         };
 
         self.drain_start.elapsed() >= global_timeout
@@ -113,6 +131,7 @@ impl ListenerDrainContext {
     pub fn get_http_drain_timeout(&self) -> Option<Duration> {
         match &self.strategy {
             DrainStrategy::Http { drain_timeout, .. } => Some(*drain_timeout),
+            DrainStrategy::Mixed { http_drain_timeout, .. } => Some(*http_drain_timeout),
             _ => None,
         }
     }
@@ -126,12 +145,36 @@ pub struct DrainSignalingManager {
     listener_drain_state: Arc<RwLock<Option<ListenerDrainState>>>,
 }
 
+impl ListenerProtocolConfig {
+    pub fn from_listener_analysis(
+        has_http_connection_manager: bool,
+        has_tcp_proxy: bool,
+        http_drain_timeout: Option<Duration>,
+    ) -> Self {
+        match (has_http_connection_manager, has_tcp_proxy) {
+            (true, true) => Self::Mixed {
+                http_drain_timeout,
+                has_tcp: true,
+                has_http: true,
+            },
+            (true, false) => Self::Http {
+                drain_timeout: http_drain_timeout,
+            },
+            (false, true) => Self::Tcp,
+            (false, false) => {
+                warn!("No HTTP connection manager or TCP proxy found in listener, defaulting to TCP draining");
+                Self::Tcp
+            }
+        }
+    }
+}
+
 impl DrainSignalingManager {
     pub fn new() -> Self {
         Self {
             drain_contexts: Arc::new(RwLock::new(HashMap::new())),
             global_drain_timeout: Duration::from_secs(600),
-            default_http_drain_timeout: Duration::from_millis(5000),
+            default_http_drain_timeout: Duration::from_secs(5),
             listener_drain_state: Arc::new(RwLock::new(None)),
         }
     }
@@ -207,17 +250,23 @@ impl DrainSignalingManager {
     pub async fn initiate_listener_drain(
         &self,
         listener_id: String,
-        is_http: bool,
-        http_drain_timeout: Option<Duration>,
+        protocol_config: ListenerProtocolConfig,
         active_connections: usize,
     ) -> Result<Arc<ListenerDrainContext>> {
-        let strategy = if is_http {
-            DrainStrategy::Http {
+        let strategy = match protocol_config {
+            ListenerProtocolConfig::Http { drain_timeout } => DrainStrategy::Http {
                 global_timeout: self.global_drain_timeout,
-                drain_timeout: http_drain_timeout.unwrap_or(self.default_http_drain_timeout),
-            }
-        } else {
-            DrainStrategy::Tcp { global_timeout: self.global_drain_timeout }
+                drain_timeout: drain_timeout.unwrap_or(self.default_http_drain_timeout),
+            },
+            ListenerProtocolConfig::Tcp => DrainStrategy::Tcp { 
+                global_timeout: self.global_drain_timeout 
+            },
+            ListenerProtocolConfig::Mixed { http_drain_timeout, has_tcp, has_http } => DrainStrategy::Mixed {
+                global_timeout: self.global_drain_timeout,
+                http_drain_timeout: http_drain_timeout.unwrap_or(self.default_http_drain_timeout),
+                tcp_connections: has_tcp,
+                http_connections: has_http,
+            },
         };
 
         let context = Arc::new(ListenerDrainContext::new(listener_id.clone(), strategy.clone(), active_connections));
@@ -326,6 +375,37 @@ impl DrainSignalingManager {
             Err(Error::new("Timeout waiting for listener drain completion"))
         }
     }
+
+    pub async fn initiate_listener_drain_from_filter_analysis(
+        &self,
+        listener_id: String,
+        filter_chains: &[FilterChain],
+        active_connections: usize,
+    ) -> Result<Arc<ListenerDrainContext>> {
+        let mut has_http = false;
+        let mut has_tcp = false;
+        let mut http_drain_timeout: Option<Duration> = None;
+
+        for filter_chain in filter_chains {
+            match &filter_chain.terminal_filter {
+                MainFilter::Http(http_config) => {
+                    has_http = true;
+                    http_drain_timeout = http_config.drain_timeout;
+                }
+                MainFilter::Tcp(_) => {
+                    has_tcp = true;
+                }
+            }
+        }
+
+        let protocol_config = ListenerProtocolConfig::from_listener_analysis(
+            has_http,
+            has_tcp,
+            http_drain_timeout,
+        );
+
+        self.initiate_listener_drain(listener_id, protocol_config, active_connections).await
+    }
 }
 
 impl Clone for DrainSignalingManager {
@@ -403,7 +483,11 @@ mod tests {
         let manager = DrainSignalingManager::new();
         assert!(!manager.has_draining_listeners().await);
 
-        let context = manager.initiate_listener_drain("test".to_string(), false, None, 1).await.unwrap();
+        let context = manager.initiate_listener_drain(
+            "test".to_string(), 
+            ListenerProtocolConfig::Tcp, 
+            1
+        ).await.unwrap();
 
         assert!(manager.has_draining_listeners().await);
         assert_eq!(manager.get_draining_listeners().await, vec!["test"]);
@@ -417,7 +501,11 @@ mod tests {
     async fn test_timeout_behavior() {
         let manager = DrainSignalingManager::with_timeouts(Duration::from_millis(50), Duration::from_millis(25));
 
-        let context = manager.initiate_listener_drain("timeout-test".to_string(), true, None, 5).await.unwrap();
+        let context = manager.initiate_listener_drain(
+            "timeout-test".to_string(), 
+            ListenerProtocolConfig::Http { drain_timeout: None }, 
+            5
+        ).await.unwrap();
 
         sleep(Duration::from_millis(10)).await;
         sleep(Duration::from_millis(60)).await;
@@ -435,4 +523,51 @@ mod tests {
             "Expected manager to no longer track the listener after timeout"
         );
     }
+
+    #[tokio::test]
+    async fn test_mixed_protocol_drain_context() {
+        let strategy = DrainStrategy::Mixed {
+            global_timeout: Duration::from_secs(600),
+            http_drain_timeout: Duration::from_secs(5),
+            tcp_connections: true,
+            http_connections: true,
+        };
+        let context = ListenerDrainContext::new("test-mixed".to_string(), strategy, 10);
+
+        assert_eq!(context.get_active_connections().await, 10);
+        assert!(!context.is_completed().await);
+        assert_eq!(context.get_http_drain_timeout(), Some(Duration::from_secs(5)));
+        assert!(!context.is_timeout_exceeded());
+    }
+
+    #[tokio::test]
+    async fn test_listener_protocol_config_analysis() {
+        let http_config = ListenerProtocolConfig::from_listener_analysis(
+            true, false, Some(Duration::from_secs(10))
+        );
+        match http_config {
+            ListenerProtocolConfig::Http { drain_timeout } => {
+                assert_eq!(drain_timeout, Some(Duration::from_secs(10)));
+            }
+            _ => panic!("Expected HTTP config"),
+        }
+
+        let tcp_config = ListenerProtocolConfig::from_listener_analysis(false, true, None);
+        match tcp_config {
+            ListenerProtocolConfig::Tcp => {},
+            _ => panic!("Expected TCP config"),
+        }
+
+        let mixed_config = ListenerProtocolConfig::from_listener_analysis(
+            true, true, Some(Duration::from_secs(3))
+        );
+        match mixed_config {
+            ListenerProtocolConfig::Mixed { http_drain_timeout, has_tcp, has_http } => {
+                assert_eq!(http_drain_timeout, Some(Duration::from_secs(3)));
+                assert!(has_tcp);
+                assert!(has_http);
+            }
+            _ => panic!("Expected Mixed config"),
+        }
+    }
 }
diff --git a/orion-lib/src/listeners/filterchain.rs b/orion-lib/src/listeners/filterchain.rs
@@ -236,7 +236,7 @@ impl FilterchainType {
                         stream,
                         hyper::service::service_fn(|req: Request<hyper::body::Incoming>| {
                             let handler_req =
-                                ExtendedRequest { request: req, downstream_metadata: downstream_metadata.clone() };
+                                ExtendedRequest { request: req, downstream_metadata: Arc::new(downstream_metadata.connection.clone()) };
                             req_handler.call(handler_req).map_err(orion_error::Error::into_inner)
                         }),
                     )
diff --git a/orion-lib/src/listeners/http_connection_manager.rs b/orion-lib/src/listeners/http_connection_manager.rs
@@ -71,10 +71,10 @@ use parking_lot::Mutex;
 use route::MatchedRequest;
 use scopeguard::defer;
 use std::collections::HashSet;
-use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
+use std::sync::atomic::AtomicUsize;
 use std::thread::ThreadId;
 use std::time::Instant;
-use std::{fmt, future::Future, result::Result as StdResult, sync::Arc};
+use std::{fmt, future::Future, result::Result as StdResult, sync::{Arc, LazyLock}};
 use tokio::sync::mpsc::Permit;
 use tokio::sync::watch;
 use upgrades as upgrade_utils;
@@ -92,7 +92,7 @@ use crate::{
     body::body_with_timeout::BodyWithTimeout,
     listeners::{
         access_log::AccessLogContext, drain_signaling::DrainSignalingManager,
-        filter_state::DownstreamConnectionMetadata, listeners_manager::ConnectionManager, rate_limiter::LocalRateLimit,
+        filter_state::{DownstreamConnectionMetadata, DownstreamMetadata}, listeners_manager::ConnectionManager, rate_limiter::LocalRateLimit,
         synthetic_http_response::SyntheticHttpResponse,
     },
     utils::http::{request_head_size, response_head_size},
@@ -102,6 +102,9 @@ use orion_tracing::http_tracer::{HttpTracer, SpanKind, SpanName};
 use orion_tracing::request_id::{RequestId, RequestIdManager};
 use orion_tracing::trace_context::TraceContext;
 
+static EMPTY_HASHMAP: LazyLock<Arc<HashMap<RouteMatch, Vec<Arc<HttpFilter>>>>> = 
+    LazyLock::new(|| Arc::new(HashMap::new()));
+
 #[derive(Debug, Clone)]
 pub struct HttpConnectionManagerBuilder {
     listener_name: Option<&'static str>,
@@ -405,7 +408,7 @@ impl HttpConnectionManager {
     }
 
     pub fn remove_route(&self) {
-        self.http_filters_per_route.swap(Arc::new(HashMap::new()));
+        self.http_filters_per_route.swap(EMPTY_HASHMAP.clone());
         let _ = self.router_sender.send_replace(None);
     }
 
@@ -418,10 +421,11 @@ impl HttpConnectionManager {
     }
 
     pub async fn start_draining(&self, drain_state: crate::listeners::drain_signaling::ListenerDrainState) {
-        if let Some(drain_timeout) = self.drain_timeout {
-            let listener_id = format!("{}-{}", self.listener_name, self.filter_chain_match_hash);
-            let _ = self.drain_signaling.initiate_listener_drain(listener_id, true, Some(drain_timeout), 0).await;
-        }
+        let listener_id = format!("{}-{}", self.listener_name, self.filter_chain_match_hash);
+        let protocol_config = crate::listeners::drain_signaling::ListenerProtocolConfig::Http { 
+            drain_timeout: self.drain_timeout 
+        };
+        let _ = self.drain_signaling.initiate_listener_drain(listener_id, protocol_config, 0).await;
 
         self.drain_signaling.start_listener_draining(drain_state).await;
     }
@@ -628,7 +632,7 @@ pub(crate) struct HttpRequestHandler {
 
 pub struct ExtendedRequest<B> {
     pub request: Request<B>,
-    pub downstream_metadata: Arc<DownstreamMetadata>,
+    pub downstream_metadata: Arc<DownstreamConnectionMetadata>,
 }
 
 #[derive(Debug)]
@@ -1174,7 +1178,7 @@ impl Service<ExtendedRequest<Incoming>> for HttpRequestHandler {
 
             //
             // 1. evaluate InitHttpContext, if logging is enabled
-            eval_http_init_context(&request, &trans_handler, downstream_metadata.server_name.as_deref());
+            eval_http_init_context(&request, &trans_handler, None);
 
             //
             // 2. create the MetricsBody, which will track the size of the request body
@@ -1329,9 +1333,14 @@ impl Service<ExtendedRequest<Incoming>> for HttpRequestHandler {
                 return Ok(response);
             };
 
+            let downstream_metadata_with_server_name = Arc::new(DownstreamMetadata::new(
+                downstream_metadata.as_ref().clone(),
+                None::<String>,
+            ));
+            
             let response = trans_handler
                 .clone()
-                .handle_transaction(route_conf, manager, permit, request, downstream_metadata)
+                .handle_transaction(route_conf, manager, permit, request, downstream_metadata_with_server_name)
                 .await;
 
             trans_handler.trace_status_code(response, listener_name)
diff --git a/orion-lib/src/listeners/lds_update.rs b/orion-lib/src/listeners/lds_update.rs
@@ -127,22 +127,13 @@ impl LdsManager {
 
             let mut listeners_guard = listeners.write().await;
             if let Some(versions) = listeners_guard.get_vec_mut(&name) {
-                let mut to_remove = Vec::new();
-                for (i, listener_info) in versions.iter().enumerate() {
-                    if listener_info.is_draining() {
-                        to_remove.push(i);
-                    }
-                }
-
-                for &index in to_remove.iter().rev() {
-                    if let Some(listener_info) = versions.get_mut(index) {
+                versions.iter_mut()
+                    .filter(|listener_info| listener_info.is_draining())
+                    .for_each(|listener_info| {
                         listener_info.handle.abort();
                         info!("LDS: Draining version of listener '{}' forcibly closed after timeout", name);
-                    }
-                }
-                for &index in to_remove.iter().rev() {
-                    versions.remove(index);
-                }
+                    });
+                versions.retain(|listener_info| !listener_info.is_draining());
                 if versions.is_empty() {
                     listeners_guard.remove(&name);
                 }
diff --git a/orion-lib/src/listeners/listener.rs b/orion-lib/src/listeners/listener.rs
diff --git a/orion-lib/src/listeners/listeners_manager.rs b/orion-lib/src/listeners/listeners_manager.rs

Original file line number	Diff line number	Diff line change
`@@ -236,7 +236,7 @@ impl FilterchainType {`
`236`	`236`	`stream,`
`237`	`237`	`hyper::service::service_fn(\|req: Request<hyper::body::Incoming>\| {`
`238`	`238`	`let handler_req =`
`239`		`- ExtendedRequest { request: req, downstream_metadata: downstream_metadata.clone() };`
	`239`	`+ ExtendedRequest { request: req, downstream_metadata: Arc::new(downstream_metadata.connection.clone()) };`
`240`	`240`	`req_handler.call(handler_req).map_err(orion_error::Error::into_inner)`
`241`	`241`	`}),`
`242`	`242`	`)`