From a8283d49c04f365a6060234a22814ebbc3360ff3 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Tue, 23 Jan 2024 14:59:21 +0000
Subject: [PATCH] add api to check whether SQ and CQ count against memlock
 rlimit

---
 tokio-epoll-uring/src/lib.rs                  |  1 +
 .../src/system/kernel_support.rs              | 56 +++++++++++++++++++
 tokio-epoll-uring/src/system/mod.rs           |  2 +
 3 files changed, 59 insertions(+)
 create mode 100644 tokio-epoll-uring/src/system/kernel_support.rs

diff --git a/tokio-epoll-uring/src/lib.rs b/tokio-epoll-uring/src/lib.rs
index cf7a826..12e87be 100644
--- a/tokio-epoll-uring/src/lib.rs
+++ b/tokio-epoll-uring/src/lib.rs
@@ -80,6 +80,7 @@ mod system;
 #[cfg(target_os = "linux")]
 pub use {
     crate::system::submission::op_fut::Error,
+    system::kernel_support,
     system::lifecycle::handle::SystemHandle,
     system::lifecycle::thread_local::{thread_local_system, Handle},
     system::lifecycle::System,
diff --git a/tokio-epoll-uring/src/system/kernel_support.rs b/tokio-epoll-uring/src/system/kernel_support.rs
new file mode 100644
index 0000000..b6102bc
--- /dev/null
+++ b/tokio-epoll-uring/src/system/kernel_support.rs
@@ -0,0 +1,56 @@
+use std::{io::Write, os::unix::process::CommandExt};
+
+/// Check whether a [`crate::System`]'s consumes memlock rlimit budget.
+///
+/// Unsafe because this function forks & uses [`std::os::unix::process::Command::pre_exec`].
+///
+/// Panics in a bunch of places.
+///
+pub unsafe fn does_system_consume_memlock_rlimit_budget() -> bool {
+    // Fork a child process.
+    // In the child process, set up a tokio-epoll-uring system.
+    // Compare the memlock rlimit usage before and after.
+
+    // poor man's bincode
+    #[repr(C, packed)]
+    struct RlimitPrePost {
+        pre: (u64, u64),
+        post: (u64, u64),
+    }
+    // Ensure `true` binary exists.
+    std::process::Command::new("true")
+        .output()
+        .expect("canont exec `true` binary");
+    let output = std::process::Command::new("true")
+        .pre_exec(|| {
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()?;
+            let pre = nix::sys::resource::getrlimit(nix::sys::resource::Resource::RLIMIT_MEMLOCK)
+                .unwrap();
+            let system = rt.block_on(crate::System::launch()).unwrap();
+            // use the memory
+            let ((), res) = rt.block_on(system.nop());
+            res.unwrap();
+            let post = nix::sys::resource::getrlimit(nix::sys::resource::Resource::RLIMIT_MEMLOCK)
+                .unwrap();
+            let output = RlimitPrePost { pre, post };
+            let output_as_slice = std::slice::from_raw_parts(
+                &output as *const _ as *const u8,
+                std::mem::size_of_val(&output),
+            );
+            std::io::stdout().write_all(output_as_slice).unwrap();
+            Ok(())
+        })
+        .output()
+        .unwrap();
+    assert_eq!(output.status.code().unwrap(), 0);
+    assert_eq!(output.stdout.len(), std::mem::size_of::<RlimitPrePost>());
+    let (head, output, tail): (_, &[RlimitPrePost], _) = output.stdout.align_to();
+    assert_eq!(head.len(), 0);
+    assert_eq!(tail.len(), 0);
+    assert_eq!(output.len(), 1);
+    let RlimitPrePost { pre, post } = output[0];
+    assert!(pre <= post, "above code can't use _less_ locked memory");
+    pre < post
+}
diff --git a/tokio-epoll-uring/src/system/mod.rs b/tokio-epoll-uring/src/system/mod.rs
index 1639bb5..65500d0 100644
--- a/tokio-epoll-uring/src/system/mod.rs
+++ b/tokio-epoll-uring/src/system/mod.rs
@@ -8,3 +8,5 @@ mod test_util;
 mod tests;
 
 pub(crate) const RING_SIZE: u32 = 128;
+
+pub mod kernel_support;