From 7bff9a9fef0929d033cef633c9caa79370fe728f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 10:57:33 +0800 Subject: [PATCH] perf: share Metric export writer and payload budgets Signed-off-by: jeremyhi --- src/common/datasource/src/parquet_writer.rs | 79 +++- src/object-store/tests/object_store_test.rs | 2 +- src/operator/src/statement/copy_table_to.rs | 253 +++++++++- src/operator/src/statement/export_database.rs | 91 ++-- .../src/statement/export_logical_tables.rs | 178 ++++--- .../statement/export_logical_tables/tests.rs | 443 +++++++++++++----- .../export_logical_tables/writers.rs | 245 ++++++++++ 7 files changed, 1067 insertions(+), 224 deletions(-) create mode 100644 src/operator/src/statement/export_logical_tables/writers.rs diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index 955df4032cb..a760de187d2 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -186,9 +186,12 @@ impl ParquetFileWriter { } async fn write_bytes(&mut self, bytes: Vec) -> Result<()> { - if !bytes.is_empty() { + let bytes = Bytes::from(bytes); + let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + // Slices retain the complete encoded allocation until its last submission. + for offset in (0..bytes.len()).step_by(chunk) { self.sink - .write(bytes) + .write(bytes.slice(offset..(offset + chunk).min(bytes.len()))) .await .context(error::WriteObjectSnafu { path: &self.path })?; } @@ -328,7 +331,9 @@ mod tests { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { self.0.write(bytes).await } - async fn close(&mut self) -> object_store::Result { + async fn close( + &mut self, + ) -> object_store::Result { self.0.close().await?; Err(object_store::Error::new( object_store::ErrorKind::Unexpected, @@ -391,6 +396,74 @@ mod tests { } } + #[tokio::test] + async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() { + use std::sync::Mutex; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + struct ObservedWriter(oio::Writer, Arc>>); + impl oio::Write for ObservedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.1.lock().unwrap().push(bytes.len()); + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let sizes = Arc::new(Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone())) + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut state = 17u64; + let values = (0..600_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as i64 + }) + .collect::>(); + let array = Arc::new(Int64Array::from(values)) as arrow::array::ArrayRef; + let batch = RecordBatch::try_from_iter([("a", array.clone()), ("b", array)]).unwrap(); + let mut writer = + ParquetFileWriter::open(batch.schema(), store.clone(), "large.parquet", 1, None) + .await + .unwrap(); + // No storage-layer chunking: observe the application's actual submissions. + writer.sink = store.writer("large.parquet").await.unwrap(); + writer.write(batch.clone(), None).await.unwrap(); + writer.finish(None).await.unwrap(); + let sizes = sizes.lock().unwrap().clone(); + let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + assert!(sizes.iter().sum::() > limit); + assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); + let actual = read(&store, "large.parquet") + .await + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(), + batch + ); + } + #[tokio::test] async fn row_and_byte_limits_split_batches() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index c1537b5b230..07b185de6cd 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -290,7 +290,7 @@ async fn test_s3_backend() -> Result<()> { .secret_access_key(&env::var("GT_S3_ACCESS_KEY")?) .region(&env::var("GT_S3_REGION")?) .bucket(&bucket); - if let Ok(endpoint) = env::var("GT_S3_ENDPOINT") { + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") { builder = builder.endpoint(&endpoint); } diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 21b82f51a3b..a415a7c3904 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -22,6 +22,7 @@ use common_datasource::file_format::csv::stream_to_csv; use common_datasource::file_format::json::stream_to_json; use common_datasource::file_format::parquet::stream_to_parquet; use common_datasource::object_store::build_backend_for_write_with_path; +use common_datasource::parquet_writer::ParquetFileWriter; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::{ @@ -32,16 +33,22 @@ use common_telemetry::{debug, tracing}; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference as DfTableReference; use datafusion_expr::LogicalPlanBuilder; +use futures::StreamExt; use object_store::ObjectStore; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::CopyTableRequest; use table::table::adapter::DfTableProviderAdapter; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result}; use crate::statement::StatementExecutor; +use crate::statement::export_logical_tables::writers::ExportWriteBudget; +use crate::statement::export_logical_tables::{ + LogicalTableExportLimits, map_writer_error, rows_within_budget, +}; // The buffer size should be greater than 5MB (minimum multipart upload size). /// Buffer size to flush data to object stores. @@ -117,6 +124,17 @@ impl StatementExecutor { table: TableRef, req: CopyTableRequest, query_ctx: QueryContextRef, + ) -> Result { + self.copy_captured_table_to_managed(table, req, query_ctx, None) + .await + } + + pub(crate) async fn copy_captured_table_to_managed( + &self, + table: TableRef, + req: CopyTableRequest, + query_ctx: QueryContextRef, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, ) -> Result { let info = table.table_info(); let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name); @@ -165,7 +183,7 @@ impl StatementExecutor { } = &req; debug!("Copy table: {table_id} to location: {location}"); - self.copy_to_file(&format, output, location, connection) + self.copy_to_file_managed(&format, output, location, connection, managed) .await } @@ -176,9 +194,25 @@ impl StatementExecutor { location: &str, connection: &HashMap, ) -> Result { - let output = output - .map_dictionary_to_values() - .context(error::BuildRecordBatchSnafu)?; + self.copy_to_file_managed(format, output, location, connection, None) + .await + } + + async fn copy_to_file_managed( + &self, + format: &Format, + output: Output, + location: &str, + connection: &HashMap, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, + ) -> Result { + let output = if managed.is_none() { + output + .map_dictionary_to_values() + .context(error::BuildRecordBatchSnafu)? + } else { + output + }; let stream = match output.data { OutputData::Stream(stream) => stream, OutputData::RecordBatches(record_batches) => record_batches.as_stream(), @@ -192,7 +226,214 @@ impl StatementExecutor { let filename = backend.object_path.context(error::UnexpectedSnafu { violated: format!("Expected filename, path: {location}"), })?; - self.stream_to_file(stream, format, backend.object_store, &filename) + if let Some((budget, token)) = managed { + stream_to_managed_parquet(stream, backend.object_store, &filename, budget, token).await + } else { + self.stream_to_file(stream, format, backend.object_store, &filename) + .await + } + } +} + +pub(crate) async fn stream_to_managed_parquet( + mut stream: SendableRecordBatchStream, + store: ObjectStore, + path: &str, + budget: &ExportWriteBudget, + token: &CancellationToken, +) -> Result { + use common_recordbatch::{ + RecordBatch, map_dictionary_to_values, map_dictionary_to_values_schema, + }; + let original = stream.schema(); + let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone()); + let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone()); + let output_schema = if json { + mapped_schema + } else { + expanded_schema.clone() + }; + let mut writer = + ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None) .await + .context(error::WriteStreamToFileSnafu { path })?; + let mut started = false; + let result = async { + let mut rows = 0; + loop { + let batch = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + batch = stream.next() => batch, + }; + let Some(batch) = batch else { break }; + let batch = batch + .context(error::BuildRecordBatchSnafu)? + .into_df_record_batch(); + let mut offset = 0; + while offset < batch.num_rows() { + let (conversion, retained) = ExportWriteBudget::conversion_budget( + &batch, + LogicalTableExportLimits::default().conversion_bytes, + )?; + let input = batch.clone(); + // JSON escaping and number formatting can expand binary JSON. + let expansion = if json { 8 } else { 1 }; + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, offset, input.num_rows(), conversion / expansion) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(2 * expansion)); + let permit = budget.reserve(reservation, token).await?; + let batch = batch.clone(); + let (original, expanded_schema, output_schema) = ( + original.clone(), + expanded_schema.clone(), + output_schema.clone(), + ); + let (batch, len) = common_runtime::spawn_blocking_global(move || { + let mut batch = RecordBatch::from_df_record_batch( + original.clone(), + batch.slice(offset, len), + ); + if expand { + batch = map_dictionary_to_values(batch, &original, &expanded_schema) + .context(error::BuildRecordBatchSnafu)?; + } + if json { + batch = map_json_type_to_string(batch, &expanded_schema, &output_schema) + .context(error::BuildRecordBatchSnafu)?; + } + let batch = batch.into_df_record_batch(); + ensure!( + batch.get_array_memory_size() <= reservation, + error::LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok::<_, error::Error>((batch, len)) + }) + .await + .context(error::JoinTaskSnafu)??; + started = true; + let write = writer.write(batch, Some(token)).await; + drop(permit); + write.map_err(|error| map_writer_error(error, path))?; + rows += len; + offset += len; + } + } + started = true; + writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, path))?; + Ok(rows) + } + .await; + if result.is_err() { + token.cancel(); + if started && let Err(error) = writer.abort().await { + common_telemetry::warn!(error; "Failed to abort ordinary export file"); + } + } + result +} + +#[cfg(test)] +mod tests { + use arrow::array::{ArrayRef, BinaryArray, DictionaryArray, Int32Array, StringArray}; + use arrow::datatypes::Int32Type; + use common_recordbatch::{RecordBatch, RecordBatches}; + use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + + use super::*; + + #[tokio::test] + async fn managed_copy_preserves_dictionary_json_and_empty_schema() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "host", + ConcreteDataType::dictionary_datatype( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ), + true, + ), + ColumnSchema::new("json", ConcreteDataType::json_datatype(), true), + ])); + let dictionary = DictionaryArray::::new( + Int32Array::from(vec![Some(0), None]), + Arc::new(StringArray::from(vec!["tag"])), + ); + let json = + datatypes::types::parse_string_to_jsonb(r#"{"value":"escaped\ntext","n":123456789}"#) + .unwrap(); + let arrays = vec![ + Arc::new(dictionary) as ArrayRef, + Arc::new(BinaryArray::from(vec![Some(json.as_slice()), None])), + ]; + let batch = + arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays) + .unwrap(); + let batch = RecordBatch::from_df_record_batch(schema.clone(), batch); + for empty in [false, true] { + let batches = RecordBatches::try_new( + schema.clone(), + if empty { vec![] } else { vec![batch.clone()] }, + ) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let managed = stream_to_managed_parquet( + batches.as_stream(), + store.clone(), + "managed.parquet", + &budget, + &token, + ) + .await + .unwrap(); + let output = Output::new_with_record_batches(batches) + .map_dictionary_to_values() + .unwrap(); + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected batches"); + }; + let stream = SendableRecordBatchMapper::new( + batches.as_stream(), + map_json_type_to_string, + map_json_type_to_string_schema, + ); + let ordinary = stream_to_parquet( + Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))), + store.clone(), + "ordinary.parquet", + WRITE_CONCURRENCY, + ) + .await + .unwrap(); + assert_eq!(managed, ordinary); + let mut outputs = Vec::new(); + for name in ["managed.parquet", "ordinary.parquet"] { + let reader = ParquetRecordBatchReaderBuilder::try_new( + store.read(name).await.unwrap().to_bytes(), + ) + .unwrap(); + let schema = reader.schema().clone(); + let values = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + outputs.push((schema, values)); + } + assert_eq!(outputs[0], outputs[1]); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } } } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index 7c5de554487..ec0d8f08b48 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -37,6 +37,7 @@ use crate::statement::database_copy::{ DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory, validate_database_export_layout, }; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error}; use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; /// A validated request-scoped selection, not a metadata snapshot or an ACL token. @@ -202,46 +203,52 @@ impl StatementExecutor { } output_files.sort(); let req = &plan.request; - let rows = run_database_export_jobs( - plan.jobs, - parse_parallelism_from_option_map(&req.with), - cancellation, - |job, token| { - let ctx = ctx.clone(); - async move { - match job { - DatabaseExportJob::Metric(unit) => self - .export_logical_tables( - &unit, - &req.location, - &req.connection, - req.time_range.as_ref(), - LogicalTableExportLimits::default(), - &token, - ctx, - ) - .await - .map(|summary| summary.rows), - DatabaseExportJob::Ordinary { table, output } => { - let info = table.table_info(); - let copy = CopyTableRequest { - catalog_name: info.catalog_name.clone(), - schema_name: info.schema_name.clone(), - table_name: info.name.clone(), - location: output.location, - with: req.with.clone(), - connection: req.connection.clone(), - pattern: None, - direction: CopyDirection::Export, - timestamp_range: req.time_range, - limit: None, - }; - self.copy_captured_table_to(table, copy, ctx).await - } + let parallelism = parse_parallelism_from_option_map(&req.with); + let budget = ExportWriteBudget::new(parallelism); + let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| { + let ctx = ctx.clone(); + let budget = budget.clone(); + async move { + match job { + DatabaseExportJob::Metric(unit) => self + .export_logical_tables_managed( + &unit, + &req.location, + &req.connection, + req.time_range.as_ref(), + LogicalTableExportLimits::default(), + &token, + ctx, + budget, + ) + .await + .map(|summary| summary.rows), + DatabaseExportJob::Ordinary { table, output } => { + let info = table.table_info(); + let copy = CopyTableRequest { + catalog_name: info.catalog_name.clone(), + schema_name: info.schema_name.clone(), + table_name: info.name.clone(), + location: output.location, + with: req.with.clone(), + connection: req.connection.clone(), + pattern: None, + direction: CopyDirection::Export, + timestamp_range: req.time_range, + limit: None, + }; + let _permit = budget.writer(&token).await?; + self.copy_captured_table_to_managed( + table, + copy, + ctx, + Some((&budget, &token)), + ) + .await } } - }, - ) + } + }) .await?; Ok(DatabaseExportSummary { rows, output_files }) } @@ -261,6 +268,7 @@ async fn run_database_export_jobs>>( loop { while first_error.is_none() && !cancellation.is_cancelled() + && !token.is_cancelled() && active.len() < parallelism.max(1) { let Some(job) = jobs.next() else { break }; @@ -284,11 +292,12 @@ async fn run_database_export_jobs>>( }; match result { Some(Ok(count)) => rows += count, - Some(Err(err)) if first_error.is_none() => { - first_error = Some(err); + Some(Err(err)) => { + if !cancellation.is_cancelled() || first_error.is_none() { + retain_error(&mut first_error, err); + } token.cancel(); } - Some(Err(err)) => common_telemetry::warn!(err; "Failed to drain database export job"), None => break, } } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index f52f4e96d90..a3f025d3627 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -19,6 +19,8 @@ //! exclusively by the export attempt. Completion and cleanup of closed files //! belong to the caller's chunk protocol, not individual file completion. +pub(crate) mod writers; + use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; @@ -54,6 +56,7 @@ use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; use crate::statement::StatementExecutor; use crate::statement::database_copy::DatabaseExportFile; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, Payload, TableWriters}; /// Export preprocessing and per-file limits. Query memory and spill remain /// governed by the query engine. @@ -103,7 +106,7 @@ pub struct LogicalTableExport { logical_tables: BTreeMap, } -struct LogicalTableProjection { +pub(crate) struct LogicalTableProjection { output: DatabaseExportFile, schema: SchemaRef, projection: Vec, @@ -308,6 +311,31 @@ impl StatementExecutor { limits: LogicalTableExportLimits, cancellation: &CancellationToken, query_ctx: QueryContextRef, + ) -> Result { + self.export_logical_tables_managed( + unit, + directory, + connection, + time_range, + limits, + &cancellation.child_token(), + query_ctx, + ExportWriteBudget::new(1), + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn export_logical_tables_managed( + &self, + unit: &LogicalTableExport, + directory: &str, + connection: &HashMap, + time_range: Option<&TimestampRange>, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + query_ctx: QueryContextRef, + budget: Arc, ) -> Result { limits.validate()?; let (store, stream) = tokio::select! { @@ -327,10 +355,11 @@ impl StatementExecutor { Ok((store, stream)) } => result?, }; - export_stream(unit, stream, &store, limits, cancellation).await + export_stream_managed(unit, stream, &store, limits, cancellation, budget).await } } +#[cfg(test)] async fn export_stream( unit: &LogicalTableExport, stream: SendableRecordBatchStream, @@ -338,19 +367,42 @@ async fn export_stream( limits: LogicalTableExportLimits, cancellation: &CancellationToken, ) -> Result { - let mut active = None; - let result = write_tables(unit, stream, store, limits, cancellation, &mut active).await; - if result.is_err() - && let Some(writer) = active - && let Err(cleanup_error) = writer - .writer - .abort() - .await - .map_err(|error| map_writer_error(error, &writer.path)) - { - common_telemetry::warn!(cleanup_error; "Failed to clean up incomplete Metric export file"); - } - result + export_stream_managed( + unit, + stream, + store, + limits, + &cancellation.child_token(), + ExportWriteBudget::new(1), + ) + .await +} + +async fn export_stream_managed( + unit: &LogicalTableExport, + stream: SendableRecordBatchStream, + store: &ObjectStore, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + budget: Arc, +) -> Result { + let mut writers = TableWriters::new(budget.clone()); + let result = write_tables( + unit, + stream, + store, + limits, + cancellation, + &mut writers, + &budget, + ) + .await; + let (summary, result) = match result { + Ok(summary) => (summary, Ok(())), + Err(error) => (LogicalTableExportSummary::default(), Err(error)), + }; + writers.drain(result, cancellation).await?; + Ok(summary) } fn check_cancelled(cancellation: &CancellationToken) -> Result<()> { @@ -367,7 +419,8 @@ async fn write_tables( store: &ObjectStore, limits: LogicalTableExportLimits, cancellation: &CancellationToken, - active: &mut Option, + writers: &mut TableWriters, + budget: &ExportWriteBudget, ) -> Result { let id_index = stream @@ -425,22 +478,19 @@ async fn write_tables( while end < batch.num_rows() && ids.value(end) == id { end += 1; } - if active.as_ref().is_some_and(|writer| writer.table_id != id) { - finish_active(active, cancellation).await?; + if writers.table_id().is_some_and(|last| last != id) { + writers.close_input(); } check_cancelled(cancellation)?; if let Some(file) = unit.logical_tables.get(&id) { - if active.is_none() { - *active = Some(ActiveWriter::open(id, file, store, limits).await?); + if writers.table_id().is_none() { + writers.open(id, file, store, limits, cancellation).await?; written.insert(id); summary.files += 1; } let projected = batch .project(&file.projection) .context(error::ProjectSchemaSnafu)?; - let writer = active.as_mut().context(error::UnexpectedSnafu { - violated: "missing logical writer", - })?; let mut offset = start; while offset < end { let (expanded, consumed) = expand_bounded_slice( @@ -449,15 +499,12 @@ async fn write_tables( offset, end, limits.conversion_bytes, + budget, + cancellation, ) .await?; check_cancelled(cancellation)?; - // Do not drop an in-flight file operation before cleanup. - writer - .writer - .write(expanded, Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; + writers.send(expanded, cancellation).await?; check_cancelled(cancellation)?; offset += consumed; summary.rows += consumed; @@ -468,12 +515,12 @@ async fn write_tables( start = end; } } - finish_active(active, cancellation).await?; + writers.close_input(); for (&id, file) in &unit.logical_tables { if !written.contains(&id) { check_cancelled(cancellation)?; - *active = Some(ActiveWriter::open(id, file, store, limits).await?); - finish_active(active, cancellation).await?; + writers.open(id, file, store, limits, cancellation).await?; + writers.close_input(); summary.files += 1; } } @@ -482,20 +529,18 @@ async fn write_tables( } struct ActiveWriter { - table_id: u32, path: String, writer: ParquetFileWriter, } impl ActiveWriter { async fn open( - id: u32, table: &LogicalTableProjection, store: &ObjectStore, limits: LogicalTableExportLimits, ) -> Result { let path = table.output.path.clone(); - let conditional = store.info().full_capability().write_with_if_not_exists; + let conditional = store.info().capability().write_with_if_not_exists; if !conditional { ensure!( !store @@ -521,39 +566,29 @@ impl ActiveWriter { ) .await .map_err(|error| map_writer_error(error, &path))?; - Ok(Self { - table_id: id, - path, - writer, - }) + Ok(Self { path, writer }) } } -async fn finish_active( - active: &mut Option, - cancellation: &CancellationToken, -) -> Result<()> { - if let Some(writer) = active.as_mut() { - writer - .writer - .finish(Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; - check_cancelled(cancellation)?; - *active = None; - } - Ok(()) -} - async fn expand_bounded_slice( batch: RecordBatch, schema: SchemaRef, start: usize, end: usize, - budget: usize, -) -> Result<(RecordBatch, usize)> { + requested: usize, + budget: &ExportWriteBudget, + cancellation: &CancellationToken, +) -> Result<(Payload, usize)> { + let (conversion, retained) = ExportWriteBudget::conversion_budget(&batch, requested)?; + let input = batch.clone(); + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, start, end, conversion) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(2)); + let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { - let len = rows_within_budget(&batch, start, end, budget)?; let slice = batch.slice(start, len); let arrays = slice .columns() @@ -562,13 +597,28 @@ async fn expand_bounded_slice( .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) .collect::>>()?; let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; - Ok((expanded, len)) + ensure!( + expanded.get_array_memory_size() <= reservation, + LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok(( + Payload { + batch: expanded, + permit, + }, + len, + )) }) .await .context(error::JoinTaskSnafu)? } -fn map_writer_error(source: common_datasource::error::Error, path: &str) -> error::Error { +pub(crate) fn map_writer_error( + source: common_datasource::error::Error, + path: &str, +) -> error::Error { match source { common_datasource::error::Error::ParquetWriteCancelled {} => { error::LogicalTableExportCancelledSnafu.build() @@ -639,12 +689,12 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { Ok(bytes.saturating_add(16)) } -fn rows_within_budget( +pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, end: usize, budget: usize, -) -> Result { +) -> Result<(usize, usize)> { let mut bytes = 0usize; let mut row = start; while row < end { @@ -663,7 +713,7 @@ fn rows_within_budget( reason: "one expanded logical row exceeds conversion byte budget" } ); - Ok(row - start) + Ok((row - start, bytes)) } #[cfg(test)] diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 4576c1e395d..74311963828 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -128,52 +128,56 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) #[tokio::test] async fn routes_across_batches_and_writes_empty_files() { - let unit = unit(); - let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let batches = vec![ - batch(vec![Some(1025), Some(1025)], vec![Some(""), None]), - batch( - vec![Some(1025), Some(1027), Some(1028)], - vec![Some("a"), Some("b"), Some("unselected")], - ), - ]; - let result = write_tables( - &unit, - stream(batches), - &store, - export_limits(), - &CancellationToken::new(), - &mut None, - ) - .await - .unwrap(); - assert_eq!( - result, - LogicalTableExportSummary { - rows: 4, - skipped_rows: 1, - files: 3 - } - ); - let (schema, cpu) = read(&store, "cpu.v1.parquet").await; - assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); - let values: Vec<_> = cpu - .iter() - .flat_map(|batch| { - batch - .column(0) - .as_string::() - .iter() - .map(|s| s.map(str::to_owned)) - }) - .collect(); - assert_eq!(values, vec![Some("".into()), None, Some("a".into())]); - let (schema, empty) = read(&store, "empty.parquet").await; - assert_eq!(schema.fields(), unit.logical_tables[&1026].schema.fields()); - assert!(empty.is_empty()); - let (schema, requests) = read(&store, "requests.parquet").await; - assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields()); - assert_eq!(requests[0].column(1).null_count(), 1); + for parallelism in [1, 4] { + let unit = unit(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batches = vec![ + batch(vec![Some(1025), Some(1025)], vec![Some(""), None]), + batch( + vec![Some(1025), Some(1027), Some(1028)], + vec![Some("a"), Some("b"), Some("unselected")], + ), + ]; + let budget = ExportWriteBudget::new(parallelism); + let result = export_stream_managed( + &unit, + stream(batches), + &store, + export_limits(), + &CancellationToken::new(), + budget.clone(), + ) + .await + .unwrap(); + assert_eq!( + result, + LogicalTableExportSummary { + rows: 4, + skipped_rows: 1, + files: 3 + } + ); + let (schema, cpu) = read(&store, "cpu.v1.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); + let values: Vec<_> = cpu + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + }) + .collect(); + assert_eq!(values, vec![Some("".into()), None, Some("a".into())]); + let (schema, empty) = read(&store, "empty.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1026].schema.fields()); + assert!(empty.is_empty()); + let (schema, requests) = read(&store, "requests.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields()); + assert_eq!(requests[0].column(1).null_count(), 1); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } } #[tokio::test] @@ -211,21 +215,16 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { ), ] { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let mut active = None; - let err = write_tables( + let err = export_stream( &unit, stream(batches), &store, limits, &CancellationToken::new(), - &mut active, ) .await .unwrap_err(); assert!(err.to_string().contains(message), "{err}"); - if let Some(writer) = active { - writer.writer.abort().await.unwrap(); - } } } @@ -233,17 +232,25 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { async fn existing_outputs_are_not_overwritten() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); store.write("cpu.v1.parquet", "keep").await.unwrap(); - let err = write_tables( + let err = export_stream( &unit(), stream(vec![batch(vec![Some(1025)], vec![None])]), &store, LogicalTableExportLimits::default(), &CancellationToken::new(), - &mut None, ) .await .unwrap_err(); - assert!(err.to_string().contains("already exists")); + if store.info().capability().write_with_if_not_exists { + assert!(matches!(err, error::Error::WriteStreamToFile { + source: common_datasource::error::Error::WriteObject { error, .. }, .. + } if error.kind() == object_store::ErrorKind::ConditionNotMatch)); + } else { + assert!(matches!( + err, + error::Error::InvalidLogicalTableExport { .. } + )); + } assert_eq!( store.read("cpu.v1.parquet").await.unwrap().to_bytes(), Bytes::from_static(b"keep") @@ -272,10 +279,10 @@ fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { ("histogram", Arc::new(histogram)), ]) .unwrap(); - assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap(), 1); + assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap().0, 1); // The dictionary and container overhead fit; the nested list elements do not. assert!(rows_within_budget(&batch, 0, 3, 4180).is_err()); - assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap(), 3); + assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap().0, 3); } #[test] @@ -440,60 +447,117 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter { } #[tokio::test] -async fn cancellation_waits_for_file_creation_before_cleanup() { - use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; - let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); - let access = - common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); - let store = build_backend_for_write( - &format!("{}/", directory.path().display()), - &HashMap::new(), - &access, - ) - .await - .unwrap(); - let started = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); - let factory: MockWriterFactory = Arc::new({ - let started = started.clone(); - let release = release.clone(); - move |_, _, inner| { - Box::new(PausedFileWriter { - inner: Some(inner), - started: started.clone(), - release: release.clone(), - }) +async fn cancellation_drains_storage_and_preserves_committed_files() { + for large in [false, true] { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); + let access = + common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); + let store = build_backend_for_write( + &format!("{}/", directory.path().display()), + &HashMap::new(), + &access, + ) + .await + .unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let started = started.clone(); + let release = release.clone(); + move |_, _, inner| { + Box::new(PausedFileWriter { + inner: Some(inner), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let cancellation = CancellationToken::new(); + let (unit, input) = if large { + let field = Field::new("value", DataType::Utf8, true); + let unit = LogicalTableExport::try_new( + table( + 1024, + "phy", + vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], + true, + ), + &[table(1025, "cpu.v1", vec![field], false)], + ) + .unwrap(); + let mut state = 17u64; + let strings = (0..512) + .map(|_| { + (0..32768) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + char::from(b' ' + (state % 95) as u8) + }) + .collect::() + }) + .collect::>(); + let input = RecordBatch::try_from_iter([ + ( + TABLE_ID, + Arc::new(UInt32Array::from(vec![1025; 512])) as ArrayRef, + ), + ("value", Arc::new(StringArray::from(strings))), + ]) + .unwrap(); + (unit, input) + } else { + (unit(), batch(vec![Some(1025)], vec![Some("a")])) + }; + let budget = ExportWriteBudget::new(1); + let mut limits = LogicalTableExportLimits::default(); + limits.writer.flush_threshold_bytes = 1024 * 1024; + let export = export_stream_managed( + &unit, + stream(vec![input]), + &store, + limits, + &cancellation, + budget.clone(), + ); + tokio::pin!(export); + tokio::select! { + result = &mut export => panic!("export completed before the file open: {result:?}"), + _ = started.notified() => {}, } - }); - let store = store.layer( - MockLayerBuilder::default() - .writer_factory(factory) - .build() - .unwrap(), - ); - let cancellation = CancellationToken::new(); - let unit = unit(); - let export = export_stream( - &unit, - stream(vec![batch(vec![Some(1025)], vec![Some("a")])]), - &store, - LogicalTableExportLimits::default(), - &cancellation, - ); - tokio::pin!(export); - tokio::select! { - result = &mut export => panic!("export completed before the file open: {result:?}"), - _ = started.notified() => {}, + if large { + assert!(budget.available().1 < 64 * 1024 * 1024); + } + let held = budget.available(); + cancellation.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + assert_eq!(budget.available().0, held.0); + if large { + // Cancelled admission releases its pending reservation; the worker + // still owns payload capacity while the storage operation is paused. + assert!(budget.available().1 < 64 * 1024 * 1024); + } + release.notify_one(); + let result = export.await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportCancelled { .. }) + )); + assert_eq!(store.exists("cpu.v1.parquet").await.unwrap(), !large); + if !large { + let (_, batches) = read(&store, "cpu.v1.parquet").await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); } - cancellation.cancel(); - assert!(futures::poll!(&mut export).is_pending()); - release.notify_one(); - let result = export.await; - assert!(matches!( - result, - Err(error::Error::LogicalTableExportCancelled { .. }) - )); - assert!(!store.exists("cpu.v1.parquet").await.unwrap()); } struct FailedAbortWriter(object_store::layers::mock::oio::Writer); @@ -574,3 +638,164 @@ async fn validates_membership_by_table_route() { } } } + +#[tokio::test] +async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() { + let input = batch(vec![Some(1025); 8192], vec![Some("shared"); 8192]); + let tiny = input.slice(0, 1); + let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])); + let projected = tiny.project(&[2]).unwrap(); + let full = input.project(&[2]).unwrap().get_array_memory_size(); + assert_eq!(projected.get_array_memory_size(), full); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let blocker = budget.reserve(64 * 1024 * 1024, &token).await.unwrap(); + let convert = expand_bounded_slice(projected, schema, 0, 1, 1024, &budget, &token); + tokio::pin!(convert); + assert!(futures::poll!(&mut convert).is_pending()); + drop(blocker); + let (payload, rows) = convert.await.unwrap(); + assert_eq!(rows, 1); + assert!( + 64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size() + ); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err()); +} + +#[tokio::test] +async fn groups_and_ordinary_files_share_writer_admission_and_drain() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use tokio::sync::Semaphore; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct CountedWriter { + inner: oio::Writer, + active: Arc, + started: Arc, + release: Arc, + } + impl Drop for CountedWriter { + fn drop(&mut self) { + self.active.fetch_sub(1, Ordering::SeqCst); + } + } + impl oio::Write for CountedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.started.add_permits(1); + self.release.acquire().await.unwrap().forget(); + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + for parallelism in [1, 4] { + for cancel in [false, true] { + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + let factory: MockWriterFactory = Arc::new({ + let (active, peak, started, release) = ( + active.clone(), + peak.clone(), + started.clone(), + release.clone(), + ); + move |_, args, inner| { + assert_eq!(args.concurrent(), 1); + let now = active.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(now, Ordering::SeqCst); + Box::new(CountedWriter { + inner, + active: active.clone(), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut a = unit(); + let mut b = unit(); + for (group, unit) in [("a", &mut a), ("b", &mut b)] { + for table in unit.logical_tables.values_mut() { + table.output.path = format!("{group}/{}", table.output.path); + } + } + let rows = || { + stream(vec![batch( + vec![Some(1025), Some(1027)], + vec![Some("a"), Some("b")], + )]) + }; + let ordinary = async { + let _permit = budget.writer(&token).await?; + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &token, + ) + .await + }; + let work = async { + tokio::join!( + export_stream_managed( + &a, + rows(), + &store, + export_limits(), + &token, + budget.clone() + ), + export_stream_managed( + &b, + rows(), + &store, + export_limits(), + &token, + budget.clone() + ), + ordinary, + ) + }; + tokio::pin!(work); + tokio::select! { + _ = started.acquire_many(parallelism as u32) => {}, + result = &mut work => panic!("completed while close paused: {result:?}"), + } + assert_eq!(budget.available().0, 0); + assert_eq!(peak.load(Ordering::SeqCst), parallelism); + if cancel { + token.cancel(); + } + assert!(futures::poll!(&mut work).is_pending()); + release.add_permits(10); + let (a, b, ordinary) = work.await; + assert_eq!(a.is_err(), cancel); + assert_eq!(b.is_err(), cancel); + assert_eq!(ordinary.is_err(), cancel); + assert_eq!(active.load(Ordering::SeqCst), 0); + assert!(peak.load(Ordering::SeqCst) <= parallelism); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs new file mode 100644 index 00000000000..c19ee4824fd --- /dev/null +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -0,0 +1,245 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow::record_batch::RecordBatch; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt}; +use object_store::ObjectStore; +use snafu::{OptionExt, ResultExt, ensure}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::statement::export_logical_tables::{ + ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled, + map_writer_error, +}; + +const PAYLOAD_BYTES: usize = 64 * 1024 * 1024; + +/// Request-wide downstream payload and writer admission, separate from query memory. +pub(crate) struct ExportWriteBudget { + writers: Arc, + bytes: Arc, +} + +impl ExportWriteBudget { + pub(crate) fn new(parallelism: usize) -> Arc { + Arc::new(Self { + writers: Arc::new(Semaphore::new(parallelism.max(1))), + bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)), + }) + } + + #[cfg(test)] + pub(crate) fn available(&self) -> (usize, usize) { + ( + self.writers.available_permits(), + self.bytes.available_permits(), + ) + } + + pub(crate) async fn writer(&self, token: &CancellationToken) -> Result { + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.writers.clone().acquire_owned() => permit.map_err(|_| error::UnexpectedSnafu { violated: "writer budget closed" }.build()), + } + } + + pub(crate) async fn reserve( + &self, + size: usize, + token: &CancellationToken, + ) -> Result { + ensure!( + size <= PAYLOAD_BYTES, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.bytes.clone().acquire_many_owned(size as u32) => permit.map_err(|_| error::UnexpectedSnafu { violated: "payload budget closed" }.build()), + } + } + + /// Includes full retained input allocations, conversion capacity growth and + /// array metadata. Shared backing is conservatively charged for each payload. + pub(crate) fn conversion_budget( + batch: &RecordBatch, + requested: usize, + ) -> Result<(usize, usize)> { + let backing = batch.get_array_memory_size(); + let overhead = batch.num_columns().saturating_mul(1024); + let available = PAYLOAD_BYTES + .saturating_sub(backing) + .saturating_sub(overhead) + / 2; + let conversion = requested.min(available); + ensure!( + conversion > 0, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + Ok((conversion, backing.saturating_add(overhead))) + } +} + +pub(crate) struct Payload { + pub(crate) batch: RecordBatch, + pub(crate) permit: OwnedSemaphorePermit, +} + +pub(crate) struct TableWriters { + current: Option<(u32, mpsc::Sender)>, + tasks: FuturesUnordered>>, + budget: Arc, +} + +impl TableWriters { + pub(crate) fn new(budget: Arc) -> Self { + Self { + current: None, + tasks: FuturesUnordered::new(), + budget, + } + } + + pub(crate) fn table_id(&self) -> Option { + self.current.as_ref().map(|(id, _)| *id) + } + + pub(crate) fn close_input(&mut self) { + self.current = None; + } + + pub(crate) async fn open( + &mut self, + id: u32, + table: &LogicalTableProjection, + store: &ObjectStore, + limits: LogicalTableExportLimits, + token: &CancellationToken, + ) -> Result<()> { + // EOF must precede acquiring the next slot, including when P is one. + self.close_input(); + while let Some(Some(result)) = self.tasks.next().now_or_never() { + result.context(error::JoinTaskSnafu)??; + } + let permit = self.budget.writer(token).await?; + let writer = ActiveWriter::open(table, store, limits).await?; + let (sender, receiver) = mpsc::channel(2); + self.current = Some((id, sender)); + let token = token.clone(); + self.tasks.push(common_runtime::spawn_global(async move { + let _permit = permit; + let guard = token.clone().drop_guard(); + let result = run_writer(writer, receiver, &token).await; + if result.is_ok() { + guard.disarm(); + } + result + })); + Ok(()) + } + + pub(crate) async fn send(&self, payload: Payload, token: &CancellationToken) -> Result<()> { + let (_, sender) = self.current.as_ref().context(error::UnexpectedSnafu { + violated: "missing logical writer", + })?; + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + result = sender.send(payload) => result.map_err(|_| error::LogicalTableExportCancelledSnafu.build()), + } + } + + pub(crate) async fn drain( + &mut self, + result: Result<()>, + token: &CancellationToken, + ) -> Result<()> { + self.close_input(); + let mut first_error = result.err(); + if first_error.is_some() { + token.cancel(); + } + while let Some(result) = self.tasks.next().await { + if let Err(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) { + retain_error(&mut first_error, err); + token.cancel(); + } + } + match first_error { + Some(err) => Err(err), + None => Ok(()), + } + } +} + +pub(crate) fn retain_error(first: &mut Option, error: error::Error) { + if first.as_ref().is_none_or(|err| { + matches!( + err, + error::Error::LogicalTableExportCancelled { .. } + | error::Error::DatabaseExportCancelled { .. } + ) + }) { + *first = Some(error); + } +} + +async fn run_writer( + mut writer: ActiveWriter, + mut receiver: mpsc::Receiver, + token: &CancellationToken, +) -> Result<()> { + let result = async { + loop { + let payload = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + payload = receiver.recv() => payload, + }; + let Some(Payload { batch, permit }) = payload else { + break; + }; + let result = writer.writer.write(batch, Some(token)).await; + drop(permit); + result.map_err(|error| map_writer_error(error, &writer.path))?; + } + writer + .writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, &writer.path))?; + check_cancelled(token) + } + .await; + if result.is_err() { + token.cancel(); + receiver.close(); + drop(receiver); + if let Err(error) = writer.writer.abort().await { + common_telemetry::warn!(error; "Failed to abort Metric export file"); + } + } + result +}