diff --git a/crates/adapterlib/src/format.rs b/crates/adapterlib/src/format.rs index db34b29dd72..ca3a51cabee 100644 --- a/crates/adapterlib/src/format.rs +++ b/crates/adapterlib/src/format.rs @@ -246,24 +246,26 @@ impl InputBuffer for Vec> { fn take_some(&mut self, n: usize) -> Option> { let mut result = Vec::new(); let mut remaining = n; - // Index of first buffer that should be preserved - let mut index = 0; + // Number of leading buffers drained of all their records, which are the + // only ones safe to drop. + let mut drained = 0; for v in self.iter_mut() { if remaining == 0 { break; } - let buf = v.take_some(remaining); - if let Some(ib) = buf { - let len = ib.len().records; - if remaining >= len { - // This buffer will be completely used - index += 1; - } - remaining = remaining.saturating_sub(len); - result.push(ib); + if let Some(buffer) = v.take_some(remaining) { + // A buffer with coarse granularity may hand back more than + // `remaining`, so use saturating subtraction. + remaining = remaining.saturating_sub(buffer.len().records); + result.push(buffer); + } + if v.len().records > 0 { + // `v` kept records, so we cannot consider it drained. + break; } + drained += 1; } - self.drain(0..index); + self.drain(0..drained); if result.is_empty() { None } else { @@ -1120,3 +1122,124 @@ impl ParseErrorInner { self.tag.clone() } } + +#[cfg(test)] +mod test { + use std::hash::Hasher; + use std::sync::{Arc, Mutex}; + + use super::{BufferSize, InputBuffer}; + + const BYTES_PER_RECORD: usize = 10; + + /// Minimal [InputBuffer] over a list of record ids. `flush` appends to a + /// shared sink so a test can tell dropped records from flushed ones. + struct TestBuffer { + records: Vec, + sink: Arc>>, + } + + impl TestBuffer { + fn boxed(records: &[u64], sink: &Arc>>) -> Box { + Box::new(Self { + records: records.to_vec(), + sink: sink.clone(), + }) + } + } + + impl InputBuffer for TestBuffer { + fn flush(&mut self) { + self.sink + .lock() + .unwrap() + .extend(std::mem::take(&mut self.records)); + } + + fn len(&self) -> BufferSize { + BufferSize { + records: self.records.len(), + bytes: self.records.len() * BYTES_PER_RECORD, + } + } + + fn hash(&self, hasher: &mut dyn Hasher) { + for record in &self.records { + hasher.write_u64(*record); + } + } + + fn take_some(&mut self, n: usize) -> Option> { + if self.records.is_empty() { + return None; + } + let n = n.min(self.records.len()); + Some(Box::new(Self { + records: self.records.drain(0..n).collect(), + sink: self.sink.clone(), + })) + } + } + + /// A take that ends inside an inner buffer must leave the remainder behind + /// rather than dropping it: the connector reports those records as received, + /// so losing them stalls `pipeline_complete` forever. + #[test] + fn take_some_keeps_partially_drained_buffer() { + let sink = Arc::new(Mutex::new(Vec::new())); + let mut nested = vec![ + TestBuffer::boxed(&[1, 2, 3], &sink), + TestBuffer::boxed(&[4, 5, 6, 7, 8], &sink), + ]; + + // 5 records spans all of the first buffer and part of the second. + let mut head = nested.take_some(5).expect("buffer is not empty"); + assert_eq!(head.len().records, 5); + assert_eq!( + InputBuffer::len(&nested).records, + 3, + "records 6..8 were dropped" + ); + + head.flush(); + nested.take_all().unwrap().flush(); + assert_eq!(*sink.lock().unwrap(), vec![1, 2, 3, 4, 5, 6, 7, 8]); + } + + /// Byte accounting must survive the same split; a leak here pins + /// `buffered_input_bytes` above zero. + #[test] + fn take_some_preserves_byte_accounting() { + let sink = Arc::new(Mutex::new(Vec::new())); + let mut nested = vec![ + TestBuffer::boxed(&[1, 2, 3], &sink), + TestBuffer::boxed(&[4, 5, 6, 7, 8], &sink), + ]; + let before = InputBuffer::len(&nested); + + let head = nested.take_some(5).unwrap(); + assert_eq!(head.len() + InputBuffer::len(&nested), before); + } + + /// Repeatedly draining in batches that do not divide the buffer sizes must + /// still deliver every record exactly once. + #[test] + fn take_some_loses_no_records_across_batches() { + for batch in 1..=9 { + let sink = Arc::new(Mutex::new(Vec::new())); + let mut nested = vec![ + TestBuffer::boxed(&[1, 2, 3], &sink), + TestBuffer::boxed(&[4, 5], &sink), + TestBuffer::boxed(&[6, 7, 8, 9], &sink), + ]; + while let Some(mut head) = nested.take_some(batch) { + head.flush(); + } + assert_eq!( + *sink.lock().unwrap(), + (1..=9).collect::>(), + "batch size {batch} lost records" + ); + } + } +} diff --git a/crates/adapters/src/test/mock_input_consumer.rs b/crates/adapters/src/test/mock_input_consumer.rs index 8d5633f7ad8..9da7be70d05 100644 --- a/crates/adapters/src/test/mock_input_consumer.rs +++ b/crates/adapters/src/test/mock_input_consumer.rs @@ -20,6 +20,12 @@ pub struct MockInputConsumerState { /// Number of times `extended` has been called since the last `reset`. pub n_extended: usize, + /// Value reported by `InputConsumer::max_batch_size`. + /// + /// Unlimited by default, so a test that wants an endpoint to split a parsed + /// buffer across steps has to lower it. + max_batch_size: usize, + /// The last error received from the endpoint since the last `reset`. pub endpoint_error: Option, @@ -36,6 +42,7 @@ impl MockInputConsumerState { Self { eoi: false, n_extended: 0, + max_batch_size: usize::MAX, endpoint_error: None, error_cb: None, transaction_in_progress: false, @@ -70,6 +77,11 @@ impl MockInputConsumer { self.state().reset(); } + /// Caps the records an endpoint may hand to the circuit per step. + pub fn set_max_batch_size(&self, max_batch_size: usize) { + self.state().max_batch_size = max_batch_size; + } + pub fn state(&self) -> MutexGuard<'_, MockInputConsumerState> { self.0.lock().unwrap() } @@ -104,7 +116,7 @@ impl InputConsumer for MockInputConsumer { } fn max_batch_size(&self) -> usize { - usize::MAX + self.state().max_batch_size } fn pipeline_fault_tolerance(&self) -> Option { diff --git a/crates/adapters/src/transport/s3.rs b/crates/adapters/src/transport/s3.rs index 305a851dc79..e6526fa1049 100644 --- a/crates/adapters/src/transport/s3.rs +++ b/crates/adapters/src/transport/s3.rs @@ -965,6 +965,7 @@ fn to_s3_config(config: &Arc) -> aws_sdk_s3::Config { #[cfg(test)] mod test { use crate::{ + preprocess::PassthroughPreprocessorFactory, test::{ ErrorCallback, MockDeZSet, MockInputConsumer, MockInputParser, mock_parser_pipeline, wait, @@ -979,6 +980,9 @@ mod test { primitives::{ByteStream, SdkBody}, types::error::builders::{NoSuchBucketBuilder, NoSuchKeyBuilder}, }; + use feldera_adapterlib::format::{MessageOrientedPreprocessedParser, Parser}; + use feldera_adapterlib::preprocess::PreprocessorFactory; + use feldera_types::preprocess::PreprocessorConfig; use feldera_types::{ config::{InputEndpointConfig, TransportConfig}, deserialize_without_context, @@ -1036,6 +1040,30 @@ mod test { } }"#; + /// Newline-delimited JSON: `lines: single` picks the greedy line splitter. + const NDJSON_SINGLE_KEY_CONFIG_STR: &str = r#" +{ + "stream": "test_input", + "transport": { + "name": "s3_input", + "config": { + "aws_access_key_id": "FAKE_ACCESS_KEY", + "aws_secret_access_key": "FAKE_SECRET", + "bucket_name": "test-bucket", + "region": "us-west-1", + "key": "obj1" + } + }, + "format": { + "name": "json", + "config": { + "update_format": "raw", + "array": false, + "lines": "single" + } + } +}"#; + /// Builds a reader over `mock`, with `on_error` already installed. /// /// The callback has to be passed in rather than registered afterwards: @@ -1051,6 +1079,29 @@ mod test { MockInputConsumer, MockInputParser, MockDeZSet, + ) { + test_setup_with(config_str, mock, on_error, Preprocess::No) + } + + /// Whether to wrap the endpoint's parser the way a connector configured with + /// a message-oriented preprocessor does, which makes `parse` return a nested + /// buffer instead of a flat one. + #[derive(Copy, Clone, PartialEq)] + enum Preprocess { + No, + Passthrough, + } + + fn test_setup_with( + config_str: &str, + mock: super::MockS3Client, + on_error: ErrorCallback, + preprocess: Preprocess, + ) -> ( + Box, + MockInputConsumer, + MockInputParser, + MockDeZSet, ) { let config: InputEndpointConfig = serde_json::from_str(config_str).unwrap(); let transport_config = config.connector_config.transport.clone(); @@ -1066,10 +1117,25 @@ mod test { ) .unwrap(); consumer.on_error(Some(on_error)); + let endpoint_parser: Box = if preprocess == Preprocess::Passthrough { + let preprocessor = PassthroughPreprocessorFactory + .create(&PreprocessorConfig { + name: "passthrough".to_string(), + message_oriented: true, + config: serde_json::Value::Null, + }) + .unwrap(); + Box::new(MessageOrientedPreprocessedParser::new( + preprocessor, + Box::new(parser.clone()), + )) + } else { + Box::new(parser.clone()) + }; let reader = Box::new(S3InputReader::new_inner( &transport_config, Box::new(consumer.clone()), - Box::new(parser.clone()), + endpoint_parser, Arc::new(mock), None, )) as Box; @@ -1133,6 +1199,74 @@ mod test { run_test(SINGLE_KEY_CONFIG_STR, mock, test_data); } + /// A preprocessor makes `parse` return a nested buffer, and the endpoint + /// hands the circuit at most `max_batch_size` records per step, so a single + /// parsed buffer has to survive being split across steps. Dropping the + /// remainder loses records the endpoint already reported as received, which + /// pins `total_input_records` above `total_processed_records` and stops the + /// pipeline from ever reporting completion. + /// + /// The format matters: newline-delimited JSON splits with `LineSplitter`, + /// which takes as many records as it can per chunk, so one nested element + /// holds many records and a batch boundary lands inside it. A CSV splitter + /// yields one record per element and never exercises the split. + #[test] + fn preprocessed_read_split_across_batches() { + const RECORDS: i64 = 100; + // Not a divisor of RECORDS, so the last batch is short too. + const MAX_BATCH_SIZE: usize = 7; + + let body = (1..=RECORDS) + .map(|i| format!("{{\"i\":{i}}}\n")) + .collect::>() + .concat(); + + let mut mock = super::MockS3Client::default(); + mock.expect_get_object_keys() + .with(eq("test-bucket"), eq(""), eq(&None), eq(&None)) + .return_once(|_, _, _, _| Ok((vec!["obj1".to_string()], None))); + mock.expect_get_object() + .with(eq("test-bucket"), eq("obj1"), eq(0), eq(&None)) + .return_once(move |_, _, _, _| { + Ok(GetObjectOutput::builder() + .body(ByteStream::from(SdkBody::from(body.as_str()))) + .build()) + }); + + let (reader, consumer, _parser, input_handle) = test_setup_with( + NDJSON_SINGLE_KEY_CONFIG_STR, + mock, + Box::new(|_, _| ()), + Preprocess::Passthrough, + ); + consumer.set_max_batch_size(MAX_BATCH_SIZE); + + reader.extend(); + wait( + || { + reader.queue(false); + input_handle.state().flushed.len() == RECORDS as usize + }, + 10000, + ) + .unwrap_or_else(|_| { + panic!( + "only {} of {RECORDS} records reached the circuit", + input_handle.state().flushed.len() + ) + }); + + let mut outputs = input_handle + .state() + .flushed + .iter() + .map(|upd| upd.unwrap_insert().clone()) + .collect::>(); + outputs.sort(); + let expected: Vec = (1..=RECORDS).map(|i| TestStruct { i }).collect(); + assert_eq!(outputs, expected); + } + #[test] fn single_object_with_prefix_read() { let mut mock = super::MockS3Client::default();