diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index c9e10233..88ae8770 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -63,8 +63,8 @@ jobs: run: | cargo run -p lance-context-wal --features lab --bin context-wal-lab -- codec --payload-bytes 1024 --rows 8 --iterations 2 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 and shared admission tests - run: cargo test -p lance-context-wal --features placement,admission --lib + - name: WAL placement, admission and IPC decoding tests + run: cargo test -p lance-context-wal --features placement,admission,ipc --lib - name: Start etcd for HA scheduler test run: | ETCD_VERSION=3.7.0 diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index 43c97b18..70b94a0f 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -56,7 +56,7 @@ jobs: - name: Clippy (lance-context, default features) run: cargo clippy --manifest-path crates/lance-context/Cargo.toml --all-targets -- -D warnings - name: Clippy (WAL experiments) - run: cargo clippy -p lance-context-wal --features lab,placement,admission --all-targets -- -D warnings + run: cargo clippy -p lance-context-wal --features lab,placement,admission,ipc --all-targets -- -D warnings typos: name: Spell Check diff --git a/crates/lance-context-wal/Cargo.toml b/crates/lance-context-wal/Cargo.toml index 58b5958f..340fa6c7 100644 --- a/crates/lance-context-wal/Cargo.toml +++ b/crates/lance-context-wal/Cargo.toml @@ -9,7 +9,8 @@ description = "Independent Rust WAL experiments owned by lance-context" [features] placement = ["dep:anyhow", "dep:serde", "dep:serde_json", "dep:sha2", "dep:uuid"] admission = ["dep:anyhow", "dep:serde"] -lab = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] +ipc = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes"] +lab = ["ipc", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] [dependencies] anyhow = { version = "1", optional = true } @@ -28,6 +29,7 @@ sha2 = { version = "0.10", optional = true } uuid = { version = "1", features = ["serde"], optional = true } [dev-dependencies] +arrow-ipc = { version = "58", features = ["zstd"] } tempfile = "3" [[bin]] diff --git a/crates/lance-context-wal/README.md b/crates/lance-context-wal/README.md index f99e6c1f..c6621cd4 100644 --- a/crates/lance-context-wal/README.md +++ b/crates/lance-context-wal/README.md @@ -3,6 +3,23 @@ Feature-gated Rust measurements for choosing a WAL payload path and durable local log engine. This crate does not change production ingestion or Lance OSS. +The optional `ipc` feature exposes `ipc::decode(Bytes, expected_schema)` for +materializers. It decodes multiple uncompressed record batches and dictionaries +while sharing binary payload ownership with the input. Numeric alignment repair, +schema parsing and Arrow metadata still allocate. The codec validates footer and +block/array-buffer ranges, native endianness, schema equality, row/count limits +and unsupported expansion before constructing arrays. Compressed batches and +dictionary deltas are rejected. +Callers still reserve input and decoding workspace within their own memory budget. +Replication does not need to enable this module or decode Arrow. The `lab` codec +measurement now uses the same validated path, including its validation overhead. +Arrow dependency panics from malformed metadata become errors in panic-unwind +builds; panic-abort builds cannot recover from dependency panics. + +```bash +cargo test -p lance-context-wal --features ipc --lib +``` + The independent `placement` feature provides `placement::Layout`, session bucket hashing and scoped `Receipt` types. A validated layout assigns each stable virtual bucket to one RF3 group, with explicit voters and StatefulSet DNS addresses. Use a diff --git a/crates/lance-context-wal/src/bin/lab.rs b/crates/lance-context-wal/src/bin/lab.rs index 733282e7..33a1820d 100644 --- a/crates/lance-context-wal/src/bin/lab.rs +++ b/crates/lance-context-wal/src/bin/lab.rs @@ -9,8 +9,7 @@ use std::time::Instant; use anyhow::{ensure, Context, Result}; use arrow_array::{Array, LargeBinaryArray, RecordBatch, StringArray, UInt64Array}; -use arrow_buffer::Buffer; -use arrow_ipc::reader::{read_footer_length, FileDecoder, FileReader}; +use arrow_ipc::reader::FileReader; use arrow_ipc::writer::FileWriter; use bytes::Bytes; use clap::{Parser, Subcommand}; @@ -131,21 +130,10 @@ fn encode(batch: &RecordBatch) -> Result> { Ok(data) } -fn decode(buffer: &Buffer, schema: arrow_schema::SchemaRef) -> Result { - let end = buffer.len().checked_sub(10).context("short IPC file")?; - let len = read_footer_length(buffer[end..].try_into()?)?; - let start = end.checked_sub(len).context("bad footer")?; - let footer = arrow_ipc::root_as_footer(&buffer[start..end]) - .map_err(|e| anyhow::anyhow!("invalid IPC footer: {e}"))?; - let decoder = FileDecoder::new(schema, footer.version()); - let blocks = footer.recordBatches().context("no batches")?; - ensure!(blocks.len() == 1, "one batch required"); - let block = blocks.get(0); - let length = block.bodyLength() as usize + block.metaDataLength() as usize; - let bytes = buffer.slice_with_length(block.offset() as usize, length); - decoder - .read_record_batch(block, &bytes)? - .context("missing batch") +fn decode(buffer: &Bytes, schema: arrow_schema::SchemaRef) -> Result { + let mut batches = lance_context_wal::ipc::decode(buffer.clone(), schema)?; + ensure!(batches.len() == 1, "one benchmark batch required"); + Ok(batches.remove(0)) } fn measured( @@ -178,15 +166,15 @@ fn measured( fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { ensure!(size > 0 && rows > 0 && iterations > 0); let batch = make_batch(rows, size)?; - let buffer = Buffer::from_vec(encode(&batch)?); + let buffer = Bytes::from(encode(&batch)?); verify_shared_decode(&buffer, &batch)?; - let reader_batch = FileReader::try_new(Cursor::new(buffer.as_slice()), None)? + let reader_batch = FileReader::try_new(Cursor::new(buffer.as_ref()), None)? .next() .context("missing reader batch")??; ensure!(reader_batch == batch, "reader round trip mismatch"); measured("encode_ipc", iterations, buffer.len(), || encode(&batch))?; measured("file_reader", iterations, buffer.len(), || { - Ok(FileReader::try_new(Cursor::new(buffer.as_slice()), None)? + Ok(FileReader::try_new(Cursor::new(buffer.as_ref()), None)? .next() .context("missing batch")??) })?; @@ -197,7 +185,7 @@ fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { let mut unaligned = Vec::with_capacity(buffer.len() + 1); unaligned.push(0); unaligned.extend_from_slice(&buffer); - let unaligned = Buffer::from_vec(unaligned).slice(1); + let unaligned = Bytes::from(unaligned).slice(1..); verify_shared_decode(&unaligned, &batch)?; measured( "unaligned_shared_buffer_decoder", @@ -226,7 +214,7 @@ fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { /// Check equality and actual payload ownership outside the measured interval. /// Numeric offset buffers can still allocate when fixing alignment. -fn verify_shared_decode(buffer: &Buffer, expected: &RecordBatch) -> Result<()> { +fn verify_shared_decode(buffer: &Bytes, expected: &RecordBatch) -> Result<()> { let decoded = decode(buffer, expected.schema())?; ensure!(&decoded == expected, "full Arrow round trip mismatch"); let bodies = decoded diff --git a/crates/lance-context-wal/src/ipc.rs b/crates/lance-context-wal/src/ipc.rs new file mode 100644 index 00000000..6e89c1c2 --- /dev/null +++ b/crates/lance-context-wal/src/ipc.rs @@ -0,0 +1,183 @@ +//! IPC file decoding for materializers over reference-counted input buffers. +//! Replication does not need this optional module or decode Arrow rows. +//! Compressed IPC is rejected until expansion is separately budgeted. +use anyhow::{ensure, Context, Result}; +use arrow_array::RecordBatch; +use arrow_buffer::Buffer; +use arrow_ipc::{ + reader::{read_footer_length, FileDecoder}, + Block, +}; +use arrow_schema::SchemaRef; +use bytes::Bytes; + +const MAX_METADATA: usize = 1024 * 1024; +const MAX_ROWS: i64 = 1_000_000; + +#[cfg(test)] +mod tests; + +/// Decode an uncompressed IPC file against the expected table schema. +/// +/// Binary values share the input's ownership; alignment repair and Arrow array +/// metadata can still allocate. Callers must bound/reserve input and decoding +/// workspace before calling. The format checks and row/count limits here do not +/// replace process-wide memory admission or a transport checksum. +/// Dictionary deltas and compressed batches are intentionally unsupported. +/// Arrow dependency panics are converted to errors in panic-unwind builds; +/// panic-abort builds cannot recover from dependency panics. +pub fn decode(body: Bytes, schema: SchemaRef) -> Result> { + // Arrow's schema conversion is infallible in its API, but malformed type + // metadata can panic even after FlatBuffers verification. All decoder state + // is local to this attempt and is discarded on unwind, never reused. + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| decode_inner(body, schema))) + .map_err(|_| anyhow::anyhow!("Arrow IPC decoder panicked on invalid metadata"))? +} + +fn decode_inner(body: Bytes, schema: SchemaRef) -> Result> { + let buffer = Buffer::from(body); + ensure!( + buffer.len() >= 16 && &buffer[..6] == b"ARROW1", + "not an IPC file" + ); + let end = buffer.len() - 10; + let footer_len = read_footer_length(buffer[end..].try_into()?)?; + ensure!(footer_len <= MAX_METADATA, "IPC footer too large"); + let start = end + .checked_sub(footer_len) + .context("invalid IPC footer range")?; + let footer = arrow_ipc::root_as_footer(&buffer[start..end]) + .map_err(|e| anyhow::anyhow!("invalid IPC footer: {e}"))?; + let ipc_schema = footer.schema().context("missing IPC schema")?; + ensure!( + ipc_schema.endianness().equals_to_target_endianness(), + "IPC endianness differs from this system" + ); + let source_schema = arrow_ipc::convert::fb_to_schema(ipc_schema); + ensure!( + source_schema == *schema, + "IPC schema differs from expected schema" + ); + let records = footer.recordBatches().context("missing record blocks")?; + ensure!( + !records.is_empty() && records.len() <= 128, + "invalid batch count" + ); + let dictionaries = footer.dictionaries(); + ensure!( + dictionaries.is_none_or(|d| d.len() <= 128), + "too many dictionary blocks" + ); + let mut blocks = Vec::new(); + if let Some(dictionaries) = dictionaries { + blocks.extend(dictionaries.iter().map(|b| (*b, true))); + } + blocks.extend(records.iter().map(|b| (*b, false))); + blocks.sort_by_key(|(b, _)| b.offset()); + let mut previous_end = 8; + let mut rows = 0; + // Validate the complete block layout and compression before decoding any + // arrays. Disallow overlapping blocks and bound logical rows, including + // null/dictionary values that could otherwise expand from tiny messages. + for (block, dictionary) in &blocks { + let offset = usize::try_from(block.offset())?; + let length = usize::try_from(block.metaDataLength())? + .checked_add(usize::try_from(block.bodyLength())?) + .context("IPC block overflow")?; + let block_end = offset.checked_add(length).context("IPC offset overflow")?; + ensure!( + offset >= previous_end && block_end <= start, + "overlapping/out-of-range IPC blocks" + ); + previous_end = block_end; + rows += validate_message(block, &buffer[offset..block_end], *dictionary)?; + ensure!(rows <= MAX_ROWS, "IPC row limit exceeded"); + } + let mut decoder = FileDecoder::new(schema, footer.version()); + // Match FileReader: load dictionary blocks before records, preserving + // dictionary replacement/delta validation in Arrow's decoder. + for (block, _) in blocks.iter().filter(|(_, dictionary)| *dictionary) { + decoder.read_dictionary(block, &slice(&buffer, block))?; + } + let mut output = Vec::new(); + for (block, _) in blocks.iter().filter(|(_, dictionary)| !dictionary) { + output.push( + decoder + .read_record_batch(block, &slice(&buffer, block))? + .context("missing IPC batch")?, + ); + } + Ok(output) +} + +fn slice(buffer: &Buffer, block: &Block) -> Buffer { + buffer.slice_with_length( + block.offset() as usize, + block.metaDataLength() as usize + block.bodyLength() as usize, + ) +} + +fn validate_message(block: &Block, bytes: &[u8], dictionary: bool) -> Result { + let metadata_len = usize::try_from(block.metaDataLength())?; + ensure!( + (8..=MAX_METADATA).contains(&metadata_len), + "invalid IPC metadata size" + ); + let prefix = if bytes.get(..4) == Some(&[255; 4]) { + 8 + } else { + 4 + }; + let metadata = bytes + .get(prefix..metadata_len) + .context("truncated IPC metadata")?; + let message = arrow_ipc::root_as_message(metadata) + .map_err(|e| anyhow::anyhow!("invalid IPC message: {e}"))?; + ensure!( + message.bodyLength() == block.bodyLength(), + "IPC body length mismatch" + ); + let batch = if dictionary { + let dict = message + .header_as_dictionary_batch() + .context("invalid dictionary block")?; + ensure!( + !dict.isDelta(), + "dictionary deltas require separate expansion accounting" + ); + dict.data().context("missing dictionary data")? + } else { + message + .header_as_record_batch() + .context("invalid record block")? + }; + ensure!( + batch.compression().is_none(), + "compressed IPC requires expansion accounting" + ); + ensure!( + (0..=MAX_ROWS).contains(&batch.length()), + "invalid IPC row count" + ); + if let Some(nodes) = batch.nodes() { + ensure!( + nodes + .iter() + .all(|node| (0..=MAX_ROWS).contains(&node.length())), + "nested IPC row limit exceeded" + ); + } + if let Some(buffers) = batch.buffers() { + for buffer in buffers { + let end = buffer + .offset() + .checked_add(buffer.length()) + .context("IPC array buffer overflow")?; + ensure!( + buffer.offset() >= 0 && buffer.length() >= 0 && end <= block.bodyLength(), + "IPC array buffer outside block body" + ); + } + } + Ok(batch.length()) +} diff --git a/crates/lance-context-wal/src/ipc/tests.rs b/crates/lance-context-wal/src/ipc/tests.rs new file mode 100644 index 00000000..085d0c3e --- /dev/null +++ b/crates/lance-context-wal/src/ipc/tests.rs @@ -0,0 +1,205 @@ +use super::*; +use arrow_array::{ + types::Int8Type, ArrayRef, DictionaryArray, Int8Array, LargeBinaryArray, NullArray, StringArray, +}; +use arrow_ipc::writer::{FileWriter, IpcWriteOptions}; +use arrow_schema::Schema; +use std::sync::Arc; + +fn fixture() -> RecordBatch { + let payload = vec![73u8; 256 * 1024]; + RecordBatch::try_from_iter([ + ( + "payload", + Arc::new(LargeBinaryArray::from(vec![ + Some(payload.as_slice()), + None, + Some(&[]), + ])) as ArrayRef, + ), + ( + "role", + Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), None, Some(1)]), + Arc::new(StringArray::from(vec!["request", "response"])), + ) + .unwrap(), + ) as ArrayRef, + ), + ]) + .unwrap() +} + +fn encode(batches: &[RecordBatch], compressed: bool) -> Result { + let mut bytes = Vec::new(); + let options = IpcWriteOptions::default() + .try_with_compression(compressed.then_some(arrow_ipc::CompressionType::ZSTD))?; + let mut writer = FileWriter::try_new_with_options(&mut bytes, &batches[0].schema(), options)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + drop(writer); + Ok(bytes.into()) +} + +#[test] +fn dictionaries_nulls_and_unaligned_binary_keep_input_ownership() -> Result<()> { + let batch = fixture(); + let encoded = encode(&[batch.clone(), batch.clone()], false)?; + for prefix in [0, 1] { + let mut bytes = vec![0; prefix]; + bytes.extend_from_slice(&encoded); + let bytes = Bytes::from(bytes).slice(prefix..); + let start = bytes.as_ptr() as usize; + let end = start + bytes.len(); + let decoded = decode(bytes, batch.schema())?; + // The caller's Bytes has been consumed. Arrow owns the backing storage. + assert_eq!(decoded, vec![batch.clone(), batch.clone()]); + for batch in &decoded { + let binary = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let pointer = binary.value_data().as_ptr() as usize; + assert!(pointer >= start && pointer + binary.value_data().len() <= end); + } + } + Ok(()) +} + +#[test] +fn rejects_truncation_wrong_schema_and_unsupported_expansion() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + for len in [0, 5, 15, bytes.len() / 2, bytes.len() - 8, bytes.len() - 1] { + assert!(decode(bytes.slice(..len), batch.schema()).is_err()); + } + assert!(decode(bytes, Arc::new(Schema::empty())).is_err()); + let compressed = encode(std::slice::from_ref(&batch), true)?; + assert!(decode(compressed, batch.schema()) + .unwrap_err() + .to_string() + .contains("compressed IPC")); + // Null rows encode almost no payload; row limits must apply before decoding. + let huge = RecordBatch::try_from_iter([( + "nulls", + Arc::new(NullArray::new(MAX_ROWS as usize + 1)) as ArrayRef, + )])?; + assert!(decode(encode(std::slice::from_ref(&huge), false)?, huge.schema()).is_err()); + assert!(decode(encode(&vec![batch.clone(); 129], false)?, batch.schema()).is_err()); + Ok(()) +} + +#[test] +fn rejects_invalid_footer_and_block_ranges_before_slicing() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let block_offset = { + let footer = arrow_ipc::root_as_footer(&bytes[end - footer_len..end]).unwrap(); + let block = footer.recordBatches().unwrap().get(0); + block as *const Block as usize - bytes.as_ptr() as usize + }; + for offset in [1i64, -1, i64::MAX] { + let mut corrupt = bytes.to_vec(); + corrupt[block_offset..block_offset + 8].copy_from_slice(&offset.to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + } + let mut corrupt = bytes.to_vec(); + corrupt[end..end + 4].copy_from_slice(&((MAX_METADATA + 1) as i32).to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + Ok(()) +} + +#[test] +fn malformed_dictionary_schema_returns_an_error() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let footer_start = end - footer_len; + let bit_width = { + let footer = arrow_ipc::root_as_footer(&bytes[footer_start..end]).unwrap(); + let field = footer.schema().unwrap().fields().unwrap().get(1); + let index = field.dictionary().unwrap().indexType().unwrap(); + footer_start + + index._tab.loc() + + usize::from(index._tab.vtable().get(arrow_ipc::Int::VT_BITWIDTH)) + }; + let mut corrupt = bytes.to_vec(); + // Structurally valid FlatBuffers can still contain invalid Arrow types. + corrupt[bit_width..bit_width + 4].copy_from_slice(&7i32.to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + Ok(()) +} + +#[test] +fn rejects_array_buffers_outside_the_record_body() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let position = { + let footer = arrow_ipc::root_as_footer(&bytes[end - footer_len..end]).unwrap(); + let block = footer.recordBatches().unwrap().get(0); + let start = block.offset() as usize; + let message = + arrow_ipc::root_as_message(&bytes[start + 8..start + block.metaDataLength() as usize]) + .unwrap(); + let buffer = message + .header_as_record_batch() + .unwrap() + .buffers() + .unwrap() + .get(0); + buffer as *const arrow_ipc::Buffer as usize - bytes.as_ptr() as usize + }; + for (offset, length) in [(-1i64, 1i64), (0, -1), (i64::MAX, 1), (0, i64::MAX)] { + let mut corrupt = bytes.to_vec(); + corrupt[position..position + 8].copy_from_slice(&offset.to_le_bytes()); + corrupt[position + 8..position + 16].copy_from_slice(&length.to_le_bytes()); + let error = decode(corrupt.into(), batch.schema()).unwrap_err(); + assert!(error.to_string().starts_with("IPC array buffer")); + } + Ok(()) +} + +#[test] +fn rejects_foreign_endianness_before_schema_conversion() { + let schema = Arc::new(Schema::empty()); + let mut builder = arrow_ipc::convert::IpcSchemaEncoder::new().schema_to_fb(&schema); + builder.reset(); + let endianness = if cfg!(target_endian = "little") { + arrow_ipc::Endianness::Big + } else { + arrow_ipc::Endianness::Little + }; + let ipc_schema = arrow_ipc::Schema::create( + &mut builder, + &arrow_ipc::SchemaArgs { + endianness, + ..Default::default() + }, + ); + let footer = arrow_ipc::Footer::create( + &mut builder, + &arrow_ipc::FooterArgs { + schema: Some(ipc_schema), + ..Default::default() + }, + ); + builder.finish(footer, None); + let footer = builder.finished_data(); + let mut bytes = b"ARROW1\0\0".to_vec(); + bytes.extend_from_slice(footer); + bytes.extend_from_slice(&(footer.len() as i32).to_le_bytes()); + bytes.extend_from_slice(b"ARROW1"); + assert!(decode(bytes.into(), schema) + .unwrap_err() + .to_string() + .contains("endianness")); +} diff --git a/crates/lance-context-wal/src/lib.rs b/crates/lance-context-wal/src/lib.rs index 5ac9820f..7ea0eee9 100644 --- a/crates/lance-context-wal/src/lib.rs +++ b/crates/lance-context-wal/src/lib.rs @@ -10,3 +10,6 @@ pub mod placement; #[cfg(feature = "admission")] pub mod admission; + +#[cfg(feature = "ipc")] +pub mod ipc;