vey-daemon: enable ebpf with ListenUdpRuntime

Author: zh-jq

Cargo.lock

@@ -53,10 +53,12 @@ libc.workspace = true
 
 [target.'cfg(target_os = "linux")'.dependencies]
 vey-journal.workspace = true
+vey-reuseport = { workspace = true, optional = true }
 
 [features]
 default = []
 jemalloc = ["dep:vey-jemalloc"]
+ebpf = ["dep:vey-reuseport"]
 event-log = ["dep:vey-fluentd"]
 register = ["vey-yaml/http", "dep:http", "dep:serde_json", "dep:vey-http"]
 quic = ["dep:quinn", "vey-types/acl-rule"]

lib/vey-daemon/Cargo.toml

@@ -131,7 +131,7 @@ impl<T: QuitAction> QuitLoop<T> {
                             Ok(None) => break 'outer,
                             Err(_) => {
                                 if is_restart {
-                                    // the new process may be failed to start, so we consume
+                                    // the new process may be failed to start, so we resume
                                     info!(
                                         "timeout to wait StartGracefulShutdown request, will resume"
                                     );
@@ -140,7 +140,7 @@ impl<T: QuitAction> QuitLoop<T> {
                                     continue 'outer;
                                 } else {
                                     // treat timeout to stop
-                                    info!("timeout to wait new request, will stop");
+                                    info!("timeout to wait new controller request, will stop");
                                     self.release_controller().await;
                                     break;
                                 }

lib/vey-daemon/src/control/quit.rs

@@ -16,6 +16,8 @@ use std::collections::VecDeque;
 use std::future::poll_fn;
 use std::io::{self, IoSlice, IoSliceMut};
 use std::net::SocketAddr;
+#[cfg(all(target_os = "linux", feature = "ebpf"))]
+use std::os::fd::AsRawFd;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -33,6 +35,9 @@ use tokio::sync::{broadcast, mpsc};
 
 use vey_io_ext::{UdpMoveRecv, UdpMoveSend, UdpSocketExt};
 use vey_io_sys::udp::{RecvMsgHdr, SendMsgHdr};
+#[cfg(all(target_os = "linux", feature = "ebpf"))]
+use vey_reuseport::udp::UdpSocketSelector;
+use vey_std_ext::core::NonZeroExt;
 use vey_types::net::{UdpConnectionTrackConfig, UdpListenConfig};
 
 use crate::listen::{ListenAliveGuard, ListenStats};
@@ -511,11 +516,21 @@ where
             }
         }
 
+        #[cfg(all(target_os = "linux", feature = "ebpf"))]
+        let mut udp_reuseport_selector = UdpSocketSelector::new(
+            rustix::process::getpid().as_raw_pid(),
+            self.server.version() as u32,
+            listen_config.address(),
+            self.conn_track.max_sessions(),
+        )?;
+
         for i in 0..instance_count {
             let mut runtime = self.create_instance();
             runtime.instance_id = i;
 
             let socket = vey_socket::udp::new_std_bind_listen(listen_config)?;
+            #[cfg(all(target_os = "linux", feature = "ebpf"))]
+            udp_reuseport_selector.add_socket(socket.as_raw_fd());
             let listen_addr = socket.local_addr()?;
             runtime.into_running(
                 socket,
@@ -524,6 +539,21 @@ where
                 server_reload_sender.subscribe(),
             );
         }
+
+        #[cfg(all(target_os = "linux", feature = "ebpf"))]
+        {
+            udp_reuseport_selector.load_and_attach()?;
+            let mut server_reload_receiver = server_reload_sender.subscribe();
+            tokio::spawn(async move {
+                while let Ok(cmd) = server_reload_receiver.recv().await {
+                    if matches!(cmd, ServerReloadCommand::QuitRuntime) {
+                        break;
+                    }
+                }
+                drop(udp_reuseport_selector);
+            });
+        }
+
         Ok(())
     }
 }
@@ -604,8 +634,10 @@ where
     ) {
         use broadcast::error::RecvError;
 
-        let mut ct_table =
-            LruCache::with_hasher(self.conn_track.max_sessions(), FixedState::with_seed(0));
+        let mut ct_table = LruCache::with_hasher(
+            self.conn_track.max_sessions().cast_usize(),
+            FixedState::with_seed(0),
+        );
         let mut rt_state = RuntimeState::new(socket, self.conn_track.send_queue_size());
 
         let mut event_recv_buf: Vec<Event> = Vec::with_capacity(EVENT_RECV_BATCH_SIZE);

lib/vey-daemon/src/listen/udp/listen.rs

@@ -6,6 +6,7 @@
 use std::fs;
 use std::io;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket};
+use std::num::NonZeroU32;
 use std::os::fd::AsRawFd;
 use std::thread;
 use std::time::Duration;
@@ -38,6 +39,8 @@ fn main() -> anyhow::Result<()> {
     println!("   UdpSocketSelector Hot-Upgrade & Fallback Test  ");
     println!("==================================================");
 
+    let max_entries = NonZeroU32::new(1024).unwrap();
+
     // 1. Check root privileges
     if unsafe { libc::getuid() } != 0 {
         println!("[ERROR] This test case must be run as root to load eBPF programs and pin maps.");
@@ -61,8 +64,8 @@ fn main() -> anyhow::Result<()> {
     );
 
     let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);
-    let mut selector_gen1 =
-        UdpSocketSelector::new(1234, 1, addr, 1024).context("failed to create Gen 1 selector")?;
+    let mut selector_gen1 = UdpSocketSelector::new(1234, 1, addr, max_entries)
+        .context("failed to create Gen 1 selector")?;
     selector_gen1.add_socket(s1_gen1.as_raw_fd());
     selector_gen1.add_socket(s2_gen1.as_raw_fd());
 
@@ -116,8 +119,8 @@ fn main() -> anyhow::Result<()> {
         s2_gen2.as_raw_fd()
     );
 
-    let mut selector_gen2 =
-        UdpSocketSelector::new(1234, 2, addr, 1024).context("failed to create Gen 2 selector")?;
+    let mut selector_gen2 = UdpSocketSelector::new(1234, 2, addr, max_entries)
+        .context("failed to create Gen 2 selector")?;
     selector_gen2.add_socket(s1_gen2.as_raw_fd());
     selector_gen2.add_socket(s2_gen2.as_raw_fd());
 

lib/vey-reuseport/src/bin/udp_reuseport_test.rs

@@ -4,11 +4,11 @@
 #include <linux/bpf.h>
 #include <bpf/bpf_helpers.h>
 
-const volatile __u32 load_pid = 0;
+const volatile __s32 load_pid = 0;
 const volatile __u32 load_generation = 0;
 
 struct socket_id {
-	__u32 pid;
+	__s32 pid;
 	__u32 generation;
 	__u32 worker;
 };
@@ -22,7 +22,7 @@ struct {
 } conn_track SEC(".maps");
 
 struct proc_info_key {
-	__u32 pid;
+	__s32 pid;
 	__u32 generation;
 };
 

lib/vey-reuseport/src/bpf/udp.bpf.c

@@ -5,6 +5,7 @@
 
 use std::io;
 use std::net::{IpAddr, SocketAddr};
+use std::num::NonZeroU32;
 use std::os::fd::AsFd;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::path::{Path, PathBuf};
@@ -25,15 +26,15 @@ const NAME_PROGRAM: &str = "udp_select_reuseport";
 #[derive(IntoBytes, Immutable)]
 #[repr(C)]
 struct SocketId {
-    pid: u32,
+    pid: i32,
     generation: u32,
     worker: u32,
 }
 
 #[derive(IntoBytes, Immutable)]
 #[repr(C)]
 struct ProcMapKey {
-    pid: u32,
+    pid: i32,
     generation: u32,
 }
 
@@ -45,14 +46,14 @@ struct ProcMapValue {
 }
 
 struct ReadOnlyData {
-    load_pid: u32,
+    load_pid: i32,
     load_generation: u32,
 }
 
 pub struct UdpSocketSelector {
     pin_dir: PathBuf,
-    conn_track_max_entries: u32,
-    pid: u32,
+    conn_track_max_entries: NonZeroU32,
+    pid: i32,
     generation: u32,
     sockets: Vec<RawFd>,
     proc_map_handle: Option<MapHandle>,
@@ -65,10 +66,10 @@ impl UdpSocketSelector {
     }
 
     pub fn new(
-        pid: u32,
+        pid: i32,
         generation: u32,
         addr: SocketAddr,
-        conn_track_max_entries: u32,
+        conn_track_max_entries: NonZeroU32,
     ) -> anyhow::Result<Self> {
         let ip = match addr.ip() {
             IpAddr::V4(ip) => ip.to_ipv6_compatible(), // IPv4 "." is not allowed in path
@@ -148,8 +149,9 @@ impl UdpSocketSelector {
             };
             match name {
                 NAME_CONN_TRACK => {
+                    let max_entries = self.conn_track_max_entries.get();
                     if let Ok(handle) = MapHandle::from_pinned_path(&conn_track_path) {
-                        if handle.max_entries() != self.conn_track_max_entries {
+                        if handle.max_entries() != max_entries {
                             warn!(
                                 "udp conn_track map {} already pinned with max entries {}, delete it first if you want to set max entries to {}",
                                 conn_track_path.display(),
@@ -165,10 +167,9 @@ impl UdpSocketSelector {
                             )
                         })?;
                     } else {
-                        map.set_max_entries(self.conn_track_max_entries)
-                            .map_err(|e| {
-                                anyhow!("failed to set max entries for conn_track map: {e}")
-                            })?;
+                        map.set_max_entries(max_entries).map_err(|e| {
+                            anyhow!("failed to set max entries for conn_track map: {e}")
+                        })?;
                     }
                 }
                 NAME_SOCKET_MAP => {

lib/vey-reuseport/src/udp.rs

@@ -8,3 +8,6 @@ pub use option::OptionExt;
 
 mod zig_zag;
 pub use zig_zag::{FromZigZag, ToZigZag};
+
+mod non_zero;
+pub use non_zero::NonZeroExt;

lib/vey-std-ext/src/core/mod.rs

@@ -0,0 +1,17 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: 2026 VEY-OSS Developers.
+ */
+
+use std::num::{NonZeroU32, NonZeroUsize};
+
+pub trait NonZeroExt {
+    fn cast_usize(&self) -> NonZeroUsize;
+}
+
+impl NonZeroExt for NonZeroU32 {
+    fn cast_usize(&self) -> NonZeroUsize {
+        let v = self.get();
+        unsafe { NonZeroUsize::new_unchecked(v as usize) }
+    }
+}

lib/vey-std-ext/src/core/non_zero.rs

@@ -5,7 +5,7 @@
  */
 
 use std::net::{IpAddr, Ipv6Addr, SocketAddr};
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU32, NonZeroUsize};
 
 use anyhow::anyhow;
 use num_traits::ToPrimitive;
@@ -204,7 +204,7 @@ impl UdpListenConfig {
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct UdpConnectionTrackConfig {
-    max_sessions: NonZeroUsize,
+    max_sessions: NonZeroU32,
     dispatch_queue_size: NonZeroUsize,
     send_queue_size: NonZeroUsize,
     batch_recv_size: NonZeroUsize,
@@ -213,7 +213,7 @@ pub struct UdpConnectionTrackConfig {
 impl Default for UdpConnectionTrackConfig {
     fn default() -> Self {
         UdpConnectionTrackConfig {
-            max_sessions: unsafe { NonZeroUsize::new_unchecked(4096) },
+            max_sessions: unsafe { NonZeroU32::new_unchecked(32768) },
             dispatch_queue_size: unsafe { NonZeroUsize::new_unchecked(32) },
             send_queue_size: unsafe { NonZeroUsize::new_unchecked(512) },
             batch_recv_size: unsafe { NonZeroUsize::new_unchecked(16) },
@@ -223,12 +223,12 @@ impl Default for UdpConnectionTrackConfig {
 
 impl UdpConnectionTrackConfig {
     #[inline]
-    pub fn max_sessions(&self) -> NonZeroUsize {
+    pub fn max_sessions(&self) -> NonZeroU32 {
         self.max_sessions
     }
 
     #[inline]
-    pub fn set_max_sessions(&mut self, max_sessions: NonZeroUsize) {
+    pub fn set_max_sessions(&mut self, max_sessions: NonZeroU32) {
         self.max_sessions = max_sessions;
     }