diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index b1b5d8c8799..8d9df700393 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -95,6 +95,8 @@ jobs: RUST_BACKTRACE: full VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + # Temporary: exercise the Compio path for the vortex-compio PR. + VORTEX_IO_RUNTIME: compio run: | python3 scripts/random-access-split.py diff --git a/Cargo.lock b/Cargo.lock index 45d447e3cd4..a27ea199935 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1562,6 +1562,165 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "compio" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b84ee96a86948d04388f3a0b8c36b9f0a6b40b3528ac0d65737e53632fb37fe" +dependencies = [ + "compio-buf", + "compio-driver", + "compio-fs", + "compio-io", + "compio-log", + "compio-quic", + "compio-runtime", +] + +[[package]] +name = "compio-buf" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e8777c3ad31ab42f8a3a4a1bd629b78f688371df9b0f528d94dfbdbe5c945c9" +dependencies = [ + "arrayvec", + "bytes", + "libc", +] + +[[package]] +name = "compio-driver" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d42d98dc890ee4db00c1e68a723391711aab6d67085880d716b72830f7c715" +dependencies = [ + "cfg-if", + "cfg_aliases", + "compio-buf", + "compio-log", + "crossbeam-queue", + "flume", + "futures-util", + "io-uring", + "io_uring_buf_ring", + "libc", + "once_cell", + "paste", + "pin-project-lite", + "polling", + "slab", + "smallvec", + "socket2", + "synchrony", + "thin-cell", + "windows-sys 0.61.2", +] + +[[package]] +name = "compio-fs" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65ee36e1acf2cec4835efe9a986c012b2462c5ef53580e4ee84ae6d5a3d8e3b3" +dependencies = [ + "cfg-if", + "cfg_aliases", + "compio-buf", + "compio-driver", + "compio-io", + "compio-runtime", + "libc", + "os_pipe", + "pin-project-lite", + "widestring", + "windows-sys 0.61.2", +] + +[[package]] +name = "compio-io" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "637522f28a64fd5f7dcceaa4ddef13fa8d8020025e8c993f7a069e237835580e" +dependencies = [ + "compio-buf", + "futures-util", + "paste", + "synchrony", +] + +[[package]] +name = "compio-log" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc4e560213c1996b618da369b7c9109564b41af9033802ae534465c4ee4e132f" +dependencies = [ + "tracing", +] + +[[package]] +name = "compio-net" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "becd7d40522c885113752a3640cba9f9d347f205b646bb3f8ff3967173a228f2" +dependencies = [ + "cfg-if", + "compio-buf", + "compio-driver", + "compio-io", + "compio-runtime", + "either", + "libc", + "once_cell", + "socket2", + "widestring", + "windows-sys 0.61.2", +] + +[[package]] +name = "compio-quic" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad9efdad81b920108b9de57148e1b9d73dc408b6d06a59ee64836dde651cf026" +dependencies = [ + "cfg_aliases", + "compio-buf", + "compio-io", + "compio-log", + "compio-net", + "compio-runtime", + "flume", + "futures-util", + "libc", + "quinn-proto", + "rustc-hash", + "rustls", + "synchrony", + "thiserror 2.0.19", + "windows-sys 0.61.2", +] + +[[package]] +name = "compio-runtime" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6c1c71f011bdd9c8f30e97d877b606505ee6d241c7782cfaed172f66acbd9cd" +dependencies = [ + "async-task", + "cfg-if", + "compio-buf", + "compio-driver", + "compio-log", + "core_affinity", + "crossbeam-queue", + "futures-util", + "libc", + "once_cell", + "pin-project-lite", + "scoped-tls", + "slab", + "socket2", + "windows-sys 0.61.2", +] + [[package]] name = "compress-bench" version = "0.1.0" @@ -1751,6 +1910,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "cpubits" version = "0.1.1" @@ -3260,6 +3430,17 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" +[[package]] +name = "flume" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" +dependencies = [ + "futures-core", + "futures-sink", + "spin", +] + [[package]] name = "fnv" version = "1.0.7" @@ -4380,6 +4561,17 @@ dependencies = [ "libc", ] +[[package]] +name = "io_uring_buf_ring" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1838759bb8c2f24cf05a35429d83145c4aa6af43f8ad38477295e12a7320a80e" +dependencies = [ + "bytes", + "io-uring", + "rustix", +] + [[package]] name = "ipnet" version = "2.12.1" @@ -5481,6 +5673,7 @@ checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" dependencies = [ "cfg-if", "generator", + "pin-utils", "scoped-tls", "tracing", "tracing-subscriber", @@ -6327,6 +6520,16 @@ dependencies = [ "hashbrown 0.14.5", ] +[[package]] +name = "os_pipe" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8fae84b431384b68627d0f9b3b1245fcf9f46f6c0e3dc902e9dce64edd1967" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "owo-colors" version = "4.3.0" @@ -6666,6 +6869,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + [[package]] name = "ping" version = "0.7.1" @@ -7284,11 +7493,14 @@ dependencies = [ "clap", "indicatif", "lance-bench", + "parking_lot", "rand 0.10.2", "rand_distr 0.6.0", "tabled", "tokio", "vortex-bench", + "vortex-compio", + "vortex-io", ] [[package]] @@ -8433,6 +8645,15 @@ dependencies = [ "spatialbench", ] +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" +dependencies = [ + "lock_api", +] + [[package]] name = "sqllogictest" version = "0.29.1" @@ -8665,6 +8886,16 @@ dependencies = [ "futures-core", ] +[[package]] +name = "synchrony" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d6d5fbc4583cf3e5eee953506f13853a42ee6b4a21a983dffffb10c765cb80" +dependencies = [ + "futures-util", + "loom", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -8846,6 +9077,12 @@ dependencies = [ "unicode-width 0.2.2", ] +[[package]] +name = "thin-cell" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4164c6c316ba9733b0ab021e7f9852c788a4b991b49c25820f1be48e1d41345b" + [[package]] name = "thiserror" version = "1.0.69" @@ -9773,6 +10010,7 @@ dependencies = [ "bitvec", "bytes", "codspeed-divan-compat", + "compio-buf", "itertools 0.14.0", "memmap2", "num-traits", @@ -9835,6 +10073,25 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-compio" +version = "0.1.0" +dependencies = [ + "compio", + "futures", + "kanal", + "oneshot", + "parking_lot", + "tempfile", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-file", + "vortex-io", + "vortex-layout", + "vortex-metrics", +] + [[package]] name = "vortex-compressor" version = "0.1.0" @@ -10907,6 +11164,12 @@ dependencies = [ "libc", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index 1484b75de11..ef659732b1e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "vortex-flatbuffers", "vortex-metrics", "vortex-io", + "vortex-compio", "vortex-cloud", "vortex-proto", "vortex-array", @@ -127,6 +128,8 @@ cc = "1.2" cfg-if = "1.0.1" chrono = "0.4.44" clap = "4.5" +compio = { version = "=0.18.0", default-features = false } +compio-buf = "=0.8.0" criterion = { package = "codspeed-criterion-compat-walltime", version = "5.0.0" } crossterm = "0.29" cudarc = { version = "0.19.0", features = [ @@ -299,6 +302,7 @@ vortex-bytebool = { version = "0.1.0", path = "./encodings/bytebool", default-fe vortex-cloud = { version = "0.1.0", path = "./vortex-cloud", default-features = false } vortex-compressor = { version = "0.1.0", path = "./vortex-compressor", default-features = false } vortex-compute = { version = "0.1.0", path = "./vortex-compute", default-features = false } +vortex-compio = { version = "0.1.0", path = "./vortex-compio", default-features = false } vortex-datafusion = { version = "0.1.0", path = "./vortex-datafusion", default-features = false } vortex-datetime-parts = { version = "0.1.0", path = "./encodings/datetime-parts", default-features = false } vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte-parts", default-features = false } diff --git a/benchmarks/random-access-bench/Cargo.toml b/benchmarks/random-access-bench/Cargo.toml index c45c1ac2bff..d2010df3f6a 100644 --- a/benchmarks/random-access-bench/Cargo.toml +++ b/benchmarks/random-access-bench/Cargo.toml @@ -19,11 +19,14 @@ anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } indicatif = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } +parking_lot = { workspace = true } rand = { workspace = true } rand_distr = { workspace = true } tabled = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-bench = { workspace = true } +vortex-compio = { workspace = true } +vortex-io = { workspace = true } [features] lance = ["dep:lance-bench"] diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 14949c85fcb..6554d5f53aa 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -20,3 +20,15 @@ reopening the file per lookup. CI drives the full matrix via ```bash cargo run -p random-access-bench --profile release_debug --features lance ``` + +Set `VORTEX_IO_RUNTIME=compio` to use completion-based local Vortex reads instead of the default +Tokio blocking-I/O path. This only changes Vortex files; Parquet and Lance retain their native I/O +runtimes. + +```bash +cargo run -p random-access-bench --profile release_debug -- \ + --formats vortex --datasets feature-vectors --patterns uniform --open-mode cached --time-limit 5 + +VORTEX_IO_RUNTIME=compio cargo run -p random-access-bench --profile release_debug -- \ + --formats vortex --datasets feature-vectors --patterns uniform --open-mode cached --time-limit 5 +``` diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 11888d34c43..3a7f5ee058f 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -2,12 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::path::PathBuf; +use std::str::FromStr; +use std::sync::Arc; +use std::sync::mpsc::sync_channel; use std::time::Duration; use std::time::Instant; use anyhow::Result; use clap::ValueEnum; use indicatif::ProgressBar; +use parking_lot::Mutex; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; @@ -16,6 +20,7 @@ use rand_distr::Exp; use vortex_bench::Engine; use vortex_bench::Format; use vortex_bench::Target; +use vortex_bench::VortexSession; use vortex_bench::create_output_writer; use vortex_bench::display::DisplayFormat; use vortex_bench::display::print_measurements_json; @@ -24,14 +29,108 @@ use vortex_bench::random_access::BenchDataset; use vortex_bench::random_access::ParquetRandomAccessor; use vortex_bench::random_access::RandomAccessor; use vortex_bench::random_access::VortexRandomAccessor; +use vortex_bench::session_with_handle; use vortex_bench::utils::constants::STORAGE_NVME; use vortex_bench::v3; +use vortex_compio::CompioFileReadAt; +use vortex_compio::CompioHandle; +use vortex_compio::CompioRuntime; +use vortex_io::runtime::BlockingRuntime; use crate::render::RandomAccessRun; use crate::render::render_random_access_table; mod render; +const IO_RUNTIME_ENV: &str = "VORTEX_IO_RUNTIME"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum IoRuntime { + Tokio, + Compio, +} + +impl IoRuntime { + fn from_env() -> Result { + match std::env::var(IO_RUNTIME_ENV) { + Ok(value) => Self::from_str(&value), + Err(std::env::VarError::NotPresent) => Ok(Self::Tokio), + Err(error) => Err(error.into()), + } + } + + fn name(self) -> &'static str { + match self { + Self::Tokio => "tokio", + Self::Compio => "compio", + } + } + + fn for_format(self, format: Format) -> Self { + if matches!(format, Format::OnDiskVortex | Format::VortexCompact) { + self + } else { + Self::Tokio + } + } +} + +impl FromStr for IoRuntime { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + match value { + "tokio" => Ok(Self::Tokio), + "compio" => Ok(Self::Compio), + _ => anyhow::bail!( + "unsupported {IO_RUNTIME_ENV} value {value:?}; expected `tokio` or `compio`" + ), + } + } +} + +#[derive(Clone)] +struct CompioHandles { + io: CompioHandle, + session: VortexSession, +} + +static COMPIO_HANDLES: Mutex> = Mutex::new(None); + +fn compio_handles() -> Result { + let mut handles = COMPIO_HANDLES.lock(); + if let Some(handles) = handles.as_ref() { + return Ok(handles.clone()); + } + + let (handle_send, handle_recv) = sync_channel(1); + std::thread::Builder::new() + .name("vortex-compio-bench".to_string()) + .spawn(move || { + let runtime = match CompioRuntime::new() { + Ok(runtime) => runtime, + Err(error) => { + drop(handle_send.send(Err(error))); + return; + } + }; + let handles = CompioHandles { + io: runtime.compio_handle(), + session: session_with_handle(runtime.handle()), + }; + if handle_send.send(Ok(handles)).is_err() { + return; + } + runtime.block_on(std::future::pending::<()>()); + })?; + + let new_handles = handle_recv + .recv() + .map_err(|error| anyhow::anyhow!("Compio benchmark runtime failed to start: {error}"))??; + *handles = Some(new_handles.clone()); + Ok(new_handles) +} + // --------------------------------------------------------------------------- // Access patterns // --------------------------------------------------------------------------- @@ -145,11 +244,14 @@ async fn benchmark_random_access( time_limit_secs: u64, storage: &str, reopen: bool, + io_runtime: IoRuntime, ) -> Result { let time_limit = Duration::from_secs(time_limit_secs); - let overall_start = Instant::now(); let mut runs = Vec::new(); - let mut accessor = open_accessor(dataset, format).await?; + let mut accessor = open_accessor(dataset, format, io_runtime).await?; + // Opening may generate or download the dataset on its first use. Keep that one-time setup out + // of the measurement budget so cold and warm invocations collect comparable sample counts. + let overall_start = Instant::now(); loop { let start = Instant::now(); @@ -162,7 +264,7 @@ async fn benchmark_random_access( } if reopen { - accessor = open_accessor(dataset, format).await?; + accessor = open_accessor(dataset, format, io_runtime).await?; } } @@ -193,20 +295,26 @@ fn display_name(dataset: &str, pattern: Option) -> String { /// Build a measurement name for a benchmark run. /// -/// For taxi (legacy), the name is `random-access/{format}-tokio-local-disk` to preserve -/// historical continuity with existing benchmark data. -/// For other datasets, includes dataset and pattern: -/// `random-access/{dataset}/{pattern}/{format}-tokio-local-disk`. -fn measurement_name(dataset: &str, pattern: Option, format: Format) -> String { +/// For taxi (legacy), the name is `random-access/{format}-{runtime}-local-disk`. +/// For other datasets, it includes the dataset and pattern: +/// `random-access/{dataset}/{pattern}/{format}-{runtime}-local-disk`. +fn measurement_name( + dataset: &str, + pattern: Option, + format: Format, + io_runtime: IoRuntime, +) -> String { + let io_runtime = io_runtime.for_format(format); let fmt = format.ext(); + let runtime = io_runtime.name(); match pattern { Some(p) => format!( - "random-access/{}/{}/{}-tokio-local-disk", + "random-access/{}/{}/{}-{runtime}-local-disk", dataset, p.name(), fmt ), - None => format!("random-access/{}-tokio-local-disk", fmt), + None => format!("random-access/{fmt}-{runtime}-local-disk"), } } @@ -233,18 +341,33 @@ fn push_v3_random_access_record(records: &mut Vec, run: &RandomAcc async fn open_accessor( dataset: &dyn BenchDataset, format: Format, + io_runtime: IoRuntime, ) -> Result> { + let io_runtime = io_runtime.for_format(format); let name = format!( - "random-access/{}/{}-tokio-local-disk", + "random-access/{}/{}-{}-local-disk", dataset.name(), - format.ext() + format.ext(), + io_runtime.name() ); match format { Format::OnDiskVortex | Format::VortexCompact => { let path = dataset.path(format).await?; - Ok(Box::new( - VortexRandomAccessor::open(path, name, format).await?, - )) + let accessor = match io_runtime { + IoRuntime::Tokio => VortexRandomAccessor::open(path, name, format).await?, + IoRuntime::Compio => { + let handles = compio_handles()?; + let reader = CompioFileReadAt::open(path, handles.io).await?; + VortexRandomAccessor::open_source_with_session( + handles.session, + Arc::new(reader), + name, + format, + ) + .await? + } + }; + Ok(Box::new(accessor)) } Format::Parquet => { let path = dataset.path(format).await?; @@ -309,6 +432,7 @@ pub async fn run(config: RunConfig) -> Result<()> { output_path, ingest_output, } = config; + let io_runtime = IoRuntime::from_env()?; let reopen_variants: &[bool] = match open_mode { OpenMode::Cached => &[false], @@ -333,7 +457,7 @@ pub async fn run(config: RunConfig) -> Result<()> { for dataset in &datasets { for format in &formats { if dataset.name() == "taxi" { - let name = measurement_name(dataset.name(), None, *format); + let name = measurement_name(dataset.name(), None, *format, io_runtime); for &reopen in reopen_variants { let bench_name = if reopen { format!("{name}-footer") @@ -349,6 +473,7 @@ pub async fn run(config: RunConfig) -> Result<()> { time_limit, STORAGE_NVME, reopen, + io_runtime, ) .await?; @@ -360,7 +485,7 @@ pub async fn run(config: RunConfig) -> Result<()> { for pattern in &patterns { let indices = generate_indices(dataset.as_ref(), *pattern); - let name = measurement_name(dataset.name(), Some(*pattern), *format); + let name = measurement_name(dataset.name(), Some(*pattern), *format, io_runtime); for &reopen in reopen_variants { let bench_name = if reopen { format!("{name}-footer") @@ -376,6 +501,7 @@ pub async fn run(config: RunConfig) -> Result<()> { time_limit, STORAGE_NVME, reopen, + io_runtime, ) .await?; @@ -412,6 +538,32 @@ pub async fn run(config: RunConfig) -> Result<()> { mod tests { use super::*; + #[test] + fn parses_io_runtime_values() -> Result<()> { + assert_eq!(IoRuntime::from_str("tokio")?, IoRuntime::Tokio); + assert_eq!(IoRuntime::from_str("compio")?, IoRuntime::Compio); + assert!(IoRuntime::from_str("unknown").is_err()); + Ok(()) + } + + #[test] + fn io_runtime_only_changes_vortex_measurement_names() { + let vortex = measurement_name( + "feature-vectors", + Some(AccessPattern::Uniform), + Format::OnDiskVortex, + IoRuntime::Compio, + ); + let parquet = measurement_name( + "feature-vectors", + Some(AccessPattern::Uniform), + Format::Parquet, + IoRuntime::Compio, + ); + assert!(vortex.ends_with("vortex-compio-local-disk")); + assert!(parquet.ends_with("parquet-tokio-local-disk")); + } + #[test] fn v3_random_access_dataset_names_match_schema_dims() { assert_eq!(v3_random_access_dataset_name("taxi", None), "taxi"); diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index bc077a64719..c426be01b5e 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -71,8 +71,9 @@ pub use output::BenchmarkOutput; pub use output::create_output_writer; use vortex::VortexSessionDefault; pub use vortex::error::vortex_panic; +use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; -use vortex::session::VortexSession; +pub use vortex::session::VortexSession; // All benchmarks run with mimalloc for consistency. #[global_allocator] @@ -84,6 +85,13 @@ pub static SESSION: LazyLock = LazyLock::new(|| { session }); +/// Create a benchmark session that schedules work on `handle`. +pub fn session_with_handle(handle: Handle) -> VortexSession { + let session = VortexSession::default().with_handle(handle); + vortex_geo::initialize(&session); + session +} + #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] pub struct Target { pub engine: Engine, diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index fa941d07362..969a200252d 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -25,7 +25,9 @@ use vortex::array::stream::ArrayStreamExt; use vortex::buffer::Buffer; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; +use vortex::io::VortexReadAt; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; use crate::Format; @@ -40,6 +42,7 @@ pub struct VortexRandomAccessor { name: String, format: Format, file: VortexFile, + session: VortexSession, } impl VortexRandomAccessor { @@ -49,7 +52,8 @@ impl VortexRandomAccessor { name: impl Into, format: Format, ) -> anyhow::Result { - let file = SESSION + let session = SESSION.clone(); + let file = session .open_options() .with_layout_reader_cache() .open_path(path.as_ref()) @@ -58,6 +62,36 @@ impl VortexRandomAccessor { name: name.into(), format, file, + session, + }) + } + + /// Open a Vortex file using a caller-provided random-access I/O implementation. + pub async fn open_source( + source: Arc, + name: impl Into, + format: Format, + ) -> anyhow::Result { + Self::open_source_with_session(SESSION.clone(), source, name, format).await + } + + /// Open a Vortex file using a caller-provided session and random-access I/O implementation. + pub async fn open_source_with_session( + session: VortexSession, + source: Arc, + name: impl Into, + format: Format, + ) -> anyhow::Result { + let file = session + .open_options() + .with_layout_reader_cache() + .open(source) + .await?; + Ok(Self { + name: name.into(), + format, + file, + session, }) } } @@ -83,7 +117,7 @@ impl RandomAccessor for VortexRandomAccessor { .await?; // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es. - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = self.session.create_execution_ctx(); let canonical = array.execute::(&mut ctx)?.into_array(); Ok(RandomAccessorRet::ArrayRef(canonical)) } diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ed9df797535..87c9ff7e66e 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -18,6 +18,7 @@ all-features = true [features] arrow = [] +compio = ["dep:compio-buf"] memmap2 = ["dep:memmap2"] serde = ["dep:serde", "serde/serde_derive"] warn-copy = ["dep:tracing"] @@ -26,6 +27,7 @@ warn-copy = ["dep:tracing"] arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } +compio-buf = { workspace = true, optional = true } itertools = { workspace = true } memmap2 = { workspace = true, optional = true } serde = { workspace = true, optional = true } diff --git a/vortex-buffer/src/compio.rs b/vortex-buffer/src/compio.rs new file mode 100644 index 00000000000..dbb134f2921 --- /dev/null +++ b/vortex-buffer/src/compio.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::MaybeUninit; + +use compio_buf::IoBuf; +use compio_buf::IoBufMut; +use compio_buf::SetLen; + +use crate::ByteBuffer; +use crate::ByteBufferMut; + +impl IoBuf for ByteBuffer { + fn as_init(&self) -> &[u8] { + self.as_slice() + } +} + +impl IoBuf for ByteBufferMut { + fn as_init(&self) -> &[u8] { + self.as_slice() + } +} + +impl IoBufMut for ByteBufferMut { + fn as_uninit(&mut self) -> &mut [MaybeUninit] { + let ptr = self.bytes.as_mut_ptr().cast::>(); + let capacity = self.capacity(); + + // SAFETY: `BytesMut` guarantees its pointer is valid for its capacity. `BufferMut` has + // the same byte capacity, and `MaybeUninit` permits both its initialized prefix and its + // uninitialized spare capacity. + unsafe { std::slice::from_raw_parts_mut(ptr, capacity) } + } +} + +impl SetLen for ByteBufferMut { + unsafe fn set_len(&mut self, len: usize) { + // SAFETY: `SetLen` has the same requirements as `BufferMut::set_len`: `len` is within the + // allocation and the newly exposed byte range has been initialized by the I/O operation. + unsafe { ByteBufferMut::set_len(self, len) } + } +} + +#[cfg(test)] +mod tests { + use compio_buf::IoBuf; + use compio_buf::IoBufMut; + use compio_buf::SetLen; + + use crate::Alignment; + use crate::ByteBufferMut; + + #[test] + fn byte_buffer_mut_supports_completion_io() { + let alignment = Alignment::new(4096); + let mut buffer = ByteBufferMut::with_capacity_aligned(4, alignment); + + assert!(buffer.as_init().is_empty()); + assert!(buffer.buf_capacity() >= 4); + + buffer.as_uninit()[..4].copy_from_slice(&[ + std::mem::MaybeUninit::new(1), + std::mem::MaybeUninit::new(2), + std::mem::MaybeUninit::new(3), + std::mem::MaybeUninit::new(4), + ]); + // SAFETY: the first four bytes were initialized immediately above. + unsafe { SetLen::set_len(&mut buffer, 4) }; + + assert_eq!(buffer.as_init(), &[1, 2, 3, 4]); + assert!(alignment.is_ptr_aligned(buffer.as_init().as_ptr())); + + let sealed = buffer.freeze(); + assert_eq!(sealed.as_init(), &[1, 2, 3, 4]); + assert_eq!(sealed.alignment(), alignment); + } +} diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..b9ada211eee 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -45,6 +45,9 @@ //! The `arrow` feature can be enabled to provide conversion functions to/from Arrow Rust buffers, //! including `arrow_buffer::Buffer`, `arrow_buffer::ScalarBuffer`, and //! `arrow_buffer::OffsetBuffer`. +//! +//! The `compio` feature implements completion-based I/O buffer traits from `compio-buf` for +//! [`ByteBuffer`] and [`ByteBufferMut`]. pub use alignment::*; pub use bit::*; @@ -61,6 +64,8 @@ mod bit; mod buffer; mod buffer_mut; mod bytes; +#[cfg(feature = "compio")] +mod compio; mod r#const; mod debug; mod dispatch; diff --git a/vortex-compio/Cargo.toml b/vortex-compio/Cargo.toml new file mode 100644 index 00000000000..0546bd60b28 --- /dev/null +++ b/vortex-compio/Cargo.toml @@ -0,0 +1,42 @@ +[package] +name = "vortex-compio" +authors = { workspace = true } +categories = { workspace = true } +description = "Compio runtime and filesystem I/O integration for Vortex" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[package.metadata.docs.rs] +all-features = true +targets = [ + "aarch64-apple-darwin", + "x86_64-unknown-linux-gnu", + "x86_64-pc-windows-msvc", +] + +[dependencies] +compio = { workspace = true, features = ["fs", "io-uring", "polling", "sync"] } +futures = { workspace = true, features = ["std"] } +kanal = { workspace = true } +oneshot = { workspace = true } +parking_lot = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true, features = ["compio"] } +vortex-error = { workspace = true } +vortex-io = { workspace = true } + +[dev-dependencies] +tempfile = { workspace = true } +vortex-file = { workspace = true } +vortex-layout = { workspace = true } +vortex-metrics = { workspace = true } + +[lints] +workspace = true diff --git a/vortex-compio/src/lib.rs b/vortex-compio/src/lib.rs new file mode 100644 index 00000000000..1a5b66264f5 --- /dev/null +++ b/vortex-compio/src/lib.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compio runtime and filesystem I/O integration for Vortex. +//! +//! [`CompioRuntime`] drives Vortex work on Compio's thread-local completion-based runtime. Create +//! one runtime per worker thread to use a thread-per-core execution model. [`CompioFileReadAt`] +//! performs positioned reads directly into aligned Vortex buffers, using io_uring on Linux when +//! available and Compio's polling driver as a fallback. + +mod read_at; +mod runtime; + +pub use read_at::*; +pub use runtime::*; diff --git a/vortex-compio/src/read_at.rs b/vortex-compio/src/read_at.rs new file mode 100644 index 00000000000..49c23ea08f5 --- /dev/null +++ b/vortex-compio/src/read_at.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::path::Path; +use std::sync::Arc; + +use compio::buf::BufResult; +use compio::buf::IntoInner; +use compio::buf::IoBuf; +use compio::io::AsyncReadAtExt; +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_io::CoalesceConfig; +use vortex_io::VortexReadAt; + +/// Default number of concurrent reads issued for a local Compio file. +pub const DEFAULT_CONCURRENCY: usize = 32; + +/// A completion-based positioned reader for a local file. +/// +/// Reads are issued through the active Compio runtime directly into aligned Vortex buffers. On +/// Linux, Compio uses io_uring by default. Files are opened without `O_DIRECT`, so reads retain the +/// operating system page cache. +#[derive(Clone)] +pub struct CompioFileReadAt { + uri: Arc, + file: Arc, + size: u64, + handle: crate::CompioHandle, +} + +impl CompioFileReadAt { + /// Open a local file for completion-based positioned reads. + /// + /// The associated [`crate::CompioRuntime`] must be driven while this future and subsequent read + /// futures are pending. + pub async fn open(path: impl AsRef, handle: crate::CompioHandle) -> VortexResult { + let path = path.as_ref().to_path_buf(); + let uri: Arc = path.to_string_lossy().into(); + let (file, size) = handle + .spawn_local(move || async move { + let file = compio::fs::File::open(path).await?; + let size = file.metadata().await?.len(); + std::io::Result::Ok((file, size)) + }) + .await?; + Ok(Self { + uri, + file: Arc::new(file), + size, + handle, + }) + } +} + +impl VortexReadAt for CompioFileReadAt { + fn uri(&self) -> Option<&Arc> { + Some(&self.uri) + } + + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::file()) + } + + fn concurrency(&self) -> usize { + DEFAULT_CONCURRENCY + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let size = self.size; + async move { Ok(size) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let file = Arc::clone(&self.file); + self.handle + .spawn_local(move || async move { + let buffer = ByteBufferMut::with_capacity_aligned(length, alignment); + // `BufferMut` may over-allocate to achieve alignment. Restrict Compio's owned view to + // the requested length so `read_exact_at` does not fill the spare capacity too. + let buffer = buffer.slice(..length); + let BufResult(result, buffer) = file.read_exact_at(buffer, offset).await; + result?; + + let buffer = buffer.into_inner(); + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .boxed() + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use futures::future::join_all; + use tempfile::NamedTempFile; + use vortex_buffer::Alignment; + use vortex_error::VortexError; + use vortex_error::VortexResult; + use vortex_io::VortexReadAt; + use vortex_io::runtime::BlockingRuntime; + + use crate::CompioFileReadAt; + use crate::CompioRuntime; + + const DATA: &[u8] = b"completion-based Vortex reads"; + + #[test] + fn reads_exact_aligned_ranges() -> VortexResult<()> { + let mut temp = NamedTempFile::new()?; + temp.write_all(DATA)?; + temp.flush()?; + + let runtime = CompioRuntime::new()?; + let handle = runtime.compio_handle(); + runtime.block_on(async move { + let reader = CompioFileReadAt::open(temp.path(), handle).await?; + assert_eq!(reader.size().await?, DATA.len() as u64); + + let alignment = Alignment::new(4096); + let buffer = reader.read_at(17, 6, alignment).await?.unwrap_host(); + assert_eq!(buffer.as_slice(), b"Vortex"); + assert_eq!(buffer.len(), 6); + assert!(alignment.is_ptr_aligned(buffer.as_ptr())); + VortexResult::Ok(()) + }) + } + + #[test] + fn supports_concurrent_positioned_reads() -> VortexResult<()> { + let mut temp = NamedTempFile::new()?; + temp.write_all(DATA)?; + temp.flush()?; + + let runtime = CompioRuntime::new()?; + let handle = runtime.compio_handle(); + runtime.block_on(async move { + let reader = CompioFileReadAt::open(temp.path(), handle).await?; + let reads = [ + reader.read_at(0, 10, Alignment::none()), + reader.read_at(17, 6, Alignment::none()), + reader.read_at(24, 5, Alignment::none()), + ]; + let buffers = join_all(reads) + .await + .into_iter() + .collect::>>()?; + + assert_eq!(buffers[0].as_host().as_slice(), b"completion"); + assert_eq!(buffers[1].as_host().as_slice(), b"Vortex"); + assert_eq!(buffers[2].as_host().as_slice(), b"reads"); + VortexResult::Ok(()) + }) + } + + #[test] + fn reports_unexpected_eof() -> VortexResult<()> { + let mut temp = NamedTempFile::new()?; + temp.write_all(DATA)?; + temp.flush()?; + + let runtime = CompioRuntime::new()?; + let handle = runtime.compio_handle(); + runtime.block_on(async move { + let reader = CompioFileReadAt::open(temp.path(), handle).await?; + let error = reader + .read_at(DATA.len() as u64 - 2, 8, Alignment::none()) + .await + .expect_err("short positioned read must fail"); + assert!(matches!( + &error, + VortexError::Io(error, _) if error.kind() == std::io::ErrorKind::UnexpectedEof + )); + VortexResult::Ok(()) + }) + } +} diff --git a/vortex-compio/src/runtime.rs b/vortex-compio/src/runtime.rs new file mode 100644 index 00000000000..e88128d4642 --- /dev/null +++ b/vortex-compio/src/runtime.rs @@ -0,0 +1,429 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::io; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; + +use compio::runtime::JoinHandle; +use futures::FutureExt; +use futures::Stream; +use futures::StreamExt; +use futures::future::Abortable; +use futures::future::BoxFuture; +use futures::future::LocalBoxFuture; +use futures::stream::BoxStream; +use parking_lot::Mutex; +use vortex_error::vortex_panic; +use vortex_io::runtime::AbortHandle; +use vortex_io::runtime::AbortHandleRef; +use vortex_io::runtime::BlockingRuntime; +use vortex_io::runtime::Executor; +use vortex_io::runtime::Handle; + +/// A Vortex runtime backed by a thread-local Compio runtime. +/// +/// The runtime does no work unless [`BlockingRuntime::block_on`] or an iterator returned by +/// [`BlockingRuntime::block_on_stream`] is being driven. For thread-per-core execution, construct +/// one `CompioRuntime` on each worker thread. +pub struct CompioRuntime { + sender: Arc, + runtime: compio::runtime::Runtime, +} + +impl CompioRuntime { + /// Create a Compio-backed Vortex runtime using Compio's default platform driver. + pub fn new() -> io::Result { + let runtime = compio::runtime::Runtime::new()?; + let sender = Arc::new(Sender::new(&runtime)); + Ok(Self { sender, runtime }) + } + + /// Return the Compio driver selected for this runtime. + pub fn driver_type(&self) -> compio::driver::DriverType { + self.runtime.driver_type() + } + + /// Return a handle for scheduling operations that must remain local to this Compio runtime. + pub fn compio_handle(&self) -> CompioHandle { + CompioHandle { + sender: Arc::clone(&self.sender), + } + } +} + +/// A thread-safe handle for scheduling local operations on a [`CompioRuntime`]. +/// +/// Unlike a Vortex [`Handle`], this handle accepts a factory rather than an already-created future. +/// This ensures Compio's `!Send` I/O futures are created and polled only on their owning runtime. +#[derive(Clone)] +pub struct CompioHandle { + sender: Arc, +} + +impl CompioHandle { + pub(crate) fn spawn_local(&self, make_future: F) -> LocalTask + where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future + 'static, + R: Send + 'static, + { + self.sender.spawn_local(make_future) + } +} + +impl BlockingRuntime for CompioRuntime { + type BlockingIterator<'a, R: 'a> = CompioBlockingIterator<'a, R>; + + fn handle(&self) -> Handle { + let executor: Arc = Arc::clone(&self.sender) as Arc; + Handle::new(Arc::downgrade(&executor)) + } + + fn block_on(&self, future: Fut) -> R + where + Fut: Future, + { + self.runtime.block_on(future) + } + + fn block_on_stream<'a, S, R>(&self, stream: S) -> Self::BlockingIterator<'a, R> + where + S: Stream + Send + 'a, + R: Send + 'a, + { + CompioBlockingIterator { + runtime: self.runtime.clone(), + stream: stream.boxed(), + } + } +} + +/// Run a future to completion on a new [`CompioRuntime`]. +/// +/// The closure receives a Vortex [`Handle`] associated with the runtime. +pub fn block_on(f: F) -> io::Result +where + F: FnOnce(Handle) -> Fut, + Fut: Future, +{ + let runtime = CompioRuntime::new()?; + let handle = runtime.handle(); + Ok(runtime.block_on(f(handle))) +} + +/// An iterator that drives a stream with a Compio runtime on the current thread. +pub struct CompioBlockingIterator<'a, T> { + runtime: compio::runtime::Runtime, + stream: BoxStream<'a, T>, +} + +impl Iterator for CompioBlockingIterator<'_, T> { + type Item = T; + + fn next(&mut self) -> Option { + self.runtime.block_on(self.stream.next()) + } +} + +struct Sender { + tasks: kanal::Sender, + runtime_id: usize, +} + +impl Sender { + fn new(runtime: &compio::runtime::Runtime) -> Self { + let (send, recv) = kanal::unbounded::(); + let runtime_id = runtime_identity(runtime); + + runtime + .spawn(async move { + while let Ok(spawn) = recv.as_async().recv().await { + spawn.schedule(); + } + }) + .detach(); + + Self { + tasks: send, + runtime_id, + } + } + + fn is_current_runtime(&self) -> bool { + compio::runtime::Runtime::try_with_current(|runtime| { + runtime_identity(runtime) == self.runtime_id + }) + .unwrap_or(false) + } + + fn schedule(&self, spawn: Spawn) { + // Vortex's coalescing driver runs on this runtime, so this is the hot path for local file + // reads. Bypass the synchronized cross-thread queue when scheduling from the owner. + if self.is_current_runtime() { + spawn.schedule(); + } else if let Err(error) = self.tasks.send(spawn) { + vortex_panic!("Compio executor missing: {error}"); + } + } + + fn send(&self, spawn: Spawn) -> AbortHandleRef { + let (task_send, task_recv) = oneshot::channel(); + let spawn = spawn.with_callback(task_send); + self.schedule(spawn); + Box::new(LazyAbortHandle { + task: Mutex::new(task_recv), + }) + } + + fn spawn_local(&self, make_future: F) -> LocalTask + where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future + 'static, + R: Send + 'static, + { + let (result_send, result_recv) = oneshot::channel(); + let (abort_handle, abort_registration) = futures::future::AbortHandle::new_pair(); + let factory = Box::new(move || { + async move { + let future = async move { make_future().await }; + let output = + Abortable::new(AssertUnwindSafe(future).catch_unwind(), abort_registration) + .await; + if let Ok(output) = output { + // The receiver may have been dropped immediately after the operation completed. + drop(result_send.send(output)); + } + } + .boxed_local() + }); + + self.schedule(Spawn::Local { factory }); + + LocalTask { + result: result_recv.into_future(), + abort_handle: Some(abort_handle), + } + } +} + +fn runtime_identity(runtime: &compio::runtime::Runtime) -> usize { + std::ptr::from_ref(&**runtime).addr() +} + +impl Executor for Sender { + fn spawn(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef { + self.send(Spawn::Future { + future, + callback: None, + }) + } + + fn spawn_io(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef { + self.send(Spawn::Future { + future, + callback: None, + }) + } + + fn spawn_cpu(&self, task: Box) -> AbortHandleRef { + self.send(Spawn::Cpu { + task, + callback: None, + }) + } + + fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { + self.send(Spawn::Blocking { + task, + callback: None, + }) + } +} + +enum Spawn { + Future { + future: BoxFuture<'static, ()>, + callback: Option>, + }, + Cpu { + task: Box, + callback: Option>, + }, + Blocking { + task: Box, + callback: Option>, + }, + Local { + factory: LocalFutureFactory, + }, +} + +type LocalFutureFactory = Box LocalBoxFuture<'static, ()> + Send + 'static>; + +impl Spawn { + fn with_callback(self, callback: oneshot::Sender) -> Self { + match self { + Spawn::Future { future, .. } => Spawn::Future { + future, + callback: Some(callback), + }, + Spawn::Cpu { task, .. } => Spawn::Cpu { + task, + callback: Some(callback), + }, + Spawn::Blocking { task, .. } => Spawn::Blocking { + task, + callback: Some(callback), + }, + Spawn::Local { .. } => { + vortex_panic!("local Compio tasks manage cancellation directly") + } + } + } + + fn schedule(self) { + if let Spawn::Local { factory } = self { + compio::runtime::spawn(factory()).detach(); + return; + } + + let (task, callback) = match self { + Spawn::Future { future, callback } => (compio::runtime::spawn(future), callback), + Spawn::Cpu { task, callback } => { + (compio::runtime::spawn(async move { task() }), callback) + } + Spawn::Blocking { task, callback } => (compio::runtime::spawn_blocking(task), callback), + Spawn::Local { .. } => unreachable!(), + }; + + let abort_handle: AbortHandleRef = Box::new(CompioAbortHandle { task: Some(task) }); + if let Some(callback) = callback { + // A failed send means the caller dropped or aborted its task before it was scheduled. + drop(callback.send(abort_handle)); + } + } +} + +type LocalTaskOutput = Result>; + +pub(crate) struct LocalTask { + result: oneshot::AsyncReceiver>, + abort_handle: Option, +} + +impl Future for LocalTask { + type Output = T; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + match ready!(this.result.poll_unpin(cx)) { + Ok(Ok(output)) => { + this.abort_handle.take(); + Poll::Ready(output) + } + Ok(Err(panic)) => { + this.abort_handle.take(); + std::panic::resume_unwind(panic) + } + Err(error) => vortex_panic!("Compio local task was cancelled: {error}"), + } + } +} + +impl Drop for LocalTask { + fn drop(&mut self) { + if let Some(abort_handle) = self.abort_handle.take() { + abort_handle.abort(); + } + } +} + +struct CompioAbortHandle { + task: Option>, +} + +impl AbortHandle for CompioAbortHandle { + fn abort(mut self: Box) { + // Dropping a Compio join handle cancels its task. + drop(self.task.take()); + } +} + +impl Drop for CompioAbortHandle { + fn drop(&mut self) { + if let Some(task) = self.task.take() { + task.detach(); + } + } +} + +struct LazyAbortHandle { + task: Mutex>, +} + +impl AbortHandle for LazyAbortHandle { + fn abort(self: Box) { + if let Ok(task) = self.task.lock().try_recv() { + task.abort(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::stream; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + + use super::CompioRuntime; + + #[test] + fn drives_spawned_work() -> VortexResult<()> { + let runtime = CompioRuntime::new()?; + let handle = runtime.handle(); + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = Arc::clone(&counter); + + let result = runtime.block_on(async move { + let task = handle.spawn_io(async move { + counter_clone.fetch_add(1, Ordering::SeqCst); + 42 + }); + task.await + }); + + assert_eq!(result, 42); + assert_eq!(counter.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[test] + fn drives_streams() -> VortexResult<()> { + let runtime = CompioRuntime::new()?; + let values = runtime + .block_on_stream(stream::iter([1, 2, 3])) + .collect::>(); + assert_eq!(values, [1, 2, 3]); + Ok(()) + } + + #[test] + fn identifies_the_owning_runtime() -> VortexResult<()> { + let runtime = CompioRuntime::new()?; + assert!(!runtime.sender.is_current_runtime()); + runtime.block_on(async { + assert!(runtime.sender.is_current_runtime()); + }); + Ok(()) + } +} diff --git a/vortex-compio/tests/file_segment_source.rs b/vortex-compio/tests/file_segment_source.rs new file mode 100644 index 00000000000..3f56eb92a1a --- /dev/null +++ b/vortex-compio/tests/file_segment_source.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(test)] +mod tests { + use std::io::Write; + use std::sync::Arc; + + use futures::future::try_join; + use tempfile::NamedTempFile; + use vortex_buffer::Alignment; + use vortex_compio::CompioFileReadAt; + use vortex_compio::CompioRuntime; + use vortex_error::VortexResult; + use vortex_file::SegmentSpec; + use vortex_file::segments::FileSegmentSource; + use vortex_file::segments::RequestMetrics; + use vortex_io::runtime::BlockingRuntime; + use vortex_layout::segments::SegmentId; + use vortex_layout::segments::SegmentSource; + use vortex_metrics::DefaultMetricsRegistry; + + const DATA: &[u8] = b"completion-based Vortex reads"; + + #[test] + fn coalescing_driver_reads_with_compio() -> VortexResult<()> { + let mut temp = NamedTempFile::new()?; + temp.write_all(DATA)?; + temp.flush()?; + + let runtime = CompioRuntime::new()?; + let compio_handle = runtime.compio_handle(); + let handle = runtime.handle(); + runtime.block_on(async move { + let reader = CompioFileReadAt::open(temp.path(), compio_handle).await?; + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + Arc::from([ + SegmentSpec { + offset: 0, + length: 10, + alignment: Alignment::none(), + }, + SegmentSpec { + offset: 17, + length: 6, + alignment: Alignment::new(4096), + }, + ]), + reader, + handle, + RequestMetrics::new(&metrics, vec![]), + ); + + let (first, second) = try_join( + source.request(SegmentId::from(0)), + source.request(SegmentId::from(1)), + ) + .await?; + assert_eq!(first.as_host().as_slice(), b"completion"); + assert_eq!(second.as_host().as_slice(), b"Vortex"); + assert!(Alignment::new(4096).is_ptr_aligned(second.as_host().as_ptr())); + VortexResult::Ok(()) + }) + } +} diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 3af33362b05..3ce3d2aed90 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -168,7 +168,7 @@ impl FileSegmentSource { // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers // join it (below) only to surface a panic raised while driving reads. - let mut task = handle.spawn(drive_fut); + let mut task = handle.spawn_io(drive_fut); let driver_panic: DriverPanic = Arc::new(Mutex::new(None)); let driver = { let driver_panic = Arc::clone(&driver_panic); @@ -381,14 +381,89 @@ impl SegmentSource for BufferSegmentSource { #[cfg(test)] mod tests { use std::panic::AssertUnwindSafe; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use futures::future::BoxFuture; + use vortex_io::runtime::AbortHandle; + use vortex_io::runtime::AbortHandleRef; + use vortex_io::runtime::Executor; use vortex_io::runtime::tokio::TokioRuntime; use vortex_layout::segments::SegmentSource; use vortex_metrics::DefaultMetricsRegistry; use super::*; + #[derive(Default)] + struct CountingExecutor { + spawn_io_count: AtomicUsize, + } + + impl Executor for CountingExecutor { + fn spawn(&self, _future: BoxFuture<'static, ()>) -> AbortHandleRef { + panic!("read driver used spawn instead of spawn_io") + } + + fn spawn_io(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef { + self.spawn_io_count.fetch_add(1, Ordering::SeqCst); + drop(future); + Box::new(NoOpAbortHandle) + } + + fn spawn_cpu(&self, _task: Box) -> AbortHandleRef { + unreachable!("read driver does not spawn CPU work") + } + + fn spawn_blocking_io(&self, _task: Box) -> AbortHandleRef { + unreachable!("read driver does not spawn blocking work") + } + } + + struct NoOpAbortHandle; + + impl AbortHandle for NoOpAbortHandle { + fn abort(self: Box) {} + } + + #[derive(Clone)] + struct EmptyReadAt; + + impl VortexReadAt for EmptyReadAt { + fn concurrency(&self) -> usize { + 1 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(0) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { Ok(BufferHandle::new_host(ByteBuffer::empty())) }.boxed() + } + } + + #[test] + fn read_driver_uses_spawn_io() { + let executor = Arc::new(CountingExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + let metrics = DefaultMetricsRegistry::default(); + + let _source = FileSegmentSource::open( + Arc::from([]), + EmptyReadAt, + handle, + RequestMetrics::new(&metrics, vec![]), + ); + + assert_eq!(executor.spawn_io_count.load(Ordering::SeqCst), 1); + } + #[derive(Clone)] struct PanickingReadAt;