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

Eeshu-Yadav · Eeshu-Yadav · commit 2da8a930c763 · 2025-09-24T02:36:03.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
@@ -16,7 +16,7 @@
 //
 
 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 +25,13 @@ 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 +53,8 @@ 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 +109,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 +121,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 +135,30 @@ 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 +234,21 @@ 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 +357,33 @@ 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 +461,8 @@ 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 +476,14 @@ 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 +501,47 @@ 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
@@ -235,8 +235,10 @@ impl FilterchainType {
                     .serve_connection_with_upgrades(
                         stream,
                         hyper::service::service_fn(|req: Request<hyper::body::Incoming>| {
-                            let handler_req =
-                                ExtendedRequest { request: req, downstream_metadata: downstream_metadata.clone() };
+                            let handler_req = 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,15 @@ 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;
@@ -91,8 +96,11 @@ use crate::{
 use crate::{
     body::body_with_timeout::BodyWithTimeout,
     listeners::{
-        access_log::AccessLogContext, drain_signaling::DrainSignalingManager,
-        filter_state::DownstreamConnectionMetadata, listeners_manager::ConnectionManager, rate_limiter::LocalRateLimit,
+        access_log::AccessLogContext,
+        drain_signaling::DrainSignalingManager,
+        filter_state::{DownstreamConnectionMetadata, DownstreamMetadata},
+        listeners_manager::ConnectionManager,
+        rate_limiter::LocalRateLimit,
         synthetic_http_response::SyntheticHttpResponse,
     },
     utils::http::{request_head_size, response_head_size},
@@ -102,6 +110,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 +416,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 +429,10 @@ 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 +639,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 +1185,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 +1340,12 @@ 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,11 @@ 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) {
-                        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.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);
+                });
+                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
@@ -44,7 +44,7 @@ use std::{
     fmt::Debug,
     net::SocketAddr,
     sync::{
-        atomic::{AtomicBool, Ordering},
+        atomic::{AtomicBool, AtomicU64, Ordering},
         Arc,
     },
 };
@@ -53,7 +53,8 @@ use tokio::{
     sync::broadcast::{self},
 };
 use tracing::{debug, info, warn};
-use uuid;
+
+static CONNECTION_COUNTER: AtomicU64 = AtomicU64::new(1);
 
 #[derive(Debug, Clone)]
 struct PartialListener {
@@ -375,7 +376,8 @@ impl Listener {
         drain_handler: Option<Arc<DefaultConnectionHandler>>,
     ) -> Result<()> {
         let shard_id = std::thread::current().id();
-        let connection_id = format!("{}:{}:{}", local_address, peer_addr, uuid::Uuid::new_v4());
+        let connection_id =
+            format!("{}:{}:{}", local_address, peer_addr, CONNECTION_COUNTER.fetch_add(1, Ordering::Relaxed));
 
         debug!("New connection {} established on listener {}", connection_id, listener_name);
 
diff --git a/orion-lib/src/listeners/listeners_manager.rs b/orion-lib/src/listeners/listeners_manager.rs
