diff --git a/docs/staging-query-leases.md b/docs/staging-query-leases.md new file mode 100644 index 000000000..71743c22a --- /dev/null +++ b/docs/staging-query-leases.md @@ -0,0 +1,13 @@ +# Staging Parquet query leases + +Queries capture published staging Parquet files as hard links before reading the remote snapshot. The native Parquet scan owns those links through planning, repartitioning and execution. Upload cleanup can remove the original pathname without invalidating an already captured scan. Files also selected from the remote manifest are excluded from the staging scan using the uploader's canonical path mapping. + +The staging filesystem must support hard links. No Parquet payload is copied into memory or onto disk. Links retain the underlying disk space until their query releases them. Capture I/O errors, including an original file disappearing before it can be linked, fail the query rather than silently omit rows. Corrupt and empty Parquet files retain the uploader's existing skip behavior. + +Lease directories use `.query-parquet-*` and contain non-Parquet lease filenames. An OS file lock identifies a live owner. Query captures reclaim unlocked orphan directories, including after a process crash; active owners and incomplete setup directories are preserved. Stream construction does no lease cleanup because it can run under the shared stream-registry write lock. Orphan data remains until that stream is next captured by a query. Normal cleanup errors are logged. A crash before `.owner.lock` publication can leave an empty setup directory, but no Parquet hard links or retained data. + +Tradeoffs: unique lease paths reduce staging footer-cache reuse, and standalone queries within the staging window perform an additional stream-metadata read. Query/Prism snapshot read or parse errors now fail the query instead of silently omitting a node's snapshot. No CPU, latency or cache-hit improvement is claimed. + +This change covers staged-file lifetime and local/remote overlap. Memory-to-Parquet publication is a separate ownership transition; it requires the conversion-visibility fix. Stream deletion and external mutation of staging directories are not coordinated by these leases. + +Canonical-path deduplication assumes a stable partition layout between upload and query planning. If the uploader captures one layout and the configuration changes before a query maps the same file, the overlap may not match. This existing configuration/identity race is outside this fix. diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index 495946cc9..74480aff2 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -2166,43 +2166,46 @@ impl Stream { /// function to validate parquet files #[cfg_attr(feature = "hotpath", hotpath::measure)] fn is_valid_parquet_file(path: &Path, stream_name: &str) -> bool { - // First check file size as a quick validation - match path.metadata() { - Ok(meta) if meta.len() < FOOTER_SIZE as u64 => { - error!( - "Invalid parquet file {path:?} detected for stream {stream_name}, size: {} bytes", - meta.len() - ); - return false; - } - Err(e) => { - error!( - "Cannot read metadata for parquet file {path:?} for stream {stream_name}: {e}" - ); - return false; - } - _ => {} // File size is adequate, continue validation - } + Self::validate_parquet_file(path, stream_name).unwrap_or_else(|error| { + error!("Cannot validate parquet file {path:?} for stream {stream_name}: {error}"); + false + }) + } - // Try to open and read the parquet file metadata to verify it's valid - match std::fs::File::open(path) { - Ok(file) => match SerializedFileReader::new(file) { - Ok(reader) => { - if reader.metadata().file_metadata().num_rows() == 0 { - error!("Invalid parquet file {path:?} for stream {stream_name}"); - false - } else { - true - } + /// The query calls this on its immutable lease, propagating access/metadata + /// failures while retaining the uploader's corrupt/empty-file skip policy. + pub(crate) fn validate_parquet_file(path: &Path, stream_name: &str) -> std::io::Result { + let file = std::fs::File::open(path)?; + let size = file.metadata()?.len(); + if size < FOOTER_SIZE as u64 { + error!( + "Invalid parquet file {path:?} detected for stream {stream_name}, size: {size} bytes" + ); + return Ok(false); + } + match SerializedFileReader::new(file) { + Ok(reader) => { + let valid = reader.metadata().file_metadata().num_rows() != 0; + if !valid { + error!("Invalid parquet file {path:?} for stream {stream_name}: no rows"); } - Err(e) => { - error!("Failed to read parquet file {path:?} for stream {stream_name}: {e}"); - false + Ok(valid) + } + Err(parquet::errors::ParquetError::External(error)) => { + // Actual filesystem read errors are capture failures too. + match error.downcast::() { + Ok(error) => Err(*error), + Err(error) => { + error!( + "Failed to read parquet file {path:?} for stream {stream_name}: {error}" + ); + Ok(false) + } } - }, - Err(e) => { - error!("Failed to open parquet file {path:?} for stream {stream_name}: {e}"); - false + } + Err(error) => { + error!("Failed to read parquet file {path:?} for stream {stream_name}: {error}"); + Ok(false) } } } diff --git a/src/query/stream_schema_provider.rs b/src/query/stream_schema_provider.rs index 4ee5ab3f5..8d7b9d7b4 100644 --- a/src/query/stream_schema_provider.rs +++ b/src/query/stream_schema_provider.rs @@ -18,6 +18,9 @@ use std::{any::Any, cmp::Reverse, collections::HashMap, ops::Bound, sync::Arc}; +pub(crate) mod staging_parquet; +use staging_parquet::capture_staging_parquet; + use arrow_array::RecordBatch; use arrow_schema::{Schema, SchemaRef}; use chrono::{DateTime, NaiveDateTime, TimeDelta, Timelike, Utc}; @@ -397,43 +400,17 @@ impl StandardTableProvider { filters: &[Expr], limit: Option, state: &dyn Session, + captured: (Vec, Vec), ) -> Result<(), DataFusionError> { - let Ok(staging) = PARSEABLE.get_stream(&self.stream, &self.tenant_id) else { - return Ok(()); - }; - - // Staging arrow exection plan - let records = staging - .recordbatches_cloned(&self.schema) - .map_err(|e| DataFusionError::Internal(format!("Error during staging exec- {e}")))?; + // Both sources were captured before the remote snapshot was read. + let (records, partitioned_files) = captured; 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(); - parquet_files.sort_by(|a, b| a.cmp(b).reverse()); - - // NOTE: We don't partition among CPUs to ensure consistent results. - // i.e. We were seeing in-consistent ordering when querying over parquets in staging. - let mut partitioned_files = Vec::with_capacity(parquet_files.len()); - for file_path in parquet_files { - let Ok(file_meta) = file_path.metadata() else { - continue; - }; - let file = PartitionedFile::new(file_path.display().to_string(), file_meta.len()); - partitioned_files.push(file) - } - - // // NOTE: There is the possibility of a parquet file being pushed to object store - // // and deleted from staging in the time it takes for datafusion to get to it. - // // Staging parquet execution plan - // let object_store_url = if let Some(tenant_id) = self.tenant_id.as_ref() { - // &format!("file://{tenant_id}/") - // } else { - // "file:///" - // }; + // Keep the native scan and its file extensions: each extension owns the + // hard-link lease through optimizer rewrites and execution streams. let object_store_url = "file:///"; self.create_parquet_physical_plan( execution_plans, @@ -756,34 +733,55 @@ impl TableProvider for StandardTableProvider { let time_partition = object_store_format.time_partition; let mut time_filters = extract_primary_filter(filters, &time_partition); - if is_within_staging_window(&time_filters) { - self.get_staging_execution_plan( - &mut execution_plans, - projection, - filters, - limit, - state, - ) - .await?; + let include_staging = is_within_staging_window(&time_filters); + // Capture before reading the snapshot used by this scan. Upload publishes + // the manifest before unlinking staging: either source will then survive. + let (staged_records, staged_files) = if include_staging { + match PARSEABLE.get_stream(&self.stream, &self.tenant_id) { + Ok(staging) => { + // Preserve memory-first capture. A future coherent staging + // API must make this handoff atomic with conversion as well. + let records = staging + .recordbatches_cloned(&self.schema) + .map_err(|error| { + DataFusionError::Execution(format!( + "Cannot capture staging record batches: {error}" + )) + })?; + let files = capture_staging_parquet(&staging.data_path)?; + (records, files) + } + Err(_) => (Vec::new(), Vec::new()), + } + } else { + (Vec::new(), Vec::new()) }; let mut merged_snapshot = Snapshot::default(); if PARSEABLE.options.mode == Mode::Query || PARSEABLE.options.mode == Mode::Prism { - let obs = PARSEABLE + let objects = PARSEABLE .metastore .get_all_stream_jsons(&self.stream, None, &self.tenant_id, false) - .await; - if let Ok(obs) = obs { - for ob in obs { - if let Ok(object_store_format) = - serde_json::from_slice::(&ob) - { - let snapshot = object_store_format.snapshot; - for manifest in snapshot.manifest_list { - merged_snapshot.manifest_list.push(manifest); - } - } - } + .await + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + for object in objects { + let format: ObjectStoreFormat = serde_json::from_slice(&object) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + merged_snapshot + .manifest_list + .extend(format.snapshot.manifest_list); } + } else if include_staging { + // The first read supplies stream settings only; its snapshot may + // predate the staging capture and must not be used for this scan. + let refreshed: ObjectStoreFormat = serde_json::from_slice( + &PARSEABLE + .metastore + .get_stream_json(&self.stream, false, &self.tenant_id, false) + .await + .map_err(|e| DataFusionError::Plan(e.to_string()))?, + ) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + merged_snapshot = refreshed.snapshot; } else { merged_snapshot = object_store_format.snapshot; } @@ -818,6 +816,46 @@ impl TableProvider for StandardTableProvider { ) .await?; + // Compare only with files actually selected for the remote scan, before + // hot-tier routing mutates that list. Pruned/limit-excluded files must not + // hide a staging copy. Use exactly the uploader's path mapping. + if include_staging { + let selected_remote = manifest_files + .iter() + .map(|file| file.file_path.as_str()) + .collect::>(); + let staged_files = + staging_parquet::exclude_uploaded(staged_files, &selected_remote, |path| { + let filename = + path.file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "Invalid staging Parquet filename: {}", + path.display() + )) + })?; + let relative = crate::storage::object_storage::stream_relative_path( + &self.stream, + filename, + &object_store_format.custom_partition, + &self.tenant_id, + ); + Ok(glob_storage + .absolute_url(relative_path::RelativePath::new(&relative)) + .to_string()) + })?; + self.get_staging_execution_plan( + &mut execution_plans, + projection, + filters, + limit, + state, + (staged_records, staged_files), + ) + .await?; + } + if manifest_files.is_empty() { return self.final_plan(execution_plans, projection); } diff --git a/src/query/stream_schema_provider/staging_parquet.rs b/src/query/stream_schema_provider/staging_parquet.rs new file mode 100644 index 000000000..0a71db8c2 --- /dev/null +++ b/src/query/stream_schema_provider/staging_parquet.rs @@ -0,0 +1,694 @@ +//! Query-owned references to immutable, published staging Parquet files. + +use std::{ + collections::HashSet, + fs, + path::{Path, PathBuf}, + sync::Arc, +}; + +use datafusion::{ + datasource::listing::PartitionedFile, + error::{DataFusionError, Result}, +}; + +/// Keep the original name for canonical upload-path comparison; scans use only +/// the leased path. Published `.parquet` files must never be rewritten in place. +pub(super) struct CapturedFile { + original: PathBuf, + file: PartitionedFile, +} + +const LEASE_PREFIX: &str = ".query-parquet-"; +const OWNER_LOCK: &str = ".owner.lock"; + +#[derive(Debug)] +struct StagingLease { + directory: Option, + _owner: Option, +} + +impl StagingLease { + fn path(&self) -> &Path { + self.directory + .as_ref() + .expect("live lease directory") + .path() + } + + fn new(directory: &Path) -> Result> { + let mut lease = Self { + directory: Some( + tempfile::Builder::new() + .prefix(LEASE_PREFIX) + .tempdir_in(directory) + .map_err(|error| capture_error(directory, error))?, + ), + _owner: None, + }; + let pending = lease.path().join(".owner.pending"); + let owner = fs::OpenOptions::new() + .read(true) + .write(true) + .create_new(true) + .open(&pending) + .map_err(|error| capture_error(directory, error))?; + owner + .try_lock() + .map_err(|error| capture_error(directory, error))?; + lease._owner = Some(owner); + // Publish the ownership file only AFTER taking its lock and BEFORE + // creating hard links. A sweeper cannot mistake setup for an orphan. + fs::rename(&pending, lease.path().join(OWNER_LOCK)) + .map_err(|error| capture_error(directory, error))?; + Ok(Arc::new(lease)) + } +} + +impl Drop for StagingLease { + fn drop(&mut self) { + if let Some(directory) = self.directory.take() { + let path = directory.path().to_path_buf(); + // Keep ownership locked during removal. close(), unlike TempDir's + // default Drop, reports cleanup failures instead of hiding them. + if let Err(error) = directory.close() { + tracing::warn!(path = %path.display(), %error, "Cannot remove staging query lease"); + } + } + } +} + +/// Runs during query capture, outside stream-registry locks. Safe while serving +/// active queries, including in another server process. +/// Only the OS lock proves a lease is abandoned; never infer it from age/PID. +/// Unpublished ownership files are skipped (setup has not created any links). +pub(crate) fn reclaim_staging_leases(directory: &Path) { + let entries = match fs::read_dir(directory) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return, + Err(error) => { + tracing::warn!(path = %directory.display(), %error, "Cannot list staging query leases"); + return; + } + }; + for entry in entries { + let result = (|| -> std::io::Result<()> { + let entry = entry?; + if !entry + .file_name() + .to_string_lossy() + .starts_with(LEASE_PREFIX) + || !entry.file_type()?.is_dir() + { + return Ok(()); + } + let path = entry.path(); + let owner_path = path.join(OWNER_LOCK); + // Avoid following links outside our private, generated directory. + if !fs::symlink_metadata(&owner_path)?.file_type().is_file() { + return Ok(()); + } + let owner = fs::OpenOptions::new() + .read(true) + .write(true) + .open(owner_path)?; + match owner.try_lock() { + Ok(()) => { + fs::remove_dir_all(&path)?; + tracing::debug!(path = %path.display(), "Removed abandoned staging query lease"); + } + Err(std::fs::TryLockError::WouldBlock) => {} + Err(std::fs::TryLockError::Error(error)) => return Err(error), + } + Ok(()) + })(); + if let Err(error) = result { + // A concurrent drop/sweep may already have removed the directory; + // a not-yet-published owner lock is intentionally skipped too. + if error.kind() != std::io::ErrorKind::NotFound { + tracing::warn!(path = %directory.display(), %error, "Cannot reclaim staging query lease"); + } + } + } +} + +fn capture_error(path: &Path, error: impl std::fmt::Display) -> DataFusionError { + DataFusionError::Execution(format!( + "Cannot capture staging Parquet {} for query: {error}", + path.display() + )) +} + +pub(super) fn capture_staging_parquet(directory: &Path) -> Result> { + reclaim_staging_leases(directory); + let entries = match fs::read_dir(directory) { + Ok(entries) => entries, + // Streams that have never staged data have no directory yet. Once a + // filename is observed, any disappearance is an error, never an omission. + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(capture_error(directory, error)), + }; + let mut paths = Vec::new(); + for entry in entries { + let path = entry + .map_err(|error| capture_error(directory, error))? + .path(); + if path + .extension() + .is_some_and(|extension| extension == "parquet") + { + paths.push(path); + } + } + paths.sort_by(|a, b| b.cmp(a)); + capture_paths(directory, paths) +} + +fn capture_paths(directory: &Path, paths: Vec) -> Result> { + if paths.is_empty() { + return Ok(Vec::new()); + } + // Same directory/filesystem as the sources. Neither this directory nor its + // `.lease` children match the upload scanner's top-level `.parquet` filter. + let lease = StagingLease::new(directory)?; + let mut files = Vec::with_capacity(paths.len()); + for (index, original) in paths.into_iter().enumerate() { + let leased = lease.path().join(format!("{index}.lease")); + // Fail closed on unsupported hard links, cross-device links, permission + // errors and upload cleanup winning this race. RAII rolls back all links. + fs::hard_link(&original, &leased).map_err(|error| capture_error(&original, error))?; + let metadata = leased + .metadata() + .map_err(|error| capture_error(&original, error))?; + if !metadata.is_file() { + return Err(capture_error(&original, "not a regular file")); + } + // Validate only the retained immutable inode. The uploader uses the + // same predicate: corrupt footers and zero-row files remain excluded, + // while stat/open/read failures fail the query clearly. + if !crate::parseable::Stream::validate_parquet_file(&leased, "query staging") + .map_err(|error| capture_error(&original, error))? + { + fs::remove_file(&leased).map_err(|error| capture_error(&original, error))?; + continue; + } + let location = object_store::path::Path::from_absolute_path(&leased) + .map_err(|error| capture_error(&original, error))?; + let file = PartitionedFile::new(location.to_string(), metadata.len()) + .with_extension(Arc::clone(&lease)); + files.push(CapturedFile { original, file }); + } + // FileScanConfig and ParquetFileReader retain PartitionedFile extensions. + // Their Arc clones keep these links alive even if the plan is dropped while + // its stream is still running. No execution wrapper or custom reader needed. + Ok(files) +} + +pub(super) fn exclude_uploaded( + files: Vec, + selected_remote: &HashSet<&str>, + canonical_path: impl Fn(&Path) -> Result, +) -> Result> { + let mut retained = Vec::with_capacity(files.len()); + for captured in files { + if !selected_remote.contains(canonical_path(&captured.original)?.as_str()) { + retained.push(captured.file); + } + } + Ok(retained) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::{ + execution::SendableRecordBatchStream, + physical_plan::{ExecutionPlan, Statistics, collect, execute_stream}, + prelude::{SessionConfig, SessionContext}, + }; + use futures_util::TryStreamExt; + use parquet::arrow::ArrowWriter; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])) + } + + fn write_file(directory: &Path, name: &str, values: Vec) -> PathBuf { + let path = directory.join(name); + let batch = + RecordBatch::try_new(schema(), vec![Arc::new(Int64Array::from(values))]).unwrap(); + let mut writer = + ArrowWriter::try_new(fs::File::create(&path).unwrap(), schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + path + } + + fn lease_directories(directory: &Path) -> Vec { + fs::read_dir(directory) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| { + path.file_name() + .unwrap() + .to_string_lossy() + .starts_with(".query-parquet-") + }) + .collect() + } + + fn take_files(files: Vec) -> Vec { + files.into_iter().map(|captured| captured.file).collect() + } + + async fn plan(context: &SessionContext, files: Vec) -> Arc { + super::super::parquet_scan_plan( + schema(), + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + vec![files], + Statistics::new_unknown(&schema()), + None, + &[], + None, + &context.state(), + 2, + ) + .await + .unwrap() + } + + fn count_sum(batches: &[RecordBatch]) -> (usize, i64) { + let count = batches.iter().map(RecordBatch::num_rows).sum(); + let sum = batches + .iter() + .map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .sum::() + }) + .sum(); + (count, sum) + } + + #[tokio::test] + async fn staging_lease_survives_original_deletion_and_native_plan_replacement() { + let directory = tempfile::tempdir().unwrap(); + let first = write_file(directory.path(), "first.parquet", vec![1, 2, 3]); + let second = write_file(directory.path(), "second.parquet", vec![10, 20]); + let context = SessionContext::new(); + let unleased_files = [&first, &second] + .into_iter() + .map(|path| { + PartitionedFile::new( + object_store::path::Path::from_absolute_path(path) + .unwrap() + .to_string(), + path.metadata().unwrap().len(), + ) + }) + .collect(); + let old_strategy = plan(&context, unleased_files).await; + let scan = plan( + &context, + take_files(capture_staging_parquet(directory.path()).unwrap()), + ) + .await; + assert!( + scan.name().contains("DataSourceExec"), + "must remain a native Parquet scan" + ); + #[allow(deprecated)] + let replacement = Arc::clone(&scan).with_new_children(vec![]).unwrap(); + drop(scan); + // Exercise the native optimizer's range splitting as well as replacement. + let mut config = datafusion::common::config::ConfigOptions::default(); + config.optimizer.repartition_file_min_size = 0; + let repartitioned = replacement + .repartitioned(4, &config) + .unwrap() + .expect("native Parquet range repartitioning remains available"); + drop(replacement); + // Deterministically reproduce upload cleanup AFTER planning, BEFORE I/O. + fs::remove_file(first).unwrap(); + fs::remove_file(second).unwrap(); + assert!( + collect(old_strategy, context.task_ctx()).await.is_err(), + "the original mutable-path strategy must reproduce file-not-found" + ); + assert_eq!( + count_sum(&collect(repartitioned, context.task_ctx()).await.unwrap()), + (5, 36) + ); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[tokio::test] + async fn staging_lease_stream_outlives_plan_and_cancellation_cleans_up() { + let directory = tempfile::tempdir().unwrap(); + let original = write_file(directory.path(), "rows.parquet", (1..=100).collect()); + let context = SessionContext::new_with_config(SessionConfig::new().with_batch_size(2)); + let scan = plan( + &context, + take_files(capture_staging_parquet(directory.path()).unwrap()), + ) + .await; + let stream: SendableRecordBatchStream = + execute_stream(Arc::clone(&scan), context.task_ctx()).unwrap(); + fs::remove_file(original).unwrap(); + drop(scan); + assert!(!lease_directories(directory.path()).is_empty()); + let batches = stream.try_collect::>().await.unwrap(); + assert_eq!(count_sum(&batches), (100, 5050)); + assert!(lease_directories(directory.path()).is_empty()); + + let original = write_file(directory.path(), "cancel.parquet", vec![7, 8, 9]); + let scan = plan( + &context, + take_files(capture_staging_parquet(directory.path()).unwrap()), + ) + .await; + let stream = execute_stream(Arc::clone(&scan), context.task_ctx()).unwrap(); + fs::remove_file(original).unwrap(); + drop(scan); + drop(stream); // Cancel without ever polling the stream. + assert!(lease_directories(directory.path()).is_empty()); + } + + #[test] + fn staging_lease_unexecuted_captures_are_independent_and_drop_cleanly() { + let directory = tempfile::tempdir().unwrap(); + let original = write_file(directory.path(), "rows.parquet", vec![1, 2]); + let first = capture_staging_parquet(directory.path()).unwrap(); + let second = capture_staging_parquet(directory.path()).unwrap(); + assert_eq!(first.len(), 1); // Lease files must never be recaptured. + assert_eq!(second.len(), 1); + assert_eq!(lease_directories(directory.path()).len(), 2); + fs::remove_file(original).unwrap(); + drop(first); + assert_eq!(lease_directories(directory.path()).len(), 1); + drop(second); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[test] + fn staging_lease_setup_failure_rolls_back_all_links_and_reports_missing_source() { + let directory = tempfile::tempdir().unwrap(); + let existing = write_file(directory.path(), "rows.parquet", vec![1]); + let missing = directory.path().join("deleted-after-listing.parquet"); + let result = capture_paths(directory.path(), vec![existing, missing]); + let error = result + .err() + .expect("must fail instead of returning incomplete rows") + .to_string(); + assert!(error.contains("deleted-after-listing.parquet"), "{error}"); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[test] + fn staging_lease_listing_error_fails_clearly() { + let directory = tempfile::tempdir().unwrap(); + let not_directory = directory.path().join("not-a-directory"); + fs::write(¬_directory, b"file").unwrap(); + let error = capture_staging_parquet(¬_directory) + .err() + .unwrap() + .to_string(); + assert!(error.contains("not-a-directory"), "{error}"); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[cfg(unix)] + #[test] + fn staging_lease_metadata_failure_rolls_back_links() { + let directory = tempfile::tempdir().unwrap(); + let existing = write_file(directory.path(), "rows.parquet", vec![1]); + let dangling = directory.path().join("dangling.parquet"); + std::os::unix::fs::symlink(directory.path().join("missing-target"), &dangling).unwrap(); + let error = capture_paths(directory.path(), vec![existing, dangling]) + .err() + .unwrap() + .to_string(); + assert!(error.contains("dangling.parquet"), "{error}"); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[tokio::test] + async fn staging_lease_invalid_footer_and_empty_parquet_keep_prior_skip_policy() { + let directory = tempfile::tempdir().unwrap(); + fs::write( + directory.path().join("corrupt.parquet"), + b"invalid footer with enough bytes", + ) + .unwrap(); + write_file(directory.path(), "empty.parquet", vec![]); + let valid = write_file(directory.path(), "valid.parquet", vec![1, 2, 3]); + let files = capture_staging_parquet(directory.path()).unwrap(); + assert_eq!(files.len(), 1, "only valid, nonempty Parquet is eligible"); + let context = SessionContext::new(); + let scan = plan(&context, take_files(files)).await; + fs::remove_file(valid).unwrap(); + assert_eq!( + count_sum(&collect(scan, context.task_ctx()).await.unwrap()), + (3, 6) + ); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[tokio::test] + async fn staging_lease_failed_plan_setup_and_unexecuted_plan_drop_clean_up() { + let directory = tempfile::tempdir().unwrap(); + write_file(directory.path(), "rows.parquet", vec![1]); + let context = SessionContext::new(); + let files = take_files(capture_staging_parquet(directory.path()).unwrap()); + let result = super::super::parquet_scan_plan( + schema(), + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + vec![files], + Statistics::new_unknown(&schema()), + None, + &[ + datafusion::prelude::col("p_timestamp").gt(datafusion::prelude::Expr::Literal( + datafusion::scalar::ScalarValue::TimestampMillisecond(Some(0), None), + None, + )), + ], + None, + &context.state(), + 2, + ) + .await; + assert!(result.is_err()); + assert!(lease_directories(directory.path()).is_empty()); + let scan = plan( + &context, + take_files(capture_staging_parquet(directory.path()).unwrap()), + ) + .await; + drop(scan); + assert!(lease_directories(directory.path()).is_empty()); + } + + #[tokio::test] + async fn staging_lease_remote_overlap_counts_each_file_once_with_canonical_partitions() { + use crate::storage::object_storage::stream_relative_path; + use crate::storage::{LocalFS, ObjectStorage}; + use relative_path::RelativePath; + let directory = tempfile::tempdir().unwrap(); + let staging = directory.path().join("staging"); + let remote = directory.path().join("remote"); + fs::create_dir(&staging).unwrap(); + fs::create_dir(&remote).unwrap(); + let name = "date=2026-10-21.hour=10.minute=30.region=us.service=api.host.data.ULID.parquet"; + let original = write_file(&staging, name, vec![1, 2, 3]); + let local_only = write_file(&staging, "local.parquet", vec![10, 20]); + let store = LocalFS::new(remote.clone()); + let tenant = Some("tenant-a".to_owned()); + let custom = Some("region,service".to_owned()); + let canonical = |path: &Path| -> Result { + let relative = stream_relative_path( + "events", + path.file_name().unwrap().to_str().unwrap(), + &custom, + &tenant, + ); + Ok(store.absolute_url(RelativePath::new(&relative)).to_string()) + }; + // Capture before publication, then publish the same immutable bytes. + let captures = capture_staging_parquet(&staging).unwrap(); + let relative = stream_relative_path("events", name, &custom, &tenant); + assert_eq!( + relative, + format!( + "tenant-a/events/date=2026-10-21/hour=10/minute=30/region=us/service=api/host.data.ULID.parquet" + ) + ); + let uploaded = remote.join(&relative); + fs::create_dir_all(uploaded.parent().unwrap()).unwrap(); + fs::copy(&original, &uploaded).unwrap(); + let uploaded_key = canonical(&original).unwrap(); + let selected = HashSet::from([uploaded_key.as_str()]); + let mut files = exclude_uploaded(captures, &selected, canonical).unwrap(); + assert_eq!(files.len(), 1); + files.push(PartitionedFile::new( + uploaded_key.clone(), + uploaded.metadata().unwrap().len(), + )); + let context = SessionContext::new(); + let scan = plan(&context, files).await; + fs::remove_file(original).unwrap(); + fs::remove_file(local_only).unwrap(); + assert_eq!( + count_sum(&collect(scan, context.task_ctx()).await.unwrap()), + (5, 36) + ); + assert!(lease_directories(&staging).is_empty()); + } + + // Invoked only by the parent test as a separate process. Killing this + // process bypasses destructors, exactly like a server crash/SIGKILL. + #[test] + fn staging_lease_process_owner_fixture() { + use std::io::{Read, Write}; + let Some(directory) = std::env::var_os("PARSEABLE_TEST_LEASE_OWNER_DIRECTORY") else { + return; + }; + let directory = PathBuf::from(directory); + let _files = capture_staging_parquet(&directory).unwrap(); + let lease = lease_directories(&directory).pop().unwrap(); + println!("\nLEASE_OWNER:{}", lease.display()); + std::io::stdout().flush().unwrap(); + std::io::stdin().read_exact(&mut [0]).unwrap(); + } + + #[test] + fn staging_lease_sweeper_preserves_live_process_and_reclaims_killed_owner() { + use std::io::BufRead; + use std::process::{Child, Command, Stdio}; + struct ChildGuard(Child); + impl Drop for ChildGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } + } + let directory = tempfile::tempdir().unwrap(); + let original = write_file(directory.path(), "rows.parquet", vec![1, 2, 3]); + let fixture = concat!(module_path!(), "::staging_lease_process_owner_fixture"); + // libtest names omit the leading crate name returned by module_path!. + let fixture = fixture.split_once("::").unwrap().1; + let child = Command::new(std::env::current_exe().unwrap()) + .args(["--exact", fixture, "--nocapture", "--test-threads=1"]) + .env("PARSEABLE_TEST_LEASE_OWNER_DIRECTORY", directory.path()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); + let mut child = ChildGuard(child); + let stdout = child.0.stdout.take().unwrap(); + let lease = std::io::BufReader::new(stdout) + .lines() + .take(20) + .find_map(|line| { + line.unwrap() + .strip_prefix("LEASE_OWNER:") + .map(PathBuf::from) + }) + .expect("child must publish its locked lease before parent sweeps"); + fs::remove_file(original).unwrap(); + reclaim_staging_leases(directory.path()); + assert!( + lease.join("0.lease").exists(), + "another process still owns this query" + ); + child.0.kill().unwrap(); + child.0.wait().unwrap(); + reclaim_staging_leases(directory.path()); + assert!( + !lease.exists(), + "OS releases the dead owner's lock; orphan must be reclaimed" + ); + } + + #[test] + fn staging_lease_constructor_defers_orphan_cleanup_until_capture() { + use crate::{cli::Options, metadata::LogStreamMetadata, parseable::Stream}; + + let directory = tempfile::tempdir().unwrap(); + let options = Arc::new(Options { + local_staging_path: directory.path().to_path_buf(), + ..Default::default() + }); + let staging = options.local_stream_data_path("deferred", &None); + fs::create_dir_all(&staging).unwrap(); + let original = write_file(&staging, "rows.parquet", vec![1, 2, 3]); + + // Model a crashed owner: retain its hard-link directory but release the + // ownership lock. No process or timing assumptions are needed. + let mut abandoned = Arc::try_unwrap(StagingLease::new(&staging).unwrap()).unwrap(); + fs::hard_link(&original, abandoned.path().join("0.lease")).unwrap(); + let orphan = abandoned.directory.take().unwrap().keep(); + drop(abandoned); + assert!(orphan.join("0.lease").exists()); + + // get_or_create constructs streams under the shared registry write + // lock. Construction must not scan/remove filesystem entries there. + let stream = Stream::new( + options, + "deferred", + LogStreamMetadata::default(), + None, + &None, + ); + assert!( + orphan.join("0.lease").exists(), + "stream construction must defer orphan cleanup" + ); + + let captured = capture_staging_parquet(&stream.data_path).unwrap(); + assert!( + !orphan.exists(), + "the first query capture must reclaim the abandoned lease" + ); + assert_eq!(captured.len(), 1); + assert!(original.exists()); + drop(captured); + assert!(lease_directories(&staging).is_empty()); + } + + #[test] + fn staging_lease_sweeper_skips_unpublished_setup() { + let directory = tempfile::tempdir().unwrap(); + let pending = tempfile::Builder::new() + .prefix(LEASE_PREFIX) + .tempdir_in(directory.path()) + .unwrap(); + fs::write(pending.path().join(".owner.pending"), []).unwrap(); + reclaim_staging_leases(directory.path()); + assert!(pending.path().exists()); + } + + #[test] + fn staging_lease_dedup_keeps_files_not_selected_by_remote_pruning_or_limit() { + let directory = tempfile::tempdir().unwrap(); + write_file(directory.path(), "rows.parquet", vec![1]); + let files = capture_staging_parquet(directory.path()).unwrap(); + let retained = exclude_uploaded(files, &HashSet::new(), |_| { + Ok("published-but-not-selected".into()) + }) + .unwrap(); + assert_eq!(retained.len(), 1); + drop(retained); + assert!(lease_directories(directory.path()).is_empty()); + } +} diff --git a/src/storage/object_storage.rs b/src/storage/object_storage.rs index 368326a67..c11404ad6 100644 --- a/src/storage/object_storage.rs +++ b/src/storage/object_storage.rs @@ -1362,7 +1362,7 @@ async fn process_schema_files( } /// Builds the stream relative path for a file -fn stream_relative_path( +pub(crate) fn stream_relative_path( stream_name: &str, filename: &str, custom_partition: &Option,