feat: Introduce io_uring for the urpc (experimental)

Cargo.toml

@@ -89,6 +89,15 @@ datafusion-postgres = { git = "https://github.com/zuston/datafusion-postgres.git
 sqlparser = "0.56.0"
 nix = { version = "0.29.0" }
 dirs = { version = "6.0.0" }
+monoio = { git = "https://github.com/bytedance/monoio.git", rev = "212fde187f3d41f65ecd35b9cac80f7e246bb58e", features = ["async-cancel",
+    "sync",
+    "bytes",
+    "iouring",
+    "legacy",
+    "macros",
+    "utils",
+"bytes",
+"debug",] }
 
 [profile.release]
 strip = true

dev/anolisos8/amd64/docker-compose.yml

@@ -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,urpc_uring --release'"

riffle-server/Cargo.toml

@@ -59,6 +59,8 @@ logforth = ["dep:logforth"]
 
 deadlock-detection = ["parking_lot/deadlock_detection"]
 
+urpc_uring = ["dep:monoio"]
+
 [dependencies]
 anyhow = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
@@ -134,6 +136,12 @@ 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 }
+monoio = { workspace = true, optional = true, features = [
+    "iouring",
+    "macros",
+    "utils",
+    "bytes",
+    "debug",] }
 
 [build-dependencies]
 tonic-build = { workspace = true }

riffle-server/src/config.rs

@@ -470,6 +470,8 @@ pub struct Config {
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 pub struct UrpcConfig {
     pub get_index_rpc_version: RpcVersion,
+    #[serde(default = "bool::default")]
+    pub io_uring_enabled: bool,
 }
 
 fn as_default_get_memory_rpc_version() -> RpcVersion {

riffle-server/src/lib.rs

@@ -66,3 +66,6 @@ pub mod mini_riffle;
 mod partition_stats;
 mod raw_io;
 mod system_libc;
+
+#[cfg(feature = "urpc_uring")]
+mod urpc_uring;

riffle-server/src/main.rs

@@ -17,7 +17,7 @@
 
 #![allow(dead_code, unused)]
 #![feature(impl_trait_in_assoc_type)]
-
+#[warn(unexpected_cfgs)]
 use crate::app_manager::{AppManager, APP_MANAGER_REF};
 use crate::common::init_global_variable;
 use crate::config::{Config, LogConfig};
@@ -111,6 +111,8 @@ pub mod ddashmap;
 
 pub mod partition_stats;
 
+#[cfg(feature = "urpc_uring")]
+pub mod urpc_uring;
 const MAX_MEMORY_ALLOCATION_SIZE_ENV_KEY: &str = "MAX_MEMORY_ALLOCATION_LIMIT_SIZE";
 
 #[derive(Parser, Debug)]

riffle-server/src/rpc.rs

@@ -26,6 +26,9 @@ use tokio::sync::broadcast::{Receiver, Sender};
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::transport::Server;
 
+#[cfg(feature = "urpc_uring")]
+use crate::urpc_uring;
+
 pub static GRPC_PARALLELISM: Lazy<NonZeroUsize> = Lazy::new(|| {
     let available_cores = std::thread::available_parallelism().unwrap();
     std::env::var("GRPC_PARALLELISM").map_or(available_cores, |v| {
@@ -78,14 +81,43 @@ impl DefaultRpcService {
             let app_manager = app_manager_ref.clone();
             let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), urpc_port as u16);
 
-            std::thread::spawn(move || {
-                core_affinity::set_for_current(core_id);
-                tokio::runtime::Builder::new_current_thread()
-                    .enable_all()
-                    .build()
-                    .unwrap()
-                    .block_on(urpc_serve(addr, shutdown(rx), app_manager));
-            });
+            let common_fn = || {
+                std::thread::spawn(move || {
+                    core_affinity::set_for_current(core_id);
+                    tokio::runtime::Builder::new_current_thread()
+                        .enable_all()
+                        .build()
+                        .unwrap()
+                        .block_on(urpc_serve(addr, shutdown(rx), app_manager));
+                });
+            };
+
+            #[cfg(feature = "urpc_uring")]
+            {
+                if config.urpc_config.is_some()
+                    && config.urpc_config.as_ref().unwrap().io_uring_enabled
+                {
+                    let rx2 = tx.subscribe();
+                    let app_manager_ref = app_manager_ref.clone();
+                    std::thread::spawn(move || {
+                        core_affinity::set_for_current(core_id);
+                        // Create monoio runtime and run urpc_serve inside it
+                        monoio::RuntimeBuilder::<monoio::FusionDriver>::new()
+                            .enable_all()
+                            .with_entries(32768)
+                            .build()
+                            .unwrap()
+                            .block_on(urpc_serve_with_monoio(addr, shutdown(rx2), app_manager_ref));
+                    });
+                } else {
+                    common_fn();
+                }
+            }
+
+            #[cfg(not(feature = "urpc_uring"))]
+            {
+                common_fn();
+            }
         }
 
         Ok(())
@@ -197,6 +229,16 @@ impl DefaultRpcService {
     }
 }
 
+#[cfg(feature = "urpc_uring")]
+async fn urpc_serve_with_monoio(
+    addr: SocketAddr,
+    shutdown: impl Future,
+    app_manager_ref: AppManagerRef,
+) {
+    let listner: monoio::net::TcpListener = monoio::net::TcpListener::bind(addr).unwrap();
+    let _ = urpc_uring::server::run(listner, shutdown, app_manager_ref).await;
+}
+
 async fn urpc_serve(addr: SocketAddr, shutdown: impl Future, app_manager_ref: AppManagerRef) {
     let sock = socket2::Socket::new(
         match addr {

riffle-server/src/store/hybrid.rs

@@ -283,24 +283,28 @@ impl HybridStore {
             _ => warm,
         };
 
-        let ctx = spill_message.ctx.clone();
-        // when throwing the data lost error, it should fast fail for this partition data.
-        let result = candidate_store
-            .spill_insert(ctx)
-            .instrument_await("inserting into the persistent store, invoking [write]")
-            .await;
-
-        match &storage_type {
-            StorageType::LOCALFILE => {
-                GAUGE_MEMORY_SPILL_TO_LOCALFILE.dec();
-            }
-            StorageType::HDFS => {
-                GAUGE_MEMORY_SPILL_TO_HDFS.dec();
+        if std::env::var("MEMORY_SPILL_IGNORE").unwrap_or("".to_string()) == "true" {
+            // ignore spill to localfile for debug mode
+        } else {
+            let ctx = spill_message.ctx.clone();
+            // when throwing the data lost error, it should fast fail for this partition data.
+            let result = candidate_store
+                .spill_insert(ctx)
+                .instrument_await("inserting into the persistent store, invoking [write]")
+                .await;
+
+            match &storage_type {
+                StorageType::LOCALFILE => {
+                    GAUGE_MEMORY_SPILL_TO_LOCALFILE.dec();
+                }
+                StorageType::HDFS => {
+                    GAUGE_MEMORY_SPILL_TO_HDFS.dec();
+                }
+                _ => {}
             }
-            _ => {}
-        }
 
-        let _ = result?;
+            let _ = result?;
+        }
 
         Ok(())
     }