From 5a1c788603e72635f7662c902608cf340f3318c2 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:30:53 -0500 Subject: [PATCH 1/2] fix(staging): retain detached rows through Parquet publication --- docs/staging-conversion-visibility.md | 9 + src/parseable/staging/writer.rs | 194 +++++++-- src/parseable/streams.rs | 539 ++++++++++++++++++++++++-- src/query/stream_schema_provider.rs | 10 +- 4 files changed, 689 insertions(+), 63 deletions(-) create mode 100644 docs/staging-conversion-visibility.md diff --git a/docs/staging-conversion-visibility.md b/docs/staging-conversion-visibility.md new file mode 100644 index 000000000..db8b3d9c3 --- /dev/null +++ b/docs/staging-conversion-visibility.md @@ -0,0 +1,9 @@ +# Memory visibility during Parquet conversion + +With memory staging enabled, detaching a disk writer no longer removes its rows from the queryable memory buffers. Those buffers follow the physical Arrow source through sealing and movement into a processing directory. Conversion releases them only when their Parquet output is published, or when the source is explicitly discarded as invalid. + +Query capture and Parquet publication share a short-lived lock. Capture clones the Arrow batches and records the local Parquet files; publication renames the completed output and releases its source buffers. Parquet encoding, remote storage requests and query execution run outside that lock. Arrow payloads are shared rather than copied. + +This extends the current-minute retention change in PR #53 to detached writers. It requires memory staging: `ENABLE_MEMORY_STAGING=false` continues to expose only published Parquet. Failed conversions or failed sealing retain their buffers until recovery, which can use more memory than the usual staging window. A poisoned publication lock makes subsequent capture fail explicitly. + +This is a local memory-to-Parquet transition guarantee. Staged-file unlinking and overlap with uploaded files are handled separately by PR #57; recovery/converter ownership is a separate concern tracked in PR #50. It does not establish an atomic snapshot across memory, local files and remote manifests. When integrating the lease change, take its file leases inside `capture_staging` before reading the refreshed remote snapshot. diff --git a/src/parseable/staging/writer.rs b/src/parseable/staging/writer.rs index 1f3f47888..416d8e4c4 100644 --- a/src/parseable/staging/writer.rs +++ b/src/parseable/staging/writer.rs @@ -21,7 +21,7 @@ use std::{ collections::{HashMap, HashSet}, fs::{self, File, OpenOptions}, io::BufWriter, - path::PathBuf, + path::{Path, PathBuf}, sync::Arc, }; @@ -132,6 +132,40 @@ impl Writer { &mut self, forced: bool, now: DateTime, + ) -> HashMap { + let flushable_disk = self.split_flushable_disk(forced, now); + self.compact_memory(); + flushable_disk + } + + /// Detaches like `take_flushable_disk_at`, but keeps each detached writer's + /// rows queryable as pending memory keyed by its part path. The caller must + /// seal the writers and follow their files with `MemWriter::rekey_pending` + /// until their rows are published in Parquet or discarded. + pub(crate) fn detach_flushable_disk_at( + &mut self, + forced: bool, + now: DateTime, + ) -> Vec { + let flushable_disk = self.split_flushable_disk(forced, now); + if let Some(mem) = self.mem.as_mut() { + for (filename, writer) in &flushable_disk { + if let Some(segment) = mem.segments.remove(filename) { + mem.pending + .entry(writer.path.clone()) + .or_default() + .push(segment); + } + } + } + self.compact_memory(); + flushable_disk.into_values().collect() + } + + fn split_flushable_disk( + &mut self, + forced: bool, + now: DateTime, ) -> HashMap { let mut flushable_disk = HashMap::new(); let old_disk = std::mem::take(&mut self.disk); @@ -142,7 +176,10 @@ impl Writer { flushable_disk.insert(filename, writer); } } + flushable_disk + } + fn compact_memory(&mut self) { if let Some(mem) = self.mem.as_mut() { // Use the selected writer keys, not a second clock/range check. // Query-only segments have no disk owner and expire at each flush. @@ -159,14 +196,21 @@ impl Writer { mem.segments.shrink_to_fit(); } } - flushable_disk } } +/// A detached part whose IPC stream is finished but not yet renamed. +pub(crate) struct SealedPart { + pub(crate) part: PathBuf, + pub(crate) arrow: PathBuf, +} + pub struct DiskWriter { inner: StreamWriter>, path: PathBuf, range: TimeRange, + // Set once `seal` has taken over finishing and renaming from Drop. + sealed: bool, } impl DiskWriter { @@ -176,20 +220,9 @@ impl DiskWriter { schema: &Schema, range: TimeRange, ) -> Result { - let mut path = path.into(); - - // a rudimentary way to ensure one parquet per arrow file - if *ONE_PARQUET_PER_ARROW { - path.set_extension(Ulid::new().to_string()); - path.add_extension(ARROW_FILE_EXTENSION); - } - path.add_extension(PART_FILE_EXTENSION); - - let file = OpenOptions::new() - .write(true) - .truncate(true) - .create(true) - .open(&path)?; + let base = path.into(); + // ONE_PARQUET_PER_ARROW is a rudimentary way to ensure one parquet per arrow file + let (file, path) = Self::create_part_file(&base, *ONE_PARQUET_PER_ARROW)?; let inner = StreamWriter::try_new_with_options( BufWriter::new(file), schema, @@ -198,7 +231,59 @@ impl DiskWriter { .unwrap(), )?; - Ok(Self { inner, path, range }) + Ok(Self { + inner, + path, + range, + sealed: false, + }) + } + + /// Exclusively creates a new `.arrows.part` file for `base`. + /// + /// Each part is owned by exactly one writer, and pending memory is keyed + /// by it. A detached writer still owns its part until sealed, a part whose + /// seal failed awaits startup recovery, and a late event for that writer key + /// may open a new writer meanwhile. Truncating or sharing such a part would + /// destroy acknowledged rows and alias their memory, so an occupied name + /// gets a per-file ULID after `.data`, which keeps the same Parquet group. + fn create_part_file(base: &Path, unique: bool) -> Result<(File, PathBuf), StagingError> { + const MAX_ATTEMPTS: usize = 8; + let mut unique = unique; + let mut attempt = 0; + loop { + let mut path = base.to_path_buf(); + if unique { + path.set_extension(Ulid::new().to_string()); + path.add_extension(ARROW_FILE_EXTENSION); + } + path.add_extension(PART_FILE_EXTENSION); + + match OpenOptions::new().write(true).create_new(true).open(&path) { + Ok(file) => return Ok((file, path)), + Err(err) + if err.kind() == std::io::ErrorKind::AlreadyExists + && attempt < MAX_ATTEMPTS => + { + unique = true; + attempt += 1; + } + Err(err) => return Err(err.into()), + } + } + } + + /// Finishes the IPC stream and chooses its `.arrows` name without renaming, + /// so callers can move memory ownership to that name before conversion can + /// see it. On error the part file is left for startup recovery. + pub(crate) fn seal(mut self) -> Result { + self.sealed = true; + let part = std::mem::take(&mut self.path); + if let Err(err) = self.inner.finish() { + return Err((part, StagingError::Arrow(err))); + } + let arrow = sealed_arrow_path(&part); + Ok(SealedPart { part, arrow }) } /// Returns whether this writer belongs to the current time range. @@ -221,24 +306,15 @@ impl DiskWriter { impl Drop for DiskWriter { /// Write the continuation bytes and mark the file as done, rename to `.data.arrows` fn drop(&mut self) { + if self.sealed { + return; + } if let Err(err) = self.inner.finish() { error!("Couldn't finish arrow file {:?}, error = {err}", self.path); return; } - // `.arrows.part` becomes `.arrows` while preserving a per-file ULID. - let mut arrow_path = self.path.with_extension(""); - // If file exists, append a random string before .date to avoid overwriting - if arrow_path.exists() { - let file_name = arrow_path.file_name().unwrap().to_string_lossy(); - let date_pos = file_name - .find(".date") - .expect("File name should contain .date"); - let random_suffix = Alphanumeric.sample_string(&mut rand::thread_rng(), 8); - let new_name = format!("{}{}", random_suffix, &file_name[date_pos..]); - arrow_path.set_file_name(new_name); - } - + let arrow_path = sealed_arrow_path(&self.path); if let Err(err) = std::fs::rename(&self.path, &arrow_path) { error!("Couldn't rename file {:?}, error = {err}", self.path); } @@ -246,6 +322,22 @@ impl Drop for DiskWriter { } } +/// `.arrows.part` becomes `.arrows` while preserving a per-file ULID. +fn sealed_arrow_path(part: &Path) -> PathBuf { + let mut arrow_path = part.with_extension(""); + // If file exists, append a random string before .date to avoid overwriting + if arrow_path.exists() { + let file_name = arrow_path.file_name().unwrap().to_string_lossy(); + let date_pos = file_name + .find(".date") + .expect("File name should contain .date"); + let random_suffix = Alphanumeric.sample_string(&mut rand::thread_rng(), 8); + let new_name = format!("{}{}", random_suffix, &file_name[date_pos..]); + arrow_path.set_file_name(new_name); + } + arrow_path +} + /// Structure to keep recordbatches in memory. /// /// Any new schema is updated in the schema map. @@ -255,6 +347,17 @@ pub struct MemWriter { // Private file-owned buffers preserve the existing public MemWriter shape. // Stream query callers use recordbatch_cloned to read every segment. segments: HashMap>, + // Segments of detached writers, keyed by the current path of the staging + // file holding their rows: the part, then its `.arrows` name, then its + // processing-directory name. Released only when that exact file is + // published in Parquet or discarded, so a reused writer key or a colliding + // physical name never releases another file's rows. These rows share the + // ingested Arrow buffers. A failed conversion leaves its sources in a + // processing directory, retried only at startup, so their rows stay in RAM + // until then and can exceed the usual one or two minutes of staging; so do + // rows of parts that could not be sealed. Rows of sources discarded as + // invalid are released, matching what is persisted. + pending: HashMap>>, schema: Schema, // for checking uniqueness of schema schema_map: HashSet, @@ -266,6 +369,7 @@ impl Default for MemWriter { fn default() -> Self { Self { segments: HashMap::new(), + pending: HashMap::new(), schema: Schema::empty(), schema_map: HashSet::default(), read_buffer: Vec::default(), @@ -303,8 +407,32 @@ impl MemWriter { Ok(()) } + /// Moves pending rows to the new path of their staging file. + pub(crate) fn rekey_pending(&mut self, from: &Path, to: &Path) { + if let Some(segments) = self.pending.remove(from) { + self.pending + .entry(to.to_path_buf()) + .or_default() + .extend(segments); + } + } + + /// Drops pending rows whose staging files were published or discarded. + pub(crate) fn release_pending(&mut self, paths: &[PathBuf]) { + if self.pending.is_empty() { + return; + } + for path in paths { + self.pending.remove(path); + } + if self.pending.is_empty() { + self.pending = HashMap::new(); + } + } + pub fn clear(&mut self) { self.segments = HashMap::new(); + self.pending = HashMap::new(); self.clear_unsegmented(); } @@ -329,7 +457,11 @@ impl MemWriter { .into_iter() .map(|rb| adapt_batch(schema.clone(), &rb)) .collect::>(); - for segment in self.segments.values() { + for segment in self + .segments + .values() + .chain(self.pending.values().flatten()) + { batches.extend(segment.recordbatch_cloned(schema)?); } Ok(batches) diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index 495946cc9..bcc990d16 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -95,7 +95,7 @@ use super::{ ForwardFilePlan, ForwardStreamReader, MergedForwardRecordReader, MergedReverseRecordReader, get_reverse_reader, }, - writer::Writer, + writer::{DiskWriter, Writer}, }, }; @@ -379,10 +379,31 @@ pub struct Stream { pub data_path: PathBuf, pub options: Arc, pub writer: Mutex, + /// Held exclusively while a Parquet output is renamed into place and the + /// pending memory of its sources is released, and shared while a query + /// captures memory and the staging Parquet set, so each staged row is seen + /// exactly once. Lock order: `publication`, then `writer`. + publication: RwLock<()>, schema_writer: Mutex<()>, pub ingestor_id: Option, + #[cfg(test)] + conversion_hook: Mutex>, } +/// Test-only pause points around sealing and Parquet publication. +#[cfg(test)] +enum ConversionPhase<'a> { + /// A detached part is finished and about to be renamed to this `.arrows` path. + BeforeSealRename(&'a Path), + /// The Parquet part is complete; publication has not started. + BeforePublish(&'a Path), + /// The Parquet file is visible; its sources' memory is not yet released. + Published, +} + +#[cfg(test)] +type ConversionHook = Arc) + Send + Sync>; + /// Startup conversion owns this immutable processing-file snapshot. Periodic /// sync may run while it is consumed because it only claims root Arrow files /// and its newly-created processing directory. @@ -685,8 +706,11 @@ impl Stream { data_path, options, writer: Mutex::new(Writer::default()), + publication: RwLock::new(()), schema_writer: Mutex::new(()), ingestor_id, + #[cfg(test)] + conversion_hook: Mutex::new(None), }) } @@ -949,6 +973,7 @@ impl Stream { /// Moves eligible arrow files to the inprocess folder and groups them by parquet path. fn move_arrow_files(&self, arrow_files: Vec, inprocess_dir: &Path) { + let mut moves = Vec::with_capacity(arrow_files.len()); for arrow_file_path in arrow_files { match arrow_file_path.metadata() { Ok(meta) if meta.len() == 0 => { @@ -956,6 +981,7 @@ impl Stream { "Invalid arrow file {:?} detected for stream {}, removing it", &arrow_file_path, self.stream_name ); + self.release_pending(std::slice::from_ref(&arrow_file_path)); remove_file(&arrow_file_path).expect("File should be removed"); } Ok(_) => { @@ -964,13 +990,7 @@ impl Stream { .file_name() .expect("Arrow file should have a name"), ); - if let Err(e) = fs::rename(&arrow_file_path, &new_path) { - error!( - "Failed to rename arrow file to inprocess directory: {} -> {}: {e}", - arrow_file_path.display(), - new_path.display() - ); - } + moves.push((arrow_file_path, new_path)); } Err(e) => { warn!( @@ -980,6 +1000,21 @@ impl Stream { } } } + + // Pending rows follow their file before conversion can observe the move. + self.rekey_pending(moves.iter().map(|(from, to)| (from, to))); + let mut failed = Vec::new(); + for (arrow_file_path, new_path) in moves { + if let Err(e) = fs::rename(&arrow_file_path, &new_path) { + error!( + "Failed to rename arrow file to inprocess directory: {} -> {}: {e}", + arrow_file_path.display(), + new_path.display() + ); + failed.push((new_path, arrow_file_path)); + } + } + self.rekey_pending(failed.iter().map(|(from, to)| (from, to))); } fn inprocess_folder(base: &Path, minute: u128) -> PathBuf { @@ -1147,6 +1182,57 @@ impl Stream { } } + /// Captures staged rows for one query consistently with Parquet + /// publication: every staged row is in exactly one of the returned memory + /// batches or the Parquet files `capture_parquet` finds in the staging + /// directory. `capture_parquet` delays publication, so it should only list + /// or lease files; it must not read or scan them. + /// + /// Without memory staging (`ENABLE_MEMORY_STAGING=false`) only published + /// Parquet is captured: unconverted rows are intentionally not queryable. + pub fn capture_staging( + &self, + schema: &Arc, + capture_parquet: impl FnOnce(&Path) -> T, + ) -> Result<(Vec, T), StagingError> { + // A publication that panicked may have renamed its Parquet without + // releasing its sources' memory; fail closed rather than double count. + let _publication = self.publication.read().map_err(|poisoned| { + StagingError::PoisonError(PoisonError::new(format!( + "Publication lock poisoned while capturing staging for stream {} - {}", + self.stream_name, poisoned + ))) + })?; + let records = self.recordbatches_cloned(schema)?; + Ok((records, capture_parquet(&self.data_path))) + } + + /// Moves pending memory to the new paths of its staging files. A poisoned + /// writer is still updated: every capture already fails closed on it. + fn rekey_pending<'a>(&self, moves: impl IntoIterator) { + let mut moves = moves.into_iter().peekable(); + if moves.peek().is_none() { + return; + } + let mut writer = self.writer.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(mem) = writer.mem.as_mut() { + for (from, to) in moves { + mem.rekey_pending(from, to); + } + } + } + + /// Releases pending memory of staging files that were published or discarded. + fn release_pending(&self, paths: &[PathBuf]) { + if paths.is_empty() { + return; + } + let mut writer = self.writer.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(mem) = writer.mem.as_mut() { + mem.release_pending(paths); + } + } + pub fn clear(&self) -> Result<(), StagingError> { if let Some(m) = self .writer @@ -1187,13 +1273,84 @@ impl Stream { ))) })?; - writer.take_flushable_disk_at(forced, clock()) + writer.detach_flushable_disk_at(forced, clock()) + }; + // Sealing I/O happens here, outside the lock + self.seal_detached(stale_writers); + Ok(()) + } + + /// Finishes detached writers and renames their parts to `.arrows`. Pending + /// memory is re-keyed to the final name before the rename makes the file + /// visible to conversion, and back if the rename fails. A part that cannot + /// be finished or renamed is left for startup recovery; its rows stay + /// queryable from memory until then. + fn seal_detached(&self, detached: Vec) { + let mut sealed = Vec::with_capacity(detached.len()); + for writer in detached { + match writer.seal() { + Ok(part) => sealed.push(part), + Err((part, err)) => error!("Couldn't finish arrow file {part:?}, error = {err}"), + } + } + self.rekey_pending(sealed.iter().map(|s| (&s.part, &s.arrow))); + let mut failed = Vec::new(); + for part in sealed { + #[cfg(test)] + self.conversion_phase(ConversionPhase::BeforeSealRename(&part.arrow)); + if let Err(err) = fs::rename(&part.part, &part.arrow) { + error!("Couldn't rename file {:?}, error = {err}", part.part); + failed.push(part); + } else { + info!("flushing {:?}", part.part); + } + } + self.rekey_pending(failed.iter().map(|s| (&s.arrow, &s.part))); + } + + /// Renames a finished Parquet output into the staging set and releases the + /// pending memory of its sources as one step with respect to query capture. + fn publish_parquet( + &self, + part_path: &Path, + parquet_path: &Path, + sources: &[PathBuf], + ) -> Result<(), StagingError> { + #[cfg(test)] + self.conversion_phase(ConversionPhase::BeforePublish(parquet_path)); + // The lock covers the rename and the memory release together. If an + // earlier publication panicked between them, captures fail closed; + // keep the sources and their memory for retry instead of publishing. + let _publication = match self.publication.write() { + Ok(guard) => guard, + Err(poisoned) => { + Self::remove_partial_parquet_file(part_path); + return Err(StagingError::PoisonError(PoisonError::new(format!( + "Publication lock poisoned while publishing {parquet_path:?} for stream {} - {}", + self.stream_name, poisoned + )))); + } }; - // DiskWriter::Drop I/O happens here, outside the lock - drop(stale_writers); + if let Err(err) = std::fs::rename(part_path, parquet_path) { + error!("Couldn't rename part file: {part_path:?} -> {parquet_path:?}, error = {err}"); + Self::remove_partial_parquet_file(part_path); + return Err(err.into()); + } + #[cfg(test)] + self.conversion_phase(ConversionPhase::Published); + self.release_pending(sources); Ok(()) } + #[cfg(test)] + fn conversion_phase(&self, phase: ConversionPhase<'_>) { + // Clone out so a hook blocking one group never blocks another's phase. + let hook = self.conversion_hook.lock().unwrap().clone(); + if let Some(hook) = hook { + hook(phase); + } + } + fn parquet_writer_props( &self, merged_schema: &Schema, @@ -1925,12 +2082,7 @@ impl Stream { } } - if let Err(err) = std::fs::rename(&part_path, &parquet_path) { - error!("Couldn't rename part file: {part_path:?} -> {parquet_path:?}, error = {err}"); - Self::remove_partial_parquet_file(&part_path); - return Err(err.into()); - } - + self.publish_parquet(&part_path, &parquet_path, &readable_arrow_files)?; self.cleanup_arrow_files_and_dir(&readable_arrow_files, tenant_id); Ok(Some(merged_schema)) } @@ -2290,6 +2442,8 @@ impl Stream { ), } } + // Discarded rows are not persisted; stop serving them from memory. + self.release_pending(arrow_files); self.cleanup_arrow_files_and_dir(arrow_files, tenant_id); } @@ -2968,6 +3122,7 @@ mod tests { use arrow_ipc::writer::StreamWriter as ArrowStreamWriter; use arrow_schema::{DataType, Field, TimeUnit}; use chrono::{NaiveDate, TimeDelta, Utc}; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use temp_dir::TempDir; use tokio::time::sleep; @@ -3090,12 +3245,11 @@ mod tests { assert_eq!(writer.disk.len(), 1); assert!(writer.disk.contains_key(&expected)); } - let retained = staging.recordbatches_cloned(&schema).unwrap(); - assert_eq!(retained.len(), 1); - assert_eq!(retained[0], next); + // The detached minute stays queryable until its Parquet is published. + assert_eq!(captured_ids(&staging, &schema), (vec![7, 7, 8], vec![])); staging.flush_with_clock(true, || end).unwrap(); assert!(staging.writer.lock().unwrap().disk.is_empty()); - assert!(staging.recordbatches_cloned(&schema).unwrap().is_empty()); + assert_eq!(captured_ids(&staging, &schema), (vec![7, 7, 8], vec![])); } #[test] @@ -3160,13 +3314,13 @@ mod tests { assert_eq!(staging.writer.lock().unwrap().disk.len(), 2); assert_eq!( ids(), - vec![2, 3], - "attached current writers must remain queryable" + vec![1, 2, 3, 4], + "attached and detached-but-unpublished writers must remain queryable" ); staging.flush_with_clock(false, || now).unwrap(); assert_eq!( ids(), - vec![2, 3], + vec![1, 2, 3, 4], "a second flush must neither lose nor duplicate rows" ); staging.clear().unwrap(); @@ -3185,7 +3339,7 @@ mod tests { .unwrap(); assert_eq!(ids(), vec![5]); staging.flush_with_clock(true, || now).unwrap(); - assert!(ids().is_empty()); + assert_eq!(ids(), vec![5]); assert!(staging.writer.lock().unwrap().disk.is_empty()); // A late append for a previously detached file starts a fresh segment. staging @@ -3197,12 +3351,343 @@ mod tests { StreamType::UserDefined, ) .unwrap(); - assert_eq!(ids(), vec![5]); + assert_eq!(ids(), vec![5, 5]); staging.flush_with_clock(false, || now).unwrap(); - assert!(ids().is_empty()); + assert_eq!(ids(), vec![5, 5]); assert!(staging.writer.lock().unwrap().disk.is_empty()); } + /// Sorted ids from one coherent capture: (memory, staging Parquet). + fn captured_ids(staging: &Stream, schema: &Arc) -> (Vec, Vec) { + fn ids(batches: impl IntoIterator) -> Vec { + let mut ids = batches + .into_iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + ids.sort_unstable(); + ids + } + // Like a query, only list under the capture; read after releasing it. + let (memory, parquet_paths) = staging + .capture_staging(schema, |dir| { + fs::read_dir(dir) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "parquet")) + .collect::>() + }) + .unwrap(); + let parquet = parquet_paths.into_iter().flat_map(|path| { + ParquetRecordBatchReaderBuilder::try_new(File::open(path).unwrap()) + .unwrap() + .build() + .unwrap() + .map(Result::unwrap) + }); + (ids(memory), ids(parquet)) + } + + /// Bounds every cross-thread wait so a failing converter cannot hang a test. + const CONVERSION_WAIT: Duration = Duration::from_secs(30); + + fn conversion_now() -> DateTime { + DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc) + } + + fn conversion_stream(dir: &TempDir, name: &str) -> (StreamRef, Arc) { + let staging = Stream::new( + Arc::new(Options { + local_staging_path: dir.path().to_path_buf(), + row_group_size: 1048576, + ..Default::default() + }), + name, + LogStreamMetadata::default(), + None, + &None, + ); + staging.writer.lock().unwrap().mem = Some(Default::default()); + // Merging several Arrow sources orders them by the leading timestamp. + let schema = Arc::new(Schema::new(vec![ + Field::new( + DEFAULT_TIMESTAMP_KEY, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("id", DataType::Int32, false), + ])); + (staging, schema) + } + + fn push_ids( + staging: &Stream, + schema: &Arc, + schema_key: &str, + ids: Vec, + timestamp: DateTime, + ) { + let times = vec![timestamp.timestamp_millis(); ids.len()]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(times)), + Arc::new(Int32Array::from(ids)), + ], + ) + .unwrap(); + staging + .push( + schema_key, + &batch, + timestamp.naive_utc(), + &HashMap::new(), + StreamType::UserDefined, + ) + .unwrap(); + } + + fn part_files(staging: &Stream) -> Vec { + let mut parts = fs::read_dir(&staging.data_path) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| { + path.extension() + .is_some_and(|ext| ext == PART_FILE_EXTENSION) + }) + .collect::>(); + parts.sort(); + parts + } + + fn conversion_test_stream(dir: &TempDir, name: &str) -> (StreamRef, Arc) { + let (staging, schema) = conversion_stream(dir, name); + let now = conversion_now(); + // Two writer keys (schema keys) for one minute share an output group; + // the third is a separate output, so publishing one group must not + // release the other. + for (ids, minute, schema_key) in [ + (vec![1, 2], 0, "schema-a"), + (vec![3], 0, "schema-b"), + (vec![4, 5, 6], 1, "schema-a"), + ] { + push_ids( + &staging, + &schema, + schema_key, + ids, + now + TimeDelta::minutes(minute), + ); + } + assert_eq!( + captured_ids(&staging, &schema), + (vec![1, 2, 3, 4, 5, 6], vec![]) + ); + // Detach every writer: rows are sealed in root Arrow files only. + staging.flush_with_clock(true, || now).unwrap(); + assert!(staging.writer.lock().unwrap().disk.is_empty()); + assert_eq!(staging.arrow_files().len(), 3); + (staging, schema) + } + + fn set_conversion_hook( + staging: &Stream, + hook: impl Fn(ConversionPhase<'_>) + Send + Sync + 'static, + ) { + *staging.conversion_hook.lock().unwrap() = Some(Arc::new(hook)); + } + + #[test] + fn conversion_visibility_rows_counted_once_across_publication() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_test_stream(&dir, "publication"); + let all = vec![1, 2, 3, 4, 5, 6]; + assert_eq!(captured_ids(&staging, &schema), (all.clone(), vec![])); + + let (paused_tx, paused_rx) = std::sync::mpsc::channel::<()>(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel::<()>(); + let (published_tx, published_rx) = std::sync::mpsc::channel::(); + let resume_rx = Mutex::new(resume_rx); + let weak = Arc::downgrade(&staging); + set_conversion_hook(&staging, move |phase| match phase { + ConversionPhase::BeforeSealRename(_) => {} + ConversionPhase::BeforePublish(_) => { + paused_tx.send(()).unwrap(); + resume_rx + .lock() + .unwrap() + .recv_timeout(CONVERSION_WAIT) + .expect("test did not resume publication"); + } + ConversionPhase::Published => { + // The Parquet is visible and its sources' memory is not yet + // released: query capture must be excluded right now. + let stream = weak.upgrade().unwrap(); + published_tx + .send(stream.publication.try_read().is_err()) + .unwrap(); + } + }); + + let converter = { + let staging = Arc::clone(&staging); + spawn(move || staging.prepare_parquet(false, true, &None)) + }; + let paused = || { + paused_rx + .recv_timeout(CONVERSION_WAIT) + .expect("conversion did not reach publication") + }; + let published = || { + published_rx + .recv_timeout(CONVERSION_WAIT) + .expect("conversion did not publish") + }; + // Each output group pauses before publication with its sources moved + // to the processing directory and its Parquet part fully written. + paused(); + assert!(staging.arrow_files().is_empty()); + assert_eq!(staging.inprocess_arrow_files().len(), 3); + assert_eq!(captured_ids(&staging, &schema), (all.clone(), vec![])); + resume_tx.send(()).unwrap(); + assert!(published(), "publication must exclude capture"); + // This capture waits for the publication above to release its sources. + // One group is published, the other is not: every row exactly once. + let (memory, parquet) = captured_ids(&staging, &schema); + assert!( + (memory == [1, 2, 3] && parquet == [4, 5, 6]) + || (memory == [4, 5, 6] && parquet == [1, 2, 3]), + "memory={memory:?} parquet={parquet:?}" + ); + paused(); + resume_tx.send(()).unwrap(); + assert!(published(), "publication must exclude capture"); + converter.join().unwrap().unwrap(); + + assert_eq!(captured_ids(&staging, &schema), (vec![], all)); + assert!(staging.inprocess_arrow_files().is_empty()); + } + + #[test] + fn conversion_visibility_failed_publication_retains_rows_for_retry() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_test_stream(&dir, "failed-publication"); + let all = vec![1, 2, 3, 4, 5, 6]; + let blockers = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&blockers); + set_conversion_hook(&staging, move |phase| { + if let ConversionPhase::BeforePublish(parquet) = phase { + // A non-empty directory at the target makes the real rename fail. + fs::create_dir(parquet).unwrap(); + fs::write(parquet.join("blocker"), b"x").unwrap(); + recorded.lock().unwrap().push(parquet.to_path_buf()); + } + }); + assert!(staging.prepare_parquet(false, true, &None).is_err()); + let blockers = std::mem::take(&mut *blockers.lock().unwrap()); + assert_eq!(blockers.len(), 2); + for blocker in blockers { + fs::remove_dir_all(blocker).unwrap(); + } + // Sources stay in processing directories and their rows stay in memory. + assert_eq!(staging.inprocess_arrow_files().len(), 3); + assert_eq!(captured_ids(&staging, &schema), (all.clone(), vec![])); + + // Startup retry converts the retained processing directories. + *staging.conversion_hook.lock().unwrap() = None; + staging.prepare_parquet(true, false, &None).unwrap(); + assert_eq!(captured_ids(&staging, &schema), (vec![], all)); + assert!(staging.inprocess_arrow_files().is_empty()); + } + + #[test] + fn conversion_visibility_late_append_after_failed_seal_rename() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_stream(&dir, "failed-seal"); + let now = conversion_now(); + push_ids(&staging, &schema, "schema-a", vec![1, 2], now); + let blockers = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&blockers); + set_conversion_hook(&staging, move |phase| { + if let ConversionPhase::BeforeSealRename(arrow) = phase { + // A non-empty directory at the target makes the real rename fail. + fs::create_dir(arrow).unwrap(); + fs::write(arrow.join("blocker"), b"x").unwrap(); + recorded.lock().unwrap().push(arrow.to_path_buf()); + } + }); + staging.flush_with_clock(true, || now).unwrap(); + *staging.conversion_hook.lock().unwrap() = None; + let blockers = std::mem::take(&mut *blockers.lock().unwrap()); + assert_eq!(blockers.len(), 1); + fs::remove_dir_all(&blockers[0]).unwrap(); + // The failed part is left for recovery; its rows stay queryable. + let failed = part_files(&staging); + assert_eq!(failed.len(), 1); + let failed_len = fs::metadata(&failed[0]).unwrap().len(); + assert!(staging.arrow_files().is_empty()); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![])); + + // A late event for the same writer key and minute must not reuse, + // truncate or share the failed part. + push_ids(&staging, &schema, "schema-a", vec![3], now); + let parts = part_files(&staging); + assert_eq!(parts.len(), 2); + assert!(parts.contains(&failed[0])); + assert_eq!(fs::metadata(&failed[0]).unwrap().len(), failed_len); + staging.flush_with_clock(true, || now).unwrap(); + assert_eq!(part_files(&staging), failed); + assert_eq!(staging.arrow_files().len(), 1); + + // Publishing the late file releases only its own rows. + staging.prepare_parquet(false, true, &None).unwrap(); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![3])); + assert_eq!(part_files(&staging), failed); + } + + #[test] + fn conversion_visibility_overlapping_detach_keeps_part_ownership() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_stream(&dir, "overlapping-detach"); + let now = conversion_now(); + push_ids(&staging, &schema, "schema-a", vec![1, 2], now); + // A flush detached this writer but has not sealed it yet. + let mut detached = staging + .writer + .lock() + .unwrap() + .detach_flushable_disk_at(true, now); + assert_eq!(detached.len(), 1); + let first = part_files(&staging); + assert_eq!(first.len(), 1); + + // A late event opens a second writer for the same key and minute, and + // an overlapping flush seals it first. + push_ids(&staging, &schema, "schema-a", vec![3], now); + assert_eq!(part_files(&staging).len(), 2); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2, 3], vec![])); + staging.flush_with_clock(true, || now).unwrap(); + assert_eq!(part_files(&staging), first); + staging.prepare_parquet(false, true, &None).unwrap(); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![3])); + + // The first flush finishes late and publishes only its own rows. + staging.seal_detached(std::mem::take(&mut detached)); + assert!(part_files(&staging).is_empty()); + staging.prepare_parquet(false, true, &None).unwrap(); + assert_eq!(captured_ids(&staging, &schema), (vec![], vec![1, 2, 3])); + } + #[test] fn parallel_metric_writer_creates_valid_ordered_row_groups() { let temp_dir = TempDir::new().unwrap(); diff --git a/src/query/stream_schema_provider.rs b/src/query/stream_schema_provider.rs index c6d673d68..c820e8b8e 100644 --- a/src/query/stream_schema_provider.rs +++ b/src/query/stream_schema_provider.rs @@ -445,17 +445,17 @@ impl StandardTableProvider { return Ok(()); }; - // Staging arrow exection plan - let records = staging - .recordbatches_cloned(&self.schema) + // Capture memory and the staging Parquet set together, so a conversion + // publishing meanwhile is seen either before or after, never both. + let (records, mut parquet_files) = staging + .capture_staging(&self.schema, |_| staging.parquet_files()) .map_err(|e| DataFusionError::Internal(format!("Error during staging exec- {e}")))?; let arrow_exec = reversed_mem_table(records, self.schema.clone())? .scan(state, projection, filters, limit) .await?; execution_plans.push(arrow_exec); - // Get a list of parquet files still in staging, order by filename - let mut parquet_files = staging.parquet_files(); + // Order the staging parquet files by filename parquet_files.sort_by(|a, b| a.cmp(b).reverse()); // NOTE: We don't partition among CPUs to ensure consistent results. From b11f99ab9558eb3d63675821978f6140438817c6 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:46:29 -0500 Subject: [PATCH 2/2] fix(staging): preserve pending ownership during part recovery --- src/parseable/staging/writer.rs | 8 ++ src/parseable/streams.rs | 136 +++++++++++++++++++++++++++++--- 2 files changed, 135 insertions(+), 9 deletions(-) diff --git a/src/parseable/staging/writer.rs b/src/parseable/staging/writer.rs index 416d8e4c4..a9190b2ef 100644 --- a/src/parseable/staging/writer.rs +++ b/src/parseable/staging/writer.rs @@ -162,6 +162,14 @@ impl Writer { flushable_disk.into_values().collect() } + /// Part files still owned by attached disk writers. + pub(crate) fn active_part_paths(&self) -> HashSet { + self.disk + .values() + .map(|writer| writer.path.clone()) + .collect() + } + fn split_flushable_disk( &mut self, forced: bool, diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index bcc990d16..25b0a067e 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -384,6 +384,12 @@ pub struct Stream { /// captures memory and the staging Parquet set, so each staged row is seen /// exactly once. Lock order: `publication`, then `writer`. publication: RwLock<()>, + /// Held shared by `flush` from detaching writers until they are sealed, + /// and exclusively by part recovery. Recovery therefore never sees a part + /// that is detached but still owned by a flush, so a part with pending + /// memory and no attached writer is one whose seal failed. Lock order: + /// `part_lifecycle`, then `writer`. + part_lifecycle: RwLock<()>, schema_writer: Mutex<()>, pub ingestor_id: Option, #[cfg(test)] @@ -707,6 +713,7 @@ impl Stream { options, writer: Mutex::new(Writer::default()), publication: RwLock::new(()), + part_lifecycle: RwLock::new(()), schema_writer: Mutex::new(()), ingestor_id, #[cfg(test)] @@ -1262,9 +1269,15 @@ impl Stream { clock: impl FnOnce() -> DateTime, ) -> Result<(), StagingError> { let _span = info_span!("flush", stream_name = %self.stream_name, forced).entered(); - // Swap out stale writers under the lock, drop them after releasing it. - // DiskWriter::Drop does I/O (IPC finish + file rename) so dropping - // outside the lock avoids blocking concurrent push() calls. + let _lifecycle = self.part_lifecycle.read().map_err(|poisoned| { + StagingError::PoisonError(PoisonError::new(format!( + "Part lifecycle lock poisoned while flushing data for stream {} - {}", + self.stream_name, poisoned + ))) + })?; + // Swap out stale writers under the lock, seal them after releasing it. + // Sealing does I/O (IPC finish + file rename), so doing it outside the + // writer lock avoids blocking concurrent push() calls. let stale_writers = { let mut writer = self.writer.lock().map_err(|poisoned| { StagingError::PoisonError(PoisonError::new(format!( @@ -2714,17 +2727,51 @@ impl Stream { } /// Recovers orphaned Arrow parts and removes rebuildable Parquet parts. + /// + /// Parts of attached writers are live and skipped. Pending memory follows + /// a recovered part to its `.arrows` name before the rename exposes it to + /// conversion, and is released when a part is discarded. fn recover_orphan_part_files(&self) { + // Fail closed: without the lock a part may still belong to a flush. + let _lifecycle = match self.part_lifecycle.write() { + Ok(guard) => guard, + Err(poisoned) => { + error!( + "Part lifecycle lock poisoned, skipping part recovery for stream {} - {poisoned}", + self.stream_name + ); + return; + } + }; let Ok(dir) = self.data_path.read_dir() else { return; }; + let candidates = dir + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.extension() + .is_some_and(|ext| ext == PART_FILE_EXTENSION) + }) + .collect::>(); + if candidates.is_empty() { + return; + } + // Listed before this snapshot: writers cannot be detached meanwhile, + // and later writers create new parts, which were not listed. + let attached = match self.writer.lock() { + Ok(writer) => writer.active_part_paths(), + Err(poisoned) => { + error!( + "Writer lock poisoned, skipping part recovery for stream {} - {poisoned}", + self.stream_name + ); + return; + } + }; - for entry in dir.flatten() { - let path = entry.path(); - if path - .extension() - .is_some_and(|ext| ext == PART_FILE_EXTENSION) - { + for path in candidates { + if !attached.contains(&path) { info!( "Found orphaned .part file: {:?} for stream {}", path, self.stream_name @@ -2737,6 +2784,7 @@ impl Stream { "Removing empty orphaned .part file: {:?} for stream {}", path, self.stream_name ); + self.release_pending(std::slice::from_ref(&path)); if let Err(e) = remove_file(&path) { error!("Failed to remove empty .part file {:?}: {e}", path); } @@ -2775,6 +2823,8 @@ impl Stream { self.stream_name, meta.len() ); + // Discarded rows are not persisted. + self.release_pending(std::slice::from_ref(&path)); if let Err(delete_err) = remove_file(&path) { error!( "Failed to remove invalid/corrupted .part file {:?}: {delete_err}", @@ -2813,11 +2863,14 @@ impl Stream { "Recovering orphaned .part file: {:?} -> {:?} for stream {}", path, arrow_path, self.stream_name ); + // A failed seal's rows follow its part. + self.rekey_pending([(&path, &arrow_path)]); if let Err(e) = std::fs::rename(&path, &arrow_path) { error!( "Failed to rename .part file {:?} to {:?}: {e}", path, arrow_path ); + self.rekey_pending([(&arrow_path, &path)]); } } } @@ -3653,6 +3706,71 @@ mod tests { staging.prepare_parquet(false, true, &None).unwrap(); assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![3])); assert_eq!(part_files(&staging), failed); + + // Startup recovery on the same Stream exposes the failed part as Arrow; + // its rows must follow it and be published exactly once. + startup_then_shutdown_conversion(&staging); + assert!(part_files(&staging).is_empty()); + assert_eq!(captured_ids(&staging, &schema), (vec![], vec![1, 2, 3])); + } + + /// Runs startup preparation and its plan, then a shutdown conversion that + /// converts every root Arrow file, as a restart followed by sync would. + fn startup_then_shutdown_conversion(staging: &StreamRef) { + staging + .prepare_startup_sync(None) + .unwrap() + .execute() + .unwrap(); + staging.prepare_parquet(false, true, &None).unwrap(); + } + + #[test] + fn conversion_visibility_startup_recovery_skips_attached_writer() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_stream(&dir, "startup-attached"); + // Accepted before startup recovery reaches this stream: the writer is + // still attached and its part is live, not an orphan. + push_ids(&staging, &schema, "schema-a", vec![1, 2], conversion_now()); + assert_eq!(part_files(&staging).len(), 1); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![])); + + startup_then_shutdown_conversion(&staging); + assert!(part_files(&staging).is_empty()); + assert!(staging.arrow_files().is_empty()); + assert_eq!(captured_ids(&staging, &schema), (vec![], vec![1, 2])); + } + + #[test] + fn conversion_visibility_recovery_discard_releases_pending_rows() { + let dir = TempDir::new().unwrap(); + let (staging, schema) = conversion_stream(&dir, "recovery-discard"); + let now = conversion_now(); + push_ids(&staging, &schema, "schema-a", vec![1, 2], now); + let blockers = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&blockers); + set_conversion_hook(&staging, move |phase| { + if let ConversionPhase::BeforeSealRename(arrow) = phase { + fs::create_dir(arrow).unwrap(); + fs::write(arrow.join("blocker"), b"x").unwrap(); + recorded.lock().unwrap().push(arrow.to_path_buf()); + } + }); + staging.flush_with_clock(true, || now).unwrap(); + *staging.conversion_hook.lock().unwrap() = None; + for blocker in std::mem::take(&mut *blockers.lock().unwrap()) { + fs::remove_dir_all(blocker).unwrap(); + } + let failed = part_files(&staging); + assert_eq!(failed.len(), 1); + assert_eq!(captured_ids(&staging, &schema), (vec![1, 2], vec![])); + + // Recovery discards a part with no complete batch; its rows are not + // persisted, so memory must stop serving them too. + fs::write(&failed[0], b"bad").unwrap(); + startup_then_shutdown_conversion(&staging); + assert!(part_files(&staging).is_empty()); + assert_eq!(captured_ids(&staging, &schema), (vec![], vec![])); } #[test]