From 8f172b643c0f0abd63be855db7cd30f9f7c10075 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 19 Jul 2026 13:18:12 -0700 Subject: [PATCH 01/21] experiment: leave only vibeio --- Cargo.lock | 396 +--------------------- Cargo.toml | 5 +- README.md | 34 +- docs/how-it-works.md | 7 +- examples/rust/Cargo.lock | 710 +++++++-------------------------------- examples/rust/Cargo.toml | 4 +- src/loop_core.rs | 5 - src/runtime.rs | 472 +++++--------------------- src/stream_transport.rs | 190 +++-------- 9 files changed, 286 insertions(+), 1537 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7d6c587..9451126 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,12 +11,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "arrayvec" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" - [[package]] name = "async-channel" version = "1.9.0" @@ -256,12 +250,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "cfg_aliases" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" - [[package]] name = "cmake" version = "0.1.58" @@ -271,164 +259,6 @@ dependencies = [ "cc", ] -[[package]] -name = "compio" -version = "0.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5cf9e29d3c2d2a37078198795631b61d2b5c9cab68191e31526fe342d742135" -dependencies = [ - "compio-buf", - "compio-driver", - "compio-io", - "compio-log", - "compio-net", - "compio-runtime", - "compio-signal", -] - -[[package]] -name = "compio-buf" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b51a2c35873865376ed4cdb6cfeb602b6cf569815f017ddd7bd8f86ad49b34c7" -dependencies = [ - "arrayvec", - "bytes", - "libc", -] - -[[package]] -name = "compio-driver" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a580d231c022f336ccb7a2e133eaf56e34bf586417176a1a8d392a4a6e9a41f4" -dependencies = [ - "bitflags 2.13.0", - "cfg_aliases", - "compio-buf", - "compio-log", - "compio-send-wrapper", - "crossbeam-queue", - "flume", - "futures-util", - "io-uring 0.7.12", - "libc", - "linux-raw-sys", - "mod_use", - "once_cell", - "pastey", - "polling", - "rustix", - "smallvec", - "socket2", - "synchrony", - "thin-cell", - "windows-sys 0.61.2", -] - -[[package]] -name = "compio-executor" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "809fc31e0180a97ea760300b9aa97d99fbb19c9831cc49cebfd5a98f44a8371b" -dependencies = [ - "compio-log", - "compio-send-wrapper", - "crossbeam-queue", - "loom", - "slotmap", -] - -[[package]] -name = "compio-io" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3a354e085c4046dd8d2d9d514cf8c15c615eb8adf6c5f41dbf7899420457e9b" -dependencies = [ - "compio-buf", - "futures-util", - "libc", - "pastey", - "pin-project-lite", - "rustix", - "synchrony", - "windows-sys 0.61.2", -] - -[[package]] -name = "compio-log" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef39fff6341af7ab6c27fae9a3887e1ec618320d43f3b9c924c7a644f693a859" -dependencies = [ - "tracing", -] - -[[package]] -name = "compio-net" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a985ea5d7580318104d44f4a3c8ae48a7ee750e9c98b68a69a869e01d4d6df0" -dependencies = [ - "compio-buf", - "compio-driver", - "compio-io", - "compio-runtime", - "either", - "futures-util", - "libc", - "once_cell", - "pin-project-lite", - "socket2", - "synchrony", - "widestring", - "windows-sys 0.61.2", -] - -[[package]] -name = "compio-runtime" -version = "0.12.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dde4b0b7de96fc8b99e089964d65e9ed5669f8ea39bc46381f7fd563dc10a753" -dependencies = [ - "compio-buf", - "compio-driver", - "compio-executor", - "compio-io", - "compio-log", - "compio-send-wrapper", - "core_affinity", - "futures-util", - "libc", - "pin-project-lite", - "scoped-tls", - "synchrony", - "windows-sys 0.61.2", -] - -[[package]] -name = "compio-send-wrapper" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f1ef3947d751725afe49ab18ce447d819cd8f4ea9b58ae94b8498eb8ad1f864" -dependencies = [ - "cfg-if", - "loom", -] - -[[package]] -name = "compio-signal" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1748b284601a5a267824b36894c7fc5e488b503f14a0a508f3b457ae018a451" -dependencies = [ - "nix", - "once_cell", - "slab", - "synchrony", - "windows-sys 0.61.2", -] - [[package]] name = "concurrent-queue" version = "2.5.0" @@ -454,17 +284,6 @@ 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 = "crossbeam-channel" version = "0.5.15" @@ -495,12 +314,6 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" -[[package]] -name = "either" -version = "1.16.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" - [[package]] name = "errno" version = "0.3.14" @@ -550,15 +363,6 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" -[[package]] -name = "flume" -version = "0.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" -dependencies = [ - "spin", -] - [[package]] name = "fs_extra" version = "1.3.0" @@ -726,23 +530,6 @@ dependencies = [ "libc", ] -[[package]] -name = "io-uring" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" -dependencies = [ - "bitflags 2.13.0", - "cfg-if", - "libc", -] - -[[package]] -name = "itoa" -version = "1.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" - [[package]] name = "jobserver" version = "0.1.34" @@ -791,21 +578,6 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" -[[package]] -name = "local-event" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76dda8459b10a8960dfae91c1c77316fc15e7caf94f9f18794126512a8dc77e5" - -[[package]] -name = "lock_api" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" -dependencies = [ - "scopeguard", -] - [[package]] name = "log" version = "0.4.33" @@ -823,10 +595,7 @@ checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" dependencies = [ "cfg-if", "generator", - "pin-utils", "scoped-tls", - "serde", - "serde_json", "tracing", "tracing-subscriber", ] @@ -858,24 +627,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "mod_use" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21909324aa58f5c284d91cac514c6d081210901dc372d7c8ea6a9d7e0406097a" - -[[package]] -name = "nix" -version = "0.31.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" -dependencies = [ - "bitflags 2.13.0", - "cfg-if", - "cfg_aliases", - "libc", -] - [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -885,16 +636,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "num_cpus" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" -dependencies = [ - "hermit-abi", - "libc", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -919,12 +660,6 @@ version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" -[[package]] -name = "pastey" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" - [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1119,7 +854,6 @@ name = "rsloop" version = "0.1.26" dependencies = [ "async-std", - "compio", "crossbeam-channel", "futures", "libc", @@ -1229,12 +963,6 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" -[[package]] -name = "scopeguard" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" - [[package]] name = "security-framework" version = "3.7.0" @@ -1258,49 +986,6 @@ dependencies = [ "libc", ] -[[package]] -name = "serde" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" -dependencies = [ - "serde_core", - "serde_derive", -] - -[[package]] -name = "serde_core" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "serde_json" -version = "1.0.150" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" -dependencies = [ - "itoa", - "memchr", - "serde", - "serde_core", - "zmij", -] - [[package]] name = "sharded-slab" version = "0.1.7" @@ -1342,15 +1027,6 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" -[[package]] -name = "slotmap" -version = "1.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdd58c3c93c3d278ca835519292445cb4b0d4dc59ccfdf7ceadaab3f8aeb4038" -dependencies = [ - "version_check", -] - [[package]] name = "smallvec" version = "1.15.2" @@ -1367,15 +1043,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "spin" -version = "0.9.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" -dependencies = [ - "lock_api", -] - [[package]] name = "subtle" version = "2.6.1" @@ -1393,33 +1060,12 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "synchrony" -version = "0.1.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "416090a4d8f6358526df5f9f65dfe28750b8b7bfd1fd8a5620f483fc4a75722c" -dependencies = [ - "event-listener 5.4.1", - "futures-util", - "local-event", - "loom", -] - [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" -[[package]] -name = "thin-cell" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d3c47fea6e344ef1ac01be266c27945ec30ae8049c3292c897f0878de5c784" -dependencies = [ - "synchrony", -] - [[package]] name = "thread_local" version = "1.1.9" @@ -1533,12 +1179,6 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ba6f5989077681266825251a52748b8c1d8a4ad098cc37e440103d0ea717fc0" -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - [[package]] name = "vibeio" version = "0.2.15" @@ -1547,7 +1187,7 @@ checksum = "63b968e30d91bc5c14d23f9af184f26ea7f0d886ab309d11b37148b244d7dfa4" dependencies = [ "crossbeam-queue", "futures-util", - "io-uring 0.6.4", + "io-uring", "libc", "mio", "oneshot", @@ -1626,34 +1266,6 @@ dependencies = [ "unicode-ident", ] -[[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" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-link" version = "0.2.1" @@ -1762,9 +1374,3 @@ name = "zeroize" version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" - -[[package]] -name = "zmij" -version = "1.0.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml index 9d1e654..37b2baf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,6 @@ profiler = ["dep:tracy-client", "profiling/profile-with-tracy"] [dependencies] async-std = "1" -compio = { version = "0.19", default-features = false, features = ["runtime", "net", "io", "signal", "time", "io-uring"] } crossbeam-channel = "0.5" futures = { version = "0.3", default-features = false, features = ["std"] } libc = "0.2" @@ -31,11 +30,11 @@ rustls-native-certs = "0.8" rustls-pemfile = "2" signal-hook = "0.4" socket2 = "0.6" +tokio = { version = "1", default-features = false, features = ["io-util"] } tracy-client = { version = "0.18", default-features = false, features = ["enable", "only-localhost", "sampling"], optional = true } +vibeio = { version = "0.2.15", default-features = false, features = ["time"] } [target.'cfg(windows)'.dependencies] -tokio = { version = "1", default-features = false, features = ["io-util"] } -vibeio = { version = "0.2.15", default-features = false } windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Networking_WinSock", "Win32_System_Pipes", "Win32_System_Threading"] } [build-dependencies] diff --git a/README.md b/README.md index dea5c25..e3f82f1 100644 --- a/README.md +++ b/README.md @@ -9,11 +9,11 @@ `rsloop` is a PyO3-based `asyncio` event loop implemented in Rust. Each `rsloop.Loop` owns a dedicated Rust runtime thread for loop coordination -and I/O work. On Linux, low-level fd watchers plus plain TCP / Unix socket -readiness, socket reads, and non-TLS server accepts are driven from that thread -through `compio` with `io_uring` support enabled. Python callbacks, tasks, and -coroutines still run on the thread that calls `run_forever()` or -`run_until_complete()` (usually the main Python thread). +and I/O work. That thread runs a `vibeio` runtime, using io_uring on Linux, +IOCP on Windows, and mio-backed readiness on other supported platforms. Plain +TCP / Unix socket reads and non-TLS server accepts run on that runtime. Python +callbacks, tasks, and coroutines still run on the thread that calls +`run_forever()` or `run_until_complete()` (usually the main Python thread). The package exposes: @@ -207,16 +207,18 @@ is backed by the lower level transport code in ## Runtime Model -Today the runtime is hybrid rather than fully single-threaded: +The runtime is centered on one `vibeio` runtime per loop: - the loop coordination thread is always the central scheduler -- on Linux, `add_reader` / `add_writer`, plain socket reads, and non-TLS socket - accept loops use the `compio` runtime on that thread +- plain TCP / Unix socket reads and non-TLS accept loops use `vibeio` on that + thread across supported platforms +- generic `add_reader` / `add_writer` descriptors use cancellable OS-poll + workers because `vibeio` does not expose arbitrary raw-descriptor registration - some transport paths still fall back to helper threads, especially TLS I/O, TLS server accept, and parts of the legacy transport write path -That means the codebase has started the move toward a single-runtime-thread I/O -model, but has not finished eliminating every helper thread yet. +The runtime dependency is now unified, but the codebase has not finished +eliminating every helper thread yet. ## Current Limitations @@ -226,7 +228,7 @@ These gaps are visible in the current implementation. CPython's OpenSSL-backed `ssl` module. In particular, encrypted private keys are not supported yet, and the fast-stream monkeypatch still falls back to stdlib helpers whenever `ssl` is enabled. TLS transport internals also still - use helper-thread paths instead of the newer runtime-thread `compio` socket + use helper-thread paths instead of the runtime-thread `vibeio` socket path. - Subprocess support still has one notable gap: `preexec_fn` remains unsupported because running arbitrary Python between @@ -240,8 +242,9 @@ These gaps are visible in the current implementation. still specific to Unix process spawning. - The transport runtime model is still in transition: plain socket reads and non-TLS accepts now run on the loop runtime thread on - Linux, but writes and TLS-heavy paths are not fully collapsed onto that same - single-threaded I/O path yet. + all supported platforms, but generic descriptor watches, writes, and + TLS-heavy paths are not fully collapsed onto that same single-threaded I/O + path yet. ## Build @@ -376,10 +379,7 @@ comparison demo. ## Acknowledgements `rsloop` builds on the Python `asyncio` model and is implemented with -[PyO3](https://pyo3.rs/) on the Rust side. The runtime and I/O work in the -current implementation rely in part on -[compio](https://github.com/compio-rs/compio). On Windows, parts of the -runtime also rely on +[PyO3](https://pyo3.rs/) on the Rust side. Runtime and socket I/O are powered by [vibeio](https://crates.io/crates/vibeio). ## License diff --git a/docs/how-it-works.md b/docs/how-it-works.md index 90a8e33..70369ea 100644 --- a/docs/how-it-works.md +++ b/docs/how-it-works.md @@ -63,14 +63,15 @@ You do not need to understand every file before using the project. For a first p ## Runtime model -The project currently uses a hybrid runtime model. +Each loop owns a `vibeio` runtime on its dedicated Rust coordination thread. The simple explanation is: - there is a dedicated Rust runtime thread for loop coordination - Python tasks and callbacks still execute on the Python side -- some I/O paths are already handled directly by the runtime thread -- some paths still use helper threads, especially around TLS and older transport code +- plain TCP / Unix reads and non-TLS accepts run directly on `vibeio` +- generic descriptor watches and some TLS, write, and older transport paths + still use helper threads That is important because it explains why the project is fast in some areas while still being a work in progress in others. diff --git a/examples/rust/Cargo.lock b/examples/rust/Cargo.lock index 2ea3031..5ce6e1b 100644 --- a/examples/rust/Cargo.lock +++ b/examples/rust/Cargo.lock @@ -2,21 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "aho-corasick" -version = "1.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" -dependencies = [ - "memchr", -] - -[[package]] -name = "arrayvec" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" - [[package]] name = "async-channel" version = "1.9.0" @@ -175,15 +160,15 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-rs" -version = "1.16.3" +version = "1.17.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" +checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" dependencies = [ "aws-lc-sys", "zeroize", @@ -191,14 +176,15 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.40.0" +version = "0.43.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" +checksum = "43103168cc76fe62678a375e722fc9cb3a0146159ac5828bc4f0dfd755c2224c" dependencies = [ "cc", "cmake", "dunce", "fs_extra", + "pkg-config", ] [[package]] @@ -209,9 +195,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.11.1" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "blocking" @@ -228,21 +214,21 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.20.2" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "cc" -version = "1.2.61" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" dependencies = [ "find-msvc-tools", "jobserver", @@ -256,12 +242,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "cfg_aliases" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" - [[package]] name = "cmake" version = "0.1.58" @@ -271,139 +251,6 @@ dependencies = [ "cc", ] -[[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-io", - "compio-log", - "compio-net", - "compio-runtime", - "compio-signal", -] - -[[package]] -name = "compio-buf" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a00d719dbd8c602ab0d25d219cbc6b517008858de7a8d6c51b4dc95aefff4dce" -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 0.7.12", - "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-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-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 = "compio-signal" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f7b5f1f32c0c44441023aa012e252de890a166f773bee49226439ac9a13f16f" -dependencies = [ - "compio-buf", - "compio-driver", - "compio-runtime", - "libc", - "once_cell", - "os_pipe", - "slab", - "windows-sys 0.61.2", -] - [[package]] name = "concurrent-queue" version = "2.5.0" @@ -429,40 +276,29 @@ 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 = "crossbeam-channel" -version = "0.5.15" +version = "0.5.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" dependencies = [ "crossbeam-utils", ] [[package]] name = "crossbeam-queue" -version = "0.3.12" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" dependencies = [ "crossbeam-utils", ] [[package]] name = "crossbeam-utils" -version = "0.8.21" +version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" [[package]] name = "dunce" @@ -470,12 +306,6 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" -[[package]] -name = "either" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" - [[package]] name = "errno" version = "0.3.14" @@ -515,9 +345,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "find-msvc-tools" @@ -525,15 +355,6 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" -[[package]] -name = "flume" -version = "0.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" -dependencies = [ - "spin", -] - [[package]] name = "fs_extra" version = "1.3.0" @@ -542,9 +363,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "futures" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" dependencies = [ "futures-channel", "futures-core", @@ -556,9 +377,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" dependencies = [ "futures-core", "futures-sink", @@ -566,15 +387,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" [[package]] name = "futures-io" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" [[package]] name = "futures-lite" @@ -591,9 +412,9 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" dependencies = [ "proc-macro2", "quote", @@ -602,21 +423,21 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ "futures-channel", "futures-core", @@ -629,21 +450,6 @@ dependencies = [ "slab", ] -[[package]] -name = "generator" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52f04ae4152da20c76fe800fa48659201d5cf627c5149ca0b707b69d7eef6cf9" -dependencies = [ - "cc", - "cfg-if", - "libc", - "log", - "rustversion", - "windows-link", - "windows-result", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -657,14 +463,13 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.3.4" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", "libc", "r-efi", - "wasip2", ] [[package]] @@ -701,47 +506,24 @@ dependencies = [ "libc", ] -[[package]] -name = "io-uring" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" -dependencies = [ - "bitflags 2.11.1", - "cfg-if", - "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 0.7.12", - "rustix", -] - [[package]] name = "jobserver" -version = "0.1.34" +version = "0.1.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" dependencies = [ - "getrandom 0.3.4", + "getrandom 0.4.3", "libc", ] [[package]] name = "js-sys" -version = "0.3.95" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", - "once_cell", "wasm-bindgen", ] @@ -754,12 +536,6 @@ dependencies = [ "log", ] -[[package]] -name = "lazy_static" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" - [[package]] name = "libc" version = "0.2.186" @@ -772,58 +548,26 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" -[[package]] -name = "lock_api" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" -dependencies = [ - "scopeguard", -] - [[package]] name = "log" -version = "0.4.29" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" dependencies = [ "value-bag", ] -[[package]] -name = "loom" -version = "0.7.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" -dependencies = [ - "cfg-if", - "generator", - "pin-utils", - "scoped-tls", - "tracing", - "tracing-subscriber", -] - -[[package]] -name = "matchers" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" -dependencies = [ - "regex-automata", -] - [[package]] name = "memchr" -version = "2.8.0" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mio" -version = "1.2.0" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "log", @@ -831,25 +575,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "nu-ansi-term" -version = "0.50.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" -dependencies = [ - "windows-sys 0.61.2", -] - -[[package]] -name = "num_cpus" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" -dependencies = [ - "hermit-abi", - "libc", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -868,28 +593,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" -[[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 = "parking" version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pin-project-lite" version = "0.2.17" @@ -913,6 +622,12 @@ dependencies = [ "futures-io", ] +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "polling" version = "3.11.0" @@ -929,33 +644,33 @@ dependencies = [ [[package]] name = "portable-atomic" -version = "1.13.1" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] [[package]] name = "profiling" -version = "1.0.17" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eb8486b569e12e2c32ad3e204dbaba5e4b5b216e9367044f25f1dba42341773" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" dependencies = [ "profiling-procmacros", ] [[package]] name = "profiling-procmacros" -version = "1.0.17" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b" +checksum = "4488a4a36b9a4ba6b9334a32a39971f77c1436ec82c38707bce707699cc3bbcb" dependencies = [ "quote", "syn", @@ -963,9 +678,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ "libc", "once_cell", @@ -977,9 +692,9 @@ dependencies = [ [[package]] name = "pyo3-async-runtimes" -version = "0.28.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e7364a95bf00e8377bbf9b0f09d7ff9715a29d8fcf93b47d1a967363b973178" +checksum = "b3ef68daa7316a3fac65e5e18b2203f010346de1c1c53456811a2624673ab046" dependencies = [ "async-std", "futures-channel", @@ -991,18 +706,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -1010,9 +725,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -1022,48 +737,30 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] [[package]] name = "quote" -version = "1.0.45" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] [[package]] name = "r-efi" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" - -[[package]] -name = "regex-automata" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" -dependencies = [ - "aho-corasick", - "memchr", - "regex-syntax", -] - -[[package]] -name = "regex-syntax" -version = "0.8.10" +version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "ring" @@ -1081,12 +778,9 @@ dependencies = [ [[package]] name = "rsloop" -version = "0.1.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d61a1709b915d770793f23399f388e8d2c3e068db4d8bbb134669e8d9500ca37" +version = "0.1.26" dependencies = [ "async-std", - "compio", "crossbeam-channel", "futures", "libc", @@ -1119,7 +813,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys", @@ -1128,9 +822,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.39" +version = "0.23.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c2c118cb077cca2822033836dfb1b975355dfb784b5e8da48f7b6c5db74e60e" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", "log", @@ -1143,9 +837,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -1164,9 +858,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "zeroize", ] @@ -1185,9 +879,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "schannel" @@ -1198,25 +892,13 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "scoped-tls" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" - -[[package]] -name = "scopeguard" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" - [[package]] name = "security-framework" version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "core-foundation", "core-foundation-sys", "libc", @@ -1233,20 +915,11 @@ dependencies = [ "libc", ] -[[package]] -name = "sharded-slab" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" -dependencies = [ - "lazy_static", -] - [[package]] name = "shlex" -version = "1.3.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "signal-hook" @@ -1274,31 +947,16 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" -[[package]] -name = "smallvec" -version = "1.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" - [[package]] name = "socket2" -version = "0.6.3" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", ] -[[package]] -name = "spin" -version = "0.9.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" -dependencies = [ - "lock_api", -] - [[package]] name = "subtle" version = "2.6.1" @@ -1307,105 +965,31 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.117" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] -[[package]] -name = "synchrony" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c174d82fd56da8214ec095cfe4568e59e5ccb49d060e70c2f98e3ba352b23e45" -dependencies = [ - "futures-util", - "loom", -] - [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" -[[package]] -name = "thin-cell" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4164c6c316ba9733b0ab021e7f9852c788a4b991b49c25820f1be48e1d41345b" - -[[package]] -name = "thread_local" -version = "1.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" -dependencies = [ - "cfg-if", -] - [[package]] name = "tokio" -version = "1.52.1" +version = "1.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67dee974fe86fd92cc45b7a95fdd2f99a36a6d7b0d431a231178d3d670bbcc6" +checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee" dependencies = [ "bytes", "pin-project-lite", ] -[[package]] -name = "tracing" -version = "0.1.44" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" -dependencies = [ - "pin-project-lite", - "tracing-core", -] - -[[package]] -name = "tracing-core" -version = "0.1.36" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" -dependencies = [ - "once_cell", - "valuable", -] - -[[package]] -name = "tracing-log" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" -dependencies = [ - "log", - "once_cell", - "tracing-core", -] - -[[package]] -name = "tracing-subscriber" -version = "0.3.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" -dependencies = [ - "matchers", - "nu-ansi-term", - "once_cell", - "regex-automata", - "sharded-slab", - "smallvec", - "thread_local", - "tracing", - "tracing-core", - "tracing-log", -] - [[package]] name = "unicode-ident" version = "1.0.24" @@ -1418,27 +1002,21 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" -[[package]] -name = "valuable" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" - [[package]] name = "value-bag" -version = "1.12.0" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ba6f5989077681266825251a52748b8c1d8a4ad098cc37e440103d0ea717fc0" +checksum = "5dd4ec1eb1d240636e354a30110a1dfcb37047169a4d9bd6d9d3469df574b5c4" [[package]] name = "vibeio" -version = "0.2.8" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b6fc74ec117ac6ea756119587875e77f738dc519141ba4aee1b35977e354b9c" +checksum = "63b968e30d91bc5c14d23f9af184f26ea7f0d886ab309d11b37148b244d7dfa4" dependencies = [ "crossbeam-queue", "futures-util", - "io-uring 0.6.4", + "io-uring", "libc", "mio", "oneshot", @@ -1453,20 +1031,11 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" -[[package]] -name = "wasip2" -version = "1.0.3+wasi-0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" -dependencies = [ - "wit-bindgen", -] - [[package]] name = "wasm-bindgen" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -1477,9 +1046,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.68" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f371d383f2fb139252e0bfac3b81b265689bf45b6874af544ffa4c975ac1ebf8" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -1487,9 +1056,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1497,9 +1066,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", @@ -1510,56 +1079,19 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] -[[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" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" -[[package]] -name = "windows-result" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" -dependencies = [ - "windows-link", -] - [[package]] name = "windows-sys" version = "0.52.0" @@ -1642,14 +1174,8 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "wit-bindgen" -version = "0.57.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" - [[package]] name = "zeroize" -version = "1.8.2" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" diff --git a/examples/rust/Cargo.toml b/examples/rust/Cargo.toml index bbe2cde..d5600e6 100644 --- a/examples/rust/Cargo.toml +++ b/examples/rust/Cargo.toml @@ -10,8 +10,8 @@ crate-type = ["cdylib"] [dependencies] async-std = "1" -pyo3 = { version = "0.28", features = ["extension-module"] } -rsloop = { version = "0.1.15" } +pyo3 = { version = "0.29", features = ["extension-module"] } +rsloop = { path = "../.." } [profile.release] opt-level = 3 diff --git a/src/loop_core.rs b/src/loop_core.rs index 5ef5804..d5a801a 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -5,7 +5,6 @@ use std::ops::DerefMut; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc; use std::sync::{Arc, Mutex}; -#[cfg(target_os = "linux")] use std::task::Waker; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -153,7 +152,6 @@ pub struct LoopCore { next_callback_id: AtomicU64, command_tx: Sender, runtime_thread: Mutex>>, - #[cfg(target_os = "linux")] runtime_waker: Mutex>, active_ready_dispatch: Mutex>, } @@ -169,7 +167,6 @@ impl LoopCore { next_callback_id: AtomicU64::new(1), command_tx, runtime_thread: Mutex::new(None), - #[cfg(target_os = "linux")] runtime_waker: Mutex::new(None), active_ready_dispatch: Mutex::new(None), }); @@ -196,7 +193,6 @@ impl LoopCore { self.command_tx .send(command) .map_err(|_| LoopCoreError::ChannelClosed)?; - #[cfg(target_os = "linux")] if let Some(waker) = self .runtime_waker .lock() @@ -747,7 +743,6 @@ impl LoopCore { ACTIVE_LOOP_TLS.with(|tls| tls.core.set(self as *const Self)); } - #[cfg(target_os = "linux")] pub(crate) fn set_runtime_waker(&self, waker: Option) { *self.runtime_waker.lock().expect("poisoned runtime waker") = waker; } diff --git a/src/runtime.rs b/src/runtime.rs index 5ddcf71..4fa4790 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,25 +1,17 @@ use std::collections::{BinaryHeap, HashMap, VecDeque}; -#[cfg(target_os = "linux")] -use std::os::fd::{FromRawFd, OwnedFd}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; -#[cfg(not(target_os = "linux"))] use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; +use std::task::{Context, Poll}; use std::thread; -use std::time::{Duration, Instant}; +use std::time::Instant; use crate::fd_ops; use crate::loop_core::{ LoopCommand, LoopCore, LoopFutureCommand, LoopIoCommand, LoopRunCommand, LoopSignalCommand, LoopTransportCommand, ReadyItem, }; -#[cfg(windows)] -use crate::windows_vibeio; -#[cfg(target_os = "linux")] -use compio::runtime::fd::PollFd; -#[cfg(target_os = "linux")] -use compio::runtime::{JoinHandle as CompioJoinHandle, Runtime as CompioRuntime}; -#[cfg(not(target_os = "linux"))] -use crossbeam_channel::RecvTimeoutError; use crossbeam_channel::{Receiver, TryRecvError}; use pyo3::prelude::*; #[cfg(unix)] @@ -28,14 +20,10 @@ use signal_hook::iterator::{Handle as SignalHandle, Signals}; mod timer_entry; use timer_entry::TimerEntry; -#[cfg(target_os = "linux")] -const IO_POLL_MAX_SLEEP: Duration = Duration::from_millis(10); - struct RuntimeDispatcher { core: Arc, command_rx: Receiver, - #[cfg(target_os = "linux")] - io_runtime: CompioRuntime, + timer_wait: Option<(Instant, vibeio::time::Sleep)>, ready_batch: VecDeque, timers: BinaryHeap, next_timer_id: u64, @@ -62,47 +50,14 @@ struct SignalWatcher { #[cfg(not(unix))] struct SignalWatcher; -#[cfg(target_os = "linux")] -type WatchTask = CompioJoinHandle<()>; - -#[cfg(all(not(target_os = "linux"), not(windows)))] -struct WatchTask { - stop: Arc, - join: thread::JoinHandle<()>, -} - -#[cfg(all(not(target_os = "linux"), not(windows)))] -impl WatchTask { - fn spawn(name: String, task: impl FnOnce(Arc) + Send + 'static) -> Self { - let stop = Arc::new(AtomicBool::new(false)); - let thread_stop = Arc::clone(&stop); - let join = thread::Builder::new() - .name(name) - .spawn(move || task(thread_stop)) - .expect("failed to spawn watch task"); - Self { stop, join } - } - - fn abort(self) { - self.stop.store(true, AtomicOrdering::Release); - let _ = self.join.join(); - } - - fn cancel(self) { - self.stop.store(true, AtomicOrdering::Release); - } -} - -#[cfg(windows)] enum WatchTask { Thread { stop: Arc, join: thread::JoinHandle<()>, }, - Vibeio(windows_vibeio::TaskHandle), + Vibeio(vibeio::JoinHandle<()>), } -#[cfg(windows)] impl WatchTask { fn spawn_thread(name: String, task: impl FnOnce(Arc) + Send + 'static) -> Self { let stop = Arc::new(AtomicBool::new(false)); @@ -120,7 +75,7 @@ impl WatchTask { stop.store(true, AtomicOrdering::Release); let _ = join.join(); } - Self::Vibeio(task) => windows_vibeio::cancel(task), + Self::Vibeio(task) => task.cancel(), } } @@ -129,30 +84,16 @@ impl WatchTask { Self::Thread { stop, .. } => { stop.store(true, AtomicOrdering::Release); } - Self::Vibeio(task) => windows_vibeio::cancel(task), + Self::Vibeio(task) => task.cancel(), } } } -#[cfg(target_os = "linux")] -#[inline] -fn abort_watch_task(task: WatchTask) { - drop(task); -} - -#[cfg(not(target_os = "linux"))] #[inline] fn abort_watch_task(task: WatchTask) { task.abort(); } -#[cfg(target_os = "linux")] -#[inline] -fn cancel_watch_task(task: WatchTask) { - drop(task); -} - -#[cfg(not(target_os = "linux"))] #[inline] fn cancel_watch_task(task: WatchTask) { task.cancel(); @@ -164,16 +105,14 @@ pub fn run_runtime_thread(core: Arc, command_rx: Receiver if tracy_client::Client::is_running() { tracy_client::set_thread_name!("rsloop-runtime"); } - #[cfg(target_os = "linux")] - let io_runtime = CompioRuntime::new().expect("failed to initialize compio runtime"); - #[cfg(target_os = "linux")] - core.set_runtime_waker(Some(io_runtime.waker())); - - let mut dispatcher = RuntimeDispatcher { + let runtime = vibeio::RuntimeBuilder::new() + .enable_timer(true) + .build() + .expect("failed to initialize vibeio runtime"); + let dispatcher = RuntimeDispatcher { core: Arc::clone(&core), command_rx, - #[cfg(target_os = "linux")] - io_runtime, + timer_wait: None, ready_batch: VecDeque::new(), timers: BinaryHeap::new(), next_timer_id: 0, @@ -184,104 +123,69 @@ pub fn run_runtime_thread(core: Arc, command_rx: Receiver accept_tasks: HashMap::new(), shutting_down: false, }; - dispatcher.run(); - #[cfg(target_os = "linux")] + runtime.block_on(dispatcher); core.set_runtime_waker(None); } -impl RuntimeDispatcher { - fn run(&mut self) { - profiling::scope!("runtime.dispatcher.run"); - #[cfg(target_os = "linux")] - { - let runtime = self.io_runtime.clone(); - runtime.enter(|| self.run_inner()); - } - - #[cfg(not(target_os = "linux"))] - self.run_inner(); - } +impl Future for RuntimeDispatcher { + type Output = (); - fn run_inner(&mut self) { - profiling::scope!("runtime.dispatcher.run_inner"); + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + profiling::scope!("runtime.dispatcher.poll"); + self.core.set_runtime_waker(Some(cx.waker().clone())); loop { if self.shutting_down { - break; + return Poll::Ready(()); } self.collect_expired_timers(); if self.active_run.is_some() && self.has_ready() { - if self.dispatch_ready_batch() { - break; - } + self.dispatch_ready_batch(); + } + + if self.drain_commands() { + return Poll::Ready(()); + } + self.collect_expired_timers(); + if self.active_run.is_some() && self.has_ready() { + self.dispatch_ready_batch(); continue; } - if self.wait_for_work() { - break; + if self.active_run.is_none() { + self.timer_wait = None; + return Poll::Pending; + } + + let Some(deadline) = self.timers.peek().map(|entry| entry.when) else { + self.timer_wait = None; + return Poll::Pending; + }; + let replace_timer = self + .timer_wait + .as_ref() + .is_none_or(|(current, _)| *current != deadline); + if replace_timer { + self.timer_wait = Some((deadline, vibeio::time::Sleep::sleep_until(deadline))); + } + let (_, sleep) = self.timer_wait.as_mut().expect("timer wait missing"); + if Pin::new(sleep).poll(cx).is_ready() { + self.timer_wait = None; + continue; } + return Poll::Pending; } } +} +impl RuntimeDispatcher { #[inline] fn has_ready(&self) -> bool { !self.ready_batch.is_empty() } - fn wait_for_work(&mut self) -> bool { - profiling::scope!("runtime.wait_for_work"); - #[cfg(target_os = "linux")] - { - self.io_runtime.run(); - - loop { - match self.command_rx.try_recv() { - Ok(command) => { - if self.handle_received_command(command) { - return true; - } - } - Err(TryRecvError::Empty) => break, - Err(TryRecvError::Disconnected) => return true, - } - } - - self.io_runtime.poll_with(self.next_wait_timeout()); - self.io_runtime.run(); - false - } - - #[cfg(not(target_os = "linux"))] - { - if self.active_run.is_none() { - return match self.command_rx.recv() { - Ok(command) => self.handle_received_command(command), - Err(_) => true, - }; - } - - match self.next_wait_timeout() { - Some(timeout) if timeout.is_zero() => false, - Some(timeout) => match self.command_rx.recv_timeout(timeout) { - Ok(command) => self.handle_received_command(command), - Err(RecvTimeoutError::Timeout) => false, - Err(RecvTimeoutError::Disconnected) => true, - }, - None => match self.command_rx.recv() { - Ok(command) => self.handle_received_command(command), - Err(_) => true, - }, - } - } - } - - fn handle_received_command(&mut self, command: LoopCommand) -> bool { - profiling::scope!("runtime.handle_received_command"); - if self.handle_command(command) { - return true; - } - + fn drain_commands(&mut self) -> bool { loop { match self.command_rx.try_recv() { Ok(command) => { @@ -295,22 +199,6 @@ impl RuntimeDispatcher { } } - #[inline] - fn next_wait_timeout(&self) -> Option { - let timeout = self.timers.peek().map(|entry| { - let now = Instant::now(); - entry.when.saturating_duration_since(now) - }); - #[cfg(target_os = "linux")] - { - Some(timeout.map_or(IO_POLL_MAX_SLEEP, |value| value.min(IO_POLL_MAX_SLEEP))) - } - #[cfg(not(target_os = "linux"))] - { - timeout - } - } - fn collect_expired_timers(&mut self) { profiling::scope!("runtime.collect_expired_timers"); if self.active_run.is_none() { @@ -468,99 +356,26 @@ impl RuntimeDispatcher { cancel_watch_task(task); } - #[cfg(target_os = "linux")] - let task = { - let sender = Arc::clone(&self.core); - self.io_runtime.spawn(async move { - let Ok(poll_fd) = poll_fd_from_raw(fd) else { - return; - }; - if poll_fd.read_ready().await.is_err() { - return; - } - if callback.cancelled() { + let sender = Arc::clone(&self.core); + let task = WatchTask::spawn_thread(format!("rsloop-reader-{fd}"), move |stop| { + loop { + if stop.load(AtomicOrdering::Acquire) { return; } - - let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); - }) - }; - - #[cfg(windows)] - let task = { - let sender = Arc::clone(&self.core); - let service_callback = Arc::clone(&callback); - match fd_ops::duplicate_tcp_stream(fd).and_then(|stream| { - if stream.peer_addr().is_err() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "listener sockets stay on the blocking watcher path", - )); - } - windows_vibeio::spawn(move || async move { - let ready_result = async { - let stream = vibeio::net::PollTcpStream::from_std(stream)?; - let mut buf = [0_u8; 1]; - stream.peek(&mut buf).await.map(|_| ()) - } - .await; - if ready_result.is_err() || service_callback.cancelled() { - return; - } - - let _ = - sender.send_command(LoopCommand::ScheduleReady(service_callback)); - }) - }) { - Ok(task) => WatchTask::Vibeio(task), - Err(_) => { - let sender = Arc::clone(&self.core); - WatchTask::spawn_thread(format!("rsloop-reader-{fd}"), move |stop| { - loop { - if stop.load(AtomicOrdering::Acquire) { - return; - } - match fd_ops::poll_fd(fd, true, false, 50) { - Ok((true, _)) => break, - Ok((false, false)) => continue, - Ok((false, true)) => continue, - Err(_) => return, - } - } - - if stop.load(AtomicOrdering::Acquire) || callback.cancelled() { - return; - } - - let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); - }) + match fd_ops::poll_fd(fd, true, false, 50) { + Ok((true, _)) => break, + Ok((false, false)) => continue, + Ok((false, true)) => continue, + Err(_) => return, } } - }; - #[cfg(all(not(target_os = "linux"), not(windows)))] - let task = { - let sender = Arc::clone(&self.core); - WatchTask::spawn(format!("rsloop-reader-{fd}"), move |stop| { - loop { - if stop.load(AtomicOrdering::Acquire) { - return; - } - match fd_ops::poll_fd(fd, true, false, 50) { - Ok((true, _)) => break, - Ok((false, false)) => continue, - Ok((false, true)) => continue, - Err(_) => return, - } - } - - if stop.load(AtomicOrdering::Acquire) || callback.cancelled() { - return; - } + if stop.load(AtomicOrdering::Acquire) || callback.cancelled() { + return; + } - let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); - }) - }; + let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); + }); self.reader_tasks.insert(fd, task); } @@ -574,28 +389,10 @@ impl RuntimeDispatcher { cancel_watch_task(task); } - #[cfg(target_os = "linux")] - let task = { - let sender = Arc::clone(&self.core); - self.io_runtime.spawn(async move { - let Ok(poll_fd) = poll_fd_from_raw(fd) else { - return; - }; - if poll_fd.write_ready().await.is_err() { - return; - } - if callback.cancelled() { - return; - } - - let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); - }) - }; - - #[cfg(not(target_os = "linux"))] - let task = { - let sender = Arc::clone(&self.core); - let watch_body = move |stop: Arc| { + let sender = Arc::clone(&self.core); + let task = WatchTask::spawn_thread( + format!("rsloop-writer-{fd}"), + move |stop: Arc| { loop { if stop.load(AtomicOrdering::Acquire) { return; @@ -613,16 +410,8 @@ impl RuntimeDispatcher { } let _ = sender.send_command(LoopCommand::ScheduleReady(callback)); - }; - #[cfg(windows)] - { - WatchTask::spawn_thread(format!("rsloop-writer-{fd}"), watch_body) - } - #[cfg(not(windows))] - { - WatchTask::spawn(format!("rsloop-writer-{fd}"), watch_body) - } - }; + }, + ); self.writer_tasks.insert(fd, task); } @@ -636,50 +425,25 @@ impl RuntimeDispatcher { cancel_watch_task(task); } - #[cfg(target_os = "linux")] - let task = self - .io_runtime - .spawn(crate::stream_transport::run_socket_reader_task( - core, reader, - )); - - #[cfg(windows)] let task = match reader { crate::stream_transport::ReaderTarget::Tcp(stream) => { - let service_core = core.clone(); - match stream.try_clone().and_then(|dup| { - windows_vibeio::spawn(move || { - crate::stream_transport::run_tcp_socket_reader_task( - service_core.clone(), - dup, - ) - }) - }) { - Ok(task) => WatchTask::Vibeio(task), - Err(_) => WatchTask::spawn_thread( - format!("rsloop-socket-reader-{fd}"), - move |stop| { - crate::stream_transport::run_socket_reader_blocking( - core, - crate::stream_transport::ReaderTarget::Tcp(stream), - stop, - ) - }, - ), - } + WatchTask::Vibeio(vibeio::spawn( + crate::stream_transport::run_tcp_socket_reader_task(core, stream), + )) + } + #[cfg(unix)] + crate::stream_transport::ReaderTarget::Unix(stream) => { + WatchTask::Vibeio(vibeio::spawn( + crate::stream_transport::run_unix_socket_reader_task(core, stream), + )) } - other => { + other @ crate::stream_transport::ReaderTarget::File(_) => { WatchTask::spawn_thread(format!("rsloop-socket-reader-{fd}"), move |stop| { crate::stream_transport::run_socket_reader_blocking(core, other, stop) }) } }; - #[cfg(all(not(target_os = "linux"), not(windows)))] - let task = WatchTask::spawn(format!("rsloop-socket-reader-{fd}"), move |stop| { - crate::stream_transport::run_socket_reader_blocking(core, reader, stop) - }); - self.reader_tasks.insert(fd, task); } LoopCommand::Io(LoopIoCommand::StopSocketReader(fd)) => { @@ -696,54 +460,9 @@ impl RuntimeDispatcher { cancel_watch_task(task); } - #[cfg(target_os = "linux")] - let task = self - .io_runtime - .spawn(crate::stream_transport::run_server_accept_task( - server, listener, - )); - - #[cfg(windows)] - let task = match listener { - crate::stream_transport::ServerListener::Tcp(listener) => { - let service_server = server.clone(); - match listener.try_clone().and_then(|dup| { - windows_vibeio::spawn(move || { - crate::stream_transport::run_tcp_accept_task( - service_server.clone(), - dup, - ) - }) - }) { - Ok(task) => WatchTask::Vibeio(task), - Err(_) => WatchTask::spawn_thread( - format!("rsloop-accept-{fd}"), - move |stop| { - crate::stream_transport::run_server_accept_blocking( - crate::stream_transport::BlockingAcceptLoop::new( - server, - crate::stream_transport::ServerListener::Tcp(listener), - stop, - ), - ) - }, - ), - } - } - #[cfg(unix)] - other => WatchTask::spawn_thread(format!("rsloop-accept-{fd}"), move |stop| { - crate::stream_transport::run_server_accept_blocking( - crate::stream_transport::BlockingAcceptLoop::new(server, other, stop), - ) - }), - }; - - #[cfg(all(not(target_os = "linux"), not(windows)))] - let task = WatchTask::spawn(format!("rsloop-accept-{fd}"), move |stop| { - crate::stream_transport::run_server_accept_blocking( - crate::stream_transport::BlockingAcceptLoop::new(server, listener, stop), - ) - }); + let task = WatchTask::Vibeio(vibeio::spawn( + crate::stream_transport::run_server_accept_task(server, listener), + )); self.accept_tasks.insert(fd, task); } @@ -768,14 +487,14 @@ impl RuntimeDispatcher { false } - fn dispatch_ready_batch(&mut self) -> bool { + fn dispatch_ready_batch(&mut self) { profiling::scope!("runtime.dispatch_ready_batch"); let Some(active_run) = self.active_run.as_ref() else { - return false; + return; }; if self.ready_batch.is_empty() { - return false; + return; } let mut pending = active_run @@ -793,7 +512,6 @@ impl RuntimeDispatcher { { let _ = active_run.wake_tx.send(()); } - false } fn finish_run(&mut self) { @@ -833,13 +551,3 @@ impl RuntimeDispatcher { } } } - -#[cfg(target_os = "linux")] -fn poll_fd_from_raw(fd: fd_ops::RawFd) -> std::io::Result> { - let dup = fd_ops::dup_raw_fd(fd)?; - let dup: i32 = dup - .try_into() - .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidInput, "fd out of range"))?; - // SAFETY: `dup_raw_fd` returned a newly owned descriptor, and `OwnedFd` takes ownership of it. - PollFd::new(unsafe { OwnedFd::from_raw_fd(dup) }) -} diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 925058b..d52bd6e 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -3,10 +3,8 @@ use std::fs; use std::io::{self, Read, Write as _}; use std::net::{Shutdown, TcpListener as StdTcpListener, TcpStream as StdTcpStream}; use std::ops::DerefMut; -#[cfg(target_os = "linux")] -use std::os::fd::OwnedFd; #[cfg(unix)] -use std::os::fd::{AsRawFd, FromRawFd}; +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd}; use std::os::raw::c_int; #[cfg(unix)] use std::os::unix::net::{UnixListener as StdUnixListener, UnixStream as StdUnixStream}; @@ -21,20 +19,16 @@ use std::sync::{Arc, Condvar, Mutex, Weak}; use std::thread; use std::time::Duration; -#[cfg(target_os = "linux")] -use compio::runtime::fd::PollFd; -#[cfg(target_os = "linux")] -use compio::time::sleep as compio_sleep; use pyo3::exceptions::{PyRuntimeError, PyTimeoutError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyByteArray, PyByteArrayMethods, PyBytes, PyDict, PySlice, PyString, PyTuple}; use pyo3_async_runtimes::TaskLocals; use rustls::{ClientConnection, ServerConnection}; use socket2::Socket; -#[cfg(windows)] use tokio::io::AsyncReadExt; -#[cfg(windows)] use vibeio::net::{PollTcpStream as VibePollTcpStream, TcpListener as VibeTcpListener}; +#[cfg(unix)] +use vibeio::net::{PollUnixStream as VibePollUnixStream, UnixListener as VibeUnixListener}; use crate::async_event::AsyncEvent; use crate::context::{ @@ -3296,7 +3290,6 @@ fn run_tcp_accept_loop(params: BlockingAcceptLoop) { } } -#[cfg(target_os = "linux")] pub(crate) async fn run_server_accept_task(server: Arc, listener: ServerListener) { profiling::scope!("stream.run_server_accept_task"); match listener { @@ -3306,72 +3299,8 @@ pub(crate) async fn run_server_accept_task(server: Arc, listener: Se } } -#[cfg(not(target_os = "linux"))] -pub(crate) fn run_server_accept_blocking(params: BlockingAcceptLoop) { - let BlockingAcceptLoop { - server, - listener, - stop, - } = params; - match listener { - ServerListener::Tcp(listener) => { - run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - } - #[cfg(unix)] - ServerListener::Unix(listener) => { - run_unix_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - } - } -} - -#[cfg(target_os = "linux")] async fn run_tcp_accept_task(server: Arc, listener: StdTcpListener) { profiling::scope!("stream.run_tcp_accept_task"); - let poll_fd = listener.try_clone().and_then(PollFd::new); - - let Ok(poll_fd) = poll_fd else { - return; - }; - - loop { - if server.is_closed() { - return; - } - - if let Err(err) = poll_fd.accept_ready().await { - report_server_io_error(&server, err, "TCP server accept failed"); - return; - } - - loop { - let accept_result = listener.accept(); - match accept_result { - Ok((stream, _addr)) => { - if !configure_accepted_tcp_stream( - &server, - &stream, - "failed to configure TCP connection", - ) { - continue; - } - schedule_accepted_transport( - &server, - AcceptedStream::Tcp(stream), - "failed to accept TCP connection", - ); - } - Err(err) if err.kind() == io::ErrorKind::WouldBlock => break, - Err(err) => { - report_server_io_error(&server, err, "TCP server accept failed"); - return; - } - } - } - } -} - -#[cfg(windows)] -pub(crate) async fn run_tcp_accept_task(server: Arc, listener: StdTcpListener) { let listener = match VibeTcpListener::from_std(listener) { Ok(listener) => listener, Err(err) => { @@ -3387,7 +3316,11 @@ pub(crate) async fn run_tcp_accept_task(server: Arc, listener: StdTc match listener.accept().await { Ok((stream, _addr)) => { + #[cfg(unix)] + let stream = unsafe { StdTcpStream::from_raw_fd(stream.into_raw_fd()) }; + #[cfg(windows)] let raw = stream.into_raw_socket(); + #[cfg(windows)] let stream = from_owned_raw_socket::(raw); if !configure_accepted_tcp_stream( &server, @@ -3457,10 +3390,14 @@ fn run_unix_accept_loop(params: BlockingAcceptLoop) { } } -#[cfg(target_os = "linux")] +#[cfg(unix)] async fn run_unix_accept_task(server: Arc, listener: StdUnixListener) { - let Ok(poll_fd) = listener.try_clone().and_then(PollFd::new) else { - return; + let listener = match VibeUnixListener::from_std(listener) { + Ok(listener) => listener, + Err(err) => { + report_server_io_error(&server, err, "Unix server accept failed"); + return; + } }; loop { @@ -3468,32 +3405,25 @@ async fn run_unix_accept_task(server: Arc, listener: StdUnixListener return; } - if let Err(err) = poll_fd.accept_ready().await { - report_server_io_error(&server, err, "Unix server accept failed"); - return; - } - - loop { - match listener.accept() { - Ok((stream, _addr)) => { - if !configure_accepted_unix_stream( - &server, - &stream, - "failed to configure Unix connection", - ) { - continue; - } - schedule_accepted_transport( - &server, - AcceptedStream::Unix(stream), - "failed to accept Unix connection", - ); - } - Err(err) if err.kind() == io::ErrorKind::WouldBlock => break, - Err(err) => { - report_server_io_error(&server, err, "Unix server accept failed"); - return; + match listener.accept().await { + Ok((stream, _addr)) => { + let stream = unsafe { StdUnixStream::from_raw_fd(stream.into_raw_fd()) }; + if !configure_accepted_unix_stream( + &server, + &stream, + "failed to configure Unix connection", + ) { + continue; } + schedule_accepted_transport( + &server, + AcceptedStream::Unix(stream), + "failed to accept Unix connection", + ); + } + Err(err) => { + report_server_io_error(&server, err, "Unix server accept failed"); + return; } } } @@ -3815,19 +3745,21 @@ fn spin_read_stream( } } -#[cfg(target_os = "linux")] -pub(crate) async fn run_socket_reader_task( +pub(crate) async fn run_tcp_socket_reader_task( core: Arc, - mut reader: ReaderTarget, + stream: StdTcpStream, ) { - profiling::scope!("stream.run_socket_reader_task"); - let Ok(poll_fd) = poll_fd_from_raw(reader.fd()) else { - core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( - "failed to attach socket reader".to_owned(), - ))); - return; + profiling::scope!("stream.run_tcp_socket_reader_task"); + let mut reader = match VibePollTcpStream::from_std(stream) { + Ok(reader) => reader, + Err(err) => { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( + err.to_string(), + ))); + return; + } }; - let mut buf = [0_u8; STREAM_READ_BUFFER_SIZE]; + let mut buf = vec![0_u8; STREAM_READ_BUFFER_SIZE]; loop { if core.is_closing() { @@ -3836,21 +3768,14 @@ pub(crate) async fn run_socket_reader_task( } while !core.is_closing() && !core.is_reading() { - compio_sleep(Duration::from_millis(1)).await; + vibeio::time::sleep(Duration::from_millis(1)).await; } if core.is_closing() { core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); return; } - if let Err(err) = poll_fd.read_ready().await { - core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( - err.to_string(), - ))); - return; - } - - match reader.read(&mut buf) { + match reader.read(&mut buf).await { Ok(0) => { core.enqueue_pending_read_event(PendingReadEvent::Eof); return; @@ -3869,13 +3794,13 @@ pub(crate) async fn run_socket_reader_task( } } -#[cfg(windows)] -pub(crate) async fn run_tcp_socket_reader_task( +#[cfg(unix)] +pub(crate) async fn run_unix_socket_reader_task( core: Arc, - stream: StdTcpStream, + stream: StdUnixStream, ) { - profiling::scope!("stream.run_tcp_socket_reader_task"); - let mut reader = match VibePollTcpStream::from_std(stream) { + profiling::scope!("stream.run_unix_socket_reader_task"); + let mut reader = match VibePollUnixStream::from_std(stream) { Ok(reader) => reader, Err(err) => { core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( @@ -3891,15 +3816,13 @@ pub(crate) async fn run_tcp_socket_reader_task( core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); return; } - while !core.is_closing() && !core.is_reading() { - thread::sleep(Duration::from_millis(1)); + vibeio::time::sleep(Duration::from_millis(1)).await; } if core.is_closing() { core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); return; } - match reader.read(&mut buf).await { Ok(0) => { core.enqueue_pending_read_event(PendingReadEvent::Eof); @@ -3919,7 +3842,6 @@ pub(crate) async fn run_tcp_socket_reader_task( } } -#[cfg(not(target_os = "linux"))] pub(crate) fn run_socket_reader_blocking( core: Arc, reader: ReaderTarget, @@ -4383,14 +4305,6 @@ fn wait_socket_ready_until( Ok(()) } -#[cfg(target_os = "linux")] -fn poll_fd_from_raw(fd: fd_ops::RawFd) -> io::Result> { - let dup = fd_ops::dup_raw_fd(fd)?; - let owned = from_owned_raw_fd::(dup) - .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err.to_string()))?; - PollFd::new(owned) -} - fn tls_io_error(err: rustls::Error) -> io::Error { io::Error::new(io::ErrorKind::InvalidData, err.to_string()) } From c14eeb1742745706ec55bf780d4f392b1b7e5643 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 19 Jul 2026 13:26:13 -0700 Subject: [PATCH 02/21] replace shell script with Python script for TLS cert generation --- justfile | 2 +- scripts/generate-test-tls-certs.sh | 97 --------------- scripts/generate_test_tls_certs.py | 188 +++++++++++++++++++++++++++++ scripts/test-supported-pythons.sh | 2 +- tests/test_tls.py | 19 ++- 5 files changed, 207 insertions(+), 101 deletions(-) delete mode 100755 scripts/generate-test-tls-certs.sh create mode 100644 scripts/generate_test_tls_certs.py diff --git a/justfile b/justfile index 7f39e38..3984386 100644 --- a/justfile +++ b/justfile @@ -1,7 +1,7 @@ set shell := ["bash", "-euo", "pipefail", "-c"] tls-test-certs outdir="tests/fixtures/tls": - ./scripts/generate-test-tls-certs.sh {{outdir}} + uv run --no-project python scripts/generate_test_tls_certs.py {{outdir}} fmt: uv run ruff format . diff --git a/scripts/generate-test-tls-certs.sh b/scripts/generate-test-tls-certs.sh deleted file mode 100755 index 33b521a..0000000 --- a/scripts/generate-test-tls-certs.sh +++ /dev/null @@ -1,97 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -outdir="${1:-tests/fixtures/tls}" - -mkdir -p "${outdir}" - -rm -f \ - "${outdir}/ca-cert.pem" \ - "${outdir}/ca-key.pem" \ - "${outdir}/cert.pem" \ - "${outdir}/key.pem" \ - "${outdir}/cert.csr" \ - "${outdir}/cert.cnf" \ - "${outdir}/ca.cnf" \ - "${outdir}/ca-cert.srl" - -cat <<'EOF' > "${outdir}/ca.cnf" -[req] -distinguished_name = dn -x509_extensions = v3_ca -prompt = no - -[dn] -CN = rsloop-test-ca - -[v3_ca] -basicConstraints = critical, CA:true -keyUsage = critical, keyCertSign, cRLSign -subjectKeyIdentifier = hash -authorityKeyIdentifier = keyid:always,issuer -EOF - -cleanup() { - rm -f \ - "${outdir}/cert.csr" \ - "${outdir}/cert.cnf" \ - "${outdir}/ca.cnf" \ - "${outdir}/ca-cert.srl" -} - -trap cleanup EXIT - -cat <<'EOF' > "${outdir}/cert.cnf" -[req] -distinguished_name = dn -req_extensions = v3_req -prompt = no - -[dn] -CN = localhost - -[v3_req] -subjectAltName = @alt_names -basicConstraints = CA:false -keyUsage = critical, digitalSignature, keyEncipherment -extendedKeyUsage = serverAuth - -[alt_names] -DNS.1 = localhost -IP.1 = 127.0.0.1 -EOF - -openssl req \ - -x509 \ - -newkey rsa:2048 \ - -sha256 \ - -nodes \ - -keyout "${outdir}/ca-key.pem" \ - -out "${outdir}/ca-cert.pem" \ - -days 3650 \ - -config "${outdir}/ca.cnf" \ - >/dev/null 2>&1 - -openssl req \ - -newkey rsa:2048 \ - -sha256 \ - -nodes \ - -keyout "${outdir}/key.pem" \ - -out "${outdir}/cert.csr" \ - -config "${outdir}/cert.cnf" \ - >/dev/null 2>&1 - -openssl x509 \ - -req \ - -in "${outdir}/cert.csr" \ - -CA "${outdir}/ca-cert.pem" \ - -CAkey "${outdir}/ca-key.pem" \ - -CAcreateserial \ - -out "${outdir}/cert.pem" \ - -days 3650 \ - -sha256 \ - -extfile "${outdir}/cert.cnf" \ - -extensions v3_req \ - >/dev/null 2>&1 - -echo "Generated TLS test certs in ${outdir}" diff --git a/scripts/generate_test_tls_certs.py b/scripts/generate_test_tls_certs.py new file mode 100644 index 0000000..e8b3c43 --- /dev/null +++ b/scripts/generate_test_tls_certs.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import os +import pathlib +import shutil +import subprocess +import sys +import textwrap + + +CA_CONFIG = """ +[req] +distinguished_name = dn +x509_extensions = v3_ca +prompt = no + +[dn] +CN = rsloop-test-ca + +[v3_ca] +basicConstraints = critical, CA:true +keyUsage = critical, keyCertSign, cRLSign +subjectKeyIdentifier = hash +authorityKeyIdentifier = keyid:always,issuer +""" + +CERT_CONFIG = """ +[req] +distinguished_name = dn +req_extensions = v3_req +prompt = no + +[dn] +CN = localhost + +[v3_req] +subjectAltName = @alt_names +basicConstraints = CA:false +keyUsage = critical, digitalSignature, keyEncipherment +extendedKeyUsage = serverAuth + +[alt_names] +DNS.1 = localhost +IP.1 = 127.0.0.1 +""" + + +def find_openssl() -> pathlib.Path: + candidates: list[str | pathlib.Path | None] = [ + os.environ.get("OPENSSL"), + shutil.which("openssl"), + ] + + if os.name == "nt": + git = shutil.which("git") + if git: + git_root = pathlib.Path(git).resolve().parent.parent + candidates.extend( + [ + git_root / "usr" / "bin" / "openssl.exe", + git_root / "mingw64" / "bin" / "openssl.exe", + git_root / "mingw32" / "bin" / "openssl.exe", + ] + ) + program_files = pathlib.Path(os.environ.get("ProgramFiles", "C:/Program Files")) + candidates.extend( + [ + program_files / "Git" / "usr" / "bin" / "openssl.exe", + program_files / "Git" / "mingw64" / "bin" / "openssl.exe", + ] + ) + + for candidate in candidates: + if candidate and pathlib.Path(candidate).is_file(): + return pathlib.Path(candidate) + + raise RuntimeError( + "OpenSSL was not found. Install OpenSSL, install Git for Windows, or set " + "the OPENSSL environment variable to the executable path." + ) + + +def run(openssl: pathlib.Path, *args: str) -> None: + try: + subprocess.run( + [str(openssl), *args], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + ) + except subprocess.CalledProcessError as exc: + raise RuntimeError(f"OpenSSL failed: {exc.stderr.strip()}") from exc + + +def generate(outdir: pathlib.Path) -> None: + outdir.mkdir(parents=True, exist_ok=True) + generated = { + name: outdir / name + for name in ( + "ca-cert.pem", + "ca-key.pem", + "cert.pem", + "key.pem", + "cert.csr", + "cert.cnf", + "ca.cnf", + "ca-cert.srl", + ) + } + for path in generated.values(): + path.unlink(missing_ok=True) + + generated["ca.cnf"].write_text( + textwrap.dedent(CA_CONFIG).lstrip(), encoding="utf-8" + ) + generated["cert.cnf"].write_text( + textwrap.dedent(CERT_CONFIG).lstrip(), encoding="utf-8" + ) + + openssl = find_openssl() + try: + run( + openssl, + "req", + "-x509", + "-newkey", + "rsa:2048", + "-sha256", + "-nodes", + "-keyout", + str(generated["ca-key.pem"]), + "-out", + str(generated["ca-cert.pem"]), + "-days", + "3650", + "-config", + str(generated["ca.cnf"]), + ) + run( + openssl, + "req", + "-newkey", + "rsa:2048", + "-sha256", + "-nodes", + "-keyout", + str(generated["key.pem"]), + "-out", + str(generated["cert.csr"]), + "-config", + str(generated["cert.cnf"]), + ) + run( + openssl, + "x509", + "-req", + "-in", + str(generated["cert.csr"]), + "-CA", + str(generated["ca-cert.pem"]), + "-CAkey", + str(generated["ca-key.pem"]), + "-CAcreateserial", + "-out", + str(generated["cert.pem"]), + "-days", + "3650", + "-sha256", + "-extfile", + str(generated["cert.cnf"]), + "-extensions", + "v3_req", + ) + finally: + for name in ("cert.csr", "cert.cnf", "ca.cnf", "ca-cert.srl"): + generated[name].unlink(missing_ok=True) + + print(f"Generated TLS test certs in {outdir}") + + +if __name__ == "__main__": + output = ( + pathlib.Path(sys.argv[1]) + if len(sys.argv) > 1 + else pathlib.Path("tests/fixtures/tls") + ) + generate(output) diff --git a/scripts/test-supported-pythons.sh b/scripts/test-supported-pythons.sh index 9322681..222864c 100755 --- a/scripts/test-supported-pythons.sh +++ b/scripts/test-supported-pythons.sh @@ -66,7 +66,7 @@ cd "$ROOT_DIR" if (( GENERATE_TLS_CERTS )); then echo "Generating TLS test certificates" - "${ROOT_DIR}/scripts/generate-test-tls-certs.sh" + uv run --no-project python "${ROOT_DIR}/scripts/generate_test_tls_certs.py" fi if (( INSTALL_PYTHONS )); then diff --git a/tests/test_tls.py b/tests/test_tls.py index 0593443..399123a 100644 --- a/tests/test_tls.py +++ b/tests/test_tls.py @@ -4,9 +4,11 @@ import importlib.util import os import pathlib +import shutil import socket import ssl import subprocess +import sys import tempfile import unittest @@ -18,8 +20,9 @@ pathlib.Path(__file__) .resolve() .parents[1] - .joinpath("scripts", "generate-test-tls-certs.sh") + .joinpath("scripts", "generate_test_tls_certs.py") ) +UV = shutil.which("uv") or "uv" def ensure_tls_fixtures() -> None: @@ -29,7 +32,19 @@ def ensure_tls_fixtures() -> None: ): return - subprocess.run([str(TLS_GENERATOR), str(TLS_FIXTURES_DIR)], check=True) + subprocess.run( + [ + UV, + "run", + "--no-project", + "--python", + sys.executable, + "python", + str(TLS_GENERATOR), + str(TLS_FIXTURES_DIR), + ], + check=True, + ) ensure_tls_fixtures() From f121141cdfff864392aa59363d76689a7ac9a159 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 19 Jul 2026 14:39:36 -0700 Subject: [PATCH 03/21] simplify listener handling in ServerCore for non-linux platforms --- src/stream_transport.rs | 88 +++++++++++++---------------------------- 1 file changed, 28 insertions(+), 60 deletions(-) diff --git a/src/stream_transport.rs b/src/stream_transport.rs index d52bd6e..6d6123d 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -1963,69 +1963,43 @@ impl ServerCore { std::mem::take(&mut state.listeners) }; - #[cfg(target_os = "linux")] - { - if self.tls.is_some() { - let mut tasks = self.accept_tasks.lock().expect("poisoned accept tasks"); - for listener in listeners { - let server = Arc::clone(self); - let task = match listener { - ServerListener::Tcp(listener) => { - WorkerThread::spawn("rsloop-tcp-accept", move |stop| { - run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - }) - } - #[cfg(unix)] - ServerListener::Unix(listener) => { - WorkerThread::spawn("rsloop-unix-accept", move |stop| { - run_unix_accept_loop(BlockingAcceptLoop::new( - server, listener, stop, - )) - }) - } - }; - tasks.push(task); - } - return; - } - - let mut accept_fds = self.accept_fds.lock().expect("poisoned accept fds"); + if self.tls.is_some() { + let mut tasks = self.accept_tasks.lock().expect("poisoned accept tasks"); for listener in listeners { - let fd = match &listener { - ServerListener::Tcp(listener) => tcp_listener_raw_fd(listener), + let server = Arc::clone(self); + let task = match listener { + ServerListener::Tcp(listener) => { + WorkerThread::spawn("rsloop-tcp-accept", move |stop| { + run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) + }) + } #[cfg(unix)] - ServerListener::Unix(listener) => unix_raw_fd(listener.as_raw_fd()), + ServerListener::Unix(listener) => { + WorkerThread::spawn("rsloop-unix-accept", move |stop| { + run_unix_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) + }) + } }; - accept_fds.push(fd); - let _ = self.loop_core.send_command(LoopCommand::Io( - LoopIoCommand::StartServerAccept { - fd, - server: Arc::clone(self), - listener, - }, - )); + tasks.push(task); } + return; } - #[cfg(not(target_os = "linux"))] - let mut tasks = self.accept_tasks.lock().expect("poisoned accept tasks"); - #[cfg(not(target_os = "linux"))] + let mut accept_fds = self.accept_fds.lock().expect("poisoned accept fds"); for listener in listeners { - let server = Arc::clone(self); - let task = match listener { - ServerListener::Tcp(listener) => { - WorkerThread::spawn("rsloop-tcp-accept", move |stop| { - run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - }) - } + let fd = match &listener { + ServerListener::Tcp(listener) => tcp_listener_raw_fd(listener), #[cfg(unix)] - ServerListener::Unix(listener) => { - WorkerThread::spawn("rsloop-unix-accept", move |stop| { - run_unix_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - }) - } + ServerListener::Unix(listener) => unix_raw_fd(listener.as_raw_fd()), }; - tasks.push(task); + accept_fds.push(fd); + let _ = + self.loop_core + .send_command(LoopCommand::Io(LoopIoCommand::StartServerAccept { + fd, + server: Arc::clone(self), + listener, + })); } } } @@ -2819,7 +2793,6 @@ pub fn spawn_tcp_transport( server.connection_opened(); } - #[cfg(target_os = "linux")] core.loop_core .send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { fd: raw_fd, @@ -2827,8 +2800,6 @@ pub fn spawn_tcp_transport( reader: ReaderTarget::Tcp(stream), })) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; - #[cfg(not(target_os = "linux"))] - spawn_reader_worker(Arc::clone(&core), ReaderTarget::Tcp(stream)); Ok(transport) } @@ -2878,7 +2849,6 @@ pub fn spawn_unix_transport( server.connection_opened(); } - #[cfg(target_os = "linux")] core.loop_core .send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { fd: raw_fd, @@ -2886,8 +2856,6 @@ pub fn spawn_unix_transport( reader: ReaderTarget::Unix(stream), })) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; - #[cfg(not(target_os = "linux"))] - spawn_reader_worker(Arc::clone(&core), ReaderTarget::Unix(stream)); Ok(transport) } From 5b7a50b8ddc44494eef169b38a301a6c1c2f9511 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 19 Jul 2026 15:17:25 -0700 Subject: [PATCH 04/21] add workload matrix benchmark for production-shaped network traffic --- benchmarks/README.md | 70 +++ benchmarks/compare_event_loops.py | 4 +- benchmarks/workload_matrix.py | 767 ++++++++++++++++++++++++++++++ justfile | 6 + 4 files changed, 845 insertions(+), 2 deletions(-) create mode 100644 benchmarks/workload_matrix.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 55cc8b1..8b55f4e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -67,3 +67,73 @@ three loops to go through the stdlib `asyncio` streams layer instead, pass: ```bash uv run --with uvloop python benchmarks/compare_event_loops.py --no-rsloop-fast-streams ``` + +## Representative workload matrix + +[`workload_matrix.py`](./workload_matrix.py) complements the microbenchmarks +with concurrent, production-shaped network traffic. It runs each loop and +scenario in a fresh subprocess and reports throughput, transferred MiB/s, +p50/p95/p99 operation latency, and peak RSS. + +The scenarios are: + +- `http_keepalive`: concurrent HTTP/1.1 keep-alive clients, configurable + response bodies, and modest application CPU work +- `tls_http`: the same workload over TLS +- `websocket_messages`: concurrent persistent RFC 6455 connections with + masked client frames and mixed message sizes +- `mixed_streams`: concurrent connections cycling through 64 B to 64 KiB + messages +- `bulk_transfer`: concurrent large transfers with `drain()` backpressure +- `idle_connections`: many idle connections activated at the same time + +Build rsloop in release mode and run the standard matrix with: + +```bash +uv run --with maturin maturin develop --release +uv run --with uvloop python benchmarks/workload_matrix.py # Unix +uv run --with winloop python benchmarks/workload_matrix.py # Windows +``` + +The default comparison is `asyncio,uvloop,rsloop` on Unix. Because uvloop is +not available on Windows, the default there is `asyncio,winloop,rsloop`. +Unavailable optional loops are reported and skipped. + +For a quick smoke run: + +```bash +uv run --with uvloop python benchmarks/workload_matrix.py \ + --warmups 0 \ + --repeat 1 \ + --concurrency 4 \ + --requests-per-connection 5 \ + --bulk-bytes 262144 \ + --idle-connections 20 \ + --idle-seconds 0.01 +``` + +Use `--json-output benchmarks/results/matrix.json` to retain raw measurements. +TLS uses the test certificates under `tests/fixtures/tls`; regenerate them +cross-platform when needed with: + +```bash +uv run --no-project python scripts/generate_test_tls_certs.py tests/fixtures/tls +``` + +For Tracy, build with the profiler feature and request an unmeasured profiling +pass before each rsloop scenario: + +```bash +uv run --with maturin maturin develop --release --features profiler +uv run --with uvloop python benchmarks/workload_matrix.py \ + --loops rsloop \ + --profile-rsloop-dir benchmarks/profiles +``` + +On Windows, replace `--with uvloop` with `--with winloop` in the quick and +profiling commands. + +Treat this matrix as a regression and trade-off tool, not a single leaderboard. +Compare throughput together with tail latency and memory, and tune concurrency, +payload sizes, application work, and connection counts to match the target +deployment. diff --git a/benchmarks/compare_event_loops.py b/benchmarks/compare_event_loops.py index bf88243..2ec37da 100644 --- a/benchmarks/compare_event_loops.py +++ b/benchmarks/compare_event_loops.py @@ -25,9 +25,9 @@ def default_loops_csv() -> str: loops = ["asyncio", "rsloop"] if sys.platform == "win32": - loops.insert(2, "winloop") + loops.insert(1, "winloop") else: - loops.insert(2, "uvloop") + loops.insert(1, "uvloop") return ",".join(loops) diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py new file mode 100644 index 0000000..6a36a82 --- /dev/null +++ b/benchmarks/workload_matrix.py @@ -0,0 +1,767 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import json +import os +import ssl +import statistics +import struct +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Awaitable, Callable + +from compare_event_loops import ( + LOOP_CHOICES, + default_loops_csv, + format_bytes, + get_peak_rss_bytes, + is_loop_available, + loop_factory_for, + maybe_wait_closed, + normalize_csv, +) + + +SCENARIO_CHOICES = ( + "http_keepalive", + "tls_http", + "websocket_messages", + "mixed_streams", + "bulk_transfer", + "idle_connections", +) +DEFAULT_SCENARIOS = ",".join(SCENARIO_CHOICES) +ROOT = Path(__file__).resolve().parent.parent +TLS_DIR = ROOT / "tests" / "fixtures" / "tls" +HTTP_REQUEST = ( + b"GET /resource HTTP/1.1\r\nHost: localhost\r\nConnection: keep-alive\r\n\r\n" +) +WEBSOCKET_KEY = "dGhlIHNhbXBsZSBub25jZQ==" +WEBSOCKET_REQUEST = ( + "GET /socket HTTP/1.1\r\n" + "Host: localhost\r\n" + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" + f"Sec-WebSocket-Key: {WEBSOCKET_KEY}\r\n" + "Sec-WebSocket-Version: 13\r\n\r\n" +).encode("ascii") +WEBSOCKET_ACCEPT = base64.b64encode( + hashlib.sha1( + (WEBSOCKET_KEY + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii") + ).digest() +).decode("ascii") +WEBSOCKET_RESPONSE = ( + "HTTP/1.1 101 Switching Protocols\r\n" + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" + f"Sec-WebSocket-Accept: {WEBSOCKET_ACCEPT}\r\n\r\n" +).encode("ascii") + + +@dataclass(frozen=True) +class MatrixResult: + loop: str + scenario: str + seconds: float + operations: int + bytes_transferred: int + latency_ms: list[float] + peak_rss_bytes: int = 0 + + @property + def ops_per_sec(self) -> float: + return self.operations / self.seconds if self.seconds else float("inf") + + @property + def mib_per_sec(self) -> float: + return ( + self.bytes_transferred / (1024 * 1024) / self.seconds + if self.seconds + else 0.0 + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run production-shaped networking workloads against asyncio event loops.", + ) + parser.add_argument( + "--loops", + default=default_loops_csv(), + help=( + "Comma-separated loops. Defaults to asyncio, uvloop, and rsloop on " + "Unix; asyncio, winloop, and rsloop on Windows." + ), + ) + parser.add_argument("--scenarios", default=DEFAULT_SCENARIOS) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--repeat", type=int, default=3) + parser.add_argument("--concurrency", type=int, default=16) + parser.add_argument("--requests-per-connection", type=int, default=50) + parser.add_argument("--http-response-size", type=int, default=4096) + parser.add_argument( + "--app-work-iterations", + type=int, + default=8, + help="SHA-256 iterations per HTTP request to model modest application work.", + ) + parser.add_argument( + "--mixed-payload-sizes", + default="64,1024,16384,65536", + help="Comma-separated payload sizes cycled by mixed_streams clients.", + ) + parser.add_argument( + "--websocket-payload-sizes", + default="32,256,4096", + help="Comma-separated binary frame sizes cycled by WebSocket clients.", + ) + parser.add_argument("--bulk-bytes", type=int, default=2 * 1024 * 1024) + parser.add_argument("--bulk-chunk-size", type=int, default=64 * 1024) + parser.add_argument("--idle-connections", type=int, default=200) + parser.add_argument("--idle-seconds", type=float, default=0.2) + parser.add_argument("--tls-dir", type=Path, default=TLS_DIR) + parser.add_argument("--json-output", type=Path) + parser.add_argument( + "--profile-rsloop-dir", + type=Path, + help="Run one unmeasured Tracy pass per rsloop scenario before measurements.", + ) + parser.add_argument("--child", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--loop", choices=LOOP_CHOICES, help=argparse.SUPPRESS) + parser.add_argument("--scenario", choices=SCENARIO_CHOICES, help=argparse.SUPPRESS) + parser.add_argument("--profile-label", help=argparse.SUPPRESS) + return parser.parse_args() + + +def positive(value: int, name: str) -> None: + if value <= 0: + raise SystemExit(f"{name} must be > 0") + + +def validate_args(args: argparse.Namespace) -> None: + if args.warmups < 0: + raise SystemExit("--warmups must be >= 0") + positive(args.repeat, "--repeat") + positive(args.concurrency, "--concurrency") + positive(args.requests_per_connection, "--requests-per-connection") + positive(args.http_response_size, "--http-response-size") + if args.app_work_iterations < 0: + raise SystemExit("--app-work-iterations must be >= 0") + positive(args.bulk_bytes, "--bulk-bytes") + positive(args.bulk_chunk_size, "--bulk-chunk-size") + positive(args.idle_connections, "--idle-connections") + if args.idle_seconds < 0: + raise SystemExit("--idle-seconds must be >= 0") + for attribute, option in ( + ("mixed_payload_sizes", "--mixed-payload-sizes"), + ("websocket_payload_sizes", "--websocket-payload-sizes"), + ): + try: + sizes = [ + int(item.strip()) + for item in getattr(args, attribute).split(",") + if item.strip() + ] + except ValueError as exc: + raise SystemExit(f"{option} must contain integers") from exc + if not sizes or any(size <= 0 for size in sizes): + raise SystemExit(f"{option} values must be > 0") + setattr(args, attribute, sizes) + + +async def close_writer(writer: asyncio.StreamWriter) -> None: + writer.close() + await maybe_wait_closed(writer) + + +async def close_server(server: asyncio.AbstractServer) -> None: + server.close() + await server.wait_closed() + + +def application_digest(iterations: int) -> bytes: + digest = b"rsloop-real-world-workload" + for _ in range(iterations): + digest = hashlib.sha256(digest).digest() + return digest + + +def tls_contexts(tls_dir: Path) -> tuple[ssl.SSLContext, ssl.SSLContext]: + cert = tls_dir / "cert.pem" + key = tls_dir / "key.pem" + ca = tls_dir / "ca-cert.pem" + missing = [path for path in (cert, key, ca) if not path.is_file()] + if missing: + names = ", ".join(str(path) for path in missing) + raise RuntimeError( + f"missing TLS fixtures: {names}. Generate them with: " + "uv run --no-project python scripts/generate_test_tls_certs.py tests/fixtures/tls" + ) + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.load_cert_chain(cert, key) + client_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=ca) + return server_context, client_context + + +async def run_http( + loop_name: str, + args: argparse.Namespace, + *, + use_tls: bool, +) -> MatrixResult: + response_body = application_digest(args.app_work_iterations) or b"x" + response_body = ( + response_body * ((args.http_response_size // len(response_body)) + 1) + )[: args.http_response_size] + response_head = ( + "HTTP/1.1 200 OK\r\n" + f"Content-Length: {len(response_body)}\r\n" + "Content-Type: application/octet-stream\r\n" + "Connection: keep-alive\r\n\r\n" + ).encode("ascii") + + async def handle( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + while True: + try: + request = await reader.readexactly(len(HTTP_REQUEST)) + except asyncio.IncompleteReadError: + return + if request != HTTP_REQUEST: + raise RuntimeError("unexpected HTTP request") + application_digest(args.app_work_iterations) + writer.write(response_head) + writer.write(response_body) + await writer.drain() + finally: + await close_writer(writer) + + server_ssl = client_ssl = None + if use_tls: + server_ssl, client_ssl = tls_contexts(args.tls_dir) + server = await asyncio.start_server(handle, "127.0.0.1", 0, ssl=server_ssl) + host, port = server.sockets[0].getsockname()[:2] + latencies: list[float] = [] + + async def client() -> int: + reader, writer = await asyncio.open_connection( + host, + port, + ssl=client_ssl, + server_hostname="localhost" if use_tls else None, + ) + transferred = 0 + try: + for _ in range(args.requests_per_connection): + started = time.perf_counter() + writer.write(HTTP_REQUEST) + await writer.drain() + await reader.readexactly(len(response_head)) + await reader.readexactly(len(response_body)) + latencies.append((time.perf_counter() - started) * 1000) + transferred += ( + len(HTTP_REQUEST) + len(response_head) + len(response_body) + ) + finally: + await close_writer(writer) + return transferred + + started = time.perf_counter() + try: + transferred = sum( + await asyncio.gather(*(client() for _ in range(args.concurrency))) + ) + finally: + await close_server(server) + scenario = "tls_http" if use_tls else "http_keepalive" + return MatrixResult( + loop_name, + scenario, + time.perf_counter() - started, + args.concurrency * args.requests_per_connection, + transferred, + latencies, + ) + + +def websocket_frame(payload: bytes, mask_key: bytes | None = None) -> bytes: + length = len(payload) + masked = mask_key is not None + if length < 126: + header = bytes((0x82, length | (0x80 if masked else 0))) + elif length <= 0xFFFF: + header = bytes((0x82, 126 | (0x80 if masked else 0))) + struct.pack( + "!H", length + ) + else: + header = bytes((0x82, 127 | (0x80 if masked else 0))) + struct.pack( + "!Q", length + ) + if mask_key is None: + return header + payload + masked_payload = bytes( + value ^ mask_key[index % 4] for index, value in enumerate(payload) + ) + return header + mask_key + masked_payload + + +async def read_websocket_frame(reader: asyncio.StreamReader) -> bytes: + first, second = await reader.readexactly(2) + if first != 0x82: + raise RuntimeError(f"unexpected WebSocket frame type: {first:#x}") + length = second & 0x7F + if length == 126: + length = struct.unpack("!H", await reader.readexactly(2))[0] + elif length == 127: + length = struct.unpack("!Q", await reader.readexactly(8))[0] + mask_key = await reader.readexactly(4) if second & 0x80 else None + payload = await reader.readexactly(length) + if mask_key is not None: + payload = bytes( + value ^ mask_key[index % 4] for index, value in enumerate(payload) + ) + return payload + + +async def run_websocket_messages( + loop_name: str, args: argparse.Namespace +) -> MatrixResult: + async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + request = await reader.readexactly(len(WEBSOCKET_REQUEST)) + if request != WEBSOCKET_REQUEST: + raise RuntimeError("unexpected WebSocket handshake") + writer.write(WEBSOCKET_RESPONSE) + await writer.drain() + while True: + try: + payload = await read_websocket_frame(reader) + except asyncio.IncompleteReadError: + return + writer.write(websocket_frame(payload)) + await writer.drain() + finally: + await close_writer(writer) + + server = await asyncio.start_server(echo, "127.0.0.1", 0) + host, port = server.sockets[0].getsockname()[:2] + latencies: list[float] = [] + + async def client(client_id: int) -> int: + reader, writer = await asyncio.open_connection(host, port) + transferred = len(WEBSOCKET_REQUEST) + len(WEBSOCKET_RESPONSE) + try: + writer.write(WEBSOCKET_REQUEST) + await writer.drain() + response = await reader.readexactly(len(WEBSOCKET_RESPONSE)) + if response != WEBSOCKET_RESPONSE: + raise RuntimeError("WebSocket upgrade failed") + for index in range(args.requests_per_connection): + size = args.websocket_payload_sizes[ + (client_id + index) % len(args.websocket_payload_sizes) + ] + payload = bytes([(client_id + index) % 251]) * size + mask_key = struct.pack("!I", (client_id << 16) ^ index ^ 0xA5A55A5A) + outbound = websocket_frame(payload, mask_key) + started = time.perf_counter() + writer.write(outbound) + await writer.drain() + response_payload = await read_websocket_frame(reader) + if response_payload != payload: + raise RuntimeError("WebSocket echo mismatch") + latencies.append((time.perf_counter() - started) * 1000) + transferred += len(outbound) + len(websocket_frame(payload)) + finally: + await close_writer(writer) + return transferred + + started = time.perf_counter() + try: + transferred = sum( + await asyncio.gather(*(client(index) for index in range(args.concurrency))) + ) + finally: + await close_server(server) + return MatrixResult( + loop_name, + "websocket_messages", + time.perf_counter() - started, + args.concurrency * args.requests_per_connection, + transferred, + latencies, + ) + + +async def run_mixed_streams(loop_name: str, args: argparse.Namespace) -> MatrixResult: + async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + while data := await reader.read(64 * 1024): + writer.write(data) + await writer.drain() + finally: + await close_writer(writer) + + server = await asyncio.start_server(echo, "127.0.0.1", 0) + host, port = server.sockets[0].getsockname()[:2] + latencies: list[float] = [] + + async def client(client_id: int) -> int: + reader, writer = await asyncio.open_connection(host, port) + transferred = 0 + try: + for index in range(args.requests_per_connection): + size = args.mixed_payload_sizes[ + (client_id + index) % len(args.mixed_payload_sizes) + ] + payload = bytes([(client_id + index) % 251]) * size + started = time.perf_counter() + writer.write(payload) + await writer.drain() + response = await reader.readexactly(size) + if response != payload: + raise RuntimeError("mixed stream echo mismatch") + latencies.append((time.perf_counter() - started) * 1000) + transferred += size * 2 + finally: + await close_writer(writer) + return transferred + + started = time.perf_counter() + try: + transferred = sum( + await asyncio.gather(*(client(index) for index in range(args.concurrency))) + ) + finally: + await close_server(server) + return MatrixResult( + loop_name, + "mixed_streams", + time.perf_counter() - started, + args.concurrency * args.requests_per_connection, + transferred, + latencies, + ) + + +async def run_bulk_transfer(loop_name: str, args: argparse.Namespace) -> MatrixResult: + chunk = b"b" * min(args.bulk_chunk_size, args.bulk_bytes) + + async def send_bulk( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + await reader.readexactly(1) + remaining = args.bulk_bytes + while remaining: + data = chunk[:remaining] + writer.write(data) + await writer.drain() + remaining -= len(data) + finally: + await close_writer(writer) + + server = await asyncio.start_server(send_bulk, "127.0.0.1", 0) + host, port = server.sockets[0].getsockname()[:2] + latencies: list[float] = [] + + async def client() -> int: + reader, writer = await asyncio.open_connection(host, port) + started = time.perf_counter() + try: + writer.write(b"!") + await writer.drain() + await reader.readexactly(args.bulk_bytes) + latencies.append((time.perf_counter() - started) * 1000) + finally: + await close_writer(writer) + return args.bulk_bytes + 1 + + started = time.perf_counter() + try: + transferred = sum( + await asyncio.gather(*(client() for _ in range(args.concurrency))) + ) + finally: + await close_server(server) + return MatrixResult( + loop_name, + "bulk_transfer", + time.perf_counter() - started, + args.concurrency, + transferred, + latencies, + ) + + +async def run_idle_connections( + loop_name: str, args: argparse.Namespace +) -> MatrixResult: + async def ping(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + while data := await reader.read(1): + writer.write(data) + await writer.drain() + finally: + await close_writer(writer) + + server = await asyncio.start_server( + ping, + "127.0.0.1", + 0, + backlog=max(100, args.idle_connections), + ) + host, port = server.sockets[0].getsockname()[:2] + connections = await asyncio.gather( + *(asyncio.open_connection(host, port) for _ in range(args.idle_connections)) + ) + await asyncio.sleep(args.idle_seconds) + latencies: list[float] = [] + + async def activate( + connection: tuple[asyncio.StreamReader, asyncio.StreamWriter], + ) -> int: + reader, writer = connection + started = time.perf_counter() + writer.write(b"p") + await writer.drain() + if await reader.readexactly(1) != b"p": + raise RuntimeError("idle connection ping mismatch") + latencies.append((time.perf_counter() - started) * 1000) + return 2 + + started = time.perf_counter() + try: + transferred = sum( + await asyncio.gather(*(activate(item) for item in connections)) + ) + finally: + await asyncio.gather(*(close_writer(writer) for _, writer in connections)) + await close_server(server) + return MatrixResult( + loop_name, + "idle_connections", + time.perf_counter() - started, + args.idle_connections, + transferred, + latencies, + ) + + +SCENARIO_RUNNERS: dict[ + str, Callable[[str, argparse.Namespace], Awaitable[MatrixResult]] +] = { + "http_keepalive": lambda loop, args: run_http(loop, args, use_tls=False), + "tls_http": lambda loop, args: run_http(loop, args, use_tls=True), + "websocket_messages": run_websocket_messages, + "mixed_streams": run_mixed_streams, + "bulk_transfer": run_bulk_transfer, + "idle_connections": run_idle_connections, +} + + +def run_with_loop(loop_name: str, awaitable: Awaitable[MatrixResult]) -> MatrixResult: + factory = loop_factory_for(loop_name) + if sys.version_info[:2] >= (3, 12): + return asyncio.run(awaitable, loop_factory=factory) + loop = factory() + try: + asyncio.set_event_loop(loop) + return loop.run_until_complete(awaitable) + finally: + asyncio.set_event_loop(None) + loop.close() + + +def child_main(args: argparse.Namespace) -> int: + awaitable = SCENARIO_RUNNERS[args.scenario](args.loop, args) + if args.profile_label: + if args.loop != "rsloop": + raise RuntimeError("Tracy profiling is only supported for rsloop") + print(f"[profile] Tracy session label: {args.profile_label}", flush=True) + import rsloop + + with rsloop.profile(): + result = run_with_loop(args.loop, awaitable) + else: + result = run_with_loop(args.loop, awaitable) + result = MatrixResult(**{**asdict(result), "peak_rss_bytes": get_peak_rss_bytes()}) + print(json.dumps(asdict(result))) + return 0 + + +def child_command( + args: argparse.Namespace, + loop_name: str, + scenario: str, + profile_label: str | None = None, +) -> list[str]: + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + "--child", + "--loop", + loop_name, + "--scenario", + scenario, + "--concurrency", + str(args.concurrency), + "--requests-per-connection", + str(args.requests_per_connection), + "--http-response-size", + str(args.http_response_size), + "--app-work-iterations", + str(args.app_work_iterations), + "--mixed-payload-sizes", + ",".join(str(size) for size in args.mixed_payload_sizes), + "--websocket-payload-sizes", + ",".join(str(size) for size in args.websocket_payload_sizes), + "--bulk-bytes", + str(args.bulk_bytes), + "--bulk-chunk-size", + str(args.bulk_chunk_size), + "--idle-connections", + str(args.idle_connections), + "--idle-seconds", + str(args.idle_seconds), + "--tls-dir", + str(args.tls_dir), + ] + if profile_label: + cmd.extend(("--profile-label", profile_label)) + return cmd + + +def run_child( + args: argparse.Namespace, + loop_name: str, + scenario: str, + profile_label: str | None = None, +) -> MatrixResult: + env = os.environ.copy() + if loop_name == "rsloop": + env["RSLOOP_USE_FAST_STREAMS"] = "1" + proc = subprocess.run( + child_command(args, loop_name, scenario, profile_label), + cwd=ROOT, + env=env, + capture_output=True, + text=True, + check=False, + ) + if proc.returncode: + raise RuntimeError( + f"{loop_name}/{scenario} failed with exit code {proc.returncode}\n" + f"stdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + ) + lines = [line for line in proc.stdout.splitlines() if line.strip()] + if not lines: + raise RuntimeError(f"{loop_name}/{scenario} produced no output") + return MatrixResult(**json.loads(lines[-1])) + + +def percentile(values: list[float], fraction: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + index = max(0, min(len(ordered) - 1, int(len(ordered) * fraction + 0.999999) - 1)) + return ordered[index] + + +def summarize(scenario: str, runs: dict[str, list[MatrixResult]]) -> None: + rows = [] + for loop_name, measured in runs.items(): + median_seconds = statistics.median(item.seconds for item in measured) + representative = min( + measured, key=lambda item: abs(item.seconds - median_seconds) + ) + rows.append( + ( + loop_name, + median_seconds, + representative.operations / median_seconds, + representative.bytes_transferred / (1024 * 1024) / median_seconds, + percentile(representative.latency_ms, 0.50), + percentile(representative.latency_ms, 0.95), + percentile(representative.latency_ms, 0.99), + int(statistics.median(item.peak_rss_bytes for item in measured)), + ) + ) + rows.sort(key=lambda row: row[1]) + print(f"\n{scenario}") + print( + f"{'loop':<10} {'median_s':>10} {'ops/s':>12} {'MiB/s':>10} " + f"{'p50_ms':>10} {'p95_ms':>10} {'p99_ms':>10} {'peak_rss':>12}" + ) + for row in rows: + print( + f"{row[0]:<10} {row[1]:>10.4f} {row[2]:>12,.0f} {row[3]:>10.1f} " + f"{row[4]:>10.3f} {row[5]:>10.3f} {row[6]:>10.3f} " + f"{format_bytes(row[7]):>12}" + ) + + +def parent_main(args: argparse.Namespace) -> int: + loops = normalize_csv(args.loops, allowed=LOOP_CHOICES, label="loops") + scenarios = normalize_csv( + args.scenarios, allowed=SCENARIO_CHOICES, label="scenarios" + ) + available = [] + for loop_name in loops: + ok, reason = is_loop_available(loop_name) + if ok: + available.append(loop_name) + else: + print(f"Skipping {loop_name}: {reason}") + if not available: + raise SystemExit("no benchmarkable loops are available") + + output: list[dict[str, object]] = [] + for scenario in scenarios: + scenario_runs: dict[str, list[MatrixResult]] = {} + for loop_name in available: + print(f"Running {scenario} on {loop_name}...") + if args.profile_rsloop_dir and loop_name == "rsloop": + args.profile_rsloop_dir.mkdir(parents=True, exist_ok=True) + label = str(args.profile_rsloop_dir / f"rsloop-{scenario}") + run_child(args, loop_name, scenario, label) + for _ in range(args.warmups): + run_child(args, loop_name, scenario) + measured = [ + run_child(args, loop_name, scenario) for _ in range(args.repeat) + ] + scenario_runs[loop_name] = measured + output.append( + { + "scenario": scenario, + "loop": loop_name, + "runs": [asdict(item) for item in measured], + } + ) + summarize(scenario, scenario_runs) + + if args.json_output: + args.json_output.parent.mkdir(parents=True, exist_ok=True) + args.json_output.write_text(json.dumps(output, indent=2), encoding="utf-8") + print(f"\nWrote raw results to {args.json_output}") + return 0 + + +def main() -> int: + args = parse_args() + validate_args(args) + if args.child: + return child_main(args) + return parent_main(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/justfile b/justfile index 3984386..1494fbc 100644 --- a/justfile +++ b/justfile @@ -1,4 +1,7 @@ set shell := ["bash", "-euo", "pipefail", "-c"] +set windows-shell := ["powershell.exe", "-NoLogo", "-NoProfile", "-Command"] + +benchmark-backend := if os() == "windows" { "winloop" } else { "uvloop" } tls-test-certs outdir="tests/fixtures/tls": uv run --no-project python scripts/generate_test_tls_certs.py {{outdir}} @@ -21,3 +24,6 @@ test-frameworks: uv run --with falcon --with uvicorn python tests/packages/falcon_test.py uv run --with quart --with hypercorn python tests/packages/quart_test.py uv run --with 'faststream[nats]' python tests/packages/faststream_test.py + +bench-real-world: + uv run --with {{benchmark-backend}} python benchmarks/workload_matrix.py From afbaf712b92f991494280eb9bb10a95ada964f07 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 19 Jul 2026 16:18:27 -0700 Subject: [PATCH 05/21] implement stream transport write handling and flush logic --- src/loop_core.rs | 18 +++++ src/loop_core/commands.rs | 2 + src/runtime.rs | 5 ++ src/stream_transport.rs | 144 +++++++++++++++++++++++++++++++++++++- tests/test_compat.py | 28 ++++++++ 5 files changed, 194 insertions(+), 3 deletions(-) diff --git a/src/loop_core.rs b/src/loop_core.rs index d5a801a..3c66c33 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -463,6 +463,10 @@ impl LoopCore { profiling::scope!("ready.stream_transport_read"); core.drain_pending_read_events_with_py(py)?; } + ReadyItem::StreamTransportWrite(core) => { + profiling::scope!("ready.stream_transport_write"); + core.flush_pending_direct_write(); + } ReadyItem::ProcessTransport(core) => { profiling::scope!("ready.process_transport"); core.drain_pending_events_with_py(py)?; @@ -792,6 +796,9 @@ impl LoopCore { ReadyItem::StreamTransportRead(core) => { LoopCommand::Transport(LoopTransportCommand::StreamRead(core)) } + ReadyItem::StreamTransportWrite(core) => { + LoopCommand::Transport(LoopTransportCommand::StreamWrite(core)) + } ReadyItem::ProcessTransport(core) => { LoopCommand::Transport(LoopTransportCommand::Process(core)) } @@ -842,6 +849,17 @@ impl LoopCore { unreachable!("local stream read enqueue preserves item kind") } }), + LoopCommand::Transport(LoopTransportCommand::StreamWrite(core)) => self + .try_enqueue_local_ready(ReadyItem::StreamTransportWrite(core)) + .or_else(|item| self.try_enqueue_active_ready(item)) + .map_err(|item| match item { + ReadyItem::StreamTransportWrite(core) => { + LoopCommand::Transport(LoopTransportCommand::StreamWrite(core)) + } + _ => { + unreachable!("local stream write enqueue preserves item kind") + } + }), LoopCommand::Transport(LoopTransportCommand::Process(core)) => self .try_enqueue_local_ready(ReadyItem::ProcessTransport(core)) .or_else(|item| self.try_enqueue_active_ready(item)) diff --git a/src/loop_core/commands.rs b/src/loop_core/commands.rs index 3304a52..d5052a2 100644 --- a/src/loop_core/commands.rs +++ b/src/loop_core/commands.rs @@ -23,6 +23,7 @@ pub enum ReadyItem { value: Py, }, StreamTransportRead(Arc), + StreamTransportWrite(Arc), ProcessTransport(Arc), ServerAccepted { server: Arc, @@ -100,6 +101,7 @@ pub enum LoopFutureCommand { pub enum LoopTransportCommand { StreamRead(Arc), + StreamWrite(Arc), Process(Arc), ServerAccepted { server: Arc, diff --git a/src/runtime.rs b/src/runtime.rs index 4fa4790..b0503fc 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -243,6 +243,11 @@ impl RuntimeDispatcher { self.ready_batch .push_back(ReadyItem::StreamTransportRead(core)); } + LoopCommand::Transport(LoopTransportCommand::StreamWrite(core)) => { + profiling::scope!("runtime.cmd.stream_transport_write"); + self.ready_batch + .push_back(ReadyItem::StreamTransportWrite(core)); + } LoopCommand::Transport(LoopTransportCommand::Process(core)) => { profiling::scope!("runtime.cmd.process_transport"); self.ready_batch diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 6d6123d..39c1f44 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -59,6 +59,11 @@ enum PendingReadEvent { const DEFAULT_WRITE_BUFFER_HIGH_WATER: usize = 64 * 1024; const DEFAULT_WRITE_BUFFER_LOW_WATER: usize = DEFAULT_WRITE_BUFFER_HIGH_WATER / 4; const MAX_PENDING_READ_COALESCE_BYTES: usize = 256 * 1024; +// Servers commonly emit a small protocol header followed immediately by a +// body. Defer only that header-sized first write for one loop turn so adjacent +// writes share a syscall; keep tiny control frames and larger payloads direct. +const SMALL_WRITE_COALESCE_MIN_BYTES: usize = 64; +const SMALL_WRITE_COALESCE_MAX_BYTES: usize = 512; const BLOCKING_POLL_INTERVAL_MS: i32 = 50; const STREAM_READ_BUFFER_SIZE: usize = 64 * 1024; @@ -126,6 +131,14 @@ impl OwnedWriteBuffer { } } + #[inline] + fn from_vec(data: Vec) -> Self { + Self { + bytes: data.into_boxed_slice(), + offset: 0, + } + } + #[inline] fn remaining(&self) -> &[u8] { &self.bytes[self.offset..] @@ -521,8 +534,11 @@ pub struct StreamTransportCore { state: Mutex, pending_read_events: Mutex>, read_events_scheduled: AtomicBool, + reading: AtomicBool, writer_tx: Sender, direct_writer: Option>, + pending_direct_write: Mutex>, + direct_write_scheduled: AtomicBool, lazy_writer: Mutex>, workers: Mutex>, // Signaled whenever `read_paused` clears or the transport starts @@ -532,6 +548,7 @@ pub struct StreamTransportCore { // The extra map is fixed at construction; cache the text-mode marker so // the per-write hot path avoids a state lock plus hash lookup. has_text_encoding: bool, + server_side: bool, } struct ServerState { @@ -1156,6 +1173,7 @@ impl StreamTransportCore { pending_data: &mut Option, fast_path: Option<&StreamReaderFastPath>, ) -> PyResult<()> { + profiling::scope!("StreamTransportCore::flush_pending_data_with_py"); let Some(data) = pending_data.take() else { return Ok(()); }; @@ -1447,6 +1465,7 @@ impl StreamTransportCore { state.reading = false; state.writable = false; drop(state); + self.reading.store(false, Ordering::Release); self.state_cv.notify_all(); } @@ -1477,6 +1496,7 @@ impl StreamTransportCore { let mut state = self.state.lock().expect("poisoned transport state"); state.read_paused = true; state.reading = false; + self.reading.store(false, Ordering::Release); } fn resume_reading(&self) { @@ -1484,11 +1504,12 @@ impl StreamTransportCore { state.read_paused = false; state.reading = true; drop(state); + self.reading.store(true, Ordering::Release); self.state_cv.notify_all(); } fn is_reading(&self) -> bool { - self.state.lock().expect("poisoned transport state").reading + self.reading.load(Ordering::Acquire) } fn wait_until_readable(&self) { @@ -1596,6 +1617,7 @@ impl StreamTransportCore { } fn try_direct_tasked_write(&self, data: &[u8]) -> io::Result { + profiling::scope!("StreamTransportCore::try_direct_tasked_write"); let Some(writer) = &self.direct_writer else { return Err(io::Error::other("not direct-tasked")); }; @@ -1606,7 +1628,9 @@ impl StreamTransportCore { // A paused reader does not consume Winsock's asynchronous reset // notification. Surface it before another small write can appear // to succeed from the local send buffer. - if let Some(err) = stream.take_error()? { + if !self.is_reading() + && let Some(err) = stream.take_error()? + { return Err(err); } @@ -1629,6 +1653,11 @@ impl StreamTransportCore { )); self.set_closing(); self.set_write_backpressure_active(false); + self.pending_direct_write + .lock() + .expect("poisoned pending direct write") + .clear(); + self.direct_write_scheduled.store(false, Ordering::Release); self.clear_write_buffer(false); let _ = self.writer_tx.send(WriterCommand::Stop); } @@ -1646,8 +1675,107 @@ impl StreamTransportCore { Ok(()) } + fn queue_recorded_write(self: &Arc, data: OwnedWriteBuffer) { + self.ensure_writer_worker(); + if self.writer_tx.send(WriterCommand::Data(data)).is_err() { + self.clear_write_buffer(false); + self.fail_write(None); + } + } + + fn stage_direct_write(self: &Arc, data: &[u8]) -> io::Result<()> { + if data.is_empty() { + return Ok(()); + } + + let should_pause = self.record_write_buffer_enqueued(data.len()); + self.pending_direct_write + .lock() + .expect("poisoned pending direct write") + .extend_from_slice(data); + if should_pause { + self.notify_pause_writing(); + } + + if !self.direct_write_scheduled.swap(true, Ordering::AcqRel) + && self + .loop_core + .send_command(LoopCommand::Transport(LoopTransportCommand::StreamWrite( + Arc::clone(self), + ))) + .is_err() + { + self.direct_write_scheduled.store(false, Ordering::Release); + self.fail_write(None); + } + Ok(()) + } + + pub(crate) fn flush_pending_direct_write(self: &Arc) { + profiling::scope!("StreamTransportCore::flush_pending_direct_write"); + self.direct_write_scheduled.store(false, Ordering::Release); + let data = std::mem::take( + self.pending_direct_write + .lock() + .expect("poisoned pending direct write") + .deref_mut(), + ); + if data.is_empty() { + return; + } + if self.is_closing() { + self.record_write_buffer_drained(data.len()); + return; + } + + match self.try_direct_tasked_write(&data) { + Ok(written) if written == data.len() => { + self.record_write_buffer_drained(written); + } + Ok(written) => { + self.record_write_buffer_drained(written); + let mut pending = OwnedWriteBuffer::from_vec(data); + pending.advance(written); + self.set_write_backpressure_active(true); + self.queue_recorded_write(pending); + } + Err(err) + if matches!( + err.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted + ) => + { + self.set_write_backpressure_active(true); + self.queue_recorded_write(OwnedWriteBuffer::from_vec(data)); + } + Err(err) => self.fail_write(Some(err)), + } + } + + fn discard_pending_direct_write(self: &Arc) { + self.direct_write_scheduled.store(false, Ordering::Release); + let discarded = { + let mut pending = self + .pending_direct_write + .lock() + .expect("poisoned pending direct write"); + let len = pending.len(); + pending.clear(); + len + }; + self.record_write_buffer_drained(discarded); + } + fn try_write_bytes(self: &Arc, data: &[u8]) -> io::Result<()> { + profiling::scope!("StreamTransportCore::try_write_bytes"); if self.direct_writer.is_some() && !self.write_backpressure_active() { + if self.direct_write_scheduled.load(Ordering::Acquire) + || (self.server_side + && data.len() > SMALL_WRITE_COALESCE_MIN_BYTES + && data.len() <= SMALL_WRITE_COALESCE_MAX_BYTES) + { + return self.stage_direct_write(data); + } match self.try_direct_tasked_write(data) { Ok(written) if written == data.len() => return Ok(()), Ok(written) => { @@ -1695,6 +1823,7 @@ impl StreamTransportCore { self: &Arc, py: Python<'_>, ) -> PyResult<(TransportSpawnContext, StreamKind)> { + self.flush_pending_direct_write(); let protocol = self.get_protocol(py); let context = self .state @@ -2006,6 +2135,7 @@ impl ServerCore { impl PyStreamTransport { pub(crate) fn write_data(&self, py: Python<'_>, data: &Bound<'_, PyAny>) -> PyResult<()> { + profiling::scope!("PyStreamTransport::write_data"); if self.core.is_closing() { return Ok(()); } @@ -2057,6 +2187,7 @@ impl PyStreamTransport { } fn close(&self) -> PyResult<()> { + self.core.flush_pending_direct_write(); self.core.set_closing(); if let Some(fd) = self.core.runtime_socket_fd() { let _ = self @@ -2083,6 +2214,7 @@ impl PyStreamTransport { } fn abort(&self) -> PyResult<()> { + self.core.discard_pending_direct_write(); self.core.set_closing(); if let Some(fd) = self.core.runtime_socket_fd() { let _ = self @@ -2117,6 +2249,7 @@ impl PyStreamTransport { "transport does not support write_eof", )); } + self.core.flush_pending_direct_write(); self.core.mark_write_eof(); if self.core.direct_writer.is_some() && !self.core.write_backpressure_active() { if let Some(writer) = &self.core.direct_writer { @@ -2596,18 +2729,24 @@ fn new_stream_transport_core( lazy_writer: Option, ) -> Arc { let has_text_encoding = parts.state.extra.contains_key("text_encoding"); + let server_side = parts.state.server.is_some(); + let reading = parts.state.reading; Arc::new(StreamTransportCore { loop_core: parts.loop_core, loop_obj: parts.loop_obj, state: Mutex::new(parts.state), pending_read_events: Mutex::new(VecDeque::new()), read_events_scheduled: AtomicBool::new(false), + reading: AtomicBool::new(reading), writer_tx, direct_writer: direct_writer.map(Mutex::new), + pending_direct_write: Mutex::new(Vec::new()), + direct_write_scheduled: AtomicBool::new(false), lazy_writer: Mutex::new(lazy_writer), workers: Mutex::new(Vec::new()), state_cv: Condvar::new(), has_text_encoding, + server_side, }) } @@ -3966,7 +4105,6 @@ fn write_stream_data_batch( } } } - if pending_command.is_none() { core.set_write_backpressure_active(false); } diff --git a/tests/test_compat.py b/tests/test_compat.py index ca06076..7c55a32 100644 --- a/tests/test_compat.py +++ b/tests/test_compat.py @@ -223,6 +223,34 @@ def connection_lost(self, exc): self.assertEqual(received, b"response-before-eof") self.assertEqual(events, ["data", "eof", "lost"]) + def test_close_flushes_coalesced_server_writes(self) -> None: + first = b"a" * 128 + second = b"b" * 128 + + async def main() -> bytes: + async def handle( + _reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + writer.write(first) + writer.write(second) + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(handle, "127.0.0.1", 0) + try: + host, port = server.sockets[0].getsockname()[:2] + reader, writer = await asyncio.open_connection(host, port) + try: + return await reader.read() + finally: + writer.close() + await writer.wait_closed() + finally: + server.close() + await server.wait_closed() + + self.assertEqual(rsloop.run(main()), first + second) + @unittest.skipUnless(sys.platform == "win32", "requires Winsock") def test_write_reports_reset_while_reading_is_paused(self) -> None: async def main() -> BaseException | None: From 09ae8fd6ccb3e83f5cb6c83371ed7efa2f2f73df Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 08:48:34 +0000 Subject: [PATCH 06/21] enhance profiling support and buffer management in stream transport --- benchmarks/workload_matrix.py | 90 ++++++++++++++++++------ python/rsloop/__init__.py | 2 + python/rsloop/_profile.py | 6 ++ src/fast_streams.rs | 25 ++++++- src/lib.rs | 2 +- src/module_init.rs | 5 +- src/profiler.rs | 7 ++ src/stream_transport.rs | 127 +++++++++++++++++++++++++--------- 8 files changed, 204 insertions(+), 60 deletions(-) diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py index 6a36a82..2ee8083 100644 --- a/benchmarks/workload_matrix.py +++ b/benchmarks/workload_matrix.py @@ -74,11 +74,19 @@ class MatrixResult: bytes_transferred: int latency_ms: list[float] peak_rss_bytes: int = 0 + connection_setup_seconds: float = 0.0 + traffic_seconds: float = 0.0 + teardown_seconds: float = 0.0 @property def ops_per_sec(self) -> float: return self.operations / self.seconds if self.seconds else float("inf") + @property + def traffic_ops_per_sec(self) -> float: + denominator = self.traffic_seconds or self.seconds + return self.operations / denominator if denominator else float("inf") + @property def mib_per_sec(self) -> float: return ( @@ -133,6 +141,11 @@ def parse_args() -> argparse.Namespace: type=Path, help="Run one unmeasured Tracy pass per rsloop scenario before measurements.", ) + parser.add_argument( + "--allow-profiler-build", + action="store_true", + help="Allow measured rsloop runs from a Tracy-enabled build.", + ) parser.add_argument("--child", action="store_true", help=argparse.SUPPRESS) parser.add_argument("--loop", choices=LOOP_CHOICES, help=argparse.SUPPRESS) parser.add_argument("--scenario", choices=SCENARIO_CHOICES, help=argparse.SUPPRESS) @@ -252,44 +265,58 @@ async def handle( host, port = server.sockets[0].getsockname()[:2] latencies: list[float] = [] - async def client() -> int: - reader, writer = await asyncio.open_connection( + async def open_client() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + connection = await asyncio.open_connection( host, port, ssl=client_ssl, server_hostname="localhost" if use_tls else None, ) + clients.append(connection) + return connection + + async def client( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> int: transferred = 0 - try: - for _ in range(args.requests_per_connection): - started = time.perf_counter() - writer.write(HTTP_REQUEST) - await writer.drain() - await reader.readexactly(len(response_head)) - await reader.readexactly(len(response_body)) - latencies.append((time.perf_counter() - started) * 1000) - transferred += ( - len(HTTP_REQUEST) + len(response_head) + len(response_body) - ) - finally: - await close_writer(writer) + for _ in range(args.requests_per_connection): + started = time.perf_counter() + writer.write(HTTP_REQUEST) + await writer.drain() + await reader.readexactly(len(response_head)) + await reader.readexactly(len(response_body)) + latencies.append((time.perf_counter() - started) * 1000) + transferred += len(HTTP_REQUEST) + len(response_head) + len(response_body) return transferred started = time.perf_counter() + clients: list[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = [] + setup_finished = traffic_finished = started try: - transferred = sum( - await asyncio.gather(*(client() for _ in range(args.concurrency))) + clients = list( + await asyncio.gather(*(open_client() for _ in range(args.concurrency))) ) + setup_finished = time.perf_counter() + transferred = sum(await asyncio.gather(*(client(*pair) for pair in clients))) + traffic_finished = time.perf_counter() finally: + await asyncio.gather( + *(close_writer(writer) for _, writer in clients), + return_exceptions=True, + ) await close_server(server) + finished = time.perf_counter() scenario = "tls_http" if use_tls else "http_keepalive" return MatrixResult( loop_name, scenario, - time.perf_counter() - started, + finished - started, args.concurrency * args.requests_per_connection, transferred, latencies, + connection_setup_seconds=setup_finished - started, + traffic_seconds=traffic_finished - setup_finished, + teardown_seconds=finished - traffic_finished, ) @@ -582,6 +609,14 @@ def run_with_loop(loop_name: str, awaitable: Awaitable[MatrixResult]) -> MatrixR def child_main(args: argparse.Namespace) -> int: + if args.loop == "rsloop" and not args.profile_label: + import rsloop + + if rsloop.profiler_compiled() and not args.allow_profiler_build: + raise RuntimeError( + "refusing to measure a Tracy-enabled rsloop build; rebuild without " + "--features profiler or pass --allow-profiler-build explicitly" + ) awaitable = SCENARIO_RUNNERS[args.scenario](args.loop, args) if args.profile_label: if args.loop != "rsloop": @@ -637,6 +672,8 @@ def child_command( ] if profile_label: cmd.extend(("--profile-label", profile_label)) + if args.allow_profiler_build: + cmd.append("--allow-profiler-build") return cmd @@ -688,24 +725,31 @@ def summarize(scenario: str, runs: dict[str, list[MatrixResult]]) -> None: loop_name, median_seconds, representative.operations / median_seconds, + representative.traffic_ops_per_sec, representative.bytes_transferred / (1024 * 1024) / median_seconds, percentile(representative.latency_ms, 0.50), percentile(representative.latency_ms, 0.95), percentile(representative.latency_ms, 0.99), int(statistics.median(item.peak_rss_bytes for item in measured)), + statistics.median(item.connection_setup_seconds for item in measured), + statistics.median(item.traffic_seconds for item in measured), + statistics.median(item.teardown_seconds for item in measured), ) ) rows.sort(key=lambda row: row[1]) print(f"\n{scenario}") print( - f"{'loop':<10} {'median_s':>10} {'ops/s':>12} {'MiB/s':>10} " - f"{'p50_ms':>10} {'p95_ms':>10} {'p99_ms':>10} {'peak_rss':>12}" + f"{'loop':<10} {'median_s':>10} {'total_ops/s':>12} {'traffic_ops/s':>14} " + f"{'MiB/s':>10} " + f"{'p50_ms':>10} {'p95_ms':>10} {'p99_ms':>10} {'peak_rss':>12} " + f"{'setup_s':>9} {'traffic_s':>10} {'close_s':>9}" ) for row in rows: print( - f"{row[0]:<10} {row[1]:>10.4f} {row[2]:>12,.0f} {row[3]:>10.1f} " - f"{row[4]:>10.3f} {row[5]:>10.3f} {row[6]:>10.3f} " - f"{format_bytes(row[7]):>12}" + f"{row[0]:<10} {row[1]:>10.4f} {row[2]:>12,.0f} {row[3]:>14,.0f} " + f"{row[4]:>10.1f} {row[5]:>10.3f} {row[6]:>10.3f} {row[7]:>10.3f} " + f"{format_bytes(row[8]):>12} {row[9]:>9.4f} {row[10]:>10.4f} " + f"{row[11]:>9.4f}" ) diff --git a/python/rsloop/__init__.py b/python/rsloop/__init__.py index 568f16f..21532a2 100644 --- a/python/rsloop/__init__.py +++ b/python/rsloop/__init__.py @@ -9,6 +9,7 @@ from ._loop_compat import Loop from ._loop_compat import __version__ from ._profile import profile +from ._profile import profiler_compiled from ._profile import profiler_running from ._profile import start_profiler from ._profile import stop_profiler @@ -25,6 +26,7 @@ "install", "new_event_loop", "profile", + "profiler_compiled", "profiler_running", "run", "start_profiler", diff --git a/python/rsloop/_profile.py b/python/rsloop/_profile.py index dbb5ec9..f124205 100644 --- a/python/rsloop/_profile.py +++ b/python/rsloop/_profile.py @@ -3,11 +3,17 @@ import contextlib as __contextlib import typing as __typing +from ._loop import profiler_compiled as __profiler_compiled from ._loop import profiler_running as __profiler_running from ._loop import start_profiler as __start_profiler from ._loop import stop_profiler as __stop_profiler +def profiler_compiled() -> bool: + """Return whether this installation was built with Tracy support.""" + return __profiler_compiled() + + def profiler_running() -> bool: return __profiler_running() diff --git a/src/fast_streams.rs b/src/fast_streams.rs index b2589c9..821c099 100644 --- a/src/fast_streams.rs +++ b/src/fast_streams.rs @@ -21,6 +21,7 @@ fn loop_create_future(py: Python<'_>, loop_obj: &Py) -> PyResult, start: usize, + retained_capacity: usize, } impl ReadBuffer { @@ -28,6 +29,7 @@ impl ReadBuffer { Self { bytes: Vec::with_capacity(capacity), start: 0, + retained_capacity: capacity, } } @@ -77,7 +79,11 @@ impl ReadBuffer { return; } if self.start == self.bytes.len() { - self.bytes.clear(); + if self.bytes.capacity() > self.retained_capacity.saturating_mul(4) { + self.bytes = Vec::with_capacity(self.retained_capacity); + } else { + self.bytes.clear(); + } self.start = 0; return; } @@ -89,6 +95,23 @@ impl ReadBuffer { } } +#[cfg(test)] +mod read_buffer_tests { + use super::ReadBuffer; + + #[test] + fn releases_oversized_allocation_after_consuming_all_data() { + let mut buffer = ReadBuffer::with_capacity(4096); + buffer.extend(&vec![0_u8; 2 * 1024 * 1024]); + assert!(buffer.bytes.capacity() >= 2 * 1024 * 1024); + + buffer.consume_all(); + + assert_eq!(buffer.len(), 0); + assert!(buffer.bytes.capacity() <= 4096); + } +} + #[derive(Clone)] enum ReadWaitKind { Any(usize), diff --git a/src/lib.rs b/src/lib.rs index 7400458..bdb74a7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,7 @@ pub use loop_core::{ LoopTransportCommand, }; pub use process_transport::{PyProcessPipeTransport, PyProcessTransport}; -pub use profiler::{profiler_running, start_profiler, stop_profiler}; +pub use profiler::{profiler_compiled, profiler_running, start_profiler, stop_profiler}; pub use python_api::{ PyLoop, asyncgen_finalizer_hook, asyncgen_firstiter_hook, future_done_stop, new_event_loop, signal_bridge, diff --git a/src/module_init.rs b/src/module_init.rs index 2ec883a..a60d332 100644 --- a/src/module_init.rs +++ b/src/module_init.rs @@ -3,8 +3,8 @@ use pyo3::prelude::*; use crate::{ PyFastStreamReader, PyFastStreamWriter, PyHandle, PyLoop, PyProcessPipeTransport, PyProcessTransport, PyServer, PyStreamTransport, PyTimerHandle, asyncgen_finalizer_hook, - asyncgen_firstiter_hook, future_done_stop, new_event_loop, open_connection, profiler_running, - signal_bridge, start_profiler, start_server, stop_profiler, + asyncgen_firstiter_hook, future_done_stop, new_event_loop, open_connection, profiler_compiled, + profiler_running, signal_bridge, start_profiler, start_server, stop_profiler, }; pub(crate) fn add_module_contents(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -45,6 +45,7 @@ fn add_stream_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { } fn add_profiler_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(profiler_compiled, m)?)?; m.add_function(wrap_pyfunction!(profiler_running, m)?)?; m.add_function(wrap_pyfunction!(start_profiler, m)?)?; m.add_function(wrap_pyfunction!(stop_profiler, m)?)?; diff --git a/src/profiler.rs b/src/profiler.rs index b8f8f59..1b061b6 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -1,3 +1,5 @@ +use pyo3::prelude::*; + #[cfg(feature = "profiler")] mod imp { use std::cell::RefCell; @@ -94,3 +96,8 @@ mod imp { } pub use imp::{profiler_running, start_profiler, stop_profiler}; + +#[pyfunction] +pub const fn profiler_compiled() -> bool { + cfg!(feature = "profiler") +} diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 39c1f44..42ac974 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -65,7 +65,10 @@ const MAX_PENDING_READ_COALESCE_BYTES: usize = 256 * 1024; const SMALL_WRITE_COALESCE_MIN_BYTES: usize = 64; const SMALL_WRITE_COALESCE_MAX_BYTES: usize = 512; const BLOCKING_POLL_INTERVAL_MS: i32 = 50; -const STREAM_READ_BUFFER_SIZE: usize = 64 * 1024; +// Keep the per-connection allocation modest. Large reads are delivered in +// several owned chunks, while the pending-event drain can still coalesce the +// slow protocol path up to MAX_PENDING_READ_COALESCE_BYTES. +const STREAM_READ_BUFFER_SIZE: usize = 16 * 1024; // After a successful read on a blocking reader worker, retry non-blocking // reads for this long before falling back to poll(). Request/response peers @@ -887,6 +890,7 @@ impl io::Write for WriterTarget { struct WorkerThread { stop: Arc, + wake: Option>, join: thread::JoinHandle<()>, } @@ -898,11 +902,28 @@ impl WorkerThread { .name(name.to_owned()) .spawn(move || task(thread_stop)) .expect("failed to spawn stream worker"); - Self { stop, join } + Self { + stop, + wake: None, + join, + } } - fn abort(self) { + fn spawn_interruptible( + name: &'static str, + wake: impl FnOnce() + Send + 'static, + task: impl FnOnce(Arc) + Send + 'static, + ) -> Self { + let mut worker = Self::spawn(name, task); + worker.wake = Some(Box::new(wake)); + worker + } + + fn abort(mut self) { self.stop.store(true, Ordering::Release); + if let Some(wake) = self.wake.take() { + wake(); + } let _ = self.join.join(); } } @@ -1015,6 +1036,19 @@ impl StreamTransportCore { match event { PendingReadEvent::Data(data) => { profiling::scope!("stream.pending.data"); + if let Some(fast_path) = fast_path.as_ref() { + if let Err(err) = fast_path.feed_data(py, &data) { + let _ = self.report_error_with_py( + py, + err, + "stream data_received callback failed", + ); + let _ = self.connection_lost_with_py(py, None); + self.read_events_scheduled.store(false, Ordering::Release); + return Ok(()); + } + continue; + } match &mut pending_data { Some(buffer) if buffer.len() + data.len() <= MAX_PENDING_READ_COALESCE_BYTES => @@ -2098,15 +2132,38 @@ impl ServerCore { let server = Arc::clone(self); let task = match listener { ServerListener::Tcp(listener) => { - WorkerThread::spawn("rsloop-tcp-accept", move |stop| { - run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - }) + let wake_addr = listener.local_addr().ok(); + WorkerThread::spawn_interruptible( + "rsloop-tcp-accept", + move || { + if let Some(addr) = wake_addr { + let _ = StdTcpStream::connect(addr); + } + }, + move |stop| { + run_tcp_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) + }, + ) } #[cfg(unix)] ServerListener::Unix(listener) => { - WorkerThread::spawn("rsloop-unix-accept", move |stop| { - run_unix_accept_loop(BlockingAcceptLoop::new(server, listener, stop)) - }) + let wake_path = listener + .local_addr() + .ok() + .and_then(|addr| addr.as_pathname().map(PathBuf::from)); + WorkerThread::spawn_interruptible( + "rsloop-unix-accept", + move || { + if let Some(path) = wake_path { + let _ = StdUnixStream::connect(path); + } + }, + move |stop| { + run_unix_accept_loop(BlockingAcceptLoop::new( + server, listener, stop, + )) + }, + ) } }; tasks.push(task); @@ -3371,6 +3428,10 @@ fn run_tcp_accept_loop(params: BlockingAcceptLoop) { } } + if stop.load(Ordering::Acquire) || server.is_closed() { + return; + } + loop { match listener.accept() { Ok((stream, _addr)) => { @@ -3471,6 +3532,10 @@ fn run_unix_accept_loop(params: BlockingAcceptLoop) { } } + if stop.load(Ordering::Acquire) || server.is_closed() { + return; + } + loop { match listener.accept() { Ok((stream, _addr)) => { @@ -3716,14 +3781,6 @@ fn tls_socket_wait_target(tls_state: &SharedTlsIoState) -> (fd_ops::RawFd, bool) (state.fd(), state.pollable()) } -fn write_tls_data(tls_state: &SharedTlsIoState, data: &[u8]) -> io::Result<()> { - let mut state = tls_state.lock().expect("poisoned tls state"); - match state.connection.writer_write_all(data) { - Ok(()) => flush_tls_io_locked(&mut state), - Err(err) => Err(err), - } -} - fn close_tls_writer(tls_state: &SharedTlsIoState) -> io::Result<()> { let mut state = tls_state.lock().expect("poisoned tls state"); let shutdown_timeout = state.shutdown_timeout; @@ -4214,16 +4271,23 @@ fn write_tls_data_batch( data: OwnedWriteBuffer, pending_command: &mut Option, ) -> bool { - if !write_one_tls_buffer(core, tls_state, &data) { + let mut buffered_len = 0; + let mut state = tls_state.lock().expect("poisoned tls state"); + if let Err(err) = state.connection.writer_write_all(data.remaining()) { + drop(state); + report_writer_io_error(core, err); return false; } - + buffered_len += data.remaining().len(); loop { match writer_rx.try_recv() { Ok(WriterCommand::Data(next)) => { - if !write_one_tls_buffer(core, tls_state, &next) { + if let Err(err) = state.connection.writer_write_all(next.remaining()) { + drop(state); + report_writer_io_error(core, err); return false; } + buffered_len += next.remaining().len(); } Ok(command) => { *pending_command = Some(command); @@ -4241,23 +4305,20 @@ fn write_tls_data_batch( } } - if pending_command.is_none() { - core.set_write_backpressure_active(false); - } - true -} - -fn write_one_tls_buffer( - core: &Arc, - tls_state: &SharedTlsIoState, - data: &OwnedWriteBuffer, -) -> bool { - let buffered_len = data.remaining().len(); - if let Err(err) = write_tls_data(tls_state, data.remaining()) { + // Feed all immediately available plaintext into rustls before flushing + // encrypted records. This turns common header/body write pairs into one + // socket-flush pass and one TLS-state lock acquisition. + if let Err(err) = flush_tls_io_locked(&mut state) { + drop(state); report_writer_io_error(core, err); return false; } + drop(state); core.record_write_buffer_drained(buffered_len); + + if pending_command.is_none() { + core.set_write_backpressure_active(false); + } true } From 0268318a06af23155a882f78ff3179338ac42cb8 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 08:51:17 +0000 Subject: [PATCH 07/21] a fix --- benchmarks/README.md | 8 +++++++- benchmarks/workload_matrix.py | 38 +++++++++++++++++++++++++++++------ 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 8b55f4e..7289fe9 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -127,9 +127,15 @@ pass before each rsloop scenario: uv run --with maturin maturin develop --release --features profiler uv run --with uvloop python benchmarks/workload_matrix.py \ --loops rsloop \ - --profile-rsloop-dir benchmarks/profiles + --profile-rsloop-dir benchmarks/profiles \ + --allow-profiler-build ``` +`--allow-profiler-build` is required because this invocation also reports +measurements from the Tracy-enabled binary. Treat those measurements as +profiling diagnostics, not as comparable release-build results. Rebuild without +`--features profiler` before collecting the normal comparison matrix. + On Windows, replace `--with uvloop` with `--with winloop` in the quick and profiling commands. diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py index 2ee8083..8e3186f 100644 --- a/benchmarks/workload_matrix.py +++ b/benchmarks/workload_matrix.py @@ -609,7 +609,18 @@ def run_with_loop(loop_name: str, awaitable: Awaitable[MatrixResult]) -> MatrixR def child_main(args: argparse.Namespace) -> int: - if args.loop == "rsloop" and not args.profile_label: + if args.profile_label: + if args.loop != "rsloop": + raise RuntimeError("Tracy profiling is only supported for rsloop") + import rsloop + + if not rsloop.profiler_compiled(): + raise RuntimeError( + "Tracy profiling was requested, but rsloop was built without profiler " + "support; rebuild with `uv run --with maturin maturin develop " + "--release --features profiler`" + ) + elif args.loop == "rsloop": import rsloop if rsloop.profiler_compiled() and not args.allow_profiler_build: @@ -617,16 +628,14 @@ def child_main(args: argparse.Namespace) -> int: "refusing to measure a Tracy-enabled rsloop build; rebuild without " "--features profiler or pass --allow-profiler-build explicitly" ) - awaitable = SCENARIO_RUNNERS[args.scenario](args.loop, args) + if args.profile_label: - if args.loop != "rsloop": - raise RuntimeError("Tracy profiling is only supported for rsloop") print(f"[profile] Tracy session label: {args.profile_label}", flush=True) - import rsloop - with rsloop.profile(): + awaitable = SCENARIO_RUNNERS[args.scenario](args.loop, args) result = run_with_loop(args.loop, awaitable) else: + awaitable = SCENARIO_RUNNERS[args.scenario](args.loop, args) result = run_with_loop(args.loop, awaitable) result = MatrixResult(**{**asdict(result), "peak_rss_bytes": get_peak_rss_bytes()}) print(json.dumps(asdict(result))) @@ -768,6 +777,23 @@ def parent_main(args: argparse.Namespace) -> int: if not available: raise SystemExit("no benchmarkable loops are available") + if args.profile_rsloop_dir: + if "rsloop" not in available: + raise SystemExit("--profile-rsloop-dir requires rsloop in --loops") + import rsloop + + if not rsloop.profiler_compiled(): + raise SystemExit( + "--profile-rsloop-dir requires a Tracy-enabled rsloop build. Run:\n" + " uv run --with maturin maturin develop --release --features profiler" + ) + if not args.allow_profiler_build: + raise SystemExit( + "this invocation profiles and then measures the same Tracy-enabled build; " + "pass --allow-profiler-build to acknowledge that its measured results are " + "not comparable to a normal release build" + ) + output: list[dict[str, object]] = [] for scenario in scenarios: scenario_runs: dict[str, list[MatrixResult]] = {} From 9ac480532798d94727967a9b6e4c1e7db1b592be Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 09:18:21 +0000 Subject: [PATCH 08/21] enhance websocket message handling and idle connection metrics in benchmarks --- benchmarks/workload_matrix.py | 92 +++++++++----- src/fast_streams.rs | 62 +++++++++- src/python_api.rs | 15 +++ src/stream_transport.rs | 221 ++++++++++++++++++++++++++-------- 4 files changed, 306 insertions(+), 84 deletions(-) diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py index 8e3186f..be373fa 100644 --- a/benchmarks/workload_matrix.py +++ b/benchmarks/workload_matrix.py @@ -77,6 +77,7 @@ class MatrixResult: connection_setup_seconds: float = 0.0 traffic_seconds: float = 0.0 teardown_seconds: float = 0.0 + idle_residency_seconds: float = 0.0 @property def ops_per_sec(self) -> float: @@ -383,48 +384,67 @@ async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> No host, port = server.sockets[0].getsockname()[:2] latencies: list[float] = [] - async def client(client_id: int) -> int: + connections: list[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = [] + + async def open_client( + client_id: int, + ) -> tuple[int, asyncio.StreamReader, asyncio.StreamWriter]: reader, writer = await asyncio.open_connection(host, port) + connections.append((reader, writer)) + writer.write(WEBSOCKET_REQUEST) + await writer.drain() + response = await reader.readexactly(len(WEBSOCKET_RESPONSE)) + if response != WEBSOCKET_RESPONSE: + raise RuntimeError("WebSocket upgrade failed") + return client_id, reader, writer + + async def client( + client_id: int, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> int: transferred = len(WEBSOCKET_REQUEST) + len(WEBSOCKET_RESPONSE) - try: - writer.write(WEBSOCKET_REQUEST) + for index in range(args.requests_per_connection): + size = args.websocket_payload_sizes[ + (client_id + index) % len(args.websocket_payload_sizes) + ] + payload = bytes([(client_id + index) % 251]) * size + mask_key = struct.pack("!I", (client_id << 16) ^ index ^ 0xA5A55A5A) + outbound = websocket_frame(payload, mask_key) + started = time.perf_counter() + writer.write(outbound) await writer.drain() - response = await reader.readexactly(len(WEBSOCKET_RESPONSE)) - if response != WEBSOCKET_RESPONSE: - raise RuntimeError("WebSocket upgrade failed") - for index in range(args.requests_per_connection): - size = args.websocket_payload_sizes[ - (client_id + index) % len(args.websocket_payload_sizes) - ] - payload = bytes([(client_id + index) % 251]) * size - mask_key = struct.pack("!I", (client_id << 16) ^ index ^ 0xA5A55A5A) - outbound = websocket_frame(payload, mask_key) - started = time.perf_counter() - writer.write(outbound) - await writer.drain() - response_payload = await read_websocket_frame(reader) - if response_payload != payload: - raise RuntimeError("WebSocket echo mismatch") - latencies.append((time.perf_counter() - started) * 1000) - transferred += len(outbound) + len(websocket_frame(payload)) - finally: - await close_writer(writer) + response_payload = await read_websocket_frame(reader) + if response_payload != payload: + raise RuntimeError("WebSocket echo mismatch") + latencies.append((time.perf_counter() - started) * 1000) + transferred += len(outbound) + len(websocket_frame(payload)) return transferred started = time.perf_counter() + setup_finished = traffic_finished = started try: - transferred = sum( - await asyncio.gather(*(client(index) for index in range(args.concurrency))) + opened = await asyncio.gather( + *(open_client(index) for index in range(args.concurrency)) ) + setup_finished = time.perf_counter() + transferred = sum(await asyncio.gather(*(client(*item) for item in opened))) + traffic_finished = time.perf_counter() finally: + await asyncio.gather( + *(close_writer(writer) for _, writer in connections), + return_exceptions=True, + ) await close_server(server) + finished = time.perf_counter() return MatrixResult( loop_name, "websocket_messages", - time.perf_counter() - started, + finished - started, args.concurrency * args.requests_per_connection, transferred, latencies, + connection_setup_seconds=setup_finished - started, + traffic_seconds=traffic_finished - setup_finished, + teardown_seconds=finished - traffic_finished, ) @@ -547,10 +567,13 @@ async def ping(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> No backlog=max(100, args.idle_connections), ) host, port = server.sockets[0].getsockname()[:2] + started = time.perf_counter() connections = await asyncio.gather( *(asyncio.open_connection(host, port) for _ in range(args.idle_connections)) ) + setup_finished = time.perf_counter() await asyncio.sleep(args.idle_seconds) + idle_finished = time.perf_counter() latencies: list[float] = [] async def activate( @@ -565,21 +588,29 @@ async def activate( latencies.append((time.perf_counter() - started) * 1000) return 2 - started = time.perf_counter() try: transferred = sum( await asyncio.gather(*(activate(item) for item in connections)) ) + traffic_finished = time.perf_counter() finally: await asyncio.gather(*(close_writer(writer) for _, writer in connections)) await close_server(server) + finished = time.perf_counter() + setup_seconds = setup_finished - started + traffic_seconds = traffic_finished - idle_finished + teardown_seconds = finished - traffic_finished return MatrixResult( loop_name, "idle_connections", - time.perf_counter() - started, + setup_seconds + traffic_seconds + teardown_seconds, args.idle_connections, transferred, latencies, + connection_setup_seconds=setup_seconds, + traffic_seconds=traffic_seconds, + teardown_seconds=teardown_seconds, + idle_residency_seconds=idle_finished - setup_finished, ) @@ -743,6 +774,7 @@ def summarize(scenario: str, runs: dict[str, list[MatrixResult]]) -> None: statistics.median(item.connection_setup_seconds for item in measured), statistics.median(item.traffic_seconds for item in measured), statistics.median(item.teardown_seconds for item in measured), + statistics.median(item.idle_residency_seconds for item in measured), ) ) rows.sort(key=lambda row: row[1]) @@ -751,14 +783,14 @@ def summarize(scenario: str, runs: dict[str, list[MatrixResult]]) -> None: f"{'loop':<10} {'median_s':>10} {'total_ops/s':>12} {'traffic_ops/s':>14} " f"{'MiB/s':>10} " f"{'p50_ms':>10} {'p95_ms':>10} {'p99_ms':>10} {'peak_rss':>12} " - f"{'setup_s':>9} {'traffic_s':>10} {'close_s':>9}" + f"{'setup_s':>9} {'traffic_s':>10} {'close_s':>9} {'idle_s':>8}" ) for row in rows: print( f"{row[0]:<10} {row[1]:>10.4f} {row[2]:>12,.0f} {row[3]:>14,.0f} " f"{row[4]:>10.1f} {row[5]:>10.3f} {row[6]:>10.3f} {row[7]:>10.3f} " f"{format_bytes(row[8]):>12} {row[9]:>9.4f} {row[10]:>10.4f} " - f"{row[11]:>9.4f}" + f"{row[11]:>9.4f} {row[12]:>8.4f}" ) diff --git a/src/fast_streams.rs b/src/fast_streams.rs index 821c099..fd4f94d 100644 --- a/src/fast_streams.rs +++ b/src/fast_streams.rs @@ -1,5 +1,6 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; use pyo3::types::{PyByteArray, PyBytes, PyDict, PyTuple}; use pyo3_async_runtimes::TaskLocals; @@ -9,6 +10,15 @@ use crate::stream_transport::{PyStreamTransport, task_locals_for_loop}; const DEFAULT_STREAM_LIMIT: usize = 65_536; +fn asyncio_iscoroutine(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { + static ISCOROUTINE: PyOnceLock> = PyOnceLock::new(); + Ok(ISCOROUTINE + .get_or_try_init(py, || -> PyResult> { + Ok(py.import("asyncio")?.getattr("iscoroutine")?.unbind()) + })? + .bind(py)) +} + /// Create a future on `loop_obj`, skipping the Python-level method dispatch /// when the loop is a native rsloop instance running on this thread. fn loop_create_future(py: Python<'_>, loop_obj: &Py) -> PyResult> { @@ -57,6 +67,18 @@ impl ReadBuffer { self.bytes.extend_from_slice(data); } + fn extend_owned(&mut self, data: Box<[u8]>) { + if data.is_empty() { + return; + } + if self.is_empty() { + self.bytes = data.into_vec(); + self.start = 0; + } else { + self.extend(&data); + } + } + fn consume(&mut self, n: usize) { self.start = (self.start + n).min(self.bytes.len()); self.compact_if_needed(); @@ -110,6 +132,18 @@ mod read_buffer_tests { assert_eq!(buffer.len(), 0); assert!(buffer.bytes.capacity() <= 4096); } + + #[test] + fn adopts_owned_data_without_copying_when_empty() { + let mut buffer = ReadBuffer::with_capacity(4096); + let data = vec![7_u8; 1024].into_boxed_slice(); + let data_ptr = data.as_ptr(); + + buffer.extend_owned(data); + + assert_eq!(buffer.unread().as_ptr(), data_ptr); + assert_eq!(buffer.unread(), &[7_u8; 1024]); + } } #[derive(Clone)] @@ -188,7 +222,9 @@ impl PyFastStreamReader { Ok(Self { loop_obj, limit, - buffer: ReadBuffer::with_capacity(limit.max(4096)), + // The flow-control limit is not an allocation target. Idle + // streams should retain only a small baseline and grow on data. + buffer: ReadBuffer::with_capacity(limit.min(4096)), waiter: None, transport: py.None(), paused: false, @@ -393,6 +429,19 @@ impl PyFastStreamReader { self.maybe_pause_transport(py) } + pub(crate) fn feed_owned_data_internal( + &mut self, + py: Python<'_>, + data: Box<[u8]>, + ) -> PyResult<()> { + if self.eof { + return Err(PyValueError::new_err("feed_data after feed_eof")); + } + self.buffer.extend_owned(data); + self.maybe_complete_waiter(py)?; + self.maybe_pause_transport(py) + } + #[inline] pub(crate) fn feed_eof_internal(&mut self, py: Python<'_>) -> PyResult<()> { self.eof = true; @@ -719,15 +768,18 @@ impl PyFastStreamProtocol { }, )?; let result = callback.call1(py, (reader.clone_ref(py), writer))?; - let asyncio = py.import("asyncio")?; - if !asyncio - .call_method1("iscoroutine", (result.clone_ref(py),))? + if !asyncio_iscoroutine(py)? + .call1((result.clone_ref(py),))? .extract::()? { return Ok(()); } - let task = loop_obj.call_method1(py, "create_task", (result,))?; + let task = + match crate::python_api::try_fast_create_task(py, &loop_obj, result.clone_ref(py))? { + Some(task) => task, + None => loop_obj.call_method1(py, "create_task", (result,))?, + }; slf.borrow_mut(py).task = task.clone_ref(py); let done_cb = Py::new( py, diff --git a/src/python_api.rs b/src/python_api.rs index 73cbec4..4a12b1e 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -452,6 +452,21 @@ pub(crate) fn try_fast_create_future( create_asyncio_future_for_running_loop(py).map(Some) } +pub(crate) fn try_fast_create_task( + py: Python<'_>, + loop_obj: &Py, + coro: Py, +) -> PyResult>> { + let Ok(pyloop) = loop_obj.bind(py).cast_exact::() else { + return Ok(None); + }; + let core = &pyloop.borrow().core; + if !core.on_runtime_thread() || core.has_task_factory() { + return Ok(None); + } + create_asyncio_task_for_running_loop(py, coro).map(Some) +} + fn create_asyncio_task_for_loop( py: Python<'_>, loop_obj: &Py, diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 42ac974..8752f2f 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -59,6 +59,8 @@ enum PendingReadEvent { const DEFAULT_WRITE_BUFFER_HIGH_WATER: usize = 64 * 1024; const DEFAULT_WRITE_BUFFER_LOW_WATER: usize = DEFAULT_WRITE_BUFFER_HIGH_WATER / 4; const MAX_PENDING_READ_COALESCE_BYTES: usize = 256 * 1024; +const MAX_READ_EVENTS_PER_DRAIN: usize = 32; +const MAX_READ_BYTES_PER_DRAIN: usize = 256 * 1024; // Servers commonly emit a small protocol header followed immediately by a // body. Defer only that header-sized first write for one loop turn so adjacent // writes share a syscall; keep tiny control frames and larger payloads direct. @@ -388,6 +390,13 @@ impl StreamReaderFastPath { } } + fn feed_owned_data(&self, py: Python<'_>, data: Box<[u8]>) -> PyResult<()> { + match self { + Self::Native { reader, .. } => reader.borrow_mut(py).feed_owned_data_internal(py, data), + Self::Generic { .. } => self.feed_data(py, &data), + } + } + fn feed_eof(&self, py: Python<'_>) -> PyResult<()> { match self { Self::Native { reader, .. } => reader.borrow_mut(py).feed_eof_internal(py), @@ -499,6 +508,7 @@ struct StreamTransportState { context: Py, context_needs_run: bool, extra: HashMap>, + lazy_socket_family: Option, closing: bool, read_paused: bool, reading: bool, @@ -593,6 +603,14 @@ enum TaskedDirectWriter { } impl TaskedDirectWriter { + fn fd(&self) -> fd_ops::RawFd { + match self { + Self::Tcp(stream) => tcp_stream_raw_fd(stream), + #[cfg(unix)] + Self::Unix(stream) => unix_raw_fd(stream.as_raw_fd()), + } + } + fn shutdown_close(&self) -> io::Result<()> { match self { Self::Tcp(stream) => shutdown_tcp_stream(stream, Shutdown::Both), @@ -650,10 +668,26 @@ enum WriterTarget { } struct LazyWriterConfig { - target: WriterTarget, + target: LazyWriterTarget, writer_rx: Receiver, } +enum LazyWriterTarget { + Tcp(fd_ops::RawFd), + #[cfg(unix)] + Unix(fd_ops::RawFd), +} + +impl LazyWriterTarget { + fn materialize(self) -> PyResult { + match self { + Self::Tcp(fd) => duplicate_configured_tcp_stream(fd).map(WriterTarget::Tcp), + #[cfg(unix)] + Self::Unix(fd) => duplicate_unix_direct_writer(fd).map(WriterTarget::Unix), + } + } +} + impl WriterTarget { fn fd(&self) -> Option { match self { @@ -959,7 +993,10 @@ impl StreamTransportCore { let Some(LazyWriterConfig { target, writer_rx }) = lazy else { return; }; - spawn_writer_worker(Arc::clone(self), target, writer_rx); + match target.materialize() { + Ok(target) => spawn_writer_worker(Arc::clone(self), target, writer_rx), + Err(err) => self.fail_write(Some(io::Error::other(err.to_string()))), + } } #[inline] @@ -1003,7 +1040,10 @@ impl StreamTransportCore { } } - pub(crate) fn drain_pending_read_events_with_py(&self, py: Python<'_>) -> PyResult<()> { + pub(crate) fn drain_pending_read_events_with_py( + self: &Arc, + py: Python<'_>, + ) -> PyResult<()> { profiling::scope!("StreamTransportCore::drain_pending_read_events_with_py"); // Snapshot the reader fast path once per drain instead of re-locking // transport state for every data event. @@ -1018,6 +1058,8 @@ impl StreamTransportCore { let fast_path = fast_path.as_ref(); let mut pending_data: Option = None; let mut drained = VecDeque::new(); + let mut drained_events = 0; + let mut drained_bytes = 0; loop { { let mut queue = self @@ -1036,8 +1078,10 @@ impl StreamTransportCore { match event { PendingReadEvent::Data(data) => { profiling::scope!("stream.pending.data"); + drained_events += 1; + drained_bytes += data.len(); if let Some(fast_path) = fast_path.as_ref() { - if let Err(err) = fast_path.feed_data(py, &data) { + if let Err(err) = fast_path.feed_owned_data(py, data) { let _ = self.report_error_with_py( py, err, @@ -1047,32 +1091,41 @@ impl StreamTransportCore { self.read_events_scheduled.store(false, Ordering::Release); return Ok(()); } - continue; - } - match &mut pending_data { - Some(buffer) - if buffer.len() + data.len() <= MAX_PENDING_READ_COALESCE_BYTES => - { - buffer.extend(data); - } - Some(_) => { - if let Err(err) = self.flush_pending_data_with_py( - py, - &mut pending_data, - fast_path, - ) { - let _ = self.report_error_with_py( + } else { + match &mut pending_data { + Some(buffer) + if buffer.len() + data.len() + <= MAX_PENDING_READ_COALESCE_BYTES => + { + buffer.extend(data); + } + Some(_) => { + if let Err(err) = self.flush_pending_data_with_py( py, - err, - "stream data_received callback failed", - ); - let _ = self.connection_lost_with_py(py, None); - self.read_events_scheduled.store(false, Ordering::Release); - return Ok(()); + &mut pending_data, + fast_path, + ) { + let _ = self.report_error_with_py( + py, + err, + "stream data_received callback failed", + ); + let _ = self.connection_lost_with_py(py, None); + self.read_events_scheduled.store(false, Ordering::Release); + return Ok(()); + } + pending_data = Some(PendingReadBuffer::Boxed(data)); } - pending_data = Some(PendingReadBuffer::Boxed(data)); + None => pending_data = Some(PendingReadBuffer::Boxed(data)), } - None => pending_data = Some(PendingReadBuffer::Boxed(data)), + } + + if (drained_events >= MAX_READ_EVENTS_PER_DRAIN + || drained_bytes >= MAX_READ_BYTES_PER_DRAIN) + && self.reschedule_pending_read_events(&mut drained) + { + self.flush_pending_data_with_py(py, &mut pending_data, fast_path)?; + return Ok(()); } } PendingReadEvent::Eof => { @@ -1175,6 +1228,29 @@ impl StreamTransportCore { } } } + + fn reschedule_pending_read_events( + self: &Arc, + drained: &mut VecDeque, + ) -> bool { + let mut queue = self + .pending_read_events + .lock() + .expect("poisoned pending read queue"); + if drained.is_empty() && queue.is_empty() { + return false; + } + + drained.append(&mut queue); + std::mem::swap(drained, queue.deref_mut()); + drop(queue); + let _ = + self.loop_core + .send_command(LoopCommand::Transport(LoopTransportCommand::StreamRead( + Arc::clone(self), + ))); + true + } } impl StreamTransportCore { @@ -1468,12 +1544,37 @@ impl StreamTransportCore { } fn get_extra(&self, py: Python<'_>, name: &str) -> Option> { - self.state - .lock() - .expect("poisoned transport state") - .extra - .get(name) - .map(|value| value.clone_ref(py)) + let (cached, lazy_socket_family, transport_closed) = { + let state = self.state.lock().expect("poisoned transport state"); + ( + state.extra.get(name).map(|value| value.clone_ref(py)), + state.lazy_socket_family, + state.closing || state.lost_called, + ) + }; + if let Some(value) = cached { + if name == "socket" && transport_closed { + let _ = value.bind(py).call_method0("close"); + } + return Some(value); + } + if lazy_socket_family.is_none() || !matches!(name, "socket" | "sockname" | "peername") { + return None; + } + + let family = lazy_socket_family?; + let fd = self + .direct_writer + .as_ref() + .map(|writer| writer.lock().expect("poisoned direct tasked writer").fd())?; + let extra = make_stream_extra(py, fd, family).ok()?; + if transport_closed && let Some(socket) = extra.get("socket") { + let _ = socket.bind(py).call_method0("close"); + } + let mut state = self.state.lock().expect("poisoned transport state"); + state.extra.extend(extra); + state.lazy_socket_family = None; + state.extra.get(name).map(|value| value.clone_ref(py)) } #[inline] @@ -2726,6 +2827,7 @@ struct StreamTransportStateConfig { io_fd: Option, runtime_socket_io: bool, extra: HashMap>, + lazy_socket_family: Option, reading: bool, writable: bool, can_write_eof: bool, @@ -2763,6 +2865,7 @@ fn stream_transport_state_parts( context, context_needs_run, extra: config.extra, + lazy_socket_family: config.lazy_socket_family, closing: false, read_paused: false, reading: config.reading, @@ -2840,6 +2943,7 @@ fn pipe_transport_core( io_fd: None, runtime_socket_io: false, extra, + lazy_socket_family: None, reading, writable, can_write_eof: writable, @@ -2951,13 +3055,11 @@ pub fn spawn_tcp_transport( server: Option>, ) -> PyResult> { let raw_fd = tcp_stream_raw_fd(&stream); - let extra = make_stream_extra(py, raw_fd, tcp_family(&stream))?; + let family = tcp_family(&stream); + let extra = HashMap::new(); let callbacks = build_protocol_callbacks(py, &spawn_context.protocol)?; spawn_context.context_needs_run &= callbacks.stream_reader_fast_path.is_none(); let direct_writer = duplicate_configured_tcp_stream(raw_fd)?; - let writer = stream - .try_clone() - .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; let (writer_tx, writer_rx) = mpsc::channel(); let parts = stream_transport_state_parts( spawn_context, @@ -2966,6 +3068,7 @@ pub fn spawn_tcp_transport( io_fd: Some(raw_fd), runtime_socket_io: true, extra, + lazy_socket_family: Some(family), reading: true, writable: true, can_write_eof: true, @@ -2978,7 +3081,7 @@ pub fn spawn_tcp_transport( writer_tx, Some(TaskedDirectWriter::Tcp(direct_writer)), Some(LazyWriterConfig { - target: WriterTarget::Tcp(writer), + target: LazyWriterTarget::Tcp(raw_fd), writer_rx, }), ); @@ -3007,13 +3110,10 @@ pub fn spawn_unix_transport( server: Option>, ) -> PyResult> { let raw_fd = unix_raw_fd(stream.as_raw_fd()); - let extra = make_stream_extra(py, raw_fd, libc::AF_UNIX)?; + let extra = HashMap::new(); let callbacks = build_protocol_callbacks(py, &spawn_context.protocol)?; spawn_context.context_needs_run &= callbacks.stream_reader_fast_path.is_none(); let direct_writer = duplicate_unix_direct_writer(raw_fd)?; - let writer = stream - .try_clone() - .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; let (writer_tx, writer_rx) = mpsc::channel(); let parts = stream_transport_state_parts( spawn_context, @@ -3022,6 +3122,7 @@ pub fn spawn_unix_transport( io_fd: Some(raw_fd), runtime_socket_io: true, extra, + lazy_socket_family: Some(libc::AF_UNIX), reading: true, writable: true, can_write_eof: true, @@ -3034,7 +3135,7 @@ pub fn spawn_unix_transport( writer_tx, Some(TaskedDirectWriter::Unix(direct_writer)), Some(LazyWriterConfig { - target: WriterTarget::Unix(writer), + target: LazyWriterTarget::Unix(raw_fd), writer_rx, }), ); @@ -3183,6 +3284,7 @@ fn spawn_tls_transport( io_fd: Some(stream_fd), runtime_socket_io: true, extra, + lazy_socket_family: None, reading: true, writable: true, can_write_eof: false, @@ -3923,7 +4025,9 @@ pub(crate) async fn run_tcp_socket_reader_task( return; } }; - let mut buf = vec![0_u8; STREAM_READ_BUFFER_SIZE]; + // `read_buf` writes into spare capacity, so idle transports reserve this + // address space without eagerly zero-filling and faulting in every page. + let mut buf = Vec::with_capacity(STREAM_READ_BUFFER_SIZE); loop { if core.is_closing() { @@ -3939,7 +4043,8 @@ pub(crate) async fn run_tcp_socket_reader_task( return; } - match reader.read(&mut buf).await { + buf.clear(); + match reader.read_buf(&mut buf).await { Ok(0) => { core.enqueue_pending_read_event(PendingReadEvent::Eof); return; @@ -3973,7 +4078,7 @@ pub(crate) async fn run_unix_socket_reader_task( return; } }; - let mut buf = vec![0_u8; STREAM_READ_BUFFER_SIZE]; + let mut buf = Vec::with_capacity(STREAM_READ_BUFFER_SIZE); loop { if core.is_closing() { @@ -3987,7 +4092,8 @@ pub(crate) async fn run_unix_socket_reader_task( core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); return; } - match reader.read(&mut buf).await { + buf.clear(); + match reader.read_buf(&mut buf).await { Ok(0) => { core.enqueue_pending_read_event(PendingReadEvent::Eof); return; @@ -4509,6 +4615,25 @@ fn shutdown_unix_stream(stream: &StdUnixStream, how: Shutdown) -> io::Result<()> fn build_protocol_callbacks(py: Python<'_>, protocol: &Py) -> PyResult { let bound = protocol.bind(py); + let stream_reader_fast_path = stream_reader_fast_path(py, bound)?; + if matches!( + &stream_reader_fast_path, + Some(StreamReaderFastPath::Native { .. }) + ) { + // Native streams handle these lifecycle and read callbacks directly + // in Rust. Only bind callbacks that remain reachable. + return Ok(ProtocolCallbacks { + connection_made: protocol.clone_ref(py), + data_received: None, + eof_received: None, + connection_lost: protocol.clone_ref(py), + pause_writing: bound.getattr(python_names::pause_writing(py))?.unbind(), + resume_writing: bound.getattr(python_names::resume_writing(py))?.unbind(), + get_buffer: None, + buffer_updated: None, + stream_reader_fast_path, + }); + } let data_received = match bound.getattr("data_received") { Ok(callback) => Some(callback.unbind()), Err(_) => None, @@ -4525,8 +4650,6 @@ fn build_protocol_callbacks(py: Python<'_>, protocol: &Py) -> PyResult Some(callback.unbind()), Err(_) => None, }; - let stream_reader_fast_path = stream_reader_fast_path(py, bound)?; - Ok(ProtocolCallbacks { connection_made: bound.getattr("connection_made")?.unbind(), data_received, From 96a221c7222cf4ebcb7be0bdb8f4df7916f1f02c Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 09:47:09 +0000 Subject: [PATCH 09/21] update client connection handling and improve data transfer logic --- benchmarks/workload_matrix.py | 94 +++++++++++++++++++--------- src/fast_streams.rs | 99 ++++++++++++++++++++++++++--- src/python_api.rs | 48 ++++++++------ src/stream_transport.rs | 114 ++++++++++++++++++---------------- 4 files changed, 246 insertions(+), 109 deletions(-) diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py index be373fa..7e39155 100644 --- a/benchmarks/workload_matrix.py +++ b/benchmarks/workload_matrix.py @@ -461,41 +461,60 @@ async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> No host, port = server.sockets[0].getsockname()[:2] latencies: list[float] = [] - async def client(client_id: int) -> int: + connections: list[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = [] + + async def open_client( + client_id: int, + ) -> tuple[int, asyncio.StreamReader, asyncio.StreamWriter]: reader, writer = await asyncio.open_connection(host, port) + connections.append((reader, writer)) + return client_id, reader, writer + + async def client( + client_id: int, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> int: transferred = 0 - try: - for index in range(args.requests_per_connection): - size = args.mixed_payload_sizes[ - (client_id + index) % len(args.mixed_payload_sizes) - ] - payload = bytes([(client_id + index) % 251]) * size - started = time.perf_counter() - writer.write(payload) - await writer.drain() - response = await reader.readexactly(size) - if response != payload: - raise RuntimeError("mixed stream echo mismatch") - latencies.append((time.perf_counter() - started) * 1000) - transferred += size * 2 - finally: - await close_writer(writer) + for index in range(args.requests_per_connection): + size = args.mixed_payload_sizes[ + (client_id + index) % len(args.mixed_payload_sizes) + ] + payload = bytes([(client_id + index) % 251]) * size + started = time.perf_counter() + writer.write(payload) + await writer.drain() + response = await reader.readexactly(size) + if response != payload: + raise RuntimeError("mixed stream echo mismatch") + latencies.append((time.perf_counter() - started) * 1000) + transferred += size * 2 return transferred started = time.perf_counter() + setup_finished = traffic_finished = started try: - transferred = sum( - await asyncio.gather(*(client(index) for index in range(args.concurrency))) + opened = await asyncio.gather( + *(open_client(index) for index in range(args.concurrency)) ) + setup_finished = time.perf_counter() + transferred = sum(await asyncio.gather(*(client(*item) for item in opened))) + traffic_finished = time.perf_counter() finally: + await asyncio.gather( + *(close_writer(writer) for _, writer in connections), + return_exceptions=True, + ) await close_server(server) + finished = time.perf_counter() return MatrixResult( loop_name, "mixed_streams", - time.perf_counter() - started, + finished - started, args.concurrency * args.requests_per_connection, transferred, latencies, + connection_setup_seconds=setup_finished - started, + traffic_seconds=traffic_finished - setup_finished, + teardown_seconds=finished - traffic_finished, ) @@ -520,32 +539,47 @@ async def send_bulk( host, port = server.sockets[0].getsockname()[:2] latencies: list[float] = [] - async def client() -> int: + connections: list[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = [] + + async def open_client() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: reader, writer = await asyncio.open_connection(host, port) + connections.append((reader, writer)) + return reader, writer + + async def client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> int: started = time.perf_counter() - try: - writer.write(b"!") - await writer.drain() - await reader.readexactly(args.bulk_bytes) - latencies.append((time.perf_counter() - started) * 1000) - finally: - await close_writer(writer) + writer.write(b"!") + await writer.drain() + await reader.readexactly(args.bulk_bytes) + latencies.append((time.perf_counter() - started) * 1000) return args.bulk_bytes + 1 started = time.perf_counter() + setup_finished = traffic_finished = started try: + opened = await asyncio.gather(*(open_client() for _ in range(args.concurrency))) + setup_finished = time.perf_counter() transferred = sum( - await asyncio.gather(*(client() for _ in range(args.concurrency))) + await asyncio.gather(*(client(*connection) for connection in opened)) ) + traffic_finished = time.perf_counter() finally: + await asyncio.gather( + *(close_writer(writer) for _, writer in connections), + return_exceptions=True, + ) await close_server(server) + finished = time.perf_counter() return MatrixResult( loop_name, "bulk_transfer", - time.perf_counter() - started, + finished - started, args.concurrency, transferred, latencies, + connection_setup_seconds=setup_finished - started, + traffic_seconds=traffic_finished - setup_finished, + teardown_seconds=finished - traffic_finished, ) diff --git a/src/fast_streams.rs b/src/fast_streams.rs index fd4f94d..be162ca 100644 --- a/src/fast_streams.rs +++ b/src/fast_streams.rs @@ -1,4 +1,5 @@ use pyo3::exceptions::PyValueError; +use pyo3::ffi; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; use pyo3::types::{PyByteArray, PyBytes, PyDict, PyTuple}; @@ -67,12 +68,12 @@ impl ReadBuffer { self.bytes.extend_from_slice(data); } - fn extend_owned(&mut self, data: Box<[u8]>) { + fn extend_owned(&mut self, data: Vec) { if data.is_empty() { return; } if self.is_empty() { - self.bytes = data.into_vec(); + self.bytes = data; self.start = 0; } else { self.extend(&data); @@ -136,7 +137,7 @@ mod read_buffer_tests { #[test] fn adopts_owned_data_without_copying_when_empty() { let mut buffer = ReadBuffer::with_capacity(4096); - let data = vec![7_u8; 1024].into_boxed_slice(); + let data = vec![7_u8; 1024]; let data_ptr = data.as_ptr(); buffer.extend_owned(data); @@ -156,6 +157,64 @@ enum ReadWaitKind { struct ReadWaiter { future: Py, kind: ReadWaitKind, + exact: Option, +} + +struct ExactReadAccumulator { + value: Py, + filled: usize, + expected: usize, +} + +impl ExactReadAccumulator { + fn new(py: Python<'_>, expected: usize) -> PyResult { + // SAFETY: The object is retained only inside this waiter and is not + // exposed to Python until all `expected` bytes have been initialized. + let ptr = unsafe { + ffi::PyBytes_FromStringAndSize(core::ptr::null(), expected as ffi::Py_ssize_t) + }; + if ptr.is_null() { + return Err(PyErr::fetch(py)); + } + let value = unsafe { + Bound::::from_owned_ptr(py, ptr) + .cast_into_unchecked::() + .unbind() + }; + Ok(Self { + value, + filled: 0, + expected, + }) + } + + fn fill_from(&mut self, buffer: &mut ReadBuffer) { + let amount = buffer.len().min(self.expected - self.filled); + if amount == 0 { + return; + } + // SAFETY: `value` has exactly `expected` bytes, remains private to + // this accumulator, and the destination range is disjoint/in-bounds. + unsafe { + let destination = ffi::PyBytes_AsString(self.value.as_ptr()) + .cast::() + .add(self.filled); + core::ptr::copy_nonoverlapping(buffer.unread().as_ptr(), destination, amount); + } + self.filled += amount; + buffer.consume(amount); + } + + fn partial(&self) -> &[u8] { + // SAFETY: The first `filled` bytes were initialized by `fill_from` and + // the returned slice is consumed immediately while `self` is alive. + unsafe { + core::slice::from_raw_parts( + ffi::PyBytes_AsString(self.value.as_ptr()).cast::(), + self.filled, + ) + } + } } #[pyclass(module = "rsloop._loop")] @@ -345,6 +404,14 @@ impl PyFastStreamReader { return Ok(()); } + if let Some(exact) = self + .waiter + .as_mut() + .and_then(|waiter| waiter.exact.as_mut()) + { + exact.fill_from(&mut self.buffer); + } + match kind { ReadWaitKind::Any(n) => { if self.buffer.is_empty() && !self.eof { @@ -356,9 +423,14 @@ impl PyFastStreamReader { Self::set_future_result_or_ignore_cancelled(py, &future, data)?; } ReadWaitKind::Exact(n) => { - if self.buffer.len() >= n { + let exact = self + .waiter + .as_ref() + .and_then(|waiter| waiter.exact.as_ref()) + .expect("exact read waiter missing accumulator"); + if exact.filled >= n { + let data = exact.value.clone_ref(py).into_any(); self.waiter = None; - let data = self.unread_bytes_object(py, n); self.maybe_resume_transport(py)?; Self::set_future_result_or_ignore_cancelled(py, &future, data)?; return Ok(()); @@ -366,8 +438,7 @@ impl PyFastStreamReader { if !self.eof { return Ok(()); } - let err = Self::incomplete_read_error(py, self.buffer.unread(), n)?; - self.buffer.consume_all(); + let err = Self::incomplete_read_error(py, exact.partial(), n)?; self.waiter = None; Self::set_future_exception_or_ignore_cancelled(py, &future, err)?; } @@ -404,10 +475,22 @@ impl PyFastStreamReader { )?; } let future = self.create_future(py)?; + let exact = match &kind { + ReadWaitKind::Exact(expected) => Some(ExactReadAccumulator::new(py, *expected)?), + _ => None, + }; self.waiter = Some(ReadWaiter { future: future.clone_ref(py), kind, + exact, }); + if let Some(exact) = self + .waiter + .as_mut() + .and_then(|waiter| waiter.exact.as_mut()) + { + exact.fill_from(&mut self.buffer); + } Ok(future) } @@ -432,7 +515,7 @@ impl PyFastStreamReader { pub(crate) fn feed_owned_data_internal( &mut self, py: Python<'_>, - data: Box<[u8]>, + data: Vec, ) -> PyResult<()> { if self.eof { return Err(PyValueError::new_err("feed_data after feed_eof")); diff --git a/src/python_api.rs b/src/python_api.rs index 4a12b1e..85db175 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -1973,8 +1973,8 @@ impl PyLoop { })? }; - let transport = Python::attach(|py| { - if let Some(ssl) = ssl.as_ref() { + let transport = if let Some(ssl) = ssl.as_ref() { + let (spawn_context, tls) = Python::attach(|py| { let tls = client_tls_settings( py, ssl.bind(py), @@ -1982,8 +1982,7 @@ impl PyLoop { ssl_handshake_timeout, ssl_shutdown_timeout, )?; - transport_from_socket_tls( - py, + Ok::<_, PyErr>(( stream_spawn_context( py, &core, @@ -1992,10 +1991,17 @@ impl PyLoop { &context, context_needs_run, ), - socket_obj, tls, - ) - } else { + )) + })?; + async_std::task::spawn_blocking(move || { + Python::attach(|py| { + transport_from_socket_tls(py, spawn_context, socket_obj, tls) + }) + }) + .await? + } else { + Python::attach(|py| { transport_from_socket( py, stream_spawn_context( @@ -2008,8 +2014,8 @@ impl PyLoop { ), socket_obj, ) - } - })?; + })? + }; Python::attach(|py| { let result = PyTuple::new(py, [transport.into_any(), protocol.clone_ref(py)])?; @@ -2270,8 +2276,8 @@ impl PyLoop { socket_obj }; - let transport = Python::attach(|py| { - if let Some(ssl) = ssl.as_ref() { + let transport = if let Some(ssl) = ssl.as_ref() { + let (spawn_context, tls) = Python::attach(|py| { let tls = client_tls_settings( py, ssl.bind(py), @@ -2279,8 +2285,7 @@ impl PyLoop { ssl_handshake_timeout, ssl_shutdown_timeout, )?; - transport_from_socket_tls( - py, + Ok::<_, PyErr>(( stream_spawn_context( py, &core, @@ -2289,10 +2294,17 @@ impl PyLoop { &context, context_needs_run, ), - socket_obj, tls, - ) - } else { + )) + })?; + async_std::task::spawn_blocking(move || { + Python::attach(|py| { + transport_from_socket_tls(py, spawn_context, socket_obj, tls) + }) + }) + .await? + } else { + Python::attach(|py| { transport_from_socket( py, stream_spawn_context( @@ -2305,8 +2317,8 @@ impl PyLoop { ), socket_obj, ) - } - })?; + })? + }; Python::attach(|py| { let result = PyTuple::new(py, [transport.into_any(), protocol.clone_ref(py)])?; diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 8752f2f..70de3f9 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -49,7 +49,7 @@ enum WriterCommand { } enum PendingReadEvent { - Data(Box<[u8]>), + Data(Vec), Eof, ConnectionLost(Option), PauseWriting, @@ -71,6 +71,8 @@ const BLOCKING_POLL_INTERVAL_MS: i32 = 50; // several owned chunks, while the pending-event drain can still coalesce the // slow protocol path up to MAX_PENDING_READ_COALESCE_BYTES. const STREAM_READ_BUFFER_SIZE: usize = 16 * 1024; +const OWNED_READ_HANDOFF_MIN_BYTES: usize = STREAM_READ_BUFFER_SIZE / 2; +const TLS_WORKER_STACK_SIZE: usize = 256 * 1024; // After a successful read on a blocking reader worker, retry non-blocking // reads for this long before falling back to poll(). Request/response peers @@ -92,38 +94,21 @@ struct OwnedWriteBuffer { offset: usize, } -enum PendingReadBuffer { - Boxed(Box<[u8]>), - Vec(Vec), -} +struct PendingReadBuffer(Vec); impl PendingReadBuffer { #[inline] fn len(&self) -> usize { - match self { - Self::Boxed(data) => data.len(), - Self::Vec(data) => data.len(), - } + self.0.len() } #[inline] fn as_slice(&self) -> &[u8] { - match self { - Self::Boxed(data) => data, - Self::Vec(data) => data, - } + &self.0 } - fn extend(&mut self, data: Box<[u8]>) { - match self { - Self::Boxed(existing) => { - let mut combined = Vec::with_capacity(existing.len() + data.len()); - combined.extend_from_slice(existing); - combined.extend_from_slice(&data); - *self = Self::Vec(combined); - } - Self::Vec(existing) => existing.extend_from_slice(&data), - } + fn extend(&mut self, data: Vec) { + self.0.extend_from_slice(&data); } } @@ -390,7 +375,7 @@ impl StreamReaderFastPath { } } - fn feed_owned_data(&self, py: Python<'_>, data: Box<[u8]>) -> PyResult<()> { + fn feed_owned_data(&self, py: Python<'_>, data: Vec) -> PyResult<()> { match self { Self::Native { reader, .. } => reader.borrow_mut(py).feed_owned_data_internal(py, data), Self::Generic { .. } => self.feed_data(py, &data), @@ -930,10 +915,21 @@ struct WorkerThread { impl WorkerThread { fn spawn(name: &'static str, task: impl FnOnce(Arc) + Send + 'static) -> Self { + Self::spawn_with_stack(name, None, task) + } + + fn spawn_with_stack( + name: &'static str, + stack_size: Option, + task: impl FnOnce(Arc) + Send + 'static, + ) -> Self { let stop = Arc::new(AtomicBool::new(false)); let thread_stop = Arc::clone(&stop); - let join = thread::Builder::new() - .name(name.to_owned()) + let mut builder = thread::Builder::new().name(name.to_owned()); + if let Some(stack_size) = stack_size { + builder = builder.stack_size(stack_size); + } + let join = builder .spawn(move || task(thread_stop)) .expect("failed to spawn stream worker"); Self { @@ -1114,9 +1110,9 @@ impl StreamTransportCore { self.read_events_scheduled.store(false, Ordering::Release); return Ok(()); } - pending_data = Some(PendingReadBuffer::Boxed(data)); + pending_data = Some(PendingReadBuffer(data)); } - None => pending_data = Some(PendingReadBuffer::Boxed(data)), + None => pending_data = Some(PendingReadBuffer(data)), } } @@ -3490,12 +3486,15 @@ pub(crate) fn spawn_accepted_transport_with_py( fn schedule_accepted_transport(server: &Arc, stream: AcceptedStream, message: &str) { if server.tls.is_some() { - let result = Python::try_attach(|py| spawn_accepted_transport_with_py(py, server, stream)); - match result { - Some(Ok(_)) => {} - Some(Err(err)) => server.report_error(err, message), - None => {} - } + let server = Arc::clone(server); + let message = message.to_owned(); + drop(async_std::task::spawn_blocking(move || { + let result = + Python::try_attach(|py| spawn_accepted_transport_with_py(py, &server, stream)); + if let Some(Err(err)) = result { + server.report_error(err, &message); + } + })); return; } @@ -3713,9 +3712,11 @@ fn spawn_reader_worker(core: Arc, reader: ReaderTarget) { fn spawn_tls_reader_worker(core: Arc, tls_state: SharedTlsIoState) { let thread_core = Arc::clone(&core); - let worker = WorkerThread::spawn("rsloop-tls-reader", move |stop| { - run_tls_reader(thread_core, tls_state, stop) - }); + let worker = WorkerThread::spawn_with_stack( + "rsloop-tls-reader", + Some(TLS_WORKER_STACK_SIZE), + move |stop| run_tls_reader(thread_core, tls_state, stop), + ); core.register_worker(worker); } @@ -3738,9 +3739,11 @@ fn spawn_tls_writer_worker( writer_rx: Receiver, ) { let thread_core = Arc::clone(&core); - let worker = WorkerThread::spawn("rsloop-tls-writer", move |stop| { - run_tls_writer(thread_core, tls_state, writer_rx, stop) - }); + let worker = WorkerThread::spawn_with_stack( + "rsloop-tls-writer", + Some(TLS_WORKER_STACK_SIZE), + move |stop| run_tls_writer(thread_core, tls_state, writer_rx, stop), + ); core.register_worker(worker); } @@ -3815,9 +3818,7 @@ fn drain_tls_plaintext_locked( Ok(0) => break, Ok(n) => { saw_data = true; - core.enqueue_pending_read_event(PendingReadEvent::Data(Box::<[u8]>::from( - &plaintext[..n], - ))); + core.enqueue_pending_read_event(PendingReadEvent::Data(plaintext[..n].to_vec())); } Err(err) if err.kind() == io::ErrorKind::WouldBlock => break, Err(err) => return Err(err.to_string()), @@ -3940,9 +3941,7 @@ fn run_stream_reader( return; } Ok(n) => { - core.enqueue_pending_read_event(PendingReadEvent::Data(Box::<[u8]>::from( - &buf[..n], - ))); + core.enqueue_pending_read_event(PendingReadEvent::Data(buf[..n].to_vec())); if !spin_window.is_zero() && !spin_read_stream(&core, &mut reader, &stop, &mut buf, spin_window) { @@ -3986,9 +3985,7 @@ fn spin_read_stream( return false; } Ok(n) => { - core.enqueue_pending_read_event(PendingReadEvent::Data(Box::<[u8]>::from( - &buf[..n], - ))); + core.enqueue_pending_read_event(PendingReadEvent::Data(buf[..n].to_vec())); if core.is_closing() || !core.is_reading() { return true; } @@ -4049,8 +4046,10 @@ pub(crate) async fn run_tcp_socket_reader_task( core.enqueue_pending_read_event(PendingReadEvent::Eof); return; } - Ok(n) => core - .enqueue_pending_read_event(PendingReadEvent::Data(Box::<[u8]>::from(&buf[..n]))), + Ok(_) => { + let data = take_async_read_data(&mut buf); + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); + } Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue, Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, Err(err) => { @@ -4098,8 +4097,10 @@ pub(crate) async fn run_unix_socket_reader_task( core.enqueue_pending_read_event(PendingReadEvent::Eof); return; } - Ok(n) => core - .enqueue_pending_read_event(PendingReadEvent::Data(Box::<[u8]>::from(&buf[..n]))), + Ok(_) => { + let data = take_async_read_data(&mut buf); + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); + } Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue, Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, Err(err) => { @@ -4112,6 +4113,13 @@ pub(crate) async fn run_unix_socket_reader_task( } } +fn take_async_read_data(buf: &mut Vec) -> Vec { + if buf.len() < OWNED_READ_HANDOFF_MIN_BYTES { + return buf.clone(); + } + std::mem::replace(buf, Vec::with_capacity(STREAM_READ_BUFFER_SIZE)) +} + pub(crate) fn run_socket_reader_blocking( core: Arc, reader: ReaderTarget, From f962982c28aeeac2df4d547e5eb712e1cac00544 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 15:27:46 +0000 Subject: [PATCH 10/21] add async-io dependency and optimize socket connection handling --- Cargo.lock | 1 + Cargo.toml | 1 + src/fd_ops.rs | 43 +++++++++++++------------- src/python_api.rs | 78 ++++++++++++++++++++++++++++++++++++++--------- 4 files changed, 88 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9451126..943c8d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -853,6 +853,7 @@ dependencies = [ name = "rsloop" version = "0.1.26" dependencies = [ + "async-io", "async-std", "crossbeam-channel", "futures", diff --git a/Cargo.toml b/Cargo.toml index 37b2baf..be420ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ default = [] profiler = ["dep:tracy-client", "profiling/profile-with-tracy"] [dependencies] +async-io = "2" async-std = "1" crossbeam-channel = "0.5" futures = { version = "0.3", default-features = false, features = ["std"] } diff --git a/src/fd_ops.rs b/src/fd_ops.rs index 54cce34..67063ff 100644 --- a/src/fd_ops.rs +++ b/src/fd_ops.rs @@ -1,9 +1,5 @@ use std::io; -#[cfg(not(windows))] -use std::thread; -#[cfg(not(windows))] -use futures::channel::oneshot; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; #[cfg(windows)] @@ -15,6 +11,7 @@ pub use self::windows::{ pub type RawFd = i64; +#[cfg(windows)] const FD_POLL_INTERVAL_MS: i32 = 50; pub fn fileobj_to_fd(_py: Python<'_>, fileobj: &Bound<'_, PyAny>) -> PyResult { @@ -159,25 +156,29 @@ async fn wait_for_interest(fd: RawFd, read: bool, write: bool) -> PyResult<()> { #[cfg(not(windows))] { - let (tx, rx) = oneshot::channel(); - thread::Builder::new() - .name(format!("rsloop-fd-wait-{fd}")) - .spawn(move || { - let result = loop { - match poll_fd(fd, read, write, FD_POLL_INTERVAL_MS) { - Ok((true, _)) if read => break Ok(()), - Ok((_, true)) if write => break Ok(()), - Ok(_) => continue, - Err(err) => break Err(err), - } - }; - let _ = tx.send(result); - }) + use std::os::fd::BorrowedFd; + + // Register the descriptor with async-io's shared, process-wide reactor + // thread instead of spawning a fresh OS thread per wait. The previous + // thread-per-wait approach dominated connection-setup latency: a burst + // of N concurrent connects spawned N `poll()` threads. async-io drives + // the same epoll/kqueue reactor smol/async-std already run, so this adds + // no extra threads and deregisters as soon as the wait resolves. + let raw = raw_fd_to_c_int(fd).map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + // SAFETY: the caller keeps `fd` open for the duration of this await + // (the owning Python socket outlives the connect/recv/send operation). + let borrowed = unsafe { BorrowedFd::borrow_raw(raw) }; + let async_fd = async_io::Async::new_nonblocking(borrowed) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; - rx.await - .map_err(|_| PyRuntimeError::new_err("fd wait worker dropped"))? - .map_err(|err| PyRuntimeError::new_err(err.to_string())) + let result = if write { + futures::future::poll_fn(|cx| async_fd.poll_writable(cx)).await + } else if read { + futures::future::poll_fn(|cx| async_fd.poll_readable(cx)).await + } else { + Ok(()) + }; + result.map_err(|err| PyRuntimeError::new_err(err.to_string())) } } diff --git a/src/python_api.rs b/src/python_api.rs index 85db175..b515d69 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -683,26 +683,36 @@ fn set_socket_bool_option_unix( } async fn connect_socket_to_address(sock: Py, address: Py) -> PyResult<()> { - let fd = Python::attach(|py| fd_ops::fileobj_to_fd(py, sock.bind(py)))?; - match Python::attach(|py| -> PyResult<()> { - sock.call_method1(py, "connect", (address.clone_ref(py),))?; - Ok(()) - }) { - Ok(()) => return Ok(()), - Err(err) => { - let retry = Python::attach(|py| fd_ops::is_retryable_socket_error(py, &err))?; - if !retry { - if Python::attach(|py| is_already_connected_socket_error(py, &err))? { - return Ok(()); + // Initiate the connect and look up the descriptor in a single GIL + // acquisition. Every `Python::attach` here runs on the async runtime's + // worker thread and contends for the GIL held by the main loop thread, so + // collapsing the fd lookup, the connect() call, and its error + // classification into one attach removes several contended handoffs per + // connection. On a non-blocking socket, connect() to a reachable peer + // returns EINPROGRESS (retryable); anything else is an immediate success or + // a hard error. + let fd = match Python::attach(|py| -> PyResult> { + let fd = fd_ops::fileobj_to_fd(py, sock.bind(py))?; + match sock.call_method1(py, "connect", (address.bind(py),)) { + Ok(_) => Ok(None), + Err(err) => { + if fd_ops::is_retryable_socket_error(py, &err)? { + Ok(Some(fd)) + } else if is_already_connected_socket_error(py, &err)? { + Ok(None) + } else { + Err(err) } - return Err(err); } } - } + })? { + Some(fd) => fd, + None => return Ok(()), + }; loop { fd_ops::wait_writable(fd).await?; - let so_error = Python::attach(|py| socket_so_error(py, &sock))?; + let so_error = connect_so_error(fd, &sock)?; if so_error == 0 { return Ok(()); } @@ -716,6 +726,46 @@ async fn connect_socket_to_address(sock: Py, address: Py) -> PyRes } } +/// Reads `SO_ERROR` for a connecting socket. On Unix this uses a direct +/// `getsockopt` so the hot connect-completion path never re-acquires the GIL +/// (or re-imports the `socket` module); Windows keeps the Python fallback. +fn connect_so_error(fd: fd_ops::RawFd, sock: &Py) -> PyResult { + #[cfg(unix)] + { + let _ = sock; + let fd: libc::c_int = fd + .try_into() + .map_err(|_| PyRuntimeError::new_err("socket file descriptor out of range"))?; + let mut value: libc::c_int = 0; + let mut len: libc::socklen_t = std::mem::size_of::() + .try_into() + .expect("socklen_t can represent c_int size"); + // SAFETY: `fd` is a live socket descriptor and `value`/`len` point to a + // correctly sized `c_int`/`socklen_t` that `getsockopt` fills in. + let result = unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&mut value as *mut libc::c_int).cast(), + &mut len, + ) + }; + if result == 0 { + Ok(value) + } else { + Err(PyErr::from(std::io::Error::last_os_error())) + } + } + + #[cfg(windows)] + { + let _ = fd; + Python::attach(|py| socket_so_error(py, sock)) + } +} + +#[cfg(windows)] fn socket_so_error(py: Python<'_>, sock: &Py) -> PyResult { let socket_mod = py.import("socket")?; sock.call_method1( From 17bd3ed820fb111cb620931dcf6bb6722f5ab5b1 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 15:41:58 +0000 Subject: [PATCH 11/21] optimize IP address resolution by adding __ipaddr_info function to reduce executor overhead --- python/rsloop/_loop_compat.py | 80 ++++++++++++++++++++++++++++++----- 1 file changed, 70 insertions(+), 10 deletions(-) diff --git a/python/rsloop/_loop_compat.py b/python/rsloop/_loop_compat.py index 1cd2b90..6d2e34d 100644 --- a/python/rsloop/_loop_compat.py +++ b/python/rsloop/_loop_compat.py @@ -1149,6 +1149,62 @@ def protocol_factory(): __asyncio.create_subprocess_shell = __create_text_subprocess_shell +__HAS_INET_PTON = hasattr(__socket, "inet_pton") +__HAS_IPv6 = hasattr(__socket, "AF_INET6") + + +def __ipaddr_info(host, port, family, socktype, proto): + # Mirror asyncio.base_events._ipaddr_info: when the host is already a + # numeric IP literal there is nothing to resolve, so build the addrinfo + # inline instead of dispatching socket.getaddrinfo() to a thread-pool + # executor. rsloop's getaddrinfo() goes through run_in_executor(), so the + # old unconditional call paid a full executor round-trip (queue hop + a + # worker thread + a future callback back to the loop) on every connection, + # even for "127.0.0.1". stdlib asyncio and uvloop both take this shortcut. + if not __HAS_INET_PTON: + return None + if proto not in (0, __socket.IPPROTO_TCP, __socket.IPPROTO_UDP) or host is None: + return None + if socktype == __socket.SOCK_STREAM: + proto = __socket.IPPROTO_TCP + elif socktype == __socket.SOCK_DGRAM: + proto = __socket.IPPROTO_UDP + else: + return None + if port is None: + port = 0 + elif isinstance(port, bytes) and port == b"": + port = 0 + elif isinstance(port, str) and port == "": + port = 0 + else: + try: + port = int(port) + except (TypeError, ValueError): + return None + if family == __socket.AF_UNSPEC: + afs = [__socket.AF_INET] + if __HAS_IPv6: + afs.append(__socket.AF_INET6) + else: + afs = [family] + if isinstance(host, bytes): + try: + host = host.decode("idna") + except UnicodeError: + return None + if "%" in host: + # Scoped IPv6 addresses need getaddrinfo() to resolve the scope id. + return None + for af in afs: + try: + __socket.inet_pton(af, host) + return af, socktype, proto, "", (host, port) + except OSError: + pass + return None + + def __interleave_addrinfos(addrinfos, first_address_family_count=1): import collections as __collections import itertools as __itertools @@ -1332,16 +1388,20 @@ async def __loop_create_connection( if sock is not None: raise ValueError("host/port and sock can not be specified at the same time") - addrinfos = await self.getaddrinfo( - host, - port, - family=family, - type=__socket.SOCK_STREAM, - proto=proto, - flags=flags, - ) - if not addrinfos: - raise OSError("getaddrinfo() returned empty list") + info = __ipaddr_info(host, port, family, __socket.SOCK_STREAM, proto) + if info is not None: + addrinfos = [info] + else: + addrinfos = await self.getaddrinfo( + host, + port, + family=family, + type=__socket.SOCK_STREAM, + proto=proto, + flags=flags, + ) + if not addrinfos: + raise OSError("getaddrinfo() returned empty list") if local_addr is not None: local_addrinfos = await self.getaddrinfo( From 3fe2c721c523b90582bc1477ff6d5afc7fa7db3f Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 16:05:50 +0000 Subject: [PATCH 12/21] implement non-blocking TCP connect handling with vibeio reactor for improved performance --- python/rsloop/_loop_compat.py | 5 +- src/fd_ops.rs | 39 ++++++++++++++ src/loop_core.rs | 95 +++++++++++++++++++++++++++++++++++ src/loop_core/commands.rs | 24 +++++++++ src/python_api.rs | 83 ++++++++++++++++++++++++++++++ src/runtime.rs | 24 +++++++++ src/stream_transport.rs | 43 ++++++++++++++++ tests/test_compat.py | 5 +- 8 files changed, 316 insertions(+), 2 deletions(-) diff --git a/python/rsloop/_loop_compat.py b/python/rsloop/_loop_compat.py index 6d2e34d..0a1692e 100644 --- a/python/rsloop/_loop_compat.py +++ b/python/rsloop/_loop_compat.py @@ -1429,7 +1429,10 @@ async def __loop_create_connection( connection_exceptions.append(attempt_exceptions) continue try: - await self.sock_connect(created_sock, address) + # Awaited directly (never wrapped in a task), so a bare + # Future is fine here; on Unix this keeps the connect + # writability wait on the loop's own reactor. + await self._sock_connect_fast(created_sock, address) break except OSError as exc: attempt_exceptions.append(exc) diff --git a/src/fd_ops.rs b/src/fd_ops.rs index 67063ff..ff0161b 100644 --- a/src/fd_ops.rs +++ b/src/fd_ops.rs @@ -193,3 +193,42 @@ fn raw_fd_to_c_int(fd: RawFd) -> io::Result { fd.try_into() .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "descriptor out of range")) } + +/// A connect() attempt that is still completing in the background. +#[cfg(unix)] +pub fn is_connect_in_progress_errno(errno: i32) -> bool { + errno == libc::EINPROGRESS || errno == libc::EALREADY || errno == libc::EWOULDBLOCK +} + +/// The socket is already connected (a benign outcome for connect()). +#[cfg(unix)] +pub fn is_already_connected_errno(errno: i32) -> bool { + errno == libc::EISCONN +} + +/// Reads the pending `SO_ERROR` for a socket via a direct `getsockopt`, so the +/// connect-completion path resolves without acquiring the GIL. +#[cfg(unix)] +pub fn socket_so_error(fd: RawFd) -> io::Result { + let fd = raw_fd_to_c_int(fd)?; + let mut value: libc::c_int = 0; + let mut len: libc::socklen_t = std::mem::size_of::() + .try_into() + .expect("socklen_t can represent c_int size"); + // SAFETY: `fd` is a socket descriptor and `value`/`len` describe a correctly + // sized `c_int`/`socklen_t` that `getsockopt` fills in. + let result = unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&mut value as *mut libc::c_int).cast(), + &mut len, + ) + }; + if result == 0 { + Ok(value) + } else { + Err(io::Error::last_os_error()) + } +} diff --git a/src/loop_core.rs b/src/loop_core.rs index 3c66c33..b13b478 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -479,6 +479,15 @@ impl LoopCore { server.report_error(err, "failed to accept connection"); } } + #[cfg(unix)] + ReadyItem::ConnectCompleted { + future, + fd, + wait_errno, + } => { + profiling::scope!("ready.connect_completed"); + self.resolve_connect_completed(py, future, fd, wait_errno)?; + } } processed_since_refill += 1; @@ -777,6 +786,58 @@ impl LoopCore { ACTIVE_LOOP_TLS.with(|tls| std::ptr::eq(tls.core.get(), self)) } + /// Resolves a TCP connect whose writability wait finished on the vibeio + /// reactor. Runs on the loop thread so the SO_ERROR check and the + /// set_result / set_exception happen with the GIL already held for the + /// whole ready batch — no per-completion GIL handoff. + #[cfg(unix)] + fn resolve_connect_completed( + &self, + py: Python<'_>, + future: Py, + fd: RawFd, + wait_errno: i32, + ) -> PyResult<()> { + let future = future.bind(py); + let done = crate::python_names::call_method0(py, future, crate::python_names::done(py))? + .bind(py) + .extract::()?; + if done { + return Ok(()); + } + + // SO_ERROR is authoritative for the connect outcome; the wait error is + // only a fallback for the rare case where SO_ERROR is already cleared. + let so_error = crate::fd_ops::socket_so_error(fd) + .unwrap_or_else(|err| err.raw_os_error().unwrap_or(libc::EBADF)); + let errno = if so_error != 0 { so_error } else { wait_errno }; + + if errno == 0 || crate::fd_ops::is_already_connected_errno(errno) { + crate::python_names::call_method1( + py, + future, + crate::python_names::set_result(py), + py.None().bind(py), + )?; + } else if crate::fd_ops::is_connect_in_progress_errno(errno) { + // Spurious writability wakeup while still connecting; re-arm. + let _ = self.send_command(LoopCommand::Io(LoopIoCommand::WatchConnect { + fd, + future: future.clone().unbind(), + })); + } else { + let message = std::io::Error::from_raw_os_error(errno).to_string(); + let oserror = pyo3::exceptions::PyOSError::new_err((errno, message)).into_value(py); + crate::python_names::call_method1( + py, + future, + crate::python_names::set_exception(py), + oserror.bind(py).as_any(), + )?; + } + Ok(()) + } + #[inline] fn try_handle_local_command(&self, command: LoopCommand) -> Result<(), LoopCommand> { match command { @@ -808,6 +869,16 @@ impl LoopCore { stream, }) } + #[cfg(unix)] + ReadyItem::ConnectCompleted { + future, + fd, + wait_errno, + } => LoopCommand::ConnectCompleted { + future, + fd, + wait_errno, + }, }), LoopCommand::ScheduleReadyHandle(handle) => self .try_enqueue_local_ready(ReadyItem::HandleCallback(handle)) @@ -885,6 +956,30 @@ impl LoopCore { unreachable!("local accepted transport enqueue preserves item kind") } }), + #[cfg(unix)] + LoopCommand::ConnectCompleted { + future, + fd, + wait_errno, + } => self + .try_enqueue_local_ready(ReadyItem::ConnectCompleted { + future, + fd, + wait_errno, + }) + .or_else(|item| self.try_enqueue_active_ready(item)) + .map_err(|item| match item { + ReadyItem::ConnectCompleted { + future, + fd, + wait_errno, + } => LoopCommand::ConnectCompleted { + future, + fd, + wait_errno, + }, + _ => unreachable!("local connect completion enqueue preserves item kind"), + }), LoopCommand::RequestStop => self .try_enqueue_local_ready(ReadyItem::Stop) .or_else(|item| self.try_enqueue_active_ready(item)) diff --git a/src/loop_core/commands.rs b/src/loop_core/commands.rs index d5052a2..403e373 100644 --- a/src/loop_core/commands.rs +++ b/src/loop_core/commands.rs @@ -29,6 +29,17 @@ pub enum ReadyItem { server: Arc, stream: AcceptedStream, }, + // A non-blocking TCP connect finished its writability wait on the vibeio + // reactor. The Python-object work (SO_ERROR check, set_result / + // set_exception) is deferred to the loop thread so the vibeio side never + // touches the GIL — many concurrent connects then drain in one GIL-held + // batch instead of one contended handoff per completion. + #[cfg(unix)] + ConnectCompleted { + future: Py, + fd: RawFd, + wait_errno: i32, + }, Stop, } @@ -44,6 +55,12 @@ pub enum LoopCommand { Io(LoopIoCommand), Future(LoopFutureCommand), Transport(LoopTransportCommand), + #[cfg(unix)] + ConnectCompleted { + future: Py, + fd: RawFd, + wait_errno: i32, + }, RequestStop, Close, } @@ -92,6 +109,13 @@ pub enum LoopIoCommand { listener: ServerListener, }, StopServerAccept(RawFd), + // Watch a connecting TCP socket for writability on the vibeio reactor, then + // report completion back to the loop thread via `ReadyItem::ConnectCompleted`. + #[cfg(unix)] + WatchConnect { + fd: RawFd, + future: Py, + }, } pub enum LoopFutureCommand { diff --git a/src/python_api.rs b/src/python_api.rs index b515d69..c9af459 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -810,6 +810,63 @@ fn is_connect_in_progress_errno(errno: i32) -> bool { errno == libc::EINPROGRESS || errno == libc::EALREADY || errno == libc::EWOULDBLOCK } +/// True for an AF_INET/AF_INET6 SOCK_STREAM socket, the sockets the vibeio +/// connect fast path handles. A type carrying extra flag bits declines the fast +/// path (falls back), which is never a correctness problem. +#[cfg(unix)] +fn socket_is_inet_stream(py: Python<'_>, sock: &Py) -> PyResult { + let socket_mod = py.import("socket")?; + let family: i32 = sock.getattr(py, "family")?.extract(py)?; + let af_inet: i32 = socket_mod.getattr("AF_INET")?.extract()?; + let af_inet6: i32 = socket_mod.getattr("AF_INET6")?.extract()?; + if family != af_inet && family != af_inet6 { + return Ok(false); + } + let sock_type: i32 = sock.getattr(py, "type")?.extract(py)?; + let sock_stream: i32 = socket_mod.getattr("SOCK_STREAM")?.extract()?; + Ok(sock_type == sock_stream) +} + +/// Initiates a non-blocking connect on the loop thread and, when it does not +/// complete synchronously, hands the writability wait to the vibeio reactor via +/// `WatchConnect`. Returns the loop Future the caller awaits. +#[cfg(unix)] +fn fast_sock_connect<'py>( + slf: &Py, + py: Python<'py>, + sock: Py, + address: Py, +) -> PyResult> { + let future = slf.call_method0(py, "create_future")?; + let future_bound = future.bind(py).clone(); + let fd = fd_ops::fileobj_to_fd(py, sock.bind(py))?; + + match sock.call_method1(py, "connect", (address.bind(py),)) { + Ok(_) => { + future_bound.call_method1("set_result", (py.None(),))?; + return Ok(future_bound); + } + Err(err) => { + if is_already_connected_socket_error(py, &err)? { + future_bound.call_method1("set_result", (py.None(),))?; + return Ok(future_bound); + } + if !fd_ops::is_retryable_socket_error(py, &err)? { + future_bound.call_method1("set_exception", (err.into_value(py),))?; + return Ok(future_bound); + } + } + } + + let core = slf.borrow(py).core.clone(); + core.send_command(LoopCommand::Io(LoopIoCommand::WatchConnect { + fd, + future: future.clone_ref(py), + })) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + Ok(future_bound) +} + fn listener_sources_from_sockets( py: Python<'_>, sockets: &[Py], @@ -2810,6 +2867,32 @@ impl PyLoop { }) } + /// Connects an INET/INET6 stream socket, returning a loop-native Future + /// (not a coroutine — awaited directly, never `create_task`ed). On Unix the + /// writability wait runs on the vibeio reactor and its completion is + /// delivered through the loop's batched, GIL-free ready queue, so many + /// concurrent connections drain in one loop iteration instead of paying a + /// per-connection async-runtime handoff. Non-Unix and non-INET sockets fall + /// back to the general `sock_connect` path. + fn _sock_connect_fast<'py>( + slf: Py, + py: Python<'py>, + sock: Py, + address: Py, + ) -> PyResult> { + #[cfg(unix)] + { + if socket_is_inet_stream(py, &sock)? { + return fast_sock_connect(&slf, py, sock, address); + } + } + let locals = Self::task_locals(py, &slf)?; + pyo3_async_runtimes::async_std::future_into_py_with_locals(py, locals, async move { + connect_socket_to_address(sock, address).await?; + Ok(Python::attach(|py| py.None())) + }) + } + #[pyo3(signature=(host, port, *, family=0, r#type=0, proto=0, flags=0))] #[expect( clippy::too_many_arguments, diff --git a/src/runtime.rs b/src/runtime.rs index b0503fc..cf11435 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -476,6 +476,30 @@ impl RuntimeDispatcher { cancel_watch_task(task); } } + #[cfg(unix)] + LoopCommand::Io(LoopIoCommand::WatchConnect { fd, future }) => { + profiling::scope!("runtime.cmd.watch_connect"); + // Detached: the task self-reports via ConnectCompleted, and a + // vibeio JoinHandle has no Drop, so dropping it detaches (does + // not cancel) the running task. + let core = Arc::clone(&self.core); + let _ = vibeio::spawn(crate::stream_transport::run_connect_watch_task( + core, fd, future, + )); + } + #[cfg(unix)] + LoopCommand::ConnectCompleted { + future, + fd, + wait_errno, + } => { + profiling::scope!("runtime.cmd.connect_completed"); + self.ready_batch.push_back(ReadyItem::ConnectCompleted { + future, + fd, + wait_errno, + }); + } LoopCommand::RequestStop => { profiling::scope!("runtime.cmd.request_stop"); self.ready_batch.push_back(ReadyItem::Stop); diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 70de3f9..72ef998 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -3559,6 +3559,49 @@ fn run_tcp_accept_loop(params: BlockingAcceptLoop) { } } +/// Waits for a connecting TCP socket to become writable on the vibeio reactor, +/// then hands the outcome back to the loop thread via `ConnectCompleted`. The +/// descriptor is duplicated so the vibeio stream owns (and closes) only the +/// dup; the connecting socket stays owned by its Python object. The dup shares +/// the same underlying socket, so its POLLOUT readiness reflects the connect +/// completing. +#[cfg(unix)] +pub(crate) async fn run_connect_watch_task( + core: Arc, + fd: fd_ops::RawFd, + future: Py, +) { + use vibeio::io::AsyncWritePoll; + + let wait_errno = 'wait: { + let Ok(raw) = c_int::try_from(fd) else { + break 'wait libc::EBADF; + }; + // SAFETY: `raw` is a live socket descriptor owned by the Python socket. + let dup = unsafe { libc::dup(raw) }; + if dup < 0 { + break 'wait io::Error::last_os_error() + .raw_os_error() + .unwrap_or(libc::EBADF); + } + // SAFETY: `dup` is a freshly duplicated, owned socket descriptor. + let std_stream = unsafe { StdTcpStream::from_raw_fd(dup) }; + match VibePollTcpStream::from_std(std_stream) { + Ok(stream) => match stream.writable().await { + Ok(()) => 0, + Err(err) => err.raw_os_error().unwrap_or(0), + }, + Err(err) => err.raw_os_error().unwrap_or(libc::EBADF), + } + }; + + let _ = core.send_command(LoopCommand::ConnectCompleted { + future, + fd, + wait_errno, + }); +} + pub(crate) async fn run_server_accept_task(server: Arc, listener: ServerListener) { profiling::scope!("stream.run_server_accept_task"); match listener { diff --git a/tests/test_compat.py b/tests/test_compat.py index 7c55a32..55df7bb 100644 --- a/tests/test_compat.py +++ b/tests/test_compat.py @@ -97,9 +97,12 @@ async def fake_sock_connect(self, sock, address): calls.append(address[1]) raise OSError(errno.ECONNREFUSED, "boom") + # create_connection() drives the connect through _sock_connect_fast + # (the loop-thread/vibeio fast path), mirroring how uvloop's + # create_connection bypasses the public sock_connect(). with mock.patch("socket.getaddrinfo", new=fake_getaddrinfo): with mock.patch.object( - rsloop.Loop, "sock_connect", new=fake_sock_connect + rsloop.Loop, "_sock_connect_fast", new=fake_sock_connect ): with self.assertRaises(OSError): await loop.create_connection( From 66b1871d64764bcb4536347cc5bb78b613ef5625 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 21:04:36 +0000 Subject: [PATCH 13/21] refactor loop core and runtime to utilize thread-local storage for I/O tasks and wake mechanisms --- src/loop_core.rs | 329 +++++++++++++++++++++++++++++--------- src/loop_core/commands.rs | 2 - src/python_api.rs | 130 ++++++++++----- src/python_names.rs | 5 + src/runtime.rs | 26 +-- src/stream_transport.rs | 82 ++++++---- 6 files changed, 401 insertions(+), 173 deletions(-) diff --git a/src/loop_core.rs b/src/loop_core.rs index b13b478..3c3dbca 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -1,11 +1,13 @@ -use std::cell::Cell; +use std::cell::{Cell, RefCell}; use std::collections::{HashMap, VecDeque}; use std::fmt; +use std::future::Future; use std::ops::DerefMut; +use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc; use std::sync::{Arc, Mutex}; -use std::task::Waker; +use std::task::{Context, Poll, Waker}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -15,9 +17,128 @@ use crate::errors::handle_callback_error; use crate::fd_ops::RawFd; use crate::runtime::run_runtime_thread; use crossbeam_channel::Sender; +use futures::task::AtomicWaker; use pyo3::prelude::*; use pyo3::types::{PyDict, PySet, PyTuple}; +thread_local! { + /// Per-loop vibeio runtime hosted on the loop thread. Keyed by `LoopCore` + /// pointer so a thread that runs several loops sequentially keeps each + /// loop's spawned I/O tasks alive across `run_until_complete` calls. The + /// runtime is `!Send`, which is why it lives in thread-local storage rather + /// than on `LoopCore`; asyncio's contract that a loop only runs on one + /// thread makes that safe. + /// `ManuallyDrop` so a runtime is only ever dropped explicitly in + /// `LoopCore::close` (on the loop thread, while vibeio's own thread-locals + /// are still alive). If a loop is never closed, its runtime is leaked at + /// thread exit rather than dropped — dropping during TLS destruction trips + /// an AccessError panic inside vibeio's `Runtime::drop`. + static LOOP_RUNTIMES: RefCell>> = + RefCell::new(HashMap::new()); + + /// Handles for cancellable I/O tasks (accept loops, socket readers) spawned + /// on the loop runtime, keyed by loop pointer then by fd. A separate + /// thread-local from `LOOP_RUNTIMES` so registering a task (`borrow_mut` + /// here) never conflicts with the park holding an immutable `LOOP_RUNTIMES` + /// borrow across `block_on`. `JoinHandle` is `!Send`, so this must live in + /// TLS on the loop thread. + static IO_TASKS: RefCell>>> = + RefCell::new(HashMap::new()); +} + +/// Cross-thread wake state for the loop thread, kept separate from `LoopCore` +/// so it stays `Ungil` (contains no Python objects) and can therefore be held +/// by a future driven under `py.detach`. `ready_pending` is the "ready queue +/// non-empty" flag the park future polls; `ready_waker` holds the loop thread's +/// task waker while it is parked. Replaces the old mpsc wake channel. +pub struct LoopWake { + ready_pending: AtomicBool, + ready_waker: AtomicWaker, +} + +impl LoopWake { + fn new() -> Self { + Self { + ready_pending: AtomicBool::new(false), + ready_waker: AtomicWaker::new(), + } + } + + /// Marks the ready queue non-empty and wakes the parked loop thread. Cheap + /// and idempotent while a wake is already pending. + #[inline] + pub fn signal(&self) { + if !self.ready_pending.swap(true, Ordering::AcqRel) { + self.ready_waker.wake(); + } + } +} + +/// The park primitive for the loop thread. `run_forever` drives this to +/// completion via `runtime.block_on` with the GIL released, so vibeio's reactor +/// (`driver.wait`) runs on the loop thread instead of a separate one. It +/// resolves when a ready item is enqueued (`LoopWake::signal`) or the +/// signal-poll timeout elapses, at which point `run_forever` re-acquires the +/// GIL and drains. +struct WaitForWake { + wake: Arc, + sleep: Pin>, +} + +impl WaitForWake { + fn new(wake: Arc, timeout: Duration) -> Self { + Self { + wake, + sleep: Box::pin(vibeio::time::Sleep::new(timeout)), + } + } +} + +/// Releases the GIL for the lifetime of the guard, re-acquiring it on drop +/// (including on unwind). This is the raw-FFI equivalent of +/// `Py_BEGIN_ALLOW_THREADS`/`Py_END_ALLOW_THREADS`. We use it instead of +/// `Python::detach` because that requires the closure to be `Ungil` (which is +/// `Send` on this build), and vibeio's runtime and timer futures are `!Send`. +/// Safe here because the GIL is released and re-acquired on the same thread and +/// no `Python`/`Py` token is moved across the boundary. +struct GilSuspend(*mut pyo3::ffi::PyThreadState); + +impl GilSuspend { + #[inline] + fn new() -> Self { + // SAFETY: called while the GIL is held; detaches this thread's state. + Self(unsafe { pyo3::ffi::PyEval_SaveThread() }) + } +} + +impl Drop for GilSuspend { + #[inline] + fn drop(&mut self) { + // SAFETY: re-attaches the thread state saved in `new` on the same thread. + unsafe { pyo3::ffi::PyEval_RestoreThread(self.0) }; + } +} + +impl Future for WaitForWake { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if self.wake.ready_pending.load(Ordering::Acquire) { + return Poll::Ready(()); + } + // Register before the second check so a wake that races with + // registration is never lost. + self.wake.ready_waker.register(cx.waker()); + if self.wake.ready_pending.load(Ordering::Acquire) { + return Poll::Ready(()); + } + if self.sleep.as_mut().poll(cx).is_ready() { + return Poll::Ready(()); + } + Poll::Pending + } +} + mod commands; pub use commands::{ LoopCommand, LoopFutureCommand, LoopIoCommand, LoopRunCommand, LoopSignalCommand, @@ -46,20 +167,6 @@ thread_local! { }; } -// Bounded busy-wait before parking on the wake channel. Cross-thread wakeups -// (transport reader workers, the runtime thread) otherwise pay a full condvar -// sleep/wake cycle, which dominates ping-pong style I/O latency. -fn wake_spin_window() -> Duration { - static WINDOW: std::sync::OnceLock = std::sync::OnceLock::new(); - *WINDOW.get_or_init(|| { - let micros = std::env::var("RSLOOP_WAKE_SPIN_US") - .ok() - .and_then(|value| value.trim().parse::().ok()) - .unwrap_or(40); - Duration::from_micros(micros.min(1_000)) - }) -} - #[derive(Debug)] pub enum LoopCoreError { Closed, @@ -102,8 +209,6 @@ pub struct FdWatch { struct ActiveReadyDispatch { pending_ready: Arc>>, - wake_tx: std::sync::mpsc::Sender<()>, - wake_pending: Arc, } pub struct LoopState { @@ -154,6 +259,9 @@ pub struct LoopCore { runtime_thread: Mutex>>, runtime_waker: Mutex>, active_ready_dispatch: Mutex>, + // Wakes the loop thread when a producer enqueues a ready item. Held in an + // Arc so the park future (`WaitForWake`) can own a clone under `py.detach`. + wake: Arc, } impl LoopCore { @@ -169,6 +277,7 @@ impl LoopCore { runtime_thread: Mutex::new(None), runtime_waker: Mutex::new(None), active_ready_dispatch: Mutex::new(None), + wake: Arc::new(LoopWake::new()), }); let thread_core = Arc::clone(&core); @@ -300,7 +409,7 @@ impl LoopCore { } #[profiling::function] - pub fn run_forever(&self, py: Python<'_>, loop_obj: Py) -> PyResult<()> { + pub fn run_forever(self: &Arc, py: Python<'_>, loop_obj: Py) -> PyResult<()> { const SIGNAL_POLL_INTERVAL: Duration = Duration::from_millis(50); { @@ -319,10 +428,8 @@ impl LoopCore { state.stopping = false; } - let (wake_tx, wake_rx) = std::sync::mpsc::channel(); - let wake_rx = Arc::new(Mutex::new(wake_rx)); let pending_ready = Arc::new(Mutex::new(VecDeque::new())); - let wake_pending = Arc::new(AtomicBool::new(false)); + self.wake.ready_pending.store(false, Ordering::Release); { let mut active_dispatch = self .active_ready_dispatch @@ -330,17 +437,31 @@ impl LoopCore { .expect("poisoned active ready dispatch"); *active_dispatch = Some(ActiveReadyDispatch { pending_ready: Arc::clone(&pending_ready), - wake_tx: wake_tx.clone(), - wake_pending: Arc::clone(&wake_pending), }); } self.send_command(LoopCommand::Run(LoopRunCommand::EnterRun { pending_ready: Arc::clone(&pending_ready), - wake_tx, - wake_pending: Arc::clone(&wake_pending), })) .map_err(|err| pyo3::exceptions::PyRuntimeError::new_err(err.to_string()))?; + // The vibeio runtime that drives I/O for this loop lives on this (the + // loop) thread. Parking below runs its reactor via `block_on`, so I/O + // readiness and Python callbacks share one thread. + let loop_runtime_key = Arc::as_ptr(self) as usize; + LOOP_RUNTIMES.with(|runtimes| { + runtimes + .borrow_mut() + .entry(loop_runtime_key) + .or_insert_with(|| { + std::mem::ManuallyDrop::new( + vibeio::RuntimeBuilder::new() + .enable_timer(true) + .build() + .expect("failed to initialize loop-thread vibeio runtime"), + ) + }); + }); + ensure_running_loop(py, &loop_obj)?; self.mark_runtime_thread(); let mut local_ready = VecDeque::new(); @@ -348,7 +469,6 @@ impl LoopCore { let mut pending_signal_error: Option = None; let mut ready_batch = VecDeque::new(); - let mut spin_next_wait = false; let run_result = loop { self.set_ready_drain_active(true); @@ -360,7 +480,7 @@ impl LoopCore { // pushing, so the pending queue only needs to be locked // when the flag is set; a hot chain of locally scheduled // callbacks otherwise skips the mutex entirely. - if wake_pending.load(Ordering::Acquire) { + if self.wake.ready_pending.load(Ordering::Acquire) { let mut pending = pending_ready.lock().expect("poisoned pending ready queue"); if !pending.is_empty() { @@ -372,7 +492,7 @@ impl LoopCore { } } if pending.is_empty() { - wake_pending.store(false, Ordering::Release); + self.wake.ready_pending.store(false, Ordering::Release); } } @@ -397,7 +517,6 @@ impl LoopCore { let item = ready_batch .pop_front() .expect("ready batch was checked as non-empty"); - spin_next_wait = true; match item { ReadyItem::Stop => { profiling::scope!("ready.stop"); @@ -514,51 +633,22 @@ impl LoopCore { continue; } - let wake_rx = Arc::clone(&wake_rx); - let spin_window = if std::mem::take(&mut spin_next_wait) { - wake_spin_window() - } else { - Duration::ZERO - }; - let spin_wake_pending = Arc::clone(&wake_pending); - match py.detach(move || { - // Busy-wait briefly for the next wakeup while work is flowing; - // this skips the condvar round trip that dominates request/ - // response latency. If nothing shows up, park as before. - if !spin_window.is_zero() { - let spin_deadline = Instant::now() + spin_window; - 'spin: loop { - for _ in 0..64 { - if spin_wake_pending.load(Ordering::Acquire) { - // Consume at most one queued wake token so - // tokens for wakeups observed via the spin - // path cannot accumulate. Missing a token here - // is fine: the drain loop re-reads the pending - // queue directly. - let _ = wake_rx.lock().expect("poisoned wake receiver").try_recv(); - return Ok(()); - } - std::hint::spin_loop(); - } - if Instant::now() >= spin_deadline { - break 'spin; - } - } - } - - wake_rx - .lock() - .expect("poisoned wake receiver") - .recv_timeout(SIGNAL_POLL_INTERVAL) - }) { - Ok(()) => {} - Err(mpsc::RecvTimeoutError::Timeout) => {} - Err(mpsc::RecvTimeoutError::Disconnected) => { - break Err(pyo3::exceptions::PyRuntimeError::new_err( - "event loop runtime terminated unexpectedly", - )); - } + // Park by driving this loop's vibeio runtime with the GIL released, + // so the reactor's `driver.wait` runs on the loop thread. Resolves + // when a ready item is enqueued (cross-thread or same-thread) or the + // signal-poll timeout elapses; either way we loop back and drain. + let wait = WaitForWake::new(Arc::clone(&self.wake), SIGNAL_POLL_INTERVAL); + { + let _gil = GilSuspend::new(); + LOOP_RUNTIMES.with(|runtimes| { + let runtimes = runtimes.borrow(); + let runtime = runtimes + .get(&loop_runtime_key) + .expect("loop runtime missing"); + runtime.block_on(wait); + }); } + let _ = py; }; self.set_ready_drain_active(false); @@ -624,6 +714,24 @@ impl LoopCore { state.closed = true; } + // Drop this loop's on-thread vibeio runtime now, while the loop thread + // (and vibeio's own thread-locals) are still alive. Letting it drop + // during thread destruction trips a TLS-access panic. `close()` runs on + // the loop thread with the loop stopped, so no `block_on` is active and + // the runtime holds no in-flight tasks yet (I/O still on the runtime + // thread in this phase). + let runtime_key = self as *const LoopCore as usize; + // Drop tracked task handles first (detaches them), then the runtime, + // whose drop tears down any still-registered tasks. + IO_TASKS.with(|tasks| { + tasks.borrow_mut().remove(&runtime_key); + }); + LOOP_RUNTIMES.with(|runtimes| { + if let Some(runtime) = runtimes.borrow_mut().remove(&runtime_key) { + drop(std::mem::ManuallyDrop::into_inner(runtime)); + } + }); + self.send_command(LoopCommand::Close)?; if let Some(handle) = self .runtime_thread @@ -760,6 +868,75 @@ impl LoopCore { *self.runtime_waker.lock().expect("poisoned runtime waker") = waker; } + /// Marks the ready queue non-empty and wakes the parked loop thread. Used by + /// cross-thread ready producers (the transitional runtime thread, signal and + /// transport workers). + #[inline] + pub(crate) fn signal_ready(&self) { + self.wake.signal(); + } + + /// Spawns a detached I/O task on this loop's on-thread vibeio runtime. Must + /// be called on the loop thread (asyncio contract). The task begins running + /// the next time the loop parks in `block_on`; its completions push ready + /// items and wake the loop **on the same thread**, with no cross-thread hop. + /// Returns `false` if the loop has no runtime yet (spawned before first run). + pub(crate) fn spawn_io(&self, future: F) -> bool + where + F: Future + 'static, + { + let key = self as *const LoopCore as usize; + LOOP_RUNTIMES.with(|runtimes| { + let runtimes = runtimes.borrow(); + match runtimes.get(&key) { + // Detach the JoinHandle: the task manages its own lifetime. + Some(runtime) => { + std::mem::drop(runtime.spawn(future)); + true + } + None => false, + } + }) + } + + /// Spawns a cancellable I/O task (accept loop / socket reader) on this loop's + /// runtime, tracked by `fd` so `stop_io_task` can cancel it. Any existing + /// task registered for `fd` is cancelled first. Must run on the loop thread. + /// Returns `false` if the loop has no runtime yet. + pub(crate) fn spawn_io_tracked(&self, fd: RawFd, future: F) -> bool + where + F: Future + 'static, + { + let key = self as *const LoopCore as usize; + let handle = + LOOP_RUNTIMES.with(|runtimes| runtimes.borrow().get(&key).map(|rt| rt.spawn(future))); + match handle { + Some(handle) => { + IO_TASKS.with(|tasks| { + if let Some(old) = tasks.borrow_mut().entry(key).or_default().insert(fd, handle) + { + old.cancel(); + } + }); + true + } + None => false, + } + } + + /// Cancels the tracked I/O task registered for `fd`, if any. Must run on the + /// loop thread. + pub(crate) fn stop_io_task(&self, fd: RawFd) { + let key = self as *const LoopCore as usize; + IO_TASKS.with(|tasks| { + if let Some(map) = tasks.borrow_mut().get_mut(&key) + && let Some(handle) = map.remove(&fd) + { + handle.cancel(); + } + }); + } + #[inline] pub(crate) fn install_local_ready_queue(&self, ready: *mut VecDeque) { ACTIVE_LOOP_TLS.with(|tls| tls.ready_queue.set(ready)); @@ -1021,9 +1198,7 @@ impl LoopCore { .lock() .expect("poisoned pending ready queue") .push_back(item); - if !dispatch.wake_pending.swap(true, Ordering::AcqRel) { - let _ = dispatch.wake_tx.send(()); - } + self.wake.signal(); Ok(()) } } diff --git a/src/loop_core/commands.rs b/src/loop_core/commands.rs index 403e373..df2ec10 100644 --- a/src/loop_core/commands.rs +++ b/src/loop_core/commands.rs @@ -68,8 +68,6 @@ pub enum LoopCommand { pub enum LoopRunCommand { EnterRun { pending_ready: Arc>>, - wake_tx: std::sync::mpsc::Sender<()>, - wake_pending: Arc, }, FinishRun { done_tx: std::sync::mpsc::Sender<()>, diff --git a/src/python_api.rs b/src/python_api.rs index c9af459..ba27795 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -810,26 +810,46 @@ fn is_connect_in_progress_errno(errno: i32) -> bool { errno == libc::EINPROGRESS || errno == libc::EALREADY || errno == libc::EWOULDBLOCK } -/// True for an AF_INET/AF_INET6 SOCK_STREAM socket, the sockets the vibeio -/// connect fast path handles. A type carrying extra flag bits declines the fast -/// path (falls back), which is never a correctness problem. +/// Attempts to initiate the connect via a direct `libc::connect` for a numeric +/// address, skipping Python's `socket.connect` (its dispatch, address parsing, +/// and — the expensive part — raising a `BlockingIOError` for EINPROGRESS on +/// every non-blocking connect). Returns `Some(errno)` when the libc path ran +/// (0 = connected immediately), or `None` when the address is not a plain +/// numeric literal and the caller must fall back to `socket.connect`. #[cfg(unix)] -fn socket_is_inet_stream(py: Python<'_>, sock: &Py) -> PyResult { - let socket_mod = py.import("socket")?; - let family: i32 = sock.getattr(py, "family")?.extract(py)?; - let af_inet: i32 = socket_mod.getattr("AF_INET")?.extract()?; - let af_inet6: i32 = socket_mod.getattr("AF_INET6")?.extract()?; - if family != af_inet && family != af_inet6 { - return Ok(false); +fn libc_connect_numeric(fd: fd_ops::RawFd, address: &Bound<'_, PyAny>) -> PyResult> { + let Ok(host_obj) = address.get_item(0) else { + return Ok(None); + }; + let Ok(host) = host_obj.extract::>() else { + return Ok(None); + }; + let Ok(ip) = host.parse::() else { + // Hostname, or scoped IPv6 ("fe80::1%eth0") std can't parse — fall back. + return Ok(None); + }; + let Ok(port) = address.get_item(1).and_then(|value| value.extract::()) else { + return Ok(None); + }; + let fd: libc::c_int = fd + .try_into() + .map_err(|_| PyRuntimeError::new_err("socket file descriptor out of range"))?; + let sockaddr = socket2::SockAddr::from(std::net::SocketAddr::new(ip, port)); + // SAFETY: `sockaddr` owns a valid `sockaddr` of the reported length; `fd` is + // the non-blocking socket the caller just built for this address family. + let rc = unsafe { libc::connect(fd, sockaddr.as_ptr().cast(), sockaddr.len()) }; + if rc == 0 { + Ok(Some(0)) + } else { + Ok(Some( + std::io::Error::last_os_error().raw_os_error().unwrap_or(0), + )) } - let sock_type: i32 = sock.getattr(py, "type")?.extract(py)?; - let sock_stream: i32 = socket_mod.getattr("SOCK_STREAM")?.extract()?; - Ok(sock_type == sock_stream) } /// Initiates a non-blocking connect on the loop thread and, when it does not -/// complete synchronously, hands the writability wait to the vibeio reactor via -/// `WatchConnect`. Returns the loop Future the caller awaits. +/// complete synchronously, hands the writability wait to the vibeio reactor on +/// this loop's own runtime. Returns the loop Future the caller awaits. #[cfg(unix)] fn fast_sock_connect<'py>( slf: &Py, @@ -837,33 +857,52 @@ fn fast_sock_connect<'py>( sock: Py, address: Py, ) -> PyResult> { - let future = slf.call_method0(py, "create_future")?; + let future = crate::python_names::call_method0( + py, + slf.bind(py).as_any(), + crate::python_names::create_future(py), + )?; let future_bound = future.bind(py).clone(); let fd = fd_ops::fileobj_to_fd(py, sock.bind(py))?; - match sock.call_method1(py, "connect", (address.bind(py),)) { - Ok(_) => { - future_bound.call_method1("set_result", (py.None(),))?; - return Ok(future_bound); - } - Err(err) => { - if is_already_connected_socket_error(py, &err)? { - future_bound.call_method1("set_result", (py.None(),))?; - return Ok(future_bound); - } - if !fd_ops::is_retryable_socket_error(py, &err)? { - future_bound.call_method1("set_exception", (err.into_value(py),))?; - return Ok(future_bound); - } - } + // Prefer a direct libc connect (numeric address); otherwise fall back to + // Python's socket.connect for hostnames / scoped IPv6. + let errno = match libc_connect_numeric(fd, address.bind(py))? { + Some(errno) => errno, + None => match sock.call_method1(py, "connect", (address.bind(py),)) { + Ok(_) => 0, + Err(err) => err + .value(py) + .getattr(crate::python_names::errno(py)) + .ok() + .and_then(|value| value.extract::().ok()) + .unwrap_or(0), + }, + }; + + if errno == 0 || is_already_connected_errno(errno) { + future_bound.call_method1("set_result", (py.None(),))?; + return Ok(future_bound); + } + if !is_connect_in_progress_errno(errno) { + let message = std::io::Error::from_raw_os_error(errno).to_string(); + let oserror = pyo3::exceptions::PyOSError::new_err((errno, message)).into_value(py); + future_bound.call_method1("set_exception", (oserror,))?; + return Ok(future_bound); } + // Connect is in progress: watch for writability on this loop's own reactor + // (loop thread), so the completion is delivered without a cross-thread wake. let core = slf.borrow(py).core.clone(); - core.send_command(LoopCommand::Io(LoopIoCommand::WatchConnect { + if !core.spawn_io(crate::stream_transport::run_connect_watch_task( + Arc::clone(&core), fd, - future: future.clone_ref(py), - })) - .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + future.clone_ref(py), + )) { + return Err(PyRuntimeError::new_err( + "event loop is not running; cannot start connect watch", + )); + } Ok(future_bound) } @@ -2880,17 +2919,22 @@ impl PyLoop { sock: Py, address: Py, ) -> PyResult> { + // Only `__loop_create_connection` calls this, and only with an + // INET/INET6 SOCK_STREAM socket it just built, so the family/type check + // (a `socket` import plus IntEnum property reads on every connection) is + // pure overhead — skip straight to the loop-thread connect on Unix. #[cfg(unix)] { - if socket_is_inet_stream(py, &sock)? { - return fast_sock_connect(&slf, py, sock, address); - } + return fast_sock_connect(&slf, py, sock, address); + } + #[cfg(not(unix))] + { + let locals = Self::task_locals(py, &slf)?; + pyo3_async_runtimes::async_std::future_into_py_with_locals(py, locals, async move { + connect_socket_to_address(sock, address).await?; + Ok(Python::attach(|py| py.None())) + }) } - let locals = Self::task_locals(py, &slf)?; - pyo3_async_runtimes::async_std::future_into_py_with_locals(py, locals, async move { - connect_socket_to_address(sock, address).await?; - Ok(Python::attach(|py| py.None())) - }) } #[pyo3(signature=(host, port, *, family=0, r#type=0, proto=0, flags=0))] diff --git a/src/python_names.rs b/src/python_names.rs index 703e051..a3c4693 100644 --- a/src/python_names.rs +++ b/src/python_names.rs @@ -29,6 +29,11 @@ pub(crate) fn create_future<'py>(py: Python<'py>) -> &'py Bound<'py, PyString> { interned(py, &NAME, "create_future") } +pub(crate) fn errno<'py>(py: Python<'py>) -> &'py Bound<'py, PyString> { + static NAME: OnceLock> = OnceLock::new(); + interned(py, &NAME, "errno") +} + pub(crate) fn done<'py>(py: Python<'py>) -> &'py Bound<'py, PyString> { static NAME: OnceLock> = OnceLock::new(); interned(py, &NAME, "done") diff --git a/src/runtime.rs b/src/runtime.rs index cf11435..440da53 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -37,8 +37,6 @@ struct RuntimeDispatcher { struct ActiveRun { pending_ready: Arc>>, - wake_tx: std::sync::mpsc::Sender<()>, - wake_pending: Arc, } #[cfg(unix)] @@ -268,17 +266,9 @@ impl RuntimeDispatcher { callback, }); } - LoopCommand::Run(LoopRunCommand::EnterRun { - pending_ready, - wake_tx, - wake_pending, - }) => { + LoopCommand::Run(LoopRunCommand::EnterRun { pending_ready }) => { profiling::scope!("runtime.cmd.enter_run"); - self.active_run = Some(ActiveRun { - pending_ready, - wake_tx, - wake_pending, - }); + self.active_run = Some(ActiveRun { pending_ready }); self.dispatch_ready_batch(); } LoopCommand::Run(LoopRunCommand::FinishRun { done_tx }) => { @@ -532,15 +522,9 @@ impl RuntimeDispatcher { .expect("poisoned pending ready queue"); pending.extend(self.ready_batch.drain(..)); drop(pending); - // Mirror the wake protocol used by try_enqueue_active_ready so the - // loop thread's spin-wait observes runtime-dispatched work (timers, - // fd watchers) too, and redundant wake tokens are not queued. - if !active_run - .wake_pending - .swap(true, std::sync::atomic::Ordering::AcqRel) - { - let _ = active_run.wake_tx.send(()); - } + // Wake the parked loop thread so runtime-dispatched work (timers, fd + // watchers) is observed; `signal_ready` coalesces redundant wakes. + self.core.signal_ready(); } fn finish_run(&mut self) { diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 72ef998..707f831 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -1976,9 +1976,7 @@ impl StreamTransportCore { self.detach_underlying_stream(py); let _ = self.writer_tx.send(WriterCommand::Stop); if let Some(fd) = self.runtime_socket_fd() { - let _ = self - .loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); + stop_socket_reader(self, fd); } for worker in self .workers @@ -2201,9 +2199,13 @@ impl ServerCore { .expect("poisoned accept fds") .drain(..) { - let _ = self - .loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StopServerAccept(fd))); + if self.loop_core.on_runtime_thread() { + self.loop_core.stop_io_task(fd); + } else { + let _ = self + .loop_core + .send_command(LoopCommand::Io(LoopIoCommand::StopServerAccept(fd))); + } } if let Some(path) = &self.cleanup_path { @@ -2276,13 +2278,23 @@ impl ServerCore { ServerListener::Unix(listener) => unix_raw_fd(listener.as_raw_fd()), }; accept_fds.push(fd); - let _ = + let server = Arc::clone(self); + // On the loop thread, host the accept loop directly on the loop's + // own runtime so accepted connections are delivered without a + // cross-thread hop. Off-thread callers fall back to the transitional + // runtime-thread command path. + if self.loop_core.on_runtime_thread() { self.loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StartServerAccept { + .spawn_io_tracked(fd, run_server_accept_task(server, listener)); + } else { + let _ = self.loop_core.send_command(LoopCommand::Io( + LoopIoCommand::StartServerAccept { fd, - server: Arc::clone(self), + server, listener, - })); + }, + )); + } } } } @@ -2344,10 +2356,7 @@ impl PyStreamTransport { self.core.flush_pending_direct_write(); self.core.set_closing(); if let Some(fd) = self.core.runtime_socket_fd() { - let _ = self - .core - .loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); + stop_socket_reader(&self.core, fd); } if self.core.direct_writer.is_none() { let _ = self.core.writer_tx.send(WriterCommand::Close); @@ -2371,10 +2380,7 @@ impl PyStreamTransport { self.core.discard_pending_direct_write(); self.core.set_closing(); if let Some(fd) = self.core.runtime_socket_fd() { - let _ = self - .core - .loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); + stop_socket_reader(&self.core, fd); } if self.core.direct_writer.is_none() { let _ = self.core.writer_tx.send(WriterCommand::Abort); @@ -3088,12 +3094,7 @@ pub fn spawn_tcp_transport( server.connection_opened(); } - core.loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { - fd: raw_fd, - core: Arc::clone(&core), - reader: ReaderTarget::Tcp(stream), - })) + spawn_socket_reader(raw_fd, Arc::clone(&core), ReaderTarget::Tcp(stream)) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; Ok(transport) } @@ -3142,12 +3143,7 @@ pub fn spawn_unix_transport( server.connection_opened(); } - core.loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { - fd: raw_fd, - core: Arc::clone(&core), - reader: ReaderTarget::Unix(stream), - })) + spawn_socket_reader(raw_fd, Arc::clone(&core), ReaderTarget::Unix(stream)) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; Ok(transport) } @@ -4051,6 +4047,32 @@ fn spin_read_stream( } } +/// Starts the socket reader for a stream transport. +/// +/// Readers stay on the transitional runtime thread. Hosting them on the loop +/// runtime was attempted but reverted: it segfaults for AF_UNIX socketpair +/// sockets and races `start_tls`, which reclaims the fd for a blocking handshake +/// while the reader still holds a non-blocking registration (EAGAIN). Reader +/// migration only benefits the traffic path (already ahead of uvloop), so the +/// risk is not worth it; accept loops — which are setup-relevant and terminate +/// cleanly on socket close — remain on the loop runtime. +fn spawn_socket_reader( + fd: fd_ops::RawFd, + core: Arc, + reader: ReaderTarget, +) -> Result<(), crate::loop_core::LoopCoreError> { + let loop_core = Arc::clone(&core.loop_core); + loop_core.send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { fd, core, reader })) +} + +/// Stops the socket reader for `fd` via the runtime-thread command path (readers +/// are hosted there; see `spawn_socket_reader`). +fn stop_socket_reader(core: &StreamTransportCore, fd: fd_ops::RawFd) { + let _ = core + .loop_core + .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); +} + pub(crate) async fn run_tcp_socket_reader_task( core: Arc, stream: StdTcpStream, From de5905ee5bc80a3f10288fc79ca906af5289d40f Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 21:51:57 +0000 Subject: [PATCH 14/21] implement bounded busy-wait before parking in block_on to reduce latency and improve performance --- src/loop_core.rs | 72 ++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 60 insertions(+), 12 deletions(-) diff --git a/src/loop_core.rs b/src/loop_core.rs index 3c3dbca..83c1786 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -119,6 +119,27 @@ impl Drop for GilSuspend { } } +/// Bounded busy-wait before parking in `block_on`. Cross-thread wakeups (reader +/// worker threads, the transitional runtime thread) otherwise pay the full +/// `driver.wait` park + interrupt round-trip, which dominates request/response +/// ping-pong latency and inflates its variance. Env-tunable via +/// `RSLOOP_WAKE_SPIN_US` (0 disables). +fn wake_spin_window() -> Duration { + static WINDOW: std::sync::OnceLock = std::sync::OnceLock::new(); + *WINDOW.get_or_init(|| { + let micros = std::env::var("RSLOOP_WAKE_SPIN_US") + .ok() + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(50); + Duration::from_micros(micros.min(1_000)) + }) +} + +// After this many consecutive spin-caught wakeups, park in `block_on` anyway so +// this loop's own runtime tasks (accept loops, connect watches) and io_uring +// completions are still serviced under sustained same-connection traffic. +const MAX_CONSECUTIVE_SPINS: u32 = 64; + impl Future for WaitForWake { type Output = (); @@ -469,6 +490,8 @@ impl LoopCore { let mut pending_signal_error: Option = None; let mut ready_batch = VecDeque::new(); + let spin_window = wake_spin_window(); + let mut consecutive_spins: u32 = 0; let run_result = loop { self.set_ready_drain_active(true); @@ -633,20 +656,45 @@ impl LoopCore { continue; } - // Park by driving this loop's vibeio runtime with the GIL released, - // so the reactor's `driver.wait` runs on the loop thread. Resolves - // when a ready item is enqueued (cross-thread or same-thread) or the - // signal-poll timeout elapses; either way we loop back and drain. - let wait = WaitForWake::new(Arc::clone(&self.wake), SIGNAL_POLL_INTERVAL); + // Wait for the next wakeup with the GIL released. First spin briefly + // to catch an imminent cross-thread wake (reader worker / runtime + // thread) in user space — this keeps request/response ping-pong + // latency low and tight. On spin timeout (or after too many + // consecutive catches, to avoid starving this loop's own runtime + // tasks) park by driving the runtime: its `driver.wait` runs here on + // the loop thread and is interrupted by a cross-thread wake. { let _gil = GilSuspend::new(); - LOOP_RUNTIMES.with(|runtimes| { - let runtimes = runtimes.borrow(); - let runtime = runtimes - .get(&loop_runtime_key) - .expect("loop runtime missing"); - runtime.block_on(wait); - }); + let mut caught = false; + if !spin_window.is_zero() { + let spin_deadline = Instant::now() + spin_window; + 'spin: loop { + for _ in 0..64 { + if self.wake.ready_pending.load(Ordering::Acquire) { + caught = true; + break 'spin; + } + std::hint::spin_loop(); + } + if Instant::now() >= spin_deadline { + break 'spin; + } + } + } + + if caught && consecutive_spins < MAX_CONSECUTIVE_SPINS { + consecutive_spins += 1; + } else { + consecutive_spins = 0; + let wait = WaitForWake::new(Arc::clone(&self.wake), SIGNAL_POLL_INTERVAL); + LOOP_RUNTIMES.with(|runtimes| { + let runtimes = runtimes.borrow(); + let runtime = runtimes + .get(&loop_runtime_key) + .expect("loop runtime missing"); + runtime.block_on(wait); + }); + } } let _ = py; }; From 697c454cfd004366b22eeb703d563ae1d20438c3 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 22:07:34 +0000 Subject: [PATCH 15/21] add inlines --- src/fd_ops.rs | 3 +++ src/loop_core.rs | 1 + src/python_api.rs | 2 ++ 3 files changed, 6 insertions(+) diff --git a/src/fd_ops.rs b/src/fd_ops.rs index ff0161b..0deeecb 100644 --- a/src/fd_ops.rs +++ b/src/fd_ops.rs @@ -196,12 +196,14 @@ fn raw_fd_to_c_int(fd: RawFd) -> io::Result { /// A connect() attempt that is still completing in the background. #[cfg(unix)] +#[inline] pub fn is_connect_in_progress_errno(errno: i32) -> bool { errno == libc::EINPROGRESS || errno == libc::EALREADY || errno == libc::EWOULDBLOCK } /// The socket is already connected (a benign outcome for connect()). #[cfg(unix)] +#[inline] pub fn is_already_connected_errno(errno: i32) -> bool { errno == libc::EISCONN } @@ -209,6 +211,7 @@ pub fn is_already_connected_errno(errno: i32) -> bool { /// Reads the pending `SO_ERROR` for a socket via a direct `getsockopt`, so the /// connect-completion path resolves without acquiring the GIL. #[cfg(unix)] +#[inline] pub fn socket_so_error(fd: RawFd) -> io::Result { let fd = raw_fd_to_c_int(fd)?; let mut value: libc::c_int = 0; diff --git a/src/loop_core.rs b/src/loop_core.rs index 83c1786..e914aea 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -143,6 +143,7 @@ const MAX_CONSECUTIVE_SPINS: u32 = 64; impl Future for WaitForWake { type Output = (); + #[inline] fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { if self.wake.ready_pending.load(Ordering::Acquire) { return Poll::Ready(()); diff --git a/src/python_api.rs b/src/python_api.rs index ba27795..c4dc541 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -802,10 +802,12 @@ fn is_already_connected_socket_error(py: Python<'_>, err: &PyErr) -> PyResult bool { errno == libc::EISCONN || errno == WSAEISCONN } +#[inline] fn is_connect_in_progress_errno(errno: i32) -> bool { errno == libc::EINPROGRESS || errno == libc::EALREADY || errno == libc::EWOULDBLOCK } From 08656b3bb195fed0286e0b6103ae632d185a20aa Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Mon, 20 Jul 2026 22:31:19 +0000 Subject: [PATCH 16/21] update asyncio task creation to pass running loop for Python 3.8/3.9 compatibility --- src/python_api.rs | 48 +++++++++++++++++++++++++++++++---- src/python_api/ffi_helpers.rs | 1 + 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/src/python_api.rs b/src/python_api.rs index c4dc541..9feeb51 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -336,6 +336,7 @@ fn call_callable_noargs(py: Python<'_>, callable: &Py) -> PyResult, callable: &Py, @@ -464,7 +465,7 @@ pub(crate) fn try_fast_create_task( if !core.on_runtime_thread() || core.has_task_factory() { return Ok(None); } - create_asyncio_task_for_running_loop(py, coro).map(Some) + create_asyncio_task_for_running_loop(py, loop_obj.bind(py), coro).map(Some) } fn create_asyncio_task_for_loop( @@ -507,8 +508,29 @@ fn create_asyncio_task_for_loop( } } -fn create_asyncio_task_for_running_loop(py: Python<'_>, coro: Py) -> PyResult> { - call_callable_onearg(py, asyncio_task_cls(py)?, coro.bind(py)) +#[inline] +fn create_asyncio_task_for_running_loop( + py: Python<'_>, + loop_obj: &Bound<'_, PyAny>, + coro: Py, +) -> PyResult> { + let task_cls = asyncio_task_cls(py)?; + #[cfg(not(Py_3_10))] + { + // On Python 3.8/3.9, `Task.__init__` resolves the loop via the slower + // `get_event_loop()` when none is passed; supplying the running loop + // (accepted there without a deprecation warning) skips that per-task + // lookup — ~10% of create_task. 3.10+ resolves the running loop cheaply, + // so the extra kwarg would only add overhead there and is skipped. + let kwargs = PyDict::new(py); + kwargs.set_item("loop", loop_obj)?; + task_cls.call(py, (coro,), Some(&kwargs)) + } + #[cfg(Py_3_10)] + { + let _ = loop_obj; + call_callable_onearg(py, task_cls, coro.bind(py)) + } } fn create_asyncio_task_with_kwargs( @@ -1721,6 +1743,22 @@ impl PyLoop { eager_start: Option, kwargs: Option>, ) -> PyResult> { + // Hot path: a bare `create_task(coro)` with no name/context/eager_start/ + // kwargs — what asyncio.Task step scheduling and gather() always hit. + // Skip the Arc clone and the (cached) kwarg-support probe those extras + // would need, and go straight to constructing the Task. + let bare = name.is_none() + && context.is_none() + && eager_start.is_none() + && kwargs.as_ref().is_none_or(|kwargs| kwargs.bind(py).is_empty()); + if bare { + let loop_ref = slf.borrow(py); + if !loop_ref.core.has_task_factory() && loop_ref.core.on_runtime_thread() { + drop(loop_ref); + return create_asyncio_task_for_running_loop(py, slf.bind(py).as_any(), coro); + } + } + let core = Arc::clone(&slf.borrow(py).core); let task_kwarg_support = asyncio_task_kwarg_support(py)?; let extra_kwargs = kwargs @@ -1732,7 +1770,7 @@ impl PyLoop { || (eager_start.is_some() && task_kwarg_support.eager_start); if !core.has_task_factory() && !has_kwargs && core.on_runtime_thread() { - return create_asyncio_task_for_running_loop(py, coro); + return create_asyncio_task_for_running_loop(py, slf.bind(py).as_any(), coro); } let loop_obj = Self::as_py_any(py, &slf); @@ -1800,7 +1838,7 @@ impl PyLoop { let trim_source_traceback = core.get_debug(); if is_current_running_loop(py, &loop_obj)? { let created = if !has_kwargs { - create_asyncio_task_for_running_loop(py, coro)? + create_asyncio_task_for_running_loop(py, loop_obj.bind(py), coro)? } else { create_asyncio_task_with_kwargs( py, diff --git a/src/python_api/ffi_helpers.rs b/src/python_api/ffi_helpers.rs index 9171b40..154382e 100644 --- a/src/python_api/ffi_helpers.rs +++ b/src/python_api/ffi_helpers.rs @@ -11,6 +11,7 @@ pub(super) fn call_noargs(py: Python<'_>, callable: &Py) -> PyResult, callable: &Py, From f8ab38d64e7ac486d9f412247cd606a82f8bb05c Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Tue, 21 Jul 2026 02:49:09 -0700 Subject: [PATCH 17/21] ensure Python sockets remain open until TLS accept workers are woken to prevent TCP failure timeout on Windows --- src/stream_transport.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 707f831..147ce23 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -2183,8 +2183,10 @@ impl ServerCore { state.listeners.clear(); } - self.close_python_sockets(); - + // Blocking TLS accept workers are woken by connecting to their listening + // address. Keep the exposed Python socket alive until after that wake: + // on Windows duplicated sockets share listener state, so closing the + // Python handle first makes the wake connect pay the TCP failure timeout. for task in self .accept_tasks .lock() @@ -2208,6 +2210,8 @@ impl ServerCore { } } + self.close_python_sockets(); + if let Some(path) = &self.cleanup_path { let _ = fs::remove_file(path); } From f8a525d1b0d1de5be5d34cc6b70b25a34f155458 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Tue, 21 Jul 2026 03:04:23 -0700 Subject: [PATCH 18/21] implement client config caching and update TLS context management --- python/rsloop/_bootstrap.py | 14 ++++++++++ src/tls.rs | 54 +++++++++++++++++++++++++++++++++++-- tests/test_tls.py | 51 +++++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 2 deletions(-) diff --git a/python/rsloop/_bootstrap.py b/python/rsloop/_bootstrap.py index 7a3ee33..635145a 100644 --- a/python/rsloop/_bootstrap.py +++ b/python/rsloop/_bootstrap.py @@ -50,9 +50,14 @@ def mark_default_verify_paths(context): context.__dict__["_rsloop_use_default_verify_paths"] = True return context + def mark_tls_config_changed(context): + state = context.__dict__ + state["_rsloop_tls_generation"] = state.get("_rsloop_tls_generation", 0) + 1 + orig_create_default_context = __ssl.create_default_context orig_load_cert_chain = context_cls.load_cert_chain orig_load_default_certs = context_cls.load_default_certs + orig_load_verify_locations = context_cls.load_verify_locations orig_set_default_verify_paths = context_cls.set_default_verify_paths def create_default_context(*args, **kwargs): @@ -78,21 +83,30 @@ def load_cert_chain(self, certfile, keyfile=None, password=None): __os.fspath(keyfile) if keyfile is not None else __os.fspath(certfile) ) self.__dict__["_rsloop_key_password"] = password_value + mark_tls_config_changed(self) return result def load_default_certs(self, *args, **kwargs): result = orig_load_default_certs(self, *args, **kwargs) mark_default_verify_paths(self) + mark_tls_config_changed(self) + return result + + def load_verify_locations(self, *args, **kwargs): + result = orig_load_verify_locations(self, *args, **kwargs) + mark_tls_config_changed(self) return result def set_default_verify_paths(self): result = orig_set_default_verify_paths(self) mark_default_verify_paths(self) + mark_tls_config_changed(self) return result __ssl.create_default_context = create_default_context context_cls.load_cert_chain = load_cert_chain context_cls.load_default_certs = load_default_certs + context_cls.load_verify_locations = load_verify_locations context_cls.set_default_verify_paths = set_default_verify_paths context_cls._rsloop_tracking_installed = True diff --git a/src/tls.rs b/src/tls.rs index 6123f71..2270910 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -1,8 +1,10 @@ use std::sync::Arc; use std::time::Duration; +use pyo3::ffi::c_str; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyCapsuleMethods, PyDict}; use rustls::client::ClientConfig; use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}; use rustls::pki_types::{CertificateDer, ServerName, UnixTime}; @@ -17,6 +19,13 @@ use material::{ const DEFAULT_HANDSHAKE_TIMEOUT_SECS: f64 = 60.0; const DEFAULT_SHUTDOWN_TIMEOUT_SECS: f64 = 30.0; +const CLIENT_CONFIG_CACHE_KEY: &str = "_rsloop_client_config_cache"; + +struct CachedClientConfig { + generation: u64, + verify_mode: i32, + config: Arc, +} pub struct ClientTlsSettings { pub config: Arc, @@ -42,12 +51,12 @@ pub fn client_tls_settings( ) -> PyResult { let ssl_context = normalize_client_ssl_context(py, ssl)?; let hostname = resolve_server_hostname(py, &ssl_context, server_hostname)?; - let config = build_client_config(py, &ssl_context)?; + let config = cached_client_config(py, &ssl_context)?; let handshake_timeout = handshake_timeout(ssl_handshake_timeout)?; let shutdown_timeout = shutdown_timeout(ssl_shutdown_timeout)?; Ok(ClientTlsSettings { - config: Arc::new(config), + config, server_name: ServerName::try_from(hostname.clone()) .map_err(|_| PyValueError::new_err(format!("invalid server_hostname: {hostname}")))?, handshake_timeout, @@ -56,6 +65,47 @@ pub fn client_tls_settings( }) } +fn cached_client_config( + py: Python<'_>, + ssl_context: &Py, +) -> PyResult> { + let context_dict = ssl_context.bind(py).getattr("__dict__")?; + let context_dict = context_dict.cast::()?; + let generation = context_dict + .get_item("_rsloop_tls_generation")? + .and_then(|value| value.extract::().ok()) + .unwrap_or(0); + let verify_mode = verify_mode_value(py, ssl_context)?; + + if let Some(value) = context_dict.get_item(CLIENT_CONFIG_CACHE_KEY)? + && let Ok(capsule) = value.cast::() + && let Ok(pointer) = capsule.pointer_checked(Some(c_str!( + "rsloop._loop.client_config_cache" + ))) + { + // SAFETY: capsules stored under CLIENT_CONFIG_CACHE_KEY are created below + // with a boxed CachedClientConfig and remain owned by the SSLContext for + // the duration of this borrow under the GIL. + let cached = unsafe { pointer.cast::().as_ref() }; + if cached.generation == generation && cached.verify_mode == verify_mode { + return Ok(Arc::clone(&cached.config)); + } + } + + let config = Arc::new(build_client_config(py, ssl_context)?); + let cache = PyCapsule::new_with_value( + py, + CachedClientConfig { + generation, + verify_mode, + config: Arc::clone(&config), + }, + c_str!("rsloop._loop.client_config_cache"), + )?; + context_dict.set_item(CLIENT_CONFIG_CACHE_KEY, cache)?; + Ok(config) +} + pub fn server_tls_settings( py: Python<'_>, ssl: &Bound<'_, PyAny>, diff --git a/tests/test_tls.py b/tests/test_tls.py index 399123a..aecae50 100644 --- a/tests/test_tls.py +++ b/tests/test_tls.py @@ -85,6 +85,57 @@ def test_create_default_context_marks_default_verify_paths(self) -> None: context = ssl.create_default_context() self.assertTrue(context.__dict__.get("_rsloop_use_default_verify_paths")) + def test_client_config_cache_reuses_and_invalidates(self) -> None: + async def main() -> tuple[bool, bool]: + async def echo( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + writer.write(await reader.readexactly(1)) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + with tempfile.TemporaryDirectory() as tmpdir: + ca_cert_path, _, _ = make_cert_files(tmpdir) + server_ctx, client_ctx = make_ssl_contexts(tmpdir) + server = await asyncio.start_server( + echo, "127.0.0.1", 0, ssl=server_ctx + ) + port = server.sockets[0].getsockname()[1] + + async def connect_once() -> None: + reader, writer = await asyncio.open_connection( + "127.0.0.1", + port, + ssl=client_ctx, + server_hostname="localhost", + ) + writer.write(b"x") + await writer.drain() + self.assertEqual(await reader.readexactly(1), b"x") + writer.close() + await writer.wait_closed() + + try: + await connect_once() + first = client_ctx.__dict__["_rsloop_client_config_cache"] + await connect_once() + reused = client_ctx.__dict__["_rsloop_client_config_cache"] is first + + client_ctx.load_verify_locations(cafile=ca_cert_path) + await connect_once() + invalidated = ( + client_ctx.__dict__["_rsloop_client_config_cache"] is not first + ) + return reused, invalidated + finally: + server.close() + await server.wait_closed() + + self.assertEqual(rsloop.run(main()), (True, True)) + def test_create_connection_and_server_tls_round_trip(self) -> None: async def main() -> tuple[str, tuple[int, ...]]: loop = asyncio.get_running_loop() From 47e2978d1ed5d60f2ab751528a34e803968a0e69 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Tue, 21 Jul 2026 04:04:34 -0700 Subject: [PATCH 19/21] fork vibeio into current repo to make rsloop faster in canceling --- Cargo.lock | 2 - Cargo.toml | 2 +- src/loop_core.rs | 16 +- src/python_api.rs | 4 +- src/stream_transport.rs | 145 ++- src/tls.rs | 12 +- vendor/vibeio/.cargo-ok | 1 + vendor/vibeio/.cargo_vcs_info.json | 6 + vendor/vibeio/Cargo.lock | 352 ++++++ vendor/vibeio/Cargo.toml | 131 ++ vendor/vibeio/Cargo.toml.orig | 61 + vendor/vibeio/README.md | 89 ++ vendor/vibeio/build.rs | 42 + vendor/vibeio/src/blocking/default.rs | 37 + vendor/vibeio/src/blocking/mod.rs | 57 + vendor/vibeio/src/builder.rs | 136 +++ vendor/vibeio/src/driver/iocp.rs | 923 ++++++++++++++ vendor/vibeio/src/driver/mio.rs | 310 +++++ vendor/vibeio/src/driver/mock.rs | 79 ++ vendor/vibeio/src/driver/mod.rs | 595 +++++++++ vendor/vibeio/src/driver/uring.rs | 636 ++++++++++ vendor/vibeio/src/executor.rs | 892 ++++++++++++++ vendor/vibeio/src/fd_inner.rs | 177 +++ vendor/vibeio/src/fs/file.rs | 764 ++++++++++++ vendor/vibeio/src/fs/metadata.rs | 254 ++++ vendor/vibeio/src/fs/mod.rs | 1416 ++++++++++++++++++++++ vendor/vibeio/src/fs/open_options.rs | 329 +++++ vendor/vibeio/src/io/buf.rs | 458 +++++++ vendor/vibeio/src/io/mod.rs | 150 +++ vendor/vibeio/src/io/pipe.rs | 324 +++++ vendor/vibeio/src/io/splice.rs | 164 +++ vendor/vibeio/src/io/stdio.rs | 266 ++++ vendor/vibeio/src/io/util.rs | 226 ++++ vendor/vibeio/src/lib.rs | 103 ++ vendor/vibeio/src/net/mod.rs | 55 + vendor/vibeio/src/net/tcp/listener.rs | 469 ++++++++ vendor/vibeio/src/net/tcp/mod.rs | 128 ++ vendor/vibeio/src/net/tcp/stream.rs | 794 ++++++++++++ vendor/vibeio/src/net/udp.rs | 1534 ++++++++++++++++++++++++ vendor/vibeio/src/net/unix/listener.rs | 145 +++ vendor/vibeio/src/net/unix/mod.rs | 149 +++ vendor/vibeio/src/net/unix/stream.rs | 568 +++++++++ vendor/vibeio/src/op/accept.rs | 683 +++++++++++ vendor/vibeio/src/op/accept_unix.rs | 162 +++ vendor/vibeio/src/op/connect.rs | 531 ++++++++ vendor/vibeio/src/op/fsync.rs | 90 ++ vendor/vibeio/src/op/hard_link.rs | 93 ++ vendor/vibeio/src/op/io_util.rs | 76 ++ vendor/vibeio/src/op/mkdir.rs | 89 ++ vendor/vibeio/src/op/mod.rs | 272 +++++ vendor/vibeio/src/op/open.rs | 95 ++ vendor/vibeio/src/op/read.rs | 306 +++++ vendor/vibeio/src/op/readat.rs | 171 +++ vendor/vibeio/src/op/readiness.rs | 51 + vendor/vibeio/src/op/readv.rs | 366 ++++++ vendor/vibeio/src/op/recv.rs | 287 +++++ vendor/vibeio/src/op/recvfrom.rs | 458 +++++++ vendor/vibeio/src/op/rename.rs | 93 ++ vendor/vibeio/src/op/send.rs | 256 ++++ vendor/vibeio/src/op/sendto.rs | 442 +++++++ vendor/vibeio/src/op/splice.rs | 127 ++ vendor/vibeio/src/op/statx.rs | 119 ++ vendor/vibeio/src/op/symlink.rs | 92 ++ vendor/vibeio/src/op/unlink.rs | 90 ++ vendor/vibeio/src/op/waitpid.rs | 233 ++++ vendor/vibeio/src/op/write.rs | 289 +++++ vendor/vibeio/src/op/writeat.rs | 165 +++ vendor/vibeio/src/op/writev.rs | 363 ++++++ vendor/vibeio/src/process/mod.rs | 1064 ++++++++++++++++ vendor/vibeio/src/process/reaper.rs | 329 +++++ vendor/vibeio/src/signal/mod.rs | 32 + vendor/vibeio/src/signal/unix.rs | 526 ++++++++ vendor/vibeio/src/signal/windows.rs | 154 +++ vendor/vibeio/src/task.rs | 124 ++ vendor/vibeio/src/time/interval.rs | 145 +++ vendor/vibeio/src/time/mod.rs | 130 ++ vendor/vibeio/src/time/sleep.rs | 209 ++++ vendor/vibeio/src/time/timeout.rs | 104 ++ vendor/vibeio/src/timer.rs | 791 ++++++++++++ vendor/vibeio/src/util/async_wrap.rs | 423 +++++++ vendor/vibeio/src/util/mod.rs | 110 ++ 81 files changed, 23124 insertions(+), 17 deletions(-) create mode 100644 vendor/vibeio/.cargo-ok create mode 100644 vendor/vibeio/.cargo_vcs_info.json create mode 100644 vendor/vibeio/Cargo.lock create mode 100644 vendor/vibeio/Cargo.toml create mode 100644 vendor/vibeio/Cargo.toml.orig create mode 100644 vendor/vibeio/README.md create mode 100644 vendor/vibeio/build.rs create mode 100644 vendor/vibeio/src/blocking/default.rs create mode 100644 vendor/vibeio/src/blocking/mod.rs create mode 100644 vendor/vibeio/src/builder.rs create mode 100644 vendor/vibeio/src/driver/iocp.rs create mode 100644 vendor/vibeio/src/driver/mio.rs create mode 100644 vendor/vibeio/src/driver/mock.rs create mode 100644 vendor/vibeio/src/driver/mod.rs create mode 100644 vendor/vibeio/src/driver/uring.rs create mode 100644 vendor/vibeio/src/executor.rs create mode 100644 vendor/vibeio/src/fd_inner.rs create mode 100644 vendor/vibeio/src/fs/file.rs create mode 100644 vendor/vibeio/src/fs/metadata.rs create mode 100644 vendor/vibeio/src/fs/mod.rs create mode 100644 vendor/vibeio/src/fs/open_options.rs create mode 100644 vendor/vibeio/src/io/buf.rs create mode 100644 vendor/vibeio/src/io/mod.rs create mode 100644 vendor/vibeio/src/io/pipe.rs create mode 100644 vendor/vibeio/src/io/splice.rs create mode 100644 vendor/vibeio/src/io/stdio.rs create mode 100644 vendor/vibeio/src/io/util.rs create mode 100644 vendor/vibeio/src/lib.rs create mode 100644 vendor/vibeio/src/net/mod.rs create mode 100644 vendor/vibeio/src/net/tcp/listener.rs create mode 100644 vendor/vibeio/src/net/tcp/mod.rs create mode 100644 vendor/vibeio/src/net/tcp/stream.rs create mode 100644 vendor/vibeio/src/net/udp.rs create mode 100644 vendor/vibeio/src/net/unix/listener.rs create mode 100644 vendor/vibeio/src/net/unix/mod.rs create mode 100644 vendor/vibeio/src/net/unix/stream.rs create mode 100644 vendor/vibeio/src/op/accept.rs create mode 100644 vendor/vibeio/src/op/accept_unix.rs create mode 100644 vendor/vibeio/src/op/connect.rs create mode 100644 vendor/vibeio/src/op/fsync.rs create mode 100644 vendor/vibeio/src/op/hard_link.rs create mode 100644 vendor/vibeio/src/op/io_util.rs create mode 100644 vendor/vibeio/src/op/mkdir.rs create mode 100644 vendor/vibeio/src/op/mod.rs create mode 100644 vendor/vibeio/src/op/open.rs create mode 100644 vendor/vibeio/src/op/read.rs create mode 100644 vendor/vibeio/src/op/readat.rs create mode 100644 vendor/vibeio/src/op/readiness.rs create mode 100644 vendor/vibeio/src/op/readv.rs create mode 100644 vendor/vibeio/src/op/recv.rs create mode 100644 vendor/vibeio/src/op/recvfrom.rs create mode 100644 vendor/vibeio/src/op/rename.rs create mode 100644 vendor/vibeio/src/op/send.rs create mode 100644 vendor/vibeio/src/op/sendto.rs create mode 100644 vendor/vibeio/src/op/splice.rs create mode 100644 vendor/vibeio/src/op/statx.rs create mode 100644 vendor/vibeio/src/op/symlink.rs create mode 100644 vendor/vibeio/src/op/unlink.rs create mode 100644 vendor/vibeio/src/op/waitpid.rs create mode 100644 vendor/vibeio/src/op/write.rs create mode 100644 vendor/vibeio/src/op/writeat.rs create mode 100644 vendor/vibeio/src/op/writev.rs create mode 100644 vendor/vibeio/src/process/mod.rs create mode 100644 vendor/vibeio/src/process/reaper.rs create mode 100644 vendor/vibeio/src/signal/mod.rs create mode 100644 vendor/vibeio/src/signal/unix.rs create mode 100644 vendor/vibeio/src/signal/windows.rs create mode 100644 vendor/vibeio/src/task.rs create mode 100644 vendor/vibeio/src/time/interval.rs create mode 100644 vendor/vibeio/src/time/mod.rs create mode 100644 vendor/vibeio/src/time/sleep.rs create mode 100644 vendor/vibeio/src/time/timeout.rs create mode 100644 vendor/vibeio/src/timer.rs create mode 100644 vendor/vibeio/src/util/async_wrap.rs create mode 100644 vendor/vibeio/src/util/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 943c8d9..aa06876 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1183,8 +1183,6 @@ checksum = "7ba6f5989077681266825251a52748b8c1d8a4ad098cc37e440103d0ea717fc0" [[package]] name = "vibeio" version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63b968e30d91bc5c14d23f9af184f26ea7f0d886ab309d11b37148b244d7dfa4" dependencies = [ "crossbeam-queue", "futures-util", diff --git a/Cargo.toml b/Cargo.toml index be420ca..fd38693 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,7 +33,7 @@ signal-hook = "0.4" socket2 = "0.6" tokio = { version = "1", default-features = false, features = ["io-util"] } tracy-client = { version = "0.18", default-features = false, features = ["enable", "only-localhost", "sampling"], optional = true } -vibeio = { version = "0.2.15", default-features = false, features = ["time"] } +vibeio = { path = "vendor/vibeio", default-features = false, features = ["time"] } [target.'cfg(windows)'.dependencies] windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Networking_WinSock", "Win32_System_Pipes", "Win32_System_Threading"] } diff --git a/src/loop_core.rs b/src/loop_core.rs index e914aea..e08647d 100644 --- a/src/loop_core.rs +++ b/src/loop_core.rs @@ -770,10 +770,14 @@ impl LoopCore { // the runtime holds no in-flight tasks yet (I/O still on the runtime // thread in this phase). let runtime_key = self as *const LoopCore as usize; - // Drop tracked task handles first (detaches them), then the runtime, - // whose drop tears down any still-registered tasks. + // Cancel tracked tasks while the runtime and its driver are alive, then + // drop the runtime after cancellation has released pending operations. IO_TASKS.with(|tasks| { - tasks.borrow_mut().remove(&runtime_key); + if let Some(handles) = tasks.borrow_mut().remove(&runtime_key) { + for (_, handle) in handles { + handle.cancel(); + } + } }); LOOP_RUNTIMES.with(|runtimes| { if let Some(runtime) = runtimes.borrow_mut().remove(&runtime_key) { @@ -962,7 +966,11 @@ impl LoopCore { match handle { Some(handle) => { IO_TASKS.with(|tasks| { - if let Some(old) = tasks.borrow_mut().entry(key).or_default().insert(fd, handle) + if let Some(old) = tasks + .borrow_mut() + .entry(key) + .or_default() + .insert(fd, handle) { old.cancel(); } diff --git a/src/python_api.rs b/src/python_api.rs index 9feeb51..1c76329 100644 --- a/src/python_api.rs +++ b/src/python_api.rs @@ -1750,7 +1750,9 @@ impl PyLoop { let bare = name.is_none() && context.is_none() && eager_start.is_none() - && kwargs.as_ref().is_none_or(|kwargs| kwargs.bind(py).is_empty()); + && kwargs + .as_ref() + .is_none_or(|kwargs| kwargs.bind(py).is_empty()); if bare { let loop_ref = slf.borrow(py); if !loop_ref.core.has_task_factory() && loop_ref.core.on_runtime_thread() { diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 147ce23..625210b 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -26,6 +26,10 @@ use pyo3_async_runtimes::TaskLocals; use rustls::{ClientConnection, ServerConnection}; use socket2::Socket; use tokio::io::AsyncReadExt; +#[cfg(windows)] +use vibeio::io::AsyncRead as VibeAsyncRead; +#[cfg(windows)] +use vibeio::net::TcpStream as VibeTcpStream; use vibeio::net::{PollTcpStream as VibePollTcpStream, TcpListener as VibeTcpListener}; #[cfg(unix)] use vibeio::net::{PollUnixStream as VibePollUnixStream, UnixListener as VibeUnixListener}; @@ -4066,7 +4070,11 @@ fn spawn_socket_reader( reader: ReaderTarget, ) -> Result<(), crate::loop_core::LoopCoreError> { let loop_core = Arc::clone(&core.loop_core); - loop_core.send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { fd, core, reader })) + loop_core.send_command(LoopCommand::Io(LoopIoCommand::StartSocketReader { + fd, + core, + reader, + })) } /// Stops the socket reader for `fd` via the runtime-thread command path (readers @@ -4077,6 +4085,7 @@ fn stop_socket_reader(core: &StreamTransportCore, fd: fd_ops::RawFd) { .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); } +#[cfg(not(windows))] pub(crate) async fn run_tcp_socket_reader_task( core: Arc, stream: StdTcpStream, @@ -4131,6 +4140,140 @@ pub(crate) async fn run_tcp_socket_reader_task( } } +#[cfg(windows)] +pub(crate) async fn run_tcp_socket_reader_task( + core: Arc, + stream: StdTcpStream, +) { + profiling::scope!("stream.run_tcp_socket_reader_task"); + + // Server-side sockets retain the readiness reader. A server transport can + // duplicate its socket for direct bulk writes, and an outstanding IOCP read + // on the original socket materially slows that path. Client-side sockets + // use completion reads, avoiding one AFD poll per request/response turn. + if core.server_side { + match VibePollTcpStream::from_std(stream) { + Ok(reader) => { + run_windows_poll_tcp_reader( + core, + reader, + Vec::with_capacity(STREAM_READ_BUFFER_SIZE), + ) + .await; + } + Err(err) => core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( + err.to_string(), + ))), + } + return; + } + + let mut reader = match VibeTcpStream::from_std(stream) { + Ok(reader) => reader, + Err(err) => { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( + err.to_string(), + ))); + return; + } + }; + let mut buf = Vec::with_capacity(STREAM_READ_BUFFER_SIZE); + + loop { + if core.is_closing() { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); + return; + } + while !core.is_closing() && !core.is_reading() { + vibeio::time::sleep(Duration::from_millis(1)).await; + } + if core.is_closing() { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); + return; + } + + buf.clear(); + let (result, returned_buf) = VibeAsyncRead::read(&mut reader, buf).await; + buf = returned_buf; + match result { + Ok(0) => { + core.enqueue_pending_read_event(PendingReadEvent::Eof); + return; + } + Ok(read) => { + let saturated = read == buf.capacity(); + let data = take_async_read_data(&mut buf); + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); + + // Completion reads win on small request/response traffic. Once + // a buffer fills, switch this connection to readiness mode so + // sustained bulk transfer does not keep an overlapped receive + // outstanding alongside the direct writer socket. + if saturated { + match reader.into_poll() { + Ok(poll_reader) => { + run_windows_poll_tcp_reader(core, poll_reader, buf).await; + } + Err(err) => core.enqueue_pending_read_event( + PendingReadEvent::ConnectionLost(Some(err.to_string())), + ), + } + return; + } + } + Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue, + Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + Err(err) => { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( + err.to_string(), + ))); + return; + } + } + } +} + +#[cfg(windows)] +async fn run_windows_poll_tcp_reader( + core: Arc, + mut reader: VibePollTcpStream, + mut buf: Vec, +) { + loop { + if core.is_closing() { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); + return; + } + while !core.is_closing() && !core.is_reading() { + vibeio::time::sleep(Duration::from_millis(1)).await; + } + if core.is_closing() { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(None)); + return; + } + + buf.clear(); + match reader.read_buf(&mut buf).await { + Ok(0) => { + core.enqueue_pending_read_event(PendingReadEvent::Eof); + return; + } + Ok(_) => { + let data = take_async_read_data(&mut buf); + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); + } + Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue, + Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + Err(err) => { + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( + err.to_string(), + ))); + return; + } + } + } +} + #[cfg(unix)] pub(crate) async fn run_unix_socket_reader_task( core: Arc, diff --git a/src/tls.rs b/src/tls.rs index 2270910..f2f089a 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -1,8 +1,8 @@ use std::sync::Arc; use std::time::Duration; -use pyo3::ffi::c_str; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::ffi::c_str; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyCapsuleMethods, PyDict}; use rustls::client::ClientConfig; @@ -65,10 +65,7 @@ pub fn client_tls_settings( }) } -fn cached_client_config( - py: Python<'_>, - ssl_context: &Py, -) -> PyResult> { +fn cached_client_config(py: Python<'_>, ssl_context: &Py) -> PyResult> { let context_dict = ssl_context.bind(py).getattr("__dict__")?; let context_dict = context_dict.cast::()?; let generation = context_dict @@ -79,9 +76,8 @@ fn cached_client_config( if let Some(value) = context_dict.get_item(CLIENT_CONFIG_CACHE_KEY)? && let Ok(capsule) = value.cast::() - && let Ok(pointer) = capsule.pointer_checked(Some(c_str!( - "rsloop._loop.client_config_cache" - ))) + && let Ok(pointer) = + capsule.pointer_checked(Some(c_str!("rsloop._loop.client_config_cache"))) { // SAFETY: capsules stored under CLIENT_CONFIG_CACHE_KEY are created below // with a boxed CachedClientConfig and remain owned by the SSLContext for diff --git a/vendor/vibeio/.cargo-ok b/vendor/vibeio/.cargo-ok new file mode 100644 index 0000000..5f8b795 --- /dev/null +++ b/vendor/vibeio/.cargo-ok @@ -0,0 +1 @@ +{"v":1} \ No newline at end of file diff --git a/vendor/vibeio/.cargo_vcs_info.json b/vendor/vibeio/.cargo_vcs_info.json new file mode 100644 index 0000000..9547fa2 --- /dev/null +++ b/vendor/vibeio/.cargo_vcs_info.json @@ -0,0 +1,6 @@ +{ + "git": { + "sha1": "170b58bd9b63d10e806d36cf22ce9ebe1cd1d033" + }, + "path_in_vcs": "vibeio" +} \ No newline at end of file diff --git a/vendor/vibeio/Cargo.lock b/vendor/vibeio/Cargo.lock new file mode 100644 index 0000000..1e719ac --- /dev/null +++ b/vendor/vibeio/Cargo.lock @@ -0,0 +1,352 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bytes" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "io-uring" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "595a0399f411a508feb2ec1e970a4a30c249351e30208960d58298de8660b0e5" +dependencies = [ + "bitflags", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "log" +version = "0.4.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oneshot" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rusty_pool" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ed36cdb20de66d89a17ea04b8883fc7a386f2cf877aaedca5005583ce4876ff" +dependencies = [ + "crossbeam-channel", + "futures", + "futures-channel", + "futures-executor", + "num_cpus", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "pin-project-lite", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "vibeio" +version = "0.2.15" +dependencies = [ + "async-channel", + "crossbeam-queue", + "futures-util", + "io-uring", + "libc", + "mio", + "once_cell", + "oneshot", + "rusty_pool", + "slab", + "tokio", + "windows-sys", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] diff --git a/vendor/vibeio/Cargo.toml b/vendor/vibeio/Cargo.toml new file mode 100644 index 0000000..2d32dee --- /dev/null +++ b/vendor/vibeio/Cargo.toml @@ -0,0 +1,131 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +name = "vibeio" +version = "0.2.15" +authors = ["Dorian Niemiec "] +build = "build.rs" +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "A high-performance, cross-platform asynchronous runtime for Rust" +documentation = "https://docs.rs/vibeio" +readme = "README.md" +keywords = [ + "async", + "io", +] +categories = [ + "network-programming", + "asynchronous", +] +license = "MIT" +repository = "https://github.com/ferronweb/vibeio" + +[package.metadata.docs.rs] +all-features = true +rustdoc-args = [ + "--cfg", + "docsrs", +] + +[features] +blocking-default = ["rusty_pool"] +default = [ + "fs", + "splice", + "time", + "pipe", + "stdio", + "signal", + "process", + "blocking-default", +] +fs = [] +pipe = [] +process = ["async-channel"] +signal = ["once_cell"] +splice = [] +stdio = [] +time = [] + +[lib] +name = "vibeio" +path = "src/lib.rs" + +[dependencies.async-channel] +version = "2.5.0" +optional = true + +[dependencies.crossbeam-queue] +version = "0.3" + +[dependencies.futures-util] +version = "0.3.32" + +[dependencies.mio] +version = "1.1.1" + +[dependencies.once_cell] +version = "1.21.3" +optional = true + +[dependencies.oneshot] +version = "0.2.1" +features = [ + "async", + "std", +] + +[dependencies.rusty_pool] +version = "0.7.0" +optional = true + +[dependencies.slab] +version = "0.4.12" + +[dependencies.tokio] +version = "1.49.0" +features = ["io-util"] + +[dev-dependencies] + +[target.'cfg(target_os = "linux")'.dependencies.io-uring] +version = "0.6" + +[target."cfg(unix)".dependencies.libc] +version = "0.2.176" + +[target."cfg(unix)".dependencies.mio] +version = "1.1.1" +features = [ + "os-poll", + "os-ext", +] + +[target."cfg(windows)".dependencies.windows-sys] +version = "0.61.2" +features = [ + "Wdk_Foundation", + "Wdk_Storage_FileSystem", + "Wdk_System_IO", + "Win32_Networking_WinSock", + "Win32_Foundation", + "Win32_Security", + "Win32_System_Console", + "Win32_System_IO", + "Win32_System_Threading", + "Win32_Storage_FileSystem", +] diff --git a/vendor/vibeio/Cargo.toml.orig b/vendor/vibeio/Cargo.toml.orig new file mode 100644 index 0000000..dff1823 --- /dev/null +++ b/vendor/vibeio/Cargo.toml.orig @@ -0,0 +1,61 @@ +[package] +name = "vibeio" +version = "0.2.15" +edition = "2021" +authors = ["Dorian Niemiec "] +description = "A high-performance, cross-platform asynchronous runtime for Rust" +license = "MIT" +readme = "../README.md" +keywords = ["async", "io"] +categories = ["network-programming", "asynchronous"] +documentation = "https://docs.rs/vibeio" +repository = "https://github.com/ferronweb/vibeio" + +[dependencies] +futures-util = "0.3.32" +mio = "1.1.1" +slab = "0.4.12" +tokio = { version = "1.49.0", features = ["io-util"] } +crossbeam-queue = "0.3" +oneshot = { version = "0.2.1", features = ["async", "std"] } +rusty_pool = { version = "0.7.0", optional = true } +once_cell = { version = "1.21.3", optional = true } +async-channel = { version = "2.5.0", optional = true } + +[target.'cfg(unix)'.dependencies] +mio = { version = "1.1.1", features = ["os-poll", "os-ext"] } +libc = "0.2.176" + +[target.'cfg(windows)'.dependencies] +windows-sys = { version = "0.61.2", features = [ + "Wdk_Foundation", + "Wdk_Storage_FileSystem", + "Wdk_System_IO", + "Win32_Networking_WinSock", + "Win32_Foundation", + "Win32_Security", + "Win32_System_Console", + "Win32_System_IO", + "Win32_System_Threading", + "Win32_Storage_FileSystem", +] } + +[target.'cfg(target_os = "linux")'.dependencies] +io-uring = "0.6" + +[dev-dependencies] + +[features] +default = ["fs", "splice", "time", "pipe", "stdio", "signal", "process", "blocking-default"] +fs = [] +splice = [] +pipe = [] +stdio = [] +process = ["async-channel"] +signal = ["once_cell"] +time = [] +blocking-default = ["rusty_pool"] + +[package.metadata.docs.rs] +all-features = true +rustdoc-args = ["--cfg", "docsrs"] diff --git a/vendor/vibeio/README.md b/vendor/vibeio/README.md new file mode 100644 index 0000000..2cb0321 --- /dev/null +++ b/vendor/vibeio/README.md @@ -0,0 +1,89 @@ +# vibeio + +A high-performance, cross-platform asynchronous runtime for Rust. + +`vibeio` provides an efficient I/O event loop that leverages the best available driver for each operating system: + +- **Linux** - uses `io_uring` for true asynchronous I/O. +- **Windows** - uses I/O Completion Ports (IOCP) for scalable I/O. +- **macOS / BSD / Others** - uses `kqueue` or `epoll` via `mio` for event notification. + +## Core features + +- **Networking** - asynchronous TCP, UDP, and Unix Domain Sockets. +- **File system** - asynchronous file operations. +- **Timers** - efficient timer and sleep functionality. +- **Signals** - handling of OS signals. +- **Process management** - spawning and managing child processes. +- **Blocking tasks** - offload CPU-intensive or blocking operations to a thread pool. + +## Concurrency model: thread-per-core + +`vibeio` is designed as a **single-threaded** runtime. To utilize multiple cores, you should employ a **thread-per-core** architecture, where a separate `Runtime` is pinned to each processor core. This approach minimizes synchronization overhead and maximizes cache locality. + +Shared state can be communicated between runtimes using message passing (e.g., channels) or shared atomic structures, but I/O resources are typically owned by the thread that created them. + +## But why not Tokio? + +Tokio is a popular asynchronous runtime for Rust, but it uses a **work-stealing** model and may introduce additional synchronization overhead when using it as a thread-per-core runtime. `vibeio` is more specialized for thread-per-core architectures that are optimized for low overhead and cache locality. + +## Getting started + +Add `vibeio` to your `Cargo.toml`: + +```toml +[dependencies] +vibeio = "0.2" +``` + +### Example: TCP echo server + +```rust +use vibeio::RuntimeBuilder; +use vibeio::net::TcpListener; + +fn main() -> std::io::Result<()> { + // 1. Build the runtime + let runtime = RuntimeBuilder::new() + .enable_timer(true) + .build()?; + + // 2. Run the main future + runtime.block_on(async { + let listener = TcpListener::bind("127.0.0.1:8080")?; + println!("Listening on 127.0.0.1:8080"); + + loop { + let (mut stream, _) = listener.accept().await?; + + vibeio::spawn(async move { + let (mut reader, mut writer) = vibeio::io::split(stream); + if let Err(e) = vibeio::io::copy(&mut reader, &mut writer).await { + eprintln!("Echo failed: {}", e); + } + }); + } + }) +} +``` + +## Feature flags + +The following features are available (most are enabled by default): + +- `fs` - enables asynchronous file system operations. +- `time` - enables time and timer functionality. +- `signal` - enables signal handling. +- `process` - enables child process management. +- `pipe` - enables pipe support. +- `stdio` - enables standard I/O support. +- `splice` - enables splice support (Linux). +- `blocking-default` - enables the default blocking thread pool. + +## "vibeio" name + +The name `vibeio` has two parts - "vibe" (from "vibe coding", as originally this project was coded with help of AI), and "io" (common suffix for Rust async runtime names). + +## License + +[MIT](./LICENSE) diff --git a/vendor/vibeio/build.rs b/vendor/vibeio/build.rs new file mode 100644 index 0000000..c0cad62 --- /dev/null +++ b/vendor/vibeio/build.rs @@ -0,0 +1,42 @@ +fn main() { + let musl = std::env::var("CARGO_CFG_TARGET_ENV").unwrap_or_default() == "musl"; + let musl_v1_2_3 = std::env::var("RUST_LIBC_UNSTABLE_MUSL_V1_2_3").is_ok(); + println!("cargo:rerun-if-env-changed=RUST_LIBC_UNSTABLE_MUSL_V1_2_3"); + println!("cargo:rustc-check-cfg=cfg(musl_v1_2_3)"); + if musl && musl_v1_2_3 { + println!("cargo:rustc-cfg=musl_v1_2_3"); + } + + println!("cargo:rustc-check-cfg=cfg(syscall_pipe2)"); + println!("cargo:rustc-check-cfg=cfg(syscall_accept4)"); + let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if [ + "linux", + "android", + "freebsd", + "illumos", + "solaris", + "emscripten", + "hurd", + "redox", + "netbsd", + "cygwin", + ] + .contains(&target_os.as_str()) + { + println!("cargo:rustc-cfg=syscall_pipe2"); + } + if [ + "freebsd", + "netbsd", + "emscripten", + "fuchsia", + "solaris", + "illumos", + "linux", + ] + .contains(&target_os.as_str()) + { + println!("cargo:rustc-cfg=syscall_accept4"); + } +} diff --git a/vendor/vibeio/src/blocking/default.rs b/vendor/vibeio/src/blocking/default.rs new file mode 100644 index 0000000..5c6fd74 --- /dev/null +++ b/vendor/vibeio/src/blocking/default.rs @@ -0,0 +1,37 @@ +use super::BlockingThreadPool; + +/// A default implementation of `BlockingThreadPool` using `rusty_pool` crate. +pub struct DefaultBlockingThreadPool { + inner: rusty_pool::ThreadPool, +} + +impl DefaultBlockingThreadPool { + /// Creates a new `DefaultBlockingThreadPool` with the default maximum number of threads. + #[inline] + pub fn new() -> Self { + Self::with_max_threads(512) + } + + /// Creates a new `DefaultBlockingThreadPool` with the specified maximum number of threads. + #[inline] + pub fn with_max_threads(num_threads: usize) -> Self { + Self { + inner: rusty_pool::Builder::new().max_size(num_threads).build(), + } + } +} + +impl BlockingThreadPool for DefaultBlockingThreadPool { + #[inline] + fn spawn(&self, task: Box) { + self.inner.execute(move || { + task(); + }); + } +} + +impl Default for DefaultBlockingThreadPool { + fn default() -> Self { + Self::new() + } +} diff --git a/vendor/vibeio/src/blocking/mod.rs b/vendor/vibeio/src/blocking/mod.rs new file mode 100644 index 0000000..26daa16 --- /dev/null +++ b/vendor/vibeio/src/blocking/mod.rs @@ -0,0 +1,57 @@ +//! Support for spawning blocking tasks in an async runtime. +//! +//! This module provides a trait for pluggable blocking thread pools and a default implementation +//! using a thread pool with automatically adjusted thread count and a work-stealing queue. + +#[cfg(feature = "blocking-default")] +mod default; + +use std::fmt; + +#[cfg(feature = "blocking-default")] +pub use default::*; + +/// Error returned when a blocking task fails to spawn. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SpawnBlockingError; + +impl fmt::Display for SpawnBlockingError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "failed to spawn blocking task") + } +} + +impl std::error::Error for SpawnBlockingError {} + +/// A trait for pluggable blocking thread pools. +/// +/// This trait allows users to provide their own implementation of a thread pool for executing +/// blocking tasks. The thread pool must be able to spawn tasks and shut down gracefully. +pub trait BlockingThreadPool: 'static { + /// Spawns a blocking task onto the thread pool. + /// + /// The task will be executed on one of the threads in the pool, and its output will be + /// returned back to `spawn_blocking`. + fn spawn(&self, task: Box); +} + +/// Spawns a blocking task onto a blocking thread pool. +/// +/// This function is a convenience wrapper around a blocking thread pool. +#[inline] +pub(crate) async fn spawn_blocking( + pool: &dyn BlockingThreadPool, + f: F, +) -> Result +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, +{ + let (tx, rx) = oneshot::async_channel::(); + let task: Box = Box::new(move || { + let _ = tx.send(f()); + }); + pool.spawn(task); + + rx.await.map_err(|_| SpawnBlockingError) +} diff --git a/vendor/vibeio/src/builder.rs b/vendor/vibeio/src/builder.rs new file mode 100644 index 0000000..18fa08d --- /dev/null +++ b/vendor/vibeio/src/builder.rs @@ -0,0 +1,136 @@ +#[cfg(feature = "blocking-default")] +use crate::blocking::DefaultBlockingThreadPool; +use crate::{blocking::BlockingThreadPool, driver::AnyDriver}; + +/// I/O driver selection for the async runtime. +/// +/// This enum allows choosing which I/O driver to use when building the runtime. +#[derive(Clone)] +pub enum DriverKind { + /// Uses the Mio driver for I/O operations (Unix only). + #[cfg(unix)] + Mio, + /// Uses the IOCP driver for completion-based I/O operations (Windows only). + #[cfg(windows)] + Iocp, + /// Uses the mock driver for testing purposes. + Mock, + /// Uses the io_uring driver (Linux only). + #[cfg(target_os = "linux")] + IoUring, + /// Uses a custom io_uring driver (Linux only). + #[cfg(target_os = "linux")] + IoUringCustom(io_uring::Builder), +} + +impl DriverKind { + /// Creates a new runtime I/O driver from this kind. + #[inline] + pub(crate) fn into_driver(self) -> Result { + match self { + #[cfg(unix)] + DriverKind::Mio => AnyDriver::new_mio(), + #[cfg(windows)] + DriverKind::Iocp => AnyDriver::new_iocp(), + DriverKind::Mock => Ok(AnyDriver::new_mock()), + #[cfg(target_os = "linux")] + DriverKind::IoUring => AnyDriver::new_uring(), + #[cfg(target_os = "linux")] + DriverKind::IoUringCustom(builder) => AnyDriver::new_uring_custom(builder), + } + } +} + +/// Builder for configuring and creating an async runtime. +/// +/// Provides a convenient way to configure the runtime's I/O driver +/// before building it. +/// +/// # Examples +/// +/// ``` +/// use vibeio::RuntimeBuilder; +/// +/// let runtime = RuntimeBuilder::new() +/// .build(); +/// ``` +pub struct RuntimeBuilder { + driver_kind: Option, + enable_timer: bool, + enable_fs_offload: bool, + blocking_pool: Option>, +} + +impl RuntimeBuilder { + /// Creates a new runtime builder with default configuration. + /// + /// By default, the builder will select the best available driver for the platform. + pub fn new() -> Self { + Self { + driver_kind: None, + enable_timer: false, + enable_fs_offload: false, + blocking_pool: None, + } + } + + /// Sets the I/O driver for the runtime. + pub fn driver(mut self, driver_kind: DriverKind) -> Self { + self.driver_kind = Some(driver_kind); + self + } + + /// Enables or disables the timer for the runtime. + /// + /// By default, the timer is disabled. + pub fn enable_timer(mut self, enable: bool) -> Self { + self.enable_timer = enable; + self + } + + /// Enables or disables the offload of file I/O to blocking threads for the runtime. + /// + /// By default, the fs offload is disabled. + pub fn enable_fs_offload(mut self, enable: bool) -> Self { + self.enable_fs_offload = enable; + self + } + + /// Sets the blocking thread pool for the runtime. + pub fn blocking_pool(mut self, blocking_pool: Box) -> Self { + self.blocking_pool = Some(blocking_pool); + self + } + + /// Sets the default blocking thread pool for the runtime with specified maximum number of threads. + #[cfg(feature = "blocking-default")] + pub fn default_blocking_pool(mut self, max_threads: usize) -> Self { + self.blocking_pool = Some(Box::new(DefaultBlockingThreadPool::with_max_threads( + max_threads, + ))); + self + } + + /// Builds the async runtime with the configured settings. + /// + /// If no driver was explicitly set, selects the best available driver for the platform. + pub fn build(self) -> Result { + let driver = if let Some(driver_kind) = self.driver_kind { + driver_kind.into_driver()? + } else { + AnyDriver::new_best()? + }; + Ok(crate::executor::Runtime::with_options( + driver, + self.enable_timer, + self.blocking_pool, + self.enable_fs_offload, + )) + } +} + +impl Default for RuntimeBuilder { + fn default() -> Self { + Self::new() + } +} diff --git a/vendor/vibeio/src/driver/iocp.rs b/vendor/vibeio/src/driver/iocp.rs new file mode 100644 index 0000000..dbf3e6d --- /dev/null +++ b/vendor/vibeio/src/driver/iocp.rs @@ -0,0 +1,923 @@ +use std::cell::RefCell; +use std::ffi::c_void; +use std::io::{self, ErrorKind}; +use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; +use std::ptr; +use std::sync::Arc; +use std::task::Waker; +use std::time::Duration; + +use mio::{Interest, Token}; +use slab::Slab; +use windows_sys::Wdk::Foundation::OBJECT_ATTRIBUTES; +use windows_sys::Wdk::Storage::FileSystem::{ + FileReplaceCompletionInformation, NtCancelIoFileEx, NtCreateFile, NtSetInformationFile, + FILE_COMPLETION_INFORMATION, FILE_OPEN, +}; +use windows_sys::Wdk::System::IO::NtDeviceIoControlFile; +use windows_sys::Win32::Foundation::{ + RtlNtStatusToDosError, ERROR_ABANDONED_WAIT_0, HANDLE, INVALID_HANDLE_VALUE, NTSTATUS, + OBJ_CASE_INSENSITIVE, UNICODE_STRING, WAIT_TIMEOUT, +}; +use windows_sys::Win32::Networking::WinSock::{ + self as WinSock, INVALID_SOCKET, SIO_BASE_HANDLE, SIO_BSP_HANDLE_POLL, SOCKET, SOCKET_ERROR, +}; +use windows_sys::Win32::Storage::FileSystem::{ + SetFileCompletionNotificationModes, FILE_SHARE_READ, FILE_SHARE_WRITE, +}; +use windows_sys::Win32::System::IO::{ + CancelIoEx, CreateIoCompletionPort, GetQueuedCompletionStatusEx, PostQueuedCompletionStatus, + IO_STATUS_BLOCK, OVERLAPPED, OVERLAPPED_ENTRY, +}; + +use crate::driver::{CompletionIoResult, Interruptor}; +use crate::{ + driver::{Driver, RegistrationMode}, + fd_inner::{InnerRawHandle, RawOsHandle}, + op::Op, +}; + +const INTERRUPT_KEY: usize = usize::MAX; +const AFD_POLL_COMPLETION_KEY: usize = usize::MAX - 1; +const IOCP_BATCH_SIZE: usize = 128; + +const STATUS_PENDING: NTSTATUS = 0x0000_0103; +const IOCTL_AFD_POLL: u32 = 0x0001_2024; +const FILE_SKIP_SET_EVENT_ON_HANDLE: u8 = 0x02; +const SYNCHRONIZE_ACCESS: u32 = 0x0010_0000; + +const AFD_POLL_RECEIVE: u32 = 0x0001; +const AFD_POLL_RECEIVE_EXPEDITED: u32 = 0x0002; +const AFD_POLL_SEND: u32 = 0x0004; +const AFD_POLL_DISCONNECT: u32 = 0x0008; +const AFD_POLL_ABORT: u32 = 0x0010; +const AFD_POLL_LOCAL_CLOSE: u32 = 0x0020; +const AFD_POLL_CONNECT: u32 = 0x0040; +const AFD_POLL_ACCEPT: u32 = 0x0080; +const AFD_POLL_CONNECT_FAIL: u32 = 0x0100; + +#[repr(C)] +#[derive(Clone, Copy, Default)] +struct AfdPollHandleInfo { + handle: HANDLE, + events: u32, + status: NTSTATUS, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct AfdPollInfo { + timeout: i64, + number_of_handles: u32, + exclusive: u32, + handles: [AfdPollHandleInfo; 1], +} + +impl AfdPollInfo { + #[inline] + fn new(socket: SOCKET, events: u32) -> Self { + Self { + timeout: i64::MAX, + number_of_handles: 1, + exclusive: 0, + handles: [AfdPollHandleInfo { + handle: socket as HANDLE, + events, + status: 0, + }], + } + } +} + +pub struct IocpInterruptor { + port: std::sync::Weak, +} + +impl Interruptor for IocpInterruptor { + #[inline] + fn interrupt(&self) { + if let Some(port) = self.port.upgrade() { + let _ = unsafe { + PostQueuedCompletionStatus( + port.as_raw_handle() as HANDLE, + 0, + INTERRUPT_KEY, + ptr::null_mut(), + ) + }; + } + } +} + +struct Completion { + waiter: Option, + completed: Option, + overlapped: Option>, + ignored_data: Option>, +} + +#[repr(C)] +struct OverlappedCtx { + overlapped: OVERLAPPED, + token: usize, +} + +#[repr(C)] +struct AfdIoStatusCtx { + io_status: IO_STATUS_BLOCK, + token: usize, +} + +struct PollOp { + registration_token: usize, + io_status: Box, + _input: Box, + _output: Box, +} + +struct PollRegistration { + socket: SOCKET, + waiter: Option, + interest: Interest, + poll_token: Option, +} + +enum HandleRegistration { + Completion, + Poll(PollRegistration), +} + +struct DriverState { + registrations: Slab, + completions: Slab, + poll_ops: Slab, +} + +pub struct IocpDriver { + port: Arc, + afd: RefCell>, + state: RefCell, +} + +impl IocpDriver { + #[inline] + pub(crate) fn new() -> Result { + let port = unsafe { + CreateIoCompletionPort(INVALID_HANDLE_VALUE as HANDLE, ptr::null_mut(), 0, 0) + }; + if port.is_null() { + return Err(io::Error::last_os_error()); + } + + Ok(Self { + port: Arc::new(unsafe { OwnedHandle::from_raw_handle(port as RawHandle) }), + afd: RefCell::new(None), + state: RefCell::new(DriverState { + registrations: Slab::with_capacity(1024), + completions: Slab::with_capacity(1024), + poll_ops: Slab::with_capacity(1024), + }), + }) + } + + #[inline] + fn update_waiter(waiter_slot: &mut Option, waker: Waker) { + if !waiter_slot + .as_ref() + .is_some_and(|waiter| waiter.will_wake(&waker)) + { + *waiter_slot = Some(waker); + } + } + + #[inline] + fn iocp_handle(&self) -> HANDLE { + self.port.as_raw_handle() as HANDLE + } + + #[inline] + fn status_is_success(status: NTSTATUS) -> bool { + status >= 0 + } + + #[inline] + fn ntstatus_to_io_error(status: NTSTATUS) -> io::Error { + let mapped = unsafe { RtlNtStatusToDosError(status) } as i32; + if mapped != 0 { + io::Error::from_raw_os_error(mapped) + } else { + io::Error::from_raw_os_error(status) + } + } + + #[inline] + fn raw_os_handle_to_windows_handle(handle: RawOsHandle) -> HANDLE { + match handle { + RawOsHandle::Socket(socket) => socket as HANDLE, + RawOsHandle::Handle(handle) => handle as HANDLE, + } + } + + #[inline] + fn raw_os_handle_to_socket(handle: RawOsHandle) -> Result { + match handle { + RawOsHandle::Socket(socket) => Ok(socket as SOCKET), + RawOsHandle::Handle(_) => Err(io::Error::new( + ErrorKind::Unsupported, + "poll mode currently supports sockets only on Windows", + )), + } + } + + #[inline] + fn duration_to_timeout_ms(timeout: Option) -> u32 { + match timeout { + Some(timeout) => timeout.as_millis().min(u32::MAX as u128) as u32, + None => u32::MAX, + } + } + + #[inline] + fn completion_result_from_entry(entry: &OVERLAPPED_ENTRY) -> i32 { + if entry.Internal == 0 { + return entry.dwNumberOfBytesTransferred as i32; + } + + let ntstatus = entry.Internal as i32; + let win32_error = unsafe { RtlNtStatusToDosError(ntstatus) } as i32; + let mapped_error = if win32_error == 0 { + ntstatus + } else { + win32_error + }; + -mapped_error + } + + #[inline] + fn interest_to_afd_events(interest: Interest) -> u32 { + let mut events = AFD_POLL_DISCONNECT | AFD_POLL_ABORT | AFD_POLL_LOCAL_CLOSE; + + if interest.is_readable() { + events |= AFD_POLL_RECEIVE | AFD_POLL_RECEIVE_EXPEDITED | AFD_POLL_ACCEPT; + } + if interest.is_writable() { + events |= AFD_POLL_SEND | AFD_POLL_CONNECT | AFD_POLL_CONNECT_FAIL; + } + + if events == (AFD_POLL_DISCONNECT | AFD_POLL_ABORT | AFD_POLL_LOCAL_CLOSE) { + events |= AFD_POLL_RECEIVE | AFD_POLL_SEND | AFD_POLL_ACCEPT | AFD_POLL_CONNECT; + } + + events + } + + #[inline] + fn get_base_socket(socket: SOCKET, ioctl: u32) -> Result { + let mut base_socket: SOCKET = INVALID_SOCKET; + let mut bytes: u32 = 0; + let result = unsafe { + WinSock::WSAIoctl( + socket, + ioctl, + ptr::null_mut(), + 0, + (&mut base_socket as *mut SOCKET).cast(), + std::mem::size_of::() as u32, + &mut bytes, + ptr::null_mut(), + None, + ) + }; + + if result == SOCKET_ERROR { + let err = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err)); + } + + Ok(base_socket) + } + + #[inline] + fn resolve_base_socket(mut socket: SOCKET) -> Result { + loop { + if let Ok(base_socket) = Self::get_base_socket(socket, SIO_BASE_HANDLE) { + return Ok(base_socket); + } + + match Self::get_base_socket(socket, SIO_BSP_HANDLE_POLL) { + Ok(base_socket) if base_socket != INVALID_SOCKET && base_socket != socket => { + socket = base_socket; + } + Ok(_) => { + return Err(io::Error::new( + ErrorKind::Other, + "failed to resolve base socket for AFD polling", + )) + } + Err(err) => return Err(err), + } + } + } + + #[inline] + fn open_afd_handle() -> Result { + let device_name = format!("\\Device\\Afd\\vibeio-{}", std::process::id()); + let mut device_name_utf16: Vec = device_name.encode_utf16().collect(); + let name_len = (device_name_utf16.len() * std::mem::size_of::()) as u16; + + let object_name = UNICODE_STRING { + Length: name_len, + MaximumLength: name_len, + Buffer: device_name_utf16.as_mut_ptr(), + }; + let object_attributes = OBJECT_ATTRIBUTES { + Length: std::mem::size_of::() as u32, + RootDirectory: ptr::null_mut(), + ObjectName: &object_name, + Attributes: OBJ_CASE_INSENSITIVE, + SecurityDescriptor: ptr::null(), + SecurityQualityOfService: ptr::null(), + }; + + let mut afd_handle: HANDLE = ptr::null_mut(); + let mut create_status = IO_STATUS_BLOCK::default(); + let status = unsafe { + NtCreateFile( + &mut afd_handle, + SYNCHRONIZE_ACCESS, + &object_attributes, + &mut create_status, + ptr::null(), + 0, + FILE_SHARE_READ | FILE_SHARE_WRITE, + FILE_OPEN, + 0, + ptr::null(), + 0, + ) + }; + + if !Self::status_is_success(status) { + return Err(Self::ntstatus_to_io_error(status)); + } + if afd_handle.is_null() { + return Err(io::Error::new( + ErrorKind::Other, + "NtCreateFile returned a null AFD handle", + )); + } + + Ok(unsafe { OwnedHandle::from_raw_handle(afd_handle as RawHandle) }) + } + + #[inline] + fn ensure_afd_handle(&self) -> Result { + { + let afd = self.afd.borrow(); + if let Some(afd) = afd.as_ref() { + return Ok(afd.as_raw_handle() as HANDLE); + } + } + + let mut afd_slot = self.afd.borrow_mut(); + if afd_slot.is_none() { + let afd = Self::open_afd_handle()?; + let afd_handle = afd.as_raw_handle() as HANDLE; + + let completion_port = unsafe { + CreateIoCompletionPort(afd_handle, self.iocp_handle(), AFD_POLL_COMPLETION_KEY, 0) + }; + if completion_port.is_null() { + return Err(io::Error::last_os_error()); + } + + if unsafe { + SetFileCompletionNotificationModes(afd_handle, FILE_SKIP_SET_EVENT_ON_HANDLE) + } == 0 + { + return Err(io::Error::last_os_error()); + } + + *afd_slot = Some(afd); + } + + Ok(afd_slot + .as_ref() + .expect("AFD handle must be initialized") + .as_raw_handle() as HANDLE) + } + + #[inline] + fn arm_poll_operation( + &self, + registration_token: usize, + socket: SOCKET, + interest: Interest, + ) -> Result<(), io::Error> { + let afd_handle = self.ensure_afd_handle()?; + let poll_events = Self::interest_to_afd_events(interest); + + let (poll_token, io_status_ptr, input_ptr, output_ptr) = { + let mut state = self.state.borrow_mut(); + + match state.registrations.get(registration_token) { + Some(HandleRegistration::Poll(registration)) => { + if registration.poll_token.is_some() { + return Ok(()); + } + } + Some(HandleRegistration::Completion) => { + return Err(io::Error::new( + ErrorKind::Unsupported, + format!( + "I/O token {} is registered for completion mode, not poll mode", + registration_token + ), + )); + } + None => { + return Err(io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + registration_token + ), + )); + } + } + + let mut io_status = Box::new(unsafe { std::mem::zeroed::() }); + let mut input = Box::new(AfdPollInfo::new(socket, poll_events)); + let mut output = Box::new(AfdPollInfo::new(socket, 0)); + + let poll_entry = state.poll_ops.vacant_entry(); + let poll_token = poll_entry.key(); + io_status.token = poll_token; + + let io_status_ptr = &mut io_status.io_status as *mut IO_STATUS_BLOCK; + let input_ptr = (&mut *input as *mut AfdPollInfo).cast(); + let output_ptr = (&mut *output as *mut AfdPollInfo).cast(); + + poll_entry.insert(PollOp { + registration_token, + io_status, + _input: input, + _output: output, + }); + if let Some(HandleRegistration::Poll(registration)) = + state.registrations.get_mut(registration_token) + { + registration.poll_token = Some(poll_token); + } + + (poll_token, io_status_ptr, input_ptr, output_ptr) + }; + + let status = unsafe { + NtDeviceIoControlFile( + afd_handle, + ptr::null_mut(), + None, + io_status_ptr as *const c_void, + io_status_ptr, + IOCTL_AFD_POLL, + input_ptr, + std::mem::size_of::() as u32, + output_ptr, + std::mem::size_of::() as u32, + ) + }; + + if Self::status_is_success(status) || status == STATUS_PENDING { + return Ok(()); + } + + let mut state = self.state.borrow_mut(); + let _ = state.poll_ops.try_remove(poll_token); + if let Some(HandleRegistration::Poll(registration)) = + state.registrations.get_mut(registration_token) + { + if registration.poll_token == Some(poll_token) { + registration.poll_token = None; + } + } + + Err(Self::ntstatus_to_io_error(status)) + } + + #[inline] + fn cancel_poll_operation(&self, poll_token: usize) { + let afd_handle = { + let afd = self.afd.borrow(); + let Some(afd) = afd.as_ref() else { + return; + }; + afd.as_raw_handle() as HANDLE + }; + + let io_status_ptr = { + let state = self.state.borrow(); + let Some(poll_op) = state.poll_ops.get(poll_token) else { + return; + }; + &poll_op.io_status.io_status as *const IO_STATUS_BLOCK + }; + + let mut cancel_status = IO_STATUS_BLOCK::default(); + let _ = unsafe { NtCancelIoFileEx(afd_handle, io_status_ptr, &mut cancel_status) }; + } + + #[inline] + fn disassociate_iocp_handle(&self, handle: &InnerRawHandle) { + let windows_handle = Self::raw_os_handle_to_windows_handle(handle.handle); + let info: FILE_COMPLETION_INFORMATION = FILE_COMPLETION_INFORMATION { + Port: std::ptr::null_mut(), + Key: std::ptr::null_mut(), + }; + let mut status = IO_STATUS_BLOCK::default(); + let _ = unsafe { + NtSetInformationFile( + windows_handle, + &mut status, + &info as *const FILE_COMPLETION_INFORMATION as *const c_void, + std::mem::size_of::() as u32, + FileReplaceCompletionInformation, + ) + }; + } + + #[inline] + fn process_entries(&self, entries: &[OVERLAPPED_ENTRY]) { + let mut state = self.state.borrow_mut(); + for entry in entries { + if entry.lpCompletionKey == INTERRUPT_KEY { + continue; + } + + if entry.lpCompletionKey == AFD_POLL_COMPLETION_KEY { + if entry.lpOverlapped.is_null() { + continue; + } + + // SAFETY: every AFD poll submission passes a pointer to AfdIoStatusCtx::io_status, + // and AfdIoStatusCtx is repr(C) with io_status as its first field. + let poll_token = unsafe { (*entry.lpOverlapped.cast::()).token }; + let registration_token = state + .poll_ops + .get(poll_token) + .map(|poll| poll.registration_token); + let _ = state.poll_ops.try_remove(poll_token); + + if let Some(registration_token) = registration_token { + if let Some(HandleRegistration::Poll(registration)) = + state.registrations.get_mut(registration_token) + { + if registration.poll_token == Some(poll_token) { + registration.poll_token = None; + } + if let Some(waiter) = registration.waiter.take() { + waiter.wake(); + } + } + } + continue; + } + + if entry.lpOverlapped.is_null() { + continue; + } + + // SAFETY: every OVERLAPPED pointer submitted by completion operations points to the + // first field of OverlappedCtx (repr(C), first field), and lives in Completion::overlapped + // until consumed here. + let completion_token = unsafe { (*entry.lpOverlapped.cast::()).token }; + if let Some(completion) = state.completions.get_mut(completion_token) { + completion.completed = Some(Self::completion_result_from_entry(entry)); + completion.overlapped = None; + if let Some(waiter) = completion.waiter.take() { + waiter.wake(); + } + if completion.ignored_data.is_some() { + state.completions.remove(completion_token); + } + } + } + } + + #[inline] + fn process_batch(&self, timeout_ms: u32) -> Result { + let mut entries = [OVERLAPPED_ENTRY::default(); IOCP_BATCH_SIZE]; + let mut entries_removed: u32 = 0; + + let success = unsafe { + GetQueuedCompletionStatusEx( + self.iocp_handle(), + entries.as_mut_ptr(), + entries.len() as u32, + &mut entries_removed, + timeout_ms, + 0, + ) + } != 0; + + if !success { + let err = io::Error::last_os_error(); + if matches!( + err.raw_os_error(), + Some(code) if code == WAIT_TIMEOUT as i32 || code == ERROR_ABANDONED_WAIT_0 as i32 + ) { + return Ok(0); + } + return Err(err); + } + + if entries_removed == 0 { + return Ok(0); + } + + self.process_entries(&entries[..entries_removed as usize]); + + Ok(entries_removed as usize) + } + + #[inline] + fn process_ready_completions(&self) -> Result<(), io::Error> { + while self.process_batch(0)? > 0 {} + Ok(()) + } +} + +impl Driver for IocpDriver { + type Interruptor = IocpInterruptor; + + #[inline] + fn flush(&self) { + match self.process_ready_completions() { + Ok(_) => {} + Err(err) if err.kind() == io::ErrorKind::Interrupted => {} + Err(err) => panic!("iocp flush failed while processing completions: {err}"), + } + } + + #[inline] + fn wait(&self, timeout: Option) { + let timeout_ms = Self::duration_to_timeout_ms(timeout); + match self.process_batch(timeout_ms) { + Ok(processed) if processed > 0 => { + if let Err(err) = self.process_ready_completions() { + panic!("iocp drain failed while waiting for I/O: {err}"); + } + } + Ok(_) => {} + Err(err) if err.kind() == io::ErrorKind::Interrupted => {} + Err(err) => panic!("iocp wait failed while waiting for I/O: {err}"), + } + } + + #[inline] + fn register_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result { + self.register_handle_with_mode(handle, interest, RegistrationMode::Completion) + } + + #[inline] + fn register_handle_with_mode( + &self, + handle: &InnerRawHandle, + interest: Interest, + mode: RegistrationMode, + ) -> Result { + match mode { + RegistrationMode::Completion => { + let token = { + let mut state = self.state.borrow_mut(); + let entry = state.registrations.vacant_entry(); + let token = Token(entry.key()); + entry.insert(HandleRegistration::Completion); + token + }; + + let source_handle = Self::raw_os_handle_to_windows_handle(handle.handle); + let completion_port = unsafe { + CreateIoCompletionPort(source_handle, self.iocp_handle(), token.0, 0) + }; + if completion_port.is_null() { + let mut state = self.state.borrow_mut(); + let _ = state.registrations.try_remove(token.0); + return Err(io::Error::last_os_error()); + } + + Ok(token) + } + RegistrationMode::Poll => { + let socket = Self::raw_os_handle_to_socket(handle.handle)?; + let base_socket = Self::resolve_base_socket(socket)?; + self.ensure_afd_handle()?; + + let mut state = self.state.borrow_mut(); + let entry = state.registrations.vacant_entry(); + let token = Token(entry.key()); + entry.insert(HandleRegistration::Poll(PollRegistration { + socket: base_socket, + waiter: None, + interest, + poll_token: None, + })); + Ok(token) + } + } + } + + #[inline] + fn reregister_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result<(), io::Error> { + let mut state = self.state.borrow_mut(); + match state.registrations.get_mut(handle.token.0) { + Some(HandleRegistration::Completion) => Ok(()), + Some(HandleRegistration::Poll(registration)) => { + registration.interest = interest; + Ok(()) + } + None => Err(io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + )), + } + } + + #[inline] + fn deregister_handle(&self, handle: &InnerRawHandle) -> Result<(), io::Error> { + let poll_token = { + let mut state = self.state.borrow_mut(); + match state.registrations.try_remove(handle.token.0) { + Some(HandleRegistration::Completion) => None, + Some(HandleRegistration::Poll(registration)) => registration.poll_token, + None => { + return Err(io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + )) + } + } + }; + + if let Some(poll_token) = poll_token { + // Poll-based I/O + self.cancel_poll_operation(poll_token); + } else { + // Completion-based I/O + self.disassociate_iocp_handle(handle); + } + + Ok(()) + } + + #[inline] + fn supports_completion(&self) -> bool { + true + } + + #[inline] + fn submit_completion(&self, op: &mut O, waker: Waker) -> CompletionIoResult + where + O: Op, + { + let (completion_token, overlapped_ptr) = { + let mut state = self.state.borrow_mut(); + let vacant_completion = state.completions.vacant_entry(); + let completion_token = vacant_completion.key(); + + let mut overlapped = Box::new(unsafe { std::mem::zeroed::() }); + overlapped.token = completion_token; + let overlapped_ptr: *mut OVERLAPPED = &mut overlapped.overlapped; + + vacant_completion.insert(Completion { + waiter: Some(waker), + completed: None, + overlapped: Some(overlapped), + ignored_data: None, + }); + + (completion_token, overlapped_ptr) + }; + + if let Err(err) = op.submit_windows(overlapped_ptr) { + let mut state = self.state.borrow_mut(); + let _ = state.completions.try_remove(completion_token); + return CompletionIoResult::SubmitErr(err); + } + + CompletionIoResult::Retry(completion_token) + } + + #[inline] + fn submit_poll( + &self, + handle: &InnerRawHandle, + waker: Waker, + interest: Interest, + ) -> Result<(), io::Error> { + let (socket, should_arm) = { + let mut state = self.state.borrow_mut(); + let registration = state.registrations.get_mut(handle.token.0).ok_or_else(|| { + io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + ) + })?; + + let HandleRegistration::Poll(registration) = registration else { + return Err(io::Error::new( + ErrorKind::Unsupported, + format!( + "I/O token {} is registered for completion mode, not poll mode", + handle.token.0 + ), + )); + }; + + registration.interest = interest; + Self::update_waiter(&mut registration.waiter, waker); + (registration.socket, registration.poll_token.is_none()) + }; + + if should_arm { + self.arm_poll_operation(handle.token.0, socket, interest)?; + } + + Ok(()) + } + + #[inline] + fn get_completion_result(&self, token: usize) -> Option { + let mut state = self.state.borrow_mut(); + let completed = state + .completions + .get(token) + .and_then(|completion| completion.completed); + if completed.is_some() { + state.completions.remove(token); + } + completed + } + + #[inline] + fn set_completion_waker(&self, token: usize, waker: Waker) { + let mut state = self.state.borrow_mut(); + if let Some(completion) = state.completions.get_mut(token) { + Self::update_waiter(&mut completion.waiter, waker); + } + } + + #[inline] + fn get_interruptor(&self) -> Self::Interruptor { + IocpInterruptor { + port: Arc::downgrade(&self.port), + } + } + + #[inline] + fn ignore_completion(&self, token: usize, data: Box) { + let mut state = self.state.borrow_mut(); + if let Some(c) = state.completions.get_mut(token) { + c.ignored_data = Some(data); + } + } + + #[inline] + fn cancel_completion(&self, token: usize, handle: RawOsHandle, data: Box) { + let overlapped = { + let mut state = self.state.borrow_mut(); + let Some(completion) = state.completions.get_mut(token) else { + return; + }; + completion.ignored_data = Some(data); + completion + .overlapped + .as_deref_mut() + .map(|ctx| &mut ctx.overlapped as *mut OVERLAPPED) + }; + + if let Some(overlapped) = overlapped { + // The completion entry and ignored_data remain registered until + // IOCP reports the aborted (or concurrently completed) operation. + let _ = + unsafe { CancelIoEx(Self::raw_os_handle_to_windows_handle(handle), overlapped) }; + } + } +} diff --git a/vendor/vibeio/src/driver/mio.rs b/vendor/vibeio/src/driver/mio.rs new file mode 100644 index 0000000..cfd18c4 --- /dev/null +++ b/vendor/vibeio/src/driver/mio.rs @@ -0,0 +1,310 @@ +use std::cell::RefCell; +use std::io::{self, ErrorKind}; +use std::os::fd::RawFd; +use std::sync::Arc; +use std::task::Waker; +use std::time::Duration; + +use mio::{Events, Interest, Poll, Registry, Token, Waker as MioWaker}; +use slab::Slab; + +use crate::driver::Interruptor; +use crate::{driver::Driver, fd_inner::InnerRawHandle}; + +pub struct MioInterruptor { + waker: std::sync::Weak, +} + +impl Interruptor for MioInterruptor { + #[inline] + fn interrupt(&self) { + if let Some(waker) = self.waker.upgrade() { + let _ = waker.wake(); + } + } +} + +struct Registration { + fd: RawFd, + waiter: Option, + interest: Interest, +} + +struct DriverState { + registrations: Slab, +} + +pub struct MioDriver { + poll: RefCell, + registry: Registry, + events: RefCell, + state: RefCell, + waker: Arc, +} + +impl MioDriver { + #[inline] + pub(crate) fn new() -> Result { + let poll = Poll::new()?; + let registry = poll.registry().try_clone()?; + let waker = MioWaker::new(®istry, Token(usize::MAX))?; + + Ok(Self { + poll: RefCell::new(poll), + registry, + events: RefCell::new(Events::with_capacity(1024)), + state: RefCell::new(DriverState { + registrations: Slab::with_capacity(1024), + }), + waker: Arc::new(waker), + }) + } + + #[inline] + fn update_waiter(waiter_slot: &mut Option, waker: Waker) { + if !waiter_slot + .as_ref() + .is_some_and(|waiter| waiter.will_wake(&waker)) + { + *waiter_slot = Some(waker); + } + } + + #[inline] + pub(crate) fn wait_timeout(&self, timeout: Option) { + let mut poll = self.poll.borrow_mut(); + let mut events = self.events.borrow_mut(); + match poll.poll(&mut events, timeout) { + Ok(_) => {} + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} // strace edge case + Err(e) => panic!("mio poll failed while waiting for I/O events: {}", e), + }; + + { + let mut state = self.state.borrow_mut(); + for event in events.iter() { + // Check if this is an interrupt event + if event.token().0 == usize::MAX { + continue; + } + + if let Some(registration) = state.registrations.get_mut(event.token().0) { + if let Some(task) = registration.waiter.take() { + task.wake(); + } + } + } + } + } +} + +impl Driver for MioDriver { + type Interruptor = MioInterruptor; + + #[inline] + fn flush(&self) { + self.wait_timeout(Some(Duration::ZERO)); + } + + #[inline] + fn wait(&self, timeout: Option) { + self.wait_timeout(timeout); + } + + #[inline] + fn get_interruptor(&self) -> Self::Interruptor { + MioInterruptor { + waker: Arc::downgrade(&self.waker), + } + } + + #[inline] + fn register_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result { + let token = { + let mut state = self.state.borrow_mut(); + let entry = state.registrations.vacant_entry(); + let token = Token(entry.key()); + entry.insert(Registration { + fd: handle.handle, + waiter: None, + interest, + }); + token + }; + + let mut source = mio::unix::SourceFd(&handle.handle); + if let Err(err) = self.registry.register(&mut source, token, interest) { + let mut state = self.state.borrow_mut(); + let _ = state.registrations.try_remove(token.0); + return Err(err); + } + + Ok(token) + } + + #[inline] + fn reregister_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result<(), io::Error> { + let mut state = self.state.borrow_mut(); + let registration = state.registrations.get_mut(handle.token.0).ok_or_else(|| { + io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + ) + })?; + + let mut source = mio::unix::SourceFd(®istration.fd); + self.registry + .reregister(&mut source, handle.token, interest)?; + registration.interest = interest; + Ok(()) + } + + #[inline] + fn deregister_handle(&self, handle: &InnerRawHandle) -> Result<(), io::Error> { + let fd = { + let state = self.state.borrow(); + let registration = state.registrations.get(handle.token.0).ok_or_else(|| { + io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + ) + })?; + registration.fd + }; + + let mut source = mio::unix::SourceFd(&fd); + self.registry.deregister(&mut source)?; + + let mut state = self.state.borrow_mut(); + let _ = state.registrations.try_remove(handle.token.0); + Ok(()) + } + + #[inline] + fn submit_poll( + &self, + handle: &InnerRawHandle, + waker: Waker, + interest: Interest, + ) -> Result<(), io::Error> { + let token = handle.token(); + + let mut state = self.state.borrow_mut(); + let registration = state.registrations.get_mut(token.0).ok_or_else(|| { + io::Error::new( + ErrorKind::NotFound, + format!("I/O token {} is not registered with this driver", token.0), + ) + })?; + + if registration.interest != interest { + // Re-register, but only if the interest has change + self.registry.reregister( + &mut mio::unix::SourceFd(®istration.fd), + token, + interest, + )?; + registration.interest = interest; + } + + Self::update_waiter(&mut registration.waiter, waker); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use super::MioDriver; + + struct TestWake { + count: AtomicUsize, + } + + impl TestWake { + #[inline] + fn new() -> Self { + Self { + count: AtomicUsize::new(0), + } + } + + #[inline] + fn wake_count(&self) -> usize { + self.count.load(Ordering::SeqCst) + } + } + + impl std::task::Wake for TestWake { + #[inline] + fn wake(self: Arc) { + self.count.fetch_add(1, Ordering::SeqCst); + } + + #[inline] + fn wake_by_ref(self: &Arc) { + self.count.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn wait_wakes_task_for_ready_token() { + use std::{ + io::Write, + os::fd::AsRawFd, + rc::Rc, + task::{Context, Poll}, + time::Duration, + }; + + use crate::{driver::AnyDriver, fd_inner::InnerRawHandle, op::ReadOp}; + + let driver = Rc::new(AnyDriver::Mio( + MioDriver::new().expect("mio driver should initialize"), + )); + let wake = Arc::new(TestWake::new()); + let waker = std::task::Waker::from(wake.clone()); + + // Since the driver already has a waker, let's use an Unix pipe instead + let (side1, mut side2) = + std::os::unix::net::UnixStream::pair().expect("failed to create pipe"); + let buffer = [0u8; 1]; + side1 + .set_nonblocking(true) + .expect("failed to set non-blocking"); + let inner_raw_handle = InnerRawHandle::new_with_driver_and_mode( + &driver, + side1.as_raw_fd(), + mio::Interest::READABLE, + crate::driver::RegistrationMode::Poll, + ) + .expect("failed to register pipe"); + let mut read_op = ReadOp::new(&inner_raw_handle, buffer); + match inner_raw_handle.poll_op(&mut Context::from_waker(&waker), &mut read_op) { + Poll::Pending => {} + Poll::Ready(Ok(_)) => panic!("unexpected success"), + Poll::Ready(Err(e)) => panic!("failed to submit operation: {}", e), + }; + + side2.write_all(b"!").expect("failed to write to pipe"); // Exact data written doesn't matter... + + driver.wait(Some(Duration::from_millis(100))); + assert_eq!(wake.wake_count(), 1); + } +} diff --git a/vendor/vibeio/src/driver/mock.rs b/vendor/vibeio/src/driver/mock.rs new file mode 100644 index 0000000..016a7b5 --- /dev/null +++ b/vendor/vibeio/src/driver/mock.rs @@ -0,0 +1,79 @@ +use std::time::Duration; + +use crate::driver::{Driver, Interruptor}; +use mio::{Interest, Token}; + +pub struct MockInterruptor { + thread: std::thread::Thread, +} + +impl Interruptor for MockInterruptor { + #[inline] + fn interrupt(&self) { + self.thread.unpark(); + } +} + +pub struct MockDriver; + +impl MockDriver { + #[inline] + pub(crate) fn new() -> Self { + MockDriver {} + } +} + +impl Driver for MockDriver { + type Interruptor = MockInterruptor; + + #[inline] + fn wait(&self, timeout: Option) { + if let Some(timeout) = timeout { + std::thread::park_timeout(timeout); + } else { + std::thread::park(); + } + } + + #[inline] + fn get_interruptor(&self) -> Self::Interruptor { + MockInterruptor { + thread: std::thread::current(), + } + } + + #[inline] + fn register_handle( + &self, + _handle: &crate::fd_inner::InnerRawHandle, + _interest: Interest, + ) -> Result { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "MockDriver does not support I/O handle registration", + )) + } + + #[inline] + fn reregister_handle( + &self, + _handle: &crate::fd_inner::InnerRawHandle, + _interest: Interest, + ) -> Result<(), std::io::Error> { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "MockDriver does not support I/O handle re-registration", + )) + } + + #[inline] + fn deregister_handle( + &self, + _handle: &crate::fd_inner::InnerRawHandle, + ) -> Result<(), std::io::Error> { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "MockDriver does not support I/O handle deregistration", + )) + } +} diff --git a/vendor/vibeio/src/driver/mod.rs b/vendor/vibeio/src/driver/mod.rs new file mode 100644 index 0000000..e6c3be7 --- /dev/null +++ b/vendor/vibeio/src/driver/mod.rs @@ -0,0 +1,595 @@ +#[cfg(windows)] +mod iocp; +#[cfg(unix)] +mod mio; +mod mock; +#[cfg(target_os = "linux")] +mod uring; + +use std::task::Waker; +use std::{io, time::Duration}; + +use ::mio::{Interest, Token}; + +#[cfg(windows)] +use crate::driver::iocp::{IocpDriver, IocpInterruptor}; +#[cfg(unix)] +use crate::driver::mio::{MioDriver, MioInterruptor}; +use crate::driver::mock::MockInterruptor; +#[cfg(target_os = "linux")] +use crate::driver::uring::{UringDriver, UringInterruptor}; +use crate::op::Op; +use crate::{driver::mock::MockDriver, fd_inner::InnerRawHandle}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RegistrationMode { + Poll, + Completion, +} + +#[derive(Debug)] +pub enum CompletionIoResult { + #[allow(dead_code)] + Ok(i32), + Retry(usize), // usize -> token + SubmitErr(std::io::Error), +} + +#[inline] +fn unsupported_completion_error() -> io::Error { + io::Error::new( + io::ErrorKind::Unsupported, + "driver does not support completion-based I/O submission", + ) +} + +#[inline] +fn unsupported_poll_error() -> io::Error { + io::Error::new( + io::ErrorKind::Unsupported, + "driver does not support poll-based I/O submission", + ) +} + +pub trait Interruptor { + /// Interrupts a waiting I/O operation. + fn interrupt(&self); +} + +pub enum AnyInterruptor { + Mock(MockInterruptor), + #[cfg(windows)] + Iocp(IocpInterruptor), + #[cfg(unix)] + Mio(MioInterruptor), + #[cfg(target_os = "linux")] + IoUring(UringInterruptor), +} + +impl AnyInterruptor { + pub(crate) fn interrupt(&self) { + match self { + AnyInterruptor::Mock(interruptor) => interruptor.interrupt(), + #[cfg(windows)] + AnyInterruptor::Iocp(interruptor) => interruptor.interrupt(), + #[cfg(unix)] + AnyInterruptor::Mio(interruptor) => interruptor.interrupt(), + #[cfg(target_os = "linux")] + AnyInterruptor::IoUring(interruptor) => interruptor.interrupt(), + } + } +} + +pub trait Driver { + type Interruptor: Interruptor; + + /// Flushes the driver's I/O. + #[inline] + fn flush(&self) {} + + /// Returns whether the executor should call `flush` after polling a task batch. + #[inline] + fn should_flush(&self) -> bool { + true + } + + /// Waits for the I/O. + fn wait(&self, timeout: Option); + + /// Registers an I/O source and returns its token. + fn register_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result; + + /// Registers an I/O source with the requested mode. + fn register_handle_with_mode( + &self, + handle: &InnerRawHandle, + interest: Interest, + _mode: RegistrationMode, + ) -> Result { + self.register_handle(handle, interest) + } + + /// Updates the interest set for a registered I/O source. + fn reregister_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result<(), std::io::Error>; + + /// Removes an I/O source from the poller. + fn deregister_handle(&self, handle: &InnerRawHandle) -> Result<(), std::io::Error>; + + /// Returns whether the driver supports completion-based I/O operations. + #[inline] + fn supports_completion(&self) -> bool { + false + } + + /// Submits a completion-based I/O operation. + #[inline] + fn submit_completion(&self, _op: &mut O, _waker: Waker) -> CompletionIoResult + where + O: Op, + { + CompletionIoResult::SubmitErr(unsupported_completion_error()) + } + + /// Re-registers interest and submits a waker for poll-based I/O. + #[inline] + fn submit_poll( + &self, + _handle: &InnerRawHandle, + _waker: Waker, + _interest: Interest, + ) -> Result<(), io::Error> { + Err(unsupported_poll_error()) + } + + /// Obtains the result for a completion-based I/O operation. + #[inline] + fn get_completion_result(&self, _token: usize) -> Option { + None + } + + /// Sets the waker for a completion-based I/O operation. + #[inline] + fn set_completion_waker(&self, _token: usize, _waker: Waker) {} + + /// Cancels a completion-based I/O operation. + #[inline] + fn ignore_completion(&self, _token: usize, _data: Box) {} + + /// Cancels a Windows completion operation while retaining its owned data + /// until the completion packet is observed. + #[cfg(windows)] + #[inline] + fn cancel_completion( + &self, + token: usize, + _handle: crate::fd_inner::RawOsHandle, + data: Box, + ) { + self.ignore_completion(token, data); + } + + /// Interrupts a waiting I/O operation. + fn get_interruptor(&self) -> Self::Interruptor; +} + +#[allow(clippy::large_enum_variant)] +pub enum AnyDriver { + Mock(MockDriver), + #[cfg(windows)] + Iocp(IocpDriver), + #[cfg(unix)] + Mio(MioDriver), + #[cfg(target_os = "linux")] + IoUring(UringDriver), +} + +impl AnyDriver { + #[cfg(unix)] + #[inline] + pub(crate) fn new_mio() -> Result { + Ok(AnyDriver::Mio(MioDriver::new()?)) + } + + #[inline] + pub(crate) fn new_mock() -> Self { + AnyDriver::Mock(MockDriver::new()) + } + + #[cfg(windows)] + #[inline] + pub(crate) fn new_iocp() -> Result { + Ok(AnyDriver::Iocp(IocpDriver::new()?)) + } + + #[cfg(target_os = "linux")] + #[inline] + pub(crate) fn new_uring_custom(builder: io_uring::Builder) -> Result { + Ok(AnyDriver::IoUring(UringDriver::new(1024, builder)?)) + } + + #[cfg(target_os = "linux")] + #[inline] + pub(crate) fn new_uring() -> Result { + let mut builder = io_uring::IoUring::builder(); + builder + .setup_single_issuer() + .setup_coop_taskrun() + .setup_taskrun_flag() + .setup_submit_all(); + if let Ok(driver) = Self::new_uring_custom(builder) { + return Ok(driver); + } + + // Fallback for older kernels + Self::new_uring_custom(io_uring::IoUring::builder()) + } + + #[inline] + pub(crate) fn new_best() -> Result { + #[cfg(target_os = "linux")] + if let Ok(driver) = Self::new_uring() { + return Ok(driver); + } + + #[cfg(unix)] + let driver = Self::new_mio()?; + #[cfg(windows)] + let driver = Self::new_iocp()?; + + Ok(driver) + } + + #[inline] + pub(crate) fn flush(&self) { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.flush(), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.flush(), + AnyDriver::Mock(driver) => driver.flush(), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.flush(), + } + } + + #[inline] + pub(crate) fn should_flush(&self) -> bool { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.should_flush(), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.should_flush(), + AnyDriver::Mock(driver) => driver.should_flush(), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.should_flush(), + } + } + + #[inline] + pub(crate) fn wait(&self, timeout: Option) { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.wait(timeout), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.wait(timeout), + AnyDriver::Mock(driver) => driver.wait(timeout), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.wait(timeout), + } + } + + #[allow(dead_code)] + #[inline] + pub(crate) fn register_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.register_handle(handle, interest), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.register_handle(handle, interest), + AnyDriver::Mock(driver) => driver.register_handle(handle, interest), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.register_handle(handle, interest), + } + } + + #[inline] + pub(crate) fn register_handle_with_mode( + &self, + handle: &InnerRawHandle, + interest: Interest, + mode: RegistrationMode, + ) -> Result { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.register_handle_with_mode(handle, interest, mode), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.register_handle_with_mode(handle, interest, mode), + AnyDriver::Mock(driver) => driver.register_handle_with_mode(handle, interest, mode), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.register_handle_with_mode(handle, interest, mode), + } + } + + #[inline] + pub(crate) fn reregister_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result<(), std::io::Error> { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.reregister_handle(handle, interest), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.reregister_handle(handle, interest), + AnyDriver::Mock(driver) => driver.reregister_handle(handle, interest), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.reregister_handle(handle, interest), + } + } + + #[inline] + pub(crate) fn deregister_handle(&self, handle: &InnerRawHandle) -> Result<(), std::io::Error> { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.deregister_handle(handle), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.deregister_handle(handle), + AnyDriver::Mock(driver) => driver.deregister_handle(handle), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.deregister_handle(handle), + } + } + + #[inline] + pub(crate) fn supports_completion(&self) -> bool { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.supports_completion(), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.supports_completion(), + AnyDriver::Mock(driver) => driver.supports_completion(), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.supports_completion(), + } + } + + #[inline] + pub(crate) fn submit_completion(&self, op: &mut O, waker: Waker) -> CompletionIoResult + where + O: Op, + { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.submit_completion(op, waker), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.submit_completion(op, waker), + AnyDriver::Mock(driver) => driver.submit_completion(op, waker), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.submit_completion(op, waker), + } + } + + #[inline] + pub(crate) fn submit_poll( + &self, + handle: &InnerRawHandle, + waker: Waker, + interest: Interest, + ) -> Result<(), io::Error> { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.submit_poll(handle, waker, interest), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.submit_poll(handle, waker, interest), + AnyDriver::Mock(driver) => driver.submit_poll(handle, waker, interest), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.submit_poll(handle, waker, interest), + } + } + + #[inline] + pub(crate) fn get_completion_result(&self, token: usize) -> Option { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.get_completion_result(token), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.get_completion_result(token), + AnyDriver::Mock(driver) => driver.get_completion_result(token), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.get_completion_result(token), + } + } + + #[inline] + pub(crate) fn set_completion_waker(&self, token: usize, waker: Waker) { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.set_completion_waker(token, waker), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.set_completion_waker(token, waker), + AnyDriver::Mock(driver) => driver.set_completion_waker(token, waker), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.set_completion_waker(token, waker), + } + } + + #[inline] + pub(crate) fn ignore_completion(&self, token: usize, data: Box) { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => driver.ignore_completion(token, data), + #[cfg(unix)] + AnyDriver::Mio(driver) => driver.ignore_completion(token, data), + AnyDriver::Mock(driver) => driver.ignore_completion(token, data), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => driver.ignore_completion(token, data), + } + } + + #[cfg(windows)] + #[inline] + pub(crate) fn cancel_completion( + &self, + token: usize, + handle: crate::fd_inner::RawOsHandle, + data: Box, + ) { + match self { + AnyDriver::Iocp(driver) => driver.cancel_completion(token, handle, data), + AnyDriver::Mock(driver) => driver.ignore_completion(token, data), + } + } + + #[inline] + pub(crate) fn get_interruptor(&self) -> AnyInterruptor { + match self { + #[cfg(windows)] + AnyDriver::Iocp(driver) => AnyInterruptor::Iocp(driver.get_interruptor()), + #[cfg(unix)] + AnyDriver::Mio(driver) => AnyInterruptor::Mio(driver.get_interruptor()), + AnyDriver::Mock(driver) => AnyInterruptor::Mock(driver.get_interruptor()), + #[cfg(target_os = "linux")] + AnyDriver::IoUring(driver) => AnyInterruptor::IoUring(driver.get_interruptor()), + } + } +} + +#[cfg(test)] +mod tests { + use super::AnyDriver; + use std::{ + future::poll_fn, + task::{Poll, Waker}, + }; + + #[cfg(unix)] + #[test] + fn test_mio_driver_interrupt_basic() { + let driver = AnyDriver::new_mio().expect("Failed to create MioDriver"); + let interruptor = driver.get_interruptor(); + + // Test that interrupt doesn't panic and can be called multiple times + interruptor.interrupt(); + interruptor.interrupt(); + interruptor.interrupt(); + } + + #[cfg(target_os = "linux")] + #[test] + fn test_uring_driver_interrupt_basic() { + let driver = AnyDriver::new_uring().expect("Failed to create UringDriver"); + let interruptor = driver.get_interruptor(); + + // Test that interrupt doesn't panic and can be called multiple times + interruptor.interrupt(); + interruptor.interrupt(); + interruptor.interrupt(); + } + + #[test] + fn test_mock_driver_interrupt_basic() { + let driver = AnyDriver::new_mock(); + let interruptor = driver.get_interruptor(); + + // Test that interrupt doesn't panic and can be called multiple times + interruptor.interrupt(); + interruptor.interrupt(); + interruptor.interrupt(); + } + + #[cfg(windows)] + #[test] + fn test_iocp_driver_interrupt_basic() { + let driver = AnyDriver::new_iocp().expect("Failed to create IocpDriver"); + let interruptor = driver.get_interruptor(); + + // Test that interrupt doesn't panic and can be called multiple times + interruptor.interrupt(); + interruptor.interrupt(); + interruptor.interrupt(); + } + + #[cfg(unix)] + #[test] + fn test_interrupt_mio() { + let runtime = crate::executor::Runtime::new( + AnyDriver::new_mio().expect("Failed to create MioDriver"), + ); + + let (tx, rx) = std::sync::mpsc::channel(); + + std::thread::spawn(move || { + let waker: Waker = rx.recv().unwrap(); + drop(rx); // Drop the receiver before waking the task + waker.wake(); + }); + + runtime.block_on(poll_fn(move |cx| { + if tx.send(cx.waker().clone()).is_ok() { + Poll::Pending + } else { + Poll::Ready(()) + } + })); + } + + #[cfg(target_os = "linux")] + #[test] + fn test_interruptor_uring() { + let runtime = crate::executor::Runtime::new( + AnyDriver::new_uring().expect("Failed to create UringDriver"), + ); + + let (tx, rx) = std::sync::mpsc::channel(); + + std::thread::spawn(move || { + let waker: Waker = rx.recv().unwrap(); + drop(rx); // Drop the receiver before waking the task + waker.wake(); + }); + + runtime.block_on(poll_fn(move |cx| { + if tx.send(cx.waker().clone()).is_ok() { + Poll::Pending + } else { + Poll::Ready(()) + } + })); + } + + #[cfg(windows)] + #[test] + fn test_interrupt_iocp() { + let runtime = crate::executor::Runtime::new( + AnyDriver::new_iocp().expect("Failed to create IocpDriver"), + ); + + let (tx, rx) = std::sync::mpsc::channel(); + + std::thread::spawn(move || { + let waker: Waker = rx.recv().unwrap(); + drop(rx); // Drop the receiver before waking the task + waker.wake(); + }); + + runtime.block_on(poll_fn(move |cx| { + if tx.send(cx.waker().clone()).is_ok() { + Poll::Pending + } else { + Poll::Ready(()) + } + })); + } +} diff --git a/vendor/vibeio/src/driver/uring.rs b/vendor/vibeio/src/driver/uring.rs new file mode 100644 index 0000000..403a9bd --- /dev/null +++ b/vendor/vibeio/src/driver/uring.rs @@ -0,0 +1,636 @@ +use std::cell::RefCell; +use std::io::{self, ErrorKind}; +use std::os::fd::RawFd; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc as StdArc; +use std::task::Waker; +use std::time::Duration; + +use io_uring::types::{SubmitArgs, Timespec}; +use io_uring::{opcode, squeue, types, IoUring}; +use mio::{Interest, Token}; +use slab::Slab; + +use crate::driver::{CompletionIoResult, Interruptor}; +use crate::{ + driver::{Driver, RegistrationMode}, + fd_inner::InnerRawHandle, +}; + +const KEY_KIND_BITS: u64 = 1; +const KEY_KIND_MASK: u64 = (1u64 << KEY_KIND_BITS) - 1; +const POLL_KEY_KIND: u8 = 0; +const COMPLETION_KEY_KIND: u8 = 1; + +pub struct UringInterruptor { + eventfd: std::sync::Weak, +} + +impl Interruptor for UringInterruptor { + #[inline] + fn interrupt(&self) { + if let Some(eventfd) = self.eventfd.upgrade() { + // Write to the eventfd to wake up the driver + let value: u64 = 1; + let _ = unsafe { + libc::write( + *eventfd, + &value as *const u64 as *const std::ffi::c_void, + std::mem::size_of::(), + ) + }; + } + } +} + +struct PollRegistration { + fd: RawFd, + poll_mask: u32, + waiter: Option, + poll_armed: bool, +} + +enum HandleRegistration { + Completion, + Poll(PollRegistration), +} + +struct Completion { + waiter: Option, + completed: Option, + ignored_data: Option>, +} + +struct DriverState { + registrations: Slab, + completions: Slab, +} + +pub struct UringDriver { + ring: RefCell, + state: RefCell, + interrupt_eventfd: Option>, + interrupt_buffer: RefCell>, + pending_submissions: AtomicBool, + ext_arg: bool, + timespec: RefCell>, +} + +impl Drop for UringDriver { + fn drop(&mut self) { + if let Some(eventfd) = self.interrupt_eventfd.take() { + // Close eventfd + unsafe { libc::close(*eventfd) }; + } + } +} + +impl UringDriver { + #[inline] + pub(crate) fn new(entries: u32, builder: io_uring::Builder) -> Result { + // Create eventfd for interruption + let eventfd = unsafe { libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_CLOEXEC) }; + if eventfd < 0 { + return Err(io::Error::last_os_error()); + } + + let ring = builder.build(entries)?; + let ext_arg = ring.params().is_feature_ext_arg(); + let driver = Self { + ring: RefCell::new(ring), + state: RefCell::new(DriverState { + registrations: Slab::with_capacity(entries as usize), + completions: Slab::with_capacity(entries as usize), + }), + interrupt_eventfd: Some(StdArc::new(eventfd)), + interrupt_buffer: RefCell::new(Box::new([0; 8])), + pending_submissions: AtomicBool::new(false), + ext_arg, + timespec: RefCell::new(Box::new(Timespec::new())), + }; + + driver.submit_interrupt(); + + Ok(driver) + } + + #[inline] + fn update_waiter(waiter_slot: &mut Option, waker: Waker) { + if !waiter_slot + .as_ref() + .is_some_and(|waiter| waiter.will_wake(&waker)) + { + *waiter_slot = Some(waker); + } + } + + #[inline] + fn encode_completion_key(token: usize) -> u64 { + ((token as u64) << KEY_KIND_BITS) | COMPLETION_KEY_KIND as u64 + } + + #[inline] + fn encode_poll_key(token: Token) -> u64 { + ((token.0 as u64) << KEY_KIND_BITS) | POLL_KEY_KIND as u64 + } + + #[inline] + fn decode_token(key: u64) -> Token { + Token((key >> KEY_KIND_BITS) as usize) + } + + #[inline] + fn decode_key_kind(key: u64) -> u8 { + (key & KEY_KIND_MASK) as u8 + } + + #[inline] + fn interest_to_poll_mask(interest: Interest) -> u32 { + let mut mask = 0; + if interest.is_readable() { + mask |= libc::POLLIN as u32; + } + if interest.is_writable() { + mask |= libc::POLLOUT as u32; + } + mask + } + + #[inline] + fn submitter_call_result(result: Result) -> Result<(), io::Error> { + match result { + Ok(_) => Ok(()), + Err(err) if err.raw_os_error() == Some(libc::EBUSY) => Ok(()), + Err(err) if err.raw_os_error() == Some(libc::ETIME) => Ok(()), // io_uring Timeout + Err(err) => Err(err), + } + } + + #[inline] + fn push_entry(&self, entry: squeue::Entry) -> Result<(), io::Error> { + let mut ring = self.ring.borrow_mut(); + + if ring.submission().is_full() { + Self::submitter_call_result(ring.submit())?; + } + + let mut sq = ring.submission(); + unsafe { + sq.push(&entry) + .map_err(|_| io::Error::other("io_uring submission queue is full"))?; + } + + self.pending_submissions.store(true, Ordering::Release); + + Ok(()) + } + + #[inline] + fn push_poll_add(&self, token: Token, fd: RawFd, poll_mask: u32) -> Result<(), io::Error> { + let entry = opcode::PollAdd::new(types::Fd(fd), poll_mask) + .build() + .user_data(Self::encode_poll_key(token)); + self.push_entry(entry) + } + + #[inline] + fn collect_completions( + &self, + wait_for_one: bool, + timeout: Option, + ) -> Result<(), io::Error> { + { + let mut ring = self.ring.borrow_mut(); + let should_submit = if wait_for_one { + true + } else { + !ring.submission().is_empty() + }; + + if should_submit { + let submit_result = if wait_for_one { + if let Some(timeout) = timeout { + if self.ext_arg { + // Linux 5.11+ + ring.submitter() + .submit_with_args(1, &SubmitArgs::new().timespec(&timeout.into())) + } else { + // Linux 5.4+ + let timespec = timeout.into(); + let mut ts_box = self.timespec.borrow_mut(); + **ts_box = timespec; + let timespec_ptr = &**ts_box as *const Timespec; + + // We must drop the borrow here so we can call push_entry (which borrows ring, but doesn't borrow timespec). + // BUT push_entry borrows ring. We currently have `ring` borrowed mutably! + // We need to drop `ring` borrow before calling `push_entry`? + // `collect_completions` has `let mut ring = self.ring.borrow_mut();` at the top of this block. + // If we call `self.push_entry`, it tries `self.ring.borrow_mut()`, which panics! + // + // FIX: We cannot call `self.push_entry` while holding `ring` borrow. + // We must push manually using the already borrowed `ring`. + + // Duplicate logic of push_entry but using existing `ring` ref? + // Or just push directly to `ring.submission()`. + + let entry = opcode::Timeout::new(timespec_ptr) + .build() + .user_data(u64::MAX - 1); + + if ring.submission().is_full() { + // We are holding borrow, so we can't call methods that borrow. + // But `ring.submit()` is on `IoUring`. + Self::submitter_call_result(ring.submit())?; + } + + { + let mut sq = ring.submission(); + unsafe { + sq.push(&entry).map_err(|_| { + io::Error::other("io_uring submission queue is full") + })?; + } + } + + ring.submit_and_wait(1) + } + } else { + ring.submit_and_wait(1) + } + } else { + ring.submit() + }; + Self::submitter_call_result(submit_result)?; + self.pending_submissions + .store(!ring.submission().is_empty(), Ordering::Release); + } else { + self.pending_submissions.store(false, Ordering::Release); + } + } + + // Drain any new completions produced by the submit above. + let need_interrupt = { + let mut ring = self.ring.borrow_mut(); + let mut state = self.state.borrow_mut(); + Self::drain_cq(&mut ring, &mut state) + }; + if need_interrupt { + self.submit_interrupt(); + } + + Ok(()) + } + + /// Drain the completion queue and wake any registered waiters. + #[inline] + fn drain_cq(ring: &mut IoUring, state: &mut DriverState) -> bool { + let mut interrupt = false; + + // Collect wakers in a small inline array to avoid heap allocation + // in the common case (0-8 completions per collect_completions call). + // Most flush/wait calls produce very few completions. + let mut fast_wakers: [Option; 8] = Default::default(); + let mut fast_count = 0; + let mut overflow_wakers: Vec = Vec::new(); + + { + let cq = ring.completion(); + + for cqe in cq { + let key = cqe.user_data(); + let result = cqe.result(); + + if key == u64::MAX { + // Task interrupted + interrupt = true; + continue; + } else if key == u64::MAX - 1 { + // Timeout (Linux <5.10) + continue; + } + + let token = Self::decode_token(key); + let key_kind = Self::decode_key_kind(key); + + if key_kind == POLL_KEY_KIND { + let waiter = match state.registrations.get_mut(token.0) { + Some(HandleRegistration::Poll(registration)) => { + registration.poll_armed = false; + registration.waiter.take() + } + _ => None, + }; + if let Some(waiter) = waiter { + if fast_count < fast_wakers.len() { + fast_wakers[fast_count] = Some(waiter); + } else { + overflow_wakers.push(waiter); + } + fast_count += 1; + } + continue; + } + + let mut remove_completion = false; + let waiter = match state.completions.get_mut(token.0) { + Some(completion) => { + completion.completed = Some(result); + remove_completion = completion.ignored_data.is_some(); + completion.waiter.take() + } + None => None, + }; + if remove_completion { + state.completions.remove(token.0); + } + if let Some(waiter) = waiter { + if fast_count < fast_wakers.len() { + fast_wakers[fast_count] = Some(waiter); + } else { + overflow_wakers.push(waiter); + } + fast_count += 1; + } + } + } + + for waker in fast_wakers.iter_mut().take(fast_count) { + if let Some(w) = waker.take() { + w.wake(); + } + } + for waker in overflow_wakers { + waker.wake(); + } + + interrupt + } + + #[inline] + fn submit_interrupt(&self) { + use io_uring::{opcode, types}; + // Submit a read operation to the eventfd to wake up the driver + let mut buffer = self.interrupt_buffer.borrow_mut(); + let entry = opcode::Read::new( + types::Fd( + *self + .interrupt_eventfd + .as_ref() + .expect("interrupt_eventfd is not initialized") + .as_ref(), + ), + buffer.as_mut_ptr(), + buffer.len() as u32, + ) + .build() + .user_data(u64::MAX); + + // We use push_entry here. It handles submission if full. + // We panic if it fails because we cannot recover (we won't be able to wake up). + if let Err(err) = self.push_entry(entry) { + panic!("io_uring: failed to submit interrupt task: {}", err); + } + } +} + +impl Driver for UringDriver { + type Interruptor = UringInterruptor; + + #[inline] + fn flush(&self) { + match self.collect_completions(false, None) { + Ok(_) => {} + Err(err) if err.kind() == io::ErrorKind::Interrupted => {} + Err(err) => panic!("io_uring submit failed while processing I/O completions: {err}"), + } + } + + #[inline] + fn should_flush(&self) -> bool { + self.pending_submissions.load(Ordering::Acquire) + } + + #[inline] + fn wait(&self, timeout: Option) { + match self.collect_completions(true, timeout) { + Ok(_) => {} + Err(err) if err.kind() == io::ErrorKind::Interrupted => {} + Err(err) => panic!("io_uring submit_and_wait failed while waiting for I/O: {err}"), + } + } + + #[inline] + fn get_interruptor(&self) -> Self::Interruptor { + UringInterruptor { + eventfd: StdArc::downgrade( + self.interrupt_eventfd + .as_ref() + .expect("interrupt_eventfd is not initialized"), + ), + } + } + + #[inline] + fn register_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result { + self.register_handle_with_mode(handle, interest, RegistrationMode::Completion) + } + + #[inline] + fn register_handle_with_mode( + &self, + handle: &InnerRawHandle, + interest: Interest, + mode: RegistrationMode, + ) -> Result { + let mut state = self.state.borrow_mut(); + let entry = state.registrations.vacant_entry(); + let token = Token(entry.key()); + + match mode { + RegistrationMode::Completion => { + entry.insert(HandleRegistration::Completion); + } + RegistrationMode::Poll => { + entry.insert(HandleRegistration::Poll(PollRegistration { + fd: handle.handle, + poll_mask: Self::interest_to_poll_mask(interest), + waiter: None, + poll_armed: false, + })); + } + } + + Ok(token) + } + + #[inline] + fn reregister_handle( + &self, + handle: &InnerRawHandle, + interest: Interest, + ) -> Result<(), io::Error> { + let mut state = self.state.borrow_mut(); + match state.registrations.get_mut(handle.token.0) { + Some(HandleRegistration::Completion) => Ok(()), + Some(HandleRegistration::Poll(registration)) => { + registration.poll_mask = Self::interest_to_poll_mask(interest); + Ok(()) + } + None => Err(io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + )), + } + } + + #[inline] + fn deregister_handle(&self, handle: &InnerRawHandle) -> Result<(), io::Error> { + { + // Cancel any pending io_uring operations for this handle + let ring = self.ring.borrow_mut(); + let _ = ring.submitter().register_sync_cancel( + Some(Timespec::new().nsec(0).sec(0)), + types::CancelBuilder::fd(types::Fd(handle.handle)), + ); + } + + let mut state = self.state.borrow_mut(); + if state.registrations.try_remove(handle.token.0).is_none() { + return Err(io::Error::new( + ErrorKind::NotFound, + format!( + "I/O token {} is not registered with this driver", + handle.token.0 + ), + )); + } + + Ok(()) + } + + #[inline] + fn supports_completion(&self) -> bool { + true + } + + #[inline] + fn submit_poll( + &self, + handle: &InnerRawHandle, + waker: Waker, + interest: Interest, + ) -> Result<(), io::Error> { + let token = handle.token(); + let poll_spec = { + let mut state = self.state.borrow_mut(); + let registration = match state.registrations.get_mut(token.0) { + Some(HandleRegistration::Poll(registration)) => registration, + Some(HandleRegistration::Completion) => { + return Err(io::Error::new( + ErrorKind::Unsupported, + format!( + "I/O token {} is registered for completion mode, not poll mode", + token.0 + ), + )); + } + None => { + return Err(io::Error::new( + ErrorKind::NotFound, + format!("I/O token {} is not registered with this driver", token.0), + )); + } + }; + + Self::update_waiter(&mut registration.waiter, waker); + let desired_mask = Self::interest_to_poll_mask(interest); + registration.poll_mask = desired_mask; + + if registration.poll_armed { + None + } else { + registration.poll_armed = true; + Some((registration.fd, desired_mask)) + } + }; + + if let Some((fd, poll_mask)) = poll_spec { + if let Err(submit_err) = self.push_poll_add(token, fd, poll_mask) { + let mut state = self.state.borrow_mut(); + if let Some(HandleRegistration::Poll(registration)) = + state.registrations.get_mut(token.0) + { + registration.poll_armed = false; + registration.waiter = None; + } + return Err(submit_err); + } + } + + Ok(()) + } + + #[inline] + fn submit_completion(&self, op: &mut O, waker: Waker) -> super::CompletionIoResult + where + O: crate::op::Op, + { + let mut state = self.state.borrow_mut(); + let vacant_completion = state.completions.vacant_entry(); + let token = vacant_completion.key(); + + // Build the SQE. If this fails, return the error. + let entry = match op.build_completion_entry(Self::encode_completion_key(token)) { + Ok(entry) => entry, + Err(err) => return CompletionIoResult::SubmitErr(err), + }; + + // Push the SQE into the submission queue. If this fails, undo the inflight + // flag and clear waiters on the registration. + if let Err(err) = self.push_entry(entry) { + return CompletionIoResult::SubmitErr(err); + } + + // Store the operation in the completions slab. + vacant_completion.insert(Completion { + waiter: Some(waker), + completed: None, + ignored_data: None, + }); + + CompletionIoResult::Retry(token) + } + + #[inline] + fn get_completion_result(&self, token: usize) -> Option { + let mut state = self.state.borrow_mut(); + let completed = state.completions.get(token).and_then(|c| c.completed); + if completed.is_some() { + state.completions.remove(token); + } + completed + } + + #[inline] + fn set_completion_waker(&self, token: usize, waker: Waker) { + let mut state = self.state.borrow_mut(); + if let Some(c) = state.completions.get_mut(token) { + Self::update_waiter(&mut c.waiter, waker); + } + } + + #[inline] + fn ignore_completion(&self, token: usize, data: Box) { + let mut state = self.state.borrow_mut(); + if let Some(c) = state.completions.get_mut(token) { + c.ignored_data = Some(data); + } + } +} diff --git a/vendor/vibeio/src/executor.rs b/vendor/vibeio/src/executor.rs new file mode 100644 index 0000000..24f129d --- /dev/null +++ b/vendor/vibeio/src/executor.rs @@ -0,0 +1,892 @@ +//! Async runtime and task execution utilities. +//! +//! This module provides the core async runtime infrastructure: +//! - `Runtime`: the main async runtime that drives futures to completion. +//! - `spawn`: spawn a task on the current runtime. +//! - `spawn_blocking`: spawn a blocking task on the thread pool. +//! - `JoinHandle`: a handle to a spawned task that can be awaited. +//! - `current_driver`: get the driver for the current runtime. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::RuntimeBuilder; +//! +//! let runtime = RuntimeBuilder::new().build().unwrap(); +//! let value = runtime.block_on(async { +//! // Spawn a task and await its result +//! let handle = vibeio::spawn(async { 42 }); +//! handle.await + 10 +//! }); +//! assert_eq!(value, 52); +//! ``` +//! +//! # Implementation notes +//! - The runtime is single-threaded and uses a work-stealing queue. +//! - Tasks are polled in batches for better performance. +//! - The runtime supports timers, blocking pools, and file I/O offloading via features. + +use std::cell::{RefCell, UnsafeCell}; +use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::rc::Rc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + +use crossbeam_queue::SegQueue; +use slab::Slab; + +#[cfg(feature = "blocking-default")] +use crate::blocking::DefaultBlockingThreadPool; +use crate::blocking::{BlockingThreadPool, SpawnBlockingError}; +use crate::driver::{AnyDriver, AnyInterruptor}; +#[cfg(feature = "process")] +use crate::process::{start_zombie_reaper, ZombieReaperMessage}; +use crate::task::Task; +#[cfg(feature = "time")] +use crate::timer::Timer; + +thread_local! { + static CURRENT_RUNTIME: RefCell>> = const { RefCell::new(None) }; +} + +/// Internal state for a spawned task. +/// +/// Stores the task's output and a waker to notify when the task completes. +struct JoinState { + output: Option, + waker: Option, + canceled: bool, + task: std::sync::Weak, +} + +#[inline] +fn update_waker_slot(waiter_slot: &mut Option, waker: &Waker) { + if !waiter_slot + .as_ref() + .is_some_and(|waiter| waiter.will_wake(waker)) + { + *waiter_slot = Some(waker.clone()); + } +} + +struct SpawnFuture { + future: F, + state: Rc>>, +} + +impl Future for SpawnFuture +where + F: Future, +{ + type Output = (); + + #[inline] + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // SAFETY: once `self` is pinned we never move `future`. + let this = unsafe { self.get_unchecked_mut() }; + + if this.state.borrow().canceled { + return Poll::Ready(()); + } + + // SAFETY: `future` is pinned together with `self`. + match unsafe { Pin::new_unchecked(&mut this.future) }.poll(cx) { + Poll::Ready(output) => { + let mut state = this.state.borrow_mut(); + state.output = Some(output); + if let Some(waker) = state.waker.take() { + waker.wake(); + } + Poll::Ready(()) + } + Poll::Pending => Poll::Pending, + } + } +} + +/// A handle to a spawned asynchronous task. +/// +/// This handle implements `Future` and can be `await`ed to retrieve the task's output. +/// It allows you to wait for a spawned task to complete and get its result. +/// +/// # Examples +/// ```ignore +/// let handle = vibeio::spawn(async { 42 }); +/// let result = handle.await; // result == 42 +/// ``` +pub struct JoinHandle { + state: Rc>>, +} + +impl JoinHandle { + /// Creates a new `JoinHandle` with the given state. + #[inline] + fn new(state: Rc>>) -> Self { + Self { state } + } + + /// Cancels the task associated with this handle. + /// + /// The task will be interrupted and not resumed. + #[inline] + pub fn cancel(self) { + let task = { + let mut state = self.state.borrow_mut(); + state.canceled = true; + state.task.upgrade() + }; + + if let Some(task) = task { + // JoinHandle is !Send, and tasks are bound to this single-threaded + // runtime, so cancellation can drop the future synchronously. This + // is important for pending overlapped I/O: dropping the operation + // initiates cancellation before the caller can reuse the socket. + task.future.borrow_mut().take(); + task.waker().wake(); + } + } +} + +impl Future for JoinHandle { + type Output = T; + + #[inline] + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut state = self.state.borrow_mut(); + if let Some(output) = state.output.take() { + Poll::Ready(output) + } else { + update_waker_slot(&mut state.waker, cx.waker()); + Poll::Pending + } + } +} + +struct BlockOnNotify { + ready: AtomicBool, + thread_id: std::thread::ThreadId, + interruptor: AnyInterruptor, + waiting: Arc, + interrupt_pending: Arc, +} + +impl BlockOnNotify { + #[inline] + fn new( + interruptor: AnyInterruptor, + waiting: Arc, + interrupt_pending: Arc, + ) -> Arc { + Arc::new(Self { + ready: AtomicBool::new(true), + thread_id: std::thread::current().id(), + interruptor, + waiting, + interrupt_pending, + }) + } + + #[inline] + fn is_ready(&self) -> bool { + self.ready.load(Ordering::Acquire) + } + + #[inline] + fn take_ready(&self) -> bool { + self.ready.swap(false, Ordering::AcqRel) + } + + #[inline] + fn wake_by_ref(&self) { + self.ready.store(true, Ordering::Release); + + if std::thread::current().id() != self.thread_id + && self.waiting.load(Ordering::Acquire) + && !self.interrupt_pending.swap(true, Ordering::AcqRel) + { + self.interruptor.interrupt(); + } + } + + #[inline] + fn waker(self: &Arc) -> Waker { + // SAFETY: the vtable methods correctly clone/drop the Arc reference count. + unsafe { Waker::from_raw(Self::raw_waker(Arc::into_raw(Arc::clone(self)) as *const ())) } + } + + #[inline] + unsafe fn raw_waker(ptr: *const ()) -> RawWaker { + RawWaker::new(ptr, &Self::VTABLE) + } + + const VTABLE: RawWakerVTable = RawWakerVTable::new( + Self::raw_waker_clone, + Self::raw_waker_wake, + Self::raw_waker_wake_by_ref, + Self::raw_waker_drop, + ); + + #[inline] + unsafe fn raw_waker_clone(ptr: *const ()) -> RawWaker { + let notify = Arc::::from_raw(ptr as *const Self); + let cloned = Arc::clone(¬ify); + let _ = Arc::into_raw(notify); + Self::raw_waker(Arc::into_raw(cloned) as *const ()) + } + + #[inline] + unsafe fn raw_waker_wake(ptr: *const ()) { + let notify = Arc::::from_raw(ptr as *const Self); + notify.wake_by_ref(); + } + + #[inline] + unsafe fn raw_waker_wake_by_ref(ptr: *const ()) { + let notify = Arc::::from_raw(ptr as *const Self); + notify.wake_by_ref(); + let _ = Arc::into_raw(notify); + } + + #[inline] + unsafe fn raw_waker_drop(ptr: *const ()) { + drop(Arc::::from_raw(ptr as *const Self)); + } +} + +struct CurrentRuntimeGuard; + +impl CurrentRuntimeGuard { + /// Enter a runtime, making it available for spawning tasks. + /// + /// Panics if called while already inside a runtime. + #[inline] + fn enter(runtime_inner: Rc) -> Self { + CURRENT_RUNTIME.with(|runtime| { + let mut runtime = runtime.borrow_mut(); + if runtime.is_some() { + panic!("can't spawn a runtime inside another runtime"); + } + + *runtime = Some(runtime_inner); + }); + + Self + } +} + +impl Drop for CurrentRuntimeGuard { + /// Exit the runtime, clearing the current runtime reference. + #[inline] + fn drop(&mut self) { + CURRENT_RUNTIME.with(|runtime| { + let mut runtime = runtime.borrow_mut(); + *runtime = None; + }); + } +} + +/// Get the I/O driver for the current runtime. +/// +/// Returns `None` if called outside a runtime context. +pub(crate) fn current_driver() -> Option> { + CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + runtime + .as_ref() + .map(|runtime_inner| runtime_inner.driver.clone()) + }) +} + +/// Get the timer for the current runtime. +/// +/// Returns `None` if called outside a runtime context or if timers are not enabled. +#[cfg(feature = "time")] +pub(crate) fn current_timer() -> Option> { + CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + runtime.as_ref().map(|runtime_inner| { + runtime_inner + .timer + .as_ref() + .expect("timer not enabled") + .clone() + }) + }) +} + +/// Get the zombie reaper channel for the current runtime. +/// +/// Returns `None` if called outside a runtime context or if process support is not enabled. +#[cfg(feature = "process")] +pub(crate) async fn current_zombie_reaper() -> Option> { + let runtime = CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + runtime.as_ref().map(|runtime_inner| runtime_inner.clone()) + })?; + let option = runtime + .zombie_reaper + .try_borrow() + .ok() + .and_then(|e| e.as_ref().cloned()); + if let Some(option) = option { + Some(option.clone()) + } else { + let reaper = runtime.spawn(start_zombie_reaper()).await; + if let Ok(mut option) = runtime.zombie_reaper.try_borrow_mut() { + *option = Some(reaper.clone()); + } + Some(reaper) + } +} + +/// Spawn a task on the current runtime. +/// +/// This function spawns the given future on the runtime and returns a `JoinHandle` +/// that can be awaited to get the task's output. +/// +/// # Panics +/// Panics if called outside a runtime context. +/// +/// # Examples +/// ```ignore +/// let handle = vibeio::spawn(async { 42 }); +/// let result = handle.await; +/// ``` +pub fn spawn(future: impl Future + 'static) -> JoinHandle +where + T: 'static, +{ + let runtime = CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + if let Some(runtime_inner) = &*runtime { + runtime_inner.clone() + } else { + panic!("can't spawn a task outside runtime"); + } + }); + + runtime.spawn(future) +} + +/// Spawn a blocking task on the thread pool. +/// +/// This function spawns the given closure on a blocking thread pool and returns +/// a future that resolves to the result. +/// +/// # Panics +/// Panics if called outside a runtime context. +/// +/// # Examples +/// ```ignore +/// let result = vibeio::spawn_blocking(|| { +/// // blocking work +/// 42 +/// }).await?; +/// ``` +pub async fn spawn_blocking(f: F) -> Result +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, +{ + let runtime = CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + if let Some(runtime_inner) = &*runtime { + runtime_inner.clone() + } else { + panic!("can't spawn a blocking task outside runtime"); + } + }); + + runtime.spawn_blocking(f).await +} + +/// Check if file I/O should be offloaded to blocking threads. +/// +/// Returns `true` if fs offload is enabled and we're inside a runtime. +#[cfg(feature = "fs")] +#[inline] +pub(crate) fn offload_fs() -> bool { + CURRENT_RUNTIME.with(|runtime| { + let runtime = runtime.borrow(); + if let Some(runtime_inner) = &*runtime { + runtime_inner.fs_offload + } else { + false + } + }) +} + +pub(crate) struct RuntimeInner { + queue: Rc>>>, + next_task: Rc>>>, + remote_queue: Arc>, + token_to_task: RefCell>>, + driver: Rc, + waiting: Arc, + interrupt_pending: Arc, + blocking_pool: Option>, + #[cfg(feature = "fs")] + fs_offload: bool, + #[cfg(feature = "time")] + timer: Option>, + #[cfg(feature = "process")] + zombie_reaper: RefCell>>, +} + +/// The async runtime that drives futures to completion. +/// +/// The runtime provides: +/// - Task spawning via `spawn()` and `block_on()`. +/// - Blocking task support via `spawn_blocking()`. +/// - Timer support (when the `time` feature is enabled). +/// - File I/O offloading (when the `fs` feature is enabled). +/// +/// # Examples +/// ```ignore +/// let runtime = RuntimeBuilder::new().build().unwrap(); +/// let value = runtime.block_on(async { 42 }); +/// assert_eq!(value, 42); +/// ``` +pub struct Runtime { + inner: Option>, +} + +impl RuntimeInner { + /// Spawn a task on this runtime. + #[inline] + pub(crate) fn spawn(&self, future: impl Future + 'static) -> JoinHandle + where + T: 'static, + { + let state = Rc::new(RefCell::new(JoinState { + output: None, + waker: None, + canceled: false, + task: std::sync::Weak::new(), + })); + let future = Box::pin(SpawnFuture { + future, + state: state.clone(), + }); + + let mut slab = self.token_to_task.borrow_mut(); + let vacant_slab_entry = slab.vacant_entry(); + #[allow(clippy::arc_with_non_send_sync)] + let task = Arc::new(Task { + future: RefCell::new(Some(future)), + queue: Rc::downgrade(&self.queue), + next_task: Rc::downgrade(&self.next_task), + remote_queue: Arc::downgrade(&self.remote_queue), + queued: AtomicBool::new(true), + thread_id: std::thread::current().id(), + interruptor: self.driver.get_interruptor(), + waiting: Arc::downgrade(&self.waiting), + interrupt_pending: Arc::downgrade(&self.interrupt_pending), + token: vacant_slab_entry.key(), + }); + state.borrow_mut().task = Arc::downgrade(&task); + vacant_slab_entry.insert(task.clone()); + + self.enqueue(task); + JoinHandle::new(state) + } + + /// Spawn a blocking task on this runtime's thread pool. + #[inline] + pub(crate) async fn spawn_blocking(&self, f: F) -> Result + where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, + { + let pool = self + .blocking_pool + .as_ref() + .expect("blocking pool not initialized"); + crate::blocking::spawn_blocking(pool.as_ref(), f).await + } + + /// Enqueue a task for polling. + #[inline] + fn enqueue(&self, task: Arc) { + // SAFETY: this runtime is single-threaded. All ready-queue mutation goes + // through runtime/task wake paths on the same thread. + unsafe { + (&mut *self.queue.get()).push_back(task); + } + } + + /// Drain ready tasks into the given batch. + #[inline] + fn drain_ready(&self, batch: &mut Vec>, mut budget: usize) { + if budget != 0 { + let slab = self.token_to_task.borrow(); + while budget != 0 { + let Some(token) = self.remote_queue.pop() else { + break; + }; + if let Some(task) = slab.get(token) { + task.mark_dequeued(); + batch.push(task.clone()); + budget -= 1; + } + } + } + + // SAFETY: this runtime is single-threaded and we only hold this mutable + // access while draining the queue before polling any task futures. + let queue = unsafe { &mut *self.queue.get() }; + while budget != 0 { + let Some(task) = queue.pop_front() else { + break; + }; + task.mark_dequeued(); + batch.push(task); + budget -= 1; + } + } + + #[inline] + fn stop_waiting(&self) { + self.waiting.store(false, Ordering::Release); + self.interrupt_pending.store(false, Ordering::Release); + } + + #[inline] + fn should_skip_wait(&self) -> bool { + if self.next_task.borrow().is_some() || !self.remote_queue.is_empty() { + return true; + } + + // SAFETY: the runtime only mutates the local ready queue on the runtime thread. + unsafe { !(&*self.queue.get()).is_empty() } + } + + /// Take the next task to run, if any. + #[inline] + fn take_next_task(&self) -> Option> { + let task = self.next_task.take(); + if let Some(task) = &task { + task.mark_dequeued(); + } + task + } +} + +impl Runtime { + /// Create a new runtime with the given driver. + /// + /// By default, this enables the timer and file I/O offload. + #[allow(dead_code)] + #[inline] + pub(crate) fn new(driver: AnyDriver) -> Self { + #[cfg(not(feature = "blocking-default"))] + let blocking_pool = None; + #[cfg(feature = "blocking-default")] + let blocking_pool: Option> = + Some(Box::new(DefaultBlockingThreadPool::new())); + Self::with_options(driver, true, blocking_pool, true) + } + + /// Create a new runtime with the given driver and options. + #[inline] + pub(crate) fn with_options( + driver: AnyDriver, + enable_timer: bool, + blocking_pool: Option>, + fs_offload: bool, + ) -> Self { + #[cfg(not(feature = "fs"))] + let _ = fs_offload; + + let ready_queue = Rc::new(UnsafeCell::new(VecDeque::with_capacity(4096))); + Runtime { + inner: Some(Rc::new(RuntimeInner { + queue: ready_queue, + next_task: Rc::new(RefCell::new(None)), + remote_queue: Arc::new(SegQueue::new()), + token_to_task: RefCell::new(Slab::with_capacity(4096)), + driver: Rc::new(driver), + waiting: Arc::new(AtomicBool::new(false)), + interrupt_pending: Arc::new(AtomicBool::new(false)), + blocking_pool, + #[cfg(feature = "fs")] + fs_offload, + #[cfg(feature = "time")] + timer: if enable_timer { + Some(Rc::new(Timer::new())) + } else { + None + }, + #[cfg(feature = "process")] + zombie_reaper: RefCell::new(None), + })), + } + } + + /// Spawn a task on this runtime. + /// + /// Returns a `JoinHandle` that can be awaited to get the task's output. + #[inline] + pub fn spawn(&self, future: impl Future + 'static) -> JoinHandle + where + T: 'static, + { + self.inner + .as_ref() + .expect("runtime has been dropped") + .spawn(future) + } + + /// Spawn a blocking task on this runtime's thread pool. + #[inline] + pub async fn spawn_blocking(&self, f: F) -> Result + where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, + { + let inner = self.inner.as_ref().expect("runtime has been dropped"); + inner.spawn_blocking(f).await + } + + /// Run the runtime and execute the given future to completion. + /// + /// This method blocks the current thread and drives the runtime until + /// the provided future completes. + #[inline] + pub fn block_on(&self, future: impl Future + 'static) -> T + where + T: 'static, + { + let inner = self.inner.as_ref().expect("runtime has been dropped"); + let _runtime_guard = CurrentRuntimeGuard::enter(inner.clone()); + + let mut future = std::pin::pin!(future); + let root_notify = BlockOnNotify::new( + inner.driver.get_interruptor(), + inner.waiting.clone(), + inner.interrupt_pending.clone(), + ); + let root_waker = root_notify.waker(); + let mut batch = Vec::with_capacity(256); + + loop { + if root_notify.take_ready() { + let mut context = Context::from_waker(&root_waker); + if let Poll::Ready(output) = future.as_mut().poll(&mut context) { + return output; + } + } + + let mut next_task_taken = false; + + batch.clear(); + + let mut budget = 256; + if let Some(next_task) = inner.take_next_task() { + batch.push(next_task); + next_task_taken = true; + budget -= 1; + } + // Always drain to fill the rest of the batch + if budget > 0 { + inner.drain_ready(&mut batch, budget); + } + + if batch.is_empty() { + if root_notify.is_ready() { + continue; + } + + inner.interrupt_pending.store(false, Ordering::Release); + inner.waiting.store(true, Ordering::Release); + + if root_notify.is_ready() || inner.should_skip_wait() { + inner.stop_waiting(); + continue; + } + + #[cfg(feature = "time")] + let (deadline, woken_up) = if let Some(timer) = inner.timer.as_ref() { + // Spin the timing wheel + timer.spin_and_get_deadline() + } else { + (None, false) + }; + + #[cfg(feature = "time")] + if woken_up { + inner.stop_waiting(); + continue; + } + + #[cfg(feature = "time")] + inner.driver.wait(deadline); + #[cfg(not(feature = "time"))] + inner.driver.wait(None); + + inner.stop_waiting(); + continue; + } + + #[cfg(feature = "time")] + if batch.len() > 64 { + if let Some(timer) = inner.timer.as_ref() { + // Spin the timing wheel to avoid starving timers + let _ = timer.spin_and_get_deadline(); + } + } + + for task in batch.drain(..) { + let mut future_slot = task.future.borrow_mut(); + if let Some(mut future) = future_slot.take() { + drop(future_slot); + let waker = task.waker(); + let mut context = Context::from_waker(&waker); + + if future.as_mut().poll(&mut context).is_pending() { + let mut future_slot = task.future.borrow_mut(); + *future_slot = Some(future); + } else { + // Future completed, remove task from token_to_task slab to prevent memory leaks + inner.token_to_task.borrow_mut().remove(task.token); + } + } else { + // Cancellation can synchronously drop the future and then + // enqueue the task so its slab entry is reclaimed here. + // Check identity in case a stale wake targets a reused token. + let should_remove = inner + .token_to_task + .borrow() + .get(task.token) + .is_some_and(|current| Arc::ptr_eq(current, &task)); + if should_remove { + inner.token_to_task.borrow_mut().remove(task.token); + } + } + } + + if !next_task_taken && inner.driver.should_flush() { + inner.driver.flush(); + } + } + } +} + +impl Drop for Runtime { + fn drop(&mut self) { + // Drop all tasks with current runtime entered + let inner = self.inner.take().expect("runtime has been dropped"); + #[cfg(feature = "process")] + if let Some(zombie_reaper) = inner.zombie_reaper.borrow_mut().take() { + zombie_reaper.close(); + } + let _runtime_guard = CurrentRuntimeGuard::enter(inner.clone()); + drop(inner); + drop(_runtime_guard); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + + struct PendingUntilDropped(Rc>); + + impl Future for PendingUntilDropped { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Pending + } + } + + impl Drop for PendingUntilDropped { + fn drop(&mut self) { + self.0.set(true); + } + } + + #[test] + fn block_on_returns_future_output() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let value = runtime.block_on(async { 42usize }); + assert_eq!(value, 42); + } + + #[test] + fn spawn_join_handle_returns_task_output() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let value = runtime.block_on(async { + let handle = spawn(async { 21usize }); + handle.await * 2 + }); + assert_eq!(value, 42); + } + + #[test] + fn runtime_spawn_returns_join_handle() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let handle = runtime.spawn(async { 7usize }); + let value = runtime.block_on(handle); + assert_eq!(value, 7); + } + + #[test] + fn cancel_drops_and_reclaims_pending_task() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let dropped = Rc::new(Cell::new(false)); + let handle = runtime.spawn(PendingUntilDropped(dropped.clone())); + + handle.cancel(); + assert!(dropped.get()); + + let mut first_poll = true; + runtime.block_on(std::future::poll_fn(move |cx| { + if first_poll { + first_poll = false; + cx.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(()) + } + })); + assert!(runtime + .inner + .as_ref() + .unwrap() + .token_to_task + .borrow() + .is_empty()); + } + + #[test] + fn block_on_repolls_root_future_after_self_wake() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let mut polled_once = false; + let value = runtime.block_on(std::future::poll_fn(move |cx| { + if polled_once { + Poll::Ready(11usize) + } else { + polled_once = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + })); + assert_eq!(value, 11); + } + + #[cfg(feature = "blocking-default")] + #[test] + fn spawn_blocking_returns_task_output() { + let runtime = crate::executor::Runtime::new(AnyDriver::new_mock()); + let value = runtime.block_on(async { + let handle = spawn_blocking(|| 21usize).await.unwrap(); + handle * 2 + }); + assert_eq!(value, 42); + } +} diff --git a/vendor/vibeio/src/fd_inner.rs b/vendor/vibeio/src/fd_inner.rs new file mode 100644 index 0000000..b7bb4e5 --- /dev/null +++ b/vendor/vibeio/src/fd_inner.rs @@ -0,0 +1,177 @@ +use std::rc::Rc; +use std::task::Poll; +use std::{io, task::Context}; + +use mio::{Interest, Token}; + +use crate::{ + driver::{AnyDriver, RegistrationMode}, + executor::current_driver, + op::Op, +}; + +#[cfg(unix)] +pub type RawOsHandle = std::os::fd::RawFd; +#[cfg(windows)] +#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)] +pub enum RawOsHandle { + Socket(std::os::windows::io::RawSocket), + Handle(std::os::windows::io::RawHandle), +} + +pub struct InnerRawHandle { + pub(crate) handle: RawOsHandle, + pub(crate) token: Token, + interest: Interest, + mode: RegistrationMode, + driver: Rc, +} + +impl InnerRawHandle { + #[inline] + pub(crate) fn new(handle: RawOsHandle, interest: Interest) -> Result { + let default_mode = if current_driver() + .as_ref() + .is_some_and(|driver| driver.supports_completion()) + { + RegistrationMode::Completion + } else { + RegistrationMode::Poll + }; + + Self::new_with_mode(handle, interest, default_mode) + } + + #[inline] + pub(crate) fn new_with_mode( + handle: RawOsHandle, + interest: Interest, + mode: RegistrationMode, + ) -> Result { + let driver = current_driver().ok_or_else(|| { + io::Error::new( + io::ErrorKind::NotConnected, + "can't register I/O handle outside runtime", + ) + })?; + Self::new_with_driver_and_mode(&driver, handle, interest, mode) + } + + #[inline] + pub(crate) fn new_with_driver_and_mode( + driver: &Rc, + handle: RawOsHandle, + interest: Interest, + mode: RegistrationMode, + ) -> Result { + let mode = if matches!(mode, RegistrationMode::Completion) && !driver.supports_completion() + { + RegistrationMode::Poll + } else { + mode + }; + let mut inner = InnerRawHandle { + handle, + token: Token(0), + interest, + mode, + driver: driver.clone(), + }; + + inner.token = driver.register_handle_with_mode(&inner, interest, mode)?; + Ok(inner) + } + + #[inline] + pub(crate) fn token(&self) -> Token { + self.token + } + + #[allow(dead_code)] + #[inline] + pub(crate) fn reregister(&self, interest: Interest) -> Result<(), io::Error> { + self.driver.reregister_handle(self, interest) + } + + #[inline] + pub(crate) fn supports_completion(&self) -> bool { + self.driver.supports_completion() + } + + #[inline] + pub(crate) fn uses_completion(&self) -> bool { + self.supports_completion() && matches!(self.mode, RegistrationMode::Completion) + } + + #[allow(dead_code)] + #[inline] + pub(crate) fn mode(&self) -> RegistrationMode { + self.mode + } + + #[inline] + pub(crate) fn rebind_mode( + &mut self, + requested_mode: RegistrationMode, + ) -> Result<(), io::Error> { + let mode = if matches!(requested_mode, RegistrationMode::Completion) + && !self.driver.supports_completion() + { + RegistrationMode::Poll + } else { + requested_mode + }; + + if self.mode == mode { + return Ok(()); + } + + self.driver.deregister_handle(self)?; + self.token = self + .driver + .register_handle_with_mode(self, self.interest, mode)?; + self.mode = mode; + Ok(()) + } + + #[inline] + pub(crate) fn poll_op( + &self, + cx: &mut Context<'_>, + op: &mut O, + ) -> Poll> + where + O: Op, + { + if self.uses_completion() { + op.poll_completion(cx, &self.driver) + } else { + op.poll_poll(cx, &self.driver) + } + } + + #[inline] + pub(crate) fn poll_op_poll( + &self, + cx: &mut Context<'_>, + op: &mut O, + ) -> Poll> + where + O: Op, + { + if self.uses_completion() { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "poll-based I/O operation called on a completion-based I/O handle", + ))); + } + op.poll_poll(cx, &self.driver) + } +} + +impl Drop for InnerRawHandle { + #[inline] + fn drop(&mut self) { + let _ = self.driver.deregister_handle(self); + } +} diff --git a/vendor/vibeio/src/fs/file.rs b/vendor/vibeio/src/fs/file.rs new file mode 100644 index 0000000..cfd0918 --- /dev/null +++ b/vendor/vibeio/src/fs/file.rs @@ -0,0 +1,764 @@ +use std::cell::RefCell; +use std::future::poll_fn; +use std::io::{self, ErrorKind}; +use std::mem::ManuallyDrop; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use mio::Interest; + +#[cfg(unix)] +use std::os::fd::{AsRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawHandle, IntoRawHandle, RawHandle}; + +use crate::fs::Metadata; +use crate::io::{iobuf_to_slice, iobufmut_to_slice, IoBuf, IoBufMut, IoBufWithCursor}; +use crate::{ + driver::RegistrationMode, + executor::current_driver, + fd_inner::InnerRawHandle, + io::{AsyncRead, AsyncWrite}, + op::{ReadAtOp, WriteAtOp}, +}; + +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; + +use crate::fs::open_options::OpenOptions; + +/// A file handle for asynchronous file I/O operations. +/// +/// This struct provides async versions of common file operations like reading, +/// writing, and syncing. It supports both io_uring completion-based I/O on Linux +/// and blocking thread pool fallback for other platforms. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::fs::File; +/// +/// // Open a file for reading +/// let file = File::open("hello.txt").await?; +/// +/// // Read from the file +/// let mut buf = [0u8; 1024]; +/// let (read, buf) = file.read_at(buf, 0).await; +/// let read = read?; +/// +/// println!("Read {} bytes", read); +/// ``` +enum FileIo { + Completion(ManuallyDrop), + Blocking, +} + +/// A file handle for asynchronous file I/O operations. +/// +/// This struct provides async versions of common file operations like reading, +/// writing, and syncing. It supports both io_uring completion-based I/O on Linux +/// and blocking thread pool fallback for other platforms. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::fs::File; +/// +/// // Open a file for reading +/// let file = File::open("hello.txt").await?; +/// +/// // Read from the file +/// let mut buf = [0u8; 1024]; +/// let (read, buf) = file.read_at(buf, 0).await; +/// let read = read?; +/// +/// println!("Read {} bytes", read); +/// ``` +pub struct File { + inner: std::fs::File, + io: FileIo, + cursor: u64, +} + +impl File { + /// Opens a file for reading. + /// + /// This is the async version of [`std::fs::File::open`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `openat` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::File::open`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - `path` does not exist + /// - The process lacks permissions to read the file + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::open("hello.txt").await?; + /// ``` + #[inline] + pub async fn open(path: impl AsRef) -> io::Result { + OpenOptions::new().read(true).open(path).await + } + + /// Opens a file for writing, creating it if it does not exist. + /// + /// This is the async version of [`std::fs::File::create`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `openat` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::File::create`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The file cannot be created + /// - The process lacks permissions to create the file + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::create("hello.txt").await?; + /// ``` + #[inline] + pub async fn create(path: impl AsRef) -> io::Result { + OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(path) + .await + } + + /// Returns a new `OpenOptions` builder. + /// + /// This is a convenience method equivalent to `OpenOptions::new()`. + #[inline] + pub fn options() -> OpenOptions { + OpenOptions::new() + } + + /// Creates a new `File` from a standard library file. + /// + /// This is a convenience method equivalent to `File::from_std_with_cursor(inner, 0)`. + #[inline] + pub fn from_std(inner: std::fs::File) -> io::Result { + Self::from_std_with_cursor(inner, 0) + } + + /// Creates a new `File` from a standard library file with a specified cursor position. + /// + /// This is an internal method used to create a `File` with a custom cursor position. + #[inline] + pub(crate) fn from_std_with_cursor(inner: std::fs::File, cursor: u64) -> io::Result { + let io = if let Some(driver) = current_driver() { + if driver.supports_completion() { + #[cfg(unix)] + let raw_handle = inner.as_raw_fd(); + #[cfg(windows)] + let raw_handle = RawOsHandle::Handle(inner.as_raw_handle()); + + match InnerRawHandle::new_with_driver_and_mode( + &driver, + raw_handle, + Interest::READABLE | Interest::WRITABLE, + RegistrationMode::Completion, + ) { + Ok(handle) => FileIo::Completion(ManuallyDrop::new(handle)), + Err(_) => FileIo::Blocking, + } + } else { + FileIo::Blocking + } + } else { + FileIo::Blocking + }; + + Ok(Self { inner, io, cursor }) + } + + /// Converts the `File` back into a standard library `std::fs::File`. + #[inline] + pub fn into_std(self) -> std::fs::File { + let mut this = ManuallyDrop::new(self); + unsafe { + if let FileIo::Completion(handle) = &mut this.io { + ManuallyDrop::drop(handle); + } + std::ptr::read(&this.inner) + } + } + + /// Returns the completion handle if this file is using io_uring completion. + #[inline] + fn completion_handle(&self) -> Option<&InnerRawHandle> { + match &self.io { + FileIo::Completion(handle) => Some(handle), + FileIo::Blocking => None, + } + } + + /// Reads bytes from the file at a specific offset. + /// + /// This method reads into the provided buffer starting at the given offset. + /// The cursor position of the file is not modified. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `readv` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to synchronous reading. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The read operation fails + /// - The offset is invalid + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::open("hello.txt").await?; + /// let mut buf = [0u8; 1024]; + /// let (read, buf) = file.read_at(buf, 0).await; + /// let read = read?; + /// ``` + #[inline] + pub async fn read_at(&self, mut buf: B, offset: u64) -> (io::Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + if let Some(handle) = self.completion_handle() { + let mut op = ReadAtOp::new(handle, buf, offset); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } else if crate::executor::offload_fs() && current_driver().is_some() { + read_at_in_blocking_pool(&self.inner, buf, offset).await + } else { + let slice = iobufmut_to_slice(&mut buf); + (read_at_blocking(&self.inner, slice, offset), buf) + } + } + + /// Reads bytes from the file at a specific offset, filling the entire buffer. + /// + /// This method reads into the provided buffer starting at the given offset, + /// ensuring the entire buffer is filled. The cursor position of the file is not modified. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `readv` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to synchronous reading. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The read operation fails + /// - The offset is invalid + /// - The file does not contain enough data to fill the buffer + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::open("hello.txt").await?; + /// let mut buf = [0u8; 1024]; + /// let (result, buf) = file.read_exact_at(buf, 0).await; + /// result?; + /// ``` + #[inline] + pub async fn read_exact_at(&self, buf: B, mut offset: u64) -> (io::Result<()>, B) { + let mut buf = IoBufWithCursor::new(buf); + while buf.buf_len() > 0 { + let (read, mut buf_returned) = self.read_at(buf, offset).await; + let read = match read { + Ok(read) => read, + Err(err) => { + return (Err(err), buf_returned.into_inner()); + } + }; + if read == 0 { + return ( + Err(io::Error::new( + ErrorKind::UnexpectedEof, + "failed to fill whole buffer", + )), + buf_returned.into_inner(), + ); + } + + offset = offset.saturating_add(read as u64); + buf_returned.advance(read); + buf = buf_returned + } + + (Ok(()), buf.into_inner()) + } + + /// Writes bytes to the file at a specific offset. + /// + /// This method writes from the provided buffer starting at the given offset. + /// The cursor position of the file is not modified. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `writev` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to synchronous writing. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The write operation fails + /// - The offset is invalid + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::create("hello.txt").await?; + /// let buf = b"Hello, world!"; + /// let (written, buf) = file.write_at(buf.to_vec(), 0).await; + /// let written = written?; + /// ``` + #[inline] + pub async fn write_at(&self, buf: B, offset: u64) -> (io::Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + if let Some(handle) = self.completion_handle() { + let mut op = WriteAtOp::new(handle, buf, offset); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } else if crate::executor::offload_fs() && current_driver().is_some() { + write_at_in_blocking_pool(&self.inner, buf, offset).await + } else { + let slice = iobuf_to_slice(&buf); + (write_at_blocking(&self.inner, slice, offset), buf) + } + } + + /// Writes bytes to the file at a specific offset, writing the entire buffer. + /// + /// This method writes from the provided buffer starting at the given offset, + /// ensuring the entire buffer is written. The cursor position of the file is not modified. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `writev` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to synchronous writing. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The write operation fails + /// - The offset is invalid + /// - The write operation fails to write the entire buffer + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::create("hello.txt").await?; + /// let buf = b"Hello, world!"; + /// let (result, buf) = file.write_exact_at(buf.to_vec(), 0).await; + /// result?; + /// ``` + #[inline] + pub async fn write_exact_at(&self, buf: B, mut offset: u64) -> (io::Result<()>, B) { + let mut buf = IoBufWithCursor::new(buf); + while buf.buf_len() > 0 { + let (written, mut buf_returned) = self.write_at(buf, offset).await; + let written = match written { + Ok(written) => written, + Err(err) => { + return (Err(err), buf_returned.into_inner()); + } + }; + if written == 0 { + return ( + Err(io::Error::new( + ErrorKind::UnexpectedEof, + "failed to write whole buffer", + )), + buf_returned.into_inner(), + ); + } + + offset = offset.saturating_add(written as u64); + buf_returned.advance(written); + buf = buf_returned; + } + + (Ok(()), buf.into_inner()) + } + + /// Synchronizes all data and metadata to disk. + /// + /// This is the async version of [`std::fs::File::sync_all`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `fsync` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::File::sync_all`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The sync operation fails + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::create("hello.txt").await?; + /// file.sync_all().await?; + /// ``` + #[inline] + pub async fn sync_all(&self) -> io::Result<()> { + if let Some(handle) = self.completion_handle() { + #[cfg(target_os = "linux")] + { + let mut op = crate::op::FsyncOp::new(handle, false); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await + } + #[cfg(not(target_os = "linux"))] + { + let _ = handle; + sync_all_in_blocking_pool(&self.inner).await + } + } else if crate::executor::offload_fs() && current_driver().is_some() { + sync_all_in_blocking_pool(&self.inner).await + } else { + sync_all_blocking(&self.inner) + } + } + + /// Synchronizes file data to disk without necessarily syncing metadata. + /// + /// This is the async version of [`std::fs::File::sync_data`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `fsync` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::File::sync_data`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The sync operation fails + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::create("hello.txt").await?; + /// file.sync_data().await?; + /// ``` + #[inline] + pub async fn sync_data(&self) -> io::Result<()> { + if let Some(handle) = self.completion_handle() { + #[cfg(target_os = "linux")] + { + let mut op = crate::op::FsyncOp::new(handle, true); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await + } + #[cfg(not(target_os = "linux"))] + { + let _ = handle; + sync_data_in_blocking_pool(&self.inner).await + } + } else if crate::executor::offload_fs() && current_driver().is_some() { + sync_data_in_blocking_pool(&self.inner).await + } else { + sync_data_blocking(&self.inner) + } + } + + /// Returns the metadata for this file. + /// + /// This is the async version of [`std::fs::File::metadata`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support and glibc/musl v1.2.3+, this uses the `statx` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::File::metadata`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The metadata operation fails + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::File; + /// + /// let file = File::open("hello.txt").await?; + /// let metadata = file.metadata().await?; + /// println!("File size: {} bytes", metadata.len()); + /// ``` + #[inline] + pub async fn metadata(&self) -> io::Result { + if let Some(handle) = self.completion_handle() { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + { + use std::ffi::CString; + + let mut op = crate::op::StatxOp::new( + handle.handle, + CString::new(b"").expect("invalid path"), + libc::AT_EMPTY_PATH, + libc::STATX_ALL, + ); + let statx = poll_fn(move |cx| handle.poll_op(cx, &mut op)).await?; + Ok(Metadata::from_statx(statx)) + } + #[cfg(not(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3))))] + { + let _ = handle; + metadata_in_blocking_pool(&self.inner).await + } + } else if crate::executor::offload_fs() && current_driver().is_some() { + metadata_in_blocking_pool(&self.inner).await + } else { + metadata_blocking(&self.inner) + } + } +} + +#[cfg(unix)] +#[inline] +fn read_at_blocking(file: &std::fs::File, buf: &mut [u8], offset: u64) -> io::Result { + use std::os::unix::fs::FileExt; + file.read_at(buf, offset) +} + +#[cfg(windows)] +#[inline] +fn read_at_blocking(file: &std::fs::File, buf: &mut [u8], offset: u64) -> io::Result { + use std::os::windows::fs::FileExt; + file.seek_read(buf, offset) +} + +#[cfg(unix)] +#[inline] +fn write_at_blocking(file: &std::fs::File, buf: &[u8], offset: u64) -> io::Result { + use std::os::unix::fs::FileExt; + file.write_at(buf, offset) +} + +#[cfg(windows)] +#[inline] +fn write_at_blocking(file: &std::fs::File, buf: &[u8], offset: u64) -> io::Result { + use std::os::windows::fs::FileExt; + file.seek_write(buf, offset) +} + +#[inline] +fn sync_all_blocking(file: &std::fs::File) -> io::Result<()> { + file.sync_all() +} + +#[inline] +fn sync_data_blocking(file: &std::fs::File) -> io::Result<()> { + file.sync_data() +} + +#[inline] +fn metadata_blocking(file: &std::fs::File) -> io::Result { + Ok(Metadata::from_std(file.metadata()?)) +} + +#[inline] +pub(crate) fn blocking_pool_io_error() -> io::Error { + io::Error::other("can't spawn blocking task for file I/O") +} + +#[inline] +async fn read_at_in_blocking_pool( + file: &std::fs::File, + buf: B, + offset: u64, +) -> (io::Result, B) { + let file = match file.try_clone() { + Ok(file) => file, + Err(e) => return (Err(e), buf), + }; + let buf = Arc::new(Mutex::new(RefCell::new(Some(buf)))); + let buf_clone = buf.clone(); + crate::spawn_blocking(move || { + let mut buf = buf_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"); + let temp_slice = iobufmut_to_slice(&mut buf); + let result = read_at_blocking(&file, temp_slice, offset); + (result, buf) + }) + .await + .unwrap_or_else(|_| { + ( + Err(blocking_pool_io_error()), + buf.try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"), + ) + }) +} + +#[inline] +async fn write_at_in_blocking_pool( + file: &std::fs::File, + buf: B, + offset: u64, +) -> (io::Result, B) { + let file = match file.try_clone() { + Ok(file) => file, + Err(e) => return (Err(e), buf), + }; + let buf = Arc::new(Mutex::new(RefCell::new(Some(buf)))); + let buf_clone = buf.clone(); + crate::spawn_blocking(move || { + let buf = buf_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"); + let temp_slice = iobuf_to_slice(&buf); + let result = write_at_blocking(&file, temp_slice, offset); + (result, buf) + }) + .await + .unwrap_or_else(|_| { + ( + Err(blocking_pool_io_error()), + buf.try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"), + ) + }) +} + +#[inline] +async fn sync_all_in_blocking_pool(file: &std::fs::File) -> io::Result<()> { + let file = file.try_clone()?; + crate::spawn_blocking(move || sync_all_blocking(&file)) + .await + .map_err(|_| blocking_pool_io_error())? +} + +#[inline] +async fn sync_data_in_blocking_pool(file: &std::fs::File) -> io::Result<()> { + let file = file.try_clone()?; + crate::spawn_blocking(move || sync_data_blocking(&file)) + .await + .map_err(|_| blocking_pool_io_error())? +} + +#[inline] +async fn metadata_in_blocking_pool(file: &std::fs::File) -> io::Result { + let file = file.try_clone()?; + crate::spawn_blocking(move || metadata_blocking(&file)) + .await + .map_err(|_| blocking_pool_io_error())? +} + +impl AsyncRead for File { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + let (read, buf) = self.read_at(buf, self.cursor).await; + if let Ok(read) = read { + self.cursor = self.cursor.saturating_add(read as u64); + } + (read, buf) + } +} + +impl AsyncWrite for File { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + let (written, buf) = self.write_at(buf, self.cursor).await; + if let Ok(written) = written { + self.cursor = self.cursor.saturating_add(written as u64); + } + (written, buf) + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } +} + +impl Drop for File { + #[inline] + fn drop(&mut self) { + unsafe { + if let FileIo::Completion(handle) = &mut self.io { + ManuallyDrop::drop(handle); + } + } + } +} + +#[cfg(unix)] +impl AsRawFd for File { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for File { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.into_std().into_raw_fd() + } +} + +#[cfg(windows)] +impl AsRawHandle for File { + #[inline] + fn as_raw_handle(&self) -> RawHandle { + self.inner.as_raw_handle() + } +} + +#[cfg(windows)] +impl IntoRawHandle for File { + #[inline] + fn into_raw_handle(self) -> RawHandle { + self.into_std().into_raw_handle() + } +} diff --git a/vendor/vibeio/src/fs/metadata.rs b/vendor/vibeio/src/fs/metadata.rs new file mode 100644 index 0000000..c92caa0 --- /dev/null +++ b/vendor/vibeio/src/fs/metadata.rs @@ -0,0 +1,254 @@ +use std::io; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +/// File metadata information. +/// +/// This type mirrors a subset of [`std::fs::Metadata`]. On supported Linux +/// targets we back it by `statx` (via io_uring), otherwise we delegate to +/// `std::fs::metadata` on a blocking thread. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support and glibc/musl v1.2.3+, this uses the `statx` syscall directly +/// for better async performance. +/// - On other platforms, this uses the standard library's `std::fs::Metadata`. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::fs; +/// +/// let metadata = fs::metadata("hello.txt").await?; +/// println!("File size: {} bytes", metadata.len()); +/// ``` +#[derive(Clone, Debug)] +pub struct Metadata { + inner: MetadataInner, +} + +#[derive(Clone, Debug)] +enum MetadataInner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + Statx(libc::statx), + Std(std::fs::Metadata), +} + +impl Metadata { + /// Creates a new `Metadata` from a standard library `std::fs::Metadata`. + #[inline] + pub(crate) fn from_std(md: std::fs::Metadata) -> Self { + Self { + inner: MetadataInner::Std(md), + } + } + + /// Creates a new `Metadata` from a `libc::statx` structure. + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + #[inline] + pub(crate) fn from_statx(st: libc::statx) -> Self { + Self { + inner: MetadataInner::Statx(st), + } + } + + /// Returns the size of the file in bytes. + #[allow(clippy::len_without_is_empty)] + #[inline] + pub fn len(&self) -> u64 { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => st.stx_size, + MetadataInner::Std(md) => md.len(), + } + } + + /// Returns the file permissions. + #[inline] + pub fn permissions(&self) -> std::fs::Permissions { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => { + use std::os::unix::fs::PermissionsExt; + std::fs::Permissions::from_mode(st.stx_mode as u32) + } + MetadataInner::Std(md) => md.permissions(), + } + } + + /// Returns the file type. + #[inline] + pub fn file_type(&self) -> FileType { + FileType { + is_dir: self.is_dir(), + is_file: self.is_file(), + is_symlink: self.is_symlink(), + } + } + + /// Returns `true` if this metadata is for a directory. + #[inline] + pub fn is_dir(&self) -> bool { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => (st.stx_mode as u32 & libc::S_IFMT) == libc::S_IFDIR, + MetadataInner::Std(md) => md.is_dir(), + } + } + + /// Returns `true` if this metadata is for a regular file. + #[inline] + pub fn is_file(&self) -> bool { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => (st.stx_mode as u32 & libc::S_IFMT) == libc::S_IFREG, + MetadataInner::Std(md) => md.is_file(), + } + } + + /// Returns `true` if this metadata is for a symbolic link. + #[inline] + pub fn is_symlink(&self) -> bool { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => (st.stx_mode as u32 & libc::S_IFMT) == libc::S_IFLNK, + MetadataInner::Std(md) => md.file_type().is_symlink(), + } + } + + /// Returns the time the file was last accessed. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The timestamp is invalid + #[inline] + pub fn accessed(&self) -> io::Result { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => statx_timestamp_to_system_time(&st.stx_atime), + MetadataInner::Std(md) => md.accessed(), + } + } + + /// Returns the time the file was created. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The timestamp is invalid + #[inline] + pub fn created(&self) -> io::Result { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => statx_timestamp_to_system_time(&st.stx_btime), + MetadataInner::Std(md) => md.created(), + } + } + + /// Returns the time the file was last modified. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The timestamp is invalid + #[inline] + pub fn modified(&self) -> io::Result { + match &self.inner { + #[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] + MetadataInner::Statx(st) => statx_timestamp_to_system_time(&st.stx_mtime), + MetadataInner::Std(md) => md.modified(), + } + } +} + +/// Converts a `libc::statx_timestamp` to a `SystemTime`. +#[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] +#[inline] +fn statx_timestamp_to_system_time(ts: &libc::statx_timestamp) -> io::Result { + let secs = ts.tv_sec; + let nanos = ts.tv_nsec; + + if secs >= 0 { + Ok(UNIX_EPOCH + Duration::new(secs as u64, nanos)) + } else { + UNIX_EPOCH + .checked_sub(Duration::new((-secs) as u64, nanos)) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid statx timestamp")) + } +} + +/// A structure representing a type of file with accessors for each file type. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::fs; +/// +/// let metadata = fs::metadata("hello.txt").await?; +/// let file_type = metadata.file_type(); +/// +/// if file_type.is_file() { +/// println!("It's a file!"); +/// } else if file_type.is_dir() { +/// println!("It's a directory!"); +/// } +/// ``` +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +pub struct FileType { + is_dir: bool, + is_file: bool, + is_symlink: bool, +} + +impl FileType { + /// Test whether this file type represents a directory. + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs; + /// + /// let metadata = fs::metadata("my_dir").await?; + /// if metadata.file_type().is_dir() { + /// println!("It's a directory!"); + /// } + /// ``` + #[inline] + pub fn is_dir(&self) -> bool { + self.is_dir + } + + /// Test whether this file type represents a regular file. + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs; + /// + /// let metadata = fs::metadata("hello.txt").await?; + /// if metadata.file_type().is_file() { + /// println!("It's a file!"); + /// } + /// ``` + #[inline] + pub fn is_file(&self) -> bool { + self.is_file + } + + /// Test whether this file type represents a symbolic link. + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs; + /// + /// let metadata = fs::metadata("link_to_file").await?; + /// if metadata.file_type().is_symlink() { + /// println!("It's a symlink!"); + /// } + /// ``` + #[inline] + pub fn is_symlink(&self) -> bool { + self.is_symlink + } +} diff --git a/vendor/vibeio/src/fs/mod.rs b/vendor/vibeio/src/fs/mod.rs new file mode 100644 index 0000000..3671a97 --- /dev/null +++ b/vendor/vibeio/src/fs/mod.rs @@ -0,0 +1,1416 @@ +//! A file system module for `vibeio`. +//! +//! This module provides async versions of common file system operations: +//! - File operations: [`File`] with async read/write methods +//! - Path operations: [`canonicalize`], [`hard_link`], [`rename`], [`remove_dir`], [`remove_file`] +//! - Directory operations: [`create_dir`], [`create_dir_all`], [`symlink_dir`], [`symlink_file`] +//! - File content helpers: [`read`], [`read_to_string`], [`write()`] +//! - Metadata: [`metadata`], [`symlink_metadata`] for file information +//! +//! Implementation notes: +//! - On Linux with io_uring support, some operations use native async syscalls (e.g. `statx`, `linkat`) +//! via the async driver. When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations either offload to a blocking thread pool +//! (if file I/O offload is enabled) or fall back to synchronous std::fs calls. +//! - The runtime must be active when calling these functions; otherwise they will panic. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::fs; +//! +//! // Write to a file +//! fs::write("hello.txt", b"Hello, world!").await?; +//! +//! // Read from a file +//! let contents = fs::read_to_string("hello.txt").await?; +//! println!("File contents: {}", contents); +//! +//! // Create a directory +//! fs::create_dir("my_dir").await?; +//! +//! ``` + +mod file; +mod metadata; +mod open_options; + +#[cfg(target_os = "linux")] +use std::ffi::CString; +use std::path::PathBuf; + +pub use file::*; +pub use metadata::*; +pub use open_options::*; + +#[cfg(windows)] +use windows_sys::Win32::Storage::FileSystem::CreateSymbolicLinkW; + +use crate::io::IoBuf; +use crate::io::{AsyncRead, AsyncWrite}; +#[cfg(target_os = "linux")] +use crate::op::HardLinkOp; +#[cfg(target_os = "linux")] +use crate::op::MkDirOp; +#[cfg(target_os = "linux")] +use crate::op::Op; +#[cfg(target_os = "linux")] +use crate::op::RenameOp; +#[cfg(target_os = "linux")] +use crate::op::SymlinkOp; +#[cfg(target_os = "linux")] +use crate::op::UnlinkOp; + +/// Creates a symbolic link to a directory on Windows. +/// +/// This is a Windows-specific helper function that uses the `CreateSymbolicLinkW` API. +/// For cross-platform symlink creation, use [`symlink_dir`] instead. +/// +/// # Platform-specific behavior +/// +/// - This function is only available on Windows. +/// - It creates a symbolic link to a directory using the Windows API. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(windows)] +pub fn windows_symlink_dir(path: String, target: String) -> std::io::Result<()> { + let path_w: Vec = path.encode_utf16().chain(std::iter::once(0)).collect(); + let target_w: Vec = target.encode_utf16().chain(std::iter::once(0)).collect(); + + unsafe { + let res = CreateSymbolicLinkW( + path_w.as_ptr(), + target_w.as_ptr(), + 1, // SYMBOLIC_LINK_FLAG_DIRECTORY + ); + if !res { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } +} + +/// Creates a symbolic link to a file on Windows. +/// +/// This is a Windows-specific helper function that uses the `CreateSymbolicLinkW` API. +/// For cross-platform symlink creation, use [`symlink_file`] instead. +/// +/// # Platform-specific behavior +/// +/// - This function is only available on Windows. +/// - It creates a symbolic link to a file using the Windows API. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(windows)] +pub fn windows_symlink_file(path: String, target: String) -> std::io::Result<()> { + let path_w: Vec = path.encode_utf16().chain(std::iter::once(0)).collect(); + let target_w: Vec = target.encode_utf16().chain(std::iter::once(0)).collect(); + + unsafe { + let res = CreateSymbolicLinkW( + path_w.as_ptr(), + target_w.as_ptr(), + 0, // SYMBOLIC_LINK_FLAG_FILE + ); + if !res { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } +} + +/// Returns the canonical form of a path with all components normalized. +/// +/// This is the async version of [`std::fs::canonicalize`]. +/// +/// # Platform-specific behavior +/// +/// - On most platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::canonicalize`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - A component in the path is not a directory +/// - The process lacks permissions to access components of the path +pub async fn canonicalize>(path: P) -> std::io::Result { + let path = path.as_ref().to_path_buf(); + if crate::executor::offload_fs() { + crate::spawn_blocking(move || path.canonicalize()) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + path.canonicalize() + } +} + +/// Reads the entire contents of a file into a vector of bytes. +/// +/// This is the async version of [`std::fs::read`]. +/// +/// # Platform-specific behavior +/// +/// - On most platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::read`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions to read the file +pub async fn read(path: impl AsRef) -> std::io::Result> { + let mut file: File = OpenOptions::new().read(true).open(path).await?; + let mut bytes = Vec::new(); + let mut buf = [0u8; 8192]; + + loop { + let (read, returned_buf) = file.read(buf).await; + let read = read?; + buf = returned_buf; + + if read == 0 { + break; + } + + let slice = + unsafe { std::slice::from_raw_parts(buf.as_buf_ptr(), buf.buf_len().min(read)) }; + bytes.extend_from_slice(slice); + } + + Ok(bytes) +} + +/// Reads the entire contents of a file into a string. +/// +/// This is the async version of [`std::fs::read_to_string`]. +/// +/// # Platform-specific behavior +/// +/// - On most platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::read_to_string`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - [`read`] fails +/// - The file contents are not valid UTF-8 +pub async fn read_to_string(path: impl AsRef) -> std::io::Result { + let bytes = read(path).await?; + String::from_utf8(bytes) + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err.utf8_error())) +} + +/// Writes a byte slice to a file, creating it if necessary. +/// +/// This is the async version of [`std::fs::write`]. +/// +/// # Platform-specific behavior +/// +/// - On most platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::write`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - The file cannot be opened for writing +/// - The write operation fails +pub async fn write( + path: impl AsRef, + contents: impl AsRef<[u8]>, +) -> std::io::Result<()> { + let mut file: File = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(path) + .await?; + + let mut slice = contents.as_ref(); + while !slice.is_empty() { + let (w, _) = file.write(slice.to_vec()).await; + let w = w?; + if w == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::WriteZero, + "failed to write whole buffer", + )); + } + slice = &slice[w..]; + } + file.flush().await +} + +/// Creates a hard link at the destination path pointing to the source. +/// +/// This is the async version of [`std::fs::hard_link`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `linkat` syscall directly. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::hard_link`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist +/// - `dst` already exists +/// - The source and destination are on different filesystems +/// - The process lacks permissions +#[cfg(target_os = "linux")] +pub async fn hard_link( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let src_cstr = CString::new(src.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let dst_cstr = CString::new(dst.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + + let driver = driver.expect("invalid driver state"); + let mut op = HardLinkOp::new(src_cstr, dst_cstr); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let src = src.to_owned(); + let dst = dst.to_owned(); + crate::spawn_blocking(move || std::fs::hard_link(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::hard_link(src, dst) + } +} + +/// Creates a hard link at the destination path pointing to the source. +/// +/// This is the async version of [`std::fs::hard_link`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux, this either offloads to a blocking thread pool +/// or falls back to [`std::fs::hard_link`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist +/// - `dst` already exists +/// - The source and destination are on different filesystems +/// - The process lacks permissions +#[cfg(not(target_os = "linux"))] +pub async fn hard_link( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let src = src.as_ref().to_owned(); + let dst = dst.as_ref().to_owned(); + crate::spawn_blocking(move || std::fs::hard_link(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::hard_link(src, dst) + } +} + +/// Creates a symbolic link to a directory. +/// +/// This is the async version of [`std::os::unix::fs::symlink`] (on Unix) or +/// [`std::os::windows::fs::symlink_dir`] (on Windows). +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `symlinkat` syscall directly. +/// - On Windows, this uses the [`windows_symlink_dir`] helper. +/// - On other Unix platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a directory +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(windows)] +pub async fn symlink_dir( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let src_str = src.to_string_lossy().to_string(); + let dst_str = dst.to_string_lossy().to_string(); + crate::spawn_blocking(move || windows_symlink_dir(src_str, dst_str)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let src_str = src.to_string_lossy().to_string(); + let dst_str = dst.to_string_lossy().to_string(); + + windows_symlink_dir(src_str, dst_str) + } +} + +/// Creates a symbolic link to a directory. +/// +/// This is the async version of [`std::os::unix::fs::symlink`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `symlinkat` syscall directly. +/// - On other Unix platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a directory +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(target_os = "linux")] +pub async fn symlink_dir( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + // On Linux with io_uring, use SymlinkOp + let src_cstr = CString::new(src.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let dst_cstr = CString::new(dst.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = SymlinkOp::new(src_cstr, dst_cstr); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let src = src.to_owned(); + let dst = dst.to_owned(); + crate::spawn_blocking(move || std::os::unix::fs::symlink(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::os::unix::fs::symlink(src, dst) + } +} + +/// Creates a symbolic link to a directory. +/// +/// This is the async version of [`std::os::unix::fs::symlink`]. +/// +/// # Platform-specific behavior +/// +/// - On other Unix platforms (not Linux or Windows), this either offloads to a +/// blocking thread pool or falls back to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a directory +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(not(any(windows, target_os = "linux")))] +pub async fn symlink_dir( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let src = src.as_ref().to_owned(); + let dst = dst.as_ref().to_owned(); + crate::spawn_blocking(move || std::os::unix::fs::symlink(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::os::unix::fs::symlink(src, dst) + } +} + +/// Creates a symbolic link to a file. +/// +/// This is the async version of [`std::os::unix::fs::symlink`] (on Unix) or +/// [`std::os::windows::fs::symlink_file`] (on Windows). +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `symlinkat` syscall directly. +/// - On Windows, this uses the [`windows_symlink_file`] helper. +/// - On other Unix platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a file +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(windows)] +pub async fn symlink_file( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let src_str = src.to_string_lossy().to_string(); + let dst_str = dst.to_string_lossy().to_string(); + crate::spawn_blocking(move || windows_symlink_file(src_str, dst_str)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let src_str = src.to_string_lossy().to_string(); + let dst_str = dst.to_string_lossy().to_string(); + + windows_symlink_file(src_str, dst_str) + } +} + +/// Creates a symbolic link to a file. +/// +/// This is the async version of [`std::os::unix::fs::symlink`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `symlinkat` syscall directly. +/// - On other Unix platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a file +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(target_os = "linux")] +pub async fn symlink_file( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + let src = src.as_ref(); + let dst = dst.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + // On Linux with io_uring, use SymlinkOp + let src_cstr = CString::new(src.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let dst_cstr = CString::new(dst.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = SymlinkOp::new(src_cstr, dst_cstr); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let src = src.to_owned(); + let dst = dst.to_owned(); + crate::spawn_blocking(move || std::os::unix::fs::symlink(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::os::unix::fs::symlink(src, dst) + } +} + +/// Creates a symbolic link to a file. +/// +/// This is the async version of [`std::os::unix::fs::symlink`]. +/// +/// # Platform-specific behavior +/// +/// - On other Unix platforms (not Linux or Windows), this either offloads to a +/// blocking thread pool or falls back to [`std::os::unix::fs::symlink`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `src` does not exist or is not a file +/// - `dst` already exists +/// - The process lacks permissions to create the symlink +/// - The platform does not support symbolic links +#[cfg(not(any(windows, target_os = "linux")))] +pub async fn symlink_file( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let src = src.as_ref().to_owned(); + let dst = dst.as_ref().to_owned(); + crate::spawn_blocking(move || std::os::unix::fs::symlink(src, dst)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::os::unix::fs::symlink(src, dst) + } +} + +/// Creates a symbolic link. +/// +/// This is a convenience function that calls [`symlink_file`]. Use this when you +/// don't know or don't care whether the source is a file or directory. +/// +/// For explicit symlink creation, use [`symlink_file`] or [`symlink_dir`] instead. +/// +/// # Platform-specific behavior +/// +/// See [`symlink_file`] for platform-specific behavior details. +/// +/// # Errors +/// +/// See [`symlink_file`] for error conditions. +pub async fn symlink( + src: impl AsRef, + dst: impl AsRef, +) -> std::io::Result<()> { + symlink_file(src, dst).await +} + +/// Renames a file or directory to a new location. +/// +/// This is the async version of [`std::fs::rename`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `renameat` syscall directly. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::rename`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `from` does not exist +/// - `to` already exists and is not overwritable +/// - The source and destination are on different filesystems +/// - The process lacks permissions +#[cfg(target_os = "linux")] +pub async fn rename( + from: impl AsRef, + to: impl AsRef, +) -> std::io::Result<()> { + let from = from.as_ref(); + let to = to.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let from_cstr = CString::new(from.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let to_cstr = CString::new(to.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = RenameOp::new(from_cstr, to_cstr); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let from = from.to_owned(); + let to = to.to_owned(); + crate::spawn_blocking(move || std::fs::rename(from, to)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::rename(from, to) + } +} + +/// Renames a file or directory to a new location. +/// +/// This is the async version of [`std::fs::rename`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux, this either offloads to a blocking thread pool +/// or falls back to [`std::fs::rename`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `from` does not exist +/// - `to` already exists and is not overwritable +/// - The source and destination are on different filesystems +/// - The process lacks permissions +#[cfg(not(target_os = "linux"))] +pub async fn rename( + from: impl AsRef, + to: impl AsRef, +) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let from = from.as_ref().to_owned(); + let to = to.as_ref().to_owned(); + crate::spawn_blocking(move || std::fs::rename(from, to)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::rename(from, to) + } +} + +/// Removes an empty directory. +/// +/// This is the async version of [`std::fs::remove_dir`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `unlinkat` syscall directly. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::remove_dir`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - `path` is not a directory +/// - The directory is not empty +/// - The process lacks permissions +#[cfg(target_os = "linux")] +pub async fn remove_dir(path: impl AsRef) -> std::io::Result<()> { + let path = path.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let path_cstr = CString::new(path.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = UnlinkOp::new(path_cstr, true); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let path = path.to_owned(); + crate::spawn_blocking(move || std::fs::remove_dir(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::remove_dir(path) + } +} + +/// Removes an empty directory. +/// +/// This is the async version of [`std::fs::remove_dir`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux, this either offloads to a blocking thread pool +/// or falls back to [`std::fs::remove_dir`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - `path` is not a directory +/// - The directory is not empty +/// - The process lacks permissions +#[cfg(not(target_os = "linux"))] +pub async fn remove_dir(path: impl AsRef) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let path = path.as_ref().to_owned(); + crate::spawn_blocking(move || std::fs::remove_dir(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::remove_dir(path) + } +} + +/// Removes a file. +/// +/// This is the async version of [`std::fs::remove_file`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `unlinkat` syscall directly. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::remove_file`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions +#[cfg(target_os = "linux")] +pub async fn remove_file(path: impl AsRef) -> std::io::Result<()> { + let path = path.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let path_cstr = CString::new(path.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = UnlinkOp::new(path_cstr, false); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let path = path.to_owned(); + crate::spawn_blocking(move || std::fs::remove_file(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::remove_file(path) + } +} + +/// Removes a file. +/// +/// This is the async version of [`std::fs::remove_file`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux, this either offloads to a blocking thread pool +/// or falls back to [`std::fs::remove_file`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions +#[cfg(not(target_os = "linux"))] +pub async fn remove_file(path: impl AsRef) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let path = path.as_ref().to_owned(); + crate::spawn_blocking(move || std::fs::remove_file(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::remove_file(path) + } +} + +/// Creates a directory. +/// +/// This is the async version of [`std::fs::create_dir`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support, this uses the `mkdirat` syscall directly. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::create_dir`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - A component in the path does not exist +/// - A component in the path is not a directory +/// - The process lacks permissions +/// - The directory already exists +#[cfg(target_os = "linux")] +pub async fn create_dir(path: impl AsRef) -> std::io::Result<()> { + let path = path.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let path_cstr = CString::new(path.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + // mode 0o777 is standard for mkdir, umask will be applied + let mut op = MkDirOp::new(path_cstr, 0o777); + std::future::poll_fn(|cx| op.poll_completion(cx, driver.as_ref())).await + } else if crate::executor::offload_fs() { + let path = path.to_owned(); + crate::spawn_blocking(move || std::fs::create_dir(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::create_dir(path) + } +} + +/// Creates a directory. +/// +/// This is the async version of [`std::fs::create_dir`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux, this either offloads to a blocking thread pool +/// or falls back to [`std::fs::create_dir`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - A component in the path does not exist +/// - A component in the path is not a directory +/// - The process lacks permissions +/// - The directory already exists +#[cfg(not(target_os = "linux"))] +pub async fn create_dir(path: impl AsRef) -> std::io::Result<()> { + if crate::executor::offload_fs() { + let path = path.as_ref().to_owned(); + crate::spawn_blocking(move || std::fs::create_dir(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } else { + std::fs::create_dir(path) + } +} + +/// Creates a new, empty directory and all its parent components if they don't exist. +/// +/// This is the async version of [`std::fs::create_dir_all`]. +/// +/// # Platform-specific behavior +/// +/// - This function internally calls [`create_dir`] for each directory component, +/// so it inherits the platform-specific behavior of that function. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - A component in the path cannot be created +/// - A component in the path is not a directory +/// - The process lacks permissions +pub async fn create_dir_all(path: impl AsRef) -> std::io::Result<()> { + let path = path.as_ref(); + let mut stack = Vec::new(); + let mut p = path; + + // Build stack of missing directories + loop { + // Try to create current path + match create_dir(p).await { + Ok(()) => break, + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + // Exists. Check if dir. + if let Ok(metadata) = metadata(p).await { + if metadata.is_dir() { + break; + } + } + return Err(e); + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + // Parent missing. + stack.push(p); + match p.parent() { + Some(parent) => p = parent, + None => return Err(e), + } + } + Err(e) => return Err(e), + } + } + + // Now create directories in stack in reverse order (top to bottom) + while let Some(p) = stack.pop() { + match create_dir(p).await { + Ok(()) => continue, + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + if let Ok(metadata) = metadata(p).await { + if metadata.is_dir() { + continue; + } + } + return Err(e); + } + Err(e) => return Err(e), + } + } + + Ok(()) +} + +/// Returns metadata about a file or directory. +/// +/// This is the async version of [`std::fs::metadata`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support and glibc/musl v1.2.3+, this uses the `statx` syscall directly +/// for better async performance. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::metadata`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions to access the path +#[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] +pub async fn metadata(path: impl AsRef) -> std::io::Result { + let path = path.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let path_cstr = CString::new(path.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = crate::op::StatxOp::new(libc::AT_FDCWD, path_cstr, 0, libc::STATX_ALL); + let statx = std::future::poll_fn(move |cx| op.poll_completion(cx, &driver)).await?; + Ok(Metadata::from_statx(statx)) + } else if crate::executor::offload_fs() { + let path = path.to_owned(); + Ok(Metadata::from_std( + crate::spawn_blocking(move || std::fs::metadata(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())??, + )) + } else { + Ok(Metadata::from_std(std::fs::metadata(path)?)) + } +} + +/// Returns metadata about a file or directory. +/// +/// This is the async version of [`std::fs::metadata`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux with glibc/musl v1.2.3+, this either offloads +/// to a blocking thread pool or falls back to [`std::fs::metadata`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions to access the path +#[cfg(not(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3))))] +pub async fn metadata(path: impl AsRef) -> std::io::Result { + if crate::executor::offload_fs() { + let path = path.as_ref().to_owned(); + Ok(Metadata::from_std( + crate::spawn_blocking(move || std::fs::metadata(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())??, + )) + } else { + Ok(Metadata::from_std(std::fs::metadata(path)?)) + } +} + +/// Returns metadata about a file or directory without following symlinks. +/// +/// This is the async version of [`std::fs::symlink_metadata`]. +/// +/// # Platform-specific behavior +/// +/// - On Linux with io_uring support and glibc/musl v1.2.3+, this uses the `statx` syscall directly +/// with `AT_SYMLINK_NOFOLLOW` flag for better async performance. +/// - On other platforms, this either offloads to a blocking thread pool or falls back +/// to [`std::fs::symlink_metadata`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions to access the path +#[cfg(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3)))] +pub async fn symlink_metadata(path: impl AsRef) -> std::io::Result { + let path = path.as_ref(); + + let driver = crate::executor::current_driver(); + if driver.as_ref().is_some_and(|d| d.supports_completion()) { + let path_cstr = CString::new(path.as_os_str().as_encoded_bytes()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Invalid path: {}", e), + ) + })?; + let driver = driver.expect("invalid driver state"); + let mut op = crate::op::StatxOp::new( + libc::AT_FDCWD, + path_cstr, + libc::AT_SYMLINK_NOFOLLOW, + libc::STATX_ALL, + ); + let statx = std::future::poll_fn(move |cx| op.poll_completion(cx, &driver)).await?; + Ok(Metadata::from_statx(statx)) + } else if crate::executor::offload_fs() { + let path = path.to_owned(); + Ok(Metadata::from_std( + crate::spawn_blocking(move || std::fs::symlink_metadata(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())??, + )) + } else { + Ok(Metadata::from_std(std::fs::symlink_metadata(path)?)) + } +} + +/// Returns metadata about a file or directory without following symlinks. +/// +/// This is the async version of [`std::fs::symlink_metadata`]. +/// +/// # Platform-specific behavior +/// +/// - On platforms other than Linux with glibc/musl v1.2.3+, this either offloads +/// to a blocking thread pool or falls back to [`std::fs::symlink_metadata`]. +/// +/// # Errors +/// +/// This function will return an error in the following situations: +/// - `path` does not exist +/// - The process lacks permissions to access the path +#[cfg(not(all(target_os = "linux", any(target_env = "gnu", musl_v1_2_3))))] +pub async fn symlink_metadata(path: impl AsRef) -> std::io::Result { + if crate::executor::offload_fs() { + let path = path.as_ref().to_owned(); + Ok(Metadata::from_std( + crate::spawn_blocking(move || std::fs::symlink_metadata(path)) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())??, + )) + } else { + Ok(Metadata::from_std(std::fs::symlink_metadata(path)?)) + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + + use crate::{ + driver::AnyDriver, + executor::Runtime, + fs::{metadata, read, read_to_string, write, File, OpenOptions}, + io::AsyncWrite, + }; + + fn unique_path(name: &str) -> PathBuf { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock should be after epoch") + .as_nanos(); + std::env::temp_dir().join(format!("vibeio_{name}_{now}.tmp")) + } + + #[test] + fn fs_read_write_helpers_work() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("helpers"); + write(&path, b"hello world") + .await + .expect("write helper should succeed"); + + let bytes = read(&path).await.expect("read helper should succeed"); + assert_eq!(bytes, b"hello world"); + + let string = read_to_string(&path) + .await + .expect("read_to_string helper should succeed"); + assert_eq!(string, "hello world"); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn file_read_at_and_write_exact_at_work() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("offset"); + let file = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .expect("open for write should succeed"); + file.write_exact_at(b"abcdef".to_vec(), 0) + .await + .0 + .expect("write_exact_at should succeed"); + + let file = File::open(&path) + .await + .expect("open for read should succeed"); + let (read, out) = file.read_exact_at([0u8; 4], 2).await; + read.expect("read_exact_at should succeed"); + assert_eq!(&out, b"cdef"); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn metadata_basic_properties() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("metadata"); + write(&path, b"test content") + .await + .expect("write should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + assert_eq!(md.len(), 12); + assert!(md.is_file()); + assert!(!md.is_dir()); + assert!(!md.is_symlink()); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn metadata_directory() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("dir"); + std::fs::create_dir(&path).expect("create_dir should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + assert!(md.is_dir()); + assert!(!md.is_file()); + + let _ = std::fs::remove_dir(path); + }); + } + + #[test] + fn metadata_timestamps() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("timestamps"); + write(&path, b"test").await.expect("write should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + + // All timestamp methods should return valid SystemTime + let accessed = md.accessed().expect("accessed should succeed"); + let modified = md.modified().expect("modified should succeed"); + let created = md.created().expect("created should succeed"); + + // Timestamps should be reasonable (not in the far future) + let now = SystemTime::now(); + assert!(accessed <= now || accessed + Duration::from_secs(1) >= now); + assert!(modified <= now || modified + Duration::from_secs(1) >= now); + assert!(created <= now || created + Duration::from_secs(1) >= now); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn metadata_permissions() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("perms"); + write(&path, b"test").await.expect("write should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + let perms = md.permissions(); + + // Should be readable and writable by owner + assert!(!perms.readonly(), "file should be writable"); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn metadata_file_type() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("type"); + write(&path, b"test").await.expect("write should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + let file_type = md.file_type(); + + assert!(file_type.is_file()); + assert!(!file_type.is_dir()); + assert!(!file_type.is_symlink()); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn metadata_empty_file() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("empty"); + let mut file = OpenOptions::new() + .write(true) + .create(true) + .open(&path) + .await + .expect("open should succeed"); + file.flush().await.expect("flush should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + assert_eq!(md.len(), 0); + assert!(md.is_file()); + + let _ = std::fs::remove_file(path); + }); + } + + #[test] + fn create_dir_works() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("create_dir"); + crate::fs::create_dir(&path) + .await + .expect("create_dir should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + assert!(md.is_dir()); + + crate::fs::remove_dir(path) + .await + .expect("remove_dir should succeed"); + }); + } + + #[test] + fn create_dir_all_works() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let base = unique_path("create_dir_all"); + let path = base.join("a/b/c"); + + crate::fs::create_dir_all(&path) + .await + .expect("create_dir_all should succeed"); + + let md = metadata(&path).await.expect("metadata should succeed"); + assert!(md.is_dir()); + + // Clean up + crate::fs::remove_dir(base.join("a/b/c")) + .await + .expect("remove_dir c"); + crate::fs::remove_dir(base.join("a/b")) + .await + .expect("remove_dir b"); + crate::fs::remove_dir(base.join("a")) + .await + .expect("remove_dir a"); + crate::fs::remove_dir(&base).await.expect("remove_dir base"); + }); + } + + #[cfg(unix)] + #[test] + fn symlink_metadata_works() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let target = unique_path("symlink_target"); + let link = unique_path("symlink"); + + // Create a file and a symlink to it + write(&target, b"test content") + .await + .expect("write should succeed"); + crate::fs::symlink_file(&target, &link) + .await + .expect("symlink_file should succeed"); + + // symlink_metadata on the symlink should return info about the symlink itself + let md = crate::fs::symlink_metadata(&link) + .await + .expect("symlink_metadata should succeed"); + assert!(md.is_symlink()); + + // metadata on the symlink should follow the link and return info about the target + let target_md = metadata(&link).await.expect("metadata should succeed"); + assert!(target_md.is_file()); + assert!(!target_md.is_symlink()); + assert_eq!(target_md.len(), 12); + + // Clean up + crate::fs::remove_file(&link) + .await + .expect("remove_file should succeed"); + crate::fs::remove_file(&target) + .await + .expect("remove_file should succeed"); + }); + } + + #[test] + fn symlink_metadata_on_regular_file() { + let runtime = Runtime::new(AnyDriver::new_mock()); + runtime.block_on(async { + let path = unique_path("regular_file"); + write(&path, b"content") + .await + .expect("write should succeed"); + + // symlink_metadata on a regular file should work the same as metadata + let md = crate::fs::symlink_metadata(&path) + .await + .expect("symlink_metadata should succeed"); + assert!(md.is_file()); + assert!(!md.is_symlink()); + assert_eq!(md.len(), 7); + + let _ = std::fs::remove_file(path); + }); + } +} diff --git a/vendor/vibeio/src/fs/open_options.rs b/vendor/vibeio/src/fs/open_options.rs new file mode 100644 index 0000000..5991a15 --- /dev/null +++ b/vendor/vibeio/src/fs/open_options.rs @@ -0,0 +1,329 @@ +use std::future::poll_fn; +use std::io::{self, ErrorKind}; +use std::path::Path; + +use crate::executor::current_driver; +use crate::op::Op; + +#[cfg(target_os = "linux")] +use crate::op::OpenOp; + +#[cfg(unix)] +use std::ffi::CString; +#[cfg(unix)] +use std::os::fd::FromRawFd; +#[cfg(unix)] +use std::os::unix::ffi::OsStrExt; + +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; + +use crate::fs::file::File; + +/// Options and flags for opening files. +/// +/// This struct provides a builder-style interface for configuring how a file +/// should be opened, similar to [`std::fs::OpenOptions`]. It supports both +/// io_uring completion-based opening on Linux and blocking thread pool fallback +/// for other platforms. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::fs::OpenOptions; +/// +/// // Open a file for reading +/// let file = OpenOptions::new() +/// .read(true) +/// .open("hello.txt") +/// .await?; +/// +/// // Create a new file for writing (truncate if exists) +/// let file = OpenOptions::new() +/// .write(true) +/// .create(true) +/// .truncate(true) +/// .open("output.txt") +/// .await?; +/// +/// ``` +#[derive(Clone, Debug)] +pub struct OpenOptions { + read: bool, + write: bool, + append: bool, + truncate: bool, + create: bool, + create_new: bool, +} + +impl OpenOptions { + /// Creates a new `OpenOptions` with default values. + /// + /// By default, all options are set to `false`. + #[inline] + pub fn new() -> Self { + Self { + read: false, + write: false, + append: false, + truncate: false, + create: false, + create_new: false, + } + } + + /// Sets whether the file should be opened for reading. + #[inline] + pub fn read(&mut self, read: bool) -> &mut Self { + self.read = read; + self + } + + /// Sets whether the file should be opened for writing. + #[inline] + pub fn write(&mut self, write: bool) -> &mut Self { + self.write = write; + self + } + + /// Sets whether the file should be opened in append mode. + #[inline] + pub fn append(&mut self, append: bool) -> &mut Self { + self.append = append; + self + } + + /// Sets whether the file should be truncated if it already exists. + #[inline] + pub fn truncate(&mut self, truncate: bool) -> &mut Self { + self.truncate = truncate; + self + } + + /// Sets whether the file should be created if it does not exist. + #[inline] + pub fn create(&mut self, create: bool) -> &mut Self { + self.create = create; + self + } + + /// Sets whether the file should be created exclusively (fails if it already exists). + #[inline] + pub fn create_new(&mut self, create_new: bool) -> &mut Self { + self.create_new = create_new; + self + } + + /// Validates the open options. + /// + /// This is an internal method used to ensure the options are valid. + #[inline] + fn validate(&self) -> io::Result<()> { + let writing = self.write || self.append; + + if !self.read && !writing { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "must enable read, write, or append access", + )); + } + if (self.truncate || self.create || self.create_new) && !writing { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "truncate/create options require write or append access", + )); + } + + Ok(()) + } + + /// Returns the initial cursor position for the file. + /// + /// If append mode is enabled, the cursor is set to the end of the file. + #[inline] + fn initial_cursor_for_append(&self, file: &std::fs::File) -> io::Result { + if self.append { + Ok(file.metadata()?.len()) + } else { + Ok(0) + } + } + + /// Opens a file with the configured options. + /// + /// This is the async version of [`std::fs::OpenOptions::open`]. + /// + /// # Platform-specific behavior + /// + /// - On Linux with io_uring support, this uses the `openat` syscall directly. + /// - On other platforms, this either offloads to a blocking thread pool or falls back + /// to [`std::fs::OpenOptions::open`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The file cannot be opened with the specified options + /// - The process lacks permissions to open the file + /// + /// # Examples + /// + /// ```ignore + /// use vibeio::fs::OpenOptions; + /// + /// let file = OpenOptions::new() + /// .read(true) + /// .open("hello.txt") + /// .await?; + /// ``` + #[inline] + pub async fn open(&self, path: impl AsRef) -> io::Result { + self.validate()?; + let path = path.as_ref(); + + let std_file = if let Some(driver) = current_driver() { + #[cfg(target_os = "linux")] + { + if driver.supports_completion() { + let mut op = self.build_open_op(path)?; + let raw = poll_fn(move |cx| op.poll(cx, &driver)).await?; + unsafe { std::fs::File::from_raw_fd(raw) } + } else if crate::offload_fs() { + self.open_in_blocking_pool(path).await? + } else { + self.open_blocking(path)? + } + } + + #[cfg(not(target_os = "linux"))] + { + let _ = driver; + if crate::offload_fs() { + self.open_in_blocking_pool(path).await? + } else { + self.open_blocking(path)? + } + } + } else { + self.open_blocking(path)? + }; + + let cursor = self.initial_cursor_for_append(&std_file)?; + File::from_std_with_cursor(std_file, cursor) + } + + /// Opens a file in the blocking thread pool. + /// + /// This is an internal method used when io_uring is not available. + #[inline] + async fn open_in_blocking_pool(&self, path: &Path) -> io::Result { + let path = path.to_path_buf(); + let read = self.read; + let write = self.write; + let append = self.append; + let truncate = self.truncate; + let create = self.create; + let create_new = self.create_new; + + crate::spawn_blocking(move || { + #[cfg(windows)] + use std::os::windows::fs::OpenOptionsExt; + #[cfg(windows)] + use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OVERLAPPED; + + let mut options = std::fs::OpenOptions::new(); + options + .read(read) + .write(write) + .append(append) + .truncate(truncate) + .create(create) + .create_new(create_new); + #[cfg(windows)] + options.attributes(FILE_FLAG_OVERLAPPED); + options.open(path) + }) + .await + .map_err(|_| crate::fs::file::blocking_pool_io_error())? + } + + /// Opens a file synchronously. + /// + /// This is an internal method used when io_uring is not available. + #[inline] + fn open_blocking(&self, path: &Path) -> io::Result { + #[cfg(windows)] + use std::os::windows::fs::OpenOptionsExt; + #[cfg(windows)] + use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OVERLAPPED; + + let mut options = std::fs::OpenOptions::new(); + options + .read(self.read) + .write(self.write) + .append(self.append) + .truncate(self.truncate) + .create(self.create) + .create_new(self.create_new); + #[cfg(windows)] + options.attributes(FILE_FLAG_OVERLAPPED); + options.open(path) + } + + /// Builds an `OpenOp` for use with io_uring. + /// + /// This is an internal method used on Linux with io_uring support. + #[cfg(target_os = "linux")] + #[inline] + fn build_open_op(&self, path: &Path) -> io::Result { + let writing = self.write || self.append; + let mut flags = match (self.read, writing) { + (true, false) => libc::O_RDONLY, + (false, true) => libc::O_WRONLY, + (true, true) => libc::O_RDWR, + (false, false) => { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "must enable read, write, or append access", + )) + } + }; + + if self.append { + flags |= libc::O_APPEND; + } + if self.truncate { + flags |= libc::O_TRUNC; + } + if self.create_new { + flags |= libc::O_CREAT | libc::O_EXCL; + } else if self.create { + flags |= libc::O_CREAT; + } + flags |= libc::O_CLOEXEC; + + let path_bytes = path.as_os_str().as_bytes(); + if path_bytes.contains(&0) { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "path contains interior NUL byte", + )); + } + + let path = CString::new(path_bytes).map_err(|_| { + io::Error::new(ErrorKind::InvalidInput, "path contains interior NUL byte") + })?; + + Ok(OpenOp::new(path, flags, 0o666)) + } +} + +impl Default for OpenOptions { + /// Returns the default `OpenOptions`. + /// + /// This is equivalent to `OpenOptions::new()`. + #[inline] + fn default() -> Self { + Self::new() + } +} diff --git a/vendor/vibeio/src/io/buf.rs b/vendor/vibeio/src/io/buf.rs new file mode 100644 index 0000000..3034a8c --- /dev/null +++ b/vendor/vibeio/src/io/buf.rs @@ -0,0 +1,458 @@ +//! Buffer traits for async I/O operations. +//! +//! This module provides traits for working with buffers in async I/O: +//! - `IoBuf` and `IoBufMut`: traits for read/write buffers. +//! - `IoVectoredBuf` and `IoVectoredBufMut`: traits for vectored I/O buffers. +//! +//! # Buffer types +//! +//! `IoBuf` is implemented for: +//! - `Vec` +//! - `String` +//! - `&'static [u8]` +//! - `&'static str` +//! - `[u8; N]` for any size `N` +//! - `Box<[u8]>` +//! +//! `IoBufMut` is implemented for: +//! - `Vec` +//! - `String` +//! - `[u8; N]` for any size `N` +//! - `Box<[u8]>` +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{AsyncRead, IoBufMut}; +//! +//! async fn read_something(reader: &mut R) { +//! let mut buf = vec![0u8; 1024]; +//! let (result, buf) = reader.read(buf).await; +//! let bytes_read = result.unwrap_or(0); +//! println!("Read {} bytes", bytes_read); +//! } +//! ``` + +use std::io::{IoSlice, IoSliceMut}; + +/// Trait for read-only buffers. +/// +/// This trait is implemented by types that can be used as buffers for +/// reading data in async I/O operations. +pub trait IoBuf: Send + 'static { + /// Returns a raw pointer to the inner buffer. + fn as_buf_ptr(&self) -> *const u8; + + /// Returns the length of the initialized part of the buffer. + fn buf_len(&self) -> usize; + + /// Returns the capacity of the buffer. + fn buf_capacity(&self) -> usize; +} + +/// Trait for mutable buffers. +/// +/// This trait extends `IoBuf` with mutable operations needed for writing. +pub trait IoBufMut: IoBuf { + /// Returns a raw mutable pointer to the inner buffer. + fn as_buf_mut_ptr(&mut self) -> *mut u8; + + /// Updates the length of the initialized part of the buffer. + /// + /// # Safety + /// + /// The caller must ensure that the given `len` does not exceed the capacity + /// of the buffer, and that the elements up to `len` have been initialized. + unsafe fn set_buf_init(&mut self, len: usize); +} + +impl IoBuf for Vec { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + self.len() + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.capacity() + } +} + +impl IoBufMut for Vec { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + self.as_mut_ptr() + } + + #[inline] + unsafe fn set_buf_init(&mut self, len: usize) { + self.set_len(len); + } +} + +impl IoBuf for String { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + self.len() + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.capacity() + } +} + +impl IoBufMut for String { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + unsafe { self.as_mut_vec().as_mut_ptr() } + } + + #[inline] + unsafe fn set_buf_init(&mut self, len: usize) { + self.as_mut_vec().set_len(len); + } +} + +impl IoBuf for &'static [u8] { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + self.len() + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.len() + } +} + +impl IoBuf for &'static str { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_bytes().as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + self.len() + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.len() + } +} + +impl IoBuf for [u8; N] { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + N + } + + #[inline] + fn buf_capacity(&self) -> usize { + N + } +} + +impl IoBufMut for [u8; N] { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + self.as_mut_ptr() + } + + unsafe fn set_buf_init(&mut self, _len: usize) {} +} + +impl IoBuf for Box<[u8]> { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.as_ptr() + } + + #[inline] + fn buf_len(&self) -> usize { + self.len() + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.len() + } +} + +impl IoBufMut for Box<[u8]> { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + self.as_mut_ptr() + } + + unsafe fn set_buf_init(&mut self, _len: usize) {} +} + +/// A buffer wrapper with a cursor for tracking progress. +pub(crate) struct IoBufWithCursor { + pub(crate) buf: I, + pub(crate) cursor: usize, +} + +impl IoBufWithCursor { + /// Create a new `IoBufWithCursor` with the given buffer. + #[inline] + pub(crate) fn new(buf: I) -> Self { + IoBufWithCursor { buf, cursor: 0 } + } + + /// Advance the cursor by `n` bytes. + #[inline] + pub(crate) fn advance(&mut self, n: usize) { + self.cursor += n; + } + + /// Consume the wrapper and return the inner buffer. + #[inline] + pub(crate) fn into_inner(self) -> I { + self.buf + } +} + +impl IoBuf for IoBufWithCursor { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + unsafe { self.buf.as_buf_ptr().add(self.cursor) } + } + + #[inline] + fn buf_len(&self) -> usize { + self.buf.buf_len() - self.cursor + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.buf.buf_capacity() - self.cursor + } +} + +impl IoBufMut for IoBufWithCursor { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + unsafe { self.buf.as_buf_mut_ptr().add(self.cursor) } + } + + unsafe fn set_buf_init(&mut self, len: usize) { + self.buf.set_buf_init(self.cursor + len); + } +} + +/// A temporary buffer for polling operations. +pub(crate) struct IoBufTemporaryPoll { + ptr: *mut u8, + len: usize, +} + +impl IoBufTemporaryPoll { + /// Create a new `IoBufTemporaryPoll` with the given pointer and length. + #[inline] + pub(crate) unsafe fn new(ptr: *mut u8, len: usize) -> Self { + Self { ptr, len } + } +} + +impl IoBuf for IoBufTemporaryPoll { + #[inline] + fn as_buf_ptr(&self) -> *const u8 { + self.ptr as *const u8 + } + + #[inline] + fn buf_len(&self) -> usize { + self.len + } + + #[inline] + fn buf_capacity(&self) -> usize { + self.len + } +} + +impl IoBufMut for IoBufTemporaryPoll { + #[inline] + fn as_buf_mut_ptr(&mut self) -> *mut u8 { + self.ptr + } + + #[inline] + unsafe fn set_buf_init(&mut self, _len: usize) {} +} + +unsafe impl Send for IoBufTemporaryPoll {} + +/// A single I/O vector entry. +pub struct IoVec { + /// Pointer to the data. + pub ptr: *mut u8, + /// Length of the data. + pub len: usize, +} + +/// Trait for vectored read buffers. +pub trait IoVectoredBuf: 'static { + /// Returns a pointer to an array of `iovec` structures and its length. + #[inline] + fn as_iovecs(&self) -> Box<[IoVec]> { + unimplemented!() + } + + /// Returns `true` if the vectored buffer is empty. + #[inline] + fn is_empty(&self) -> bool { + self.as_iovecs().is_empty() + } +} + +/// Trait for vectored write buffers. +pub trait IoVectoredBufMut: IoVectoredBuf { + /// Returns a mutable pointer to an array of `iovec` structures and its length. + #[inline] + fn as_iovecs_mut(&mut self) -> Box<[IoVec]> { + unimplemented!() + } +} + +#[cfg(unix)] +impl IoVectoredBuf for Vec { + #[inline] + fn as_iovecs(&self) -> Box<[IoVec]> { + let mut iovecs = Box::new_uninit_slice(self.len()); + for (index, iovec) in self.iter().enumerate() { + iovecs[index].write(IoVec { + ptr: iovec.iov_base as *mut u8, + len: iovec.iov_len, + }); + } + + unsafe { iovecs.assume_init() } + } + + #[inline] + fn is_empty(&self) -> bool { + self.is_empty() + } +} + +#[cfg(unix)] +impl IoVectoredBufMut for Vec { + #[inline] + fn as_iovecs_mut(&mut self) -> Box<[IoVec]> { + self.as_iovecs() + } +} + +#[cfg(unix)] +impl IoVectoredBuf for Box<[libc::iovec]> { + #[inline] + fn as_iovecs(&self) -> Box<[IoVec]> { + let mut iovecs = Box::new_uninit_slice(self.len()); + for (index, iovec) in self.iter().enumerate() { + iovecs[index].write(IoVec { + ptr: iovec.iov_base as *mut u8, + len: iovec.iov_len, + }); + } + + unsafe { iovecs.assume_init() } + } +} + +#[cfg(unix)] +impl IoVectoredBufMut for Box<[libc::iovec]> { + #[inline] + fn as_iovecs_mut(&mut self) -> Box<[IoVec]> { + self.as_iovecs() + } +} + +/// A temporary vectored buffer for polling operations. +pub(crate) struct IoVectoredBufTemporaryPoll { + pub(crate) iovecs: Vec<(*mut u8, usize)>, +} + +impl IoVectoredBufTemporaryPoll { + /// Create a new `IoVectoredBufTemporaryPoll` from immutable slices. + #[inline] + pub(crate) unsafe fn new(iovecs: &[IoSlice<'_>]) -> Self { + let iovecs = iovecs + .iter() + .map(|iovec| (iovec.as_ptr() as *mut u8, iovec.len())) + .collect(); + Self { iovecs } + } + + /// Create a new `IoVectoredBufTemporaryPoll` from mutable slices. + #[allow(dead_code)] + #[inline] + pub(crate) unsafe fn new_mut(iovecs: &mut [IoSliceMut<'_>]) -> Self { + let iovecs = iovecs + .iter_mut() + .map(|iovec| (iovec.as_mut_ptr(), iovec.len())) + .collect(); + Self { iovecs } + } +} + +impl IoVectoredBuf for IoVectoredBufTemporaryPoll { + #[inline] + fn as_iovecs(&self) -> Box<[IoVec]> { + let mut iovecs = Box::new_uninit_slice(self.iovecs.len()); + for (index, iovec) in self.iovecs.iter().enumerate() { + iovecs[index].write(IoVec { + ptr: iovec.0, + len: iovec.1, + }); + } + + unsafe { iovecs.assume_init() } + } + + #[inline] + fn is_empty(&self) -> bool { + self.iovecs.is_empty() + } +} + +impl IoVectoredBufMut for IoVectoredBufTemporaryPoll { + #[inline] + fn as_iovecs_mut(&mut self) -> Box<[IoVec]> { + self.as_iovecs() + } +} + +#[inline] +pub(crate) fn iobuf_to_slice(buf: &impl IoBuf) -> &[u8] { + unsafe { std::slice::from_raw_parts(buf.as_buf_ptr(), buf.buf_len()) } +} + +#[inline] +pub(crate) fn iobufmut_to_slice(buf: &mut impl IoBufMut) -> &mut [u8] { + unsafe { std::slice::from_raw_parts_mut(buf.as_buf_mut_ptr(), buf.buf_len()) } +} diff --git a/vendor/vibeio/src/io/mod.rs b/vendor/vibeio/src/io/mod.rs new file mode 100644 index 0000000..862af50 --- /dev/null +++ b/vendor/vibeio/src/io/mod.rs @@ -0,0 +1,150 @@ +//! Async I/O traits and utilities. +//! +//! This module provides the core async I/O abstractions used by vibeio: +//! - `AsyncRead` and `AsyncWrite`: async traits for reading and writing. +//! - `AsyncReadPoll` and `AsyncWritePoll`: poll-based async read/write readiness interfaces. +//! - `IoBuf` and `IoBufMut`: buffer traits for async I/O operations. +//! - `pipe()`: create async-aware pipe endpoints. +//! - `splice()` and `sendfile_exact()`: zero-copy I/O operations (Linux only). +//! - `copy()`: copy data between async I/O objects. +//! - `split()`: split I/O objects into independent read/write halves. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{AsyncRead, AsyncWrite}; +//! +//! async fn echo(reader: &mut R, writer: &mut W) { +//! let mut buf = vec![0u8; 1024]; +//! loop { +//! let (read, buf) = reader.read(buf).await; +//! let read = read?; +//! if read == 0 { +//! break; +//! } +//! let (written, buf) = writer.write(buf).await; +//! written?; +//! } +//! } +//! ``` +//! +//! # Implementation notes +//! - The `AsyncRead` and `AsyncWrite` traits are similar to tokio's but return +//! `(Result, Buffer)` to support buffer reuse. +//! - The `IoBuf` trait is implemented for common buffer types like `Vec`, +//! `String`, byte arrays, and `Box<[u8]>`. +//! - On Linux with the `splice` feature, `splice()` and `sendfile_exact()` +//! provide zero-copy I/O. + +#![allow(async_fn_in_trait)] + +mod buf; +#[cfg(all(unix, feature = "pipe"))] +mod pipe; +#[cfg(all(target_os = "linux", feature = "splice"))] +mod splice; +#[cfg(feature = "stdio")] +mod stdio; +mod util; + +use crate::fd_inner::InnerRawHandle; + +pub use self::buf::*; +#[cfg(all(unix, feature = "pipe"))] +pub use self::pipe::*; +#[cfg(all(target_os = "linux", feature = "splice"))] +pub use self::splice::*; +#[cfg(feature = "stdio")] +pub use self::stdio::*; +pub use self::util::*; + +use std::io::{self, ErrorKind}; + +/// Async trait for reading data. +/// +/// This trait is similar to `tokio::io::AsyncRead` but returns a tuple +/// `(Result, Buffer)` to support buffer reuse. +pub trait AsyncRead { + /// Read data into the buffer, returning the number of bytes read and the buffer. + async fn read(&mut self, buf: B) -> (Result, B); + + /// Read data into vectored buffers. + #[inline] + async fn read_vectored( + &mut self, + bufs: V, + ) -> (Result, V) { + ( + Err(io::Error::new( + ErrorKind::Unsupported, + "readv not implemented", + )), + bufs, + ) + } +} + +/// Async trait for writing data. +/// +/// This trait is similar to `tokio::io::AsyncWrite` but returns a tuple +/// `(Result, Buffer)` to support buffer reuse. +pub trait AsyncWrite { + /// Write data from the buffer, returning the number of bytes written and the buffer. + async fn write(&mut self, buf: B) -> (Result, B); + + /// Write data from vectored buffers. + #[inline] + async fn write_vectored(&mut self, bufs: V) -> (Result, V) { + ( + Err(io::Error::new( + ErrorKind::Unsupported, + "writev not implemented", + )), + bufs, + ) + } + + /// Flush the writer. + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } +} + +/// Trait to get an inner raw handle reference. +pub trait AsInnerRawHandle<'a> { + /// Returns a reference to the inner raw handle. + #[allow(private_interfaces)] + fn as_inner_raw_handle(&'a self) -> &'a crate::fd_inner::InnerRawHandle; +} + +impl<'a> AsInnerRawHandle<'a> for InnerRawHandle { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + self + } +} + +/// Trait to poll for read readiness. +pub trait AsyncReadPoll { + /// Polls to check if the reader is readable. + fn poll_readable(&self, cx: &mut std::task::Context) -> std::task::Poll>; + + /// Returns a future that resolves when the reader becomes readable. + #[inline] + async fn readable(&self) -> Result<(), io::Error> { + std::future::poll_fn(|cx| self.poll_readable(cx)).await + } +} + +/// Trait to poll for write readiness. +pub trait AsyncWritePoll { + /// Polls to check if the writer is writable. + fn poll_writable(&self, cx: &mut std::task::Context) -> std::task::Poll>; + + /// Returns a future that resolves when the writer becomes writable. + #[inline] + async fn writable(&self) -> Result<(), io::Error> { + std::future::poll_fn(|cx| self.poll_writable(cx)).await + } +} diff --git a/vendor/vibeio/src/io/pipe.rs b/vendor/vibeio/src/io/pipe.rs new file mode 100644 index 0000000..c64df11 --- /dev/null +++ b/vendor/vibeio/src/io/pipe.rs @@ -0,0 +1,324 @@ +//! Async pipe utilities. +//! +//! This module provides async-aware pipe endpoints: +//! - `pipe()`: create a pair of async-aware pipe endpoints. +//! - `Pipe`: a pipe endpoint for async I/O. +//! - `PollPipe`: a variant that uses readiness-based polling. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{AsyncRead, AsyncWrite, pipe}; +//! +//! async fn pipe_example() { +//! let (mut reader, mut writer) = pipe().unwrap(); +//! +//! writer.write(b"hello").await.0.unwrap(); +//! let mut buf = [0u8; 5]; +//! reader.read(buf).await.0.unwrap(); +//! assert_eq!(&buf, b"hello"); +//! } +//! ``` + +use std::future::poll_fn; +use std::io::{self, IoSlice}; +use std::mem::{ManuallyDrop, MaybeUninit}; +use std::os::fd::OwnedFd; +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use mio::Interest; +use tokio::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite, ReadBuf}; + +use crate::io::{ + AsInnerRawHandle, IoBuf, IoBufMut, IoBufTemporaryPoll, IoVectoredBuf, IoVectoredBufMut, + IoVectoredBufTemporaryPoll, +}; +use crate::op::{ReadOp, ReadvOp, WriteOp, WritevOp}; +use crate::{ + driver::RegistrationMode, + fd_inner::InnerRawHandle, + io::{AsyncRead, AsyncWrite}, +}; + +fn pipe_inner() -> std::io::Result<(OwnedFd, OwnedFd)> { + let mut fds: [RawFd; 2] = [-1; 2]; + let result = unsafe { libc::pipe(fds.as_mut_ptr()) }; + if result == -1 { + return Err(std::io::Error::last_os_error()); + } + Ok((unsafe { OwnedFd::from_raw_fd(fds[0]) }, unsafe { + OwnedFd::from_raw_fd(fds[1]) + })) +} + +/// Create a new async-aware pipe. +/// +/// Returns a tuple of `(reader, writer)` pipe endpoints. +pub fn pipe() -> std::io::Result<(Pipe, Pipe)> { + let (read, write) = pipe_inner()?; + Ok(( + Pipe::from_std_with_mode(read, RegistrationMode::Completion)?, + Pipe::from_std_with_mode(write, RegistrationMode::Completion)?, + )) +} + +/// A pipe endpoint that can use either completion or readiness-based I/O. +pub struct Pipe { + inner: OwnedFd, + handle: ManuallyDrop, +} + +/// A poll-only variant that always uses readiness-based operations. +pub struct PollPipe { + stream: Pipe, +} + +impl Pipe { + /// Create a `Pipe` from a standard library `OwnedFd` with the given registration mode. + #[inline] + pub(crate) fn from_std_with_mode( + inner: OwnedFd, + mode: RegistrationMode, + ) -> Result { + #[cfg(unix)] + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + inner.as_raw_fd(), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + let flags = unsafe { libc::fcntl(inner.as_raw_fd(), libc::F_GETFL) }; + if flags != -1 { + let mut new_flags = flags | libc::O_NONBLOCK; + if handle.uses_completion() { + new_flags &= !libc::O_NONBLOCK; + } + unsafe { libc::fcntl(inner.as_raw_fd(), libc::F_SETFL, new_flags) }; + } + Ok(Self { inner, handle }) + } + + /// Convert this `Pipe` to a `PollPipe` for readiness-based operations. + #[inline] + pub fn into_poll(self) -> Result { + let mut stream = self; + stream.handle.rebind_mode(RegistrationMode::Poll)?; + let flags = unsafe { libc::fcntl(stream.inner.as_raw_fd(), libc::F_GETFL) }; + if flags != -1 { + let mut new_flags = flags | libc::O_NONBLOCK; + if stream.handle.uses_completion() { + new_flags &= !libc::O_NONBLOCK; + } + unsafe { libc::fcntl(stream.inner.as_raw_fd(), libc::F_SETFL, new_flags) }; + } + Ok(PollPipe { stream }) + } +} + +impl PollPipe { + /// Convert this `PollPipe` back to an adaptive `Pipe`. + #[inline] + pub fn into_adaptive(self) -> Pipe { + self.stream + } + + /// Convert this `PollPipe` to a completion-based `Pipe`. + #[inline] + pub fn into_completion(self) -> Result { + let mut stream = self.stream; + stream.handle.rebind_mode(RegistrationMode::Completion)?; + let flags = unsafe { libc::fcntl(stream.inner.as_raw_fd(), libc::F_GETFL) }; + if flags != -1 { + let mut new_flags = flags | libc::O_NONBLOCK; + if stream.handle.uses_completion() { + new_flags &= !libc::O_NONBLOCK; + } + unsafe { libc::fcntl(stream.inner.as_raw_fd(), libc::F_SETFL, new_flags) }; + } + Ok(stream) + } +} + +impl AsRawFd for Pipe { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +impl AsRawFd for PollPipe { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.stream.inner.as_raw_fd() + } +} + +impl IntoRawFd for Pipe { + #[inline] + fn into_raw_fd(self) -> RawFd { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its fd ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_fd() + } + } +} + +impl IntoRawFd for PollPipe { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.stream.into_raw_fd() + } +} + +impl<'a> AsInnerRawHandle<'a> for Pipe { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + &self.handle + } +} + +impl<'a> AsInnerRawHandle<'a> for PollPipe { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + self.stream.as_inner_raw_handle() + } +} + +impl AsyncRead for Pipe { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = ReadOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn read_vectored( + &mut self, + bufs: B, + ) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = ReadvOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl TokioAsyncRead for PollPipe { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let this = self.get_mut(); + // Equivalent to .assume_init_mut() in Rust 1.93.0+ + let unfilled = unsafe { &mut *(buf.unfilled_mut() as *mut [MaybeUninit] as *mut [u8]) }; + let buf_temp = unsafe { IoBufTemporaryPoll::new(unfilled.as_mut_ptr(), unfilled.len()) }; + let mut op = ReadOp::new(&this.stream.handle, buf_temp); + match this.stream.handle.poll_op_poll(cx, &mut op) { + Poll::Ready(Ok(read)) => { + unsafe { + buf.assume_init(read); + } + buf.advance(read); + Poll::Ready(Ok(())) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncWrite for Pipe { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = WriteOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } + + #[inline] + async fn write_vectored(&mut self, bufs: B) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = WritevOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl TokioAsyncWrite for PollPipe { + #[inline] + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let buf = unsafe { IoBufTemporaryPoll::new(buf.as_ptr() as *mut u8, buf.len()) }; + let mut op = WriteOp::new(&this.stream.handle, buf); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + if bufs.is_empty() { + return Poll::Ready(Ok(0)); + } + let this = self.get_mut(); + let bufs = unsafe { IoVectoredBufTemporaryPoll::new(bufs) }; + let mut op = WritevOp::new(&this.stream.handle, bufs); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn is_write_vectored(&self) -> bool { + true + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + #[inline] + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl Drop for Pipe { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/io/splice.rs b/vendor/vibeio/src/io/splice.rs new file mode 100644 index 0000000..c012f7b --- /dev/null +++ b/vendor/vibeio/src/io/splice.rs @@ -0,0 +1,164 @@ +//! Zero-copy I/O utilities using `splice` and `sendfile`. +//! +//! This module provides async-aware zero-copy I/O operations: +//! - `splice()`: transfer data between file descriptors without copying to userspace. +//! - `splice_exact()`: transfer exactly `len` bytes using `splice`. +//! - `sendfile_exact()`: transfer data from a file to a socket using a pipe as an intermediary. +//! +//! These operations are only available on Linux with the `splice` feature enabled. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{splice, AsyncRead, AsyncWrite, pipe}; +//! +//! async fn zero_copy_example() { +//! let (reader, writer) = pipe().unwrap(); +//! let mut file = std::fs::File::open("data.txt").unwrap(); +//! +//! // Transfer 1024 bytes from file to pipe +//! let n = splice(&file, &writer, 1024).await.unwrap(); +//! println!("Spliced {} bytes", n); +//! } +//! ``` + +use std::{ + mem::ManuallyDrop, + os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd}, +}; + +use mio::Interest; + +use crate::{fd_inner::InnerRawHandle, io::AsInnerRawHandle, op::SpliceOp}; + +/// Transfer data from one file descriptor to another using `splice`. +/// +/// This function uses the kernel's `splice` system call to transfer data +/// between file descriptors without copying to userspace. +pub async fn splice<'a, 'b>( + from: &'a impl AsRawFd, + to: &'b impl AsInnerRawHandle<'b>, + len: usize, +) -> Result { + let from_handle = unsafe { BorrowedFd::borrow_raw(from.as_raw_fd()) }; + let to_handle = to.as_inner_raw_handle(); + + let mut op = SpliceOp::new(from_handle, to_handle, len); + let result = std::future::poll_fn(move |cx| to_handle.poll_op(cx, &mut op)).await; + result +} + +/// Transfer exactly `len` bytes from one file descriptor to another using `splice`. +/// +/// This function calls `splice()` repeatedly until `len` bytes have been transferred +/// or EOF is reached. +pub async fn splice_exact<'a, 'b>( + from: &'a impl AsRawFd, + to: &'b impl AsInnerRawHandle<'b>, + len: u64, +) -> Result { + let mut total = 0; + while total < len { + let n = splice(from, to, (len - total).min(usize::MAX as u64) as usize).await?; + if n == 0 { + break; + } + total += n as u64; + } + + Ok(total) +} + +/// Transfer data from a file to a socket using `sendfile` semantics. +/// +/// This function implements `sendfile`-like behavior using `splice` with an +/// intermediate pipe, allowing data to be transferred from a regular file to +/// a socket without copying to userspace. +pub async fn sendfile_exact<'a, 'b>( + from: &'a impl AsRawFd, + to: &'b impl AsInnerRawHandle<'b>, + len: u64, +) -> Result { + // splice() requires at least one of the file descriptors to be a pipe. + // Therefore, we need to create a pipe and use it as the destination. + let mut fds: [RawFd; 2] = [0; 2]; + if unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC) } != 0 { + // pipe2 failed, can't continue + return Err(std::io::Error::last_os_error()); + } + let pipe_reader = unsafe { OwnedFd::from_raw_fd(fds[0]) }; + let pipe_writer = unsafe { OwnedFd::from_raw_fd(fds[1]) }; + + // We only need to poll the pipe writer for writability. + let pipe_writer_handle = WriteOwnedFd::new(pipe_writer)?; + + let mut total_from_file = 0; + let mut total_to_socket = 0; + let mut file_eof = false; + while (total_from_file < len && !file_eof) || total_to_socket < total_from_file { + let splice_from_file_len = (len - total_from_file).min(usize::MAX as u64) as usize; + if !file_eof && splice_from_file_len > 0 { + let n = splice(from, &pipe_writer_handle, splice_from_file_len).await?; + if n == 0 { + file_eof = true; + } else { + total_from_file += n as u64; + } + } + + let splice_to_socket_len = + (total_from_file - total_to_socket).min(usize::MAX as u64) as usize; + if splice_to_socket_len > 0 { + let n = splice(&pipe_reader, to, splice_to_socket_len).await?; + if n == 0 { + break; + } + total_to_socket += n as u64; + } + } + + drop(pipe_reader); + drop(pipe_writer_handle); + + Ok(total_to_socket) +} + +struct WriteOwnedFd { + _writer: OwnedFd, + handle: ManuallyDrop, +} + +impl WriteOwnedFd { + fn new(writer: OwnedFd) -> std::io::Result { + let handle = + ManuallyDrop::new(InnerRawHandle::new(writer.as_raw_fd(), Interest::WRITABLE)?); + if !handle.uses_completion() { + // Set the pipe write side to non-blocking mode. + let flags = unsafe { libc::fcntl(writer.as_raw_fd(), libc::F_GETFL) }; + if flags != -1 { + unsafe { libc::fcntl(writer.as_raw_fd(), libc::F_SETFL, flags | libc::O_NONBLOCK) }; + } + } + Ok(Self { + _writer: writer, + handle, + }) + } +} + +impl<'a> AsInnerRawHandle<'a> for WriteOwnedFd { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + &self.handle + } +} + +impl Drop for WriteOwnedFd { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/io/stdio.rs b/vendor/vibeio/src/io/stdio.rs new file mode 100644 index 0000000..4062711 --- /dev/null +++ b/vendor/vibeio/src/io/stdio.rs @@ -0,0 +1,266 @@ +//! Standard I/O utilities for stdin, stdout, and stderr. +//! +//! This module provides async-aware wrappers around standard I/O streams: +//! - `Stdin`: async stdin reader. +//! - `Stdout`: async stdout writer. +//! - `Stderr`: async stderr writer. +//! +//! These types use a blocking thread pool for I/O operations when inside +//! a runtime context. Outside a runtime, they fall back to synchronous I/O. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{AsyncRead, AsyncWrite, stdin, stdout}; +//! +//! async fn echo() { +//! let mut stdin = stdin(); +//! let mut stdout = stdout(); +//! let mut buf = vec![0u8; 1024]; +//! loop { +//! let (read, buf) = stdin.read(buf).await; +//! let read = read?; +//! if read == 0 { +//! break; +//! } +//! let (written, buf) = stdout.write(buf).await; +//! written?; +//! } +//! } +//! ``` + +use std::cell::RefCell; +use std::io::{self, Read, Write}; +use std::sync::{Arc, Mutex}; + +use crate::executor::current_driver; +use crate::io::{iobuf_to_slice, iobufmut_to_slice, AsyncRead, AsyncWrite, IoBuf, IoBufMut}; + +/// Async-aware stdin reader. +#[derive(Debug, Default, Clone, Copy)] +pub struct Stdin { + _private: (), +} + +/// Async-aware stdout writer. +#[derive(Debug, Default, Clone, Copy)] +pub struct Stdout { + _private: (), +} + +/// Async-aware stderr writer. +#[derive(Debug, Default, Clone, Copy)] +pub struct Stderr { + _private: (), +} + +/// Get an async-aware stdin reader. +#[inline] +pub fn stdin() -> Stdin { + Stdin { _private: () } +} + +/// Get an async-aware stdout writer. +#[inline] +pub fn stdout() -> Stdout { + Stdout { _private: () } +} + +/// Get an async-aware stderr writer. +#[inline] +pub fn stderr() -> Stderr { + Stderr { _private: () } +} + +#[inline] +fn read_stdin_blocking(buf: &mut [u8]) -> io::Result { + let mut stdin = std::io::stdin(); + stdin.read(buf) +} + +#[inline] +fn write_stdout_blocking(buf: &[u8]) -> io::Result { + let mut stdout = std::io::stdout(); + stdout.write(buf) +} + +#[inline] +fn write_stderr_blocking(buf: &[u8]) -> io::Result { + let mut stderr = std::io::stderr(); + stderr.write(buf) +} + +#[inline] +fn flush_stdout_blocking() -> io::Result<()> { + let mut stdout = std::io::stdout(); + stdout.flush() +} + +#[inline] +fn flush_stderr_blocking() -> io::Result<()> { + let mut stderr = std::io::stderr(); + stderr.flush() +} + +#[inline] +fn blocking_pool_io_error() -> io::Error { + io::Error::other("can't spawn blocking task for stdio I/O") +} + +#[inline] +async fn read_in_blocking_pool(buf: B) -> (io::Result, B) { + let buf = Arc::new(Mutex::new(RefCell::new(Some(buf)))); + let buf_clone = buf.clone(); + crate::spawn_blocking(move || { + let mut buf = buf_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"); + let temp_slice = iobufmut_to_slice(&mut buf); + let result = read_stdin_blocking(temp_slice); + (result, buf) + }) + .await + .unwrap_or_else(|_| { + ( + Err(blocking_pool_io_error()), + buf.try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"), + ) + }) +} + +#[inline] +async fn write_stdout_in_blocking_pool(buf: B) -> (io::Result, B) { + let buf = Arc::new(Mutex::new(RefCell::new(Some(buf)))); + let buf_clone = buf.clone(); + crate::spawn_blocking(move || { + let buf = buf_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"); + let temp_slice = iobuf_to_slice(&buf); + let result = write_stdout_blocking(temp_slice); + (result, buf) + }) + .await + .unwrap_or_else(|_| { + ( + Err(blocking_pool_io_error()), + buf.try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"), + ) + }) +} + +#[inline] +async fn write_stderr_in_blocking_pool(buf: B) -> (io::Result, B) { + let buf = Arc::new(Mutex::new(RefCell::new(Some(buf)))); + let buf_clone = buf.clone(); + crate::spawn_blocking(move || { + let buf = buf_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"); + let temp_slice = iobuf_to_slice(&buf); + let result = write_stderr_blocking(temp_slice); + (result, buf) + }) + .await + .unwrap_or_else(|_| { + ( + Err(blocking_pool_io_error()), + buf.try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("buf is none"), + ) + }) +} + +#[inline] +async fn flush_stdout_in_blocking_pool() -> io::Result<()> { + crate::spawn_blocking(flush_stdout_blocking) + .await + .map_err(|_| blocking_pool_io_error())? +} + +#[inline] +async fn flush_stderr_in_blocking_pool() -> io::Result<()> { + crate::spawn_blocking(flush_stderr_blocking) + .await + .map_err(|_| blocking_pool_io_error())? +} + +impl AsyncRead for Stdin { + #[inline] + async fn read(&mut self, mut buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + if current_driver().is_some() { + read_in_blocking_pool(buf).await + } else { + let slice = iobufmut_to_slice(&mut buf); + (read_stdin_blocking(slice), buf) + } + } +} + +impl AsyncWrite for Stdout { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + if current_driver().is_some() { + write_stdout_in_blocking_pool(buf).await + } else { + let slice = iobuf_to_slice(&buf); + (write_stdout_blocking(slice), buf) + } + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + if current_driver().is_some() { + flush_stdout_in_blocking_pool().await + } else { + flush_stdout_blocking() + } + } +} + +impl AsyncWrite for Stderr { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + if current_driver().is_some() { + write_stderr_in_blocking_pool(buf).await + } else { + let slice = iobuf_to_slice(&buf); + (write_stderr_blocking(slice), buf) + } + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + if current_driver().is_some() { + flush_stderr_in_blocking_pool().await + } else { + flush_stderr_blocking() + } + } +} diff --git a/vendor/vibeio/src/io/util.rs b/vendor/vibeio/src/io/util.rs new file mode 100644 index 0000000..729965a --- /dev/null +++ b/vendor/vibeio/src/io/util.rs @@ -0,0 +1,226 @@ +//! Async I/O utilities. +//! +//! This module provides utility functions for async I/O operations: +//! - `copy()`: copy data from a reader to a writer. +//! - `split()`: split an I/O object into independent read/write halves. +//! - `copy_bidirectional()`: copy data in both directions between two I/O objects. +//! +//! # Examples +//! +//! ```ignore +//! use vibeio::io::{AsyncRead, AsyncWrite, copy}; +//! +//! async fn copy_example(reader: &mut R, writer: &mut W) { +//! let bytes_copied = copy(reader, writer).await?; +//! println!("Copied {} bytes", bytes_copied); +//! } +//! ``` + +use std::io; +use std::sync::Arc; + +use futures_util::lock::Mutex as AsyncMutex; + +use super::{AsyncRead, AsyncWrite}; +use crate::io::{IoBuf, IoBufMut, IoBufWithCursor}; + +/// Copy data from a reader to a writer. +/// +/// This function reads from `reader` and writes to `writer` until EOF is reached. +/// Returns the number of bytes copied. +pub async fn copy(reader: &mut R, writer: &mut W) -> Result +where + R: AsyncRead + ?Sized, + W: AsyncWrite + ?Sized, +{ + let mut buffer = vec![0u8; 8192]; + buffer.clear(); + let mut copied = 0u64; + + loop { + let (read, returned_buf) = reader.read(buffer).await; + let read = read?; + + if read == 0 { + break; + } + + let mut cursor_buf = IoBufWithCursor::new(returned_buf); + while cursor_buf.buf_len() > 0 { + let (w, mut returned_buf) = writer.write(cursor_buf).await; + let w = w?; + if w == 0 { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write whole buffer", + )); + } + returned_buf.advance(w); + cursor_buf = returned_buf; + } + + buffer = cursor_buf.into_inner(); + unsafe { + buffer.set_buf_init(0); + } // reset + copied = copied.saturating_add(read as u64); + } + + writer.flush().await?; + Ok(copied) +} + +/// Owned read half for a split I/O object. +/// +/// The halves share ownership of the inner object via an `Arc>`. +pub struct ReadHalf { + inner: Arc>, +} + +/// Owned write half for a split I/O object. +pub struct WriteHalf { + inner: Arc>, +} + +/// Split an object implementing both `AsyncRead` and `AsyncWrite` into two +/// independently usable halves. +/// +/// The halves share ownership of the original object via an `Arc>` +/// so they may be used concurrently in async contexts. +/// +/// Note: this is a simple, owned split helper — it clones an `Arc` around +/// a mutex protecting the whole I/O object. It does not provide lock-free +/// simultaneous read/write on the underlying object; callers still need to +/// tolerate possible contention on the mutex. +pub fn split(io: T) -> (ReadHalf, WriteHalf) +where + T: AsyncRead + AsyncWrite + 'static, +{ + let inner = Arc::new(AsyncMutex::new(io)); + ( + ReadHalf { + inner: inner.clone(), + }, + WriteHalf { inner }, + ) +} + +impl ReadHalf +where + T: AsyncRead + AsyncWrite + 'static, +{ + /// Consume the half and return the shared inner `Arc>`. + pub fn into_inner(self) -> Arc> { + self.inner + } +} + +impl WriteHalf +where + T: AsyncRead + AsyncWrite + 'static, +{ + /// Consume the half and return the shared inner `Arc>`. + pub fn into_inner(self) -> Arc> { + self.inner + } +} + +impl AsyncRead for ReadHalf +where + T: AsyncRead + AsyncWrite + 'static, +{ + async fn read(&mut self, buf: B) -> (Result, B) { + let mut guard = self.inner.lock().await; + // Forward the call to the underlying object. + (*guard).read(buf).await + } +} + +impl AsyncWrite for WriteHalf +where + T: AsyncRead + AsyncWrite + 'static, +{ + async fn write(&mut self, buf: B) -> (Result, B) { + let mut guard = self.inner.lock().await; + (*guard).write(buf).await + } + + async fn flush(&mut self) -> Result<(), io::Error> { + let mut guard = self.inner.lock().await; + (*guard).flush().await + } +} + +impl AsyncRead for Box { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + (**self).read(buf).await + } +} + +impl AsyncRead for &mut R { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + (**self).read(buf).await + } +} + +impl AsyncWrite for Box { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + (**self).write(buf).await + } + + #[inline] + async fn flush(&mut self) -> Result<(), std::io::Error> { + (**self).flush().await + } +} + +impl AsyncWrite for &mut W { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + (**self).write(buf).await + } + + #[inline] + async fn flush(&mut self) -> Result<(), std::io::Error> { + (**self).flush().await + } +} + +/// Copy data in both directions between two I/O objects that implement both +/// `AsyncRead` and `AsyncWrite`. +/// +/// This function takes ownership of both objects, splits them into read/write +/// halves (so the copies may proceed concurrently), and runs two `copy` +/// operations in parallel: +/// - bytes read from `a` are written to `b` +/// - bytes read from `b` are written to `a` +/// +/// Returns a tuple `(a_to_b, b_to_a)` with the number of bytes copied in each +/// direction. The function returns an error if either direction returns an +/// error. +pub async fn copy_bidirectional(a: A, b: B) -> Result<(u64, u64), io::Error> +where + A: AsyncRead + AsyncWrite + 'static, + B: AsyncRead + AsyncWrite + 'static, +{ + // Split both objects into independent read/write halves. + let (mut a_r, mut a_w) = split(a); + let (mut b_r, mut b_w) = split(b); + + // Create the two copy futures. They borrow disjoint halves, so creating + // both futures is allowed. + let f1 = copy(&mut a_r, &mut b_w); + let f2 = copy(&mut b_r, &mut a_w); + + // Run both copies concurrently and await their results. + let (res1, res2) = futures_util::future::join(f1, f2).await; + + // Propagate any errors; on success return the number of bytes copied for + // each direction. + let n1 = res1?; + let n2 = res2?; + Ok((n1, n2)) +} diff --git a/vendor/vibeio/src/lib.rs b/vendor/vibeio/src/lib.rs new file mode 100644 index 0000000..13d1050 --- /dev/null +++ b/vendor/vibeio/src/lib.rs @@ -0,0 +1,103 @@ +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! # vibeio +//! +//! A high-performance, cross-platform asynchronous runtime for Rust. +//! +//! `vibeio` provides an efficient I/O event loop that leverages the best available driver for each operating system: +//! +//! - **Linux** - uses `io_uring` for true asynchronous I/O. +//! - **Windows** - uses I/O Completion Ports (IOCP) for scalable I/O. +//! - **macOS / BSD / Others** - uses `kqueue` or `epoll` via `mio` for event notification. +//! +//! ## Core features +//! +//! - **Networking** - asynchronous TCP, UDP, and Unix Domain Sockets. +//! - **File system** - asynchronous file operations. +//! - **Timers** - efficient timer and sleep functionality. +//! - **Signals** - handling of OS signals. +//! - **Process management** - spawning and managing child processes. +//! - **Blocking tasks** - offload CPU-intensive or blocking operations to a thread pool. +//! +//! ## Concurrency model: thread-per-core +//! +//! `vibeio` is designed as a **single-threaded** runtime. To utilize multiple cores, you should employ a **thread-per-core** architecture, where a separate `Runtime` is pinned to each processor core. This approach minimizes synchronization overhead and maximizes cache locality. +//! +//! Shared state can be communicated between runtimes using message passing (e.g., channels) or shared atomic structures, but I/O resources are typically owned by the thread that created them. +//! +//! ## Getting started +//! +//! Add `vibeio` to your `Cargo.toml`: +//! +//! ```toml +//! [dependencies] +//! vibeio = "0.2" +//! ``` +//! +//! ### Example: TCP echo server +//! +//! ```rust,no_run +//! use vibeio::RuntimeBuilder; +//! use vibeio::net::TcpListener; +//! +//! fn main() -> std::io::Result<()> { +//! // 1. Build the runtime +//! let runtime = RuntimeBuilder::new() +//! .enable_timer(true) +//! .build()?; +//! +//! // 2. Run the main future +//! runtime.block_on(async { +//! let listener = TcpListener::bind("127.0.0.1:8080")?; +//! println!("Listening on 127.0.0.1:8080"); +//! +//! loop { +//! let (mut stream, _) = listener.accept().await?; +//! +//! vibeio::spawn(async move { +//! let (mut reader, mut writer) = vibeio::io::split(stream); +//! if let Err(e) = vibeio::io::copy(&mut reader, &mut writer).await { +//! eprintln!("Echo failed: {}", e); +//! } +//! }); +//! } +//! }) +//! } +//! ``` +//! +//! ## Feature flags +//! +//! The following features are available (most are enabled by default): +//! +//! - `fs` - enables asynchronous file system operations. +//! - `time` - enables time and timer functionality. +//! - `signal` - enables signal handling. +//! - `process` - enables child process management. +//! - `pipe` - enables pipe support. +//! - `stdio` - enables standard I/O support. +//! - `splice` - enables splice support (Linux). +//! - `blocking-default` - enables the default blocking thread pool. + +pub mod blocking; +mod builder; +mod driver; +mod executor; +mod fd_inner; +#[cfg(feature = "fs")] +pub mod fs; +pub mod io; +pub mod net; +mod op; +#[cfg(feature = "process")] +pub mod process; +#[cfg(feature = "signal")] +pub mod signal; +mod task; +#[cfg(feature = "time")] +pub mod time; +#[cfg(feature = "time")] +mod timer; +pub mod util; + +pub use crate::builder::*; +pub use crate::executor::*; diff --git a/vendor/vibeio/src/net/mod.rs b/vendor/vibeio/src/net/mod.rs new file mode 100644 index 0000000..1e4fba6 --- /dev/null +++ b/vendor/vibeio/src/net/mod.rs @@ -0,0 +1,55 @@ +//! A networking module for `vibeio`. +//! +//! This module provides async versions of common networking operations: +//! - TCP: [`TcpListener`], [`TcpStream`], [`PollTcpStream`] +//! - UDP: [`UdpSocket`] +//! - Unix domain sockets: [`UnixListener`], [`UnixStream`], [`PollUnixStream`] +//! +//! Implementation notes: +//! - On Linux with io_uring support, some operations use native async syscalls (e.g. `accept4`, `sendto`) +//! via the async driver. When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations either offload to a blocking thread pool +//! or fall back to synchronous std::net calls. +//! - The runtime must be active when calling these functions; otherwise they will panic. +//! +//! # Examples +//! +//! ## TCP Server +//! +//! ```ignore +//! use vibeio::net::TcpListener; +//! +//! let listener = TcpListener::bind("127.0.0.1:8080").await?; +//! loop { +//! let (stream, addr) = listener.accept().await?; +//! println!("Connection from: {}", addr); +//! } +//! ``` +//! +//! ## UDP Client +//! +//! ```ignore +//! use vibeio::net::UdpSocket; +//! +//! let socket = UdpSocket::bind("127.0.0.1:0").await?; +//! socket.connect("127.0.0.1:9000").await?; +//! socket.send(b"hello").await?; +//! ``` +//! +//! ## Unix Domain Socket +//! +//! ```ignore +//! use vibeio::net::UnixStream; +//! +//! let stream = UnixStream::connect("/tmp/mysocket").await?; +//! ``` + +mod tcp; +mod udp; +#[cfg(unix)] +mod unix; + +pub use tcp::*; +pub use udp::*; +#[cfg(unix)] +pub use unix::*; diff --git a/vendor/vibeio/src/net/tcp/listener.rs b/vendor/vibeio/src/net/tcp/listener.rs new file mode 100644 index 0000000..5eb10c9 --- /dev/null +++ b/vendor/vibeio/src/net/tcp/listener.rs @@ -0,0 +1,469 @@ +//! TCP listener types for async I/O. +//! +//! This module provides: +//! - [`TcpListener`]: An async TCP listener that can use either completion-based or poll-based I/O. +//! +//! # Implementation details +//! +//! - On Linux with io_uring support, TCP operations use native async syscalls via the async driver. +//! - When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations fall back to synchronous std::net calls. +//! - The runtime must be active when calling these types' methods; otherwise they will panic. + +use std::future::poll_fn; +use std::io; +use std::mem::ManuallyDrop; +use std::net::{SocketAddr, TcpListener as StdTcpListener, ToSocketAddrs}; +#[cfg(unix)] +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::Networking::WinSock::{ + self as WinSock, AF_INET, AF_INET6, IPPROTO_IPV6, IPV6_V6ONLY, SOCKADDR, SOCKADDR_IN, + SOCKADDR_IN6, SOCKADDR_STORAGE, SOCK_STREAM, SOMAXCONN, WSADATA, +}; + +use crate::op::AcceptOp; +use crate::{fd_inner::InnerRawHandle, net::TcpStream}; + +#[cfg(unix)] +fn socket_addr_to_raw( + address: SocketAddr, +) -> (libc::c_int, libc::sockaddr_storage, libc::socklen_t) { + match address { + SocketAddr::V4(address) => { + let sockaddr = libc::sockaddr_in { + sin_family: libc::AF_INET as libc::sa_family_t, + sin_port: address.port().to_be(), + sin_addr: libc::in_addr { + s_addr: u32::from_ne_bytes(address.ip().octets()), + }, + sin_zero: [0; 8], + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin_len: 0, + }; + + let mut storage = std::mem::MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + libc::AF_INET, + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + SocketAddr::V6(address) => { + let sockaddr = libc::sockaddr_in6 { + sin6_family: libc::AF_INET6 as libc::sa_family_t, + sin6_port: address.port().to_be(), + sin6_flowinfo: address.flowinfo(), + sin6_addr: libc::in6_addr { + s6_addr: address.ip().octets(), + }, + sin6_scope_id: address.scope_id(), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin6_len: 0, + }; + + let mut storage = std::mem::MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + libc::AF_INET6, + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + } +} + +#[cfg(unix)] +fn bind_one(address: SocketAddr) -> Result { + let (domain, raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let socket_fd = unsafe { libc::socket(domain, libc::SOCK_STREAM, 0) }; + if socket_fd == -1 { + return Err(io::Error::last_os_error()); + } + + let reuse_addr: libc::c_int = 1; + let reuse_addr_result = unsafe { + libc::setsockopt( + socket_fd, + libc::SOL_SOCKET, + libc::SO_REUSEADDR, + (&reuse_addr as *const libc::c_int).cast(), + std::mem::size_of_val(&reuse_addr) as libc::socklen_t, + ) + }; + if reuse_addr_result == -1 { + let err = io::Error::last_os_error(); + let _ = unsafe { libc::close(socket_fd) }; + return Err(err); + } + + if domain == libc::AF_INET6 { + let ipv6_only: libc::c_int = 0; + let ipv6_only_result = unsafe { + libc::setsockopt( + socket_fd, + libc::IPPROTO_IPV6, + libc::IPV6_V6ONLY, + (&ipv6_only as *const libc::c_int).cast(), + std::mem::size_of_val(&ipv6_only) as libc::socklen_t, + ) + }; + if ipv6_only_result == -1 { + let err = io::Error::last_os_error(); + let _ = unsafe { libc::close(socket_fd) }; + return Err(err); + } + } + + let bind_result = unsafe { + libc::bind( + socket_fd, + (&raw_addr as *const libc::sockaddr_storage).cast::(), + raw_addr_len, + ) + }; + if bind_result == -1 { + let err = io::Error::last_os_error(); + let _ = unsafe { libc::close(socket_fd) }; + return Err(err); + } + + let listen_result = unsafe { libc::listen(socket_fd, libc::SOMAXCONN) }; + if listen_result == -1 { + let err = io::Error::last_os_error(); + let _ = unsafe { libc::close(socket_fd) }; + return Err(err); + } + + Ok(unsafe { StdTcpListener::from_raw_fd(socket_fd) }) +} + +#[cfg(windows)] +fn socket_addr_to_raw(address: SocketAddr) -> (i32, SOCKADDR_STORAGE, i32) { + match address { + SocketAddr::V4(address) => { + let mut sockaddr = SOCKADDR_IN::default(); + sockaddr.sin_family = AF_INET; + sockaddr.sin_port = address.port().to_be(); + sockaddr.sin_addr.S_un.S_addr = u32::from_ne_bytes(address.ip().octets()); + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + + ( + AF_INET as i32, + storage, + std::mem::size_of::() as i32, + ) + } + SocketAddr::V6(address) => { + let mut sockaddr = SOCKADDR_IN6::default(); + sockaddr.sin6_family = AF_INET6; + sockaddr.sin6_port = address.port().to_be(); + sockaddr.sin6_flowinfo = address.flowinfo(); + sockaddr.sin6_addr.u.Byte = address.ip().octets(); + sockaddr.Anonymous.sin6_scope_id = address.scope_id(); + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN6 as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + + ( + AF_INET6 as i32, + storage, + std::mem::size_of::() as i32, + ) + } + } +} + +#[cfg(windows)] +fn bind_one(address: SocketAddr) -> Result { + // 0x202 = MAKEWORD(2, 2) + let mut wsadata = WSADATA::default(); + if unsafe { WinSock::WSAStartup(0x202, &mut wsadata as *mut WSADATA) } != 0 { + return Err(io::Error::last_os_error()); + } + + let (domain, raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let socket = unsafe { WinSock::socket(domain, SOCK_STREAM, 0) }; + if socket == WinSock::INVALID_SOCKET { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + if domain == AF_INET6 as i32 { + let ipv6_only: i32 = 0; + let ipv6_only_result = unsafe { + WinSock::setsockopt( + socket, + IPPROTO_IPV6 as i32, + IPV6_V6ONLY as i32, + (&ipv6_only as *const i32).cast(), + std::mem::size_of_val(&ipv6_only) as i32, + ) + }; + if ipv6_only_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + let _ = unsafe { WinSock::closesocket(socket) }; + return Err(io::Error::from_raw_os_error(err_code)); + } + } + + let bind_result = unsafe { + WinSock::bind( + socket, + (&raw_addr as *const SOCKADDR_STORAGE).cast::(), + raw_addr_len, + ) + }; + if bind_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + let _ = unsafe { WinSock::closesocket(socket) }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + let listen_result = unsafe { WinSock::listen(socket, SOMAXCONN as i32) }; + if listen_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + let _ = unsafe { WinSock::closesocket(socket) }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + Ok(unsafe { StdTcpListener::from_raw_socket(socket as RawSocket) }) +} + +/// An async TCP listener that can use either completion-based or poll-based I/O. +/// +/// This is the async version of [`std::net::TcpListener`]. +/// +/// # Implementation details +/// +/// - On Linux with io_uring support, TCP operations use native async syscalls via the async driver. +/// - When io_uring completion is available, operations complete directly. +/// - For platforms without native async support, operations fall back to synchronous std::net calls. +/// - The runtime must be active when calling these methods; otherwise they will panic. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::TcpListener; +/// +/// let listener = TcpListener::bind("127.0.0.1:8080").await?; +/// loop { +/// let (stream, addr) = listener.accept().await?; +/// println!("Connection from: {}", addr); +/// } +/// ``` +pub struct TcpListener { + inner: StdTcpListener, + handle: ManuallyDrop, +} + +impl TcpListener { + /// Creates a new `TcpListener` which will be bound to the specified address. + /// + /// This is the async version of [`std::net::TcpListener::bind`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - The address is already in use + /// - The process lacks permissions to bind to the address + /// - The runtime is not active + #[inline] + pub fn bind(address: impl ToSocketAddrs) -> Result { + let addresses = address.to_socket_addrs()?; + let mut last_error = None; + for address in addresses { + match bind_one(address) { + Ok(inner) => return Self::from_std(inner), + Err(err) => last_error = Some(err), + } + } + + Err(last_error + .unwrap_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no addresses"))) + } + + /// Creates a new `TcpListener` from a standard library `TcpListener`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: std::net::TcpListener) -> Result { + #[cfg(unix)] + let handle = ManuallyDrop::new(InnerRawHandle::new(inner.as_raw_fd(), Interest::READABLE)?); + #[cfg(windows)] + let handle = ManuallyDrop::new(InnerRawHandle::new( + crate::fd_inner::RawOsHandle::Socket(inner.as_raw_socket()), + Interest::READABLE, + )?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Creates a new `TcpListener` from a standard library `TcpListener` in poll mode. + /// + /// This could be useful when using cloned `TcpListener` on Windows. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[cfg(windows)] + #[inline] + pub fn from_std_poll(inner: std::net::TcpListener) -> Result { + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + crate::fd_inner::RawOsHandle::Socket(inner.as_raw_socket()), + Interest::READABLE, + crate::driver::RegistrationMode::Poll, + )?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Returns the local address of this listener. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not bound. + #[inline] + pub fn local_addr(&self) -> Result { + self.inner.local_addr() + } + + /// Accepts a new incoming connection from this listener. + /// + /// This is the async version of [`std::net::TcpListener::accept`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The listener is not bound to an address + /// - The runtime is not active + #[inline] + pub async fn accept(&self) -> Result<(TcpStream, SocketAddr), io::Error> { + let mut op = AcceptOp::new(&self.handle); + let (raw, address) = poll_fn(move |cx| self.handle.poll_op(cx, &mut op)).await?; + // Recreate a std TcpStream from the raw fd and convert it into our async TcpStream. + // If conversion fails, the std TcpStream will be dropped and the fd closed. + #[cfg(unix)] + let std_stream = unsafe { std::net::TcpStream::from_raw_fd(raw) }; + #[cfg(windows)] + let crate::fd_inner::RawOsHandle::Socket(raw) = raw + else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid raw handle", + )); + }; + #[cfg(windows)] + let std_stream = unsafe { std::net::TcpStream::from_raw_socket(raw) }; + match TcpStream::from_std(std_stream) { + Ok(stream) => Ok((stream, address)), + Err(err) => Err(err), + } + } +} + +#[cfg(unix)] +impl AsRawFd for TcpListener { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for TcpListener { + #[inline] + fn into_raw_fd(self) -> RawFd { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its fd ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_fd() + } + } +} + +#[cfg(windows)] +impl AsRawSocket for TcpListener { + #[inline] + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } +} + +#[cfg(windows)] +impl IntoRawSocket for TcpListener { + #[inline] + fn into_raw_socket(self) -> RawSocket { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its socket ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_socket() + } + } +} + +impl Drop for TcpListener { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/net/tcp/mod.rs b/vendor/vibeio/src/net/tcp/mod.rs new file mode 100644 index 0000000..b12d1ad --- /dev/null +++ b/vendor/vibeio/src/net/tcp/mod.rs @@ -0,0 +1,128 @@ +mod listener; +mod stream; + +pub use listener::*; +pub use stream::*; + +#[cfg(test)] +mod tests { + use std::io::{self as std_io}; + use std::net::{Shutdown, SocketAddr}; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use crate::io::{AsyncRead, AsyncWrite}; + use crate::{driver::AnyDriver, executor::spawn}; + + use super::{PollTcpStream, TcpListener, TcpStream}; + + #[inline] + fn try_bind_listener(address: SocketAddr) -> Option { + match TcpListener::bind(address) { + Ok(listener) => Some(listener), + Err(err) if err.kind() == std_io::ErrorKind::PermissionDenied => None, + Err(err) => panic!("listener should bind: {err}"), + } + } + + #[test] + fn tcp_listener_and_stream_exchange_data() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + #[cfg(windows)] + AnyDriver::new_iocp().expect("iocp driver should initialize"), + ); + runtime.block_on(async { + let address = "127.0.0.1:0" + .parse::() + .expect("address should parse"); + let Some(listener) = try_bind_listener(address) else { + return; + }; + let server_address = listener + .local_addr() + .expect("listener should expose address"); + + let server = spawn(async move { + let (mut stream, _) = listener.accept().await?; + let buffer = [0u8; 4]; + let (read, buffer) = stream.read(buffer).await; + let read = read?; + assert_eq!(&buffer[..read], b"ping"); + stream.write(b"pong".to_vec()).await.0?; + stream.shutdown(Shutdown::Both)?; + Ok::<(), std_io::Error>(()) + }); + + let mut client = TcpStream::connect(server_address) + .await + .expect("client should connect"); + client + .write(b"ping".to_vec()) + .await + .0 + .expect("client should write"); + let response = [0u8; 4]; + let (read, response) = client.read(response).await; + let read = read.expect("client should read"); + assert_eq!(&response[..read], b"pong"); + assert_eq!( + client + .peer_addr() + .expect("peer address should be available"), + server_address + ); + client + .shutdown(Shutdown::Both) + .expect("shutdown should succeed"); + + server.await.expect("server task should complete"); + }); + } + + #[test] + fn poll_tcp_stream_uses_readiness_path() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + #[cfg(windows)] + AnyDriver::new_iocp().expect("iocp driver should initialize"), + ); + runtime.block_on(async { + let address = "127.0.0.1:0" + .parse::() + .expect("address should parse"); + let Some(listener) = try_bind_listener(address) else { + return; + }; + let server_address = listener + .local_addr() + .expect("listener should expose address"); + + let server = spawn(async move { + let (mut stream, _) = listener.accept().await?; + let received = [0u8; 4]; + let (read, received) = stream.read(received).await; + let read = read?; + assert_eq!(&received[..read], b"mio!"); + stream.write(b"ok".to_vec()).await.0?; + Ok::<(), std_io::Error>(()) + }); + + let mut client = PollTcpStream::connect(server_address) + .await + .expect("client should connect"); + AsyncWriteExt::write_all(&mut client, b"mio!") + .await + .expect("tokio write_all should succeed"); + let mut response = [0u8; 2]; + AsyncReadExt::read_exact(&mut client, &mut response) + .await + .expect("tokio read_exact should succeed"); + assert_eq!(&response, b"ok"); + + server.await.expect("server task should complete"); + }); + } +} diff --git a/vendor/vibeio/src/net/tcp/stream.rs b/vendor/vibeio/src/net/tcp/stream.rs new file mode 100644 index 0000000..2105158 --- /dev/null +++ b/vendor/vibeio/src/net/tcp/stream.rs @@ -0,0 +1,794 @@ +//! TCP stream types for async I/O. +//! +//! This module provides: +//! - [`TcpStream`]: An async TCP stream that can use either completion-based or poll-based I/O. +//! - [`PollTcpStream`]: A poll-only variant that always uses readiness-based operations. +//! +//! # Implementation details +//! +//! - On Linux with io_uring support, TCP operations use native async syscalls via the async driver. +//! - When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations fall back to synchronous std::net calls. +//! - The runtime must be active when calling these types' methods; otherwise they will panic. + +use std::cell::RefCell; +use std::future::poll_fn; +use std::io::{self, IoSlice}; +use std::mem::{ManuallyDrop, MaybeUninit}; +use std::net::{Shutdown, SocketAddr, ToSocketAddrs}; +#[cfg(unix)] +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use mio::Interest; +use tokio::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite, ReadBuf}; + +#[cfg(windows)] +use windows_sys::Win32::Networking::WinSock::{ + self, AF_INET, AF_INET6, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6, SOCKADDR_STORAGE, SOCK_STREAM, + WSADATA, +}; + +use crate::io::{ + AsInnerRawHandle, AsyncReadPoll, AsyncWritePoll, IoBuf, IoBufMut, IoBufTemporaryPoll, + IoVectoredBuf, IoVectoredBufMut, IoVectoredBufTemporaryPoll, +}; +use crate::op::{ConnectOp, ReadOp, ReadinessOp, ReadvOp, RecvOp, WriteOp, WritevOp}; +use crate::{ + driver::RegistrationMode, + fd_inner::InnerRawHandle, + io::{AsyncRead, AsyncWrite}, +}; + +#[cfg(unix)] +fn socket_addr_to_raw( + address: SocketAddr, +) -> (libc::c_int, libc::sockaddr_storage, libc::socklen_t) { + match address { + SocketAddr::V4(address) => { + let sockaddr = libc::sockaddr_in { + sin_family: libc::AF_INET as libc::sa_family_t, + sin_port: address.port().to_be(), + sin_addr: libc::in_addr { + s_addr: u32::from_ne_bytes(address.ip().octets()), + }, + sin_zero: [0; 8], + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + libc::AF_INET, + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + SocketAddr::V6(address) => { + let sockaddr = libc::sockaddr_in6 { + sin6_family: libc::AF_INET6 as libc::sa_family_t, + sin6_port: address.port().to_be(), + sin6_flowinfo: address.flowinfo(), + sin6_addr: libc::in6_addr { + s6_addr: address.ip().octets(), + }, + sin6_scope_id: address.scope_id(), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin6_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + libc::AF_INET6, + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + } +} + +#[cfg(windows)] +fn socket_addr_to_raw(address: SocketAddr) -> (i32, SOCKADDR_STORAGE, i32) { + match address { + SocketAddr::V4(address) => { + let mut sockaddr = SOCKADDR_IN::default(); + sockaddr.sin_family = AF_INET; + sockaddr.sin_port = address.port().to_be(); + sockaddr.sin_addr.S_un.S_addr = u32::from_ne_bytes(address.ip().octets()); + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + ( + AF_INET as _, + storage, + std::mem::size_of::() as i32, + ) + } + SocketAddr::V6(address) => { + let mut sockaddr = SOCKADDR_IN6::default(); + sockaddr.sin6_family = AF_INET6; + sockaddr.sin6_port = address.port().to_be(); + sockaddr.sin6_flowinfo = address.flowinfo(); + sockaddr.sin6_addr.u.Byte = address.ip().octets(); + sockaddr.Anonymous.sin6_scope_id = address.scope_id() as u32; + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN6 as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + ( + AF_INET6 as _, + storage, + std::mem::size_of::() as i32, + ) + } + } +} + +#[cfg(unix)] +fn new_socket( + address: SocketAddr, +) -> Result<(std::net::TcpStream, libc::sockaddr_storage, libc::socklen_t), io::Error> { + let (domain, raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let socket_fd = unsafe { libc::socket(domain, libc::SOCK_STREAM, 0) }; + if socket_fd == -1 { + return Err(io::Error::last_os_error()); + } + let stream = unsafe { std::net::TcpStream::from_raw_fd(socket_fd.into_raw_fd()) }; + Ok((stream, raw_addr, raw_addr_len)) +} + +#[cfg(windows)] +fn new_socket( + address: SocketAddr, +) -> Result<(std::net::TcpStream, SOCKADDR_STORAGE, i32), io::Error> { + // 0x202 = MAKEWORD(2, 2) + let mut wsadata = WSADATA::default(); + if unsafe { WinSock::WSAStartup(0x202, &mut wsadata as *mut WSADATA) } != 0 { + return Err(io::Error::last_os_error()); + } + let (domain, raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let socket = unsafe { WinSock::socket(domain, SOCK_STREAM, 0) }; + if socket == WinSock::INVALID_SOCKET { + let err = io::Error::last_os_error(); + let _ = unsafe { WinSock::WSACleanup() }; + return Err(err); + } + let stream = unsafe { std::net::TcpStream::from_raw_socket(socket as _) }; + Ok((stream, raw_addr, raw_addr_len)) +} + +/// An async TCP stream that can use either completion-based or poll-based I/O. +/// +/// This is the async version of [`std::net::TcpStream`]. +/// +/// # Implementation details +/// +/// - On Linux with io_uring support, TCP operations use native async syscalls via the async driver. +/// - When io_uring completion is available, operations complete directly. +/// - For platforms without native async support, operations fall back to synchronous std::net calls. +/// - The runtime must be active when calling these methods; otherwise they will panic. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::TcpStream; +/// +/// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; +/// stream.write(b"hello").await.0?; +/// let mut buf = [0u8; 1024]; +/// let (read, buf) = stream.read(buf).await; +/// let read = read?; +/// ``` +pub struct TcpStream { + inner: std::net::TcpStream, + handle: ManuallyDrop, +} + +/// A poll-only variant that always uses readiness-based operations. +/// +/// This type is useful when you want to ensure readiness-based I/O is used, +/// for example when integrating with other readiness-based systems. +/// +/// # Implementation details +/// +/// - Always uses readiness-based I/O via `mio`. +/// - Can be converted to [`TcpStream`] with adaptive or completion mode. +pub struct PollTcpStream { + stream: TcpStream, + write_ready: RefCell, + read_ready: RefCell, +} + +impl TcpStream { + /// Connects to the specified address. + /// + /// This is the async version of [`std::net::TcpStream::connect`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - Connection refused + /// - Network unreachable + /// - The runtime is not active + #[inline] + pub async fn connect(address: impl ToSocketAddrs) -> Result { + let addresses = address.to_socket_addrs()?; + let mut last_error = None; + for address in addresses { + match Self::connect_one(address).await { + Ok(stream) => return Ok(stream), + Err(err) => last_error = Some(err), + } + } + Err(last_error + .unwrap_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no addresses"))) + } + + #[inline] + async fn connect_one(address: SocketAddr) -> Result { + let (inner, raw_addr, raw_addr_len) = new_socket(address)?; + let stream = Self::from_std(inner)?; + + #[cfg(unix)] + let raw_addr_ptr = (&raw_addr as *const libc::sockaddr_storage).cast::(); + #[cfg(windows)] + let raw_addr_ptr = (&raw_addr as *const SOCKADDR_STORAGE).cast::(); + let handle = &stream.handle; + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await?; + + Ok(stream) + } + + /// Returns the local address of this connection. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn local_addr(&self) -> Result { + self.inner.local_addr() + } + + /// Returns the remote address of this connection. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn peer_addr(&self) -> Result { + self.inner.peer_addr() + } + + /// Returns the current state of the TCP_NODELAY option for this socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn nodelay(&self) -> Result { + self.inner.nodelay() + } + + /// Sets the value of the TCP_NODELAY option for this socket. + /// + /// When set, this disables the Nagle algorithm, which means that small + /// packets are sent immediately rather than being buffered. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn set_nodelay(&self, nodelay: bool) -> Result<(), io::Error> { + self.inner.set_nodelay(nodelay) + } + + /// Shuts down the connection. + #[inline] + pub fn shutdown(&self, how: Shutdown) -> Result<(), io::Error> { + match self.inner.shutdown(how) { + Ok(()) => Ok(()), + Err(e) if e.kind() == io::ErrorKind::NotConnected => Ok(()), // macOS-specific behavior + Err(e) => Err(e), + } + } + + /// Peeks at data from the socket without removing it from the buffer. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub async fn peek(&self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = RecvOp::new_peek(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Creates a new `TcpStream` from a standard library `TcpStream`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: std::net::TcpStream) -> Result { + Self::from_std_with_mode(inner, RegistrationMode::Completion) + } + + /// Creates a new `TcpStream` from a standard library `TcpStream` with a specific registration mode. + #[inline] + pub(crate) fn from_std_with_mode( + inner: std::net::TcpStream, + mode: RegistrationMode, + ) -> Result { + #[cfg(unix)] + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + inner.as_raw_fd(), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + #[cfg(windows)] + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + crate::fd_inner::RawOsHandle::Socket(inner.as_raw_socket()), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Converts this stream into a poll-only variant. + /// + /// The returned `PollTcpStream` will always use readiness-based I/O. + #[inline] + pub fn into_poll(self) -> Result { + let mut stream = self; + stream.handle.rebind_mode(RegistrationMode::Poll)?; + stream + .inner + .set_nonblocking(!stream.handle.uses_completion())?; + Ok(PollTcpStream { + stream, + write_ready: RefCell::new(false), + read_ready: RefCell::new(false), + }) + } +} + +impl PollTcpStream { + /// Connects to the specified address using poll-based I/O. + #[inline] + pub async fn connect(address: impl ToSocketAddrs) -> Result { + let addresses = address.to_socket_addrs()?; + let mut last_error = None; + for address in addresses { + match Self::connect_one(address).await { + Ok(stream) => return Ok(stream), + Err(err) => last_error = Some(err), + } + } + Err(last_error + .unwrap_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no addresses"))) + } + + #[inline] + async fn connect_one(address: SocketAddr) -> Result { + let (inner, raw_addr, raw_addr_len) = new_socket(address)?; + let stream = Self::from_std(inner)?; + + #[cfg(unix)] + let raw_addr_ptr = (&raw_addr as *const libc::sockaddr_storage).cast::(); + #[cfg(windows)] + let raw_addr_ptr = (&raw_addr as *const SOCKADDR_STORAGE).cast::(); + + let handle = &stream.stream.handle; + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await?; + + Ok(stream) + } + + /// Creates a new `PollTcpStream` from a standard library `TcpStream`. + #[inline] + pub fn from_std(inner: std::net::TcpStream) -> Result { + Ok(Self { + stream: TcpStream::from_std_with_mode(inner, RegistrationMode::Poll)?, + write_ready: RefCell::new(false), + read_ready: RefCell::new(false), + }) + } + + /// Converts this poll stream into an adaptive `TcpStream`. + #[inline] + pub fn into_adaptive(self) -> TcpStream { + self.stream + } + + /// Converts this poll stream into a completion-based `TcpStream`. + #[inline] + pub fn into_completion(self) -> Result { + let mut stream = self.stream; + stream.handle.rebind_mode(RegistrationMode::Completion)?; + stream + .inner + .set_nonblocking(!stream.handle.uses_completion())?; + Ok(stream) + } + + /// Returns the local address of this connection. + #[inline] + pub fn local_addr(&self) -> Result { + self.stream.local_addr() + } + + /// Returns the remote address of this connection. + #[inline] + pub fn peer_addr(&self) -> Result { + self.stream.peer_addr() + } + + /// Returns the current state of the TCP_NODELAY option for this socket. + #[inline] + pub fn nodelay(&self) -> Result { + self.stream.nodelay() + } + + /// Sets the value of the TCP_NODELAY option for this socket. + #[inline] + pub fn set_nodelay(&self, nodelay: bool) -> Result<(), io::Error> { + self.stream.set_nodelay(nodelay) + } + + /// Shuts down the connection. + #[inline] + pub fn shutdown(&self, how: Shutdown) -> Result<(), io::Error> { + self.stream.shutdown(how) + } + + /// Peeks at data from the socket without removing it from the buffer. + /// + /// This method uses readiness-based I/O and is compatible with `tokio::io`. + #[inline] + pub async fn peek(&self, buf: &mut [u8]) -> Result { + let handle = &self.stream.handle; + let buf = unsafe { IoBufTemporaryPoll::new(buf.as_mut_ptr(), buf.len()) }; + let mut op = RecvOp::new_peek(handle, buf); + poll_fn(move |cx| handle.poll_op_poll(cx, &mut op)).await + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_readable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.read_ready.borrow() { + let result = io(); + if result.is_err() { + *self.read_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "read not ready")) + } + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_writable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.write_ready.borrow() { + let result = io(); + if result.is_err() { + *self.write_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "write not ready")) + } + } +} + +#[cfg(unix)] +impl AsRawFd for TcpStream { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl AsRawFd for PollTcpStream { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.stream.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for TcpStream { + #[inline] + fn into_raw_fd(self) -> RawFd { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its fd ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_fd() + } + } +} + +#[cfg(unix)] +impl IntoRawFd for PollTcpStream { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.stream.into_raw_fd() + } +} + +#[cfg(windows)] +impl AsRawSocket for TcpStream { + #[inline] + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } +} + +#[cfg(windows)] +impl IntoRawSocket for TcpStream { + #[inline] + fn into_raw_socket(self) -> RawSocket { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its socket ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_socket() + } + } +} + +#[cfg(windows)] +impl AsRawSocket for PollTcpStream { + #[inline] + fn as_raw_socket(&self) -> RawSocket { + self.stream.as_raw_socket() + } +} + +#[cfg(windows)] +impl IntoRawSocket for PollTcpStream { + #[inline] + fn into_raw_socket(self) -> RawSocket { + self.stream.into_raw_socket() + } +} + +impl<'a> AsInnerRawHandle<'a> for TcpStream { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + &self.handle + } +} + +impl<'a> AsInnerRawHandle<'a> for PollTcpStream { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + self.stream.as_inner_raw_handle() + } +} + +impl AsyncRead for TcpStream { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = ReadOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn read_vectored( + &mut self, + bufs: B, + ) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = ReadvOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl TokioAsyncRead for PollTcpStream { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let this = self.get_mut(); + // Equivalent to .assume_init_mut() in Rust 1.93.0+ + let unfilled = unsafe { &mut *(buf.unfilled_mut() as *mut [MaybeUninit] as *mut [u8]) }; + let buf_temp = unsafe { IoBufTemporaryPoll::new(unfilled.as_mut_ptr(), unfilled.len()) }; + let mut op = ReadOp::new(&this.stream.handle, buf_temp); + match this.stream.handle.poll_op_poll(cx, &mut op) { + Poll::Ready(Ok(read)) => { + unsafe { + buf.assume_init(read); + } + buf.advance(read); + Poll::Ready(Ok(())) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncWrite for TcpStream { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = WriteOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } + + #[inline] + async fn write_vectored(&mut self, bufs: B) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = WritevOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl TokioAsyncWrite for PollTcpStream { + #[inline] + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let buf = unsafe { IoBufTemporaryPoll::new(buf.as_ptr() as *mut u8, buf.len()) }; + let mut op = WriteOp::new(&this.stream.handle, buf); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + if bufs.is_empty() { + return Poll::Ready(Ok(0)); + } + let this = self.get_mut(); + let bufs = unsafe { IoVectoredBufTemporaryPoll::new(bufs) }; + let mut op = WritevOp::new(&this.stream.handle, bufs); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn is_write_vectored(&self) -> bool { + true + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + #[inline] + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.get_mut().shutdown(Shutdown::Write)) + } +} + +impl AsyncReadPoll for PollTcpStream { + #[inline] + fn poll_readable(&self, cx: &mut std::task::Context) -> std::task::Poll> { + if *self.read_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .stream + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_readable(&self.stream.handle))?; + *self.read_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl AsyncWritePoll for PollTcpStream { + #[inline] + fn poll_writable(&self, cx: &mut std::task::Context) -> std::task::Poll> { + if *self.write_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .stream + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_writable(&self.stream.handle))?; + *self.write_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl Drop for TcpStream { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/net/udp.rs b/vendor/vibeio/src/net/udp.rs new file mode 100644 index 0000000..6929fab --- /dev/null +++ b/vendor/vibeio/src/net/udp.rs @@ -0,0 +1,1534 @@ +//! UDP socket types for async I/O. +//! +//! This module provides: +//! - [`UdpSocket`]: An async UDP socket that can use either completion-based or poll-based I/O. +//! +//! # Implementation details +//! +//! - On Linux with io_uring support, UDP operations use native async syscalls via the async driver. +//! - When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations fall back to synchronous std::net calls. +//! - The runtime must be active when calling these types' methods; otherwise they will panic. + +use std::cell::RefCell; +use std::future::poll_fn; +use std::io; +use std::mem::ManuallyDrop; +#[cfg(unix)] +use std::mem::MaybeUninit; +use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs, UdpSocket as StdUdpSocket}; +#[cfg(unix)] +use std::os::fd::{AsRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawSocket, IntoRawSocket, RawSocket}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::Networking::WinSock::{ + AF_INET, AF_INET6, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6, SOCKADDR_STORAGE, +}; + +use crate::driver::RegistrationMode; +use crate::fd_inner::InnerRawHandle; +use crate::io::{ + AsInnerRawHandle, AsyncReadPoll, AsyncWritePoll, IoBuf, IoBufMut, IoBufTemporaryPoll, +}; +use crate::op::{ConnectOp, ReadinessOp, RecvOp, RecvfromOp, SendOp, SendtoOp}; + +#[cfg(unix)] +#[inline] +fn socket_addr_to_raw(address: SocketAddr) -> (libc::sockaddr_storage, libc::socklen_t) { + match address { + SocketAddr::V4(address) => { + let sockaddr = libc::sockaddr_in { + sin_family: libc::AF_INET as libc::sa_family_t, + sin_port: address.port().to_be(), + sin_addr: libc::in_addr { + s_addr: u32::from_ne_bytes(address.ip().octets()), + }, + sin_zero: [0; 8], + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + SocketAddr::V6(address) => { + let sockaddr = libc::sockaddr_in6 { + sin6_family: libc::AF_INET6 as libc::sa_family_t, + sin6_port: address.port().to_be(), + sin6_flowinfo: address.flowinfo(), + sin6_addr: libc::in6_addr { + s6_addr: address.ip().octets(), + }, + sin6_scope_id: address.scope_id(), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin6_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + } +} + +#[cfg(windows)] +#[inline] +fn socket_addr_to_raw(address: SocketAddr) -> (SOCKADDR_STORAGE, i32) { + match address { + SocketAddr::V4(address) => { + let mut sockaddr = SOCKADDR_IN::default(); + sockaddr.sin_family = AF_INET; + sockaddr.sin_port = address.port().to_be(); + sockaddr.sin_addr.S_un.S_addr = u32::from_ne_bytes(address.ip().octets()); + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + (storage, std::mem::size_of::() as i32) + } + SocketAddr::V6(address) => { + let mut sockaddr = SOCKADDR_IN6::default(); + sockaddr.sin6_family = AF_INET6; + sockaddr.sin6_port = address.port().to_be(); + sockaddr.sin6_flowinfo = address.flowinfo(); + sockaddr.sin6_addr.u.Byte = address.ip().octets(); + sockaddr.Anonymous.sin6_scope_id = address.scope_id() as u32; + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN6 as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + (storage, std::mem::size_of::() as i32) + } + } +} + +#[cfg(unix)] +#[inline] +async fn connect_one(handle: &InnerRawHandle, address: SocketAddr) -> Result<(), io::Error> { + let (raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let raw_addr_ptr = (&raw_addr as *const libc::sockaddr_storage).cast::(); + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await +} + +#[cfg(windows)] +#[inline] +async fn connect_one( + handle: &mut InnerRawHandle, + socket: &mut StdUdpSocket, + address: SocketAddr, +) -> Result<(), io::Error> { + // Since ConnectEx (used by `ConnectOp`) requires the socket to be connection-oriented, + // we use poll-based I/O instead of completion-based I/O on Windows. + let old_registration_mode = handle.mode(); + handle.rebind_mode(RegistrationMode::Poll)?; + socket.set_nonblocking(true)?; + let (raw_addr, raw_addr_len) = socket_addr_to_raw(address); + let raw_addr_ptr = (&raw_addr as *const SOCKADDR_STORAGE).cast::(); + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + drop(op); + handle.rebind_mode(old_registration_mode)?; + socket.set_nonblocking(!handle.uses_completion())?; + result +} + +/// An async UDP socket that can use either completion-based or poll-based I/O. +/// +/// This is the async version of [`std::net::UdpSocket`]. +/// +/// # Implementation details +/// +/// - On Linux with io_uring support, UDP operations use native async syscalls via the async driver. +/// - When io_uring completion is available, operations complete directly. +/// - For platforms without native async support, operations fall back to synchronous std::net calls. +/// - The runtime must be active when calling these methods; otherwise they will panic. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::UdpSocket; +/// +/// let socket = UdpSocket::bind("127.0.0.1:0").await?; +/// socket.connect("127.0.0.1:9000").await?; +/// socket.send(b"hello").await?; +/// ``` +pub struct UdpSocket { + inner: StdUdpSocket, + handle: ManuallyDrop, +} + +impl UdpSocket { + /// Creates a new `UdpSocket` which will be bound to the specified address. + /// + /// This is the async version of [`std::net::UdpSocket::bind`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - The address is already in use + /// - The process lacks permissions to bind to the address + /// - The runtime is not active + #[inline] + pub fn bind(address: impl ToSocketAddrs) -> Result { + let inner = StdUdpSocket::bind(address)?; + Self::from_std(inner) + } + + /// Creates a new `UdpSocket` from a standard library `UdpSocket`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: StdUdpSocket) -> Result { + Self::from_std_with_mode(inner, RegistrationMode::Completion) + } + + /// Creates a new `UdpSocket` from a standard library `UdpSocket` with a specific registration mode. + #[inline] + pub(crate) fn from_std_with_mode( + inner: StdUdpSocket, + mode: RegistrationMode, + ) -> Result { + #[cfg(unix)] + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + inner.as_raw_fd(), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + #[cfg(windows)] + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + crate::fd_inner::RawOsHandle::Socket(inner.as_raw_socket()), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Converts this socket into a poll-only variant. + /// + /// The returned `PollUdpSocket` will always use readiness-based I/O. + #[inline] + pub fn into_poll(self) -> Result { + let mut socket = self; + socket.handle.rebind_mode(RegistrationMode::Poll)?; + socket + .inner + .set_nonblocking(!socket.handle.uses_completion())?; + Ok(PollUdpSocket { + socket, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } + + /// Converts this `UdpSocket` into the standard library `UdpSocket`. + #[inline] + pub fn into_std(self) -> StdUdpSocket { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration + // handle manually and move out the inner socket. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner) + } + } + + /// Returns the local address of this socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not bound. + #[inline] + pub fn local_addr(&self) -> Result { + self.inner.local_addr() + } + + /// Returns the remote address of this socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn peer_addr(&self) -> Result { + self.inner.peer_addr() + } + + /// Connects this UDP socket to a remote address. + /// + /// This is the async version of [`std::net::UdpSocket::connect`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - Connection fails + /// - The runtime is not active + #[inline] + pub async fn connect(&mut self, address: impl ToSocketAddrs) -> Result<(), io::Error> { + let addresses = address.to_socket_addrs()?; + let mut last_error = None; + for address in addresses { + #[cfg(unix)] + let connect_one_result = connect_one(&self.handle, address).await; + #[cfg(windows)] + let connect_one_result = connect_one(&mut self.handle, &mut self.inner, address).await; + match connect_one_result { + Ok(()) => return Ok(()), + Err(err) => last_error = Some(err), + } + } + Err(last_error + .unwrap_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no addresses"))) + } + + /// Receives a single datagram message. + /// + /// This is the async version of [`std::net::UdpSocket::recv`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The socket is not bound + /// - The runtime is not active + #[inline] + pub async fn recv(&self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = RecvOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Receives a single datagram message, returning the sender's address. + /// + /// This is the async version of [`std::net::UdpSocket::recv_from`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The socket is not bound + /// - The runtime is not active + #[inline] + pub async fn recv_from( + &self, + buf: B, + ) -> (Result<(usize, SocketAddr), io::Error>, B) { + let handle = &self.handle; + let mut op = RecvfromOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Sends data on a connected socket. + /// + /// This is the async version of [`std::net::UdpSocket::send`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The socket is not connected + /// - The runtime is not active + #[inline] + pub async fn send(&self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = SendOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Sends data to the specified address. + /// + /// This is the async version of [`std::net::UdpSocket::send_to`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - The send operation fails + /// - The runtime is not active + #[inline] + pub async fn send_to( + &self, + mut buf: B, + address: impl ToSocketAddrs, + ) -> (Result, B) { + let addresses = match address.to_socket_addrs() { + Ok(addresses) => addresses, + Err(err) => return (Err(err), buf), + }; + let mut last_error = None; + + for address in addresses { + let handle = &self.handle; + let mut op = SendtoOp::new(handle, buf, address); + match poll_fn(|cx| handle.poll_op(cx, &mut op)).await { + Ok(sent) => return (Ok(sent), op.take_bufs()), + Err(err) => { + buf = op.take_bufs(); + last_error = Some(err); + } + } + } + + ( + Err(last_error + .unwrap_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no addresses"))), + buf, + ) + } + + /// Receives data without removing it from the socket's receive queue. + /// + /// This is the async version of [`std::net::UdpSocket::peek`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The socket is not bound + /// - The runtime is not active + #[inline] + pub async fn peek(&self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = RecvOp::new_peek(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Receives data without removing it from the socket's receive queue, + /// returning the sender's address. + /// + /// This is the async version of [`std::net::UdpSocket::peek_from`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The socket is not bound + /// - The runtime is not active + #[inline] + pub async fn peek_from( + &self, + buf: B, + ) -> (Result<(usize, SocketAddr), io::Error>, B) { + let handle = &self.handle; + let mut op = RecvfromOp::new_peek(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + /// Returns a new `UdpSocket` that shares the same underlying file descriptor. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be cloned. + #[inline] + pub fn try_clone(&self) -> Result { + Self::from_std(self.inner.try_clone()?) + } + + /// Sets the broadcast flag. + /// + /// When set, the socket can send broadcast packets. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_broadcast(&self, broadcast: bool) -> Result<(), io::Error> { + self.inner.set_broadcast(broadcast) + } + + /// Returns the current value of the broadcast flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn broadcast(&self) -> Result { + self.inner.broadcast() + } + + /// Sets the time-to-live (TTL) value. + /// + /// This controls how many hops a packet can traverse before being discarded. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_ttl(&self, ttl: u32) -> Result<(), io::Error> { + self.inner.set_ttl(ttl) + } + + /// Returns the current TTL value. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn ttl(&self) -> Result { + self.inner.ttl() + } + + /// Sets the multicast loop flag for IPv4. + /// + /// When set, multicast packets are looped back to the local socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_loop_v4(&self, multicast_loop_v4: bool) -> Result<(), io::Error> { + self.inner.set_multicast_loop_v4(multicast_loop_v4) + } + + /// Returns the current IPv4 multicast loop flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_loop_v4(&self) -> Result { + self.inner.multicast_loop_v4() + } + + /// Sets the multicast TTL for IPv4. + /// + /// This controls how many hops multicast packets can traverse. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_ttl_v4(&self, multicast_ttl_v4: u32) -> Result<(), io::Error> { + self.inner.set_multicast_ttl_v4(multicast_ttl_v4) + } + + /// Returns the current IPv4 multicast TTL. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_ttl_v4(&self) -> Result { + self.inner.multicast_ttl_v4() + } + + /// Sets the multicast loop flag for IPv6. + /// + /// When set, multicast packets are looped back to the local socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_loop_v6(&self, multicast_loop_v6: bool) -> Result<(), io::Error> { + self.inner.set_multicast_loop_v6(multicast_loop_v6) + } + + /// Returns the current IPv6 multicast loop flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_loop_v6(&self) -> Result { + self.inner.multicast_loop_v6() + } + + /// Joins a multicast group for IPv4. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn join_multicast_v4( + &self, + multiaddr: &Ipv4Addr, + interface: &Ipv4Addr, + ) -> Result<(), io::Error> { + self.inner.join_multicast_v4(multiaddr, interface) + } + + /// Joins a multicast group for IPv6. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn join_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> Result<(), io::Error> { + self.inner.join_multicast_v6(multiaddr, interface) + } + + /// Leaves a multicast group for IPv4. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn leave_multicast_v4( + &self, + multiaddr: &Ipv4Addr, + interface: &Ipv4Addr, + ) -> Result<(), io::Error> { + self.inner.leave_multicast_v4(multiaddr, interface) + } + + /// Leaves a multicast group for IPv6. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn leave_multicast_v6( + &self, + multiaddr: &Ipv6Addr, + interface: u32, + ) -> Result<(), io::Error> { + self.inner.leave_multicast_v6(multiaddr, interface) + } + + /// Takes the pending error from the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn take_error(&self) -> Result, io::Error> { + self.inner.take_error() + } + + /// Sets the read timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_read_timeout(&self, dur: Option) -> Result<(), io::Error> { + self.inner.set_read_timeout(dur) + } + + /// Sets the write timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_write_timeout(&self, dur: Option) -> Result<(), io::Error> { + self.inner.set_write_timeout(dur) + } + + /// Returns the read timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn read_timeout(&self) -> Result, io::Error> { + self.inner.read_timeout() + } + + /// Returns the write timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn write_timeout(&self) -> Result, io::Error> { + self.inner.write_timeout() + } +} + +#[cfg(unix)] +impl AsRawFd for UdpSocket { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for UdpSocket { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.into_std().into_raw_fd() + } +} + +#[cfg(windows)] +impl AsRawSocket for UdpSocket { + #[inline] + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } +} + +#[cfg(windows)] +impl IntoRawSocket for UdpSocket { + #[inline] + fn into_raw_socket(self) -> RawSocket { + self.into_std().into_raw_socket() + } +} + +impl Drop for UdpSocket { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} + +impl<'a> AsInnerRawHandle<'a> for UdpSocket { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + &self.handle + } +} + +/// A poll-based UDP socket that always uses readiness-based I/O. +/// +/// This is the poll-only counterpart to [`UdpSocket`], similar to how +/// [`PollTcpStream`](crate::net::PollTcpStream) relates to +/// [`TcpStream`](crate::net::TcpStream). +/// +/// All I/O operations on this type use readiness-based (poll) I/O, +/// regardless of whether the runtime supports completion-based I/O. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::UdpSocket; +/// +/// let socket = UdpSocket::bind("127.0.0.1:0")?; +/// let poll_socket = socket.into_poll()?; +/// ``` +pub struct PollUdpSocket { + socket: UdpSocket, + read_ready: RefCell, + write_ready: RefCell, +} + +impl PollUdpSocket { + /// Creates a new `PollUdpSocket` bound to the specified address. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - The address is already in use + /// - The process lacks permissions to bind to the address + /// - The runtime is not active + #[inline] + pub fn bind(address: impl ToSocketAddrs) -> Result { + let inner = StdUdpSocket::bind(address)?; + Self::from_std(inner) + } + + /// Creates a new `PollUdpSocket` from a standard library `UdpSocket`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: StdUdpSocket) -> Result { + Ok(Self { + socket: UdpSocket::from_std_with_mode(inner, RegistrationMode::Poll)?, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } + + /// Converts this poll socket into an adaptive `UdpSocket`. + #[inline] + pub fn into_adaptive(self) -> UdpSocket { + self.socket + } + + /// Converts this poll socket into a completion-based `UdpSocket`. + /// + /// # Errors + /// + /// This function will return an error if the runtime does not support completion-based I/O. + #[inline] + pub fn into_completion(self) -> Result { + let mut socket = self.socket; + socket.handle.rebind_mode(RegistrationMode::Completion)?; + socket + .inner + .set_nonblocking(!socket.handle.uses_completion())?; + Ok(socket) + } + + /// Connects this UDP socket to a remote address. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - DNS resolution fails + /// - Connection fails + /// - The runtime is not active + #[inline] + pub async fn connect(&mut self, address: impl ToSocketAddrs) -> Result<(), io::Error> { + self.socket.connect(address).await + } + + /// Returns the local address of this socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not bound. + #[inline] + pub fn local_addr(&self) -> Result { + self.socket.local_addr() + } + + /// Returns the remote address of this socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn peer_addr(&self) -> Result { + self.socket.peer_addr() + } + + /// Receives a single datagram message. + /// + /// This is the poll-based version of [`UdpSocket::recv`]. + #[inline] + pub async fn recv(&self, buf: B) -> (Result, B) { + self.socket.recv(buf).await + } + + /// Receives a single datagram message, returning the sender's address. + /// + /// This is the poll-based version of [`UdpSocket::recv_from`]. + #[inline] + pub async fn recv_from( + &self, + buf: B, + ) -> (Result<(usize, SocketAddr), io::Error>, B) { + self.socket.recv_from(buf).await + } + + /// Sends data on a connected socket. + /// + /// This is the poll-based version of [`UdpSocket::send`]. + #[inline] + pub async fn send(&self, buf: B) -> (Result, B) { + self.socket.send(buf).await + } + + /// Sends data to the specified address. + /// + /// This is the poll-based version of [`UdpSocket::send_to`]. + #[inline] + pub async fn send_to( + &self, + buf: B, + address: impl ToSocketAddrs, + ) -> (Result, B) { + self.socket.send_to(buf, address).await + } + + /// Receives data without removing it from the socket's receive queue. + /// + /// This is the poll-based version of [`UdpSocket::peek`]. + #[inline] + pub async fn peek(&self, buf: B) -> (Result, B) { + self.socket.peek(buf).await + } + + /// Receives data without removing it from the socket's receive queue, + /// returning the sender's address. + /// + /// This is the poll-based version of [`UdpSocket::peek_from`]. + #[inline] + pub async fn peek_from( + &self, + buf: B, + ) -> (Result<(usize, SocketAddr), io::Error>, B) { + self.socket.peek_from(buf).await + } + + /// Returns a new `PollUdpSocket` that shares the same underlying file descriptor. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be cloned. + #[inline] + pub fn try_clone(&self) -> Result { + Ok(Self { + socket: self.socket.try_clone()?, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } + + /// Sets the broadcast flag. + /// + /// When set, the socket can send broadcast packets. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_broadcast(&self, broadcast: bool) -> Result<(), io::Error> { + self.socket.set_broadcast(broadcast) + } + + /// Returns the current value of the broadcast flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn broadcast(&self) -> Result { + self.socket.broadcast() + } + + /// Sets the time-to-live (TTL) value. + /// + /// This controls how many hops a packet can traverse before being discarded. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_ttl(&self, ttl: u32) -> Result<(), io::Error> { + self.socket.set_ttl(ttl) + } + + /// Returns the current TTL value. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn ttl(&self) -> Result { + self.socket.ttl() + } + + /// Sets the multicast loop flag for IPv4. + /// + /// When set, multicast packets are looped back to the local socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_loop_v4(&self, multicast_loop_v4: bool) -> Result<(), io::Error> { + self.socket.set_multicast_loop_v4(multicast_loop_v4) + } + + /// Returns the current IPv4 multicast loop flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_loop_v4(&self) -> Result { + self.socket.multicast_loop_v4() + } + + /// Sets the multicast TTL for IPv4. + /// + /// This controls how many hops multicast packets can traverse. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_ttl_v4(&self, multicast_ttl_v4: u32) -> Result<(), io::Error> { + self.socket.set_multicast_ttl_v4(multicast_ttl_v4) + } + + /// Returns the current IPv4 multicast TTL. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_ttl_v4(&self) -> Result { + self.socket.multicast_ttl_v4() + } + + /// Sets the multicast loop flag for IPv6. + /// + /// When set, multicast packets are looped back to the local socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_multicast_loop_v6(&self, multicast_loop_v6: bool) -> Result<(), io::Error> { + self.socket.set_multicast_loop_v6(multicast_loop_v6) + } + + /// Returns the current IPv6 multicast loop flag. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn multicast_loop_v6(&self) -> Result { + self.socket.multicast_loop_v6() + } + + /// Joins a multicast group for IPv4. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn join_multicast_v4( + &self, + multiaddr: &Ipv4Addr, + interface: &Ipv4Addr, + ) -> Result<(), io::Error> { + self.socket.join_multicast_v4(multiaddr, interface) + } + + /// Joins a multicast group for IPv6. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn join_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> Result<(), io::Error> { + self.socket.join_multicast_v6(multiaddr, interface) + } + + /// Leaves a multicast group for IPv4. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn leave_multicast_v4( + &self, + multiaddr: &Ipv4Addr, + interface: &Ipv4Addr, + ) -> Result<(), io::Error> { + self.socket.leave_multicast_v4(multiaddr, interface) + } + + /// Leaves a multicast group for IPv6. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn leave_multicast_v6( + &self, + multiaddr: &Ipv6Addr, + interface: u32, + ) -> Result<(), io::Error> { + self.socket.leave_multicast_v6(multiaddr, interface) + } + + /// Takes the pending error from the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn take_error(&self) -> Result, io::Error> { + self.socket.take_error() + } + + /// Sets the read timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_read_timeout(&self, dur: Option) -> Result<(), io::Error> { + self.socket.set_read_timeout(dur) + } + + /// Sets the write timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be modified. + #[inline] + pub fn set_write_timeout(&self, dur: Option) -> Result<(), io::Error> { + self.socket.set_write_timeout(dur) + } + + /// Returns the read timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn read_timeout(&self) -> Result, io::Error> { + self.socket.read_timeout() + } + + /// Returns the write timeout for the socket. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket cannot be queried. + #[inline] + pub fn write_timeout(&self) -> Result, io::Error> { + self.socket.write_timeout() + } + + /// Polls to receive a single datagram message from the socket. + /// + /// This is the poll-based counterpart to [`UdpSocket::recv`]. + #[inline] + pub fn poll_recv( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_mut_ptr(), buf.len()) }; + let mut op = RecvOp::new(handle, buf_temp); + handle.poll_op_poll(cx, &mut op) + } + + /// Polls to receive a single datagram message, returning the sender's address. + /// + /// This is the poll-based counterpart to [`UdpSocket::recv_from`]. + #[inline] + pub fn poll_recv_from( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_mut_ptr(), buf.len()) }; + let mut op = RecvfromOp::new(handle, buf_temp); + handle.poll_op_poll(cx, &mut op) + } + + /// Polls to send data on a connected socket. + /// + /// This is the poll-based counterpart to [`UdpSocket::send`]. + #[inline] + pub fn poll_send( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_ptr() as *mut u8, buf.len()) }; + let mut op = SendOp::new(handle, buf_temp); + handle.poll_op_poll(cx, &mut op) + } + + /// Polls to send data to the specified address. + /// + /// This is the poll-based counterpart to [`UdpSocket::send_to`]. + #[inline] + pub fn poll_send_to( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + target: SocketAddr, + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_ptr() as *mut u8, buf.len()) }; + let mut op = SendtoOp::new(handle, buf_temp, target); + handle.poll_op_poll(cx, &mut op) + } + + /// Polls to peek at data from the socket without removing it. + /// + /// This is the poll-based counterpart to [`UdpSocket::peek`]. + #[inline] + pub fn poll_peek( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_mut_ptr(), buf.len()) }; + let mut op = RecvOp::new_peek(handle, buf_temp); + handle.poll_op_poll(cx, &mut op) + } + + /// Polls to peek at data from the socket without removing it, + /// returning the sender's address. + /// + /// This is the poll-based counterpart to [`UdpSocket::peek_from`]. + #[inline] + pub fn poll_peek_from( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let this = self.get_mut(); + let handle = &this.socket.handle; + let buf_temp = unsafe { IoBufTemporaryPoll::new(buf.as_mut_ptr(), buf.len()) }; + let mut op = RecvfromOp::new_peek(handle, buf_temp); + handle.poll_op_poll(cx, &mut op) + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_readable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.read_ready.borrow() { + let result = io(); + if result.is_err() { + *self.read_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "read not ready")) + } + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_writable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.write_ready.borrow() { + let result = io(); + if result.is_err() { + *self.write_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "write not ready")) + } + } +} + +impl AsyncReadPoll for PollUdpSocket { + #[inline] + fn poll_readable(&self, cx: &mut Context<'_>) -> Poll> { + if *self.read_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .socket + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_readable(&self.socket.handle))?; + *self.read_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl AsyncWritePoll for PollUdpSocket { + #[inline] + fn poll_writable(&self, cx: &mut Context<'_>) -> Poll> { + if *self.write_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .socket + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_writable(&self.socket.handle))?; + *self.write_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl<'a> AsInnerRawHandle<'a> for PollUdpSocket { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + self.socket.as_inner_raw_handle() + } +} + +#[cfg(unix)] +impl AsRawFd for PollUdpSocket { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.socket.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for PollUdpSocket { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.socket.into_std().into_raw_fd() + } +} + +#[cfg(windows)] +impl AsRawSocket for PollUdpSocket { + #[inline] + fn as_raw_socket(&self) -> RawSocket { + self.socket.inner.as_raw_socket() + } +} + +#[cfg(windows)] +impl IntoRawSocket for PollUdpSocket { + #[inline] + fn into_raw_socket(self) -> RawSocket { + self.socket.into_std().into_raw_socket() + } +} + +#[cfg(test)] +mod tests { + use std::io::{self as std_io}; + use std::net::SocketAddr; + use std::pin::Pin; + + use crate::driver::AnyDriver; + + use super::{PollUdpSocket, UdpSocket}; + + #[inline] + fn try_bind_udp(address: SocketAddr) -> Option { + match UdpSocket::bind(address) { + Ok(socket) => Some(socket), + Err(err) if err.kind() == std_io::ErrorKind::PermissionDenied => None, + Err(err) => panic!("udp socket should bind: {err}"), + } + } + + #[inline] + fn try_bind_poll_udp(address: SocketAddr) -> Option { + match PollUdpSocket::bind(address) { + Ok(socket) => Some(socket), + Err(err) if err.kind() == std_io::ErrorKind::PermissionDenied => None, + Err(err) => panic!("poll udp socket should bind: {err}"), + } + } + + #[test] + fn udp_send_recv_and_peek_variants_work() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + #[cfg(windows)] + AnyDriver::new_iocp().expect("iocp driver should initialize"), + ); + runtime.block_on(async { + let address = "127.0.0.1:0" + .parse::() + .expect("address should parse"); + + let Some(mut server) = try_bind_udp(address) else { + return; + }; + let Some(mut client) = try_bind_udp(address) else { + return; + }; + + let server_addr = server.local_addr().expect("server local_addr should work"); + let client_addr = client.local_addr().expect("client local_addr should work"); + + let sent = client + .send_to(b"ping".to_vec(), server_addr) + .await + .0 + .expect("send_to should succeed"); + assert_eq!(sent, 4); + + let peek_from_buf = vec![0u8; 16]; + let (data, peek_from_buf) = server.peek_from(peek_from_buf).await; + let (peeked, from_peek) = data.expect("peek_from should succeed"); + assert_eq!(&peek_from_buf[..peeked], b"ping"); + assert_eq!(from_peek, client_addr); + + let recv_from_buf = vec![0u8; 16]; + let (data, recv_from_buf) = server.recv_from(recv_from_buf).await; + let (read, from_read) = data.expect("recv_from should succeed"); + assert_eq!(&recv_from_buf[..read], b"ping"); + assert_eq!(from_read, client_addr); + + server + .connect(client_addr) + .await + .expect("server connect should work"); + client + .connect(server_addr) + .await + .expect("client connect should work"); + + let sent = client + .send(b"echo".to_vec()) + .await + .0 + .expect("send should succeed"); + assert_eq!(sent, 4); + + let peek_buf = vec![0u8; 16]; + let (peeked, peek_buf) = server.peek(peek_buf).await; + let peeked = peeked.expect("peek should succeed"); + assert_eq!(&peek_buf[..peeked], b"echo"); + + let recv_buf = vec![0u8; 16]; + let (read, recv_buf) = server.recv(recv_buf).await; + let read = read.expect("recv should succeed"); + assert_eq!(&recv_buf[..read], b"echo"); + }); + } + + #[test] + fn poll_udp_send_recv_and_peek_variants_work() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + #[cfg(windows)] + AnyDriver::new_iocp().expect("iocp driver should initialize"), + ); + runtime.block_on(async { + let address = "127.0.0.1:0" + .parse::() + .expect("address should parse"); + + let Some(server) = try_bind_poll_udp(address) else { + return; + }; + let Some(mut client) = try_bind_poll_udp(address) else { + return; + }; + + let server_addr = server.local_addr().expect("server local_addr should work"); + let client_addr = client.local_addr().expect("client local_addr should work"); + + // poll_send_to + let sent = Pin::new(&mut client) + .send_to(b"ping".to_vec(), server_addr) + .await + .0 + .expect("send_to should succeed"); + assert_eq!(sent, 4); + + // poll_peek_from + let peek_from_buf = vec![0u8; 16]; + let (data, peek_from_buf) = server.peek_from(peek_from_buf).await; + let (peeked, from_peek) = data.expect("peek_from should succeed"); + assert_eq!(&peek_from_buf[..peeked], b"ping"); + assert_eq!(from_peek, client_addr); + + // poll_recv_from + let recv_from_buf = vec![0u8; 16]; + let (data, recv_from_buf) = server.recv_from(recv_from_buf).await; + let (read, from_read) = data.expect("recv_from should succeed"); + assert_eq!(&recv_from_buf[..read], b"ping"); + assert_eq!(from_read, client_addr); + + // Connect for connected send/recv + let mut server = server; + server + .connect(client_addr) + .await + .expect("server connect should work"); + client + .connect(server_addr) + .await + .expect("client connect should work"); + + // poll_send + let sent = Pin::new(&mut client) + .send(b"echo".to_vec()) + .await + .0 + .expect("send should succeed"); + assert_eq!(sent, 4); + + // poll_peek + let peek_buf = vec![0u8; 16]; + let (peeked, peek_buf) = server.peek(peek_buf).await; + let peeked = peeked.expect("peek should succeed"); + assert_eq!(&peek_buf[..peeked], b"echo"); + + // poll_recv + let recv_buf = vec![0u8; 16]; + let (read, recv_buf) = server.recv(recv_buf).await; + let read = read.expect("recv should succeed"); + assert_eq!(&recv_buf[..read], b"echo"); + }); + } + + #[test] + fn poll_udp_into_poll_roundtrip() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + #[cfg(windows)] + AnyDriver::new_iocp().expect("iocp driver should initialize"), + ); + runtime.block_on(async { + let address = "127.0.0.1:0" + .parse::() + .expect("address should parse"); + + let socket = UdpSocket::bind(address).expect("bind should work"); + let poll_socket = socket.into_poll().expect("into_poll should work"); + let _adaptive = poll_socket.into_adaptive(); + }); + } +} diff --git a/vendor/vibeio/src/net/unix/listener.rs b/vendor/vibeio/src/net/unix/listener.rs new file mode 100644 index 0000000..f7633b5 --- /dev/null +++ b/vendor/vibeio/src/net/unix/listener.rs @@ -0,0 +1,145 @@ +//! Unix domain socket listener types for async I/O. +//! +//! This module provides: +//! - [`UnixListener`]: An async Unix domain socket listener. +//! +//! # Implementation details +//! +//! - Unix domain sockets use native async syscalls via the async driver when available. +//! - When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations fall back to synchronous std::os::unix::net calls. +//! - The runtime must be active when calling these types' methods; otherwise they will panic. + +use std::future::poll_fn; +use std::io; +use std::mem::ManuallyDrop; +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::net::{ + SocketAddr, UnixListener as StdUnixListener, UnixStream as StdUnixStream, +}; +use std::path::Path; + +use mio::Interest; + +use crate::fd_inner::InnerRawHandle; +use crate::net::UnixStream; +use crate::op::AcceptUnixOp; + +/// An async Unix domain socket listener. +/// +/// This is the async version of [`std::os::unix::net::UnixListener`]. +/// +/// # Implementation details +/// +/// - Unix domain sockets use native async syscalls via the async driver when available. +/// - When io_uring completion is available, operations complete directly. +/// - For platforms without native async support, operations fall back to synchronous std::os::unix::net calls. +/// - The runtime must be active when calling these methods; otherwise they will panic. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::UnixListener; +/// +/// let listener = UnixListener::bind("/tmp/mysocket").await?; +/// loop { +/// let (stream, addr) = listener.accept().await?; +/// println!("Connection from: {:?}", addr); +/// } +/// ``` +pub struct UnixListener { + inner: StdUnixListener, + handle: ManuallyDrop, +} + +impl UnixListener { + /// Creates a new `UnixListener` which will be bound to the specified path. + /// + /// This is the async version of [`std::os::unix::net::UnixListener::bind`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The path does not exist + /// - The path is too long + /// - The path is already in use + /// - The process lacks permissions + /// - The runtime is not active + #[inline] + pub fn bind(path: impl AsRef) -> Result { + let inner = StdUnixListener::bind(path)?; + Self::from_std(inner) + } + + /// Creates a new `UnixListener` from a standard library `UnixListener`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: StdUnixListener) -> Result { + let handle = ManuallyDrop::new(InnerRawHandle::new(inner.as_raw_fd(), Interest::READABLE)?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Returns the local address of this listener. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not bound. + #[inline] + pub fn local_addr(&self) -> Result { + self.inner.local_addr() + } + + /// Accepts a new incoming connection from this listener. + /// + /// This is the async version of [`std::os::unix::net::UnixListener::accept`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The listener is not bound to an address + /// - The runtime is not active + #[inline] + pub async fn accept(&self) -> Result<(UnixStream, SocketAddr), io::Error> { + let mut op = AcceptUnixOp::new(&self.handle); + let raw = poll_fn(move |cx| self.handle.poll_op(cx, &mut op)).await?; + let std_stream = unsafe { StdUnixStream::from_raw_fd(raw) }; + let address = std_stream.peer_addr()?; + let stream = UnixStream::from_std(std_stream)?; + Ok((stream, address)) + } +} + +impl AsRawFd for UnixListener { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +impl IntoRawFd for UnixListener { + #[inline] + fn into_raw_fd(self) -> RawFd { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its fd ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_fd() + } + } +} + +impl Drop for UnixListener { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/net/unix/mod.rs b/vendor/vibeio/src/net/unix/mod.rs new file mode 100644 index 0000000..abef6c0 --- /dev/null +++ b/vendor/vibeio/src/net/unix/mod.rs @@ -0,0 +1,149 @@ +mod listener; +mod stream; + +pub use listener::*; +pub use stream::*; + +#[cfg(test)] +mod tests { + use std::fs; + use std::io::{self as std_io}; + use std::net::Shutdown; + use std::path::{Path, PathBuf}; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::{SystemTime, UNIX_EPOCH}; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use crate::io::{AsyncRead, AsyncWrite}; + use crate::net::PollUnixStream; + use crate::{driver::AnyDriver, executor::spawn}; + + use super::{UnixListener, UnixStream}; + + fn unique_socket_path(name: &str) -> PathBuf { + static NEXT_ID: AtomicU64 = AtomicU64::new(0); + let id = NEXT_ID.fetch_add(1, Ordering::Relaxed); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + std::env::temp_dir().join(format!("vibeio-{name}-{nanos}-{id}.sock")) + } + + fn cleanup_socket(path: &Path) { + match fs::remove_file(path) { + Ok(()) => {} + Err(err) if err.kind() == std_io::ErrorKind::NotFound => {} + Err(err) => panic!("failed to cleanup socket {}: {err}", path.display()), + } + } + + #[inline] + fn try_bind_listener(path: &Path) -> Option { + match UnixListener::bind(path) { + Ok(listener) => Some(listener), + Err(err) if err.kind() == std_io::ErrorKind::PermissionDenied => None, + Err(err) => panic!("listener should bind: {err}"), + } + } + + #[test] + fn unix_listener_and_stream_exchange_data() { + let runtime = crate::executor::Runtime::new( + AnyDriver::new_mio().expect("mio driver should initialize"), + ); + runtime.block_on(async { + let socket_path = unique_socket_path("exchange"); + cleanup_socket(&socket_path); + + let Some(listener) = try_bind_listener(&socket_path) else { + return; + }; + assert_eq!( + listener + .local_addr() + .expect("listener local addr should be available") + .as_pathname(), + Some(socket_path.as_path()) + ); + + let server = spawn(async move { + let (mut stream, _) = listener.accept().await?; + let buffer = [0u8; 4]; + let (read, buffer) = stream.read(buffer).await; + let read = read?; + assert_eq!(&buffer[..read], b"ping"); + stream.write(b"pong".to_vec()).await.0?; + stream.shutdown(Shutdown::Both)?; + Ok::<(), std_io::Error>(()) + }); + + let mut client = UnixStream::connect(&socket_path) + .await + .expect("client should connect"); + client + .write(b"ping".to_vec()) + .await + .0 + .expect("client should write"); + let response = [0u8; 4]; + let (read, response) = client.read(response).await; + let read = read.expect("client should read"); + assert_eq!(&response[..read], b"pong"); + #[cfg(not(target_os = "macos"))] + assert_eq!( + client + .peer_addr() + .expect("peer address should be available") + .as_pathname(), + Some(socket_path.as_path()) + ); + client + .shutdown(Shutdown::Both) + .expect("shutdown should succeed"); + + server.await.expect("server task should complete"); + cleanup_socket(&socket_path); + }); + } + + #[test] + fn poll_unix_stream_uses_readiness_path() { + let runtime = crate::executor::Runtime::new( + #[cfg(unix)] + AnyDriver::new_mio().expect("mio driver should initialize"), + ); + runtime.block_on(async { + let socket_path = unique_socket_path("vectored"); + cleanup_socket(&socket_path); + let Some(listener) = try_bind_listener(&socket_path) else { + return; + }; + + let server = spawn(async move { + let (mut stream, _) = listener.accept().await?; + let received = [0u8; 4]; + let (read, received) = stream.read(received).await; + let read = read?; + assert_eq!(&received[..read], b"mio!"); + stream.write(b"ok".to_vec()).await.0?; + Ok::<(), std_io::Error>(()) + }); + + let mut client = PollUnixStream::connect(socket_path) + .await + .expect("client should connect"); + AsyncWriteExt::write_all(&mut client, b"mio!") + .await + .expect("tokio write_all should succeed"); + let mut response = [0u8; 2]; + AsyncReadExt::read_exact(&mut client, &mut response) + .await + .expect("tokio read_exact should succeed"); + assert_eq!(&response, b"ok"); + + server.await.expect("server task should complete"); + }); + } +} diff --git a/vendor/vibeio/src/net/unix/stream.rs b/vendor/vibeio/src/net/unix/stream.rs new file mode 100644 index 0000000..1677986 --- /dev/null +++ b/vendor/vibeio/src/net/unix/stream.rs @@ -0,0 +1,568 @@ +//! Unix domain socket stream types for async I/O. +//! +//! This module provides: +//! - [`UnixStream`]: An async Unix domain socket stream that can use either completion-based or poll-based I/O. +//! - [`PollUnixStream`]: A poll-only variant that always uses readiness-based operations. +//! +//! # Implementation details +//! +//! - Unix domain sockets use native async syscalls via the async driver when available. +//! - When io_uring completion is available, operations complete directly. +//! - For platforms without native async support, operations fall back to synchronous std::os::unix::net calls. +//! - The runtime must be active when calling these types' methods; otherwise they will panic. + +use std::cell::RefCell; +use std::future::poll_fn; +use std::io::{self, IoSlice}; +use std::mem::{ManuallyDrop, MaybeUninit}; +use std::net::Shutdown; +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::net::{SocketAddr, UnixStream as StdUnixStream}; +use std::path::Path; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use mio::Interest; +use tokio::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite, ReadBuf}; + +use crate::io::{ + AsInnerRawHandle, AsyncReadPoll, AsyncWritePoll, IoBuf, IoBufMut, IoBufTemporaryPoll, + IoVectoredBuf, IoVectoredBufMut, IoVectoredBufTemporaryPoll, +}; +use crate::op::{ConnectOp, ReadOp, ReadinessOp, ReadvOp, WriteOp, WritevOp}; +use crate::{ + driver::RegistrationMode, + fd_inner::InnerRawHandle, + io::{AsyncRead, AsyncWrite}, +}; + +#[inline] +fn socket_addr_to_raw(path: &Path) -> Result<(libc::sockaddr_un, libc::socklen_t), io::Error> { + let bytes = path.as_os_str().as_bytes(); + if bytes.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "empty socket path", + )); + } + if bytes.contains(&0) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "socket path contains interior NUL byte", + )); + } + + let mut sockaddr = unsafe { MaybeUninit::::zeroed().assume_init() }; + sockaddr.sun_family = libc::AF_UNIX as libc::sa_family_t; + + let max_path_len = sockaddr.sun_path.len(); + if bytes.len() >= max_path_len { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "socket path is too long", + )); + } + + for (index, byte) in bytes.iter().copied().enumerate() { + sockaddr.sun_path[index] = byte as libc::c_char; + } + + let addr_len = + (std::mem::offset_of!(libc::sockaddr_un, sun_path) + bytes.len() + 1) as libc::socklen_t; + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + { + sockaddr.sun_len = addr_len as libc::sa_family_t; + } + + Ok((sockaddr, addr_len)) +} + +#[inline] +fn new_socket( + path: &Path, +) -> Result<(StdUnixStream, libc::sockaddr_un, libc::socklen_t), io::Error> { + let (raw_addr, raw_addr_len) = socket_addr_to_raw(path)?; + let socket_fd = unsafe { libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0) }; + if socket_fd == -1 { + return Err(io::Error::last_os_error()); + } + let stream = unsafe { StdUnixStream::from_raw_fd(socket_fd) }; + Ok((stream, raw_addr, raw_addr_len)) +} + +/// An async Unix domain socket stream that can use either completion-based or poll-based I/O. +/// +/// This is the async version of [`std::os::unix::net::UnixStream`]. +/// +/// # Implementation details +/// +/// - Unix domain sockets use native async syscalls via the async driver when available. +/// - When io_uring completion is available, operations complete directly. +/// - For platforms without native async support, operations fall back to synchronous std::os::unix::net calls. +/// - The runtime must be active when calling these methods; otherwise they will panic. +/// +/// # Examples +/// +/// ```ignore +/// use vibeio::net::UnixStream; +/// +/// let mut stream = UnixStream::connect("/tmp/mysocket").await?; +/// stream.write(b"hello").await.0?; +/// let mut buf = [0u8; 1024]; +/// let (read, buf) = stream.read(buf).await; +/// let read = read?; +/// ``` +pub struct UnixStream { + inner: StdUnixStream, + handle: ManuallyDrop, +} + +/// A poll-only variant that always uses readiness-based operations. +/// +/// This type is useful when you want to ensure readiness-based I/O is used, +/// for example when integrating with other readiness-based systems. +/// +/// # Implementation details +/// +/// - Always uses readiness-based I/O via `mio`. +/// - Can be converted to [`UnixStream`] with adaptive or completion mode. +pub struct PollUnixStream { + stream: UnixStream, + read_ready: RefCell, + write_ready: RefCell, +} + +impl UnixStream { + /// Connects to the specified Unix domain socket path. + /// + /// This is the async version of [`std::os::unix::net::UnixStream::connect`]. + /// + /// # Errors + /// + /// This function will return an error in the following situations: + /// - The path does not exist + /// - The path is too long + /// - Connection refused + /// - The runtime is not active + #[inline] + pub async fn connect(path: impl AsRef) -> Result { + let (inner, raw_addr, raw_addr_len) = new_socket(path.as_ref())?; + let stream = Self::from_std(inner)?; + + let raw_addr_ptr = (&raw_addr as *const libc::sockaddr_un).cast::(); + let handle = &stream.handle; + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await?; + + Ok(stream) + } + + /// Returns the local address of this connection. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn local_addr(&self) -> Result { + self.inner.local_addr() + } + + /// Returns the remote address of this connection. + /// + /// # Errors + /// + /// This function will return an error if the underlying socket is not connected. + #[inline] + pub fn peer_addr(&self) -> Result { + self.inner.peer_addr() + } + + /// Shuts down the connection. + #[inline] + pub fn shutdown(&self, how: Shutdown) -> Result<(), io::Error> { + match self.inner.shutdown(how) { + Ok(()) => Ok(()), + Err(e) if e.kind() == io::ErrorKind::NotConnected => Ok(()), // macOS-specific behavior + Err(e) => Err(e), + } + } + + /// Creates a new `UnixStream` from a standard library `UnixStream`. + /// + /// # Errors + /// + /// This function will return an error if registration with the async driver fails. + #[inline] + pub fn from_std(inner: StdUnixStream) -> Result { + Self::from_std_with_mode(inner, RegistrationMode::Completion) + } + + /// Creates a new `UnixStream` from a standard library `UnixStream` with a specific registration mode. + #[inline] + pub(crate) fn from_std_with_mode( + inner: StdUnixStream, + mode: RegistrationMode, + ) -> Result { + let handle = ManuallyDrop::new(InnerRawHandle::new_with_mode( + inner.as_raw_fd(), + Interest::READABLE | Interest::WRITABLE, + mode, + )?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(Self { inner, handle }) + } + + /// Converts this stream into a poll-only variant. + /// + /// The returned `PollUnixStream` will always use readiness-based I/O. + #[inline] + pub fn into_poll(self) -> Result { + let mut stream = self; + stream.handle.rebind_mode(RegistrationMode::Poll)?; + stream + .inner + .set_nonblocking(!stream.handle.uses_completion())?; + Ok(PollUnixStream { + stream, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } +} + +impl PollUnixStream { + /// Connects to the specified Unix domain socket path using poll-based I/O. + #[inline] + pub async fn connect(path: impl AsRef) -> Result { + let (inner, raw_addr, raw_addr_len) = new_socket(path.as_ref())?; + let stream = Self::from_std(inner)?; + + let raw_addr_ptr = (&raw_addr as *const libc::sockaddr_un).cast::(); + let handle = &stream.stream.handle; + let mut op = ConnectOp::new(handle, raw_addr_ptr, raw_addr_len); + poll_fn(move |cx| handle.poll_op(cx, &mut op)).await?; + + Ok(stream) + } + + /// Creates a new `PollUnixStream` from a standard library `UnixStream`. + #[inline] + pub fn from_std(inner: StdUnixStream) -> Result { + Ok(Self { + stream: UnixStream::from_std_with_mode(inner, RegistrationMode::Poll)?, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } + + /// Converts this poll stream into an adaptive `UnixStream`. + #[inline] + pub fn into_adaptive(self) -> UnixStream { + self.stream + } + + /// Converts this poll stream into a completion-based `UnixStream`. + #[inline] + pub fn into_completion(self) -> Result { + let mut stream = self.stream; + stream.handle.rebind_mode(RegistrationMode::Completion)?; + stream + .inner + .set_nonblocking(!stream.handle.uses_completion())?; + Ok(stream) + } + + /// Returns the local address of this connection. + #[inline] + pub fn local_addr(&self) -> Result { + self.stream.local_addr() + } + + /// Returns the remote address of this connection. + #[inline] + pub fn peer_addr(&self) -> Result { + self.stream.peer_addr() + } + + /// Shuts down the connection. + #[inline] + pub fn shutdown(&self, how: Shutdown) -> Result<(), io::Error> { + self.stream.shutdown(how) + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_readable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.read_ready.borrow() { + let result = io(); + if result.is_err() { + *self.read_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "read not ready")) + } + } + + /// Tries to perform an I/O operation on the socket, returning an error if it is not ready. + #[inline] + pub fn try_io_writable(&self, io: Io) -> io::Result + where + Io: FnOnce() -> io::Result, + { + if *self.write_ready.borrow() { + let result = io(); + if result.is_err() { + *self.write_ready.borrow_mut() = false; + } + result + } else { + Err(io::Error::new(io::ErrorKind::WouldBlock, "write not ready")) + } + } +} + +impl AsRawFd for PollUnixStream { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.stream.inner.as_raw_fd() + } +} + +impl IntoRawFd for PollUnixStream { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.stream.into_raw_fd() + } +} + +impl TokioAsyncRead for PollUnixStream { + #[inline] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let this = self.get_mut(); + // Equivalent to .assume_init_mut() in Rust 1.93.0+ + let unfilled = unsafe { &mut *(buf.unfilled_mut() as *mut [MaybeUninit] as *mut [u8]) }; + let buf_temp = unsafe { IoBufTemporaryPoll::new(unfilled.as_mut_ptr(), unfilled.len()) }; + let mut op = ReadOp::new(&this.stream.handle, buf_temp); + match this.stream.handle.poll_op_poll(cx, &mut op) { + Poll::Ready(Ok(read)) => { + unsafe { + buf.assume_init(read); + } + buf.advance(read); + Poll::Ready(Ok(())) + } + Poll::Ready(Err(err)) => Poll::Ready(Err(err)), + Poll::Pending => Poll::Pending, + } + } +} + +impl TokioAsyncWrite for PollUnixStream { + #[inline] + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let buf = unsafe { IoBufTemporaryPoll::new(buf.as_ptr() as *mut u8, buf.len()) }; + let mut op = WriteOp::new(&this.stream.handle, buf); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + if bufs.is_empty() { + return Poll::Ready(Ok(0)); + } + let this = self.get_mut(); + let bufs = unsafe { IoVectoredBufTemporaryPoll::new(bufs) }; + let mut op = WritevOp::new(&this.stream.handle, bufs); + this.stream.handle.poll_op_poll(cx, &mut op) + } + + #[inline] + fn is_write_vectored(&self) -> bool { + true + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + #[inline] + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.get_mut().shutdown(Shutdown::Write)) + } +} + +impl UnixStream { + /// Creates a new `UnixStream` that is ready to use with readiness-based I/O. + /// + /// This is useful when you want to create a Unix stream that uses poll-based I/O. + #[inline] + pub fn from_std_poll(inner: StdUnixStream) -> Result { + let handle = ManuallyDrop::new(InnerRawHandle::new( + inner.as_raw_fd(), + Interest::READABLE | Interest::WRITABLE, + )?); + inner.set_nonblocking(!handle.uses_completion())?; + Ok(PollUnixStream { + stream: Self { inner, handle }, + read_ready: RefCell::new(false), + write_ready: RefCell::new(false), + }) + } +} + +impl AsRawFd for UnixStream { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +impl<'a> AsInnerRawHandle<'a> for UnixStream { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + &self.handle + } +} + +impl<'a> AsInnerRawHandle<'a> for PollUnixStream { + #[inline] + fn as_inner_raw_handle(&'a self) -> &'a InnerRawHandle { + self.stream.as_inner_raw_handle() + } +} + +impl IntoRawFd for UnixStream { + #[inline] + fn into_raw_fd(self) -> RawFd { + let mut this = ManuallyDrop::new(self); + + // Safety: `this` will not be dropped, so we must drop the registration handle manually. + // We then move out the inner std stream and transfer its fd ownership to the caller. + unsafe { + ManuallyDrop::drop(&mut this.handle); + std::ptr::read(&this.inner).into_raw_fd() + } + } +} + +impl AsyncRead for UnixStream { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = ReadOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn read_vectored( + &mut self, + bufs: B, + ) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = ReadvOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl AsyncWrite for UnixStream { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + let handle = &self.handle; + let mut op = WriteOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } + + #[inline] + async fn write_vectored(&mut self, bufs: B) -> (Result, B) { + if bufs.is_empty() { + return (Ok(0), bufs); + } + let handle = &self.handle; + let mut op = WritevOp::new(handle, bufs); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + (result, op.take_bufs()) + } +} + +impl AsyncReadPoll for PollUnixStream { + #[inline] + fn poll_readable(&self, cx: &mut std::task::Context) -> std::task::Poll> { + if *self.read_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .stream + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_readable(&self.stream.handle))?; + *self.read_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl AsyncWritePoll for PollUnixStream { + #[inline] + fn poll_writable(&self, cx: &mut std::task::Context) -> std::task::Poll> { + if *self.write_ready.borrow() { + return Poll::Ready(Ok(())); + } + let poll = self + .stream + .handle + .poll_op_poll(cx, &mut ReadinessOp::new_writable(&self.stream.handle))?; + *self.write_ready.borrow_mut() = true; + poll.map(Ok) + } +} + +impl Drop for UnixStream { + #[inline] + fn drop(&mut self) { + // Safety: The struct is dropped after the handle is dropped. + unsafe { + ManuallyDrop::drop(&mut self.handle); + } + } +} diff --git a/vendor/vibeio/src/op/accept.rs b/vendor/vibeio/src/op/accept.rs new file mode 100644 index 0000000..209b004 --- /dev/null +++ b/vendor/vibeio/src/op/accept.rs @@ -0,0 +1,683 @@ +#[cfg(windows)] +use std::ffi::c_void; +use std::io; +#[cfg(unix)] +use std::mem::{self, MaybeUninit}; +use std::net::SocketAddr; +#[cfg(unix)] +use std::os::fd::RawFd; +#[cfg(windows)] +use std::ptr; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Networking::WinSock::{ + self as WinSock, AF_INET, AF_INET6, INVALID_SOCKET, IPPROTO_TCP, SOCKADDR, SOCKADDR_IN, + SOCKADDR_IN6, SOCKADDR_STORAGE, SOCKET, SOCK_STREAM, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT, + WSAID_ACCEPTEX, WSAID_GETACCEPTEXSOCKADDRS, WSA_FLAG_OVERLAPPED, WSA_IO_PENDING, + }, + System::IO::OVERLAPPED, +}; + +use crate::driver::CompletionIoResult; +use crate::fd_inner::{InnerRawHandle, RawOsHandle}; +use crate::op::Op; +use crate::{current_driver, driver::AnyDriver}; + +#[cfg(unix)] +fn set_cloexec(fd: RawFd) -> Result<(), io::Error> { + // set FD_CLOEXEC on file descriptor flags + let fdflags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; + if fdflags == -1 { + return Err(io::Error::last_os_error()); + } + if fdflags & libc::FD_CLOEXEC == 0 { + let result = unsafe { libc::fcntl(fd, libc::F_SETFD, fdflags | libc::FD_CLOEXEC) }; + if result == -1 { + return Err(io::Error::last_os_error()); + } + } + + Ok(()) +} + +#[cfg(unix)] +fn sockaddr_storage_to_socketaddr( + storage: &libc::sockaddr_storage, +) -> Result { + // Determine family. ss_family field is platform-dependent type; cast to c_uchar then to c_int for comparison. + let family = storage.ss_family as libc::c_int; + + if family == libc::AF_INET { + let addr_in: &libc::sockaddr_in = + unsafe { &*(storage as *const _ as *const libc::sockaddr_in) }; + let port = u16::from_be(addr_in.sin_port); + // s_addr is in network byte order + let ip_u32 = u32::from_be(addr_in.sin_addr.s_addr); + let ip = std::net::Ipv4Addr::from(ip_u32); + Ok(SocketAddr::V4(std::net::SocketAddrV4::new(ip, port))) + } else if family == libc::AF_INET6 { + let addr_in6: &libc::sockaddr_in6 = + unsafe { &*(storage as *const _ as *const libc::sockaddr_in6) }; + let port = u16::from_be(addr_in6.sin6_port); + let ip = std::net::Ipv6Addr::from(addr_in6.sin6_addr.s6_addr); + Ok(SocketAddr::V6(std::net::SocketAddrV6::new( + ip, + port, + addr_in6.sin6_flowinfo, + addr_in6.sin6_scope_id, + ))) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported socket family", + )) + } +} + +#[cfg(windows)] +fn sockaddr_storage_to_socketaddr(storage: &SOCKADDR_STORAGE) -> Result { + // Determine family. ss_family field is platform-dependent type; cast to c_uchar then to c_int for comparison. + let family = storage.ss_family; + + if family == AF_INET { + let addr_in: &SOCKADDR_IN = unsafe { &*(storage as *const _ as *const SOCKADDR_IN) }; + let port = u16::from_be(addr_in.sin_port); + // s_addr is in network byte order + let ip_u32 = u32::from_be(unsafe { addr_in.sin_addr.S_un.S_addr }); + let ip = std::net::Ipv4Addr::from(ip_u32); + Ok(SocketAddr::V4(std::net::SocketAddrV4::new(ip, port))) + } else if family == AF_INET6 { + let addr_in6: &SOCKADDR_IN6 = unsafe { &*(storage as *const _ as *const SOCKADDR_IN6) }; + let port = u16::from_be(addr_in6.sin6_port); + let ip = std::net::Ipv6Addr::from(unsafe { addr_in6.sin6_addr.u.Byte }); + Ok(SocketAddr::V6(std::net::SocketAddrV6::new( + ip, + port, + addr_in6.sin6_flowinfo, + unsafe { addr_in6.Anonymous.sin6_scope_id }, + ))) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported socket family", + )) + } +} + +#[cfg(windows)] +fn load_accept_ex(socket: SOCKET) -> Result { + let mut bytes_returned: u32 = 0; + let mut accept_ex: WinSock::LPFN_ACCEPTEX = None; + let mut guid = WSAID_ACCEPTEX; + + let ioctl_result = unsafe { + WinSock::WSAIoctl( + socket, + WinSock::SIO_GET_EXTENSION_FUNCTION_POINTER, + (&mut guid as *mut _) as *mut c_void, + std::mem::size_of_val(&guid) as u32, + (&mut accept_ex as *mut _) as *mut c_void, + std::mem::size_of_val(&accept_ex) as u32, + &mut bytes_returned, + ptr::null_mut(), + None, + ) + }; + + if ioctl_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + if accept_ex.is_none() { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "AcceptEx extension function is unavailable", + )); + } + + Ok(accept_ex) +} + +#[cfg(windows)] +fn load_get_accept_ex_sockaddrs( + socket: SOCKET, +) -> Result { + let mut bytes_returned: u32 = 0; + let mut get_accept_ex_sockaddrs: WinSock::LPFN_GETACCEPTEXSOCKADDRS = None; + let mut guid = WSAID_GETACCEPTEXSOCKADDRS; + + let ioctl_result = unsafe { + WinSock::WSAIoctl( + socket, + WinSock::SIO_GET_EXTENSION_FUNCTION_POINTER, + (&mut guid as *mut _) as *mut c_void, + std::mem::size_of_val(&guid) as u32, + (&mut get_accept_ex_sockaddrs as *mut _) as *mut c_void, + std::mem::size_of_val(&get_accept_ex_sockaddrs) as u32, + &mut bytes_returned, + ptr::null_mut(), + None, + ) + }; + + if ioctl_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + if get_accept_ex_sockaddrs.is_none() { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "GetAcceptExSockaddrs extension function is unavailable", + )); + } + + Ok(get_accept_ex_sockaddrs) +} + +#[cfg(windows)] +fn listener_socket_family(listener_socket: SOCKET) -> Result { + let mut addr = SOCKADDR_IN6::default(); + let mut addr_len = std::mem::size_of::() as i32; + let result = unsafe { + WinSock::getsockname( + listener_socket, + (&mut addr as *mut SOCKADDR_IN6).cast::(), + &mut addr_len, + ) + }; + + if result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + Ok(addr.sin6_family as i32) +} + +#[cfg(windows)] +fn create_accept_socket(listener_socket: SOCKET) -> Result { + let family = listener_socket_family(listener_socket)?; + if family != AF_INET as i32 && family != AF_INET6 as i32 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "unsupported listening socket family for AcceptEx", + )); + } + + let accept_socket = unsafe { + WinSock::WSASocketW( + family, + SOCK_STREAM as i32, + IPPROTO_TCP as i32, + ptr::null_mut(), + 0, + WSA_FLAG_OVERLAPPED, + ) + }; + if accept_socket == INVALID_SOCKET { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + Ok(accept_socket) +} + +#[cfg(windows)] +fn set_accept_context(listener_socket: SOCKET, accepted_socket: SOCKET) -> Result<(), io::Error> { + let result = unsafe { + WinSock::setsockopt( + accepted_socket, + SOL_SOCKET as i32, + SO_UPDATE_ACCEPT_CONTEXT as i32, + (&listener_socket as *const SOCKET).cast(), + std::mem::size_of::() as i32, + ) + }; + if result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + Ok(()) +} + +#[cfg(windows)] +const ACCEPTEX_ADDR_LEN: usize = std::mem::size_of::() + 16; +#[cfg(windows)] +const ACCEPTEX_OUTPUT_BUFFER_LEN: usize = ACCEPTEX_ADDR_LEN; + +pub struct AcceptOp<'a> { + handle: &'a InnerRawHandle, + #[cfg(windows)] + accept_ex: Option, + #[cfg(windows)] + get_accept_ex_sockaddrs: Option, + #[cfg(windows)] + accept_socket: Option, + #[cfg(windows)] + bytes_received: u32, + #[cfg(windows)] + accept_output_buffer: Option>, + completion_token: Option, +} + +impl<'a> AcceptOp<'a> { + #[inline] + pub fn new(handle: &'a InnerRawHandle) -> Self { + Self { + handle, + #[cfg(windows)] + accept_ex: None, + #[cfg(windows)] + get_accept_ex_sockaddrs: None, + #[cfg(windows)] + accept_socket: None, + #[cfg(windows)] + bytes_received: 0, + #[cfg(windows)] + accept_output_buffer: None, + completion_token: None, + } + } +} + +impl Op for AcceptOp<'_> { + type Output = (RawOsHandle, SocketAddr); + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + #[cfg(unix)] + { + #[cfg(syscall_accept4)] + let accepted_fd = unsafe { + libc::accept4( + self.handle.handle, + std::ptr::null_mut(), + std::ptr::null_mut(), + libc::SOCK_CLOEXEC | libc::SOCK_NONBLOCK, + ) + }; + #[cfg(not(syscall_accept4))] + let accepted_fd = unsafe { + libc::accept( + self.handle.handle, + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + if accepted_fd == -1 { + let error = io::Error::last_os_error(); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::READABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + let fd = accepted_fd as RawFd; + + // On non-Linux Unix, set close-on-exec + non-blocking manually. + // Linux accept4() above already set these atomically. + #[cfg(not(syscall_accept4))] + if let Err(err) = set_cloexec(fd) { + return Poll::Ready(Err(err)); + } + + // Obtain peer address via getpeername into a sockaddr_storage + let mut peer = MaybeUninit::::zeroed(); + let mut peer_len = mem::size_of::() as libc::socklen_t; + let getpeername_result = unsafe { + libc::getpeername( + fd, + peer.as_mut_ptr().cast::(), + &mut peer_len, + ) + }; + + if getpeername_result == -1 { + let error = io::Error::last_os_error(); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::READABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + let peer = unsafe { peer.assume_init() }; + let address = sockaddr_storage_to_socketaddr(&peer)?; + Poll::Ready(Ok((fd as RawOsHandle, address))) + } + + #[cfg(windows)] + { + let RawOsHandle::Socket(listener_socket) = self.handle.handle else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid raw handle", + ))); + }; + + let accepted_socket = unsafe { + WinSock::accept(listener_socket as SOCKET, ptr::null_mut(), ptr::null_mut()) + }; + if accepted_socket == INVALID_SOCKET { + let error = io::Error::from_raw_os_error(unsafe { WinSock::WSAGetLastError() }); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::READABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + let mut peer = SOCKADDR_STORAGE::default(); + let mut peer_len = std::mem::size_of::() as i32; + let getpeername_result = unsafe { + WinSock::getpeername( + accepted_socket, + (&mut peer as *mut SOCKADDR_STORAGE).cast::(), + &mut peer_len, + ) + }; + if getpeername_result == WinSock::SOCKET_ERROR { + let error = io::Error::from_raw_os_error(unsafe { WinSock::WSAGetLastError() }); + unsafe { WinSock::closesocket(accepted_socket) }; + return Poll::Ready(Err(error)); + } + + let address = match sockaddr_storage_to_socketaddr(&peer) { + Ok(address) => address, + Err(err) => { + unsafe { WinSock::closesocket(accepted_socket) }; + return Poll::Ready(Err(err)); + } + }; + + Poll::Ready(Ok(( + RawOsHandle::Socket(accepted_socket as std::os::windows::io::RawSocket), + address, + ))) + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + #[cfg(windows)] + if let Some(accept_socket) = self.accept_socket.take() { + unsafe { WinSock::closesocket(accept_socket) }; + } + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + + #[cfg(unix)] + { + let fd = result as RawFd; + + // Ensure close-on-exec for the accepted fd (io_uring may have already set them) + if let Err(err) = set_cloexec(fd) { + return Poll::Ready(Err(err)); + } + + // Get peer address via getpeername + let mut peer = MaybeUninit::::zeroed(); + let mut peer_len = mem::size_of::() as libc::socklen_t; + let getpeername_result = unsafe { + libc::getpeername( + fd, + peer.as_mut_ptr().cast::(), + &mut peer_len, + ) + }; + + if getpeername_result == -1 { + return Poll::Ready(Err(io::Error::last_os_error())); + } + + let peer = unsafe { peer.assume_init() }; + let address = sockaddr_storage_to_socketaddr(&peer)?; + Poll::Ready(Ok((fd as RawOsHandle, address))) + } + + #[cfg(windows)] + { + use windows_sys::Win32::Storage::FileSystem::ACCESS_DELETE; + + let RawOsHandle::Socket(listener_socket) = self.handle.handle else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidInput, + "AcceptEx can be used only with listening sockets", + ))); + }; + + let Some(accept_socket) = self.accept_socket.take() else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::Other, + "AcceptEx completion missing accepted socket", + ))); + }; + + if let Err(err) = set_accept_context(listener_socket as SOCKET, accept_socket) { + unsafe { WinSock::closesocket(accept_socket) }; + return Poll::Ready(Err(err)); + } + + let peer = match self.accept_output_buffer.take() { + Some(buf) => buf, + None => { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::Other, + "AcceptEx completion missing peer address", + ))); + } + }; + + let mut local_sockaddr: *mut SOCKADDR_STORAGE = std::ptr::null_mut(); + let mut local_sockaddr_len: i32 = 0; + let mut remote_sockaddr: *mut SOCKADDR_STORAGE = std::ptr::null_mut(); + let mut remote_sockaddr_len: i32 = 0; + + let get_accept_ex_sockaddrs = self.get_accept_ex_sockaddrs.ok_or_else(|| { + io::Error::new( + io::ErrorKind::Unsupported, + "GetAcceptExSockaddrs extension function is unavailable", + ) + })?; + let Some(get_accept_ex_sockaddrs_fn) = get_accept_ex_sockaddrs else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::Unsupported, + "GetAcceptExSockaddrs extension function is unavailable", + ))); + }; + + let _ = unsafe { + get_accept_ex_sockaddrs_fn( + peer.as_ptr() as *const c_void, + 0, + 0, + ACCEPTEX_ADDR_LEN as _, + &mut local_sockaddr as *mut *mut SOCKADDR_STORAGE as *mut *mut SOCKADDR, + &mut local_sockaddr_len as *mut i32, + &mut remote_sockaddr as *mut *mut SOCKADDR_STORAGE as *mut *mut SOCKADDR, + &mut remote_sockaddr_len as *mut i32, + ) + }; + + if remote_sockaddr.is_null() { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::Other, + "can't obtain remote socket address", + ))); + } + + let address = match sockaddr_storage_to_socketaddr( + &(unsafe { std::ptr::read_unaligned(remote_sockaddr as *const SOCKADDR_STORAGE) }), + ) { + Ok(address) => address, + Err(err) => { + unsafe { WinSock::closesocket(accept_socket) }; + return Poll::Ready(Err(err)); + } + }; + + return Poll::Ready(Ok((RawOsHandle::Socket(accept_socket as _), address))); + } + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let RawOsHandle::Socket(listener_socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "AcceptEx can be used only with listening sockets", + )); + }; + let listener_socket = listener_socket as SOCKET; + + if self.accept_ex.is_none() { + self.accept_ex = Some(load_accept_ex(listener_socket)?); + } + if self.get_accept_ex_sockaddrs.is_none() { + self.get_accept_ex_sockaddrs = Some(load_get_accept_ex_sockaddrs(listener_socket)?); + } + let accept_ex = self.accept_ex.ok_or_else(|| { + io::Error::new( + io::ErrorKind::Unsupported, + "AcceptEx extension function is unavailable", + ) + })?; + + if self.accept_socket.is_none() { + self.accept_socket = Some(create_accept_socket(listener_socket)?); + } + let accept_socket = self + .accept_socket + .expect("accept_socket must be initialized"); + if self.accept_output_buffer.is_none() { + self.accept_output_buffer = + Some(vec![0u8; ACCEPTEX_OUTPUT_BUFFER_LEN].into_boxed_slice()); + } + let accept_output_buffer = self + .accept_output_buffer + .as_mut() + .expect("accept_output_buffer must be initialized"); + + let Some(accept_ex_fn) = accept_ex else { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "AcceptEx extension function is unavailable", + )); + }; + + let accept_result = unsafe { + accept_ex_fn( + listener_socket, + accept_socket, + accept_output_buffer.as_mut_ptr().cast::(), + 0, + 0, + ACCEPTEX_ADDR_LEN as u32, + &mut self.bytes_received, + overlapped, + ) + }; + if accept_result != 0 { + return Ok(()); + } + + let err_code = unsafe { WinSock::WSAGetLastError() }; + if err_code == WSA_IO_PENDING { + Ok(()) + } else { + if let Some(socket) = self.accept_socket.take() { + unsafe { WinSock::closesocket(socket) }; + } + Err(io::Error::from_raw_os_error(err_code)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::Accept::new( + types::Fd(self.handle.handle), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + .flags(libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for AcceptOp<'_> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + #[cfg(windows)] + if let Some(socket) = self.accept_socket.take() { + unsafe { WinSock::closesocket(socket) }; + } + } +} diff --git a/vendor/vibeio/src/op/accept_unix.rs b/vendor/vibeio/src/op/accept_unix.rs new file mode 100644 index 0000000..257622e --- /dev/null +++ b/vendor/vibeio/src/op/accept_unix.rs @@ -0,0 +1,162 @@ +use std::io; +use std::os::fd::RawFd; +use std::task::{Context, Poll}; + +use mio::Interest; + +use crate::current_driver; +use crate::driver::{AnyDriver, CompletionIoResult}; +use crate::fd_inner::InnerRawHandle; +use crate::op::Op; + +#[inline] +fn set_cloexec(fd: RawFd) -> Result<(), io::Error> { + // set FD_CLOEXEC on file descriptor flags + let fdflags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; + if fdflags == -1 { + return Err(io::Error::last_os_error()); + } + if fdflags & libc::FD_CLOEXEC == 0 { + let result = unsafe { libc::fcntl(fd, libc::F_SETFD, fdflags | libc::FD_CLOEXEC) }; + if result == -1 { + return Err(io::Error::last_os_error()); + } + } + + Ok(()) +} + +pub struct AcceptUnixOp<'a> { + handle: &'a InnerRawHandle, + completion_token: Option, +} + +impl<'a> AcceptUnixOp<'a> { + #[inline] + pub fn new(handle: &'a InnerRawHandle) -> Self { + Self { + handle, + completion_token: None, + } + } +} + +impl Op for AcceptUnixOp<'_> { + type Output = RawFd; + + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + #[cfg(syscall_accept4)] + let accepted_fd = unsafe { + libc::accept4( + self.handle.handle, + std::ptr::null_mut(), + std::ptr::null_mut(), + libc::SOCK_CLOEXEC | libc::SOCK_NONBLOCK, + ) + }; + #[cfg(not(syscall_accept4))] + let accepted_fd = unsafe { + libc::accept( + self.handle.handle, + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + }; + if accepted_fd == -1 { + let error = io::Error::last_os_error(); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::READABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + let fd = accepted_fd as RawFd; + // On non-Linux Unix, set close-on-exec manually. + // Linux accept4() above already set it atomically. + #[cfg(not(syscall_accept4))] + if let Err(err) = set_cloexec(fd) { + return Poll::Ready(Err(err)); + } + + Poll::Ready(Ok(fd)) + } + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + + let fd = result as RawFd; + if let Err(err) = set_cloexec(fd) { + return Poll::Ready(Err(err)); + } + + Poll::Ready(Ok(fd)) + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::Accept::new( + types::Fd(self.handle.handle), + std::ptr::null_mut(), + std::ptr::null_mut(), + ) + .build() + .user_data(user_data); + Ok(entry) + } +} + +impl Drop for AcceptUnixOp<'_> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/connect.rs b/vendor/vibeio/src/op/connect.rs new file mode 100644 index 0000000..09282f2 --- /dev/null +++ b/vendor/vibeio/src/op/connect.rs @@ -0,0 +1,531 @@ +#[cfg(windows)] +use std::ffi::c_void; +use std::io; +#[cfg(unix)] +use std::mem; +#[cfg(unix)] +use std::mem::MaybeUninit; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::Networking::WinSock::{ + self as WinSock, AF_INET, AF_INET6, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6, SOCKADDR_STORAGE, + SOCKET, SOCKET_ERROR, SOL_SOCKET, SO_UPDATE_CONNECT_CONTEXT, WSAEADDRINUSE, WSAEALREADY, + WSAEINPROGRESS, WSAEINVAL, WSAENOTCONN, WSAEWOULDBLOCK, WSAID_CONNECTEX, WSA_IO_PENDING, +}; +#[cfg(windows)] +use windows_sys::Win32::System::IO::OVERLAPPED; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::op::Op; + +#[cfg(unix)] +fn start_nonblocking_connect( + fd: std::os::fd::RawFd, + raw_addr: *const libc::sockaddr, + raw_addr_len: libc::socklen_t, +) -> Result<(), io::Error> { + let connect_result = unsafe { libc::connect(fd, raw_addr, raw_addr_len) }; + + if connect_result == -1 { + let err = io::Error::last_os_error(); + if !matches!( + err.raw_os_error(), + Some(libc::EINPROGRESS) | Some(libc::EWOULDBLOCK) | Some(libc::EALREADY) + ) { + return Err(err); + } + } + + Ok(()) +} + +#[cfg(windows)] +fn start_nonblocking_connect( + socket: std::os::windows::io::RawSocket, + raw_addr: *const SOCKADDR, + raw_addr_len: i32, +) -> Result<(), io::Error> { + let connect_result = unsafe { WinSock::connect(socket as SOCKET, raw_addr, raw_addr_len) }; + + if connect_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + if !matches!(err_code, WSAEINPROGRESS | WSAEWOULDBLOCK | WSAEALREADY) { + return Err(io::Error::from_raw_os_error(err_code)); + } + } + + Ok(()) +} + +#[cfg(windows)] +fn ensure_connectex_bound(socket: SOCKET, addr: *const SOCKADDR) -> Result<(), io::Error> { + let family = unsafe { (*addr).sa_family as i32 }; + let bind_result = match family { + x if x == AF_INET as i32 => { + let local = SOCKADDR_IN { + sin_family: AF_INET as u16, + ..Default::default() + }; + unsafe { + WinSock::bind( + socket, + (&local as *const SOCKADDR_IN).cast::(), + std::mem::size_of::() as i32, + ) + } + } + x if x == AF_INET6 as i32 => { + let local = SOCKADDR_IN6 { + sin6_family: AF_INET6 as u16, + ..Default::default() + }; + unsafe { + WinSock::bind( + socket, + (&local as *const SOCKADDR_IN6).cast::(), + std::mem::size_of::() as i32, + ) + } + } + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "unsupported socket family for ConnectEx", + )) + } + }; + + if bind_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + if !matches!(err_code, WSAEINVAL | WSAEADDRINUSE) { + return Err(io::Error::from_raw_os_error(err_code)); + } + } + + Ok(()) +} + +#[cfg(windows)] +fn load_connect_ex(socket: SOCKET) -> Result { + let mut bytes_returned: u32 = 0; + let mut connect_ex: WinSock::LPFN_CONNECTEX = None; + let mut guid = WSAID_CONNECTEX; + + let ioctl_result = unsafe { + WinSock::WSAIoctl( + socket, + WinSock::SIO_GET_EXTENSION_FUNCTION_POINTER, + (&mut guid as *mut _) as *mut c_void, + std::mem::size_of_val(&guid) as u32, + (&mut connect_ex as *mut _) as *mut c_void, + std::mem::size_of_val(&connect_ex) as u32, + &mut bytes_returned, + std::ptr::null_mut(), + None, + ) + }; + + if ioctl_result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + + if connect_ex.is_none() { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "ConnectEx extension function is unavailable", + )); + } + + Ok(connect_ex) +} + +#[cfg(windows)] +fn set_connect_context(socket: SOCKET) -> Result<(), io::Error> { + let result = unsafe { + WinSock::setsockopt( + socket, + SOL_SOCKET as i32, + SO_UPDATE_CONNECT_CONTEXT as i32, + std::ptr::null(), + 0, + ) + }; + if result == WinSock::SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + return Err(io::Error::from_raw_os_error(err_code)); + } + Ok(()) +} + +pub struct ConnectOp<'a> { + handle: &'a InnerRawHandle, + #[cfg(unix)] + addr: (*const libc::sockaddr, libc::socklen_t), + #[cfg(windows)] + addr: (*const SOCKADDR, i32), + #[cfg(windows)] + connect_ex: Option, + #[cfg(windows)] + completion_bound: bool, + completion_token: Option, + #[cfg(any(unix, windows))] + poll_connect_started: bool, +} + +impl<'a> ConnectOp<'a> { + #[cfg(unix)] + #[inline] + pub fn new( + handle: &'a InnerRawHandle, + addr: *const libc::sockaddr, + addrlen: libc::socklen_t, + ) -> Self { + Self { + handle, + addr: (addr, addrlen), + completion_token: None, + #[cfg(windows)] + connect_ex: None, + #[cfg(windows)] + completion_bound: false, + poll_connect_started: false, + } + } + + #[cfg(windows)] + #[inline] + pub fn new(handle: &'a InnerRawHandle, addr: *const SOCKADDR, addrlen: i32) -> Self { + Self { + handle, + addr: (addr, addrlen), + completion_token: None, + connect_ex: None, + completion_bound: false, + poll_connect_started: false, + } + } +} + +impl Op for ConnectOp<'_> { + type Output = (); + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + if !self.poll_connect_started { + #[cfg(unix)] + let handle = self.handle.handle; + #[cfg(windows)] + let crate::fd_inner::RawOsHandle::Socket(handle) = self.handle.handle + else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid raw handle", + ))); + }; + + if let Err(err) = start_nonblocking_connect(handle, self.addr.0, self.addr.1) { + return Poll::Ready(Err(err)); + }; + + self.poll_connect_started = true; + } + + #[cfg(unix)] + { + let mut socket_error: libc::c_int = 0; + let mut socket_error_len = mem::size_of::() as libc::socklen_t; + let getsockopt_result = unsafe { + libc::getsockopt( + self.handle.handle, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&mut socket_error as *mut libc::c_int).cast(), + &mut socket_error_len, + ) + }; + if getsockopt_result == -1 { + let error = io::Error::last_os_error(); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + if socket_error != 0 { + if matches!( + socket_error, + libc::EINPROGRESS | libc::EALREADY | libc::EWOULDBLOCK + ) { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(io::Error::from_raw_os_error(socket_error))); + } + + let mut peer = MaybeUninit::::zeroed(); + let mut peer_len = mem::size_of::() as libc::socklen_t; + let getpeername_result = unsafe { + libc::getpeername( + self.handle.handle, + peer.as_mut_ptr().cast::(), + &mut peer_len, + ) + }; + + if getpeername_result == -1 { + let err = io::Error::last_os_error(); + if matches!( + err.raw_os_error(), + Some(libc::EINPROGRESS) + | Some(libc::EALREADY) + | Some(libc::EWOULDBLOCK) + | Some(libc::ENOTCONN) + ) { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + + return Poll::Ready(Err(err)); + } + + Poll::Ready(Ok(())) + } + + #[cfg(windows)] + { + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid raw handle", + ))); + }; + let socket = socket as SOCKET; + + let mut socket_error: i32 = 0; + let mut socket_error_len = std::mem::size_of::() as i32; + let getsockopt_result = unsafe { + WinSock::getsockopt( + socket, + SOL_SOCKET as i32, + WinSock::SO_ERROR as i32, + (&mut socket_error as *mut i32).cast(), + &mut socket_error_len, + ) + }; + if getsockopt_result == SOCKET_ERROR { + let error = io::Error::from_raw_os_error(unsafe { WinSock::WSAGetLastError() }); + if error.kind() == io::ErrorKind::WouldBlock { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(error)); + } + + if socket_error != 0 { + if matches!(socket_error, WSAEINPROGRESS | WSAEALREADY | WSAEWOULDBLOCK) { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + return Poll::Ready(Err(io::Error::from_raw_os_error(socket_error))); + } + + let mut peer = SOCKADDR_STORAGE::default(); + let mut peer_len = std::mem::size_of::() as i32; + let getpeername_result = unsafe { + WinSock::getpeername( + socket, + (&mut peer as *mut SOCKADDR_STORAGE).cast::(), + &mut peer_len, + ) + }; + + if getpeername_result == SOCKET_ERROR { + let err_code = unsafe { WinSock::WSAGetLastError() }; + if matches!( + err_code, + WSAEINPROGRESS | WSAEALREADY | WSAEWOULDBLOCK | WSAENOTCONN + ) { + if let Err(err) = + driver.submit_poll(self.handle, cx.waker().clone(), Interest::WRITABLE) + { + return Poll::Ready(Err(err)); + } + return Poll::Pending; + } + + return Poll::Ready(Err(io::Error::from_raw_os_error(err_code))); + } + + Poll::Ready(Ok(())) + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + + #[cfg(windows)] + { + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidInput, + "ConnectEx can be used only with socket handles", + ))); + }; + + if let Err(err) = set_connect_context(socket as SOCKET) { + return Poll::Ready(Err(err)); + } + } + + Poll::Ready(Ok(())) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "ConnectEx can be used only with socket handles", + )); + }; + + let socket = socket as SOCKET; + + if !self.completion_bound { + ensure_connectex_bound(socket, self.addr.0)?; + self.completion_bound = true; + } + + let connect_ex = if let Some(connect_ex) = self.connect_ex { + connect_ex + } else { + let connect_ex = load_connect_ex(socket)?; + self.connect_ex = Some(connect_ex); + connect_ex + }; + + let Some(connect_ex_fn) = connect_ex else { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "ConnectEx extension function is unavailable", + )); + }; + + let connect_result = unsafe { + connect_ex_fn( + socket, + self.addr.0, + self.addr.1, + std::ptr::null(), + 0, + std::ptr::null_mut(), + overlapped, + ) + }; + + if connect_result != 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(err)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let (addr, addrlen) = self.addr; + + let entry = opcode::Connect::new(types::Fd(self.handle.handle), addr, addrlen) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for ConnectOp<'_> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/fsync.rs b/vendor/vibeio/src/op/fsync.rs new file mode 100644 index 0000000..eb664cc --- /dev/null +++ b/vendor/vibeio/src/op/fsync.rs @@ -0,0 +1,90 @@ +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +use crate::op::Op; + +pub struct FsyncOp<'a> { + handle: &'a InnerRawHandle, + data_only: bool, + completion_token: Option, +} + +impl<'a> FsyncOp<'a> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, data_only: bool) -> Self { + Self { + handle, + data_only, + completion_token: None, + } + } +} + +impl Op for FsyncOp<'_> { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let mut entry_builder = opcode::Fsync::new(types::Fd(self.handle.handle)); + if self.data_only { + entry_builder = entry_builder.flags(types::FsyncFlags::DATASYNC); + } + let entry = entry_builder.build().user_data(user_data); + + Ok(entry) + } +} + +impl Drop for FsyncOp<'_> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/hard_link.rs b/vendor/vibeio/src/op/hard_link.rs new file mode 100644 index 0000000..71e27a8 --- /dev/null +++ b/vendor/vibeio/src/op/hard_link.rs @@ -0,0 +1,93 @@ +use std::ffi::CString; +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct HardLinkOp { + old_path: CString, + new_path: CString, + completion_token: Option, +} + +impl HardLinkOp { + #[inline] + pub fn new(old_path: CString, new_path: CString) -> Self { + Self { + old_path, + new_path, + completion_token: None, + } + } +} + +impl Op for HardLinkOp { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::LinkAt::new( + types::Fd(libc::AT_FDCWD), + self.old_path.as_ptr(), + types::Fd(libc::AT_FDCWD), + self.new_path.as_ptr(), + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for HardLinkOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/io_util.rs b/vendor/vibeio/src/op/io_util.rs new file mode 100644 index 0000000..66347d0 --- /dev/null +++ b/vendor/vibeio/src/op/io_util.rs @@ -0,0 +1,76 @@ +use std::io; +use std::task::{Context, Poll}; + +use mio::Interest; + +use crate::driver::AnyDriver; +use crate::fd_inner::InnerRawHandle; + +pub(crate) enum CompletionBuffer { + Inline(B), + Boxed(Box), +} + +impl CompletionBuffer { + #[inline] + pub(crate) fn new(buf: B, stable: bool) -> Self { + if stable { + Self::Boxed(Box::new(buf)) + } else { + Self::Inline(buf) + } + } + + #[inline] + pub(crate) fn as_ref(&self) -> &B { + match self { + Self::Inline(buf) => buf, + Self::Boxed(buf) => buf.as_ref(), + } + } + + #[inline] + pub(crate) fn as_mut(&mut self) -> &mut B { + match self { + Self::Inline(buf) => buf, + Self::Boxed(buf) => buf.as_mut(), + } + } + + #[inline] + pub(crate) fn into_inner(self) -> B { + match self { + Self::Inline(buf) => buf, + Self::Boxed(buf) => *buf, + } + } + + #[inline] + pub(crate) fn into_stable_box(self) -> Box { + match self { + Self::Inline(buf) => Box::new(buf), + Self::Boxed(buf) => buf, + } + } +} + +#[inline] +pub(crate) fn poll_result_or_wait( + result: io::Result, + handle: &InnerRawHandle, + cx: &mut Context<'_>, + driver: &AnyDriver, + interest: Interest, +) -> Poll> { + match result { + Ok(value) => Poll::Ready(Ok(value)), + Err(err) if err.kind() == io::ErrorKind::WouldBlock => { + if let Err(submit_err) = driver.submit_poll(handle, cx.waker().clone(), interest) { + Poll::Ready(Err(submit_err)) + } else { + Poll::Pending + } + } + Err(err) => Poll::Ready(Err(err)), + } +} diff --git a/vendor/vibeio/src/op/mkdir.rs b/vendor/vibeio/src/op/mkdir.rs new file mode 100644 index 0000000..7570c9c --- /dev/null +++ b/vendor/vibeio/src/op/mkdir.rs @@ -0,0 +1,89 @@ +use std::ffi::CString; +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct MkDirOp { + path: CString, + mode: libc::mode_t, + completion_token: Option, +} + +impl MkDirOp { + #[inline] + pub fn new(path: CString, mode: libc::mode_t) -> Self { + Self { + path, + mode, + completion_token: None, + } + } +} + +impl Op for MkDirOp { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::MkDirAt::new(types::Fd(libc::AT_FDCWD), self.path.as_ptr()) + .mode(self.mode) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for MkDirOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/mod.rs b/vendor/vibeio/src/op/mod.rs new file mode 100644 index 0000000..20948cd --- /dev/null +++ b/vendor/vibeio/src/op/mod.rs @@ -0,0 +1,272 @@ +mod accept; +#[cfg(unix)] +mod accept_unix; +mod connect; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod fsync; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod hard_link; +mod io_util; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod mkdir; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod open; +mod read; +#[cfg(feature = "fs")] +mod readat; +mod readiness; +mod readv; +mod recv; +mod recvfrom; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod rename; +mod send; +mod sendto; +#[cfg(all(target_os = "linux", feature = "splice"))] +mod splice; +#[cfg(all( + target_os = "linux", + any(target_env = "gnu", musl_v1_2_3), + feature = "fs" +))] +mod statx; // musl libc 1.1.x doesn't support statx +#[cfg(all(target_os = "linux", feature = "fs"))] +mod symlink; +#[cfg(all(target_os = "linux", feature = "fs"))] +mod unlink; +#[cfg(all(target_os = "linux", feature = "process"))] +mod waitpid; +mod write; +#[cfg(feature = "fs")] +mod writeat; +mod writev; + +use std::io; +use std::task::Context; +use std::task::Poll; + +pub use accept::AcceptOp; +#[cfg(unix)] +pub use accept_unix::AcceptUnixOp; +pub use connect::ConnectOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use fsync::FsyncOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use hard_link::HardLinkOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use mkdir::MkDirOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use open::OpenOp; +pub use read::ReadOp; +#[cfg(feature = "fs")] +pub use readat::ReadAtOp; +pub use readiness::ReadinessOp; +pub use readv::ReadvOp; +pub use recv::RecvOp; +pub use recvfrom::RecvfromOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use rename::RenameOp; +pub use send::SendOp; +pub use sendto::SendtoOp; +#[cfg(all(target_os = "linux", feature = "splice"))] +pub use splice::SpliceOp; +#[cfg(all( + target_os = "linux", + any(target_env = "gnu", musl_v1_2_3), + feature = "fs" +))] +pub use statx::StatxOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use symlink::SymlinkOp; +#[cfg(all(target_os = "linux", feature = "fs"))] +pub use unlink::UnlinkOp; +#[cfg(all(target_os = "linux", feature = "process"))] +pub use waitpid::WaitPidOp; +pub use write::WriteOp; +#[cfg(feature = "fs")] +pub use writeat::WriteAtOp; +pub use writev::WritevOp; + +use crate::driver::AnyDriver; + +pub trait Op { + /// I/O operation return type + type Output; + + /// Polls the operation for readiness (poll-based I/O). + #[inline] + fn poll_poll( + &mut self, + _cx: &mut Context<'_>, + _driver: &AnyDriver, + ) -> Poll> { + Poll::Ready(Err(io::Error::new( + io::ErrorKind::Unsupported, + "operation does not support poll-based submission", + ))) + } + + /// Polls the operation for readiness (completion-based I/O). + #[inline] + fn poll_completion( + &mut self, + _cx: &mut Context<'_>, + _driver: &AnyDriver, + ) -> Poll> { + Poll::Ready(Err(io::Error::new( + io::ErrorKind::Unsupported, + "operation does not support completion-based submission", + ))) + } + + /// Polls the operation for readiness (automatically determined I/O). + #[inline] + fn poll(&mut self, cx: &mut Context<'_>, driver: &AnyDriver) -> Poll> { + if driver.supports_completion() { + self.poll_completion(cx, driver) + } else { + self.poll_poll(cx, driver) + } + } + + /// Builds an io_uring submission entry for this operation. Returns the + /// constructed SQE. + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + _user_data: u64, + ) -> Result { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "operation does not support completion-based submission", + )) + } + + /// Submits a Windows overlapped I/O operation for this operation. + #[cfg(windows)] + #[inline] + fn submit_windows( + &mut self, + _overlapped: *mut windows_sys::Win32::System::IO::OVERLAPPED, + ) -> Result<(), io::Error> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "operation does not support completion-based submission", + )) + } +} + +#[cfg(all(test, target_os = "linux"))] +mod vectored_uring_tests { + use std::future::poll_fn; + + use mio::Interest; + + use crate::op::{ReadvOp, WritevOp}; + use crate::{driver::AnyDriver, fd_inner::InnerRawHandle}; + + #[test] + fn io_uring_vectored_read_write_pipe() { + // Create a runtime with an io_uring driver and run the test inside it so + // current_driver() is available for InnerRawHandle::new. + let driver = AnyDriver::new_uring().expect("failed to create uring driver"); + let runtime = crate::executor::Runtime::new(driver); + + runtime.block_on(async { + // create a pipe (pair of fds) + let mut fds: [libc::c_int; 2] = [0, 0]; + #[cfg(syscall_pipe2)] + let res = + unsafe { libc::pipe2(fds.as_mut_ptr() as *mut libc::c_int, libc::O_NONBLOCK) }; + #[cfg(not(syscall_pipe2))] + let res = unsafe { libc::pipe(fds.as_mut_ptr() as *mut libc::c_int) }; + assert_eq!(res, 0, "pipe() failed"); + + let rfd = fds[0]; + let wfd = fds[1]; + + // Register both ends with the runtime. Since the current driver supports + // completion and the InnerRawHandle default chooses completion mode, these + // handles will use the completion path (io_uring) where available. + let rhandle = InnerRawHandle::new(rfd, Interest::READABLE) + .expect("failed to create reader InnerRawHandle"); + let whandle = InnerRawHandle::new(wfd, Interest::WRITABLE) + .expect("failed to create writer InnerRawHandle"); + + // Prepare vectored write buffers + let a = b"hello "; + let b = b"world!"; + let bufs = vec![ + libc::iovec { + iov_base: a.as_ptr() as *mut libc::c_void, + iov_len: a.len(), + }, + libc::iovec { + iov_base: b.as_ptr() as *mut libc::c_void, + iov_len: b.len(), + }, + ]; + let total_len = a.len() + b.len(); + + // Submit vectored write. poll_writev will choose completion-path (io_uring) + // for this handle because the driver supports completions. + let whandle_ref = &whandle; + let mut writev_op = WritevOp::new(whandle_ref, bufs); + let write_res = poll_fn(move |cx| whandle_ref.poll_op(cx, &mut writev_op)).await; + let written = write_res.expect("writev failed"); + assert!( + written > 0, + "expected writev to write at least one byte (wrote 0)" + ); + + // Prepare vectored read buffers (split sizes arbitrarily) + let mut dst1 = vec![0u8; 3]; // will receive "hel" + let mut dst2 = vec![0u8; total_len - 3]; // rest + let rd_bufs = vec![ + libc::iovec { + iov_base: dst1.as_mut_ptr() as *mut libc::c_void, + iov_len: dst1.len(), + }, + libc::iovec { + iov_base: dst2.as_mut_ptr() as *mut libc::c_void, + iov_len: dst2.len(), + }, + ]; + + // Read using vectored read. poll_readv will choose completion-path when available. + let rhandle_ref = &rhandle; + let mut readv_op = ReadvOp::new(rhandle_ref, rd_bufs); + let read_res = poll_fn(move |cx| rhandle_ref.poll_op(cx, &mut readv_op)).await; + let read = read_res.expect("readv failed"); + + // We expect to read at least as many bytes as were written (pipe semantics + // on local write -> read without closing may give the bytes). + // It's possible for writev to be partial; just check we read some data and + // that the bytes we did read match the written prefix. + assert!(read > 0, "expected to read at least one byte"); + + // Reconstruct the received bytes and compare with written prefix. + let mut received = Vec::with_capacity(read); + received.extend_from_slice(&dst1[..dst1.len().min(read)]); + if read > dst1.len() { + let rem = read - dst1.len(); + received.extend_from_slice(&dst2[..rem]); + } + + // The bytes we received should be a prefix of the concatenation of a+b. + let expected = [a.as_slice(), b.as_slice()].concat(); + assert_eq!( + received, + expected[..received.len()], + "received bytes don't match expected prefix" + ); + + // Close fds to avoid leaking + unsafe { + let _ = libc::close(rfd); + let _ = libc::close(wfd); + } + }); + } +} diff --git a/vendor/vibeio/src/op/open.rs b/vendor/vibeio/src/op/open.rs new file mode 100644 index 0000000..48a7295 --- /dev/null +++ b/vendor/vibeio/src/op/open.rs @@ -0,0 +1,95 @@ +use std::ffi::CString; +use std::io; +use std::os::fd::RawFd; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub type OpenRawHandle = RawFd; + +pub struct OpenOp { + path: CString, + flags: i32, + mode: libc::mode_t, + completion_token: Option, +} + +impl OpenOp { + #[inline] + pub fn new(path: CString, flags: i32, mode: libc::mode_t) -> Self { + Self { + path, + flags, + mode, + completion_token: None, + } + } +} + +impl Op for OpenOp { + type Output = OpenRawHandle; + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(result as RawFd)) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::OpenAt::new(types::Fd(libc::AT_FDCWD), self.path.as_ptr()) + .flags(self.flags) + .mode(self.mode) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for OpenOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/read.rs b/vendor/vibeio/src/op/read.rs new file mode 100644 index 0000000..b895f99 --- /dev/null +++ b/vendor/vibeio/src/op/read.rs @@ -0,0 +1,306 @@ +use std::io; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_HANDLE_EOF, ERROR_IO_PENDING, HANDLE}, + Networking::WinSock::{self, SOCKET, WSABUF, WSA_IO_PENDING}, + Storage::FileSystem::ReadFile, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::op::io_util::{poll_result_or_wait, CompletionBuffer}; +use crate::op::Op; + +#[cfg(windows)] +#[inline] +fn socket_read(socket: SOCKET, buf: &mut [u8]) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let len = u32::try_from(buf.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_mut_ptr().cast(), + }; + let mut bytes: u32 = 0; + let mut flags: u32 = 0; + + let recv_result = unsafe { + WinSock::WSARecv( + socket, + &mut wsabuf, + 1, + &mut bytes, + &mut flags, + std::ptr::null_mut(), + None, + ) + }; + if recv_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +use crate::io::IoBufMut; + +pub struct ReadOp<'a, B: IoBufMut> { + handle: &'a InnerRawHandle, + buf: Option>, + completion_token: Option, + #[cfg(windows)] + socket_buf: Option>, +} + +impl<'a, B: IoBufMut> ReadOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + socket_buf: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for ReadOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_mut().unwrap().as_mut(); + + #[cfg(unix)] + let result = { + let read = unsafe { + libc::read( + self.handle.handle, + buf.as_buf_mut_ptr().cast::(), + buf.buf_capacity(), + ) + }; + if read == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(read as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => { + let slice = unsafe { + std::slice::from_raw_parts_mut(buf.as_buf_mut_ptr(), buf.buf_capacity()) + }; + socket_read(socket as SOCKET, slice) + } + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based read currently supports sockets only on Windows", + )), + }; + + match poll_result_or_wait(result, self.handle, cx, driver, Interest::READABLE) { + Poll::Ready(Ok(read)) => { + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok(read)) + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + #[cfg(windows)] + if -result == ERROR_HANDLE_EOF as i32 { + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(0) }; + return Poll::Ready(Ok(0)); + } + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let read = result as usize; + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok(read)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_mut().unwrap().as_mut(); + match self.handle.handle { + RawOsHandle::Socket(socket) => { + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let wsabuf = self.socket_buf.get_or_insert_with(|| { + Box::new(WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }) + }); + wsabuf.len = read_len; + wsabuf.buf = buf.as_buf_mut_ptr().cast(); + + let mut flags: u32 = 0; + let recv_result = unsafe { + WinSock::WSARecv( + socket as SOCKET, + wsabuf.as_mut() as *mut WSABUF, + 1, + std::ptr::null_mut(), + &mut flags, + overlapped, + None, + ) + }; + + if recv_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(err)) + } + } + RawOsHandle::Handle(handle) => { + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows file I/O", + ) + })?; + + let read_result = unsafe { + ReadFile( + handle as HANDLE, + buf.as_buf_mut_ptr().cast(), + read_len, + std::ptr::null_mut(), + overlapped, + ) + }; + + if read_result != 0 { + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + Ok(()) + } else { + Err(err) + } + } + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_mut().unwrap().as_mut(); + let entry = opcode::Read::new( + types::Fd(self.handle.handle), + buf.as_buf_mut_ptr(), + (buf.buf_capacity()) as _, + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for ReadOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(windows)] + let completion_state = self.socket_buf.take(); + #[cfg(not(windows))] + let completion_state = (); + + let ignored_data = Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )); + #[cfg(windows)] + driver.cancel_completion(completion_token, self.handle.handle, ignored_data); + #[cfg(not(windows))] + driver.ignore_completion(completion_token, ignored_data); + } + } + } +} diff --git a/vendor/vibeio/src/op/readat.rs b/vendor/vibeio/src/op/readat.rs new file mode 100644 index 0000000..72d5a60 --- /dev/null +++ b/vendor/vibeio/src/op/readat.rs @@ -0,0 +1,171 @@ +use std::io; +use std::task::{Context, Poll}; + +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_IO_PENDING, HANDLE}, + Storage::FileSystem::ReadFile, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBufMut; +use crate::op::io_util::CompletionBuffer; +use crate::op::Op; + +pub struct ReadAtOp<'a, B: IoBufMut> { + handle: &'a InnerRawHandle, + buf: Option>, + offset: u64, + completion_token: Option, +} + +impl<'a, B: IoBufMut> ReadAtOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B, offset: u64) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + offset, + completion_token: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for ReadAtOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let read = result as usize; + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok(read)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_mut().unwrap().as_mut(); + let RawOsHandle::Handle(handle) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "ReadAtOp expects a file handle, not a socket", + )); + }; + + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows file I/O", + ) + })?; + + unsafe { + (*overlapped).Anonymous.Anonymous.Offset = self.offset as u32; + (*overlapped).Anonymous.Anonymous.OffsetHigh = (self.offset >> 32) as u32; + } + + let read_result = unsafe { + ReadFile( + handle as HANDLE, + buf.as_buf_mut_ptr().cast(), + read_len, + std::ptr::null_mut(), + overlapped, + ) + }; + + if read_result != 0 { + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + Ok(()) + } else { + Err(err) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_mut().unwrap().as_mut(); + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for io_uring", + ) + })?; + + let entry = opcode::Read::new( + types::Fd(self.handle.handle), + buf.as_buf_mut_ptr(), + read_len, + ) + .offset(self.offset) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for ReadAtOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + if let Some(buf) = self.buf.take() { + driver.ignore_completion(completion_token, Box::new(buf.into_stable_box())); + } else { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } + } +} diff --git a/vendor/vibeio/src/op/readiness.rs b/vendor/vibeio/src/op/readiness.rs new file mode 100644 index 0000000..627c21b --- /dev/null +++ b/vendor/vibeio/src/op/readiness.rs @@ -0,0 +1,51 @@ +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::fd_inner::InnerRawHandle; +use crate::op::Op; + +pub struct ReadinessOp<'a> { + handle: &'a InnerRawHandle, + is_writable: bool, +} + +impl<'a> ReadinessOp<'a> { + #[inline] + pub fn new_readable(handle: &'a InnerRawHandle) -> Self { + Self { + handle, + is_writable: false, + } + } + + #[inline] + pub fn new_writable(handle: &'a InnerRawHandle) -> Self { + Self { + handle, + is_writable: true, + } + } +} + +impl Op for ReadinessOp<'_> { + type Output = (); + + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + driver.submit_poll( + self.handle, + cx.waker().clone(), + if self.is_writable { + mio::Interest::WRITABLE + } else { + mio::Interest::READABLE + }, + )?; + Poll::Pending + } +} diff --git a/vendor/vibeio/src/op/readv.rs b/vendor/vibeio/src/op/readv.rs new file mode 100644 index 0000000..65142b2 --- /dev/null +++ b/vendor/vibeio/src/op/readv.rs @@ -0,0 +1,366 @@ +use std::io; +#[cfg(unix)] +use std::mem::MaybeUninit; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_IO_PENDING, HANDLE}, + Networking::WinSock::{self as WinSock, SOCKET, WSABUF, WSA_IO_PENDING}, + Storage::FileSystem::ReadFile, + System::IO::OVERLAPPED, +}; + +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +#[cfg(unix)] +use crate::io::IoVec; +use crate::op::io_util::poll_result_or_wait; +use crate::op::Op; +use crate::{current_driver, driver::AnyDriver}; +use crate::{driver::CompletionIoResult, io::IoVectoredBufMut}; + +/// Converts a slice of `IoSlice` to a system iovec buffer. +#[cfg(unix)] +#[inline] +fn iovec_to_system(bufs: &mut [IoVec]) -> Box<[libc::iovec]> { + let mut iovecs_maybeuninit: Box<[MaybeUninit]> = Box::new_uninit_slice(bufs.len()); + for (index, s) in bufs.iter_mut().enumerate() { + let iov = libc::iovec { + iov_base: s.ptr.cast::(), + iov_len: s.len, + }; + iovecs_maybeuninit[index].write(iov); + } + // SAFETY: The boxed slice would have all values initialized after interating over original array + unsafe { iovecs_maybeuninit.assume_init() } +} + +#[cfg(windows)] +#[inline] +fn socket_read_vectored(socket: SOCKET, bufs: &mut B) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let iovecs = bufs.as_iovecs_mut(); + let mut wsabufs = Vec::with_capacity(iovecs.len()); + for iovec in iovecs { + let len = u32::try_from(iovec.len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "writev buffer is too large for Windows socket I/O", + ) + })?; + wsabufs.push(WSABUF { + len, + buf: iovec.ptr as *mut _, + }); + } + + let mut bytes: u32 = 0; + let mut flags: u32 = 0; + let recv_result = unsafe { + WinSock::WSARecv( + socket, + wsabufs.as_mut_ptr(), + wsabufs.len() as u32, + &mut bytes, + &mut flags, + std::ptr::null_mut(), + None, + ) + }; + if recv_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +pub struct ReadvOp<'a, B: IoVectoredBufMut> { + handle: &'a InnerRawHandle, + bufs: Option, + completion_token: Option, + #[cfg(windows)] + completion_wsabufs: Option>, + #[cfg(windows)] + completion_staging: Option>, + #[cfg(target_os = "linux")] + completion_system_iovecs: Option>, +} + +impl<'a, B: IoVectoredBufMut> ReadvOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, bufs: B) -> Self { + Self { + handle, + bufs: Some(bufs), + completion_token: None, + #[cfg(windows)] + completion_wsabufs: None, + #[cfg(windows)] + completion_staging: None, + #[cfg(target_os = "linux")] + completion_system_iovecs: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.bufs.take().unwrap() + } +} + +impl Op for ReadvOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let bufs = self.bufs.as_mut().unwrap(); + #[cfg(unix)] + let result = { + let mut iovecs = iovec_to_system(&mut bufs.as_iovecs_mut()); + let read = + unsafe { libc::readv(self.handle.handle, iovecs.as_mut_ptr(), iovecs.len() as _) }; + if read == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(read as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => socket_read_vectored(socket as SOCKET, bufs), + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based readv currently supports sockets only on Windows", + )), + }; + + poll_result_or_wait(result, self.handle, cx, driver, Interest::READABLE) + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + #[cfg(windows)] + { + self.completion_wsabufs = None; + self.completion_staging = None; + } + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + + #[cfg(windows)] + { + let bufs = self.bufs.as_mut().unwrap(); + if let Some(staging) = self.completion_staging.take() { + let mut src_offset = 0usize; + let mut remaining = result as usize; + let iovecs = bufs.as_iovecs_mut(); + for dst in iovecs { + if remaining == 0 { + break; + } + let chunk = remaining.min(dst.len); + if chunk == 0 { + continue; + } + + let dst_slice = unsafe { std::slice::from_raw_parts_mut(dst.ptr, dst.len) }; + dst_slice[..chunk].copy_from_slice(&staging[src_offset..src_offset + chunk]); + src_offset += chunk; + remaining -= chunk; + } + } + self.completion_wsabufs = None; + } + + Poll::Ready(Ok(result as usize)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let bufs = self.bufs.as_mut().unwrap(); + match self.handle.handle { + RawOsHandle::Socket(socket) => { + let iovecs = bufs.as_iovecs(); + let mut wsabufs = Vec::with_capacity(iovecs.len()); + for iovec in iovecs { + let len = u32::try_from(iovec.len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "writev buffer is too large for Windows socket I/O", + ) + })?; + wsabufs.push(WSABUF { + len, + buf: iovec.ptr as *mut _, + }); + } + + let mut wsabufs = wsabufs.into_boxed_slice(); + let mut flags = 0u32; + let recv_result = unsafe { + WinSock::WSARecv( + socket as SOCKET, + wsabufs.as_mut_ptr(), + wsabufs.len() as u32, + std::ptr::null_mut(), + &mut flags, + overlapped, + None, + ) + }; + + if recv_result == 0 { + self.completion_wsabufs = Some(wsabufs); + self.completion_staging = None; + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + self.completion_wsabufs = Some(wsabufs); + self.completion_staging = None; + Ok(()) + } else { + self.completion_wsabufs = None; + self.completion_staging = None; + Err(io::Error::from_raw_os_error(err)) + } + } + RawOsHandle::Handle(handle) => { + let iovecs = bufs.as_iovecs(); + let total_len = (iovecs.iter()).try_fold(0usize, |acc, iovec| { + acc.checked_add(iovec.len).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "writev buffer length overflow") + }) + })?; + let total_len_u32 = u32::try_from(total_len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "readv total length is too large for Windows file I/O", + ) + })?; + + let mut staging = vec![0u8; total_len]; + let read_result = unsafe { + ReadFile( + handle as HANDLE, + staging.as_mut_ptr().cast(), + total_len_u32, + std::ptr::null_mut(), + overlapped, + ) + }; + + if read_result != 0 { + self.completion_wsabufs = None; + self.completion_staging = Some(staging); + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + self.completion_wsabufs = None; + self.completion_staging = Some(staging); + Ok(()) + } else { + self.completion_wsabufs = None; + self.completion_staging = None; + Err(err) + } + } + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let bufs = self.bufs.as_mut().unwrap(); + + // Build a temporary iovec array for the syscall. + let mut iovecs = if let Some(iovecs) = self.completion_system_iovecs.take() { + iovecs + } else { + iovec_to_system(&mut bufs.as_iovecs_mut()) + }; + + let entry = opcode::Readv::new( + types::Fd(self.handle.handle), + iovecs.as_mut_ptr(), + iovecs.len() as _, + ) + .build() + .user_data(user_data); + + // Store the iovec array for the completion, because it needs to be kept alive until the + // completion is ready. + self.completion_system_iovecs = Some(iovecs); + + Ok(entry) + } +} + +impl Drop for ReadvOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = current_driver() { + #[cfg(target_os = "linux")] + let bufs = self.completion_system_iovecs.take(); + #[cfg(windows)] + let bufs = self.completion_wsabufs.take(); + #[cfg(not(any(target_os = "linux", windows)))] + let bufs = (); + + driver.ignore_completion(completion_token, Box::new((bufs, self.bufs.take()))); + } + } + } +} diff --git a/vendor/vibeio/src/op/recv.rs b/vendor/vibeio/src/op/recv.rs new file mode 100644 index 0000000..ee80cb1 --- /dev/null +++ b/vendor/vibeio/src/op/recv.rs @@ -0,0 +1,287 @@ +use std::io; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Networking::WinSock::{self, MSG_PEEK, SOCKET, WSABUF, WSA_IO_PENDING}, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBufMut; +use crate::op::io_util::{poll_result_or_wait, CompletionBuffer}; +use crate::op::Op; + +#[cfg(windows)] +#[inline] +fn socket_recv(socket: SOCKET, buf: &mut [u8], peek: bool) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{ + self as WinSock, MSG_PEEK, SOCKET_ERROR, WSABUF, + }; + + let len = u32::try_from(buf.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_mut_ptr().cast(), + }; + let mut bytes: u32 = 0; + let mut flags: u32 = if peek { MSG_PEEK as u32 } else { 0 }; + + let recv_result = unsafe { + WinSock::WSARecv( + socket, + &mut wsabuf, + 1, + &mut bytes, + &mut flags, + std::ptr::null_mut(), + None, + ) + }; + if recv_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +pub struct RecvOp<'a, B: IoBufMut> { + handle: &'a InnerRawHandle, + buf: Option>, + completion_token: Option, + #[cfg(windows)] + socket_buf: Option>, + peek: bool, +} + +impl<'a, B: IoBufMut> RecvOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + socket_buf: None, + peek: false, + } + } + + pub fn new_peek(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + socket_buf: None, + peek: true, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for RecvOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_mut().unwrap().as_mut(); + + #[cfg(unix)] + let result = { + let read = unsafe { + libc::recv( + self.handle.handle, + buf.as_buf_mut_ptr().cast::(), + buf.buf_capacity(), + if self.peek { libc::MSG_PEEK } else { 0 }, + ) + }; + if read == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(read as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => { + let slice = unsafe { + std::slice::from_raw_parts_mut(buf.as_buf_mut_ptr(), buf.buf_capacity()) + }; + socket_recv(socket as SOCKET, slice, self.peek) + } + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based recv currently supports sockets only on Windows", + )), + }; + match poll_result_or_wait(result, self.handle, cx, driver, Interest::READABLE) { + Poll::Ready(Ok(read)) => { + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok(read)) + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let read = result as usize; + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok(read)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_mut().unwrap().as_mut(); + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "WSARecv can be used only with listening sockets", + )); + }; + + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let wsabuf = self.socket_buf.get_or_insert_with(|| { + Box::new(WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }) + }); + wsabuf.len = read_len; + wsabuf.buf = buf.as_buf_mut_ptr().cast(); + + let mut flags: u32 = if self.peek { MSG_PEEK as u32 } else { 0 }; + let recv_result = unsafe { + WinSock::WSARecv( + socket as SOCKET, + wsabuf.as_mut() as *mut WSABUF, + 1, + std::ptr::null_mut(), + &mut flags, + overlapped, + None, + ) + }; + + if recv_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(err)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_mut().unwrap().as_mut(); + let entry = opcode::Recv::new( + types::Fd(self.handle.handle), + buf.as_buf_mut_ptr(), + (buf.buf_capacity()) as _, + ) + .flags(if self.peek { libc::MSG_PEEK } else { 0 }) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for RecvOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(windows)] + let completion_state = self.socket_buf.take(); + #[cfg(not(windows))] + let completion_state = (); + + let ignored_data = Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )); + #[cfg(windows)] + driver.cancel_completion(completion_token, self.handle.handle, ignored_data); + #[cfg(not(windows))] + driver.ignore_completion(completion_token, ignored_data); + } + } + } +} diff --git a/vendor/vibeio/src/op/recvfrom.rs b/vendor/vibeio/src/op/recvfrom.rs new file mode 100644 index 0000000..0d323c9 --- /dev/null +++ b/vendor/vibeio/src/op/recvfrom.rs @@ -0,0 +1,458 @@ +use std::io; +#[cfg(unix)] +use std::mem::MaybeUninit; +use std::net::SocketAddr; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Networking::WinSock::{ + self as WinSock, AF_INET, AF_INET6, MSG_PEEK, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6, + SOCKADDR_STORAGE, SOCKET, WSABUF, WSA_IO_PENDING, + }, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBufMut; +use crate::op::io_util::CompletionBuffer; +use crate::op::Op; + +#[cfg(unix)] +#[inline] +fn sockaddr_storage_to_socketaddr( + storage: &libc::sockaddr_storage, +) -> Result { + let family = storage.ss_family as libc::c_int; + + if family == libc::AF_INET { + let addr_in: &libc::sockaddr_in = + unsafe { &*(storage as *const _ as *const libc::sockaddr_in) }; + let port = u16::from_be(addr_in.sin_port); + let ip_u32 = u32::from_be(addr_in.sin_addr.s_addr); + let ip = std::net::Ipv4Addr::from(ip_u32); + Ok(SocketAddr::V4(std::net::SocketAddrV4::new(ip, port))) + } else if family == libc::AF_INET6 { + let addr_in6: &libc::sockaddr_in6 = + unsafe { &*(storage as *const _ as *const libc::sockaddr_in6) }; + let port = u16::from_be(addr_in6.sin6_port); + let ip = std::net::Ipv6Addr::from(addr_in6.sin6_addr.s6_addr); + Ok(SocketAddr::V6(std::net::SocketAddrV6::new( + ip, + port, + addr_in6.sin6_flowinfo, + addr_in6.sin6_scope_id, + ))) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported socket family", + )) + } +} + +#[cfg(windows)] +#[inline] +fn sockaddr_storage_to_socketaddr(storage: &SOCKADDR_STORAGE) -> Result { + let family = storage.ss_family; + + if family == AF_INET { + let addr_in: &SOCKADDR_IN = unsafe { &*(storage as *const _ as *const SOCKADDR_IN) }; + let port = u16::from_be(addr_in.sin_port); + let ip_u32 = u32::from_be(unsafe { addr_in.sin_addr.S_un.S_addr }); + let ip = std::net::Ipv4Addr::from(ip_u32); + Ok(SocketAddr::V4(std::net::SocketAddrV4::new(ip, port))) + } else if family == AF_INET6 { + let addr_in6: &SOCKADDR_IN6 = unsafe { &*(storage as *const _ as *const SOCKADDR_IN6) }; + let port = u16::from_be(addr_in6.sin6_port); + let ip = std::net::Ipv6Addr::from(unsafe { addr_in6.sin6_addr.u.Byte }); + Ok(SocketAddr::V6(std::net::SocketAddrV6::new( + ip, + port, + addr_in6.sin6_flowinfo, + unsafe { addr_in6.Anonymous.sin6_scope_id }, + ))) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "unsupported socket family", + )) + } +} + +#[cfg(windows)] +#[inline] +fn socket_recvfrom(socket: SOCKET, buf: &mut [u8], peek: bool) -> io::Result<(usize, SocketAddr)> { + use windows_sys::Win32::Networking::WinSock::{ + self as WinSock, MSG_PEEK, SOCKADDR, SOCKADDR_STORAGE, SOCKET_ERROR, WSABUF, + }; + + let len = u32::try_from(buf.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_mut_ptr().cast(), + }; + let mut bytes: u32 = 0; + let mut flags: u32 = if peek { MSG_PEEK as u32 } else { 0 }; + let mut addr = SOCKADDR_STORAGE::default(); + let mut addr_len = std::mem::size_of::() as i32; + + let recv_result = unsafe { + WinSock::WSARecvFrom( + socket, + &mut wsabuf, + 1, + &mut bytes, + &mut flags, + (&mut addr as *mut SOCKADDR_STORAGE).cast::(), + &mut addr_len, + std::ptr::null_mut(), + None, + ) + }; + if recv_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + let address = sockaddr_storage_to_socketaddr(&addr)?; + Ok((bytes as usize, address)) +} + +#[cfg(windows)] +struct RecvfromWindowsCompletion { + socket_buf: WSABUF, + addr: SOCKADDR_STORAGE, + addr_len: i32, + flags: u32, +} + +#[cfg(target_os = "linux")] +struct RecvfromLinuxCompletion { + addr: libc::sockaddr_storage, + iovec: libc::iovec, + msghdr: libc::msghdr, +} + +pub struct RecvfromOp<'a, B: IoBufMut> { + handle: &'a InnerRawHandle, + buf: Option>, + completion_token: Option, + #[cfg(windows)] + completion_state: Option>, + #[cfg(target_os = "linux")] + completion_state: Option>, + peek: bool, +} + +impl<'a, B: IoBufMut> RecvfromOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + completion_state: None, + #[cfg(target_os = "linux")] + completion_state: None, + peek: false, + } + } + + #[inline] + pub fn new_peek(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + completion_state: None, + #[cfg(target_os = "linux")] + completion_state: None, + peek: true, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for RecvfromOp<'_, B> { + type Output = (usize, SocketAddr); + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_mut().unwrap().as_mut(); + + #[cfg(unix)] + let result = { + let mut addr = MaybeUninit::::zeroed(); + let mut addr_len = std::mem::size_of::() as libc::socklen_t; + let read = unsafe { + libc::recvfrom( + self.handle.handle, + buf.as_buf_mut_ptr().cast::(), + buf.buf_capacity(), + if self.peek { libc::MSG_PEEK } else { 0 }, + addr.as_mut_ptr().cast::(), + &mut addr_len, + ) + }; + + if read == -1 { + Err(io::Error::last_os_error()) + } else { + let address = sockaddr_storage_to_socketaddr(unsafe { &addr.assume_init() })?; + Ok((read as usize, address)) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => { + let slice = unsafe { + std::slice::from_raw_parts_mut(buf.as_buf_mut_ptr(), buf.buf_capacity()) + }; + socket_recvfrom(socket as SOCKET, slice, self.peek) + } + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based recvfrom currently supports sockets only on Windows", + )), + }; + + match result { + Ok((read, address)) => { + unsafe { buf.set_buf_init(read) }; + Poll::Ready(Ok((read, address))) + } + Err(err) if err.kind() == io::ErrorKind::WouldBlock => { + match driver.submit_poll(self.handle, cx.waker().clone(), Interest::READABLE) { + Ok(_) => Poll::Pending, + Err(err) => Poll::Ready(Err(err)), + } + } + Err(err) => Poll::Ready(Err(err)), + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let read = result as usize; + + #[cfg(target_os = "linux")] + { + let address = self + .completion_state + .as_ref() + .ok_or_else(|| io::Error::other("recvfrom completion missing source address")) + .and_then(|state| sockaddr_storage_to_socketaddr(&state.addr)); + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(read) }; + Poll::Ready(address.map(|address| (read, address))) + } + + #[cfg(all(unix, not(target_os = "linux")))] + { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::Unsupported, + "completion-based recvfrom is unsupported on this Unix platform", + ))); + } + + #[cfg(windows)] + { + let address = self + .completion_state + .as_ref() + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::Other, + "recvfrom completion missing source address", + ) + }) + .and_then(|state| sockaddr_storage_to_socketaddr(&state.addr)); + let buf = self.buf.as_mut().unwrap().as_mut(); + unsafe { buf.set_buf_init(read) }; + return Poll::Ready(address.map(|address| (read, address))); + } + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_mut().unwrap().as_mut(); + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "WSARecvFrom can be used only with sockets", + )); + }; + + let read_len = u32::try_from(buf.buf_capacity()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "read buffer is too large for Windows socket I/O", + ) + })?; + + let completion = self.completion_state.get_or_insert_with(|| { + Box::new(RecvfromWindowsCompletion { + socket_buf: WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }, + addr: SOCKADDR_STORAGE::default(), + addr_len: 0, + flags: 0, + }) + }); + completion.socket_buf.len = read_len; + completion.socket_buf.buf = buf.as_buf_mut_ptr().cast(); + completion.addr_len = std::mem::size_of::() as i32; + completion.flags = if self.peek { MSG_PEEK as u32 } else { 0 }; + + let recv_result = unsafe { + WinSock::WSARecvFrom( + socket as SOCKET, + &mut completion.socket_buf as *mut WSABUF, + 1, + std::ptr::null_mut(), + &mut completion.flags, + (&mut completion.addr as *mut SOCKADDR_STORAGE).cast::(), + &mut completion.addr_len, + overlapped, + None, + ) + }; + + if recv_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + self.completion_state = None; + Err(io::Error::from_raw_os_error(err)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_mut().unwrap().as_mut(); + let completion = self.completion_state.get_or_insert_with(|| { + Box::new(RecvfromLinuxCompletion { + addr: unsafe { std::mem::zeroed() }, + iovec: libc::iovec { + iov_base: std::ptr::null_mut(), + iov_len: 0, + }, + msghdr: unsafe { std::mem::zeroed() }, + }) + }); + completion.addr = unsafe { std::mem::zeroed() }; + completion.iovec = libc::iovec { + iov_base: buf.as_buf_mut_ptr().cast::(), + iov_len: buf.buf_capacity(), + }; + completion.msghdr = unsafe { std::mem::zeroed::() }; + completion.msghdr.msg_name = + &mut completion.addr as *mut libc::sockaddr_storage as *mut libc::c_void; + completion.msghdr.msg_namelen = + std::mem::size_of::() as libc::socklen_t; + completion.msghdr.msg_iov = &mut completion.iovec as *mut libc::iovec; + completion.msghdr.msg_iovlen = 1; + completion.msghdr.msg_control = std::ptr::null_mut(); + completion.msghdr.msg_controllen = 0; + completion.msghdr.msg_flags = 0; + + let entry = opcode::RecvMsg::new( + types::Fd(self.handle.handle), + &mut completion.msghdr as *mut libc::msghdr, + ) + .flags(if self.peek { libc::MSG_PEEK as u32 } else { 0 }) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for RecvfromOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(any(windows, target_os = "linux"))] + let completion_state = self.completion_state.take(); + #[cfg(not(any(windows, target_os = "linux")))] + let completion_state = (); + + driver.ignore_completion( + completion_token, + Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )), + ); + } + } + } +} diff --git a/vendor/vibeio/src/op/rename.rs b/vendor/vibeio/src/op/rename.rs new file mode 100644 index 0000000..a53e90f --- /dev/null +++ b/vendor/vibeio/src/op/rename.rs @@ -0,0 +1,93 @@ +use std::ffi::CString; +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct RenameOp { + old_path: CString, + new_path: CString, + completion_token: Option, +} + +impl RenameOp { + #[inline] + pub fn new(old_path: CString, new_path: CString) -> Self { + Self { + old_path, + new_path, + completion_token: None, + } + } +} + +impl Op for RenameOp { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::RenameAt::new( + types::Fd(libc::AT_FDCWD), + self.old_path.as_ptr(), + types::Fd(libc::AT_FDCWD), + self.new_path.as_ptr(), + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for RenameOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/send.rs b/vendor/vibeio/src/op/send.rs new file mode 100644 index 0000000..0e4fad0 --- /dev/null +++ b/vendor/vibeio/src/op/send.rs @@ -0,0 +1,256 @@ +use std::io; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Networking::WinSock::{self as WinSock, SOCKET, WSABUF, WSA_IO_PENDING}, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBuf; +use crate::op::io_util::{poll_result_or_wait, CompletionBuffer}; +use crate::op::Op; + +#[cfg(windows)] +#[inline] +fn socket_send(socket: SOCKET, buf: &B) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_buf_ptr().cast_mut().cast(), + }; + let mut bytes: u32 = 0; + + let send_result = unsafe { + WinSock::WSASend( + socket, + &mut wsabuf, + 1, + &mut bytes, + 0, + std::ptr::null_mut(), + None, + ) + }; + if send_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +pub struct SendOp<'a, B: IoBuf> { + handle: &'a InnerRawHandle, + buf: Option>, + completion_token: Option, + #[cfg(windows)] + socket_buf: Option>, +} + +impl<'a, B: IoBuf> SendOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + socket_buf: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for SendOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_ref().unwrap().as_ref(); + + #[cfg(unix)] + let result = { + let written = unsafe { + libc::send( + self.handle.handle, + buf.as_buf_ptr().cast::(), + buf.buf_len(), + 0, + ) + }; + if written == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(written as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => socket_send(socket as SOCKET, buf), + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based send currently supports sockets only on Windows", + )), + }; + + match poll_result_or_wait(result, self.handle, cx, driver, Interest::WRITABLE) { + Poll::Ready(Ok(written)) => Poll::Ready(Ok(written)), + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let written = result as usize; + Poll::Ready(Ok(written)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_ref().unwrap().as_ref(); + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "WSASend can be used only with sockets", + )); + }; + + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let wsabuf = self.socket_buf.get_or_insert_with(|| { + Box::new(WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }) + }); + wsabuf.len = write_len; + wsabuf.buf = buf.as_buf_ptr().cast_mut().cast(); + + let send_result = unsafe { + WinSock::WSASend( + socket as SOCKET, + wsabuf.as_mut() as *mut WSABUF, + 1, + std::ptr::null_mut(), + 0, + overlapped, + None, + ) + }; + + if send_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(err)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_ref().unwrap().as_ref(); + let entry = opcode::Send::new( + types::Fd(self.handle.handle), + buf.as_buf_ptr(), + buf.buf_len() as _, + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for SendOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(windows)] + let completion_state = self.socket_buf.take(); + #[cfg(not(windows))] + let completion_state = (); + + driver.ignore_completion( + completion_token, + Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )), + ); + } + } + } +} diff --git a/vendor/vibeio/src/op/sendto.rs b/vendor/vibeio/src/op/sendto.rs new file mode 100644 index 0000000..abb2a7d --- /dev/null +++ b/vendor/vibeio/src/op/sendto.rs @@ -0,0 +1,442 @@ +use std::io; +#[cfg(unix)] +use std::mem::MaybeUninit; +use std::net::SocketAddr; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Networking::WinSock::{ + self as WinSock, AF_INET, AF_INET6, SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6, SOCKADDR_STORAGE, + SOCKET, WSABUF, WSA_IO_PENDING, + }, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBuf; +use crate::op::io_util::{poll_result_or_wait, CompletionBuffer}; +use crate::op::Op; + +#[cfg(unix)] +#[inline] +fn socket_addr_to_raw(address: SocketAddr) -> (libc::sockaddr_storage, libc::socklen_t) { + match address { + SocketAddr::V4(address) => { + let sockaddr = libc::sockaddr_in { + sin_family: libc::AF_INET as libc::sa_family_t, + sin_port: address.port().to_be(), + sin_addr: libc::in_addr { + s_addr: u32::from_ne_bytes(address.ip().octets()), + }, + sin_zero: [0; 8], + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + SocketAddr::V6(address) => { + let sockaddr = libc::sockaddr_in6 { + sin6_family: libc::AF_INET6 as libc::sa_family_t, + sin6_port: address.port().to_be(), + sin6_flowinfo: address.flowinfo(), + sin6_addr: libc::in6_addr { + s6_addr: address.ip().octets(), + }, + sin6_scope_id: address.scope_id(), + #[cfg(any( + target_os = "macos", + target_os = "ios", + target_os = "freebsd", + target_os = "openbsd", + target_os = "dragonfly", + target_os = "netbsd", + target_os = "haiku", + target_os = "aix", + ))] + sin6_len: 0, + }; + + let mut storage = MaybeUninit::::zeroed(); + unsafe { + storage + .as_mut_ptr() + .cast::() + .write(sockaddr); + ( + storage.assume_init(), + std::mem::size_of::() as libc::socklen_t, + ) + } + } + } +} + +#[cfg(windows)] +#[inline] +fn socket_addr_to_raw(address: SocketAddr) -> (SOCKADDR_STORAGE, i32) { + match address { + SocketAddr::V4(address) => { + let mut sockaddr = SOCKADDR_IN::default(); + sockaddr.sin_family = AF_INET; + sockaddr.sin_port = address.port().to_be(); + sockaddr.sin_addr.S_un.S_addr = u32::from_ne_bytes(address.ip().octets()); + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + (storage, std::mem::size_of::() as i32) + } + SocketAddr::V6(address) => { + let mut sockaddr = SOCKADDR_IN6::default(); + sockaddr.sin6_family = AF_INET6; + sockaddr.sin6_port = address.port().to_be(); + sockaddr.sin6_flowinfo = address.flowinfo(); + sockaddr.sin6_addr.u.Byte = address.ip().octets(); + sockaddr.Anonymous.sin6_scope_id = address.scope_id() as u32; + + let mut storage = SOCKADDR_STORAGE::default(); + unsafe { + std::ptr::copy_nonoverlapping( + &sockaddr as *const SOCKADDR_IN6 as *const u8, + &mut storage as *mut SOCKADDR_STORAGE as *mut u8, + std::mem::size_of::(), + ); + } + (storage, std::mem::size_of::() as i32) + } + } +} + +#[cfg(windows)] +#[inline] +fn socket_sendto(socket: SOCKET, buf: &B, addr: SocketAddr) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_buf_ptr().cast_mut().cast(), + }; + let (raw_addr, raw_addr_len) = socket_addr_to_raw(addr); + let mut bytes: u32 = 0; + + let send_result = unsafe { + WinSock::WSASendTo( + socket, + &mut wsabuf, + 1, + &mut bytes, + 0, + (&raw_addr as *const SOCKADDR_STORAGE).cast::(), + raw_addr_len, + std::ptr::null_mut(), + None, + ) + }; + if send_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +#[cfg(windows)] +struct SendtoWindowsCompletion { + socket_buf: WSABUF, + addr: SOCKADDR_STORAGE, + addr_len: i32, +} + +#[cfg(target_os = "linux")] +struct SendtoLinuxCompletion { + addr: libc::sockaddr_storage, + addr_len: libc::socklen_t, + iovec: libc::iovec, + msghdr: libc::msghdr, +} + +pub struct SendtoOp<'a, B: IoBuf> { + handle: &'a InnerRawHandle, + buf: Option>, + addr: SocketAddr, + completion_token: Option, + #[cfg(windows)] + completion_state: Option>, + #[cfg(target_os = "linux")] + completion_state: Option>, +} + +impl<'a, B: IoBuf> SendtoOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B, addr: SocketAddr) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + addr, + completion_token: None, + #[cfg(windows)] + completion_state: None, + #[cfg(target_os = "linux")] + completion_state: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for SendtoOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_ref().unwrap().as_ref(); + + #[cfg(unix)] + let result = { + let (raw_addr, raw_addr_len) = socket_addr_to_raw(self.addr); + let written = unsafe { + libc::sendto( + self.handle.handle, + buf.as_buf_ptr().cast::(), + buf.buf_len(), + 0, + (&raw_addr as *const libc::sockaddr_storage).cast::(), + raw_addr_len, + ) + }; + if written == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(written as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => socket_sendto(socket as SOCKET, buf, self.addr), + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based sendto currently supports sockets only on Windows", + )), + }; + + match poll_result_or_wait(result, self.handle, cx, driver, Interest::WRITABLE) { + Poll::Ready(Ok(written)) => Poll::Ready(Ok(written)), + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let written = result as usize; + Poll::Ready(Ok(written)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_ref().unwrap().as_ref(); + let RawOsHandle::Socket(socket) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "WSASendTo can be used only with sockets", + )); + }; + + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let (raw_addr, raw_addr_len) = socket_addr_to_raw(self.addr); + let completion = self.completion_state.get_or_insert_with(|| { + Box::new(SendtoWindowsCompletion { + socket_buf: WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }, + addr: SOCKADDR_STORAGE::default(), + addr_len: 0, + }) + }); + completion.socket_buf.len = write_len; + completion.socket_buf.buf = buf.as_buf_ptr().cast_mut().cast(); + completion.addr = raw_addr; + completion.addr_len = raw_addr_len; + + let send_result = unsafe { + WinSock::WSASendTo( + socket as SOCKET, + &mut completion.socket_buf as *mut WSABUF, + 1, + std::ptr::null_mut(), + 0, + (&completion.addr as *const SOCKADDR_STORAGE).cast::(), + completion.addr_len, + overlapped, + None, + ) + }; + + if send_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + self.completion_state = None; + Err(io::Error::from_raw_os_error(err)) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let (raw_addr, raw_addr_len) = socket_addr_to_raw(self.addr); + let buf = self.buf.as_ref().unwrap().as_ref(); + let completion = self.completion_state.get_or_insert_with(|| { + Box::new(SendtoLinuxCompletion { + addr: unsafe { std::mem::zeroed() }, + addr_len: 0, + iovec: libc::iovec { + iov_base: std::ptr::null_mut(), + iov_len: 0, + }, + msghdr: unsafe { std::mem::zeroed() }, + }) + }); + completion.addr = raw_addr; + completion.addr_len = raw_addr_len; + completion.iovec = libc::iovec { + iov_base: buf.as_buf_ptr().cast_mut().cast::(), + iov_len: buf.buf_len(), + }; + + completion.msghdr = unsafe { std::mem::zeroed::() }; + completion.msghdr.msg_name = + &mut completion.addr as *mut libc::sockaddr_storage as *mut libc::c_void; + completion.msghdr.msg_namelen = completion.addr_len; + completion.msghdr.msg_iov = &mut completion.iovec as *mut libc::iovec; + completion.msghdr.msg_iovlen = 1; + completion.msghdr.msg_control = std::ptr::null_mut(); + completion.msghdr.msg_controllen = 1; + completion.msghdr.msg_flags = 1; + + let entry = opcode::SendMsg::new( + types::Fd(self.handle.handle), + &completion.msghdr as *const libc::msghdr, + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for SendtoOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(any(windows, target_os = "linux"))] + let completion_state = self.completion_state.take(); + #[cfg(not(any(windows, target_os = "linux")))] + let completion_state = (); + + driver.ignore_completion( + completion_token, + Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )), + ); + } + } + } +} diff --git a/vendor/vibeio/src/op/splice.rs b/vendor/vibeio/src/op/splice.rs new file mode 100644 index 0000000..88c6bef --- /dev/null +++ b/vendor/vibeio/src/op/splice.rs @@ -0,0 +1,127 @@ +use std::io; +use std::os::fd::{AsRawFd, BorrowedFd}; +use std::task::{Context, Poll}; + +use mio::Interest; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +use crate::op::io_util::poll_result_or_wait; +use crate::op::Op; + +pub struct SpliceOp<'a> { + fd_in: BorrowedFd<'a>, + fd_out: &'a InnerRawHandle, + len: usize, + completion_token: Option, +} + +impl<'a> SpliceOp<'a> { + #[inline] + pub fn new(fd_in: BorrowedFd<'a>, fd_out: &'a InnerRawHandle, len: usize) -> Self { + Self { + fd_in, + fd_out, + len, + completion_token: None, + } + } +} + +impl Op for SpliceOp<'_> { + type Output = usize; + + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = { + let returned = unsafe { + libc::splice( + self.fd_in.as_raw_fd(), + std::ptr::null_mut(), + self.fd_out.handle, + std::ptr::null_mut(), + self.len, + libc::SPLICE_F_NONBLOCK, + ) + }; + if returned == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(returned as usize) + } + }; + + poll_result_or_wait(result, self.fd_out, cx, driver, Interest::WRITABLE) + } + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(result as usize)) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::Splice::new( + types::Fd(self.fd_in.as_raw_fd()), + -1, + types::Fd(self.fd_out.handle), + -1, + self.len as u32, + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for SpliceOp<'_> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/statx.rs b/vendor/vibeio/src/op/statx.rs new file mode 100644 index 0000000..ce4dbc7 --- /dev/null +++ b/vendor/vibeio/src/op/statx.rs @@ -0,0 +1,119 @@ +use std::ffi::CString; +use std::io; +use std::mem::MaybeUninit; +use std::task::{Context, Poll}; + +use crate::current_driver; +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct StatxOp { + dirfd: libc::c_int, + pathname: CString, + flags: libc::c_int, + mask: libc::c_uint, + statxbuf: Option>>, + completion_token: Option, +} + +impl StatxOp { + #[inline] + pub fn new( + dirfd: libc::c_int, + pathname: CString, + flags: libc::c_int, + mask: libc::c_uint, + ) -> Self { + Self { + dirfd, + pathname, + flags, + mask, + statxbuf: None, + completion_token: None, + } + } +} + +impl Op for StatxOp { + type Output = libc::statx; + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + // SAFETY: kernel fills the statx struct on success. + let statxbuf = self.statxbuf.take().expect("statxbuf is None"); + let st = unsafe { *statxbuf.assume_init() }; + Poll::Ready(Ok(st)) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let mut statxbuf = if let Some(statxbuf) = self.statxbuf.take() { + statxbuf + } else { + Box::new_uninit() + }; + + let entry = opcode::Statx::new( + types::Fd(self.dirfd), + self.pathname.as_ptr(), + statxbuf.as_mut_ptr() as *mut types::statx, + ) + .flags(self.flags as _) + .mask(self.mask) + .build() + .user_data(user_data); + + self.statxbuf = Some(statxbuf); + + Ok(entry) + } +} + +impl Drop for StatxOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = current_driver() { + let statxbuf = self.statxbuf.take(); + driver.ignore_completion(completion_token, Box::new(statxbuf)); + } + } + } +} diff --git a/vendor/vibeio/src/op/symlink.rs b/vendor/vibeio/src/op/symlink.rs new file mode 100644 index 0000000..552a0f4 --- /dev/null +++ b/vendor/vibeio/src/op/symlink.rs @@ -0,0 +1,92 @@ +use std::ffi::CString; +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct SymlinkOp { + old_path: CString, + new_path: CString, + completion_token: Option, +} + +impl SymlinkOp { + #[inline] + pub fn new(old_path: CString, new_path: CString) -> Self { + Self { + old_path, + new_path, + completion_token: None, + } + } +} + +impl Op for SymlinkOp { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let entry = opcode::SymlinkAt::new( + types::Fd(libc::AT_FDCWD), + self.old_path.as_ptr(), + self.new_path.as_ptr(), + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for SymlinkOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/unlink.rs b/vendor/vibeio/src/op/unlink.rs new file mode 100644 index 0000000..de1ee45 --- /dev/null +++ b/vendor/vibeio/src/op/unlink.rs @@ -0,0 +1,90 @@ +use std::ffi::CString; +use std::io; +use std::task::{Context, Poll}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::op::Op; + +pub struct UnlinkOp { + path: CString, + completion_token: Option, + is_dir: bool, +} + +impl UnlinkOp { + #[inline] + pub fn new(path: CString, is_dir: bool) -> Self { + Self { + path, + completion_token: None, + is_dir, + } + } +} + +impl Op for UnlinkOp { + type Output = (); + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + + if result < 0 { + Poll::Ready(Err(io::Error::from_raw_os_error(-result))) + } else { + Poll::Ready(Ok(())) + } + } + + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + // AT_Unlink flag is passed as the flags parameter to unlinkat + let entry = opcode::UnlinkAt::new(types::Fd(libc::AT_FDCWD), self.path.as_ptr()) + .flags(if self.is_dir { libc::AT_REMOVEDIR } else { 0 }) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for UnlinkOp { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } +} diff --git a/vendor/vibeio/src/op/waitpid.rs b/vendor/vibeio/src/op/waitpid.rs new file mode 100644 index 0000000..29140f5 --- /dev/null +++ b/vendor/vibeio/src/op/waitpid.rs @@ -0,0 +1,233 @@ +use std::io; +use std::os::fd::RawFd; +use std::task::{Context, Poll}; + +use mio::Interest; + +use crate::driver::AnyDriver; +use crate::fd_inner::InnerRawHandle; +use crate::op::Op; + +/// Wraps a raw pidfd and closes it on drop. +struct OwnedPidFd(RawFd); + +impl Drop for OwnedPidFd { + #[inline] + fn drop(&mut self) { + if self.0 >= 0 { + unsafe { + libc::close(self.0); + } + } + } +} + +/// State machine for the pidfd-based waitpid operation. +enum WaitPidState { + /// Initial state: we have the child PID but haven't opened the pidfd yet. + Init { pid: libc::pid_t }, + /// The pidfd is open and registered with the driver; waiting for readability. + Polling { + pid: libc::pid_t, + pidfd: OwnedPidFd, + handle: InnerRawHandle, + }, + /// Terminal state after the result has been consumed. + Done, +} + +/// An async operation that waits for a child process to exit using Linux's +/// `pidfd_open(2)`. The pidfd is registered with the I/O driver so the +/// executor is woken only when the child actually terminates — no signals, +/// no busy-polling. +pub struct WaitPidOp { + state: WaitPidState, +} + +impl WaitPidOp { + /// Create a new `WaitPidOp` for the given child PID. + #[inline] + pub fn new(pid: u32) -> Self { + Self { + state: WaitPidState::Init { + pid: pid as libc::pid_t, + }, + } + } + + /// Open a pidfd for `pid` via the `pidfd_open` syscall (Linux 5.3+). + #[inline] + fn open_pidfd(pid: libc::pid_t) -> io::Result { + let fd = unsafe { libc::syscall(libc::SYS_pidfd_open, pid, 0 as libc::c_uint) }; + if fd < 0 { + Err(io::Error::last_os_error()) + } else { + let raw = fd as RawFd; + // Set non-blocking + close-on-exec atomically via a single fcntl + // call chain. O_NONBLOCK is required so that read() on the pidfd + // returns EAGAIN instead of blocking when the child hasn't exited + // yet, and FD_CLOEXEC prevents leaking across exec. + let flags = unsafe { libc::fcntl(raw, libc::F_GETFL) }; + if flags != -1 { + unsafe { + libc::fcntl(raw, libc::F_SETFL, flags | libc::O_NONBLOCK); + } + } + let fd_flags = unsafe { libc::fcntl(raw, libc::F_GETFD) }; + if fd_flags != -1 { + unsafe { + libc::fcntl(raw, libc::F_SETFD, fd_flags | libc::FD_CLOEXEC); + } + } + Ok(raw) + } + } + + /// Reap the child with `waitpid` after the pidfd signalled readability. + #[inline] + fn reap(pid: libc::pid_t) -> io::Result { + let mut status: libc::c_int = 0; + let rc = unsafe { libc::waitpid(pid, &mut status, libc::WNOHANG) }; + if rc < 0 { + return Err(io::Error::last_os_error()); + } + if rc == 0 { + // Shouldn't happen after pidfd is readable, but be defensive. + return Err(io::Error::new( + io::ErrorKind::WouldBlock, + "child has not exited yet", + )); + } + Ok(status) + } +} + +impl Op for WaitPidOp { + type Output = i32; // raw waitpid status + + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + loop { + match &self.state { + WaitPidState::Init { pid } => { + let pid = *pid; + let raw_fd = Self::open_pidfd(pid)?; + let pidfd = OwnedPidFd(raw_fd); + let handle = InnerRawHandle::new_with_mode( + raw_fd, + Interest::READABLE, + crate::driver::RegistrationMode::Poll, + )?; + self.state = WaitPidState::Polling { pid, pidfd, handle }; + // Fall through to the Polling arm. + } + WaitPidState::Polling { + pid, + pidfd: _, + handle, + } => { + // Try a non-blocking read on the pidfd. If it's readable the + // child has exited. + let mut buf = [0u8; 8]; + let n = unsafe { + libc::read( + handle.handle, + buf.as_mut_ptr().cast::(), + buf.len(), + ) + }; + if n < 0 { + let err = io::Error::last_os_error(); + if err.kind() == io::ErrorKind::WouldBlock { + // Not ready yet — register for readability and pend. + if let Err(submit_err) = + driver.submit_poll(handle, cx.waker().clone(), Interest::READABLE) + { + return Poll::Ready(Err(submit_err)); + } + return Poll::Pending; + } + // EBADF or similar — try reaping anyway + } + + let pid = *pid; + let status = Self::reap(pid); + self.state = WaitPidState::Done; + return Poll::Ready(status); + } + WaitPidState::Done => { + return Poll::Ready(Err(io::Error::other("WaitPidOp already completed"))); + } + } + } + } + + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + // We need to take ownership to transition states. Use a temporary Done. + match std::mem::replace(&mut self.state, WaitPidState::Done) { + WaitPidState::Init { pid } => { + let raw_fd = Self::open_pidfd(pid)?; + let pidfd = OwnedPidFd(raw_fd); + // For io_uring we register in Poll mode so we can use PollAdd + // to wait for readability on the pidfd. + let handle = InnerRawHandle::new_with_mode( + raw_fd, + Interest::READABLE, + crate::driver::RegistrationMode::Poll, + )?; + // Submit the poll via the driver (which uses io_uring PollAdd + // on uring, or mio on poll-based). + if let Err(submit_err) = + driver.submit_poll(&handle, cx.waker().clone(), Interest::READABLE) + { + self.state = WaitPidState::Polling { pid, pidfd, handle }; + return Poll::Ready(Err(submit_err)); + } + self.state = WaitPidState::Polling { pid, pidfd, handle }; + Poll::Pending + } + WaitPidState::Polling { pid, pidfd, handle } => { + // Check if the pidfd is readable (child exited). + let mut buf = [0u8; 8]; + let n = unsafe { + libc::read( + handle.handle, + buf.as_mut_ptr().cast::(), + buf.len(), + ) + }; + if n < 0 { + let err = io::Error::last_os_error(); + if err.kind() == io::ErrorKind::WouldBlock { + // Re-arm poll. + if let Err(submit_err) = + driver.submit_poll(&handle, cx.waker().clone(), Interest::READABLE) + { + self.state = WaitPidState::Polling { pid, pidfd, handle }; + return Poll::Ready(Err(submit_err)); + } + self.state = WaitPidState::Polling { pid, pidfd, handle }; + return Poll::Pending; + } + // Other error — try reaping anyway. + } + + let status = Self::reap(pid); + drop(handle); + drop(pidfd); + self.state = WaitPidState::Done; + Poll::Ready(status) + } + WaitPidState::Done => Poll::Ready(Err(io::Error::other("WaitPidOp already completed"))), + } + } +} diff --git a/vendor/vibeio/src/op/write.rs b/vendor/vibeio/src/op/write.rs new file mode 100644 index 0000000..4a07b06 --- /dev/null +++ b/vendor/vibeio/src/op/write.rs @@ -0,0 +1,289 @@ +use std::io; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_IO_PENDING, HANDLE}, + Networking::WinSock::{self, SOCKET, WSABUF, WSA_IO_PENDING}, + Storage::FileSystem::WriteFile, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::op::io_util::{poll_result_or_wait, CompletionBuffer}; +use crate::op::Op; + +#[cfg(windows)] +#[inline] +fn socket_write(socket: SOCKET, buf: &[u8]) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let len = u32::try_from(buf.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let mut wsabuf = WSABUF { + len, + buf: buf.as_ptr().cast_mut().cast(), + }; + let mut bytes: u32 = 0; + + let send_result = unsafe { + WinSock::WSASend( + socket, + &mut wsabuf, + 1, + &mut bytes, + 0, + std::ptr::null_mut(), + None, + ) + }; + if send_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +use crate::io::IoBuf; + +pub struct WriteOp<'a, B: IoBuf> { + handle: &'a InnerRawHandle, + buf: Option>, + completion_token: Option, + #[cfg(windows)] + socket_buf: Option>, +} + +impl<'a, B: IoBuf> WriteOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + completion_token: None, + #[cfg(windows)] + socket_buf: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for WriteOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let buf = self.buf.as_ref().unwrap().as_ref(); + + #[cfg(unix)] + let result = { + let written = unsafe { + libc::write( + self.handle.handle, + buf.as_buf_ptr().cast::(), + buf.buf_len(), + ) + }; + if written == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(written as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => { + let slice = unsafe { std::slice::from_raw_parts(buf.as_buf_ptr(), buf.buf_len()) }; + socket_write(socket as SOCKET, slice) + } + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based write currently supports sockets only on Windows", + )), + }; + + match poll_result_or_wait(result, self.handle, cx, driver, Interest::WRITABLE) { + Poll::Ready(Ok(written)) => Poll::Ready(Ok(written)), + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + Poll::Ready(Ok(result as usize)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let buf = self.buf.as_ref().unwrap().as_ref(); + match self.handle.handle { + RawOsHandle::Socket(socket) => { + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows socket I/O", + ) + })?; + + let wsabuf = self.socket_buf.get_or_insert_with(|| { + Box::new(WSABUF { + len: 0, + buf: std::ptr::null_mut(), + }) + }); + wsabuf.len = write_len; + wsabuf.buf = buf.as_buf_ptr() as *mut _; + + let send_result = unsafe { + WinSock::WSASend( + socket as SOCKET, + wsabuf.as_mut() as *mut WSABUF, + 1, + std::ptr::null_mut(), + 0, + overlapped, + None, + ) + }; + + if send_result == 0 { + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(err)) + } + } + RawOsHandle::Handle(handle) => { + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows file I/O", + ) + })?; + + let write_result = unsafe { + WriteFile( + handle as HANDLE, + buf.as_buf_ptr().cast(), + write_len, + std::ptr::null_mut(), + overlapped, + ) + }; + + if write_result != 0 { + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + Ok(()) + } else { + Err(err) + } + } + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_ref().unwrap().as_ref(); + let entry = opcode::Write::new( + types::Fd(self.handle.handle), + buf.as_buf_ptr(), + buf.buf_len() as _, + ) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for WriteOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + #[cfg(windows)] + let completion_state = self.socket_buf.take(); + #[cfg(not(windows))] + let completion_state = (); + + driver.ignore_completion( + completion_token, + Box::new(( + completion_state, + self.buf.take().map(CompletionBuffer::into_stable_box), + )), + ); + } + } + } +} diff --git a/vendor/vibeio/src/op/writeat.rs b/vendor/vibeio/src/op/writeat.rs new file mode 100644 index 0000000..62a83fe --- /dev/null +++ b/vendor/vibeio/src/op/writeat.rs @@ -0,0 +1,165 @@ +use std::io; +use std::task::{Context, Poll}; + +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_IO_PENDING, HANDLE}, + Storage::FileSystem::WriteFile, + System::IO::OVERLAPPED, +}; + +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +use crate::io::IoBuf; +use crate::op::io_util::CompletionBuffer; +use crate::op::Op; + +pub struct WriteAtOp<'a, B: IoBuf> { + handle: &'a InnerRawHandle, + buf: Option>, + offset: u64, + completion_token: Option, +} + +impl<'a, B: IoBuf> WriteAtOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, buf: B, offset: u64) -> Self { + Self { + handle, + buf: Some(CompletionBuffer::new(buf, handle.uses_completion())), + offset, + completion_token: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.buf.take().unwrap().into_inner() + } +} + +impl Op for WriteAtOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + let written = result as usize; + Poll::Ready(Ok(written)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let RawOsHandle::Handle(handle) = self.handle.handle else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "WriteAtOp expects a file handle, not a socket", + )); + }; + + let buf = self.buf.as_ref().unwrap().as_ref(); + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for Windows file I/O", + ) + })?; + + unsafe { + (*overlapped).Anonymous.Anonymous.Offset = self.offset as u32; + (*overlapped).Anonymous.Anonymous.OffsetHigh = (self.offset >> 32) as u32; + } + + let write_result = unsafe { + WriteFile( + handle as HANDLE, + buf.as_buf_ptr().cast(), + write_len, + std::ptr::null_mut(), + overlapped, + ) + }; + + if write_result != 0 { + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + Ok(()) + } else { + Err(err) + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let buf = self.buf.as_ref().unwrap().as_ref(); + let write_len = u32::try_from(buf.buf_len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "write buffer is too large for io_uring", + ) + })?; + + let entry = opcode::Write::new(types::Fd(self.handle.handle), buf.as_buf_ptr(), write_len) + .offset(self.offset) + .build() + .user_data(user_data); + + Ok(entry) + } +} + +impl Drop for WriteAtOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = crate::current_driver() { + if let Some(buf) = self.buf.take() { + driver.ignore_completion(completion_token, Box::new(buf.into_stable_box())); + } else { + driver.ignore_completion(completion_token, Box::new(())); + } + } + } + } +} diff --git a/vendor/vibeio/src/op/writev.rs b/vendor/vibeio/src/op/writev.rs new file mode 100644 index 0000000..c4b503c --- /dev/null +++ b/vendor/vibeio/src/op/writev.rs @@ -0,0 +1,363 @@ +use std::io::{self}; +use std::task::{Context, Poll}; + +use mio::Interest; +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_IO_PENDING, HANDLE}, + Networking::WinSock::{self as WinSock, SOCKET, WSABUF, WSA_IO_PENDING}, + Storage::FileSystem::WriteFile, + System::IO::OVERLAPPED, +}; + +use crate::current_driver; +use crate::driver::AnyDriver; +use crate::driver::CompletionIoResult; +use crate::fd_inner::InnerRawHandle; +#[cfg(windows)] +use crate::fd_inner::RawOsHandle; +#[cfg(unix)] +use crate::io::IoVec; +use crate::io::IoVectoredBuf; +use crate::op::io_util::poll_result_or_wait; +use crate::op::Op; + +/// Converts a slice of `IoSlice` to a system iovec buffer. +#[cfg(unix)] +#[inline] +fn iovec_to_system(bufs: &[IoVec]) -> Box<[libc::iovec]> { + use std::mem::MaybeUninit; + + let mut iovecs_maybeuninit: Box<[MaybeUninit]> = Box::new_uninit_slice(bufs.len()); + for (index, s) in bufs.iter().enumerate() { + let iov = libc::iovec { + iov_base: s.ptr as *mut libc::c_void, + iov_len: s.len, + }; + iovecs_maybeuninit[index].write(iov); + } + // SAFETY: The boxed slice would have all values initialized after interating over original array + unsafe { iovecs_maybeuninit.assume_init() } +} + +#[cfg(windows)] +#[inline] +fn socket_write_vectored(socket: SOCKET, bufs: &B) -> io::Result { + use windows_sys::Win32::Networking::WinSock::{self as WinSock, SOCKET_ERROR, WSABUF}; + + let iovecs = bufs.as_iovecs(); + let mut wsabufs = Vec::with_capacity(iovecs.len()); + for iovec in iovecs { + let len = u32::try_from(iovec.len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "writev buffer is too large for Windows socket I/O", + ) + })?; + wsabufs.push(WSABUF { + len, + buf: iovec.ptr as *mut _, + }); + } + + let mut bytes: u32 = 0; + let send_result = unsafe { + WinSock::WSASend( + socket, + wsabufs.as_mut_ptr(), + wsabufs.len() as u32, + &mut bytes, + 0, + std::ptr::null_mut(), + None, + ) + }; + if send_result == SOCKET_ERROR { + return Err(io::Error::from_raw_os_error(unsafe { + WinSock::WSAGetLastError() + })); + } + + Ok(bytes as usize) +} + +pub struct WritevOp<'a, B: IoVectoredBuf> { + handle: &'a InnerRawHandle, + bufs: Option, + completion_token: Option, + #[cfg(windows)] + completion_wsabufs: Option>, + #[cfg(windows)] + completion_staging: Option>, + #[cfg(target_os = "linux")] + completion_system_iovecs: Option>, +} + +impl<'a, B: IoVectoredBuf> WritevOp<'a, B> { + #[inline] + pub fn new(handle: &'a InnerRawHandle, bufs: B) -> Self { + Self { + handle, + bufs: Some(bufs), + completion_token: None, + #[cfg(windows)] + completion_wsabufs: None, + #[cfg(windows)] + completion_staging: None, + #[cfg(target_os = "linux")] + completion_system_iovecs: None, + } + } + + #[inline] + pub fn take_bufs(mut self) -> B { + self.bufs.take().unwrap() + } +} + +impl Op for WritevOp<'_, B> { + type Output = usize; + + #[cfg(any(unix, windows))] + #[inline] + fn poll_poll( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let bufs = self.bufs.as_ref().unwrap(); + + #[cfg(unix)] + let result = { + let iovecs = bufs.as_iovecs(); + let iovecs_system = iovec_to_system(&iovecs); + let written = unsafe { + libc::writev( + self.handle.handle, + iovecs_system.as_ptr(), + iovecs_system.len() as libc::c_int, + ) + }; + if written == -1 { + Err(io::Error::last_os_error()) + } else { + Ok(written as usize) + } + }; + + #[cfg(windows)] + let result = match self.handle.handle { + RawOsHandle::Socket(socket) => socket_write_vectored(socket as SOCKET, bufs), + RawOsHandle::Handle(_) => Err(io::Error::new( + io::ErrorKind::Unsupported, + "poll-based writev currently supports sockets only on Windows", + )), + }; + + match poll_result_or_wait(result, self.handle, cx, driver, Interest::WRITABLE) { + Poll::Ready(Ok(written)) => Poll::Ready(Ok(written)), + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + + #[cfg(any(unix, windows))] + #[inline] + fn poll_completion( + &mut self, + cx: &mut Context<'_>, + driver: &AnyDriver, + ) -> Poll> { + let result = if let Some(completion_token) = self.completion_token { + // Get the completion result + match driver.get_completion_result(completion_token) { + Some(result) => { + self.completion_token = None; + result + } + None => { + // The completion is not ready yet + driver.set_completion_waker(completion_token, cx.waker().clone()); + return Poll::Pending; + } + } + } else { + // Submit the op + match driver.submit_completion(self, cx.waker().clone()) { + CompletionIoResult::Ok(result) => result, + CompletionIoResult::Retry(token) => { + self.completion_token = Some(token); + return Poll::Pending; + } + CompletionIoResult::SubmitErr(err) => return Poll::Ready(Err(err)), + } + }; + if result < 0 { + #[cfg(windows)] + { + self.completion_wsabufs = None; + self.completion_staging = None; + } + return Poll::Ready(Err(io::Error::from_raw_os_error(-result))); + } + + #[cfg(windows)] + { + self.completion_wsabufs = None; + self.completion_staging = None; + } + + Poll::Ready(Ok(result as usize)) + } + + #[cfg(windows)] + #[inline] + fn submit_windows(&mut self, overlapped: *mut OVERLAPPED) -> Result<(), io::Error> { + let bufs = self.bufs.as_ref().unwrap(); + match self.handle.handle { + RawOsHandle::Socket(socket) => { + let iovecs = bufs.as_iovecs(); + let mut wsabufs = Vec::with_capacity(iovecs.len()); + for iovec in iovecs { + let len = u32::try_from(iovec.len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "writev buffer is too large for Windows socket I/O", + ) + })?; + wsabufs.push(WSABUF { + len, + buf: iovec.ptr as *mut _, + }); + } + + let mut wsabufs = wsabufs.into_boxed_slice(); + let send_result = unsafe { + WinSock::WSASend( + socket as SOCKET, + wsabufs.as_mut_ptr(), + wsabufs.len() as u32, + std::ptr::null_mut(), + 0, + overlapped, + None, + ) + }; + + if send_result == 0 { + self.completion_wsabufs = Some(wsabufs); + self.completion_staging = None; + return Ok(()); + } + + let err = unsafe { WinSock::WSAGetLastError() }; + if err == WSA_IO_PENDING { + self.completion_wsabufs = Some(wsabufs); + self.completion_staging = None; + Ok(()) + } else { + self.completion_wsabufs = None; + self.completion_staging = None; + Err(io::Error::from_raw_os_error(err)) + } + } + RawOsHandle::Handle(handle) => { + let iovecs = bufs.as_iovecs(); + let total_len = (iovecs.iter()).try_fold(0usize, |acc, iovec| { + acc.checked_add(iovec.len).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "writev buffer length overflow") + }) + })?; + let total_len_u32 = u32::try_from(total_len).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "writev total length is too large for Windows file I/O", + ) + })?; + + let mut staging = Vec::with_capacity(total_len); + for iovec in iovecs { + let slice = unsafe { std::slice::from_raw_parts(iovec.ptr, iovec.len) }; + staging.extend_from_slice(slice); + } + + let write_result = unsafe { + WriteFile( + handle as HANDLE, + staging.as_ptr().cast(), + total_len_u32, + std::ptr::null_mut(), + overlapped, + ) + }; + + if write_result != 0 { + self.completion_wsabufs = None; + self.completion_staging = Some(staging); + return Ok(()); + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_IO_PENDING as i32) { + self.completion_wsabufs = None; + self.completion_staging = Some(staging); + Ok(()) + } else { + self.completion_wsabufs = None; + self.completion_staging = None; + Err(err) + } + } + } + } + + #[cfg(target_os = "linux")] + #[inline] + fn build_completion_entry( + &mut self, + user_data: u64, + ) -> Result { + use io_uring::{opcode, types}; + + let bufs = self.bufs.as_ref().unwrap(); + + // Build a temporary iovec array for the syscall. + let iovecs = if let Some(iovecs) = self.completion_system_iovecs.take() { + iovecs + } else { + iovec_to_system(&bufs.as_iovecs()) + }; + + let entry = opcode::Writev::new( + types::Fd(self.handle.handle), + iovecs.as_ptr(), + iovecs.len() as _, + ) + .build() + .user_data(user_data); + + // Store the iovec array for the completion, because it needs to be kept alive until the + // completion is ready. + self.completion_system_iovecs = Some(iovecs); + + Ok(entry) + } +} + +impl Drop for WritevOp<'_, B> { + #[inline] + fn drop(&mut self) { + if let Some(completion_token) = self.completion_token { + if let Some(driver) = current_driver() { + #[cfg(target_os = "linux")] + let bufs = self.completion_system_iovecs.take(); + #[cfg(windows)] + let bufs = self.completion_wsabufs.take(); + #[cfg(not(any(target_os = "linux", windows)))] + let bufs = (); + + driver.ignore_completion(completion_token, Box::new((bufs, self.bufs.take()))); + } + } + } +} diff --git a/vendor/vibeio/src/process/mod.rs b/vendor/vibeio/src/process/mod.rs new file mode 100644 index 0000000..57f339b --- /dev/null +++ b/vendor/vibeio/src/process/mod.rs @@ -0,0 +1,1064 @@ +//! Process utilities for spawning and managing child processes. +//! +//! This module provides async-aware wrappers around `std::process::Command` and +//! `std::process::Child`, allowing you to spawn and interact with child processes +//! without blocking the executor. +//! +//! Key types: +//! - `Command`: an async-aware builder for spawning child processes. +//! - `Child`: represents a running child process with async `wait()`, `kill()`, etc. +//! - `ChildStdin`, `ChildStdout`, `ChildStderr`: async-aware stdio streams. +//! +//! Implementation notes: +//! - On Unix, the module uses `mio`/`io_uring` drivers to register child process +//! file descriptors for async I/O when possible. Falls back to a blocking pool +//! when the driver is unavailable or registration fails. +//! - A background `ZombieReaper` task runs per runtime to reap exited child processes +//! and avoid leaving zombies. This is started automatically when the runtime is +//! created. +//! - Calling any process API outside a runtime will panic (matching the library's +//! general behavior for runtime-only APIs). + +mod reaper; + +use reaper::ZombieReaper; +pub(crate) use reaper::{start_zombie_reaper, ZombieReaperMessage}; + +use std::cell::RefCell; +use std::ffi::OsStr; +#[cfg(unix)] +use std::future::poll_fn; +use std::io::{self, Read, Write}; +use std::sync::{Arc, Mutex}; + +#[cfg(unix)] +use mio::Interest; + +use std::mem::ManuallyDrop; +#[cfg(unix)] +use std::os::fd::{AsRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawHandle, IntoRawHandle, RawHandle}; + +#[cfg(unix)] +use crate::driver::RegistrationMode; +use crate::executor::current_driver; +#[cfg(unix)] +use crate::fd_inner::InnerRawHandle; +use crate::io::{iobuf_to_slice, iobufmut_to_slice, AsyncRead, AsyncWrite, IoBuf, IoBufMut}; +#[cfg(unix)] +use crate::op::{ReadOp, WriteOp}; + +pub use std::process::{ExitStatus, Output, Stdio}; + +#[cfg(unix)] +enum ChildIo { + Async(ManuallyDrop), + Blocking, +} + +#[cfg(windows)] +enum ChildIo { + Blocking, +} + +#[inline] +fn stdio_closed_error() -> io::Error { + io::Error::new(io::ErrorKind::BrokenPipe, "child stdio is closed") +} + +#[inline] +fn command_consumed_error() -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, "command has been consumed") +} + +#[inline] +fn child_consumed_error() -> io::Error { + io::Error::new( + io::ErrorKind::InvalidInput, + "child process has been consumed", + ) +} + +#[inline] +fn blocking_pool_io_error() -> io::Error { + io::Error::other("can't spawn blocking task for process I/O") +} + +#[cfg(unix)] +#[inline] +fn configure_nonblocking(fd: RawFd, uses_completion: bool) { + // On Linux, pipe2() already sets O_NONBLOCK at creation time, so + // we only need to adjust flags when the mode doesn't match. + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags == -1 { + return; + } + + let mut new_flags = flags | libc::O_NONBLOCK; + if uses_completion { + new_flags &= !libc::O_NONBLOCK; + } + if new_flags != flags { + unsafe { + libc::fcntl(fd, libc::F_SETFL, new_flags); + } + } +} + +#[cfg(unix)] +#[inline] +fn make_child_io(fd: RawFd, interest: Interest) -> ChildIo { + if let Some(driver) = current_driver() { + match InnerRawHandle::new_with_driver_and_mode( + &driver, + fd, + interest, + RegistrationMode::Completion, + ) { + Ok(handle) => { + let handle = ManuallyDrop::new(handle); + configure_nonblocking(fd, handle.uses_completion()); + ChildIo::Async(handle) + } + Err(_) => ChildIo::Blocking, + } + } else { + ChildIo::Blocking + } +} + +#[inline] +async fn read_in_blocking_pool(inner: R, buf: B) -> (io::Result, R, B) +where + R: Read + Send + 'static, + B: IoBufMut, +{ + let shared = Arc::new(Mutex::new(RefCell::new(Some((inner, buf))))); + let shared_clone = shared.clone(); + let result = crate::spawn_blocking(move || { + let (mut inner, mut buf) = shared_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner/buf is none"); + let temp_slice = iobufmut_to_slice(&mut buf); + let result = inner.read(temp_slice); + (result, inner, buf) + }) + .await; + + match result { + Ok(result) => result, + Err(_) => { + let (inner, buf) = shared + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner/buf is none"); + (Err(blocking_pool_io_error()), inner, buf) + } + } +} + +#[inline] +async fn write_in_blocking_pool(inner: W, buf: B) -> (io::Result, W, B) +where + W: Write + Send + 'static, + B: IoBuf, +{ + let shared = Arc::new(Mutex::new(RefCell::new(Some((inner, buf))))); + let shared_clone = shared.clone(); + let result = crate::spawn_blocking(move || { + let (mut inner, buf) = shared_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner/buf is none"); + let temp_slice = iobuf_to_slice(&buf); + let result = inner.write(temp_slice); + (result, inner, buf) + }) + .await; + + match result { + Ok(result) => result, + Err(_) => { + let (inner, buf) = shared + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner/buf is none"); + (Err(blocking_pool_io_error()), inner, buf) + } + } +} + +/// Async-aware child process stdin stream. +/// +/// This type wraps `std::process::ChildStdin` and implements `AsyncWrite` +/// to allow writing to a child process's standard input without blocking +/// the executor. +/// +/// # Examples +/// ```ignore +/// let mut child = Command::new("cat").stdin(Stdio::piped()).spawn()?; +/// child.stdin.as_mut().unwrap().write_all(b"hello\n").await?; +/// ``` +pub struct ChildStdin { + inner: Option, + #[allow(dead_code)] + io: ChildIo, +} + +/// Async-aware child process stdout stream. +/// +/// This type wraps `std::process::ChildStdout` and implements `AsyncRead` +/// to allow reading from a child process's standard output without blocking +/// the executor. +/// +/// # Examples +/// ```ignore +/// let mut child = Command::new("echo").stdout(Stdio::piped()).spawn()?; +/// let mut buf = vec![0u8; 64]; +/// let (result, buf) = child.stdout.as_mut().unwrap().read(buf).await; +/// ``` +pub struct ChildStdout { + inner: Option, + #[allow(dead_code)] + io: ChildIo, +} + +/// Async-aware child process stderr stream. +/// +/// This type wraps `std::process::ChildStderr` and implements `AsyncRead` +/// to allow reading from a child process's standard error without blocking +/// the executor. +/// +/// # Examples +/// ```ignore +/// let mut child = Command::new("sh").stderr(Stdio::piped()).spawn()?; +/// let mut buf = vec![0u8; 64]; +/// let (result, buf) = child.stderr.as_mut().unwrap().read(buf).await; +/// ``` +pub struct ChildStderr { + inner: Option, + #[allow(dead_code)] + io: ChildIo, +} + +impl ChildStdin { + /// Create a new `ChildStdin` from a standard library `ChildStdin`. + #[inline] + pub(crate) fn from_std(inner: std::process::ChildStdin) -> io::Result { + #[cfg(unix)] + let io = make_child_io(inner.as_raw_fd(), Interest::WRITABLE); + #[cfg(windows)] + let io = ChildIo::Blocking; + + Ok(Self { + inner: Some(inner), + io, + }) + } + + /// Consume this `ChildStdin` and return the underlying `std::process::ChildStdin`. + #[inline] + pub fn into_std(self) -> std::process::ChildStdin { + #[cfg(not(unix))] + let this = ManuallyDrop::new(self); + #[cfg(unix)] + let mut this = ManuallyDrop::new(self); + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut this.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + let inner = unsafe { std::ptr::read(&this.inner) }; + inner.expect("child stdin is already taken") + } + + #[inline] + fn drop_handle(&mut self) { + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut self.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + } +} + +impl ChildStdout { + /// Create a new `ChildStdout` from a standard library `ChildStdout`. + #[inline] + pub(crate) fn from_std(inner: std::process::ChildStdout) -> io::Result { + #[cfg(unix)] + let io = make_child_io(inner.as_raw_fd(), Interest::READABLE); + #[cfg(windows)] + let io = ChildIo::Blocking; + + Ok(Self { + inner: Some(inner), + io, + }) + } + + /// Consume this `ChildStdout` and return the underlying `std::process::ChildStdout`. + #[inline] + pub fn into_std(self) -> std::process::ChildStdout { + #[cfg(not(unix))] + let this = ManuallyDrop::new(self); + #[cfg(unix)] + let mut this = ManuallyDrop::new(self); + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut this.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + let inner = unsafe { std::ptr::read(&this.inner) }; + inner.expect("child stdout is already taken") + } + + #[inline] + fn drop_handle(&mut self) { + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut self.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + } +} + +impl ChildStderr { + /// Create a new `ChildStderr` from a standard library `ChildStderr`. + #[inline] + pub(crate) fn from_std(inner: std::process::ChildStderr) -> io::Result { + #[cfg(unix)] + let io = make_child_io(inner.as_raw_fd(), Interest::READABLE); + #[cfg(windows)] + let io = ChildIo::Blocking; + + Ok(Self { + inner: Some(inner), + io, + }) + } + + /// Consume this `ChildStderr` and return the underlying `std::process::ChildStderr`. + #[inline] + pub fn into_std(self) -> std::process::ChildStderr { + #[cfg(not(unix))] + let this = ManuallyDrop::new(self); + #[cfg(unix)] + let mut this = ManuallyDrop::new(self); + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut this.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + let inner = unsafe { std::ptr::read(&this.inner) }; + inner.expect("child stderr is already taken") + } + + #[inline] + fn drop_handle(&mut self) { + #[cfg(unix)] + if let ChildIo::Async(handle) = &mut self.io { + unsafe { + ManuallyDrop::drop(handle); + } + } + } +} + +impl Drop for ChildStdin { + #[inline] + fn drop(&mut self) { + self.drop_handle(); + } +} + +impl Drop for ChildStdout { + #[inline] + fn drop(&mut self) { + self.drop_handle(); + } +} + +impl Drop for ChildStderr { + #[inline] + fn drop(&mut self) { + self.drop_handle(); + } +} + +impl AsyncWrite for ChildStdin { + #[inline] + async fn write(&mut self, buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + #[cfg(unix)] + if let ChildIo::Async(handle) = &self.io { + let mut op = WriteOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + return (result, op.take_bufs()); + } + + if current_driver().is_some() { + let inner = match self.inner.take() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let (result, inner, buf) = write_in_blocking_pool(inner, buf).await; + self.inner = Some(inner); + (result, buf) + } else { + let inner = match self.inner.as_mut() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let temp_slice = iobuf_to_slice(&buf); + (inner.write(temp_slice), buf) + } + } + + #[inline] + async fn flush(&mut self) -> Result<(), io::Error> { + #[cfg(unix)] + if let ChildIo::Async(_) = &self.io { + return Ok(()); + } + + if current_driver().is_some() { + let inner = match self.inner.take() { + Some(inner) => inner, + None => return Err(stdio_closed_error()), + }; + let shared = Arc::new(Mutex::new(RefCell::new(Some(inner)))); + let shared_clone = shared.clone(); + let result = crate::spawn_blocking(move || { + let mut inner = shared_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner is none"); + let flush_result = inner.flush(); + (flush_result, inner) + }) + .await; + + match result { + Ok((flush_result, inner)) => { + self.inner = Some(inner); + flush_result + } + Err(_) => { + let inner = shared + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("inner is none"); + self.inner = Some(inner); + Err(blocking_pool_io_error()) + } + } + } else { + let inner = self.inner.as_mut().ok_or_else(stdio_closed_error)?; + inner.flush() + } + } +} + +impl AsyncRead for ChildStdout { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + #[cfg(unix)] + if let ChildIo::Async(handle) = &self.io { + let mut op = ReadOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + return (result, op.take_bufs()); + } + + if current_driver().is_some() { + let inner = match self.inner.take() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let (result, inner, buf) = read_in_blocking_pool(inner, buf).await; + self.inner = Some(inner); + (result, buf) + } else { + let inner = match self.inner.as_mut() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let mut buf = buf; + let temp_slice = iobufmut_to_slice(&mut buf); + (inner.read(temp_slice), buf) + } + } +} + +impl AsyncRead for ChildStderr { + #[inline] + async fn read(&mut self, buf: B) -> (Result, B) { + if buf.buf_len() == 0 { + return (Ok(0), buf); + } + + #[cfg(unix)] + if let ChildIo::Async(handle) = &self.io { + let mut op = ReadOp::new(handle, buf); + let result = poll_fn(|cx| handle.poll_op(cx, &mut op)).await; + return (result, op.take_bufs()); + } + + if current_driver().is_some() { + let inner = match self.inner.take() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let (result, inner, buf) = read_in_blocking_pool(inner, buf).await; + self.inner = Some(inner); + (result, buf) + } else { + let inner = match self.inner.as_mut() { + Some(inner) => inner, + None => return (Err(stdio_closed_error()), buf), + }; + let mut buf = buf; + let temp_slice = iobufmut_to_slice(&mut buf); + (inner.read(temp_slice), buf) + } + } +} + +#[cfg(unix)] +impl AsRawFd for ChildStdin { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner + .as_ref() + .expect("child stdin is already taken") + .as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for ChildStdin { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.into_std().into_raw_fd() + } +} + +#[cfg(unix)] +impl AsRawFd for ChildStdout { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner + .as_ref() + .expect("child stdout is already taken") + .as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for ChildStdout { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.into_std().into_raw_fd() + } +} + +#[cfg(unix)] +impl AsRawFd for ChildStderr { + #[inline] + fn as_raw_fd(&self) -> RawFd { + self.inner + .as_ref() + .expect("child stderr is already taken") + .as_raw_fd() + } +} + +#[cfg(unix)] +impl IntoRawFd for ChildStderr { + #[inline] + fn into_raw_fd(self) -> RawFd { + self.into_std().into_raw_fd() + } +} + +#[cfg(windows)] +impl AsRawHandle for ChildStdin { + #[inline] + fn as_raw_handle(&self) -> RawHandle { + self.inner + .as_ref() + .expect("child stdin is already taken") + .as_raw_handle() + } +} + +#[cfg(windows)] +impl IntoRawHandle for ChildStdin { + #[inline] + fn into_raw_handle(self) -> RawHandle { + self.into_std().into_raw_handle() + } +} + +#[cfg(windows)] +impl AsRawHandle for ChildStdout { + #[inline] + fn as_raw_handle(&self) -> RawHandle { + self.inner + .as_ref() + .expect("child stdout is already taken") + .as_raw_handle() + } +} + +#[cfg(windows)] +impl IntoRawHandle for ChildStdout { + #[inline] + fn into_raw_handle(self) -> RawHandle { + self.into_std().into_raw_handle() + } +} + +#[cfg(windows)] +impl AsRawHandle for ChildStderr { + #[inline] + fn as_raw_handle(&self) -> RawHandle { + self.inner + .as_ref() + .expect("child stderr is already taken") + .as_raw_handle() + } +} + +#[cfg(windows)] +impl IntoRawHandle for ChildStderr { + #[inline] + fn into_raw_handle(self) -> RawHandle { + self.into_std().into_raw_handle() + } +} + +/// Async-aware wrapper around `std::process::Child`. +/// +/// This type provides async methods to interact with a running child process: +/// - `wait()`: asynchronously wait for the process to exit. +/// - `kill()`: kill the process. +/// - `try_wait()`: non-blocking check if the process has exited. +/// - `stdin`, `stdout`, `stderr`: async streams for stdio. +/// +/// # Examples +/// ```ignore +/// let mut child = Command::new("echo").arg("hello").spawn()?; +/// let status = child.wait().await?; +/// println!("exit status: {}", status); +/// ``` +pub struct Child { + inner: Option, + id: u32, + pub stdin: Option, + pub stdout: Option, + pub stderr: Option, + reaper: ZombieReaper, +} + +impl Child { + /// Create a new `Child` from a standard library `Child`. + #[inline] + pub(crate) fn from_std(mut child: std::process::Child) -> io::Result { + let stdin = child.stdin.take().map(ChildStdin::from_std).transpose()?; + let stdout = child.stdout.take().map(ChildStdout::from_std).transpose()?; + let stderr = child.stderr.take().map(ChildStderr::from_std).transpose()?; + let id = child.id(); + + Ok(Self { + inner: Some(child), + id, + stdin, + stdout, + stderr, + reaper: ZombieReaper::new(), + }) + } + + /// Returns the OS-assigned process identifier. + #[inline] + pub fn id(&self) -> u32 { + self.id + } + + #[inline] + fn inner_mut(&mut self) -> io::Result<&mut std::process::Child> { + self.inner.as_mut().ok_or_else(child_consumed_error) + } + + #[inline] + fn take_inner(&mut self) -> io::Result { + self.inner.take().ok_or_else(child_consumed_error) + } + + /// Force kill the process. + #[inline] + pub fn kill(&mut self) -> io::Result<()> { + self.inner_mut()?.kill() + } + + /// Asynchronously wait for the process to exit. + /// + /// This method returns a future that resolves to the process's exit status. + /// The future completes when the process has fully exited and been reaped. + #[inline] + pub async fn wait(&mut self) -> io::Result { + let _ = self.stdin.take(); // Similarly to std::process::Child::wait + let child = self.take_inner()?; + self.reaper.wait(child).await + } + + /// Check if the process has exited without blocking. + /// + /// Returns `Ok(Some(status))` if the process has exited, `Ok(None)` if it + /// is still running, or an error if checking the status fails. + #[inline] + pub fn try_wait(&mut self) -> io::Result> { + self.inner_mut()?.try_wait() + } +} + +impl Drop for Child { + #[inline] + fn drop(&mut self) { + let _ = self.stdin.take(); // Similarly to std::process::Child::wait + if let Some(child) = self.inner.take() { + self.reaper.reap_on_drop(child); + } + } +} + +/// Async-aware builder for spawning child processes. +/// +/// This type wraps `std::process::Command` and provides async versions of +/// the spawn methods: +/// - `spawn()`: spawn the process and return a `Child`. +/// - `status()`: run the process to completion and return its exit status. +/// - `output()`: run the process to completion and return its output. +/// +/// # Examples +/// ```ignore +/// let status = Command::new("echo") +/// .arg("hello") +/// .status() +/// .await?; +/// println!("exit status: {}", status); +/// ``` +pub struct Command { + inner: Option, +} + +impl Command { + /// Create a new `Command` for the given program. + #[inline] + pub fn new(program: impl AsRef) -> Self { + Self { + inner: Some(std::process::Command::new(program)), + } + } + + #[inline] + fn inner_mut(&mut self) -> &mut std::process::Command { + self.inner.as_mut().expect("command has been consumed") + } + + /// Add an argument to pass to the program. + #[inline] + pub fn arg(&mut self, arg: impl AsRef) -> &mut Self { + self.inner_mut().arg(arg); + self + } + + /// Add multiple arguments to pass to the program. + #[inline] + pub fn args(&mut self, args: I) -> &mut Self + where + I: IntoIterator, + S: AsRef, + { + self.inner_mut().args(args); + self + } + + /// Set an environment variable for the process. + #[inline] + pub fn env(&mut self, key: K, val: V) -> &mut Self + where + K: AsRef, + V: AsRef, + { + self.inner_mut().env(key, val); + self + } + + /// Set multiple environment variables for the process. + #[inline] + pub fn envs(&mut self, vars: I) -> &mut Self + where + I: IntoIterator, + K: AsRef, + V: AsRef, + { + self.inner_mut().envs(vars); + self + } + + /// Remove an environment variable for the process. + #[inline] + pub fn env_remove>(&mut self, key: K) -> &mut Self { + self.inner_mut().env_remove(key); + self + } + + /// Clear all environment variables for the process. + #[inline] + pub fn env_clear(&mut self) -> &mut Self { + self.inner_mut().env_clear(); + self + } + + /// Set the working directory for the process. + #[inline] + pub fn current_dir(&mut self, dir: impl AsRef) -> &mut Self { + self.inner_mut().current_dir(dir); + self + } + + /// Configure the standard input for the process. + #[inline] + pub fn stdin(&mut self, cfg: Stdio) -> &mut Self { + self.inner_mut().stdin(cfg); + self + } + + /// Configure the standard output for the process. + #[inline] + pub fn stdout(&mut self, cfg: Stdio) -> &mut Self { + self.inner_mut().stdout(cfg); + self + } + + /// Configure the standard error for the process. + #[inline] + pub fn stderr(&mut self, cfg: Stdio) -> &mut Self { + self.inner_mut().stderr(cfg); + self + } + + /// Spawn the process and return a `Child` handle. + #[inline] + pub fn spawn(&mut self) -> io::Result { + let child = self + .inner + .as_mut() + .ok_or_else(command_consumed_error)? + .spawn()?; + Child::from_std(child) + } + + /// Run the process to completion and return its exit status. + /// + /// This is an async version of `std::process::Command::status`. + #[inline] + pub async fn status(&mut self) -> io::Result { + if current_driver().is_some() { + let inner = self.inner.take().ok_or_else(command_consumed_error)?; + let shared = Arc::new(Mutex::new(RefCell::new(Some(inner)))); + let shared_clone = shared.clone(); + let result = crate::spawn_blocking(move || { + let mut cmd = shared_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("command is none"); + let status = cmd.status(); + (status, cmd) + }) + .await; + + match result { + Ok((status, cmd)) => { + self.inner = Some(cmd); + status + } + Err(_) => { + let cmd = shared + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("command is none"); + self.inner = Some(cmd); + Err(blocking_pool_io_error()) + } + } + } else { + self.inner_mut().status() + } + } + + /// Run the process to completion and return its output. + /// + /// This is an async version of `std::process::Command::output`. + #[inline] + pub async fn output(&mut self) -> io::Result { + if current_driver().is_some() { + let inner = self.inner.take().ok_or_else(command_consumed_error)?; + let shared = Arc::new(Mutex::new(RefCell::new(Some(inner)))); + let shared_clone = shared.clone(); + let result = crate::spawn_blocking(move || { + let mut cmd = shared_clone + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("command is none"); + let output = cmd.output(); + (output, cmd) + }) + .await; + + match result { + Ok((output, cmd)) => { + self.inner = Some(cmd); + output + } + Err(_) => { + let cmd = shared + .try_lock() + .ok() + .and_then(|rc| rc.take()) + .expect("command is none"); + self.inner = Some(cmd); + Err(blocking_pool_io_error()) + } + } + } else { + self.inner_mut().output() + } + } + + /// Get a mutable reference to the underlying `std::process::Command`. + #[inline] + pub fn as_std(&mut self) -> &mut std::process::Command { + self.inner_mut() + } + + /// Consume this `Command` and return the underlying `std::process::Command`. + #[inline] + pub fn into_std(mut self) -> std::process::Command { + self.inner.take().expect("command has been consumed") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::AnyDriver; + use crate::executor::Runtime; + use crate::io::{AsyncRead, AsyncWrite, IoBufWithCursor}; + + fn make_runtime() -> Runtime { + Runtime::new(AnyDriver::new_best().expect("driver should initialize")) + } + + async fn write_all(writer: &mut W, buf: &[u8]) -> io::Result<()> { + let mut cursor = IoBufWithCursor::new(buf.to_vec()); + while cursor.buf_len() > 0 { + let (result, mut next) = writer.write(cursor).await; + let written = result?; + if written == 0 { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write whole buffer", + )); + } + next.advance(written); + cursor = next; + } + Ok(()) + } + + async fn read_line(reader: &mut R) -> io::Result { + let mut output = Vec::new(); + loop { + let (result, buf) = reader.read(vec![0u8; 64]).await; + let read = result?; + if read == 0 { + break; + } + output.extend_from_slice(&buf[..read]); + if output.contains(&b'\n') { + break; + } + } + Ok(String::from_utf8_lossy(&output) + .trim_end_matches(&['\r', '\n'][..]) + .to_string()) + } + + #[test] + fn command_spawn_stdio_roundtrip() { + make_runtime().block_on(async { + let mut cmd = if cfg!(windows) { + let mut cmd = Command::new("cmd"); + cmd.args([ + "/V:ON", + "/C", + "set /p line= & echo out:!line!& echo err:!line!>&2", + ]); + cmd + } else { + let mut cmd = Command::new("sh"); + cmd.args(["-c", "read line; echo out:$line; echo err:$line 1>&2"]); + cmd + }; + + cmd.stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + + let mut child = cmd.spawn().expect("spawn should succeed"); + let mut stdin = child.stdin.take().expect("stdin should be piped"); + let mut stdout = child.stdout.take().expect("stdout should be piped"); + let mut stderr = child.stderr.take().expect("stderr should be piped"); + + write_all(&mut stdin, b"hello\n") + .await + .expect("write to stdin"); + drop(stdin); + + let out_line = read_line(&mut stdout).await.expect("read stdout"); + let err_line = read_line(&mut stderr).await.expect("read stderr"); + + assert_eq!(out_line, "out:hello"); + assert_eq!(err_line, "err:hello"); + + let status = child.wait().await.expect("wait succeeds"); + assert!(status.success()); + }); + } +} diff --git a/vendor/vibeio/src/process/reaper.rs b/vendor/vibeio/src/process/reaper.rs new file mode 100644 index 0000000..cef1cfc --- /dev/null +++ b/vendor/vibeio/src/process/reaper.rs @@ -0,0 +1,329 @@ +use std::io::{self}; +use std::process::ExitStatus; +use std::task::{Context, Poll, Waker}; + +use futures_util::FutureExt; +#[cfg(windows)] +use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE; + +#[cfg(windows)] +use std::os::windows::io::AsRawHandle; + +use crate::current_zombie_reaper; + +pub(crate) struct ZombieReaper; + +pub(crate) type ZombieReaperMessage = ( + std::process::Child, + Option>>, +); + +/// Zombie reaper process that waits on child processes asynchronously. +impl ZombieReaper { + /// Creates a new zombie reaper instance. + #[inline] + pub(crate) fn new() -> Self { + Self + } + + /// Waits on a child process asynchronously. + #[inline] + pub(crate) async fn wait(&self, mut child: std::process::Child) -> io::Result { + if let Some(reaper_send) = current_zombie_reaper().await { + // Send the child to the zombie reaper to wait on it asynchronously. + let (sender, recver) = oneshot::async_channel(); + let _ = reaper_send.send((child, Some(sender))).await; + recver + .await + .map_err(|_| std::io::Error::other("zombie reaper error"))? + } else { + child.wait() + } + } + + /// Reaps a child process on drop, waiting asynchronously if possible. + #[inline] + pub(crate) fn reap_on_drop(&self, mut child: std::process::Child) { + if let Ok(Some(_)) = child.try_wait() { + return; + } + if let Poll::Ready(Some(reaper_send)) = + Box::pin(current_zombie_reaper()).poll_unpin(&mut Context::from_waker(Waker::noop())) + { + // Send the child to the zombie reaper, so it can wait on it asynchronously. + let (sender, _) = oneshot::async_channel(); + let _ = reaper_send.try_send((child, Some(sender))); + } else { + std::thread::spawn(move || { + let _ = child.wait(); + }); + } + } +} + +#[inline] +pub(crate) async fn start_zombie_reaper() -> async_channel::Sender { + let (tx, rx) = async_channel::unbounded(); + crate::spawn(zombie_reaper_fn(rx)); + tx +} + +// --------------------------------------------------------------------------- +// Windows reaper implementation +// --------------------------------------------------------------------------- + +#[cfg(windows)] +struct WaitContext { + child: std::process::Child, + sender: Option>>, + wait_handle: windows_sys::Win32::Foundation::HANDLE, +} + +// SAFETY: `wait_handle` isn't used after the callback returns, so it's safe to send across threads. +#[cfg(windows)] +unsafe impl Send for WaitContext {} + +#[cfg(windows)] +unsafe extern "system" fn wait_callback(ctx: *mut std::ffi::c_void, _timed_out: bool) { + let mut ctx = Box::from_raw(ctx as *mut WaitContext); + let result = ctx.child.wait(); + if let Some(sender) = ctx.sender.take() { + let _ = sender.send(result); + } + // WT_EXECUTEONLYONCE unregisters the wait after this callback returns. +} + +#[inline] +#[cfg(windows)] +async fn zombie_reaper_fn(rx: async_channel::Receiver) { + use windows_sys::Win32::System::Threading::{ + RegisterWaitForSingleObject, INFINITE, WT_EXECUTEONLYONCE, + }; + + while let Ok((mut child, sender)) = rx.recv().await { + // Fast path: child may already have exited. + match child.try_wait() { + Ok(Some(status)) => { + if let Some(sender) = sender { + let _ = sender.send(Ok(status)); + } + continue; + } + Ok(None) => {} + Err(err) => { + if let Some(sender) = sender { + let _ = sender.send(Err(err)); + } + continue; + } + } + + let ctx = Box::new(WaitContext { + child, + sender, + wait_handle: INVALID_HANDLE_VALUE, + }); + let process_handle = ctx.child.as_raw_handle(); + let ctx_ptr = Box::into_raw(ctx); + + let ok = unsafe { + RegisterWaitForSingleObject( + &mut (*ctx_ptr).wait_handle, + process_handle, + Some(wait_callback), + ctx_ptr as *mut std::ffi::c_void, + INFINITE, + WT_EXECUTEONLYONCE, + ) + }; + + if ok == 0 { + // Fall back to a dedicated thread if registration fails. + let mut ctx = unsafe { Box::from_raw(ctx_ptr) }; + std::thread::spawn(move || { + let result = ctx.child.wait(); + if let Some(sender) = ctx.sender.take() { + let _ = sender.send(result); + } + }); + } + } +} + +#[inline] +#[cfg(unix)] +async fn zombie_reaper_fn(rx: async_channel::Receiver) { + // On Linux, prefer the pidfd-based reaper (no signals needed, works with + // both mio/epoll and io_uring drivers). Fall back to the generic Unix + // implementation if pidfd_open is unavailable (kernel < 5.3). + #[cfg(target_os = "linux")] + { + if pidfd_available() { + return zombie_reaper_fn_linux_pidfd(rx).await; + } + } + + zombie_reaper_fn_unix(rx).await +} + +// --------------------------------------------------------------------------- +// Linux pidfd-based reaper (kernel ≥ 5.3) +// --------------------------------------------------------------------------- + +/// Probe whether `pidfd_open` is supported on this kernel. +#[cfg(target_os = "linux")] +#[inline] +fn pidfd_available() -> bool { + // Try opening a pidfd for our own PID — it will succeed on 5.3+ and fail + // with ENOSYS on older kernels. + let fd = unsafe { libc::syscall(libc::SYS_pidfd_open, libc::getpid(), 0 as libc::c_uint) }; + if fd >= 0 { + unsafe { + libc::close(fd as libc::c_int); + } + true + } else { + let err = io::Error::last_os_error(); + // ENOSYS → syscall not available; anything else means it *is* available + // but the specific call failed for another reason (shouldn't happen for + // our own PID, but be safe). + err.raw_os_error() != Some(libc::ENOSYS) + } +} + +/// Convert a raw `waitpid` status into a `std::process::ExitStatus`. +#[cfg(target_os = "linux")] +#[inline] +fn exit_status_from_raw(raw: i32) -> ExitStatus { + use std::os::unix::process::ExitStatusExt; + ExitStatus::from_raw(raw) +} + +#[cfg(target_os = "linux")] +#[inline] +async fn zombie_reaper_fn_linux_pidfd(rx: async_channel::Receiver) { + use std::future::poll_fn; + + use crate::op::{Op, WaitPidOp}; + + loop { + let msg = rx.recv().await; + let Ok((mut child, sender)) = msg else { + // Channel closed — runtime is shutting down. + break; + }; + + // Fast path: child may already have exited. + match child.try_wait() { + Ok(Some(status)) => { + if let Some(sender) = sender { + let _ = sender.send(Ok(status)); + } + continue; + } + Ok(None) => {} + Err(err) => { + if let Some(sender) = sender { + let _ = sender.send(Err(err)); + } + continue; + } + } + + let pid = child.id(); + + // Spawn a lightweight task per child that waits on the pidfd. + // This lets us handle many children concurrently without blocking + // the reaper loop. + crate::spawn(async move { + let mut op = WaitPidOp::new(pid); + let result = poll_fn(|cx| { + // The WaitPidOp uses poll-based I/O internally (pidfd registered + // in Poll mode), so we always go through poll_poll regardless of + // whether the driver supports completions. + op.poll(cx, &crate::executor::current_driver().expect("no driver")) + }) + .await; + + let status = result.map(exit_status_from_raw); + if let Some(sender) = sender { + let _ = sender.send(status); + } + }); + } +} + +// --------------------------------------------------------------------------- +// Generic Unix reaper implementations (non-Linux or kernel < 5.3) +// --------------------------------------------------------------------------- + +#[inline] +#[cfg(all(unix, feature = "signal"))] +async fn zombie_reaper_fn_unix(rx: async_channel::Receiver) { + use futures_util::future::Either; + + let mut signal = crate::signal::Signal::new(crate::signal::SignalKind::child()) + .expect("cannot create signal handler for zombie reaper"); + let mut processes = Vec::new(); + loop { + let select = futures_util::future::select(Box::pin(rx.recv()), Box::pin(signal.recv())); + match select.await { + Either::Left((process, _)) => { + let Ok(mut process) = process else { + break; + }; + let try_wait = process.0.try_wait(); + match try_wait { + Ok(Some(exit_code)) => { + if let Some(sender) = process.1 { + let _ = sender.send(Ok(exit_code)); + } + } + Ok(None) => processes.push(process), + Err(err) => { + if let Some(sender) = process.1 { + let _ = sender.send(Err(err)); + } + } + } + } + Either::Right((signal, _)) => { + if signal.is_err() { + break; + }; + for mut process in processes.split_off(0) { + let try_wait = process.0.try_wait(); + match try_wait { + Ok(Some(exit_code)) => { + if let Some(sender) = process.1 { + let _ = sender.send(Ok(exit_code)); + } + } + Ok(None) => processes.push(process), + Err(err) => { + if let Some(sender) = process.1 { + let _ = sender.send(Err(err)); + } + } + } + } + } + } + } +} + +#[inline] +#[cfg(all(unix, not(feature = "signal")))] +async fn zombie_reaper_fn_unix(rx: async_channel::Receiver) { + while let Ok(mut msg) = rx.recv().await { + crate::spawn(async move { + crate::spawn_blocking(move || { + let result = msg.0.wait(); + if let Some(sender) = msg.1 { + let _ = sender.send(result); + } + }) + .await + }); + } +} diff --git a/vendor/vibeio/src/signal/mod.rs b/vendor/vibeio/src/signal/mod.rs new file mode 100644 index 0000000..c1401d9 --- /dev/null +++ b/vendor/vibeio/src/signal/mod.rs @@ -0,0 +1,32 @@ +//! Async signal utilities integrated with `vibeio`. +//! +//! This module provides cross-platform signal handling: +//! - `ctrl_c()` provides a future that resolves when Ctrl-C is received. +//! - On Unix, `Signal` and `signal()` allow listening for specific signals. +//! +//! # Examples +//! ```ignore +//! // Wait for Ctrl-C (cross-platform) +//! let _ = vibeio::signal::ctrl_c().await?; +//! +//! // Wait for SIGTERM (Unix only) +//! # #[cfg(unix)] +//! let mut sig = vibeio::signal::signal(vibeio::signal::SignalKind::terminate())?; +//! sig.recv().await?; +//! ``` +//! +//! # Implementation notes +//! - On Unix, signals are delivered via a dedicated dispatch thread that reads +//! from a pipe and wakes registered wakers. +//! - On Windows, Ctrl-C is handled via `SetConsoleCtrlHandler`. +//! - Multiple listeners can register for the same signal; all will be woken. + +#[cfg(unix)] +mod unix; +#[cfg(windows)] +mod windows; + +#[cfg(unix)] +pub use unix::{ctrl_c, signal, CtrlC, Signal, SignalKind}; +#[cfg(windows)] +pub use windows::{ctrl_c, CtrlC}; diff --git a/vendor/vibeio/src/signal/unix.rs b/vendor/vibeio/src/signal/unix.rs new file mode 100644 index 0000000..8b35366 --- /dev/null +++ b/vendor/vibeio/src/signal/unix.rs @@ -0,0 +1,526 @@ +//! Unix signal handling implementation for `vibeio`. +//! +//! This module provides async signal handling for Unix systems using a +//! dedicated dispatch thread and pipe-based communication. +//! +//! # Implementation details +//! - A background thread reads from a pipe connected to signal handlers. +//! - Signal handlers write the signal number to the pipe. +//! - The dispatch thread wakes registered wakers for received signals. +//! - Multiple listeners for the same signal share the same handler. + +use std::collections::HashMap; +use std::future::Future; +use std::io; +use std::os::fd::RawFd; +use std::pin::Pin; +use std::sync::atomic::{AtomicI32, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll, Waker}; + +use futures_util::future::poll_fn; +use once_cell::sync::OnceCell; + +/// Unix signal kind wrapper. +/// +/// Represents a Unix signal number. Common signal kinds are provided as +/// convenience methods: +/// - `SignalKind::interrupt()` - SIGINT (Ctrl-C) +/// - `SignalKind::terminate()` - SIGTERM +/// - `SignalKind::hangup()` - SIGHUP +/// - `SignalKind::quit()` - SIGQUIT +/// - `SignalKind::user_defined1()` - SIGUSR1 +/// - `SignalKind::user_defined2()` - SIGUSR2 +/// - `SignalKind::child()` - SIGCHLD +/// - `SignalKind::alarm()` - SIGALRM +/// - `SignalKind::pipe()` - SIGPIPE +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)] +pub struct SignalKind(libc::c_int); + +impl SignalKind { + /// Create a new `SignalKind` from a raw signal number. + #[inline] + pub const fn new(raw: libc::c_int) -> Self { + Self(raw) + } + + /// Return the raw signal number. + #[inline] + pub const fn as_raw(self) -> libc::c_int { + self.0 + } + + /// SIGINT - interrupt signal (Ctrl-C). + #[inline] + pub const fn interrupt() -> Self { + Self(libc::SIGINT) + } + + /// SIGTERM - termination signal. + #[inline] + pub const fn terminate() -> Self { + Self(libc::SIGTERM) + } + + /// SIGHUP - hangup signal. + #[inline] + pub const fn hangup() -> Self { + Self(libc::SIGHUP) + } + + /// SIGQUIT - quit signal. + #[inline] + pub const fn quit() -> Self { + Self(libc::SIGQUIT) + } + + /// SIGUSR1 - user-defined signal 1. + #[inline] + pub const fn user_defined1() -> Self { + Self(libc::SIGUSR1) + } + + /// SIGUSR2 - user-defined signal 2. + #[inline] + pub const fn user_defined2() -> Self { + Self(libc::SIGUSR2) + } + + /// SIGCHLD - child process terminated or stopped. + #[inline] + pub const fn child() -> Self { + Self(libc::SIGCHLD) + } + + /// SIGALRM - alarm clock signal. + #[inline] + pub const fn alarm() -> Self { + Self(libc::SIGALRM) + } + + /// SIGPIPE - write to pipe with no readers. + #[inline] + pub const fn pipe() -> Self { + Self(libc::SIGPIPE) + } +} + +impl From for libc::c_int { + #[inline] + fn from(kind: SignalKind) -> libc::c_int { + kind.0 + } +} + +struct SignalState { + counter: AtomicUsize, + wakers: Mutex>, +} + +struct RegisteredSignal { + state: Arc, + refs: usize, + prev_action: libc::sigaction, +} + +struct Registry { + signals: Mutex>, +} + +static REGISTRY: OnceCell> = OnceCell::new(); +static SIGNAL_WRITE_FD: AtomicI32 = AtomicI32::new(-1); + +/// Async signal listener for a specific Unix signal. +/// +/// This type listens for occurrences of a specific signal. Multiple `Signal` +/// instances can be created for the same signal kind; all will be woken when +/// the signal is received. +/// +/// # Examples +/// ```ignore +/// let mut sig = Signal::new(SignalKind::terminate())?; +/// sig.recv().await?; // Wait for SIGTERM +/// ``` +pub struct Signal { + kind: SignalKind, + state: Arc, + last_seen: usize, +} + +impl Signal { + /// Register for a Unix signal. + /// + /// Creates a new signal listener for the given signal kind. If this is the + /// first listener for this signal, the signal handler will be installed. + pub fn new(kind: SignalKind) -> io::Result { + let state = register_signal(kind)?; + let last_seen = state.counter.load(Ordering::Acquire); + Ok(Self { + kind, + state, + last_seen, + }) + } + + /// Returns the signal kind being listened to. + #[inline] + pub fn kind(&self) -> SignalKind { + self.kind + } + + /// Wait for the next occurrence of the signal. + /// + /// This method returns a future that resolves when the signal is received. + /// Multiple listeners for the same signal will all be woken on each signal. + pub async fn recv(&mut self) -> io::Result<()> { + poll_fn(|cx| self.poll_recv(cx)).await + } + + fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + let current = self.state.counter.load(Ordering::Acquire); + if current != self.last_seen { + self.last_seen = current; + return Poll::Ready(Ok(())); + } + + register_waker(&self.state, cx.waker()); + Poll::Pending + } +} + +impl Drop for Signal { + fn drop(&mut self) { + unregister_signal(self.kind); + } +} + +/// Convenience builder for Unix signals. +/// +/// This is a wrapper around `Signal::new()` that provides a more ergonomic API. +#[inline] +pub fn signal(kind: SignalKind) -> io::Result { + Signal::new(kind) +} + +/// Cross-platform Ctrl-C future (Unix implementation uses SIGINT). +/// +/// This type provides a future that resolves when Ctrl-C (SIGINT) is received. +/// On Unix, this is implemented as a `Signal` for `SignalKind::interrupt()`. +pub struct CtrlC { + signal: Signal, +} + +impl CtrlC { + /// Create a new Ctrl-C listener. + pub fn new() -> io::Result { + Ok(Self { + signal: Signal::new(SignalKind::interrupt())?, + }) + } +} + +impl Future for CtrlC { + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // SAFETY: CtrlC is not self-referential. + let this = unsafe { self.get_unchecked_mut() }; + this.signal.poll_recv(cx) + } +} + +/// Cross-platform Ctrl-C support. +/// +/// Returns a future that resolves when Ctrl-C is received. +#[inline] +pub fn ctrl_c() -> io::Result { + CtrlC::new() +} + +fn register_waker(state: &SignalState, waker: &Waker) { + let mut wakers = state.wakers.lock().unwrap(); + if let Some(existing) = wakers.iter_mut().find(|existing| existing.will_wake(waker)) { + *existing = waker.clone(); + } else { + wakers.push(waker.clone()); + } +} + +fn register_signal(kind: SignalKind) -> io::Result> { + let registry = registry()?; + let mut signals = registry.signals.lock().unwrap(); + if let Some(entry) = signals.get_mut(&kind.0) { + entry.refs += 1; + return Ok(entry.state.clone()); + } + + let prev_action = unsafe { install_handler(kind.0)? }; + let state = Arc::new(SignalState { + counter: AtomicUsize::new(0), + wakers: Mutex::new(Vec::new()), + }); + + signals.insert( + kind.0, + RegisteredSignal { + state: state.clone(), + refs: 1, + prev_action, + }, + ); + + Ok(state) +} + +fn unregister_signal(kind: SignalKind) { + let registry = match registry() { + Ok(registry) => registry, + Err(_) => return, + }; + + let prev_action = { + let mut signals = registry.signals.lock().unwrap(); + let entry = match signals.get_mut(&kind.0) { + Some(entry) => entry, + None => return, + }; + + if entry.refs > 1 { + entry.refs -= 1; + return; + } + + let prev = std::mem::replace(&mut entry.prev_action, unsafe { std::mem::zeroed() }); + signals.remove(&kind.0); + prev + }; + + unsafe { + let _ = restore_handler(kind.0, &prev_action); + } +} + +fn registry() -> io::Result<&'static Arc> { + REGISTRY.get_or_try_init(init_registry) +} + +fn init_registry() -> io::Result> { + let (read_fd, write_fd) = create_pipe()?; + SIGNAL_WRITE_FD.store(write_fd, Ordering::Release); + + let registry = Arc::new(Registry { + signals: Mutex::new(HashMap::new()), + }); + + start_dispatch_thread(read_fd, Arc::clone(®istry))?; + Ok(registry) +} + +fn start_dispatch_thread(read_fd: RawFd, registry: Arc) -> io::Result<()> { + std::thread::Builder::new() + .name("vibeio-signal-dispatch".to_string()) + .spawn(move || dispatch_loop(read_fd, registry)) + .map_err(io::Error::other)?; + Ok(()) +} + +fn dispatch_loop(read_fd: RawFd, registry: Arc) { + let mut buf = [0u8; 128]; + let mut pending = Vec::with_capacity(4); + loop { + let n = unsafe { libc::read(read_fd, buf.as_mut_ptr().cast::(), buf.len()) }; + if n == 0 { + continue; + } + if n < 0 { + let err = io::Error::last_os_error(); + match err.kind() { + io::ErrorKind::Interrupted => continue, + io::ErrorKind::WouldBlock => { + std::thread::sleep(std::time::Duration::from_millis(10)); + continue; + } + _ => continue, + } + } + + let n = n as usize; + for byte in &buf[..n] { + pending.push(*byte); + if pending.len() == 4 { + let mut raw = [0u8; 4]; + raw.copy_from_slice(&pending); + pending.clear(); + let signum = i32::from_ne_bytes(raw) as libc::c_int; + dispatch_signal(®istry, signum); + } + } + } +} + +fn dispatch_signal(registry: &Registry, signum: libc::c_int) { + let state = { + let signals = registry.signals.lock().unwrap(); + signals.get(&signum).map(|entry| entry.state.clone()) + }; + + if let Some(state) = state { + state.counter.fetch_add(1, Ordering::Release); + let wakers = { + let mut wakers = state.wakers.lock().unwrap(); + std::mem::take(&mut *wakers) + }; + for waker in wakers { + waker.wake(); + } + } +} + +extern "C" fn signal_handler(signum: libc::c_int) { + let fd = SIGNAL_WRITE_FD.load(Ordering::Relaxed); + if fd < 0 { + return; + } + + let bytes = signum.to_ne_bytes(); + unsafe { + let _ = libc::write(fd, bytes.as_ptr().cast::(), bytes.len()); + } +} + +unsafe fn install_handler(signum: libc::c_int) -> io::Result { + let mut action: libc::sigaction = std::mem::zeroed(); + action.sa_sigaction = signal_handler as extern "C" fn(libc::c_int) as usize; + action.sa_flags = libc::SA_RESTART; + libc::sigemptyset(&mut action.sa_mask); + + let mut prev: libc::sigaction = std::mem::zeroed(); + let rc = libc::sigaction(signum, &action, &mut prev); + if rc == -1 { + return Err(io::Error::last_os_error()); + } + Ok(prev) +} + +unsafe fn restore_handler(signum: libc::c_int, prev: &libc::sigaction) -> io::Result<()> { + let rc = libc::sigaction(signum, prev, std::ptr::null_mut()); + if rc == -1 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +fn create_pipe() -> io::Result<(RawFd, RawFd)> { + let mut fds = [0; 2]; + #[cfg(syscall_pipe2)] + let rc = unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC | libc::O_NONBLOCK) }; + #[cfg(not(syscall_pipe2))] + let rc = unsafe { libc::pipe(fds.as_mut_ptr()) }; + if rc == -1 { + return Err(io::Error::last_os_error()); + } + + // On non-Linux, set close-on-exec and non-blocking manually. + #[cfg(not(syscall_pipe2))] + { + if let Err(err) = set_cloexec(fds[0]).and_then(|_| set_cloexec(fds[1])) { + unsafe { + let _ = libc::close(fds[0]); + let _ = libc::close(fds[1]); + } + return Err(err); + } + + if let Err(err) = set_nonblock(fds[1]) { + unsafe { + let _ = libc::close(fds[0]); + let _ = libc::close(fds[1]); + } + return Err(err); + } + } + + Ok((fds[0], fds[1])) +} + +#[cfg(not(target_os = "linux"))] +fn set_cloexec(fd: RawFd) -> io::Result<()> { + let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; + if flags == -1 { + return Err(io::Error::last_os_error()); + } + let rc = unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) }; + if rc == -1 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +#[cfg(not(target_os = "linux"))] +fn set_nonblock(fd: RawFd) -> io::Result<()> { + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags == -1 { + return Err(io::Error::last_os_error()); + } + let rc = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) }; + if rc == -1 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::AnyDriver; + use std::time::Duration; + + #[cfg(feature = "time")] + async fn await_signal_with_timeout( + fut: impl Future>, + ) -> io::Result<()> { + crate::time::timeout(Duration::from_secs(1), fut) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "signal timeout"))? + } + + #[cfg(not(feature = "time"))] + async fn await_signal_with_timeout( + fut: impl Future>, + ) -> io::Result<()> { + fut.await + } + + fn spawn_signal_after_delay(signum: libc::c_int) { + let pid = unsafe { libc::getpid() }; + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(10)); + unsafe { + libc::kill(pid, signum); + } + }); + } + + #[test] + fn signal_recv_unblocks() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + let result = rt.block_on(async { + let mut sig = signal(SignalKind::user_defined1())?; + spawn_signal_after_delay(SignalKind::user_defined1().as_raw()); + await_signal_with_timeout(sig.recv()).await + }); + assert!(result.is_ok()); + } + + #[test] + fn ctrl_c_unblocks_on_sigint() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + let result = rt.block_on(async { + let ctrlc = ctrl_c()?; + + spawn_signal_after_delay(SignalKind::interrupt().as_raw()); + await_signal_with_timeout(ctrlc).await + }); + assert!(result.is_ok()); + } +} diff --git a/vendor/vibeio/src/signal/windows.rs b/vendor/vibeio/src/signal/windows.rs new file mode 100644 index 0000000..8bd5037 --- /dev/null +++ b/vendor/vibeio/src/signal/windows.rs @@ -0,0 +1,154 @@ +//! Windows signal handling implementation for `vibeio`. +//! +//! This module provides Ctrl-C support for Windows systems using +//! `SetConsoleCtrlHandler`. +//! +//! # Implementation details +//! - Ctrl-C events are handled via the Windows console control handler. +//! - The handler updates a counter and wakes registered wakers. +//! - Only Ctrl-C is supported on Windows (no arbitrary signals). + +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll, Waker}; + +use once_cell::sync::OnceCell; +use windows_sys::Win32::System::Console::{SetConsoleCtrlHandler, CTRL_C_EVENT}; + +struct CtrlCState { + counter: AtomicUsize, + wakers: Mutex>, +} + +static CTRL_C_STATE: OnceCell> = OnceCell::new(); + +/// Cross-platform Ctrl-C future (Windows implementation). +/// +/// This type provides a future that resolves when Ctrl-C is received. +/// On Windows, this is implemented via `SetConsoleCtrlHandler`. +pub struct CtrlC { + state: Arc, + last_seen: usize, +} + +impl CtrlC { + /// Create a new Ctrl-C listener. + pub fn new() -> io::Result { + let state = ctrl_c_state()?.clone(); + let last_seen = state.counter.load(Ordering::Acquire); + Ok(Self { state, last_seen }) + } + + fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + let current = self.state.counter.load(Ordering::Acquire); + if current != self.last_seen { + self.last_seen = current; + return Poll::Ready(Ok(())); + } + + register_waker(&self.state, cx.waker()); + Poll::Pending + } +} + +impl Future for CtrlC { + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // SAFETY: CtrlC is not self-referential. + let this = unsafe { self.get_unchecked_mut() }; + this.poll_recv(cx) + } +} + +/// Cross-platform Ctrl-C support. +/// +/// Returns a future that resolves when Ctrl-C is received. +#[inline] +pub fn ctrl_c() -> io::Result { + CtrlC::new() +} + +fn register_waker(state: &CtrlCState, waker: &Waker) { + let mut wakers = state.wakers.lock().unwrap(); + if let Some(existing) = wakers.iter_mut().find(|existing| existing.will_wake(waker)) { + *existing = waker.clone(); + } else { + wakers.push(waker.clone()); + } +} + +fn ctrl_c_state() -> io::Result<&'static Arc> { + CTRL_C_STATE.get_or_try_init(|| { + let state = Arc::new(CtrlCState { + counter: AtomicUsize::new(0), + wakers: Mutex::new(Vec::new()), + }); + + let ok = unsafe { SetConsoleCtrlHandler(Some(ctrl_c_handler), 1) }; + if ok == 0 { + return Err(io::Error::last_os_error()); + } + + Ok(state) + }) +} + +unsafe extern "system" fn ctrl_c_handler(ctrl_type: u32) -> i32 { + if ctrl_type == CTRL_C_EVENT { + if let Some(state) = CTRL_C_STATE.get() { + state.counter.fetch_add(1, Ordering::Release); + let wakers = { + let mut wakers = state.wakers.lock().unwrap(); + std::mem::take(&mut *wakers) + }; + for waker in wakers { + waker.wake(); + } + } + return 1; + } + 0 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::AnyDriver; + use std::time::Duration; + + #[cfg(feature = "time")] + async fn await_ctrl_c_with_timeout( + fut: impl Future>, + ) -> io::Result<()> { + crate::time::timeout(Duration::from_secs(1), fut) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "ctrl-c timeout"))? + } + + #[cfg(not(feature = "time"))] + async fn await_ctrl_c_with_timeout( + fut: impl Future>, + ) -> io::Result<()> { + fut.await + } + + #[test] + fn ctrl_c_unblocks_on_handler() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + let result = rt.block_on(async { + let ctrlc = ctrl_c()?; + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(10)); + unsafe { + let _ = ctrl_c_handler(CTRL_C_EVENT); + } + }); + await_ctrl_c_with_timeout(ctrlc).await + }); + assert!(result.is_ok()); + } +} diff --git a/vendor/vibeio/src/task.rs b/vendor/vibeio/src/task.rs new file mode 100644 index 0000000..71fee74 --- /dev/null +++ b/vendor/vibeio/src/task.rs @@ -0,0 +1,124 @@ +use std::cell::{RefCell, UnsafeCell}; +use std::collections::VecDeque; +use std::rc::Weak; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{RawWaker, RawWakerVTable, Waker}; + +use crossbeam_queue::SegQueue; +use futures_util::future::LocalBoxFuture; + +use crate::driver::AnyInterruptor; + +pub struct Task { + pub future: RefCell>>, + pub queue: Weak>>>, + pub next_task: Weak>>>, + pub remote_queue: std::sync::Weak>, + pub interruptor: AnyInterruptor, + pub queued: AtomicBool, + pub thread_id: std::thread::ThreadId, + pub token: usize, + pub waiting: std::sync::Weak, + pub interrupt_pending: std::sync::Weak, +} + +impl Task { + #[inline] + pub fn waker(self: &Arc) -> Waker { + // SAFETY: the vtable methods correctly clone/drop the Arc reference count. + unsafe { Waker::from_raw(Self::raw_waker(Arc::into_raw(Arc::clone(self)) as *const ())) } + } + + #[inline] + unsafe fn raw_waker(ptr: *const ()) -> RawWaker { + RawWaker::new(ptr, &Self::VTABLE) + } + + const VTABLE: RawWakerVTable = RawWakerVTable::new( + Self::raw_waker_clone, + Self::raw_waker_wake, + Self::raw_waker_wake_by_ref, + Self::raw_waker_drop, + ); + + #[inline] + unsafe fn raw_waker_clone(ptr: *const ()) -> RawWaker { + let task = Arc::::from_raw(ptr as *const Self); + let cloned = Arc::clone(&task); + let _ = Arc::into_raw(task); + Self::raw_waker(Arc::into_raw(cloned) as *const ()) + } + + #[inline] + unsafe fn raw_waker_wake(ptr: *const ()) { + let task = Arc::::from_raw(ptr as *const Self); + Self::enqueue_if_needed(&task); + } + + #[inline] + unsafe fn raw_waker_wake_by_ref(ptr: *const ()) { + let task = Arc::::from_raw(ptr as *const Self); + Self::enqueue_if_needed(&task); + let _ = Arc::into_raw(task); + } + + #[inline] + unsafe fn raw_waker_drop(ptr: *const ()) { + drop(Arc::::from_raw(ptr as *const Self)); + } + + #[inline] + fn enqueue_if_needed(task: &Arc) { + if std::thread::current().id() == task.thread_id { + if !task.queued.swap(true, Ordering::Relaxed) { + let mut pushed_next = false; + if let Some(next_task) = task.next_task.upgrade() { + let mut next_task = next_task.borrow_mut(); + if next_task.is_none() { + *next_task = Some(Arc::clone(task)); + pushed_next = true; + } + } + if !pushed_next { + if let Some(queue) = task.queue.upgrade() { + // SAFETY: the runtime is single-threaded and only mutates the ready + // queue from that thread. We also never hold a mutable queue borrow + // while polling task futures, so re-entrant wakes do not alias. + unsafe { + (&mut *queue.get()).push_back(Arc::clone(task)); + } + } + } + } + return; + } + + if !task.queued.swap(true, Ordering::Relaxed) { + if let Some(remote_queue) = task.remote_queue.upgrade() { + remote_queue.push(task.token); + } + } + + // Interrupt the driver if it's waiting + if task + .waiting + .upgrade() + .is_some_and(|waiting| waiting.load(Ordering::Acquire)) + { + let should_interrupt = task + .interrupt_pending + .upgrade() + .is_none_or(|pending| !pending.swap(true, Ordering::AcqRel)); + if should_interrupt { + // Interrupt the driver if the waker is not on the same thread as the runtime + task.interruptor.interrupt(); + } + } + } + + #[inline] + pub fn mark_dequeued(&self) { + self.queued.store(false, Ordering::Relaxed); + } +} diff --git a/vendor/vibeio/src/time/interval.rs b/vendor/vibeio/src/time/interval.rs new file mode 100644 index 0000000..0f16db5 --- /dev/null +++ b/vendor/vibeio/src/time/interval.rs @@ -0,0 +1,145 @@ +//! Interval for periodic tick scheduling with drift compensation. + +use std::time::{Duration, Instant}; + +use super::sleep::Sleep; + +/// How to handle missed ticks for `Interval`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MissedTickBehavior { + /// Skip missed ticks and schedule the next tick at the next future multiple. + Skip, + /// Return the number of missed ticks so the caller can run catch-up loop. + CatchUp, +} + +/// Interval provides a simple API to await periodic ticks while maintaining +/// a steady cadence (accounting for drift): each tick tries to occur at an +/// integer multiple of the original period relative to the interval start. +/// +/// Example: +/// ```ignore +/// let mut interval = Interval::new(Duration::from_millis(100)); +/// loop { +/// let ticks = interval.tick().await; +/// for _ in 0..ticks { /* handle tick */ } +/// } +/// ``` +pub struct Interval { + period: Duration, + /// The next absolute deadline for the interval. If `None`, the next tick + /// will be scheduled relative to the current time. + pub next_deadline: Option, + /// Behavior when ticks are missed. + missed_tick_behavior: MissedTickBehavior, +} + +impl Interval { + #[inline] + pub fn new(period: Duration) -> Self { + Self { + period, + next_deadline: None, + missed_tick_behavior: MissedTickBehavior::Skip, + } + } + + /// Configure how missed ticks are handled. + #[inline] + pub fn set_missed_tick_behavior(&mut self, behavior: MissedTickBehavior) { + self.missed_tick_behavior = behavior; + } + + /// Reset the interval schedule so the next tick is computed relative to + /// the time when `tick()` is next called (useful when you want to restart + /// the cadence). + #[inline] + pub fn reset(&mut self) { + self.next_deadline = None; + } + + /// Await the next tick. Returns the number of ticks that should be processed: + /// - For `MissedTickBehavior::Skip` this will be `1`. + /// - For `MissedTickBehavior::CatchUp` this may be `> 1` if several periods were missed. + pub async fn tick(&mut self) -> u64 { + let now = Instant::now(); + + // Determine base next (the previous next_deadline or now+period) + let base_next = self.next_deadline.unwrap_or_else(|| now + self.period); + + match self.missed_tick_behavior { + MissedTickBehavior::Skip => { + // Advance target forward until it's in the future. + let mut target = base_next; + if target <= now { + if self.period.as_nanos() == 0 { + target = now; + } else { + // Advance by whole periods until target > now. + // Use a loop; the number of iterations is typically small. + let mut cnt: u64 = 0; + while target <= now { + target += self.period; + cnt = cnt.saturating_add(1); + // Safety: in pathological cases we break after huge iterations, + // but such large loops are unlikely in practice. + if cnt == u64::MAX { + break; + } + } + } + } + + let remaining = target.saturating_duration_since(Instant::now()); + if remaining.as_millis() == 0 { + // For immediate remaining, yield once to avoid busy-looping. + Sleep::new_with_zero_behavior(remaining, super::sleep::ZeroBehavior::Yield) + .await; + } else { + Sleep::new(remaining).await; + } + + // Schedule next deadline for subsequent tick + self.next_deadline = Some(target + self.period); + 1 + } + MissedTickBehavior::CatchUp => { + if base_next > now { + // Not missed yet: sleep until base_next and return 1 + let remaining = base_next.saturating_duration_since(Instant::now()); + if remaining.as_millis() == 0 { + Sleep::new_with_zero_behavior(remaining, super::sleep::ZeroBehavior::Yield) + .await; + } else { + Sleep::new(remaining).await; + } + self.next_deadline = Some(base_next + self.period); + 1 + } else { + // We missed one or more ticks. Compute how many. + if self.period.as_nanos() == 0 { + // If period is zero, avoid infinite loops; treat as a single tick. + self.next_deadline = Some(Instant::now() + self.period); + return 1; + } + + let elapsed = now.duration_since(base_next); + let missed = (elapsed.as_nanos() / self.period.as_nanos()) as u64 + 1; + + // Advance base_next by `missed` periods to compute the next deadline. + let mut new_next = base_next; + // loop to add period missed times + for _ in 0..missed { + new_next += self.period; + } + + // Set next_deadline to the instant after the catch-up window. + self.next_deadline = Some(new_next + self.period); + + // Return the number of missed ticks so caller can catch up. + missed + } + } + } + } +} diff --git a/vendor/vibeio/src/time/mod.rs b/vendor/vibeio/src/time/mod.rs new file mode 100644 index 0000000..d30d939 --- /dev/null +++ b/vendor/vibeio/src/time/mod.rs @@ -0,0 +1,130 @@ +//! A small `time` utility module for `vibeio`. +//! +//! This module provides: +//! - `Sleep`: a Future that completes after a duration. +//! - `Interval`: a convenience type with a `tick().await` method to await periodic ticks. +//! - `timeout`: a function / `Timeout` future that races a future against a timeout. +//! +//! Implementation notes: +//! - The runtime's `Timer` driver (in `crate::timer`) is accessed through +//! `crate::executor::current_timer()` which returns an `Rc` when called +//! from inside a runtime. Calling these time utilities outside a runtime will +//! panic (matching the library's general behavior for runtime-only APIs). +//! - The `Timer` driver accepts a `Waker` and returns an optional `TimerHandle`. +//! We store the `TimerHandle` and cancel it if the Sleep is dropped before firing. + +mod interval; +mod sleep; +mod timeout; + +// Re-export public types and functions +pub use interval::{Interval, MissedTickBehavior}; +pub use sleep::{Sleep, ZeroBehavior}; +pub use timeout::{timeout, Timeout, TimeoutError}; + +/// Convenience builder: returns a `Sleep` future. +#[inline] +pub fn sleep(duration: std::time::Duration) -> Sleep { + Sleep::new(duration) +} + +/// Convenience builder allowing zero-behavior control for tiny durations. +#[inline] +pub fn sleep_with_zero_behavior(duration: std::time::Duration, behavior: ZeroBehavior) -> Sleep { + Sleep::new_with_zero_behavior(duration, behavior) +} + +/// Convenience builder: returns an `Interval`. +#[inline] +pub fn interval(period: std::time::Duration) -> Interval { + Interval::new(period) +} + +/// Convenience builder: returns a `Sleep` that completes at the provided absolute `Instant`. +#[inline] +pub fn sleep_until(deadline: std::time::Instant) -> Sleep { + Sleep::sleep_until(deadline) +} + +/// Convenience async function that awaits `future` but returns an error if it +/// does not complete before the absolute `deadline` Instant. +#[inline] +pub async fn timeout_at( + deadline: std::time::Instant, + future: impl std::future::Future, +) -> Result { + let dur = deadline.saturating_duration_since(std::time::Instant::now()); + timeout(dur, future).await +} + +/// Convenience builder: returns an `Interval` with the first tick scheduled to +/// complete at `first_tick_instant` and subsequent ticks every `period`. +#[inline] +pub fn interval_at( + first_tick_instant: std::time::Instant, + period: std::time::Duration, +) -> Interval { + let mut iv = Interval::new(period); + iv.next_deadline = Some(first_tick_instant); + iv +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::AnyDriver; + use std::time::{Duration, Instant}; + + #[test] + fn sleep_completes() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + rt.block_on(async { + Sleep::new(Duration::from_millis(1)).await; + }); + } + + #[test] + fn timeout_expires() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + let res = rt.block_on(async { + let never = async { + futures_util::future::pending::<()>().await; + }; + timeout(Duration::from_millis(1), never).await + }); + assert!(res.is_err()); + } + + #[test] + fn timeout_succeeds_if_future_completes() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + let res = rt.block_on(async { timeout(Duration::from_secs(1), async { 123usize }).await }); + assert_eq!(res, Ok(123usize)); + } + + #[test] + fn interval_ticks_skip() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + rt.block_on(async { + let mut interval = Interval::new(Duration::from_millis(1)); + // two ticks should complete quickly + let _ = interval.tick().await; + let _ = interval.tick().await; + }); + } + + #[test] + fn interval_catchup_returns_multiple() { + let rt = crate::executor::Runtime::new(AnyDriver::new_mock()); + rt.block_on(async { + let mut interval = Interval::new(Duration::from_millis(1)); + interval.set_missed_tick_behavior(MissedTickBehavior::CatchUp); + // Initialise baseline + let _ = interval.tick().await; + // Artificially push next_deadline into the past to simulate missed ticks + interval.next_deadline = Some(Instant::now() - Duration::from_millis(10)); + let missed = interval.tick().await; + assert!(missed >= 1); + }); + } +} diff --git a/vendor/vibeio/src/time/sleep.rs b/vendor/vibeio/src/time/sleep.rs new file mode 100644 index 0000000..1cf5d65 --- /dev/null +++ b/vendor/vibeio/src/time/sleep.rs @@ -0,0 +1,209 @@ +//! Sleep future and zero-duration behavior configuration. + +use std::{ + cell::Cell, + future::Future, + pin::Pin, + rc::Rc, + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +use crate::executor::current_timer; +use crate::timer::{Timer, TimerHandle}; + +/// Behavior for zero-duration sleeps (duration < 1 ms). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ZeroBehavior { + /// Complete immediately (default). + Immediate, + /// Yield to the scheduler once, then complete on the following poll. + Yield, +} + +/// Sleep is a future that completes after the given `Duration`. +/// +/// Usage: +/// ```ignore +/// Sleep::new(Duration::from_millis(100)).await; +/// ``` +pub struct Sleep { + /// The timer handle returned by the timer driver when the timer was scheduled. + /// `None` means we haven't scheduled yet. + handle: Option, + /// Whether the timer has already fired / completed. + fired: Cell, + /// For zero-duration sleeps we may schedule a one-shot yield; track whether + /// that yield has been scheduled so the subsequent poll completes. + yield_scheduled: Cell, + /// How to behave for durations below the timer resolution (less than 1ms). + zero_behavior: ZeroBehavior, + /// The absolute deadline when this sleep should complete. + deadline: Instant, + /// The timer used to schedule the sleep. + timer: Option>, +} + +impl Sleep { + /// Create a new Sleep instance for the provided `duration`. + #[inline] + pub fn new(duration: Duration) -> Self { + Self { + handle: None, + fired: Cell::new(false), + yield_scheduled: Cell::new(false), + zero_behavior: ZeroBehavior::Immediate, + deadline: Instant::now() + duration, + timer: None, + } + } + + /// Create a Sleep with custom behavior for zero-length waits. + #[inline] + pub fn new_with_zero_behavior(duration: Duration, zero_behavior: ZeroBehavior) -> Self { + Self { + handle: None, + fired: Cell::new(false), + yield_scheduled: Cell::new(false), + zero_behavior, + deadline: Instant::now() + duration, + timer: None, + } + } + + /// Create a Sleep that completes at the specified absolute `deadline`. + /// + /// This is a convenience constructor that computes the relative duration + /// from `Instant::now()` to the provided `deadline`. + #[inline] + pub fn sleep_until(deadline: Instant) -> Self { + Self { + handle: None, + fired: Cell::new(false), + yield_scheduled: Cell::new(false), + zero_behavior: ZeroBehavior::Immediate, + deadline, + timer: None, + } + } + + /// Reset the sleep to a new absolute `deadline` (`Instant`). + /// + /// If a timer was previously scheduled, cancel it. The timer will be + /// rescheduled on the next `poll`. This allows reusing a `Sleep` value + /// to implement steady intervals or dynamic timeout adjustments using + /// absolute deadlines instead of relative durations. + #[inline] + pub fn reset(&mut self, deadline: Instant) { + // Compute relative duration until deadline and update state. + self.deadline = deadline; + self.fired.set(false); + self.yield_scheduled.set(false); + + // If there was an outstanding handle, cancel it so the timer won't + // keep the old waker alive. + if let Some(handle) = self.handle.take() { + if let Some(timer_rc) = current_timer() { + timer_rc.cancel(handle); + } + } + } +} + +impl Future for Sleep { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // Safety: Sleep is !self-referential; it's safe to get a mutable reference. + let this = unsafe { self.get_unchecked_mut() }; + + if this.fired.get() { + return Poll::Ready(()); + } + + if this.handle.is_none() { + // Schedule with runtime timer. + let timer_rc = this.timer.get_or_insert_with(|| { + current_timer().expect("Sleep::poll called outside of runtime") + }); + + // The timer driver expects a task `Waker`. We clone the task waker here. + let waker = cx.waker().clone(); + match timer_rc.submit(this.deadline, waker) { + Some(handle) => { + this.handle = Some(handle); + // Not fired yet. + Poll::Pending + } + None => { + // Timer driver woke us immediately (duration rounded to 0 or similar). + match this.zero_behavior { + ZeroBehavior::Immediate => { + this.fired.set(true); + Poll::Ready(()) + } + ZeroBehavior::Yield => { + // If we haven't scheduled the one-shot yield yet, schedule it + // by waking ourselves and return Pending. On the subsequent + // poll we will observe `yield_scheduled` and complete. + if !this.yield_scheduled.replace(true) { + cx.waker().wake_by_ref(); + Poll::Pending + } else { + this.fired.set(true); + Poll::Ready(()) + } + } + } + } + } + } else { + // We were previously scheduled, and now we've been polled again. + // The runtime's timer driver will wake the task by calling the task + // waker when the timer expires. + + // Check if the deadline has actually been reached + if Instant::now() >= this.deadline { + this.fired.set(true); + // Drop the handle (we'll also attempt to cancel it in Drop if it remains). + this.handle = None; + Poll::Ready(()) + } else { + // Spurious wakeup, we need to wait more. + // The timer might have woken us up early, or another waker woke the task. + // Re-register the waker to ensure we get woken up again. + if let Some(handle) = this.handle.take() { + if let Some(timer_rc) = current_timer() { + timer_rc.cancel(handle); + let waker = cx.waker().clone(); + match timer_rc.submit(this.deadline, waker) { + Some(handle) => { + this.handle = Some(handle); + // Not fired yet. + return Poll::Pending; + } + None => { + // Timer driver woke us immediately (duration rounded to 0 or similar). + this.fired.set(true); + return Poll::Ready(()); + } + } + } + } + Poll::Pending + } + } + } +} + +impl Drop for Sleep { + fn drop(&mut self) { + // If we still have an outstanding timer handle, cancel it so the timer + // won't hold onto our waker. + if let Some(handle) = self.handle.take() { + if let Some(timer_rc) = self.timer.take() { + timer_rc.cancel(handle); + } + } + } +} diff --git a/vendor/vibeio/src/time/timeout.rs b/vendor/vibeio/src/time/timeout.rs new file mode 100644 index 0000000..c7bf763 --- /dev/null +++ b/vendor/vibeio/src/time/timeout.rs @@ -0,0 +1,104 @@ +//! Timeout future that races an inner future against a duration. + +use std::{ + fmt, + future::Future, + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; + +use super::sleep::Sleep; + +/// Error returned when a `timeout` expires. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TimeoutError; + +impl fmt::Display for TimeoutError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "operation timed out") + } +} + +impl std::error::Error for TimeoutError {} + +/// Timeout future that races the provided future against a timeout duration. +/// +/// If the inner future completes first, `Timeout` yields `Ok(T)`. +/// If the timeout elapses first, `Timeout` yields `Err(TimeoutError)` and the +/// inner future is dropped. +pub struct Timeout { + future: F, + sleep: Sleep, + /// If true, the timeout has already fired (and future should be treated as timed out). + timed_out: bool, +} + +impl Timeout { + /// Create a new `Timeout` future. + #[inline] + pub fn new(future: F, duration: Duration) -> Self { + Self { + future, + sleep: Sleep::new(duration), + timed_out: false, + } + } +} + +impl Future for Timeout +where + F: Future, +{ + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // Safety: we don't move the fields out; we only create pinned projections. + let this = unsafe { self.get_unchecked_mut() }; + + if this.timed_out { + return Poll::Ready(Err(TimeoutError)); + } + + // First, poll the inner future. If it completes, return the value and + // allow `sleep` to be dropped (which cancels the timer). + // We need a pinned projection of `future`. + let mut future_pin = unsafe { Pin::new_unchecked(&mut this.future) }; + match Future::poll(future_pin.as_mut(), cx) { + Poll::Ready(output) => return Poll::Ready(Ok(output)), + Poll::Pending => {} + } + + // Next, poll the timeout sleep. If it is ready, mark timed out and + // return Err. Otherwise remain Pending. + let mut sleep_pin = unsafe { Pin::new_unchecked(&mut this.sleep) }; + match sleep_pin.as_mut().poll(cx) { + Poll::Ready(()) => { + this.timed_out = true; + // Dropping `future` happens naturally when this Timeout is dropped + // by the executor or after we return. Return Err now. + Poll::Ready(Err(TimeoutError)) + } + Poll::Pending => Poll::Pending, + } + } +} + +/// Convenience async function that awaits `future` but returns an error if it +/// does not complete within `duration`. +/// +/// Example: +/// ```ignore +/// let res = timeout(Duration::from_secs(1), some_future).await; +/// match res { +/// Ok(v) => { /* completed in time */ } +/// Err(_) => { /* timed out */ } +/// } +/// ``` +#[inline] +pub async fn timeout( + duration: Duration, + future: impl Future, +) -> Result { + Timeout::new(future, duration).await +} diff --git a/vendor/vibeio/src/timer.rs b/vendor/vibeio/src/timer.rs new file mode 100644 index 0000000..3be6294 --- /dev/null +++ b/vendor/vibeio/src/timer.rs @@ -0,0 +1,791 @@ +use std::{ + cell::RefCell, + cmp::Reverse, + collections::BinaryHeap, + task::Waker, + time::{Duration, Instant}, +}; + +use slab::Slab; + +/// Number of levels in the hierarchical timing wheel +const NUM_LEVELS: usize = 7; +/// Number of slots per level +const SLOTS_PER_LEVEL: usize = 64; + +/// Unique identifier for a timer +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct TimerHandle { + slab_index: usize, + generation: u64, +} + +/// Entry stored in the slab for a timer +struct TimerEntry { + waker: Waker, + expiration_tick: u64, + generation: u64, + level: usize, + slot: usize, +} + +/// A hierarchical hashed timing wheel with 7 levels and 64 slots each. +/// One tick is equivalent to 1 millisecond. +/// +/// The timing wheel uses a hierarchical design where: +/// - Level 0: 64 slots, each representing 1ms (covers 0-64ms) +/// - Level 1: 64 slots, each representing 64ms (covers 64ms-4.1s) +/// - Level 2: 64 slots, each representing 4.1s (covers 4.1s-4.4min) +/// - Level 3: 64 slots, each representing 4.4min (covers 4.4min-4.7h) +/// - Level 4: 64 slots, each representing 4.7h (covers 4.7h-12.6days) +/// - Level 5: 64 slots, each representing 12.6days (covers 12.6days-2.2years) +/// - Level 6: 64 slots, each representing 2.2years (covers 2.2years-~139years) +struct TimingWheel { + /// Slab storage for all timer entries + entries: Slab, + /// The wheels for each level + wheels: [Vec>; NUM_LEVELS], + /// Current tick count (milliseconds since creation) + current_tick: u64, + /// Generation counter for handling timer reuse + generation_counter: u64, + /// The minimum expiration tick among all entries (for O(1) nearest_wakeup) + min_expiration_tick: Option, + /// Binary heap of (expiration_tick, slab_index, generation) for O(log n) + /// min-expiration tracking instead of O(n) scan on remove/expire. + expiration_heap: BinaryHeap>, +} + +impl TimingWheel { + #[inline] + pub fn new() -> Self { + Self { + entries: Slab::new(), + wheels: std::array::from_fn(|_| vec![Vec::new(); SLOTS_PER_LEVEL]), + current_tick: 0, + generation_counter: 0, + min_expiration_tick: None, + expiration_heap: BinaryHeap::new(), + } + } + + /// Check if the wheel is empty + #[inline] + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Get the current tick + #[inline] + pub fn now(&self) -> u64 { + self.current_tick + } + + /// Get the nearest wakeup tick (absolute), if any timer is scheduled + /// This is amortized O(1) because stale heap entries are lazily cleaned. + #[inline] + pub fn nearest_wakeup(&mut self) -> Option { + self.refresh_min_expiration(); + self.min_expiration_tick.and_then(std::num::NonZeroU64::new) + } + + /// Update the minimum expiration tick after an insert + #[inline] + fn update_min_on_insert(&mut self, expiration_tick: u64) { + self.min_expiration_tick = Some( + self.min_expiration_tick + .map_or(expiration_tick, |min| min.min(expiration_tick)), + ); + } + + /// Lazily refresh the minimum expiration tick by popping stale entries + /// from the binary heap. Amortized O(1) — only the heap top is inspected. + #[inline] + fn refresh_min_expiration(&mut self) { + while let Some(&Reverse((_tick, idx, gen))) = self.expiration_heap.peek() { + if let Some(entry) = self.entries.get(idx) { + if entry.generation == gen { + self.min_expiration_tick = Some(entry.expiration_tick); + return; + } + } + self.expiration_heap.pop(); + } + self.min_expiration_tick = None; + } + + /// Insert a timer that expires at the given tick + #[inline] + pub fn insert(&mut self, waker: Waker, expiration_tick: u64) -> TimerHandle { + let delay = expiration_tick.saturating_sub(self.current_tick); + let (level, slot) = self.calculate_level_and_slot(delay); + + self.generation_counter += 1; + let generation = self.generation_counter; + + let slab_index = self.entries.insert(TimerEntry { + waker, + expiration_tick, + generation, + level, + slot, + }); + + self.wheels[level][slot].push(slab_index); + self.expiration_heap + .push(Reverse((expiration_tick, slab_index, generation))); + self.update_min_on_insert(expiration_tick); + + TimerHandle { + slab_index, + generation, + } + } + + /// Remove a timer by its handle + #[inline] + pub fn remove(&mut self, handle: TimerHandle) { + if let Some(entry) = self.entries.get(handle.slab_index) { + if entry.generation == handle.generation { + let level = entry.level; + let slot = entry.slot; + let slab_index = handle.slab_index; + + // Remove from the wheel slot + if let Some(pos) = self.wheels[level][slot] + .iter() + .position(|&idx| idx == slab_index) + { + self.wheels[level][slot].remove(pos); + } + + // Remove from slab + self.entries.remove(slab_index); + self.refresh_min_expiration(); + } + } + } + + /// Advance the timing wheel by the given number of ticks + /// Returns a vector of wakers that have expired + /// This is optimized to only process slots that have timers, not every tick + pub fn advance(&mut self, ticks: u64) -> Vec { + if ticks == 0 { + return Vec::new(); + } + + let mut expired_wakers = Vec::new(); + let start_tick = self.current_tick; + self.current_tick += ticks; + + // First, cascade timers from higher levels down to level 0 + // This must happen before processing level 0 slots + self.cascade_all_levels(&mut expired_wakers); + + // Process all level 0 slots that fall within the range + // Level 0 has 64 slots, each representing 1ms + let start_slot = (start_tick as usize) & (SLOTS_PER_LEVEL - 1); + let end_slot = (self.current_tick as usize) & (SLOTS_PER_LEVEL - 1); + + if ticks >= SLOTS_PER_LEVEL as u64 { + // We've wrapped around at least once, process all slots + for slot in 0..SLOTS_PER_LEVEL { + self.process_slot_at_level_0(slot, &mut expired_wakers); + } + } else if start_slot <= end_slot { + // No wrap around, process slots in range + for slot in start_slot..=end_slot { + self.process_slot_at_level_0(slot, &mut expired_wakers); + } + } else { + // Wrapped around, process from start to end, then 0 to end + for slot in start_slot..SLOTS_PER_LEVEL { + self.process_slot_at_level_0(slot, &mut expired_wakers); + } + for slot in 0..=end_slot { + self.process_slot_at_level_0(slot, &mut expired_wakers); + } + } + + // Refresh the minimum expiration time after processing. + // The binary heap lazily cleans stale entries in O(log n). + self.refresh_min_expiration(); + + expired_wakers + } + + /// Cascade all timers from higher levels down through the hierarchy + /// This processes all slots at each level and moves timers to appropriate lower levels + fn cascade_all_levels(&mut self, expired_wakers: &mut Vec) { + // Process from highest level down to level 1 + for level in (1..NUM_LEVELS).rev() { + self.cascade_level(level, expired_wakers); + } + } + + /// Cascade all timers from a specific level down to lower levels + fn cascade_level(&mut self, level: usize, expired_wakers: &mut Vec) { + // Take all slots at this level + for slot in 0..SLOTS_PER_LEVEL { + let timers_to_cascade: Vec = std::mem::take(&mut self.wheels[level][slot]); + + for slab_index in timers_to_cascade { + if let Some(entry) = self.entries.get(slab_index) { + if entry.expiration_tick <= self.current_tick { + // Timer has expired + let entry = self.entries.remove(slab_index); + expired_wakers.push(entry.waker); + } else { + // Re-calculate the appropriate level and slot based on remaining delay + let delay = entry.expiration_tick - self.current_tick; + let (new_level, new_slot) = self.calculate_level_and_slot(delay); + + if let Some(entry) = self.entries.get_mut(slab_index) { + entry.level = new_level; + entry.slot = new_slot; + } + self.wheels[new_level][new_slot].push(slab_index); + } + } + } + } + } + + /// Process a specific slot at level 0 + fn process_slot_at_level_0(&mut self, slot: usize, expired_wakers: &mut Vec) { + let timers_to_process: Vec = std::mem::take(&mut self.wheels[0][slot]); + + for slab_index in timers_to_process { + if let Some(entry) = self.entries.get(slab_index) { + if entry.expiration_tick <= self.current_tick { + // Timer has expired, remove and wake + let entry = self.entries.remove(slab_index); + expired_wakers.push(entry.waker); + } else { + // Timer hasn't expired yet, re-insert at appropriate level + let delay = entry.expiration_tick - self.current_tick; + let (new_level, new_slot) = self.calculate_level_and_slot(delay); + + if let Some(entry) = self.entries.get_mut(slab_index) { + entry.level = new_level; + entry.slot = new_slot; + } + self.wheels[new_level][new_slot].push(slab_index); + } + } + } + } + + /// Calculate the appropriate level and slot for a given delay + #[inline] + fn calculate_level_and_slot(&self, delay: u64) -> (usize, usize) { + for level in 0..NUM_LEVELS { + let level_shift = 6 * level; + + if level == NUM_LEVELS - 1 { + // Top level: use the highest bits + let slot = ((delay >> level_shift) & (SLOTS_PER_LEVEL as u64 - 1)) as usize; + return (level, slot); + } + + // Check if delay fits in this level + let max_for_level = ((SLOTS_PER_LEVEL as u64) << level_shift) - 1; + if delay <= max_for_level || level == NUM_LEVELS - 1 { + let slot = ((delay >> level_shift) & (SLOTS_PER_LEVEL as u64 - 1)) as usize; + return (level, slot); + } + } + + // Fallback to top level + let level_shift = 6 * (NUM_LEVELS - 1); + let slot = ((delay >> level_shift) & (SLOTS_PER_LEVEL as u64 - 1)) as usize; + (NUM_LEVELS - 1, slot) + } +} + +impl Default for TimingWheel { + fn default() -> Self { + Self::new() + } +} + +pub struct Timer { + // The timer wheel will overflow after ~139 years of waiting, + // but running for that long is very unlikely. + wheel: RefCell, + instant: RefCell, +} + +impl Timer { + #[inline] + pub fn new() -> Self { + Self { + wheel: RefCell::new(TimingWheel::new()), + instant: RefCell::new(Instant::now()), + } + } + + #[inline] + pub fn submit(&self, deadline: Instant, waker: Waker) -> Option { + let millis = deadline + .saturating_duration_since(*self.instant.borrow()) + .as_millis() as u64; + if millis < 1 { + // If the duration is less than 1 millisecond, wake the task immediately + // One tick is equivalent to 1 millisecond. + waker.wake(); + return None; + } + let mut wheel = self.wheel.borrow_mut(); + let now = wheel.now(); + Some(wheel.insert(waker, now + millis)) + } + + #[inline] + pub fn cancel(&self, handle: TimerHandle) { + let mut wheel = self.wheel.borrow_mut(); + wheel.remove(handle); + } + + #[inline] + pub fn spin_and_get_deadline(&self) -> (Option, bool) { + let mut instant = self.instant.borrow_mut(); + let mut wheel = self.wheel.borrow_mut(); + let mut woken_up = false; + if !wheel.is_empty() { + // Advance the timer wheel, but only if not empty + for waker in wheel.advance(instant.elapsed().as_millis() as u64) { + // Wake the pending tasks + waker.wake(); + woken_up = true; + } + } + + *instant = Instant::now(); + + // The deadline is absolute, so we need to subtract the current time to get the relative deadline + // If there's a deadline, we need to ensure it's at least 1ms to avoid busy looping + let now = wheel.now(); + ( + wheel + .nearest_wakeup() + .map(|deadline| Duration::from_millis(deadline.get().saturating_sub(now).max(1))), + woken_up, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + use std::task::{RawWaker, RawWakerVTable, Waker}; + use std::time::Duration; + + fn mock_waker(counter: Arc>) -> Waker { + fn clone(data: *const ()) -> RawWaker { + RawWaker::new(data, &VTABLE) + } + fn wake(data: *const ()) { + let counter = unsafe { &*(data as *const Mutex) }; + let mut lock = counter.lock().unwrap(); + *lock += 1; + } + fn wake_by_ref(data: *const ()) { + wake(data) + } + fn drop(_: *const ()) {} + + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + + let ptr = Arc::into_raw(counter) as *const (); + unsafe { Waker::from_raw(RawWaker::new(ptr, &VTABLE)) } + } + + #[test] + fn test_timer_submit_and_deadline() { + let timer = Timer::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Submit a timer for 50 ms + let _handle = timer.submit(Instant::now() + Duration::from_millis(50), waker); + + // The deadline should be Some and close to 50 ms + let (deadline, _woken_up) = timer.spin_and_get_deadline(); + assert!(deadline.is_some()); + let ms = deadline.unwrap().as_millis(); + assert!(ms <= 50 && ms > 0, "Deadline should be <= 50ms, got {}", ms); + } + + #[test] + fn test_timer_cancel() { + let timer = Timer::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Submit a timer and then cancel it + let handle = timer + .submit(Instant::now() + Duration::from_millis(100), waker) + .expect("Failed to submit timer"); + timer.cancel(handle); + + // After cancel, deadline should be None + let (deadline, _woken_up) = timer.spin_and_get_deadline(); + assert!(deadline.is_none()); + } + + #[test] + fn test_timer_wake() { + let timer = Timer::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Submit a timer for 1 ms + let _handle = timer.submit(Instant::now() + Duration::from_millis(1), waker); + + // Wait a bit and spin the wheel + std::thread::sleep(Duration::from_millis(5)); + timer.spin_and_get_deadline(); + + // The waker should have been called + let count = counter.lock().unwrap(); + assert!( + *count >= 1, + "Waker should have been called or at least not panic" + ); + } + + // ==================== TimingWheel Tests ==================== + + #[test] + fn test_timing_wheel_empty() { + let mut wheel = TimingWheel::new(); + assert!(wheel.is_empty()); + assert_eq!(wheel.now(), 0); + assert!(wheel.nearest_wakeup().is_none()); + } + + #[test] + fn test_timing_wheel_insert_and_expire_single() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Insert timer for 10 ticks + let _handle = wheel.insert(waker, 10); + assert!(!wheel.is_empty()); + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(10).unwrap()) + ); + + // Advance to tick 9 - should not expire + let expired = wheel.advance(9); + assert!(expired.is_empty()); + assert_eq!(wheel.now(), 9); + + // Advance 1 more tick - should expire + let expired = wheel.advance(1); + assert_eq!(expired.len(), 1); + assert!(wheel.is_empty()); + assert!(wheel.nearest_wakeup().is_none()); + } + + #[test] + fn test_timing_wheel_insert_and_expire_immediate() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Insert timer that's already expired (expiration <= current_tick) + let _handle = wheel.insert(waker, 0); + // Note: nearest_wakeup returns None for tick 0 because NonZeroU64 can't represent 0 + // But the timer is still in the wheel + assert!(!wheel.is_empty()); + + // Advance 1 tick - should expire immediately + let expired = wheel.advance(1); + assert_eq!(expired.len(), 1); + } + + #[test] + fn test_timing_wheel_cancel() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + let handle = wheel.insert(waker, 100); + assert!(!wheel.is_empty()); + + wheel.remove(handle); + assert!(wheel.is_empty()); + assert!(wheel.nearest_wakeup().is_none()); + } + + #[test] + fn test_timing_wheel_cancel_wrong_generation() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + let handle1 = wheel.insert(waker, 100); + wheel.remove(handle1); + + // Re-insert with same slab_index but different generation + let waker2 = mock_waker(counter.clone()); + let handle2 = wheel.insert(waker2, 200); + + // Old handle should not cancel the new timer + wheel.remove(handle1); + assert!(!wheel.is_empty()); + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(200).unwrap()) + ); + + // New handle should work + wheel.remove(handle2); + assert!(wheel.is_empty()); + } + + #[test] + fn test_timing_wheel_multiple_timers_same_slot() { + let mut wheel = TimingWheel::new(); + let counter1 = Arc::new(Mutex::new(0)); + let counter2 = Arc::new(Mutex::new(0)); + let counter3 = Arc::new(Mutex::new(0)); + + // Insert 3 timers for the same expiration tick + let _h1 = wheel.insert(mock_waker(counter1.clone()), 50); + let _h2 = wheel.insert(mock_waker(counter2.clone()), 50); + let _h3 = wheel.insert(mock_waker(counter3.clone()), 50); + + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(50).unwrap()) + ); + + // Advance to expiration + let expired = wheel.advance(50); + assert_eq!(expired.len(), 3); + + // Wake the expired timers + for waker in expired { + waker.wake(); + } + + // All wakers should have been called + assert_eq!(*counter1.lock().unwrap(), 1); + assert_eq!(*counter2.lock().unwrap(), 1); + assert_eq!(*counter3.lock().unwrap(), 1); + } + + #[test] + fn test_timing_wheel_multiple_timers_different_levels() { + let mut wheel = TimingWheel::new(); + let counter1 = Arc::new(Mutex::new(0)); + let counter2 = Arc::new(Mutex::new(0)); + let counter3 = Arc::new(Mutex::new(0)); + + // Level 0: 1-63ms + let _h1 = wheel.insert(mock_waker(counter1.clone()), 10); + // Level 1: 64-4095ms + let _h2 = wheel.insert(mock_waker(counter2.clone()), 100); + // Level 2: 4096-262143ms + let _h3 = wheel.insert(mock_waker(counter3.clone()), 5000); + + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(10).unwrap()) + ); + + // Advance to first expiration + let expired = wheel.advance(10); + assert_eq!(expired.len(), 1); + expired.into_iter().for_each(|w| w.wake()); + assert_eq!(*counter1.lock().unwrap(), 1); + + // Advance to second expiration (90 more ticks) + let expired = wheel.advance(90); + assert_eq!(expired.len(), 1); + expired.into_iter().for_each(|w| w.wake()); + assert_eq!(*counter2.lock().unwrap(), 1); + + // Advance to third expiration (4900 more ticks) + let expired = wheel.advance(4900); + assert_eq!(expired.len(), 1); + expired.into_iter().for_each(|w| w.wake()); + assert_eq!(*counter3.lock().unwrap(), 1); + } + + #[test] + fn test_timing_wheel_large_advance() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Insert timer for 1000 ticks + let _handle = wheel.insert(waker, 1000); + + // Advance all at once (should be efficient, not 1000 iterations) + let expired = wheel.advance(1000); + assert_eq!(expired.len(), 1); + assert_eq!(wheel.now(), 1000); + } + + #[test] + fn test_timing_wheel_wrap_around_level_0() { + let mut wheel = TimingWheel::new(); + let counter1 = Arc::new(Mutex::new(0)); + let counter2 = Arc::new(Mutex::new(0)); + + // Insert timer at slot 63 (near wrap) + let _h1 = wheel.insert(mock_waker(counter1.clone()), 63); + // Insert timer that wraps around + let _h2 = wheel.insert(mock_waker(counter2.clone()), 65); + + // Advance past wrap point + let expired = wheel.advance(65); + assert_eq!(expired.len(), 2); + } + + #[test] + fn test_timing_wheel_nearest_wakeup_updates() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + + // Insert timers in non-sorted order + let _h1 = wheel.insert(mock_waker(counter.clone()), 100); + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(100).unwrap()) + ); + + let _h2 = wheel.insert(mock_waker(counter.clone()), 50); + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(50).unwrap()) + ); + + let _h3 = wheel.insert(mock_waker(counter.clone()), 200); + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(50).unwrap()) + ); + + // Expire the 50 tick timer + wheel.advance(50); + // Now 100 should be the nearest + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(100).unwrap()) + ); + } + + #[test] + fn test_timing_wheel_boundary_values() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + + // Test boundary: exactly 63 (last slot of level 0) + let _h1 = wheel.insert(mock_waker(counter.clone()), 63); + let expired = wheel.advance(63); + assert_eq!(expired.len(), 1); + + // Test boundary: exactly 64 (first slot of level 1) + let _h2 = wheel.insert(mock_waker(counter.clone()), 64); + let expired = wheel.advance(64); + assert_eq!(expired.len(), 1); + + // Test boundary: 4095 (last slot of level 1) + let _h3 = wheel.insert(mock_waker(counter.clone()), 4095); + let expired = wheel.advance(4095); + assert_eq!(expired.len(), 1); + + // Test boundary: 4096 (first slot of level 2) + let _h4 = wheel.insert(mock_waker(counter.clone()), 4096); + let expired = wheel.advance(4096); + assert_eq!(expired.len(), 1); + } + + #[test] + fn test_timing_wheel_zero_advance() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + let _handle = wheel.insert(waker, 100); + let expired = wheel.advance(0); + assert!(expired.is_empty()); + assert_eq!(wheel.now(), 0); + } + + #[test] + fn test_timing_wheel_cascade_from_higher_levels() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + + // Insert timer at level 2 (4096+) + let _handle = wheel.insert(mock_waker(counter.clone()), 5000); + + // Advance past the expiration - cascade should bring it down through levels + let expired = wheel.advance(5000); + assert_eq!(expired.len(), 1); + } + + #[test] + fn test_timing_wheel_many_timers() { + let mut wheel = TimingWheel::new(); + let mut handles = Vec::new(); + let counter = Arc::new(Mutex::new(0)); + + // Insert 100 timers at various intervals + for i in 1..=100 { + let waker = mock_waker(counter.clone()); + let handle = wheel.insert(waker, i * 10); + handles.push(handle); + } + + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(10).unwrap()) + ); + + // Cancel every other timer + for (i, handle) in handles.iter().enumerate() { + if i % 2 == 0 { + wheel.remove(*handle); + } + } + + // Advance and check that only 50 timers expire + let expired = wheel.advance(1000); + assert_eq!(expired.len(), 50); + } + + #[test] + fn test_timing_wheel_reinsert_on_early_wake() { + let mut wheel = TimingWheel::new(); + let counter = Arc::new(Mutex::new(0)); + let waker = mock_waker(counter.clone()); + + // Insert timer for 100 ticks + let _handle = wheel.insert(waker, 100); + + // Advance only 50 ticks - timer should be re-inserted at appropriate level + let expired = wheel.advance(50); + assert!(expired.is_empty()); + assert!(!wheel.is_empty()); + + // The timer should still be tracked + assert_eq!( + wheel.nearest_wakeup(), + Some(std::num::NonZeroU64::new(100).unwrap()) + ); + + // Advance remaining 50 ticks + let expired = wheel.advance(50); + assert_eq!(expired.len(), 1); + } +} diff --git a/vendor/vibeio/src/util/async_wrap.rs b/vendor/vibeio/src/util/async_wrap.rs new file mode 100644 index 0000000..c47d42f --- /dev/null +++ b/vendor/vibeio/src/util/async_wrap.rs @@ -0,0 +1,423 @@ +//! Async I/O wrapper for interoperability with tokio traits. +//! +//! This module provides `AsyncWrap`, a type that adapts `vibeio`'s `AsyncRead` +//! and `AsyncWrite` traits to the `tokio::io` traits. This enables using +//! `vibeio` types with tokio-based libraries that expect the tokio I/O traits. +//! +//! The wrapper buffers read operations and ensures write operations complete +//! fully (like `write_all`), making it suitable for bridging async runtimes. +//! +//! # Implementation notes +//! - Read operations are buffered with a 4KB buffer size. +//! - Write operations are completed fully before returning, splitting the +//! buffer if necessary. +//! - Concurrent operations are rejected with an error. +//! - The wrapper is `Unpin` regardless of the inner type. + +use std::{ + future::Future, + mem::MaybeUninit, + pin::Pin, + task::{Context, Poll}, +}; + +use futures_util::FutureExt; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +type Buffer = Box<[u8]>; +const BUFFER_SIZE: usize = 4096; + +/// A wrapper that adapts `vibeio`'s `AsyncRead`/`AsyncWrite` to `tokio::io` traits. +/// +/// This type bridges the gap between `vibeio`'s async I/O traits and tokio's +/// `AsyncRead`/`AsyncWrite` traits, allowing `vibeio` types to be used with +/// tokio-based libraries. +/// +/// # Examples +/// ```ignore +/// use tokio::io::{AsyncReadExt, AsyncWriteExt}; +/// use vibeio::util::AsyncWrap; +/// +/// // Wrap a vibeio async reader +/// let mut reader = some_vibeio_reader(); +/// let mut wrap = AsyncWrap::new(reader); +/// +/// let mut buf = Vec::new(); +/// wrap.read_to_end(&mut buf).await?; // tokio method +/// ``` +pub struct AsyncWrap { + inner: Option, + read_buf: Option<(Buffer, usize, usize)>, + #[allow(clippy::type_complexity)] + read_fut: Option, Buffer, T)>>>>, + #[allow(clippy::type_complexity)] + write_fut: Option, T)>>>>, + #[allow(clippy::type_complexity)] + flush_fut: Option, T)>>>>, +} + +impl AsyncWrap { + /// Create a new `AsyncWrap` wrapping the given inner value. + #[inline] + pub fn new(inner: T) -> Self { + Self { + inner: Some(inner), + read_buf: None, + read_fut: None, + write_fut: None, + flush_fut: None, + } + } +} + +impl AsyncRead for AsyncWrap +where + T: crate::io::AsyncRead + 'static, +{ + #[inline] + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let this = self.get_mut(); + + if this.read_fut.is_none() { + let buf_read = this.read_buf.take(); + if let Some((buf_read, advanced, n)) = buf_read { + let unfilled = + unsafe { &mut *(buf.unfilled_mut() as *mut [MaybeUninit] as *mut [u8]) }; + let copy_len = (n - advanced).min(unfilled.len()); + unsafe { + std::ptr::copy_nonoverlapping( + buf_read.as_ptr().add(advanced), + unfilled.as_mut_ptr(), + copy_len, + ) + }; + unsafe { buf.assume_init(copy_len) }; + buf.advance(copy_len); + if advanced + copy_len < n { + this.read_buf = Some((buf_read, advanced + copy_len, n)); + } + return Poll::Ready(Ok(())); + } + let buf = { + let slice = Buffer::new_uninit_slice(BUFFER_SIZE); + // SAFETY: u8 is a primitive type and has no padding, so + // uninitialized memory is safe to read as a `Buffer`. + unsafe { slice.assume_init() } + }; + let Some(mut inner) = this.inner.take() else { + return Poll::Ready(Err(std::io::Error::other( + "another operation is already in progress", + ))); + }; + let fut = Box::pin(async move { + let (read, buf) = crate::io::AsyncRead::read(&mut inner, buf).await; + (read, buf, inner) + }); + this.read_fut = Some(fut); + } + let read_fut = this.read_fut.as_mut().expect("read_fut is None"); + let (read, buf_read, inner) = futures_util::ready!(read_fut.poll_unpin(cx)); + this.read_fut = None; + this.inner = Some(inner); + match read { + Ok(n) => { + // Put buf_read into buf + let unfilled = + unsafe { &mut *(buf.unfilled_mut() as *mut [MaybeUninit] as *mut [u8]) }; + let copy_len = n.min(unfilled.len()); + unsafe { + std::ptr::copy_nonoverlapping( + buf_read.as_ptr(), + unfilled.as_mut_ptr(), + copy_len, + ) + }; + unsafe { buf.assume_init(copy_len) }; + buf.advance(copy_len); + if copy_len < n { + this.read_buf = Some((buf_read, copy_len, n)); + } + Poll::Ready(Ok(())) + } + Err(e) => Poll::Ready(Err(e)), + } + } +} + +impl AsyncWrite for AsyncWrap +where + T: crate::io::AsyncWrite + 'static, +{ + #[inline] + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + if buf.is_empty() { + return Poll::Ready(Ok(0)); + } + + let this = self.get_mut(); + + if this.write_fut.is_none() { + let buf: Vec = buf.into(); + let Some(mut inner) = this.inner.take() else { + return Poll::Ready(Err(std::io::Error::other( + "another operation is already in progress", + ))); + }; + let fut = Box::pin(async move { + // write_all + let mut buf = buf; + let mut total_written = 0; + while !buf.is_empty() { + let (written, mut buf_written) = + crate::io::AsyncWrite::write(&mut inner, buf).await; + match written { + Ok(n) => { + total_written += n; + buf = buf_written.split_off(n); + } + Err(e) => { + return (Err(e), inner); + } + } + } + (Ok(total_written), inner) + }); + this.write_fut = Some(fut); + } + let write_fut = this.write_fut.as_mut().expect("read_fut is None"); + let (written, inner) = futures_util::ready!(write_fut.poll_unpin(cx)); + this.write_fut = None; + this.inner = Some(inner); + Poll::Ready(written) + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + if this.flush_fut.is_none() { + let Some(mut inner) = this.inner.take() else { + return Poll::Ready(Err(std::io::Error::other( + "another operation is already in progress", + ))); + }; + let fut = Box::pin(async move { + let flush = crate::io::AsyncWrite::flush(&mut inner).await; + (flush, inner) + }); + this.flush_fut = Some(fut); + } + let flush_fut = this.flush_fut.as_mut().expect("read_fut is None"); + let (flush, inner) = futures_util::ready!(flush_fut.poll_unpin(cx)); + this.flush_fut = None; + this.inner = Some(inner); + Poll::Ready(flush) + } + + #[inline] + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + // No-op + Poll::Ready(Ok(())) + } +} + +impl Unpin for AsyncWrap {} + +#[cfg(test)] +mod tests { + use std::io; + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + use std::task::{Context, Poll}; + + use futures_util::task::noop_waker; + use tokio::io::{ + AsyncRead as TokioAsyncRead, AsyncReadExt, AsyncWrite as TokioAsyncWrite, AsyncWriteExt, + ReadBuf, + }; + + use super::AsyncWrap; + use crate::io::{AsyncRead, AsyncWrite, IoBuf, IoBufMut}; + + struct CountingReader { + data: Vec, + offset: usize, + reads: Arc, + } + + impl CountingReader { + fn new(data: &[u8], reads: Arc) -> Self { + Self { + data: data.to_vec(), + offset: 0, + reads, + } + } + } + + impl AsyncRead for CountingReader { + async fn read(&mut self, mut buf: B) -> (Result, B) { + self.reads.fetch_add(1, Ordering::SeqCst); + if self.offset >= self.data.len() { + return (Ok(0), buf); + } + + let remaining = self.data.len() - self.offset; + let cap = buf.buf_capacity(); + let read_len = remaining.min(cap); + + unsafe { + let ptr = buf.as_buf_mut_ptr(); + std::ptr::copy_nonoverlapping(self.data[self.offset..].as_ptr(), ptr, read_len); + buf.set_buf_init(read_len); + } + + self.offset += read_len; + (Ok(read_len), buf) + } + } + + struct WriterState { + data: Vec, + writes: usize, + flushed: bool, + } + + struct ChunkedWriter { + state: Arc>, + chunk_size: usize, + } + + impl ChunkedWriter { + fn new(state: Arc>, chunk_size: usize) -> Self { + Self { state, chunk_size } + } + } + + impl AsyncWrite for ChunkedWriter { + async fn write(&mut self, buf: B) -> (Result, B) { + let len = buf.buf_len(); + if len == 0 { + return (Ok(0), buf); + } + + let write_len = len.min(self.chunk_size.max(1)); + let slice = unsafe { std::slice::from_raw_parts(buf.as_buf_ptr(), write_len) }; + + let mut guard = self.state.lock().expect("lock writer state"); + guard.writes += 1; + guard.data.extend_from_slice(slice); + + (Ok(write_len), buf) + } + + async fn flush(&mut self) -> Result<(), io::Error> { + let mut guard = self.state.lock().expect("lock writer state"); + guard.flushed = true; + Ok(()) + } + } + + struct PendingIo; + + impl AsyncRead for PendingIo { + async fn read(&mut self, buf: B) -> (Result, B) { + let _ = futures_util::future::pending::<()>().await; + (Ok(0), buf) + } + } + + impl AsyncWrite for PendingIo { + async fn write(&mut self, buf: B) -> (Result, B) { + (Ok(0), buf) + } + + async fn flush(&mut self) -> Result<(), io::Error> { + Ok(()) + } + } + + #[test] + fn async_wrap_read_buffers_leftover() { + let runtime = crate::executor::Runtime::new(crate::driver::AnyDriver::new_mock()); + runtime.block_on(async { + let reads = Arc::new(AtomicUsize::new(0)); + let reader = CountingReader::new(b"abcdefghij", reads.clone()); + let mut wrap = AsyncWrap::new(reader); + + let mut buf1 = [0u8; 3]; + let n1 = wrap.read(&mut buf1).await.expect("read should succeed"); + assert_eq!(n1, 3); + assert_eq!(&buf1[..n1], b"abc"); + + let mut buf2 = [0u8; 4]; + let n2 = wrap.read(&mut buf2).await.expect("read should succeed"); + assert_eq!(n2, 4); + assert_eq!(&buf2[..n2], b"defg"); + + let mut buf3 = [0u8; 4]; + let n3 = wrap.read(&mut buf3).await.expect("read should succeed"); + assert_eq!(n3, 3); + assert_eq!(&buf3[..n3], b"hij"); + + assert_eq!(reads.load(Ordering::SeqCst), 1); + }); + } + + #[test] + fn async_wrap_write_writes_all_and_flushes() { + let runtime = crate::executor::Runtime::new(crate::driver::AnyDriver::new_mock()); + runtime.block_on(async { + let state = Arc::new(Mutex::new(WriterState { + data: Vec::new(), + writes: 0, + flushed: false, + })); + let writer = ChunkedWriter::new(state.clone(), 2); + let mut wrap = AsyncWrap::new(writer); + + let payload = b"hello world"; + let n = wrap.write(payload).await.expect("write should succeed"); + assert_eq!(n, payload.len()); + wrap.flush().await.expect("flush should succeed"); + + let guard = state.lock().expect("lock writer state"); + assert_eq!(guard.data, payload); + assert!(guard.flushed); + assert!(guard.writes > 1); + }); + } + + #[test] + fn async_wrap_rejects_concurrent_operations() { + let mut wrap = AsyncWrap::new(PendingIo); + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + + let mut buf = [0u8; 8]; + let mut read_buf = ReadBuf::new(&mut buf); + let poll = Pin::new(&mut wrap).poll_read(&mut cx, &mut read_buf); + assert!(matches!(poll, Poll::Pending)); + + let poll_write: Poll> = Pin::new(&mut wrap).poll_write(&mut cx, b"hi"); + match poll_write { + Poll::Ready(Err(err)) => { + assert_eq!(err.kind(), io::ErrorKind::Other); + } + _ => panic!("expected concurrent write to return an error"), + } + } +} diff --git a/vendor/vibeio/src/util/mod.rs b/vendor/vibeio/src/util/mod.rs new file mode 100644 index 0000000..fd17d4d --- /dev/null +++ b/vendor/vibeio/src/util/mod.rs @@ -0,0 +1,110 @@ +//! Utility types and functions for `vibeio`. +//! +//! This module provides supporting infrastructure for the library: +//! - `AsyncWrap`: a wrapper that adapts `AsyncRead`/`AsyncWrite` implementations +//! to the `tokio::io` traits, enabling interoperability with tokio-based code. +//! - `supports_completion`: check if the current driver supports completion-based I/O. +//! - `supports_io_uring`: check if the system supports io_uring with required operations. + +mod async_wrap; + +pub use async_wrap::*; + +use crate::current_driver; +#[cfg(target_os = "linux")] +use std::sync::OnceLock; + +/// Check if the current driver supports completion-based I/O operations. +/// +/// This function returns `true` when the runtime is using a driver that can +/// handle completion-based operations (e.g., io_uring with completion queues). +#[inline] +pub fn supports_completion() -> bool { + current_driver().is_some_and(|driver| driver.supports_completion()) +} + +/// Check if the system supports io_uring with required operations. +/// +/// This function performs a runtime probe to verify that io_uring is available +/// and supports the operations needed by the library (accept, connect, poll, +/// timeout, read, write, etc.). On non-Linux platforms, this always returns `false`. +/// +/// The result is cached for the lifetime of the program. +#[inline] +pub fn supports_io_uring() -> bool { + #[cfg(target_os = "linux")] + { + static SUPPORTED: OnceLock = OnceLock::new(); + *SUPPORTED.get_or_init(detect_io_uring_support) + } + + #[cfg(not(target_os = "linux"))] + { + false + } +} + +#[cfg(target_os = "linux")] +fn detect_io_uring_support() -> bool { + use io_uring::opcode; + + let ring = match build_probe_ring() { + Ok(ring) => ring, + Err(_) => return false, + }; + + let mut probe = io_uring::Probe::new(); + if ring.submitter().register_probe(&mut probe).is_err() { + return false; + } + + let mut required_ops = Vec::new(); + required_ops.extend_from_slice(&[ + opcode::Accept::CODE, + opcode::Connect::CODE, + opcode::PollAdd::CODE, + opcode::Timeout::CODE, + opcode::Read::CODE, + opcode::Readv::CODE, + opcode::Recv::CODE, + opcode::RecvMsg::CODE, + opcode::Send::CODE, + opcode::SendMsg::CODE, + opcode::Write::CODE, + opcode::Writev::CODE, + ]); + + #[cfg(feature = "fs")] + required_ops.extend_from_slice(&[ + opcode::OpenAt::CODE, + opcode::Fsync::CODE, + opcode::MkDirAt::CODE, + opcode::RenameAt::CODE, + opcode::LinkAt::CODE, + opcode::SymlinkAt::CODE, + opcode::UnlinkAt::CODE, + ]); + + #[cfg(all(feature = "fs", any(target_env = "gnu", musl_v1_2_3)))] + required_ops.push(opcode::Statx::CODE); + + #[cfg(feature = "splice")] + required_ops.push(opcode::Splice::CODE); + + required_ops.iter().all(|op| probe.is_supported(*op)) +} + +#[cfg(target_os = "linux")] +fn build_probe_ring() -> std::io::Result { + let mut builder = io_uring::IoUring::builder(); + builder + .setup_single_issuer() + .setup_coop_taskrun() + .setup_taskrun_flag() + .setup_submit_all(); + + match builder.build(2) { + Ok(ring) => Ok(ring), + Err(_) => io_uring::IoUring::new(2), + } +} From 1ac8047b07dc4a95ce89cf6c2243bd54594f1c99 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Tue, 21 Jul 2026 06:09:08 -0700 Subject: [PATCH 20/21] update StopSocketReader command to include done_tx and enhance server poll handling on Windows --- Cargo.toml | 2 +- src/loop_core/commands.rs | 5 +- src/runtime.rs | 3 +- src/stream_transport.rs | 169 ++++++++++++++++++++++++++++++++++---- 4 files changed, 159 insertions(+), 20 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index fd38693..d74cfca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,7 @@ tracy-client = { version = "0.18", default-features = false, features = ["enable vibeio = { path = "vendor/vibeio", default-features = false, features = ["time"] } [target.'cfg(windows)'.dependencies] -windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Networking_WinSock", "Win32_System_Pipes", "Win32_System_Threading"] } +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Networking_WinSock", "Win32_System_IO", "Win32_System_Pipes", "Win32_System_Threading"] } [build-dependencies] pyo3-build-config = "0.29" diff --git a/src/loop_core/commands.rs b/src/loop_core/commands.rs index df2ec10..cd540ff 100644 --- a/src/loop_core/commands.rs +++ b/src/loop_core/commands.rs @@ -100,7 +100,10 @@ pub enum LoopIoCommand { core: Arc, reader: ReaderTarget, }, - StopSocketReader(RawFd), + StopSocketReader { + fd: RawFd, + done_tx: std::sync::mpsc::Sender<()>, + }, StartServerAccept { fd: RawFd, server: Arc, diff --git a/src/runtime.rs b/src/runtime.rs index 440da53..080d3bb 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -441,10 +441,11 @@ impl RuntimeDispatcher { self.reader_tasks.insert(fd, task); } - LoopCommand::Io(LoopIoCommand::StopSocketReader(fd)) => { + LoopCommand::Io(LoopIoCommand::StopSocketReader { fd, done_tx }) => { if let Some(task) = self.reader_tasks.remove(&fd) { cancel_watch_task(task); } + let _ = done_tx.send(()); } LoopCommand::Io(LoopIoCommand::StartServerAccept { fd, diff --git a/src/stream_transport.rs b/src/stream_transport.rs index 625210b..3ace490 100644 --- a/src/stream_transport.rs +++ b/src/stream_transport.rs @@ -33,6 +33,11 @@ use vibeio::net::TcpStream as VibeTcpStream; use vibeio::net::{PollTcpStream as VibePollTcpStream, TcpListener as VibeTcpListener}; #[cfg(unix)] use vibeio::net::{PollUnixStream as VibePollUnixStream, UnixListener as VibeUnixListener}; +#[cfg(windows)] +use windows_sys::Win32::{ + Foundation::{ERROR_OPERATION_ABORTED, HANDLE}, + System::IO::CancelIoEx, +}; use crate::async_event::AsyncEvent; use crate::context::{ @@ -75,6 +80,10 @@ const BLOCKING_POLL_INTERVAL_MS: i32 = 50; // several owned chunks, while the pending-event drain can still coalesce the // slow protocol path up to MAX_PENDING_READ_COALESCE_BYTES. const STREAM_READ_BUFFER_SIZE: usize = 16 * 1024; +#[cfg(windows)] +const SERVER_POLL_READER_WRITE_THRESHOLD: usize = STREAM_READ_BUFFER_SIZE; +#[cfg(windows)] +const SERVER_POLL_READER_TINY_TRIGGER_MAX_BYTES: usize = 16; const OWNED_READ_HANDOFF_MIN_BYTES: usize = STREAM_READ_BUFFER_SIZE / 2; const TLS_WORKER_STACK_SIZE: usize = 256 * 1024; @@ -541,6 +550,10 @@ pub struct StreamTransportCore { direct_writer: Option>, pending_direct_write: Mutex>, direct_write_scheduled: AtomicBool, + #[cfg(windows)] + server_poll_reader_requested: AtomicBool, + #[cfg(windows)] + server_poll_reader_ready: AtomicBool, lazy_writer: Mutex>, workers: Mutex>, // Signaled whenever `read_paused` clears or the transport starts @@ -551,6 +564,7 @@ pub struct StreamTransportCore { // the per-write hot path avoids a state lock plus hash lookup. has_text_encoding: bool, server_side: bool, + native_stream_reader: bool, } struct ServerState { @@ -963,6 +977,48 @@ impl WorkerThread { } impl StreamTransportCore { + #[cfg(windows)] + #[inline] + fn server_poll_reader_requested(&self) -> bool { + self.server_side && self.server_poll_reader_requested.load(Ordering::Acquire) + } + + #[cfg(windows)] + fn request_server_poll_reader(&self) { + if !self.server_side + || !self.native_stream_reader + || self + .server_poll_reader_requested + .swap(true, Ordering::AcqRel) + { + return; + } + + let fd = self.state.lock().expect("poisoned transport state").io_fd; + if let Some(fd) = fd { + // Wake a server reader currently blocked in WSARecv. The operation + // still completes through vibeio's IOCP driver with + // ERROR_OPERATION_ABORTED; the reader handles that result by + // rebinding the same socket to readiness mode. Cancel before the + // direct write so a pending receive cannot throttle the duplicate + // writer socket during a bulk response. + let _ = unsafe { CancelIoEx(fd as HANDLE, std::ptr::null()) }; + } + } + + #[cfg(windows)] + fn mark_server_poll_reader_ready(self: &Arc) { + if !self.server_side { + return; + } + self.server_poll_reader_ready.store(true, Ordering::Release); + if self.direct_write_scheduled.load(Ordering::Acquire) { + let _ = self.loop_core.send_command(LoopCommand::Transport( + LoopTransportCommand::StreamWrite(Arc::clone(self)), + )); + } + } + fn close_extra_socket_with_py(&self, py: Python<'_>) { let socket = self .state @@ -1848,6 +1904,12 @@ impl StreamTransportCore { pub(crate) fn flush_pending_direct_write(self: &Arc) { profiling::scope!("StreamTransportCore::flush_pending_direct_write"); + #[cfg(windows)] + if self.server_poll_reader_requested() + && !self.server_poll_reader_ready.load(Ordering::Acquire) + { + return; + } self.direct_write_scheduled.store(false, Ordering::Release); let data = std::mem::take( self.pending_direct_write @@ -1903,6 +1965,16 @@ impl StreamTransportCore { fn try_write_bytes(self: &Arc, data: &[u8]) -> io::Result<()> { profiling::scope!("StreamTransportCore::try_write_bytes"); + #[cfg(windows)] + if data.len() >= SERVER_POLL_READER_WRITE_THRESHOLD { + self.request_server_poll_reader(); + if self.server_poll_reader_requested() + && !self.server_poll_reader_ready.load(Ordering::Acquire) + { + return self.stage_direct_write(data); + } + } + if self.direct_writer.is_some() && !self.write_backpressure_active() { if self.direct_write_scheduled.load(Ordering::Acquire) || (self.server_side @@ -2900,6 +2972,10 @@ fn new_stream_transport_core( ) -> Arc { let has_text_encoding = parts.state.extra.contains_key("text_encoding"); let server_side = parts.state.server.is_some(); + let native_stream_reader = matches!( + &parts.state.callbacks.stream_reader_fast_path, + Some(StreamReaderFastPath::Native { .. }) + ); let reading = parts.state.reading; Arc::new(StreamTransportCore { loop_core: parts.loop_core, @@ -2912,11 +2988,16 @@ fn new_stream_transport_core( direct_writer: direct_writer.map(Mutex::new), pending_direct_write: Mutex::new(Vec::new()), direct_write_scheduled: AtomicBool::new(false), + #[cfg(windows)] + server_poll_reader_requested: AtomicBool::new(false), + #[cfg(windows)] + server_poll_reader_ready: AtomicBool::new(false), lazy_writer: Mutex::new(lazy_writer), workers: Mutex::new(Vec::new()), state_cv: Condvar::new(), has_text_encoding, server_side, + native_stream_reader, }) } @@ -4080,9 +4161,20 @@ fn spawn_socket_reader( /// Stops the socket reader for `fd` via the runtime-thread command path (readers /// are hosted there; see `spawn_socket_reader`). fn stop_socket_reader(core: &StreamTransportCore, fd: fd_ops::RawFd) { - let _ = core + let (done_tx, done_rx) = mpsc::channel(); + if core .loop_core - .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader(fd))); + .send_command(LoopCommand::Io(LoopIoCommand::StopSocketReader { + fd, + done_tx, + })) + .is_ok() + { + // start_tls duplicates the socket before reaching this point. Do not + // let the TLS handshake use that duplicate until the runtime-thread + // reader has been dropped, otherwise it can consume handshake bytes. + let _ = done_rx.recv(); + } } #[cfg(not(windows))] @@ -4147,11 +4239,11 @@ pub(crate) async fn run_tcp_socket_reader_task( ) { profiling::scope!("stream.run_tcp_socket_reader_task"); - // Server-side sockets retain the readiness reader. A server transport can - // duplicate its socket for direct bulk writes, and an outstanding IOCP read - // on the original socket materially slows that path. Client-side sockets - // use completion reads, avoiding one AFD poll per request/response turn. - if core.server_side { + // Native fast streams start in completion mode for low-latency + // request/response traffic. Custom protocols retain readiness mode because + // either endpoint may call start_tls(), which reclaims the socket + // synchronously. + if !core.native_stream_reader { match VibePollTcpStream::from_std(stream) { Ok(reader) => { run_windows_poll_tcp_reader( @@ -4168,6 +4260,8 @@ pub(crate) async fn run_tcp_socket_reader_task( return; } + // Native fast-stream readers use completion mode. A native server switches + // to readiness mode when a large response begins. let mut reader = match VibeTcpStream::from_std(stream) { Ok(reader) => reader, Err(err) => { @@ -4192,6 +4286,19 @@ pub(crate) async fn run_tcp_socket_reader_task( return; } + if core.server_poll_reader_requested() { + match reader.into_poll() { + Ok(poll_reader) => { + core.mark_server_poll_reader_ready(); + run_windows_poll_tcp_reader(core, poll_reader, buf).await; + } + Err(err) => core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost( + Some(err.to_string()), + )), + } + return; + } + buf.clear(); let (result, returned_buf) = VibeAsyncRead::read(&mut reader, buf).await; buf = returned_buf; @@ -4203,26 +4310,54 @@ pub(crate) async fn run_tcp_socket_reader_task( Ok(read) => { let saturated = read == buf.capacity(); let data = take_async_read_data(&mut buf); - core.enqueue_pending_read_event(PendingReadEvent::Data(data)); - - // Completion reads win on small request/response traffic. Once - // a buffer fills, switch this connection to readiness mode so - // sustained bulk transfer does not keep an overlapped receive - // outstanding alongside the direct writer socket. - if saturated { + let tiny_server_trigger = + core.server_side && read <= SERVER_POLL_READER_TINY_TRIGGER_MAX_BYTES; + + // Completion reads win on normal request/response traffic. A + // full inbound buffer indicates sustained transfer; a tiny + // server command commonly triggers a large outbound response. + // Rebind before delivering either event so the protocol cannot + // begin bulk writes while an overlapped receive is outstanding. + if saturated || tiny_server_trigger { + if core.server_side { + core.server_poll_reader_requested + .store(true, Ordering::Release); + } match reader.into_poll() { Ok(poll_reader) => { + core.mark_server_poll_reader_ready(); + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); run_windows_poll_tcp_reader(core, poll_reader, buf).await; } - Err(err) => core.enqueue_pending_read_event( - PendingReadEvent::ConnectionLost(Some(err.to_string())), - ), + Err(err) => { + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); + core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost( + Some(err.to_string()), + )); + } } return; } + + core.enqueue_pending_read_event(PendingReadEvent::Data(data)); } Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue, Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + Err(err) + if core.server_poll_reader_requested() + && err.raw_os_error() == Some(ERROR_OPERATION_ABORTED as i32) => + { + match reader.into_poll() { + Ok(poll_reader) => { + core.mark_server_poll_reader_ready(); + run_windows_poll_tcp_reader(core, poll_reader, buf).await; + } + Err(err) => core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost( + Some(err.to_string()), + )), + } + return; + } Err(err) => { core.enqueue_pending_read_event(PendingReadEvent::ConnectionLost(Some( err.to_string(), From 8acccffffb3cf3d073fa7cac6ab635158894da53 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Tue, 21 Jul 2026 18:37:10 +0300 Subject: [PATCH 21/21] fix some issues with tls generation and benchmark on macos --- benchmarks/workload_matrix.py | 35 ++++++++++++++++++++++++++++++ scripts/generate_test_tls_certs.py | 10 ++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/benchmarks/workload_matrix.py b/benchmarks/workload_matrix.py index 7e39155..b8c4e14 100644 --- a/benchmarks/workload_matrix.py +++ b/benchmarks/workload_matrix.py @@ -207,6 +207,38 @@ def application_digest(iterations: int) -> bytes: return digest +def ensure_idle_connection_capacity(connection_count: int) -> None: + if os.name == "nt": + return + + try: + import resource + except ImportError: + return + + # The client and server live in the same process. Account for both steady- + # state sockets and descriptors held transiently while the runtime registers + # each endpoint, then leave room for the listener and interpreter runtime. + required = connection_count * 4 + 64 + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft >= required: + return + if hard < required: + raise RuntimeError( + f"--idle-connections {connection_count} requires a file-descriptor " + f"limit of at least {required}, but the hard limit is {hard}; lower " + "--idle-connections or raise the hard limit" + ) + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (required, hard)) + except (OSError, ValueError) as exc: + raise RuntimeError( + f"could not raise the file-descriptor soft limit from {soft} to " + f"{required}; lower --idle-connections or run with `ulimit -n " + f"{required}`" + ) from exc + + def tls_contexts(tls_dir: Path) -> tuple[ssl.SSLContext, ssl.SSLContext]: cert = tls_dir / "cert.pem" key = tls_dir / "key.pem" @@ -674,6 +706,9 @@ def run_with_loop(loop_name: str, awaitable: Awaitable[MatrixResult]) -> MatrixR def child_main(args: argparse.Namespace) -> int: + if args.scenario == "idle_connections": + ensure_idle_connection_capacity(args.idle_connections) + if args.profile_label: if args.loop != "rsloop": raise RuntimeError("Tracy profiling is only supported for rsloop") diff --git a/scripts/generate_test_tls_certs.py b/scripts/generate_test_tls_certs.py index e8b3c43..9fe2bc9 100644 --- a/scripts/generate_test_tls_certs.py +++ b/scripts/generate_test_tls_certs.py @@ -39,6 +39,14 @@ keyUsage = critical, digitalSignature, keyEncipherment extendedKeyUsage = serverAuth +[v3_cert] +subjectAltName = @alt_names +basicConstraints = CA:false +keyUsage = critical, digitalSignature, keyEncipherment +extendedKeyUsage = serverAuth +subjectKeyIdentifier = hash +authorityKeyIdentifier = keyid:always,issuer + [alt_names] DNS.1 = localhost IP.1 = 127.0.0.1 @@ -170,7 +178,7 @@ def generate(outdir: pathlib.Path) -> None: "-extfile", str(generated["cert.cnf"]), "-extensions", - "v3_req", + "v3_cert", ) finally: for name in ("cert.csr", "cert.cnf", "ca.cnf", "ca-cert.srl"):