bench uring

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

bench-vortex/Cargo.toml

@@ -18,6 +18,7 @@ workspace = true
 
 [features]
 lance = ["dep:lance", "dep:lance-encoding"]
+uring = ["vortex/uring"]
 
 [dependencies]
 lance = { version = "0.39.0", optional = true }

bench-vortex/src/bin/query_bench.rs

@@ -4,6 +4,7 @@
 use std::path::PathBuf;
 
 use bench_vortex::IdempotentPath as _;
+use bench_vortex::RuntimeChoice;
 use bench_vortex::Target;
 use bench_vortex::benchmark_driver::DriverConfig;
 use bench_vortex::benchmark_driver::run_benchmark;
@@ -62,6 +63,10 @@ struct CommonArgs {
     #[arg(short, long)]
     threads: Option<usize>,
 
+    /// Runtime to use for Vortex operations.
+    #[arg(long, value_enum)]
+    runtime: Option<RuntimeChoice>,
+
     #[arg(short, long)]
     verbose: bool,
 
@@ -283,6 +288,7 @@ fn validate_scale_factor(val: &str) -> Result<String, String> {
 
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
+    bench_vortex::configure_runtime(args.command.common().runtime);
     match args.command {
         Commands::ClickBench(clickbench_args) => run_clickbench(clickbench_args),
         Commands::TpcH(tpch_args) => run_tpch(tpch_args),
@@ -293,6 +299,23 @@ fn main() -> anyhow::Result<()> {
     }
 }
 
+trait CommandRuntime {
+    fn common(&self) -> &CommonArgs;
+}
+
+impl CommandRuntime for Commands {
+    fn common(&self) -> &CommonArgs {
+        match self {
+            Commands::ClickBench(args) => &args.common,
+            Commands::TpcH(args) => &args.common,
+            Commands::TpcDS(args) => &args.common,
+            Commands::StatPopGen(args) => &args.common,
+            Commands::Fineweb(args) => &args.common,
+            Commands::GhArchive(args) => &args.common,
+        }
+    }
+}
+
 fn run_clickbench(args: ClickBenchArgs) -> anyhow::Result<()> {
     setup_logging_and_tracing(args.common.verbose, args.common.tracing)?;
 

bench-vortex/src/lib.rs

@@ -8,6 +8,7 @@ use std::clone::Clone;
 use std::fmt::Display;
 use std::str::FromStr;
 use std::sync::LazyLock;
+use std::sync::OnceLock;
 
 use clap::ValueEnum;
 use itertools::Itertools;
@@ -55,8 +56,37 @@ use vortex::session::VortexSession;
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
-pub static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_tokio());
+pub static SESSION: LazyLock<VortexSession> = LazyLock::new(make_session);
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum, Serialize)]
+#[serde(rename_all = "kebab-case")]
+pub enum RuntimeChoice {
+    Tokio,
+    #[cfg(all(feature = "uring", target_os = "linux"))]
+    Uring,
+}
+
+static RUNTIME_OVERRIDE: OnceLock<RuntimeChoice> = OnceLock::new();
+
+/// Configure the runtime used by benchmarks. Must be called before `SESSION` is first accessed.
+pub fn configure_runtime(choice: Option<RuntimeChoice>) {
+    if let Some(choice) = choice {
+        let _ = RUNTIME_OVERRIDE.set(choice);
+    }
+}
+
+fn make_session() -> VortexSession {
+    match RUNTIME_OVERRIDE.get().copied().unwrap_or(RuntimeChoice::Tokio) {
+        RuntimeChoice::Tokio => VortexSession::default().with_tokio(),
+        #[cfg(all(feature = "uring", target_os = "linux"))]
+        RuntimeChoice::Uring => {
+            use vortex::io::runtime::uring::PerCoreUringPool;
+            static URING_POOL: OnceLock<PerCoreUringPool> = OnceLock::new();
+            let pool = URING_POOL.get_or_init(|| PerCoreUringPool::new(None));
+            VortexSession::default().with_handle(pool.handle())
+        }
+    }
+}
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize)]
 pub struct Target {

vortex/Cargo.toml

@@ -83,6 +83,7 @@ serde = [
 ]
 # This feature enabled unstable encodings for which we don't guarantee stability.
 unstable_encodings = ["vortex-btrblocks/unstable_encodings"]
+uring = ["vortex-io/uring"]
 
 [[bench]]
 name = "single_encoding_throughput"