From 326bd2cc4484702260f2003697fbd1d4dcf9d20c Mon Sep 17 00:00:00 2001 From: RyanStewart <47729789+RyanJamesStewart@users.noreply.github.com> Date: Tue, 2 Jun 2026 08:23:18 -0700 Subject: [PATCH] codec: use libc::memchr for LinesCodec delimiter scan (#8141) --- tokio-util/Cargo.toml | 5 +- tokio-util/src/codec/lines_codec.rs | 4 +- tokio-util/src/util/memchr.rs | 102 ++++++++++++++++++++++++++++ tokio-util/src/util/mod.rs | 2 + 4 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 tokio-util/src/util/memchr.rs diff --git a/tokio-util/Cargo.toml b/tokio-util/Cargo.toml index 84cb3cd22..24392c7cb 100644 --- a/tokio-util/Cargo.toml +++ b/tokio-util/Cargo.toml @@ -25,7 +25,7 @@ full = ["codec", "compat", "io-util", "time", "net", "rt", "join-map"] net = ["tokio/net"] compat = ["futures-io"] -codec = [] +codec = ["libc"] time = ["tokio/time", "slab"] io = [] io-util = ["io", "tokio/rt", "tokio/io-util"] @@ -46,6 +46,9 @@ slab = { version = "0.4.4", optional = true } # Backs `DelayQueue` tracing = { version = "0.1.29", default-features = false, features = ["std"], optional = true } hashbrown = { version = "0.15.0", default-features = false, optional = true } +[target.'cfg(unix)'.dependencies] +libc = { version = "0.2.168", optional = true } # Backs the LinesCodec delimiter scan via libc::memchr + [dev-dependencies] tokio = { version = "1.0.0", features = ["full"] } tokio-test = "0.4.0" diff --git a/tokio-util/src/codec/lines_codec.rs b/tokio-util/src/codec/lines_codec.rs index cc1ac4ebd..294086f31 100644 --- a/tokio-util/src/codec/lines_codec.rs +++ b/tokio-util/src/codec/lines_codec.rs @@ -115,9 +115,7 @@ impl Decoder for LinesCodec { // there's no max_length set, we'll read to the end of the buffer. let read_to = cmp::min(self.max_length.saturating_add(1), buf.len()); - let newline_offset = buf[self.next_index..read_to] - .iter() - .position(|b| *b == b'\n'); + let newline_offset = crate::util::memchr::memchr(b'\n', &buf[self.next_index..read_to]); match (self.is_discarding, newline_offset) { (true, Some(offset)) => { diff --git a/tokio-util/src/util/memchr.rs b/tokio-util/src/util/memchr.rs new file mode 100644 index 000000000..92296b1fb --- /dev/null +++ b/tokio-util/src/util/memchr.rs @@ -0,0 +1,102 @@ +//! Search for a byte in a byte array using libc. +//! +//! When nothing pulls in libc, then just use a trivial implementation. Note +//! that we only depend on libc on unix. + +#[cfg(not(all(unix, feature = "libc")))] +fn memchr_inner(needle: u8, haystack: &[u8]) -> Option { + haystack.iter().position(|val| needle == *val) +} + +#[cfg(all(unix, feature = "libc"))] +fn memchr_inner(needle: u8, haystack: &[u8]) -> Option { + let start = haystack.as_ptr(); + + // SAFETY: `start` is valid for `haystack.len()` bytes. + let ptr = (unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) }) + .cast::() + .cast_const(); + + if ptr.is_null() { + None + } else { + // SAFETY: `ptr` will always be in bounds, since libc guarantees that the ptr will either + // be to an element inside the array or the ptr will be null + // since the ptr is in bounds the offset must also always be non null + // and there can't be more than isize::MAX elements inside an array + // as rust guarantees that the maximum number of bytes a allocation + // may occupy is isize::MAX + unsafe { + // TODO(MSRV 1.87): When bumping MSRV, switch to `ptr.byte_offset_from_unsigned(start)`. + Some(usize::try_from(ptr.offset_from(start)).unwrap_unchecked()) + } + } +} + +pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option { + let index = memchr_inner(needle, haystack)?; + + // SAFETY: `memchr_inner` returns Some(index) and in that case index must point to an element in haystack + // or `memchr_inner` None which is guarded by the `?` operator above + // therefore the index must **always** point to an element in the array + // and so this indexing operation is safe + // TODO(MSRV 1.81): When bumping MSRV, switch to `std::hint::assert_unchecked(haystack.get(..=index).is_some());` + unsafe { + if haystack.get(..=index).is_none() { + std::hint::unreachable_unchecked() + } + } + + Some(index) +} + +#[cfg(test)] +mod tests { + use super::memchr; + + #[test] + fn memchr_test() { + let haystack = b"123abc456\0\xffabc\n"; + + assert_eq!(memchr(b'1', haystack), Some(0)); + assert_eq!(memchr(b'2', haystack), Some(1)); + assert_eq!(memchr(b'3', haystack), Some(2)); + assert_eq!(memchr(b'4', haystack), Some(6)); + assert_eq!(memchr(b'5', haystack), Some(7)); + assert_eq!(memchr(b'6', haystack), Some(8)); + assert_eq!(memchr(b'7', haystack), None); + assert_eq!(memchr(b'a', haystack), Some(3)); + assert_eq!(memchr(b'b', haystack), Some(4)); + assert_eq!(memchr(b'c', haystack), Some(5)); + assert_eq!(memchr(b'd', haystack), None); + assert_eq!(memchr(b'A', haystack), None); + assert_eq!(memchr(0, haystack), Some(9)); + assert_eq!(memchr(0xff, haystack), Some(10)); + assert_eq!(memchr(0xfe, haystack), None); + assert_eq!(memchr(1, haystack), None); + assert_eq!(memchr(b'\n', haystack), Some(14)); + assert_eq!(memchr(b'\r', haystack), None); + } + + #[test] + fn memchr_all() { + let mut arr = Vec::new(); + for b in 0..=255 { + arr.push(b); + } + for b in 0..=255 { + assert_eq!(memchr(b, &arr), Some(b as usize)); + } + arr.reverse(); + for b in 0..=255 { + assert_eq!(memchr(b, &arr), Some(255 - b as usize)); + } + } + + #[test] + fn memchr_empty() { + for b in 0..=255 { + assert_eq!(memchr(b, b""), None); + } + } +} diff --git a/tokio-util/src/util/mod.rs b/tokio-util/src/util/mod.rs index aaba542c2..837a1ced4 100644 --- a/tokio-util/src/util/mod.rs +++ b/tokio-util/src/util/mod.rs @@ -1,4 +1,6 @@ mod maybe_dangling; +#[cfg(feature = "codec")] +pub(crate) mod memchr; #[cfg(any(feature = "io", feature = "codec"))] mod poll_buf;