listeners: preserve distinct sockets across graceful upgrades

pingora-core/src/listeners/l4.rs

@@ -16,6 +16,7 @@
 use log::debug;
 use log::warn;
 use pingora_error::{
+    Error,
     ErrorType::{AcceptError, BindError},
     OrErr, Result,
 };
@@ -47,14 +48,14 @@ use crate::server::ListenFds;
 #[cfg(unix)]
 use std::sync::LazyLock;
 
-/// Per-address async lock map for serializing the check-bind-insert sequence
+/// Per-key async lock map for serializing the check-bind-insert sequence
 /// in [`ListenerEndpointBuilder::listen`].
 ///
 /// With `ListenFds` using a synchronous `parking_lot::Mutex`, the lock cannot
 /// be held across `bind().await`. This global map ensures that only one task at
 /// a time can be in the process of looking up, binding, and inserting a given
-/// address — preventing two concurrent callers from both seeing "not found" and
-/// racing to bind the same address.
+/// key — preventing two concurrent callers from both seeing "not found" and
+/// racing to bind the same key.
 #[cfg(unix)]
 static BIND_LOCKS: LazyLock<flurry::HashMap<String, Arc<tokio::sync::Mutex<()>>>> =
     LazyLock::new(flurry::HashMap::new);
@@ -89,6 +90,11 @@ impl ServerAddress {
             _ => None,
         }
     }
+
+    pub(crate) fn fd_transfer_key(&self, id: Option<&str>) -> String {
+        let addr = self.as_ref();
+        id.map_or_else(|| addr.to_string(), |id| format!("{addr}#id={id}"))
+    }
 }
 
 /// TCP socket configuration options, this is used for setting options on
@@ -307,6 +313,8 @@ pub struct ListenerEndpoint {
 #[derive(Default)]
 pub struct ListenerEndpointBuilder {
     listen_addr: Option<ServerAddress>,
+    #[cfg(unix)]
+    fd_transfer_key: Option<String>,
     #[cfg(feature = "connection_filter")]
     connection_filter: Option<Arc<dyn ConnectionFilter>>,
 }
@@ -315,6 +323,8 @@ impl ListenerEndpointBuilder {
     pub fn new() -> ListenerEndpointBuilder {
         Self {
             listen_addr: None,
+            #[cfg(unix)]
+            fd_transfer_key: None,
             #[cfg(feature = "connection_filter")]
             connection_filter: None,
         }
@@ -325,6 +335,12 @@ impl ListenerEndpointBuilder {
         self
     }
 
+    #[cfg(unix)]
+    pub(crate) fn fd_transfer_key(&mut self, key: String) -> &mut Self {
+        self.fd_transfer_key = Some(key);
+        self
+    }
+
     #[cfg(feature = "connection_filter")]
     pub fn connection_filter(&mut self, filter: Arc<dyn ConnectionFilter>) -> &mut Self {
         self.connection_filter = Some(filter);
@@ -338,36 +354,56 @@ impl ListenerEndpointBuilder {
             .expect("Tried to listen with no addr specified");
 
         let listener = if let Some(fds_table) = fds {
-            let addr_str = listen_addr.as_ref();
-
-            // Acquire a per-address async lock so that only one task at a
+            let key = self
+                .fd_transfer_key
+                .as_deref()
+                .unwrap_or_else(|| listen_addr.as_ref());
+            // Acquire a per-key async lock so that only one task at a
             // time can go through the check-bind-insert sequence for a given
-            // address. The flurry guard is dropped before the await so its
+            // key. The flurry guard is dropped before the await so its
             // !Send pointer does not cross an await point.
-            let addr_lock = {
+            let key_lock = {
                 let guard = BIND_LOCKS.pin();
-                match guard.get(addr_str) {
+                match guard.get(key) {
                     Some(existing) => existing.clone(),
                     None => {
-                        let new_lock = Arc::new(tokio::sync::Mutex::new(()));
-                        match guard.try_insert(addr_str.to_string(), new_lock.clone()) {
+                        let lock = Arc::new(tokio::sync::Mutex::new(()));
+                        match guard.try_insert(key.to_string(), lock.clone()) {
                             Ok(inserted) => inserted.clone(),
                             Err(e) => e.current.clone(),
                         }
                     }
                 }
             };
-            let _guard = addr_lock.lock().await;
-
-            let existing_fd = fds_table.lock().get(addr_str).copied();
+            let _guard = key_lock.lock().await;
+
+            let existing_fd = {
+                let fds = fds_table.lock();
+                if fds.is_local(key) {
+                    return Error::e_explain(
+                        BindError,
+                        format!("duplicate listener transfer identity {key}"),
+                    );
+                }
+                fds.get(key).copied()
+            };
 
             if let Some(fd) = existing_fd {
-                from_raw_fd(&listen_addr, fd)?
+                let listener = from_raw_fd(&listen_addr, fd)?;
+                fds_table.lock().mark_local(key);
+                listener
             } else {
                 let listener = bind(&listen_addr).await?;
-                fds_table
+                if fds_table
                     .lock()
-                    .add(addr_str.to_string(), listener.as_raw_fd());
+                    .try_add(key.to_string(), listener.as_raw_fd())
+                    .is_err()
+                {
+                    return Error::e_explain(
+                        BindError,
+                        format!("duplicate listener transfer identity {key}"),
+                    );
+                }
                 listener
             }
         } else {