feat(s2n-quic-xdp): add umem hugepage option (#2145)

ollie-etl · web-flow · commit 2d4df73eddf3 · 2024-03-14T08:14:14.000-06:00
diff --git a/tools/xdp/s2n-quic-xdp/src/mmap.rs b/tools/xdp/s2n-quic-xdp/src/mmap.rs
@@ -19,6 +19,12 @@ pub struct Mmap {
     len: usize,
 }
 
+#[derive(Debug)]
+pub enum Options {
+    Huge,
+    Fd(RawFd),
+}
+
 /// Safety: Mmap pointer can be sent between threads
 unsafe impl Send for Mmap {}
 
@@ -28,8 +34,12 @@ unsafe impl Sync for Mmap {}
 impl Mmap {
     /// Creates a new mmap'd region, with an optional file descriptor.
     #[inline]
-    pub fn new(len: usize, offset: usize, fd: Option<RawFd>) -> Result<Self> {
-        let addr = mmap(len, offset, fd)?;
+    pub fn new(len: usize, offset: usize, flags: Option<Options>) -> Result<Self> {
+        let addr = match flags {
+            Some(Options::Huge) => mmap(len, offset, None, true),
+            Some(Options::Fd(fd)) => mmap(len, offset, Some(fd), false),
+            None => mmap(len, offset, None, false),
+        }?;
         Ok(Self { addr, len })
     }
 
diff --git a/tools/xdp/s2n-quic-xdp/src/ring.rs b/tools/xdp/s2n-quic-xdp/src/ring.rs
@@ -3,7 +3,7 @@
 
 use crate::{
     if_xdp::{MmapOffsets, RingFlags, RingOffsetV2, RxTxDescriptor, UmemDescriptor},
-    mmap::Mmap,
+    mmap::{self, Mmap},
     socket, syscall,
 };
 use core::{fmt, mem::size_of, ptr::NonNull};
@@ -94,7 +94,7 @@ macro_rules! impl_producer {
             // Use the hard-coded offset of the ring type
             let offset = MmapOffsets::$offset;
 
-            let area = Mmap::new(len, offset, Some(socket.as_raw_fd()))?;
+            let area = Mmap::new(len, offset, Some(mmap::Options::Fd(socket.as_raw_fd())))?;
 
             let (cursor, flags) = unsafe {
                 // Safety: `area` lives as long as `cursor`
@@ -193,7 +193,7 @@ macro_rules! impl_consumer {
             // Use the hard-coded offset of the ring type
             let offset = MmapOffsets::$offset;
 
-            let area = Mmap::new(len, offset, Some(socket.as_raw_fd()))?;
+            let area = Mmap::new(len, offset, Some(mmap::Options::Fd(socket.as_raw_fd())))?;
 
             let (cursor, flags) = unsafe {
                 // Safety: `area` lives as long as `cursor`
diff --git a/tools/xdp/s2n-quic-xdp/src/syscall.rs b/tools/xdp/s2n-quic-xdp/src/syscall.rs
@@ -195,12 +195,15 @@ pub fn busy_poll<Fd: AsRawFd>(fd: &Fd) -> Result<u32> {
 ///
 /// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L273).
 #[inline]
-pub fn mmap(len: usize, offset: usize, fd: Option<RawFd>) -> Result<NonNull<c_void>> {
-    let flags = if fd.is_some() {
+pub fn mmap(len: usize, offset: usize, fd: Option<RawFd>, huge: bool) -> Result<NonNull<c_void>> {
+    let mut flags = if fd.is_some() {
         libc::MAP_SHARED | libc::MAP_POPULATE
     } else {
         libc::MAP_PRIVATE | libc::MAP_ANONYMOUS
     };
+    if huge {
+        flags |= libc::MAP_HUGETLB;
+    }
 
     // See:
     // * Fill https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L273
diff --git a/tools/xdp/s2n-quic-xdp/src/umem.rs b/tools/xdp/s2n-quic-xdp/src/umem.rs
@@ -3,7 +3,7 @@
 
 use crate::{
     if_xdp::{RxTxDescriptor, UmemDescriptor, UmemFlags, UmemReg},
-    mmap::Mmap,
+    mmap::{self, Mmap},
     syscall, Result,
 };
 use core::ptr::NonNull;
@@ -22,6 +22,8 @@ pub struct Builder {
     pub frame_headroom: u32,
     /// The flags for the Umem
     pub flags: UmemFlags,
+    /// Back the umem with a hugepage
+    pub hugepage: bool,
 }
 
 impl Default for Builder {
@@ -31,14 +33,20 @@ impl Default for Builder {
             frame_count: 1024,
             frame_headroom: 0,
             flags: Default::default(),
+            hugepage: false,
         }
     }
 }
 
 impl Builder {
     pub fn build(self) -> Result<Umem> {
         let len = self.frame_size as usize * self.frame_count as usize;
-        let area = Mmap::new(len, 0, None)?;
+        let options = if self.hugepage {
+            Some(mmap::Options::Huge)
+        } else {
+            None
+        };
+        let area = Mmap::new(len, 0, options)?;
         let area = Arc::new(area);
         let mem = area.addr().cast();
 
