diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index c467ec1..712a695 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -65,7 +65,7 @@ jobs: cargo run -p lance-context-wal --features lab --bin context-wal-lab -- disk --directory "$RUNNER_TEMP/wal-disk-smoke" --payload-bytes 1024 --batch 4 --writers 2 --groups-per-writer 2 --rounds 8 - name: WAL placement, admission, record, contract and IPC decoding tests run: | - cargo test -p lance-context-wal --features placement,admission,ipc,record,arrow-contract --lib + cargo test -p lance-context-wal --features placement,admission,ipc,record,arrow-contract,vectored-io --lib cargo test -p lance-context-wal --features arrow-contract,serde_json/preserve_order --lib ingest_contract::arrow - name: WAL experiment process lifecycle tests run: cargo test -p lance-context-wal --features runner --test runner --bin context-wal-resource diff --git a/Cargo.lock b/Cargo.lock index c5219a2..18bcd82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4201,6 +4201,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "tempfile", + "tokio", "uuid", ] diff --git a/crates/lance-context-wal/Cargo.toml b/crates/lance-context-wal/Cargo.toml index 5552ed3..76d7b41 100644 --- a/crates/lance-context-wal/Cargo.toml +++ b/crates/lance-context-wal/Cargo.toml @@ -7,6 +7,7 @@ publish = false description = "Independent Rust WAL experiments owned by lance-context" [features] +vectored-io = ["dep:bytes", "dep:tokio"] runner = ["dep:anyhow", "dep:clap", "dep:libc", "dep:serde", "dep:serde_json", "dep:sha2"] contract = ["dep:serde", "dep:sha2"] arrow-contract = ["contract", "dep:arrow-schema", "arrow-schema/serde", "dep:serde_json"] @@ -17,6 +18,7 @@ ipc = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "de lab = ["ipc", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] [dependencies] +tokio = { version = "1", features = ["io-util"], optional = true } anyhow = { version = "1", optional = true } arrow-array = { version = "58", optional = true } arrow-buffer = { version = "58", optional = true } @@ -34,6 +36,7 @@ sha2 = { version = "0.10", optional = true } uuid = { version = "1", features = ["serde"], optional = true } [dev-dependencies] +tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util"] } bincode = "1.3" arrow-ipc = { version = "58", features = ["zstd"] } tempfile = "3" diff --git a/crates/lance-context-wal/src/lib.rs b/crates/lance-context-wal/src/lib.rs index 04bf18d..c76e49f 100644 --- a/crates/lance-context-wal/src/lib.rs +++ b/crates/lance-context-wal/src/lib.rs @@ -19,3 +19,6 @@ pub mod record; #[cfg(feature = "contract")] pub mod ingest_contract; + +#[cfg(feature = "vectored-io")] +pub mod transport_io; diff --git a/crates/lance-context-wal/src/transport_io.rs b/crates/lance-context-wal/src/transport_io.rs new file mode 100644 index 0000000..7cf62bb --- /dev/null +++ b/crates/lance-context-wal/src/transport_io.rs @@ -0,0 +1,133 @@ +//! Bounded scatter/gather writes over borrowed payload owners. +//! +//! This preserves the existing length/header/body wire bytes. It does not make +//! writes atomic: errors and cancellation can leave a partial exchange, so the +//! protocol must discard that connection. Counts describe socket acceptance, +//! never remote receipt or durability. Callers enforce protocol size limits and +//! must obtain a prepared-write grant before calling write_bodies. + +use bytes::Bytes; +use std::io::{self, IoSlice}; +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +const MAX_SLICES: usize = 64; + +/// Write a big-endian u32 header length, the header, then all bodies in order. +/// The optional cumulative counter includes only successfully written body bytes, +/// including partial writes before an error or cancellation. Counter overflow +/// and an unrepresentable header length are rejected before writing anything. +pub async fn write_frame( + writer: &mut W, + header: &[u8], + bodies: &[Bytes], + body_bytes: Option<&mut u64>, +) -> io::Result<()> { + let size = u32::try_from(header.len()) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "header length exceeds u32"))?; + write_parts(writer, &size.to_be_bytes(), header, bodies, body_bytes).await +} + +/// Write only bodies, for the second stage of an already granted exchange. +/// Interrupted writes are propagated, with earlier successful bytes retained +/// in the counter, just like other I/O errors. +pub async fn write_bodies( + writer: &mut W, + bodies: &[Bytes], + body_bytes: Option<&mut u64>, +) -> io::Result<()> { + write_parts(writer, &[], &[], bodies, body_bytes).await +} + +fn part<'a>(prefix: &'a [u8], header: &'a [u8], bodies: &'a [Bytes], index: usize) -> &'a [u8] { + match index { + 0 => prefix, + 1 => header, + _ => &bodies[index - 2], + } +} + +async fn write_parts( + writer: &mut W, + prefix: &[u8], + header: &[u8], + bodies: &[Bytes], + mut body_bytes: Option<&mut u64>, +) -> io::Result<()> { + let invalid = || io::Error::new(io::ErrorKind::InvalidInput, "frame byte count overflow"); + let parts = bodies.len().checked_add(2).ok_or_else(invalid)?; + let total_body = bodies.iter().try_fold(0u64, |total, body| { + total + .checked_add(u64::try_from(body.len()).map_err(|_| invalid())?) + .ok_or_else(invalid) + })?; + if let Some(count) = body_bytes.as_ref() { + count.checked_add(total_body).ok_or_else(invalid)?; + } + let mut framing_left = prefix.len().checked_add(header.len()).ok_or_else(invalid)?; + let (mut index, mut offset) = (0, 0); + while index < parts { + // Only descriptors are temporary; payloads are never concatenated or + // cloned. A fixed stack array also bounds very fragmented body lists. + let mut slices: [IoSlice<'_>; MAX_SLICES] = std::array::from_fn(|_| IoSlice::new(&[])); + let (mut count, mut available) = (0, 0usize); + for next in index..parts { + let bytes = part(prefix, header, bodies, next); + let bytes = if next == index { + &bytes[offset..] + } else { + bytes + }; + if bytes.is_empty() { + continue; + } + available = available.checked_add(bytes.len()).ok_or_else(invalid)?; + slices[count] = IoSlice::new(bytes); + count += 1; + if count == MAX_SLICES { + break; + } + } + if count == 0 { + return Ok(()); + } + let scalar = count == 1 || !writer.is_write_vectored(); + let supplied = if scalar { slices[0].len() } else { available }; + let written = if scalar { + writer.write(&slices[0]).await? + } else { + writer.write_vectored(&slices[..count]).await? + }; + if written == 0 { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "writer made no progress", + )); + } + if written > supplied { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "writer exceeded supplied slices", + )); + } + let payload = written.saturating_sub(framing_left); + framing_left = framing_left.saturating_sub(written); + if let Some(count) = body_bytes.as_mut() { + **count += payload as u64; // Full body count was checked before I/O. + } + let mut remaining = written; + while remaining > 0 { + let bytes = part(prefix, header, bodies, index).len() - offset; + if remaining < bytes { + offset += remaining; + break; + } + remaining -= bytes; + index += 1; + offset = 0; + } + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/crates/lance-context-wal/src/transport_io/tests.rs b/crates/lance-context-wal/src/transport_io/tests.rs new file mode 100644 index 0000000..f5c2150 --- /dev/null +++ b/crates/lance-context-wal/src/transport_io/tests.rs @@ -0,0 +1,310 @@ +use super::*; +use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll, Waker}; +use tokio::io::AsyncReadExt; +use tokio::net::{TcpListener, TcpStream}; + +enum Step { + Limit(usize), + Error(io::ErrorKind), + Pending, + Overreport, +} + +struct Writer { + vectored: bool, + steps: VecDeque, + bytes: Vec, + calls: Vec>, + vector_calls: usize, +} + +impl Writer { + fn new(vectored: bool, steps: impl IntoIterator) -> Self { + Self { + vectored, + steps: steps.into_iter().collect(), + bytes: Vec::new(), + calls: Vec::new(), + vector_calls: 0, + } + } + + fn accept(&mut self, slices: &[IoSlice<'_>]) -> Poll> { + self.calls.push( + slices + .iter() + .map(|s| (s.as_ptr() as usize, s.len())) + .collect(), + ); + let available = slices.iter().map(|s| s.len()).sum::(); + let limit = match self.steps.pop_front().unwrap_or(Step::Limit(usize::MAX)) { + Step::Limit(limit) => limit.min(available), + Step::Error(kind) => return Poll::Ready(Err(kind.into())), + Step::Pending => return Poll::Pending, + Step::Overreport => return Poll::Ready(Ok(available + 1)), + }; + let mut left = limit; + for slice in slices { + let size = left.min(slice.len()); + self.bytes.extend_from_slice(&slice[..size]); + left -= size; + } + Poll::Ready(Ok(limit)) + } +} + +impl AsyncWrite for Writer { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + self.accept(&[IoSlice::new(bytes)]) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + slices: &[IoSlice<'_>], + ) -> Poll> { + assert!( + self.vectored, + "scalar writer must not receive vectored calls" + ); + self.vector_calls += 1; + self.accept(slices) + } + + fn is_write_vectored(&self) -> bool { + self.vectored + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +fn wire(header: &[u8], bodies: &[Bytes]) -> Vec { + let mut bytes = (header.len() as u32).to_be_bytes().to_vec(); + bytes.extend_from_slice(header); + for body in bodies { + bytes.extend_from_slice(body); + } + bytes +} + +#[tokio::test] +async fn short_writes_preserve_every_boundary_and_body_counter() { + let header = b"header"; + let bodies = [ + Bytes::new(), + Bytes::from_static(b"abc"), + Bytes::new(), + Bytes::from_static(b"defgh"), + ]; + let expected = wire(header, &bodies); + for vectored in [false, true] { + for size in 1..=expected.len() { + let mut writer = Writer::new(vectored, (0..expected.len()).map(|_| Step::Limit(size))); + let mut count = 17; + write_frame(&mut writer, header, &bodies, Some(&mut count)) + .await + .unwrap(); + assert_eq!(writer.bytes, expected, "vectored={vectored}, size={size}"); + assert_eq!(count, 25); + assert_eq!(writer.vector_calls > 0, vectored); + } + } +} + +#[tokio::test] +async fn fragmented_payloads_borrow_original_storage_with_bounded_descriptors() { + let bodies: Vec<_> = (0..193) + .map(|i| { + if i % 7 == 0 { + Bytes::new() + } else { + Bytes::from(vec![i as u8; i % 13 + 1]) + } + }) + .collect(); + let mut writer = Writer::new(true, []); + let mut count = 0; + write_frame(&mut writer, b"h", &bodies, Some(&mut count)) + .await + .unwrap(); + assert_eq!(writer.bytes, wire(b"h", &bodies)); + assert_eq!(count, bodies.iter().map(|b| b.len() as u64).sum::()); + assert!(writer.calls.len() > 2); + assert!(writer.calls.iter().all(|call| call.len() <= MAX_SLICES)); + assert_eq!(writer.calls[0].len(), MAX_SLICES); + // The accepting writer sees the original body pointers, not a concatenated + // temporary. Bytes owners stay usable after the exchange. + for body in bodies.iter().filter(|b| !b.is_empty()) { + assert!(writer + .calls + .iter() + .flatten() + .any(|&(ptr, len)| ptr == body.as_ptr() as usize && len == body.len())); + } +} + +#[tokio::test] +async fn empty_parts_and_body_only_exchange() { + let mut writer = Writer::new(true, []); + let mut count = 11; + write_bodies(&mut writer, &[Bytes::new(), Bytes::new()], Some(&mut count)) + .await + .unwrap(); + assert!(writer.calls.is_empty()); + write_frame(&mut writer, &[], &[], Some(&mut count)) + .await + .unwrap(); + assert_eq!(writer.bytes, [0; 4]); + assert_eq!(count, 11); + let bodies = [ + Bytes::new(), + Bytes::from_static(b"abc"), + Bytes::from_static(b"def"), + Bytes::new(), + ]; + write_bodies(&mut writer, &bodies, Some(&mut count)) + .await + .unwrap(); + assert_eq!(&writer.bytes[4..], b"abcdef"); + assert_eq!(count, 17); +} + +#[tokio::test] +async fn errors_retain_only_accepted_body_bytes() { + let header = b"abc"; + let bodies = [Bytes::from_static(b"def"), Bytes::from_static(b"ghijk")]; + let expected = wire(header, &bodies); + for kind in [io::ErrorKind::BrokenPipe, io::ErrorKind::Interrupted] { + for progress in 0..expected.len() { + let steps = (0..progress) + .map(|_| Step::Limit(1)) + .chain([Step::Error(kind)]); + let mut writer = Writer::new(true, steps); + let mut count = 19; + let error = write_frame(&mut writer, header, &bodies, Some(&mut count)) + .await + .unwrap_err(); + assert_eq!(error.kind(), kind); + assert_eq!(writer.bytes, expected[..progress]); + assert_eq!(count, 19 + progress.saturating_sub(7) as u64); + } + } +} + +#[tokio::test] +async fn zero_progress_and_invalid_writer_counts_fail_without_fabricating_progress() { + for vectored in [false, true] { + for (step, kind) in [ + (Step::Limit(0), io::ErrorKind::WriteZero), + (Step::Overreport, io::ErrorKind::InvalidData), + ] { + let mut writer = Writer::new(vectored, [step]); + let mut count = 3; + let error = write_frame( + &mut writer, + b"header", + &[Bytes::from_static(b"body")], + Some(&mut count), + ) + .await + .unwrap_err(); + assert_eq!(error.kind(), kind); + assert_eq!(count, 3); + assert!(writer.bytes.is_empty()); + } + } +} + +#[tokio::test] +async fn counter_overflow_is_rejected_before_framing_or_payload_io() { + for frame in [false, true] { + let mut writer = Writer::new(true, []); + let mut count = u64::MAX - 2; + let bodies = [Bytes::from_static(b"abc")]; + let result = if frame { + write_frame(&mut writer, b"head", &bodies, Some(&mut count)).await + } else { + write_bodies(&mut writer, &bodies, Some(&mut count)).await + }; + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidInput); + assert!(writer.calls.is_empty()); + assert_eq!(count, u64::MAX - 2); + } +} + +#[test] +fn cancellation_keeps_partial_body_accounting() { + let bodies = [Bytes::from_static(b"payload")]; + for frame in [false, true] { + let mut writer = Writer::new( + true, + [Step::Limit(if frame { 9 } else { 3 }), Step::Pending], + ); + let mut count = 20; + { + let future = async { + if frame { + write_frame(&mut writer, b"hd", &bodies, Some(&mut count)).await + } else { + write_bodies(&mut writer, &bodies, Some(&mut count)).await + } + }; + let mut future = std::pin::pin!(future); + assert!(future + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending()); + } + assert_eq!(count, 23); + assert!(writer.bytes.ends_with(b"pay")); + } +} + +#[tokio::test] +async fn tcp_frames_match_legacy_reader_across_multiple_exchanges() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let bodies: Vec<_> = (0..130).map(|i| Bytes::from(vec![i as u8; 4097])).collect(); + let expected: Vec<_> = bodies.iter().flat_map(|b| b.iter().copied()).collect(); + let reader = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + for _ in 0..3 { + assert_eq!(stream.read_u32().await.unwrap(), 6); + let mut header = [0; 6]; + stream.read_exact(&mut header).await.unwrap(); + assert_eq!(&header, b"legacy"); + let mut received = vec![0; expected.len()]; + stream.read_exact(&mut received).await.unwrap(); + assert_eq!(received, expected); + // Exercise the reverse direction with the existing scalar framing. + stream.write_u32(2).await.unwrap(); + stream.write_all(b"ok").await.unwrap(); + } + }); + let mut stream = TcpStream::connect(address).await.unwrap(); + let mut count = 0; + for _ in 0..3 { + write_frame(&mut stream, b"legacy", &bodies, Some(&mut count)) + .await + .unwrap(); + assert_eq!(stream.read_u32().await.unwrap(), 2); + let mut reply = [0; 2]; + stream.read_exact(&mut reply).await.unwrap(); + assert_eq!(&reply, b"ok"); + } + assert_eq!(count, (3 * 130 * 4097) as u64); + reader.await.unwrap(); +}