feat(localfile): Add support of io_uring engine (experimental) (#549)

Author: zuston

Cargo.lock

@@ -89,6 +89,7 @@ datafusion-postgres = { git = "https://github.com/zuston/datafusion-postgres.git
 sqlparser = "0.56.0"
 nix = { version = "0.29.0" }
 dirs = { version = "6.0.0" }
+io-uring = { version = "0.7.0" }
 
 [profile.release]
 strip = true

Cargo.toml

@@ -12,4 +12,4 @@ services:
       - ./../../../target-docker:/R1/target:rw
     #    environment:
     #      RUSTFLAGS: "-C target-cpu=skylake"
-    command: "bash -c 'source ~/.bashrc && cd /R1 && cargo build --features hdrs,logforth,memory-prof --release'"
+    command: "bash -c 'source ~/.bashrc && cd /R1 && cargo build --features hdrs,logforth,memory-prof,io-uring --release'"

dev/anolisos8/amd64/docker-compose.yml

@@ -59,6 +59,8 @@ logforth = ["dep:logforth"]
 
 deadlock-detection = ["parking_lot/deadlock_detection"]
 
+io-uring = ["dep:io-uring"]
+
 [dependencies]
 anyhow = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
@@ -134,6 +136,7 @@ tikv-jemalloc-sys= { workspace = true, features = ["stats", "profiling", "unpref
 tikv-jemallocator= { workspace = true, features = ["profiling", "unprefixed_malloc_on_supported_platforms"], optional = true }
 jemalloc_pprof= { workspace = true, features = ["symbolize", "flamegraph"], optional = true }
 nix= { workspace = true, optional = true }
+io-uring = { workspace = true, optional = true }
 
 [build-dependencies]
 tonic-build = { workspace = true }

riffle-server/Cargo.toml

@@ -139,6 +139,18 @@ pub struct LocalfileStoreConfig {
     pub read_io_sendfile_enable: bool,
 
     pub read_ahead_options: Option<ReadAheadConfig>,
+
+    pub io_uring_options: Option<IoUringConfig>,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+pub struct IoUringConfig {
+    #[serde(default = "as_default_io_uring_threads")]
+    pub threads: usize,
+}
+
+fn as_default_io_uring_threads() -> usize {
+    1
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
@@ -253,6 +265,7 @@ impl LocalfileStoreConfig {
             io_limiter: None,
             read_io_sendfile_enable: as_default_read_io_sendfile_enable(),
             read_ahead_options: None,
+            io_uring_options: None,
         }
     }
 }

riffle-server/src/config.rs

@@ -111,6 +111,9 @@ pub enum WorkerError {
 
     #[error("HDFS client initialization failed")]
     HDFS_CLIENT_INIT_FAILED,
+
+    #[error("raw io error code: {0}")]
+    RAW_IO_ERR(i32),
 }
 
 impl From<AcquireError> for WorkerError {

riffle-server/src/error.rs

@@ -18,7 +18,6 @@
 #![allow(dead_code, unused)]
 #![feature(impl_trait_in_assoc_type)]
 extern crate core;
-
 pub mod app_manager;
 pub mod await_tree;
 pub mod client_configs;

riffle-server/src/lib.rs

@@ -215,7 +215,7 @@ mod tests {
     use std::thread;
     use std::time::Duration;
 
-    fn create_runtime(pool_size: usize, name: &str) -> Arc<Runtime> {
+    pub fn create_runtime(pool_size: usize, name: &str) -> Arc<Runtime> {
         let runtime = Builder::default()
             .worker_threads(pool_size as usize)
             .thread_name(name)

riffle-server/src/runtime/mod.rs

@@ -21,6 +21,8 @@ use crate::store::local::layers::{Handler, OperatorBuilder};
 use crate::store::local::options::WriteOptions;
 use crate::store::local::read_options::ReadOptions;
 use crate::store::local::sync_io::SyncLocalIO;
+#[cfg(feature = "io-uring")]
+use crate::store::local::uring_io::UringIoEngineBuilder;
 use crate::store::local::{DiskStat, FileStat, LocalDiskStorage, LocalIO};
 use crate::store::DataBytes;
 use crate::util;
@@ -95,7 +97,24 @@ impl LocalDiskDelegator {
             Some(read_capacity.as_u64() as usize),
         );
 
+        #[cfg(feature = "io-uring")]
+        info!("Binary compiled with the io-uring feature enabled.");
+
+        #[cfg(feature = "io-uring")]
+        let mut operator_builder = if let Some(cfg) = &config.io_uring_options {
+            info!("io-uring engine is activated!");
+            OperatorBuilder::new(Arc::new(Box::new(
+                UringIoEngineBuilder::new()
+                    .build(underlying_io_handler)
+                    .unwrap(),
+            )))
+        } else {
+            OperatorBuilder::new(Arc::new(Box::new(underlying_io_handler)))
+        };
+
+        #[cfg(not(feature = "io-uring"))]
         let mut operator_builder = OperatorBuilder::new(Arc::new(Box::new(underlying_io_handler)));
+
         if let Some(conf) = config.io_limiter.as_ref() {
             operator_builder = operator_builder.layer(ThrottleLayer::new(
                 &runtime_manager.localfile_write_runtime,

riffle-server/src/store/local/delegator.rs

@@ -36,6 +36,9 @@ pub mod options;
 pub mod read_options;
 pub mod sync_io;
 
+#[cfg(feature = "io-uring")]
+pub mod uring_io;
+
 pub struct FileStat {
     pub content_length: u64,
 }
@@ -79,3 +82,34 @@ impl Default for LocalfileStoreStat {
         Self { stats: vec![] }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use bytes::BytesMut;
+
+    #[test]
+    fn test_bytes_as_mut_ptr() {
+        struct RawBuf<T> {
+            ptr: T,
+            len: usize,
+        }
+
+        // 1. use bytesMut to write
+        let mut bytes = BytesMut::zeroed(10);
+        let ptr = bytes.as_mut_ptr();
+        let raw_buf = RawBuf { ptr, len: 10 };
+        unsafe {
+            for i in 0..bytes.len() {
+                *raw_buf.ptr.add(i) = 1;
+            }
+        }
+        for &b in bytes.iter() {
+            assert_eq!(b, 1);
+        }
+
+        // 2. use bytes to read
+        let mut bytes = bytes.freeze();
+        let prt = bytes.as_ptr();
+        let raw_buf = RawBuf { ptr, len: 10 };
+    }
+}