From 5f91fa9589bc7ee06459a97b06c3145d0653d39e Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 7 Sep 2026 19:29:45 +0800 Subject: [PATCH] refactor(flow): delegate durable batching execution to extensions Signed-off-by: discord9 --- src/flow/src/batching_mode.rs | 10 +- .../src/batching_mode/batching_execution.rs | 46 ++ src/flow/src/batching_mode/checkpoint.rs | 15 +- src/flow/src/batching_mode/engine.rs | 428 +++++++++-------- src/flow/src/batching_mode/frontend_client.rs | 4 +- src/flow/src/batching_mode/persistence.rs | 141 ------ src/flow/src/batching_mode/state.rs | 78 +-- src/flow/src/batching_mode/task.rs | 255 ++++------ src/flow/src/batching_mode/task/ckpt.rs | 18 - src/flow/src/batching_mode/task/inc.rs | 2 +- src/flow/src/batching_mode/task/test.rs | 448 ++++++++---------- src/flow/src/batching_mode/utils.rs | 4 +- src/flow/src/lib.rs | 15 +- src/flow/src/server.rs | 7 +- 14 files changed, 608 insertions(+), 863 deletions(-) create mode 100644 src/flow/src/batching_mode/batching_execution.rs delete mode 100644 src/flow/src/batching_mode/persistence.rs diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index 60e7969156..bd884f8fa0 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -28,15 +28,15 @@ pub enum IncrementalMode { SequenceRange, } -mod checkpoint; +pub(crate) mod batching_execution; +pub(crate) mod checkpoint; pub(crate) mod engine; mod eval_schedule; pub(crate) mod frontend_client; -pub(crate) mod persistence; -mod state; +pub(crate) mod state; mod table_creator; -mod task; -mod time_window; +pub(crate) mod task; +pub(crate) mod time_window; pub(crate) mod utils; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/src/flow/src/batching_mode/batching_execution.rs b/src/flow/src/batching_mode/batching_execution.rs new file mode 100644 index 0000000000..944101935a --- /dev/null +++ b/src/flow/src/batching_mode/batching_execution.rs @@ -0,0 +1,46 @@ +// 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. + +//! Optional execution collaborator for batching tasks. + +use std::sync::Arc; + +use query::QueryEngineRef; +use table::TableRef; + +use crate::Result; +use crate::batching_mode::frontend_client::FrontendClient; +use crate::batching_mode::task::{BatchingTask, ExecuteOnceOutcome}; + +#[async_trait::async_trait] +pub trait BatchingExecution: Send + Sync + 'static { + async fn execute_once( + &self, + task: &BatchingTask, + engine: &QueryEngineRef, + frontend: &Arc, + max_window_cnt: Option, + ) -> ExecuteOnceOutcome; +} + +#[async_trait::async_trait] +pub trait BatchingExecutionFactory: Send + Sync + 'static { + async fn create( + &self, + task: &BatchingTask, + sink: TableRef, + engine: &QueryEngineRef, + frontend: &Arc, + ) -> Result>>; +} diff --git a/src/flow/src/batching_mode/checkpoint.rs b/src/flow/src/batching_mode/checkpoint.rs index f650f4cd5a..b7a0d9f9ea 100644 --- a/src/flow/src/batching_mode/checkpoint.rs +++ b/src/flow/src/batching_mode/checkpoint.rs @@ -87,13 +87,6 @@ pub(super) enum FlowCheckpointDecision { participating_regions: usize, watermarks: usize, }, - /// The full repair requested by persistence completed and proved a new - /// checkpoint. This is distinct from an ordinary full snapshot because it - /// is the only decision that clears the repair request. - CompletedFullRepair { - participating_regions: usize, - watermarks: usize, - }, /// FullSnapshot stayed in full snapshot mode because a scoped base repair /// found additional dirty windows that may be concurrent with the returned /// high watermark. These windows must be repaired under the fixed high @@ -121,7 +114,7 @@ impl FlowCheckpointDecision { checkpoint_mode_label(CheckpointMode::FullSnapshot) } Self::AdvancedIncremental { .. } => checkpoint_mode_label(CheckpointMode::Incremental), - Self::CompletedFullRepair { .. } | Self::ContinuedFencedRepair { .. } => { + Self::ContinuedFencedRepair { .. } => { // Fenced repair and completion of a requested full repair are // FullSnapshot sub-states, not third top-level modes. checkpoint_mode_label(CheckpointMode::FullSnapshot) @@ -134,9 +127,9 @@ impl FlowCheckpointDecision { pub(super) fn decision_label(self) -> &'static str { match self { - Self::AdvancedFromFullSnapshot { .. } - | Self::AdvancedIncremental { .. } - | Self::CompletedFullRepair { .. } => CHECKPOINT_DECISION_ADVANCE, + Self::AdvancedFromFullSnapshot { .. } | Self::AdvancedIncremental { .. } => { + CHECKPOINT_DECISION_ADVANCE + } Self::ContinuedFencedRepair { .. } => CHECKPOINT_DECISION_CONTINUE_REPAIR, Self::FallbackToFullSnapshot { .. } => CHECKPOINT_DECISION_FALLBACK, } diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index 4c26b6c4b3..3f3c092c7b 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -49,7 +49,6 @@ use tokio::sync::{RwLock, oneshot}; use crate::batching_mode::BatchingModeOptions; use crate::batching_mode::eval_schedule::EvalSchedule; use crate::batching_mode::frontend_client::FrontendClient; -use crate::batching_mode::persistence::{FrontendBatchingQueryExecutor, PersistenceContext}; use crate::batching_mode::state::DirtyTimeWindows; use crate::batching_mode::task::{BatchingTask, TaskArgs}; use crate::batching_mode::time_window::{TimeWindowExpr, find_time_window_expr}; @@ -76,7 +75,7 @@ pub struct BatchingEngine { /// Batching mode options for control how batching mode query works /// pub(crate) batch_opts: Arc, - persistence_factory: Option, + execution_factory: Option>, } #[derive(Default)] @@ -133,7 +132,7 @@ impl BatchingEngine { catalog_manager: CatalogManagerRef, batch_opts: BatchingModeOptions, ) -> Self { - Self::new_with_persistence( + Self::new_with_execution( frontend_client, query_engine, flow_metadata_manager, @@ -144,14 +143,14 @@ impl BatchingEngine { ) } - pub fn new_with_persistence( + pub fn new_with_execution( frontend_client: Arc, query_engine: QueryEngineRef, flow_metadata_manager: FlowMetadataManagerRef, table_meta: TableMetadataManagerRef, catalog_manager: CatalogManagerRef, batch_opts: BatchingModeOptions, - persistence_factory: Option, + execution_factory: Option>, ) -> Self { Self { runtime: Default::default(), @@ -161,7 +160,7 @@ impl BatchingEngine { catalog_manager, query_engine, batch_opts: Arc::new(batch_opts), - persistence_factory, + execution_factory, } } @@ -745,10 +744,6 @@ impl BatchingEngine { exact_sequence_range_required, )?; - let task_inner = task.clone(); - let engine = self.query_engine.clone(); - let frontend = self.frontend_client.clone(); - if task.config.exact_sequence_range_required { ensure!( task.sequence_range_capable().await?, @@ -760,64 +755,19 @@ impl BatchingEngine { ); } - // Create the sink before configuring persistence. A persistence-backed sink may - // contain ordinary metadata columns supplied by `begin_attempt`, so strict plan/schema - // validation is deferred to execution only when persistence is actually created. Flows - // without a created collaborator keep the existing creation-time validation. + let engine = self.query_engine.clone(); + let frontend = self.frontend_client.clone(); let table = task.check_or_create_sink_table(&engine, &frontend).await?; - - let persistence = if let Some(factory) = &self.persistence_factory { - let table_info = table.table_info(); - let meta = &table_info.meta; - let effective_mode = if task.config.exact_sequence_range_required - || (task.config.batch_opts.experimental_enable_incremental_read - && task - .sequence_range_capable() - .await - .is_ok_and(|capable| capable)) - { - crate::IncrementalMode::SequenceRange - } else { - crate::IncrementalMode::MemtableOnly - }; - let context = PersistenceContext { - flow_id, - sink: crate::batching_mode::persistence::SinkLayout { - table_id: table_info.table_id(), - table_name: task.config.sink_table_name.clone(), - engine: meta.engine.clone(), - append: Self::table_options_enable_append_mode(&meta.options.extra_options), - merge_mode: if Self::table_options_enable_merge_mode_last_non_null( - &meta.options.extra_options, - ) { - Some("last_non_null".to_string()) - } else { - None - }, - columns: meta - .schema - .column_schemas() - .iter() - .map(|column| crate::BatchingMetadataColumn { - name: column.name.clone(), - data_type: column.data_type.clone(), - nullable: column.is_nullable(), - }) - .collect(), - ordered_primary_key_indices: meta.primary_key_indices.clone(), - time_index: meta.schema.timestamp_index(), - }, - executor: Arc::new(FrontendBatchingQueryExecutor::new(frontend.clone())), - incremental_mode: effective_mode, - }; - factory.create(context).await? + let execution = if let Some(factory) = &self.execution_factory { + factory.create(&task, table, &engine, &frontend).await? } else { None }; - if persistence.is_none() { + if execution.is_none() { task.validate_sink_table_schema(&engine).await?; } - task.set_persistence(persistence).await?; + let task = task.with_execution(execution); + let task_inner = task.clone(); let (start_tx, start_rx) = oneshot::channel(); @@ -1151,7 +1101,7 @@ impl FlowEngine for BatchingEngine { mod tests { use api::v1::flow::{DirtyWindowRequest, TimeRange}; use catalog::RegisterTableRequest; - use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; + use catalog::memory::MemoryCatalogManager; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::key::table_route::TableRouteValue; @@ -1163,11 +1113,10 @@ mod tests { use datatypes::vectors::{TimestampMillisecondVector, UInt32Vector, VectorRef}; use query::options::QueryOptions; use session::context::QueryContext; + use tokio::sync::Notify; use super::*; - use crate::batching_mode::persistence::{ - BatchingAttempt, BatchingPersistence, Factory, FactoryPlugin, RestoreOutcome, - }; + use crate::ExecuteOnceOutcome; use crate::test_utils::create_test_query_engine; struct DropNotify(Option>); @@ -1180,50 +1129,13 @@ mod tests { } } - struct TestPersistenceFactory { - create_persistence: bool, - } - - struct TestPersistence; - - #[async_trait::async_trait] - impl BatchingPersistence for TestPersistence { - async fn restore(&self) -> crate::Result { - Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::new())) - } - - async fn begin_attempt(&self) -> crate::Result { - Ok(BatchingAttempt::default()) - } - - async fn persist( - &self, - _attempt: BatchingAttempt, - _validated_checkpoints: BTreeMap, - ) -> crate::Result<()> { - Ok(()) - } - } - - #[async_trait::async_trait] - impl Factory for TestPersistenceFactory { - async fn create( - &self, - _context: PersistenceContext, - ) -> crate::Result>> { - Ok(self - .create_persistence - .then_some(Arc::new(TestPersistence) as Arc)) - } - } - async fn new_test_engine() -> BatchingEngine { let kv_backend = Arc::new(MemoryKvBackend::new()); let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone())); table_meta.init().await.unwrap(); let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend)); - let catalog_manager = new_memory_catalog_manager().unwrap(); let query_engine = create_test_query_engine(); + let catalog_manager = query_engine.engine_state().catalog_manager().clone(); let (frontend_client, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default()); @@ -1237,8 +1149,8 @@ mod tests { ) } - async fn new_test_engine_with_persistence( - persistence_factory: Option, + async fn new_test_engine_with_execution( + execution_factory: Option>, ) -> BatchingEngine { let kv_backend = Arc::new(MemoryKvBackend::new()); let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone())); @@ -1249,14 +1161,14 @@ mod tests { let (frontend_client, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default()); - let engine = BatchingEngine::new_with_persistence( + let engine = BatchingEngine::new_with_execution( Arc::new(frontend_client), query_engine, flow_meta, table_meta, catalog_manager, BatchingModeOptions::default(), - persistence_factory, + execution_factory, ); engine .table_meta @@ -1270,8 +1182,8 @@ mod tests { engine } - fn register_sink_with_schema(engine: &BatchingEngine, name: &str, extended: bool) { - let mut columns = vec![ + fn register_sink_with_schema(engine: &BatchingEngine, name: &str) { + let schema = Arc::new(Schema::new(vec![ ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), false), ColumnSchema::new( "ts", @@ -1279,35 +1191,53 @@ mod tests { false, ) .with_time_index(true), - ]; - let mut vectors: Vec = vec![ - Arc::new(UInt32Vector::from_slice([1_u32])), - Arc::new(TimestampMillisecondVector::from_slice([0_i64])), - ]; - if extended { - columns.push(ColumnSchema::new( - "checkpoint", - ConcreteDataType::uint32_datatype(), - false, - )); - vectors.push(Arc::new(UInt32Vector::from_slice([1_u32]))); - } - let schema = Arc::new(Schema::new(columns)); - let recordbatch = RecordBatch::new(schema, vectors).unwrap(); - let table = table::test_util::MemTable::table(name, recordbatch); - let request = RegisterTableRequest { - catalog: "greptime".to_string(), - schema: "public".to_string(), - table_name: name.to_string(), - table_id: 9000, - table, - }; + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(UInt32Vector::from_slice([1_u32])) as VectorRef, + Arc::new(TimestampMillisecondVector::from_slice([0_i64])) as VectorRef, + ], + ) + .unwrap(); engine .catalog_manager .as_any() .downcast_ref::() .unwrap() - .register_table_sync(request) + .register_table_sync(RegisterTableRequest { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table_name: name.to_string(), + table_id: 9000, + table: table::test_util::MemTable::table(name, recordbatch), + }) + .unwrap(); + } + + fn register_number_only_sink(engine: &BatchingEngine, name: &str) { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::uint32_datatype(), + false, + )])); + let recordbatch = RecordBatch::new( + schema, + vec![Arc::new(UInt32Vector::from_slice([1_u32])) as VectorRef], + ) + .unwrap(); + engine + .catalog_manager + .as_any() + .downcast_ref::() + .unwrap() + .register_table_sync(RegisterTableRequest { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table_name: name.to_string(), + table_id: 9001, + table: table::test_util::MemTable::table(name, recordbatch), + }) .unwrap(); } @@ -1332,61 +1262,192 @@ mod tests { } } - #[tokio::test] - async fn test_installed_persistence_factory_none_still_validates_sink_schema() { - let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new( - TestPersistenceFactory { - create_persistence: false, - }, - )))) - .await; - register_sink_with_schema(&engine, "sink_factory_none", true); + struct TestExecution { + manual_calls: std::sync::atomic::AtomicUsize, + } - let result = engine - .create_flow_inner(flow_create_args(1, "sink_factory_none")) - .await; + #[async_trait::async_trait] + impl crate::BatchingExecution for TestExecution { + async fn execute_once( + &self, + task: &BatchingTask, + _engine: &QueryEngineRef, + _frontend: &Arc, + _max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + if task + .state + .read() + .unwrap() + .query_ctx + .extension(query::options::FLOW_SCHEDULED_TIME_MILLIS) + .is_none() + { + self.manual_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + ExecuteOnceOutcome { + new_query: None, + result: Ok(None), + } + } + } - assert!( - result.is_err(), - "ordinary sink validation must reject mismatch" - ); - assert!(!engine.flow_exist_inner(1).await); + struct TestExecutionFactory { + entered: Option>, + release: Option>, + result: std::sync::Mutex>>>>, + } + + #[async_trait::async_trait] + impl crate::BatchingExecutionFactory for TestExecutionFactory { + async fn create( + &self, + _task: &BatchingTask, + _sink: table::TableRef, + _engine: &QueryEngineRef, + _frontend: &Arc, + ) -> crate::Result>> { + if let Some(entered) = &self.entered { + entered.notify_one(); + } + if let Some(release) = &self.release { + release.notified().await; + } + self.result + .lock() + .unwrap() + .take() + .expect("execution factory should only be called once") + } } #[tokio::test] - async fn test_matching_persistence_factory_allows_extended_sink_schema() { - let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new( - TestPersistenceFactory { - create_persistence: true, - }, - )))) - .await; - register_sink_with_schema(&engine, "sink_factory_some", true); + async fn test_execution_factory_finishes_before_task_publication() { + const GATE_TIMEOUT: Duration = Duration::from_secs(1); - let result = engine - .create_flow_inner(flow_create_args(2, "sink_factory_some")) + let execution = Arc::new(TestExecution { + manual_calls: Default::default(), + }); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let factory = Arc::new(TestExecutionFactory { + entered: Some(entered.clone()), + release: Some(release.clone()), + result: std::sync::Mutex::new(Some(Ok(Some(execution.clone())))), + }); + let engine = Arc::new(new_test_engine_with_execution(Some(factory)).await); + register_sink_with_schema(&engine, "factory_sink"); + + let entered_wait = entered.notified(); + let mut args = flow_create_args(6, "factory_sink"); + args.eval_interval = Some(86_400); + let mut create = tokio::spawn({ + let engine = engine.clone(); + async move { engine.create_flow_inner(args).await } + }); + if tokio::time::timeout(GATE_TIMEOUT, entered_wait) + .await + .is_err() + { + release.notify_one(); + if tokio::time::timeout(GATE_TIMEOUT, &mut create) + .await + .is_err() + { + create.abort(); + } + panic!("execution factory should be entered"); + } + let unpublished = { + let runtime = engine.runtime.read().await; + !runtime.tasks.contains_key(&6) && !runtime.shutdown_txs.contains_key(&6) + }; + + release.notify_one(); + assert!(unpublished); + let created = tokio::time::timeout(GATE_TIMEOUT, &mut create) + .await + .expect("flow creation should finish after factory release"); + created.unwrap().unwrap(); + assert!(engine.flow_exist_inner(6).await); + let calls_before_flush = execution + .manual_calls + .load(std::sync::atomic::Ordering::SeqCst); + assert_eq!(engine.flush_flow_inner(6).await.unwrap(), 0); + assert_eq!( + execution + .manual_calls + .load(std::sync::atomic::Ordering::SeqCst), + calls_before_flush + 1 + ); + engine.remove_flow_inner(6).await.unwrap(); + } + + #[tokio::test] + async fn test_execution_factory_result_controls_sink_validation_and_publication() { + let no_factory_engine = new_test_engine_with_execution(None).await; + register_number_only_sink(&no_factory_engine, "no_factory_sink"); + assert!( + no_factory_engine + .create_flow_inner(flow_create_args(7, "no_factory_sink")) + .await + .is_err() + ); + assert!(!no_factory_engine.flow_exist_inner(7).await); + + let declined_engine = + new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory { + entered: None, + release: None, + result: std::sync::Mutex::new(Some(Ok(None))), + }))) + .await; + register_number_only_sink(&declined_engine, "declined_factory_sink"); + assert!( + declined_engine + .create_flow_inner(flow_create_args(8, "declined_factory_sink")) + .await + .is_err() + ); + assert!(!declined_engine.flow_exist_inner(8).await); + + let accepted_engine = + new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory { + entered: None, + release: None, + result: std::sync::Mutex::new(Some(Ok(Some(Arc::new(TestExecution { + manual_calls: Default::default(), + }))))), + }))) + .await; + register_number_only_sink(&accepted_engine, "accepted_factory_sink"); + accepted_engine + .create_flow_inner(flow_create_args(9, "accepted_factory_sink")) .await .unwrap(); + assert!(accepted_engine.flow_exist_inner(9).await); + accepted_engine.remove_flow_inner(9).await.unwrap(); - assert_eq!(Some(2), result); - assert!(engine.flow_exist_inner(2).await); - engine.remove_flow_inner(2).await.unwrap(); - } - - #[tokio::test] - async fn test_no_persistence_factory_still_validates_sink_schema() { - let engine = new_test_engine_with_persistence(None).await; - register_sink_with_schema(&engine, "sink_no_factory", true); - - let result = engine - .create_flow_inner(flow_create_args(3, "sink_no_factory")) - .await; - + let error_engine = new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory { + entered: None, + release: None, + result: std::sync::Mutex::new(Some( + UnexpectedSnafu { + reason: "test execution factory failure".to_string(), + } + .fail(), + )), + }))) + .await; + register_sink_with_schema(&error_engine, "error_factory_sink"); assert!( - result.is_err(), - "ordinary sink validation must reject mismatch" + error_engine + .create_flow_inner(flow_create_args(10, "error_factory_sink")) + .await + .is_err() ); - assert!(!engine.flow_exist_inner(3).await); + assert!(!error_engine.flow_exist_inner(10).await); } #[tokio::test] @@ -1429,13 +1490,8 @@ mod tests { #[tokio::test] async fn test_forged_query_context_does_not_enable_exact_sequence_range() { - let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new( - TestPersistenceFactory { - create_persistence: true, - }, - )))) - .await; - register_sink_with_schema(&engine, "forged_query_context", true); + let engine = new_test_engine_with_execution(None).await; + register_sink_with_schema(&engine, "forged_query_context"); let mut args = flow_create_args(4, "forged_query_context"); let mut query_ctx = QueryContext::arc().as_ref().clone(); query_ctx.set_extension("__old_forged_required_extension", "true"); @@ -1449,7 +1505,7 @@ mod tests { #[tokio::test] async fn test_exact_sequence_range_capability_is_checked_before_task_startup() { - let engine = new_test_engine_with_persistence(None).await; + let engine = new_test_engine_with_execution(None).await; let mut args = flow_create_args(5, "exact_requires_capability"); args.flow_options.insert( FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY.to_string(), diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 10ce3b6820..b81223a1ea 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -463,7 +463,7 @@ impl FrontendClient { /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. - pub(crate) async fn query_with_terminal_metrics( + pub async fn query_with_terminal_metrics( &self, catalog: &str, schema: &str, @@ -666,7 +666,7 @@ fn wrap_standalone_output_with_terminal_metrics( /// Describe a peer of frontend #[derive(Debug, Default, Clone)] -pub(crate) enum PeerDesc { +pub enum PeerDesc { /// The query failed before a frontend peer was selected. #[default] Unknown, diff --git a/src/flow/src/batching_mode/persistence.rs b/src/flow/src/batching_mode/persistence.rs deleted file mode 100644 index 1241e6e324..0000000000 --- a/src/flow/src/batching_mode/persistence.rs +++ /dev/null @@ -1,141 +0,0 @@ -// 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. - -//! Optional collaborator boundary for batching state. -//! -//! The batching engine only deals in the small, typed values described here. It -//! does not interpret or persist any of them. - -use std::collections::{BTreeMap, HashMap}; -use std::sync::Arc; - -use common_error::ext::BoxedError; -use common_query::Output; -use datafusion_common::ScalarValue; -use datatypes::prelude::ConcreteDataType; -use store_api::storage::TableId; - -use crate::Result; -use crate::batching_mode::IncrementalMode; - -/// SQL execution exposed to a persistence implementation for reading and -/// writing internal state. The batching engine does not interpret this SQL; -/// the persistence implementation owns the SQL and its schema. -#[async_trait::async_trait] -pub trait BatchingQueryExecutor: Send + Sync + 'static { - async fn execute_sql(&self, catalog: &str, schema: &str, sql: &str) -> Result; -} - -/// Describes the validated sink table available to a persistence collaborator. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct SinkLayout { - pub table_id: TableId, - pub table_name: [String; 3], - pub engine: String, - pub append: bool, - pub merge_mode: Option, - pub columns: Vec, - pub ordered_primary_key_indices: Vec, - pub time_index: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct BatchingMetadataColumn { - pub name: String, - pub data_type: ConcreteDataType, - pub nullable: bool, -} - -/// Result of restoring a task's opaque state. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RestoreOutcome { - TrustedCheckpoint(BTreeMap), - FullRepair, -} - -/// One serialized execution attempt. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct BatchingAttempt { - pub ordinary_values: BTreeMap, -} - -#[async_trait::async_trait] -pub trait BatchingPersistence: Send + Sync + 'static { - async fn restore(&self) -> Result; - async fn begin_attempt(&self) -> Result; - async fn persist( - &self, - attempt: BatchingAttempt, - validated_checkpoints: BTreeMap, - ) -> Result<()>; -} - -/// Context supplied to the typed persistence factory for one batching flow. -#[derive(Clone)] -pub struct PersistenceContext { - pub flow_id: crate::FlowId, - pub incremental_mode: IncrementalMode, - pub sink: SinkLayout, - pub executor: Arc, -} - -/// Typed factory implementation used by the batching engine. -#[async_trait::async_trait] -pub trait Factory: Send + Sync + 'static { - async fn create( - &self, - context: PersistenceContext, - ) -> Result>>; -} - -/// Typed plugin wrapper around a persistence factory. -#[derive(Clone)] -pub struct FactoryPlugin(pub Arc); - -impl std::ops::Deref for FactoryPlugin { - type Target = dyn Factory; - - fn deref(&self) -> &Self::Target { - self.0.as_ref() - } -} - -/// Query executor backed by the existing flownode frontend client. -pub(crate) struct FrontendBatchingQueryExecutor { - client: Arc, -} - -impl FrontendBatchingQueryExecutor { - pub(crate) fn new(client: Arc) -> Self { - Self { client } - } -} - -#[async_trait::async_trait] -impl BatchingQueryExecutor for FrontendBatchingQueryExecutor { - async fn execute_sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { - let mut peer = None; - let request = api::v1::QueryRequest { - query: Some(api::v1::query_request::Query::Sql(sql.to_string())), - }; - self.client - .query_with_terminal_metrics(catalog, schema, request, &[], &HashMap::new(), &mut peer) - .await - .map(|output| output.into_output()) - .map_err(|err| crate::Error::External { - source: BoxedError::new(err), - location: snafu::location!(), - }) - } -} diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index 84d4d7b022..2d339a4041 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -35,19 +35,6 @@ use crate::metrics::{ }; use crate::{Error, FlowId}; -#[derive(Debug, Clone)] -pub(crate) struct TaskStateCheckpointSnapshot { - pub(crate) checkpoint_mode: CheckpointMode, - pub(crate) checkpoints: BTreeMap, - pub(crate) pending_fenced_repair: Option, - pub(crate) full_repair_required: bool, - pub(crate) dirty_time_windows: DirtyTimeWindows, - pub(crate) last_update_time: Instant, - pub(crate) last_query_duration: Duration, - pub(crate) last_exec_time_millis: Option, - pub(crate) exec_state: ExecState, -} - pub struct TaskState { /// Query context pub(crate) query_ctx: QueryContextRef, @@ -64,7 +51,6 @@ pub struct TaskState { pub(crate) dirty_time_windows: DirtyTimeWindows, checkpoint_mode: CheckpointMode, pending_fenced_repair: Option, - full_repair_required: bool, /// Region id -> last consumed watermark sequence. Incremental scans use /// this as the next lower sequence bound for each source region. checkpoints: BTreeMap, @@ -97,7 +83,6 @@ impl TaskState { dirty_time_windows, checkpoint_mode: CheckpointMode::FullSnapshot, pending_fenced_repair: None, - full_repair_required: false, checkpoints: Default::default(), incremental_disabled: false, exec_state: ExecState::Idle, @@ -140,6 +125,11 @@ impl TaskState { self.start_time_millis } + /// Pending dirty work, without permitting mutation of the task state. + pub fn dirty_time_windows(&self) -> &DirtyTimeWindows { + &self.dirty_time_windows + } + pub fn checkpoint_mode(&self) -> CheckpointMode { self.checkpoint_mode } @@ -148,64 +138,6 @@ impl TaskState { &self.checkpoints } - pub(crate) fn checkpoint_snapshot(&self) -> TaskStateCheckpointSnapshot { - TaskStateCheckpointSnapshot { - checkpoint_mode: self.checkpoint_mode, - checkpoints: self.checkpoints.clone(), - pending_fenced_repair: self.pending_fenced_repair.clone(), - full_repair_required: self.full_repair_required, - dirty_time_windows: self.dirty_time_windows.clone(), - last_update_time: self.last_update_time, - last_query_duration: self.last_query_duration, - last_exec_time_millis: self.last_exec_time_millis, - exec_state: self.exec_state.clone(), - } - } - - pub(crate) fn restore_checkpoint_snapshot(&mut self, snapshot: TaskStateCheckpointSnapshot) { - let live_dirty = self.dirty_time_windows.clone(); - self.checkpoint_mode = snapshot.checkpoint_mode; - self.checkpoints = snapshot.checkpoints; - self.pending_fenced_repair = snapshot.pending_fenced_repair; - self.full_repair_required = snapshot.full_repair_required; - self.dirty_time_windows = snapshot.dirty_time_windows; - self.dirty_time_windows.add_dirty_windows(&live_dirty); - self.last_update_time = snapshot.last_update_time; - self.last_query_duration = snapshot.last_query_duration; - self.last_exec_time_millis = snapshot.last_exec_time_millis; - self.exec_state = snapshot.exec_state; - } - - pub(crate) fn request_full_repair(&mut self) { - self.full_repair_required = true; - } - - pub(crate) fn full_repair_required(&self) -> bool { - self.full_repair_required - } - - pub(crate) fn commit_checkpoint_candidate( - &mut self, - candidate: &TaskStateCheckpointSnapshot, - successful_full_repair: bool, - ) { - let live_dirty = self.dirty_time_windows.clone(); - self.checkpoint_mode = candidate.checkpoint_mode; - self.checkpoints = candidate.checkpoints.clone(); - self.pending_fenced_repair = candidate.pending_fenced_repair.clone(); - self.full_repair_required = if successful_full_repair { - false - } else { - candidate.full_repair_required - }; - self.dirty_time_windows = candidate.dirty_time_windows.clone(); - self.dirty_time_windows.add_dirty_windows(&live_dirty); - self.last_update_time = candidate.last_update_time; - self.last_query_duration = candidate.last_query_duration; - self.last_exec_time_millis = candidate.last_exec_time_millis; - self.exec_state = candidate.exec_state.clone(); - } - /// Returns the in-progress fenced repair, if the task is repairing dirty /// windows under a frozen full-snapshot high watermark. pub fn pending_fenced_repair(&self) -> Option<&FencedRepair> { diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index a0e92f03c6..5fa80799a9 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -28,7 +28,7 @@ use datafusion::datasource::DefaultTableSource; use datafusion::sql::unparser::expr_to_sql; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::utils::quote_identifier; -use datafusion_common::{DFSchemaRef, TableReference}; +use datafusion_common::{DFSchemaRef, ScalarValue, TableReference}; use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; use datatypes::schema::Schema; use query::QueryEngineRef; @@ -49,15 +49,13 @@ use crate::batching_mode::BatchingModeOptions; use crate::batching_mode::checkpoint::checkpoint_mode_label; use crate::batching_mode::eval_schedule::{EvalSchedule, select_due_scheduled_times}; use crate::batching_mode::frontend_client::{FrontendClient, PeerDesc}; -use crate::batching_mode::persistence::RestoreOutcome; use crate::batching_mode::state::{ - CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, TaskStateCheckpointSnapshot, - to_df_literal, + CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, to_df_literal, }; use crate::batching_mode::table_creator::{QueryType, create_table_with_expr}; use crate::batching_mode::time_window::TimeWindowExpr; use crate::batching_mode::utils::{ - AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema, + AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan, }; use crate::df_optimizer::apply_df_optimizer; @@ -211,8 +209,7 @@ pub struct BatchingTask { /// window restoration for this flow. Without this, a manual flush and the /// background loop can process the same checkpoint range concurrently. execution_lock: Arc>, - persistence: - Arc>>>, + execution: Option>, } /// Arguments for creating batching task @@ -286,17 +283,16 @@ pub enum DirtyRestore { /// TODO(discord9): Full-query runs only need a dirty bool flag. Refactor /// the unscoped path to stop reusing `DirtyTimeWindows` for this signal. Unscoped(DirtyTimeWindows), - FullRepair(DirtyTimeWindows), } -struct ExecuteOnceOutcome { - new_query: Option, +pub struct ExecuteOnceOutcome { + pub new_query: Option, /// Execution result of the generated insert plan. /// /// `Ok(Some((affected_rows, elapsed)))` means a query was executed. /// `Ok(None)` means no query was generated because there was no dirty signal. /// `Err(_)` means plan generation or execution failed. - result: Result, Error>, + pub result: Result, Error>, } impl BatchingTask { @@ -353,31 +349,16 @@ impl BatchingTask { }), state: Arc::new(RwLock::new(state)), execution_lock: Arc::new(Mutex::new(())), - persistence: Arc::new(RwLock::new(None)), + execution: None, }) } - pub(crate) async fn set_persistence( - &self, - persistence: Option>, - ) -> Result<(), Error> { - if let Some(persistence) = &persistence { - let outcome = persistence.restore().await?; - let mut state = self.state.write().unwrap(); - match outcome { - RestoreOutcome::TrustedCheckpoint(checkpoints) => { - if !checkpoints.is_empty() { - state.advance_checkpoints(checkpoints.into_iter().collect()); - } - } - RestoreOutcome::FullRepair => { - state.mark_full_snapshot(); - state.request_full_repair(); - } - } - } - *self.persistence.write().unwrap() = persistence; - Ok(()) + pub(crate) fn with_execution( + mut self, + execution: Option>, + ) -> Self { + self.execution = execution; + self } pub fn last_execution_time_millis(&self) -> Option { @@ -473,33 +454,53 @@ impl BatchingTask { pub async fn validate_sink_table_schema( &self, engine: &QueryEngineRef, + ) -> Result, Error> { + self.validate_sink_table_schema_with_values(engine, &BTreeMap::new()) + .await + } + + pub async fn validate_sink_table_schema_with_values( + &self, + engine: &QueryEngineRef, + values: &BTreeMap, ) -> Result, Error> { let (table, _) = get_table_info_df_schema( self.config.catalog_manager.clone(), self.config.sink_table_name.clone(), ) .await?; - self.validate_sink_table_schema_with_table(engine, table) + self.validate_sink_table_schema_with_table_and_values(engine, table, values) .await } - async fn validate_sink_table_schema_with_table( + pub(crate) async fn validate_sink_table_schema_with_table( &self, engine: &QueryEngineRef, table: TableRef, + ) -> Result, Error> { + self.validate_sink_table_schema_with_table_and_values(engine, table, &BTreeMap::new()) + .await + } + + async fn validate_sink_table_schema_with_table_and_values( + &self, + engine: &QueryEngineRef, + table: TableRef, + values: &BTreeMap, ) -> Result, Error> { let table_meta = &table.table_info().meta; let merge_mode_last_non_null = is_merge_mode_last_non_null(&table_meta.options.extra_options); let primary_key_indices = table_meta.primary_key_indices.clone(); let query_ctx = self.state.read().unwrap().query_ctx.clone(); - gen_plan_with_matching_schema( + gen_plan_with_matching_schema_and_values( &self.config.query, query_ctx, engine.clone(), table_meta.schema.clone(), &primary_key_indices, merge_mode_last_non_null, + Some(values), ) .await .map(|_| table_meta.schema.clone()) @@ -511,15 +512,16 @@ impl BatchingTask { engine: &QueryEngineRef, max_window_cnt: Option, ) -> Result, Error> { - self.gen_insert_plan_unlocked_with_attempt(engine, max_window_cnt, None) + self.gen_insert_plan_with_values_unlocked(engine, max_window_cnt, &BTreeMap::new(), false) .await } - async fn gen_insert_plan_unlocked_with_attempt( + pub async fn gen_insert_plan_with_values_unlocked( &self, engine: &QueryEngineRef, max_window_cnt: Option, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, + values: &BTreeMap, + force_full_snapshot: bool, ) -> Result, Error> { let (table, df_schema) = get_table_info_df_schema( self.config.catalog_manager.clone(), @@ -533,13 +535,14 @@ impl BatchingTask { let primary_key_indices = table_meta.primary_key_indices.clone(); let new_query = self - .gen_query_with_time_window_with_attempt( + .gen_query_with_time_window_with_values( engine.clone(), &table.table_info().meta.schema, &primary_key_indices, merge_mode_last_non_null, max_window_cnt, - attempt, + values, + force_full_snapshot, ) .await?; @@ -621,54 +624,6 @@ impl BatchingTask { Ok(()) } - /// Executes the insert plan. Caller must reach this through the serialized path. - async fn persist_checkpoint_candidate( - &self, - snapshot: TaskStateCheckpointSnapshot, - candidate: TaskStateCheckpointSnapshot, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, - dirty_restore: DirtyRestore, - ) -> Result<(), Error> { - let persistence = self.persistence.read().unwrap().as_ref().cloned(); - let Some(persistence) = persistence else { - self.state.write().unwrap().commit_checkpoint_candidate( - &candidate, - matches!(dirty_restore, DirtyRestore::FullRepair(_)), - ); - return Ok(()); - }; - let Some(attempt) = attempt else { - self.state.write().unwrap().commit_checkpoint_candidate( - &candidate, - matches!(dirty_restore, DirtyRestore::FullRepair(_)), - ); - return Ok(()); - }; - - let successful_full_repair = matches!(dirty_restore, DirtyRestore::FullRepair(_)); - if let Err(source) = persistence - .persist(attempt.clone(), candidate.checkpoints.clone()) - .await - { - let mut state = self.state.write().unwrap(); - state.restore_checkpoint_snapshot(snapshot); - match dirty_restore { - DirtyRestore::Scoped(filter) => state.restore_scoped_windows(&filter), - DirtyRestore::Unscoped(windows) | DirtyRestore::FullRepair(windows) => { - state.dirty_time_windows.add_dirty_windows(&windows) - } - } - state.mark_full_snapshot(); - state.request_full_repair(); - return Err(source); - } - self.state - .write() - .unwrap() - .commit_checkpoint_candidate(&candidate, successful_full_repair); - Ok(()) - } - async fn execute_logical_plan_unlocked( &self, engine: &QueryEngineRef, @@ -676,7 +631,6 @@ impl BatchingTask { plan: &LogicalPlan, dirty_restore: &DirtyRestore, coverage: &QueryCoverage, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result, Error> { let flow_id = self.config.flow_id; let Some((res, elapsed)) = self @@ -690,13 +644,7 @@ impl BatchingTask { let decision = { let mut state = self.state.write().unwrap(); let reason = Self::query_failure_reason(err, coverage); - Self::apply_query_failure_to_state( - &mut state, - elapsed, - coverage, - reason, - attempt.is_some(), - ) + Self::apply_query_failure_to_state(&mut state, elapsed, coverage, reason) }; if let Some(decision) = decision { Self::record_checkpoint_decision(flow_id, decision); @@ -705,42 +653,16 @@ impl BatchingTask { let res = res?; let (affected_rows, _) = res.output.extract_rows_and_cost(); - let (decision, checkpoint_txn) = { + let decision = { let mut state = self.state.write().unwrap(); - let snapshot = state.checkpoint_snapshot(); - let repair_required = snapshot.full_repair_required; - let decision = Self::apply_query_result_to_state( - &mut state, - &res, - elapsed, - coverage, - repair_required, - ); - let eligible = matches!( - decision, - crate::batching_mode::checkpoint::FlowCheckpointDecision::AdvancedFromFullSnapshot { .. } - | crate::batching_mode::checkpoint::FlowCheckpointDecision::AdvancedIncremental { .. } - | crate::batching_mode::checkpoint::FlowCheckpointDecision::CompletedFullRepair { .. } - ); - if eligible { - let candidate = state.checkpoint_snapshot(); - state.restore_checkpoint_snapshot(snapshot.clone()); - Some((snapshot, candidate)) - } else { - None - } - .map_or((decision, None), |txn| (decision, Some(txn))) + Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage) }; - if let Some((snapshot, candidate)) = checkpoint_txn { - self.persist_checkpoint_candidate(snapshot, candidate, attempt, dirty_restore.clone()) - .await?; - } Self::record_checkpoint_decision(flow_id, decision); Ok(Some((affected_rows, elapsed))) } - async fn execute_plan_unlocked( + pub async fn execute_plan_unlocked( &self, engine: &QueryEngineRef, frontend_client: &Arc, @@ -946,10 +868,10 @@ impl BatchingTask { /// Restore dirty windows consumed by a failed query so they are retried on /// the next execution. /// - fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) { + pub fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) { match dirty_restore { DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter), - DirtyRestore::Unscoped(dirty_windows) | DirtyRestore::FullRepair(dirty_windows) => self + DirtyRestore::Unscoped(dirty_windows) => self .state .write() .unwrap() @@ -1027,7 +949,7 @@ impl BatchingTask { dirty_windows_to_restore: DirtyTimeWindows, retention_filter: Option<(&str, Timestamp, &'static str)>, coverage: QueryCoverage, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, + values: &BTreeMap, ) -> Result { let mut plan = self.restore_unscoped_dirty_windows_on_err( &dirty_windows_to_restore, @@ -1038,7 +960,7 @@ impl BatchingTask { sink_table_schema, primary_key_indices, allow_partial, - attempt.map(|a| &a.ordinary_values), + Some(values), ) .await, )?; @@ -1066,11 +988,7 @@ impl BatchingTask { Ok(PlanInfo { plan, - dirty_restore: if self.state.read().unwrap().full_repair_required() { - DirtyRestore::FullRepair(dirty_windows_to_restore) - } else { - DirtyRestore::Unscoped(dirty_windows_to_restore) - }, + dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore), coverage, }) } @@ -1087,7 +1005,7 @@ impl BatchingTask { allow_partial: bool, retention_filter: Option<(&str, Timestamp, &'static str)>, coverage: QueryCoverage, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, + values: &BTreeMap, ) -> Result, Error> { let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal(); if !is_dirty { @@ -1104,7 +1022,7 @@ impl BatchingTask { dirty_windows_to_restore, retention_filter, coverage, - attempt, + values, ) .await .map(Some) @@ -1441,23 +1359,22 @@ impl BatchingTask { frontend_client: &Arc, max_window_cnt: Option, ) -> ExecuteOnceOutcome { - let persistence = self.persistence.read().unwrap().as_ref().cloned(); - let attempt = match persistence { - Some(state) => match state.begin_attempt().await { - Ok(attempt) => Some(attempt), - Err(source) => { - return ExecuteOnceOutcome { - new_query: None, - result: Err(source), - }; - } - }, - None => None, - }; - let new_query = match self - .gen_insert_plan_unlocked_with_attempt(engine, max_window_cnt, attempt.as_ref()) + if let Some(execution) = &self.execution { + return execution + .execute_once(self, engine, frontend_client, max_window_cnt) + .await; + } + self.execute_once_default_unlocked(engine, frontend_client, max_window_cnt) .await - { + } + + async fn execute_once_default_unlocked( + &self, + engine: &QueryEngineRef, + frontend_client: &Arc, + max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + let new_query = match self.gen_insert_plan_unlocked(engine, max_window_cnt).await { Ok(new_query) => new_query, Err(err) => { return ExecuteOnceOutcome { @@ -1479,7 +1396,6 @@ impl BatchingTask { &new_query.plan, &new_query.dirty_restore, &new_query.coverage, - attempt.as_ref(), ) .await; match res { @@ -1614,25 +1530,27 @@ impl BatchingTask { allow_partial: bool, max_window_cnt: Option, ) -> Result, Error> { - self.gen_query_with_time_window_with_attempt( + self.gen_query_with_time_window_with_values( engine, sink_table_schema, primary_key_indices, allow_partial, max_window_cnt, - None, + &BTreeMap::new(), + false, ) .await } - async fn gen_query_with_time_window_with_attempt( + async fn gen_query_with_time_window_with_values( &self, engine: QueryEngineRef, sink_table_schema: &Arc, primary_key_indices: &[usize], allow_partial: bool, max_window_cnt: Option, - attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, + values: &BTreeMap, + force_full_snapshot: bool, ) -> Result, Error> { let query_ctx = self.state.read().unwrap().query_ctx.clone(); let start = SystemTime::now(); @@ -1654,15 +1572,20 @@ impl BatchingTask { .map(|expr| expr.eval(low_bound)) .transpose()?; - let full_repair = self.state.read().unwrap().full_repair_required(); - if full_repair { + if force_full_snapshot { let detached = self.state.write().unwrap().dirty_time_windows.detach(); let retention_filter = self.config.expire_after.and_then(|_| { self.config.time_window_expr.as_ref().and_then(|expr| { expr.eval(low_bound) .ok() .and_then(|(lower, _)| lower) - .map(|lower| (expr.column_name.as_str(), lower, "full repair retention")) + .map(|lower| { + ( + expr.column_name.as_str(), + lower, + "forced full snapshot retention", + ) + }) }) }); return self @@ -1675,7 +1598,7 @@ impl BatchingTask { detached, retention_filter, QueryCoverage::UnfilteredFull, - attempt, + values, ) .await .map(Some); @@ -1713,7 +1636,7 @@ impl BatchingTask { dirty_windows_to_restore, None, QueryCoverage::UnfilteredFull, - attempt, + values, ) .await?; @@ -1770,7 +1693,7 @@ impl BatchingTask { allow_partial, retention_filter, QueryCoverage::IncrementalDelta, - attempt, + values, ) .await; } @@ -1818,9 +1741,7 @@ impl BatchingTask { sink_table_schema.clone(), primary_key_indices.to_vec(), allow_partial, - attempt - .map(|a| a.ordinary_values.clone()) - .unwrap_or_default(), + values.clone(), ); let plan = self.restore_scoped_dirty_windows_on_err( diff --git a/src/flow/src/batching_mode/task/ckpt.rs b/src/flow/src/batching_mode/task/ckpt.rs index 27e867fa4c..c6b8503c32 100644 --- a/src/flow/src/batching_mode/task/ckpt.rs +++ b/src/flow/src/batching_mode/task/ckpt.rs @@ -58,7 +58,6 @@ impl BatchingTask { elapsed: Duration, coverage: &QueryCoverage, reason: FlowQueryFallbackReason, - persistence_backed: bool, ) -> Option { state.after_query_exec(elapsed, false); let checkpoint_mode = state.checkpoint_mode(); @@ -80,9 +79,6 @@ impl BatchingTask { if checkpoint_mode == CheckpointMode::Incremental { state.mark_full_snapshot(); - if persistence_backed && matches!(coverage, QueryCoverage::IncrementalDelta) { - state.request_full_repair(); - } } Some(FlowCheckpointDecision::FallbackToFullSnapshot { previous_mode: checkpoint_mode, @@ -100,7 +96,6 @@ impl BatchingTask { res: &OutputWithMetrics, elapsed: Duration, coverage: &QueryCoverage, - started_full_repair: bool, ) -> FlowCheckpointDecision { state.after_query_exec(elapsed, true); let checkpoint_mode = state.checkpoint_mode(); @@ -201,11 +196,6 @@ impl BatchingTask { previous_mode: CheckpointMode::FullSnapshot, reason: FlowQueryFallbackReason::IncrementalDisabled, } - } else if started_full_repair { - FlowCheckpointDecision::CompletedFullRepair { - participating_regions: participating_region_count, - watermarks: watermark_count, - } } else { FlowCheckpointDecision::AdvancedFromFullSnapshot { participating_regions: participating_region_count, @@ -276,14 +266,6 @@ impl BatchingTask { "Flow {flow_id} switched to incremental mode after full snapshot, participating_regions={participating_regions}, watermarks={watermarks}" ); } - FlowCheckpointDecision::CompletedFullRepair { - participating_regions, - watermarks, - } => { - info!( - "Flow {flow_id} completed full repair, participating_regions={participating_regions}, watermarks={watermarks}" - ); - } FlowCheckpointDecision::AdvancedIncremental { participating_regions, watermarks, diff --git a/src/flow/src/batching_mode/task/inc.rs b/src/flow/src/batching_mode/task/inc.rs index 1a6a24ac30..a87a14b215 100644 --- a/src/flow/src/batching_mode/task/inc.rs +++ b/src/flow/src/batching_mode/task/inc.rs @@ -73,7 +73,7 @@ impl BatchingTask { /// table that cannot be resolved, is not the mito engine, or lacks the /// option — this returns `false` so the caller keeps the historical /// `memtable_only` mode instead of upgrading. - pub(crate) async fn sequence_range_capable(&self) -> Result { + pub async fn sequence_range_capable(&self) -> Result { for name in &self.config.source_table_names { let table = match self .config diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 299bc1896f..5ac695a530 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,10 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; -use std::pin::Pin; -use std::task::{Context, Poll}; -use bytes::Bytes; use catalog::memory::MemoryCatalogManager; use catalog::{DeregisterTableRequest, RegisterTableRequest}; use client::OutputWithMetrics; @@ -25,12 +22,12 @@ use common_error::ext::BoxedError; use common_error::mock::MockError; use common_error::status_code::StatusCode; use common_query::Output; +use common_recordbatch::RecordBatch; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; -use common_recordbatch::{RecordBatch, RecordBatchStream}; use datatypes::data_type::ConcreteDataType as CDT; use datatypes::schema::ColumnSchema; use datatypes::vectors::{ - TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, + StringVector, TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, }; use pretty_assertions::assert_eq; use query::options::{ @@ -49,7 +46,6 @@ use crate::batching_mode::checkpoint::{ FlowCheckpointDecision, FlowQueryFallbackReason, }; use crate::batching_mode::eval_schedule::{FlowMissedTickPolicy, FlowScheduleConfig}; -use crate::batching_mode::persistence::{BatchingAttempt, BatchingPersistence, RestoreOutcome}; use crate::batching_mode::state::CheckpointMode; use crate::batching_mode::time_window::find_time_window_expr; use crate::test_utils::create_test_query_engine; @@ -61,6 +57,68 @@ fn incremental_batch_opts() -> Arc { }) } +struct CountingExecution { + calls: std::sync::atomic::AtomicUsize, + active: std::sync::atomic::AtomicUsize, + max_active: std::sync::atomic::AtomicUsize, +} + +#[async_trait::async_trait] +impl crate::BatchingExecution for CountingExecution { + async fn execute_once( + &self, + _task: &BatchingTask, + _engine: &QueryEngineRef, + _frontend: &Arc, + _max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let active = self + .active + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + + 1; + self.max_active + .fetch_max(active, std::sync::atomic::Ordering::SeqCst); + tokio::task::yield_now().await; + self.active + .fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + ExecuteOnceOutcome { + new_query: None, + result: Ok(None), + } + } +} + +#[tokio::test] +async fn test_execution_delegate_dispatch_is_serialized() { + let TestTaskParts { + task, query_engine, .. + } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink") + .await; + let execution = Arc::new(CountingExecution { + calls: Default::default(), + active: Default::default(), + max_active: Default::default(), + }); + let task = task.with_execution(Some(execution.clone())); + let (frontend, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default()); + let frontend = Arc::new(frontend); + + let first = task.execute_once_serialized(&query_engine, &frontend, None); + let second = task.execute_once_serialized(&query_engine, &frontend, None); + let (first, second) = tokio::join!(first, second); + assert_eq!(first.unwrap(), None); + assert_eq!(second.unwrap(), None); + assert_eq!(execution.calls.load(std::sync::atomic::Ordering::SeqCst), 2); + assert_eq!( + execution + .max_active + .load(std::sync::atomic::Ordering::SeqCst), + 1, + "the existing execution_lock must span delegate execution" + ); +} + async fn new_test_task_and_plan_with_missing_sink() -> (BatchingTask, LogicalPlan) { new_test_task_engine_and_plan_with_query( "SELECT number, ts FROM numbers_with_ts", @@ -554,6 +612,35 @@ fn register_twe_sink(query_engine: &QueryEngineRef, table_name: &str, table_id: memory_catalog.register_table_sync(request).unwrap(); } +fn register_twe_sink_with_metadata(query_engine: &QueryEngineRef, table_name: &str, table_id: u32) { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("number", CDT::uint32_datatype(), false), + ColumnSchema::new("time_window", CDT::timestamp_millisecond_datatype(), false) + .with_time_index(true), + ColumnSchema::new("metadata", CDT::string_datatype(), false), + ])); + let columns: Vec = vec![ + Arc::new(UInt32Vector::from_slice([1_u32])), + Arc::new(TimestampMillisecondVector::from_slice([0_i64])), + Arc::new(StringVector::from_slice(&["existing"])), + ]; + let recordbatch = RecordBatch::new(schema, columns).unwrap(); + let table = MemTable::table(table_name, recordbatch); + let request = RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id, + table, + }; + let catalog_manager = query_engine.engine_state().catalog_manager(); + let memory_catalog = catalog_manager + .as_any() + .downcast_ref::() + .unwrap(); + memory_catalog.register_table_sync(request).unwrap(); +} + fn register_scheduled_now_sink(query_engine: &QueryEngineRef, table_name: &str, table_id: u32) { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new("ts", CDT::timestamp_nanosecond_datatype(), false).with_time_index(true), @@ -944,7 +1031,6 @@ fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, - false, ); assert_eq!( @@ -976,7 +1062,6 @@ fn test_apply_query_result_to_state_stays_full_snapshot_when_incremental_disable &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, - false, ); // Should NOT claim advancement to incremental; should fallback with correct reason. @@ -1008,7 +1093,6 @@ fn test_apply_query_result_to_state_rejects_unproved_watermark() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, - false, ); assert_eq!( @@ -1034,7 +1118,6 @@ fn test_apply_query_result_to_state_reports_missing_watermark() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, - false, ); assert_eq!( @@ -1065,7 +1148,6 @@ fn test_apply_query_result_to_state_advances_incremental_subset() { &result, std::time::Duration::from_millis(1), &QueryCoverage::IncrementalDelta, - false, ); assert_eq!( @@ -1101,7 +1183,6 @@ fn test_scoped_base_repair_with_dirty_backlog_starts_fenced_repair_from_full_sna &result, std::time::Duration::from_millis(1), &QueryCoverage::ScopedBaseRepair, - false, ); assert_eq!( @@ -1161,7 +1242,6 @@ fn test_fenced_repair_chunk_with_pending_windows_stays_full_snapshot() { &output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]), std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - false, ); assert_eq!( @@ -1218,7 +1298,6 @@ fn test_continued_fenced_repair_uses_pending_snapshot_not_later_live_dirty() { &output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]), std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - false, ); assert_eq!( decision, @@ -1256,7 +1335,6 @@ fn test_final_fenced_repair_chunk_advances_to_high() { &output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]), std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high: high.clone() }, - false, ); assert_eq!( @@ -1331,7 +1409,6 @@ fn test_fenced_repair_chunk_watermark_mismatch_restores_pending_but_consumes_inf &output_with_region_watermarks([(1_u64, Some(11_u64)), (2_u64, Some(20_u64))]), std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - false, ); assert_eq!( @@ -1375,7 +1452,6 @@ async fn test_fenced_repair_mismatch_next_plan_is_scoped_base_repair() { &output_with_region_watermarks([(1_u64, Some(11_u64)), (2_u64, Some(20_u64))]), std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - false, ); assert_eq!( decision, @@ -1422,7 +1498,6 @@ fn test_apply_query_failure_to_state_falls_back_from_incremental() { std::time::Duration::from_millis(1), &QueryCoverage::IncrementalDelta, FlowQueryFallbackReason::IncrementalQueryFailure, - false, ); assert_eq!( @@ -1450,7 +1525,6 @@ fn test_apply_query_failure_to_state_records_full_snapshot_failure() { std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, FlowQueryFallbackReason::QueryFailure, - false, ); assert_eq!( @@ -1485,11 +1559,11 @@ fn test_query_failure_reason_distinguishes_fenced_repair_stale_fence() { let generic_err = flow_error_with_status(StatusCode::Unexpected); assert_eq!( BatchingTask::query_failure_reason(&generic_err, &QueryCoverage::ScopedBaseRepair), - FlowQueryFallbackReason::QueryFailure + FlowQueryFallbackReason::QueryFailure, ); assert_eq!( BatchingTask::query_failure_reason(&generic_err, &QueryCoverage::IncrementalDelta), - FlowQueryFallbackReason::IncrementalQueryFailure + FlowQueryFallbackReason::IncrementalQueryFailure, ); } @@ -1542,7 +1616,6 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() { std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, FlowQueryFallbackReason::SnapshotFenceExpired, - false, ); assert_eq!( decision, @@ -1605,7 +1678,6 @@ fn test_fenced_repair_transient_non_stale_failure_retries_same_high() { std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high: high.clone() }, FlowQueryFallbackReason::QueryFailure, - false, ); assert_eq!( @@ -2047,6 +2119,44 @@ async fn test_full_snapshot_seeding_applies_expire_after_retention_filter() { ); } +#[tokio::test] +async fn test_forced_full_snapshot_keeps_retention_without_dirty_windows() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window", + ) + .await; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .expire_after = Some(expire_after_for_retention_filter_test()); + let sink_schema = aggregate_time_window_sink_schema(); + + let plan = task + .gen_query_with_time_window_with_values( + query_engine, + &sink_schema, + &[], + false, + None, + &BTreeMap::new(), + true, + ) + .await + .unwrap() + .expect("forced full snapshot must not require a dirty signal"); + + assert!(matches!(plan.coverage, QueryCoverage::UnfilteredFull)); + assert!(task.state.read().unwrap().dirty_time_windows.is_empty()); + let plan_text = plan.plan.to_string(); + assert!( + plan_text.contains("Filter: ts >= TimestampMillisecond("), + "forced full snapshot retains the configured retention predicate: {plan_text}" + ); +} + #[tokio::test] async fn test_incremental_plan_does_not_add_dirty_window_filter() { let TestTaskParts { @@ -2165,7 +2275,6 @@ async fn test_successful_incremental_checkpoint_fallback_consumes_unscoped_dirty &result, std::time::Duration::from_millis(1), &plan_info.coverage, - false, ) }; assert_eq!( @@ -2423,7 +2532,6 @@ async fn test_unsafe_incremental_plan_skip_restores_dirty_without_query() { &dml_plan, &dirty_restore, &QueryCoverage::IncrementalDelta, - None, ) .await .unwrap(); @@ -2575,229 +2683,6 @@ async fn test_auto_created_sql_aggregate_sink_reaches_incremental_safe() { ); } -struct TestIncrementalPersistence; - -#[async_trait::async_trait] -impl BatchingPersistence for TestIncrementalPersistence { - async fn restore(&self) -> crate::Result { - Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::from([ - (1_u64, 10_u64), - (2_u64, 20_u64), - ]))) - } - - async fn begin_attempt(&self) -> crate::Result { - Ok(BatchingAttempt::default()) - } - - async fn persist( - &self, - _attempt: BatchingAttempt, - _validated_checkpoints: BTreeMap, - ) -> crate::Result<()> { - Ok(()) - } -} - -#[derive(Debug, PartialEq, Eq)] -struct ModeledSink { - a: u64, - b: u64, -} - -struct StatefulRepairHandler { - sink: Arc>, - query_engine: QueryEngineRef, -} - -struct RepairOutputStream { - schema: Arc, - metrics: RecordBatchMetrics, -} - -impl futures::Stream for RepairOutputStream { - type Item = common_recordbatch::error::Result; - - fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(None) - } -} - -impl RecordBatchStream for RepairOutputStream { - fn name(&self) -> &str { - "StatefulRepairHandler" - } - - fn schema(&self) -> datatypes::schema::SchemaRef { - self.schema.clone() - } - - fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { - None - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone()) - } -} - -#[async_trait::async_trait] -impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError - for StatefulRepairHandler -{ - async fn do_query( - &self, - query: api::v1::greptime_request::Request, - ctx: QueryContextRef, - ) -> std::result::Result { - let api::v1::greptime_request::Request::Query(api::v1::QueryRequest { - query: Some(api::v1::query_request::Query::InsertIntoPlan(insert)), - }) = query - else { - panic!("expected an InsertIntoPlan request"); - }; - - let incremental_mode = ctx.extension(FLOW_INCREMENTAL_MODE); - if incremental_mode.is_some() { - assert_eq!(incremental_mode, Some(FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY)); - let after_seqs = ctx - .extension(FLOW_INCREMENTAL_AFTER_SEQS) - .expect("incremental mode must carry sequence bounds"); - assert_eq!( - serde_json::from_str::>(after_seqs).unwrap(), - BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]) - ); - - // This models the boundary failure: the incremental A+B delta is - // sent to the sink, but only A is applied before the sink errors. - let mut sink = self.sink.lock().unwrap(); - assert_eq!((sink.a, sink.b), (10, 20)); - sink.a = 11; - return Err(BoxedError::new(MockError::new(StatusCode::Unexpected))); - } - - assert!( - ctx.extension(FLOW_INCREMENTAL_AFTER_SEQS).is_none(), - "full repair must not carry incremental sequence bounds" - ); - let decoder = self - .query_engine - .engine_context(ctx.clone()) - .new_plan_decoder() - .unwrap(); - let catalog = Arc::new( - catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( - self.query_engine.engine_state().catalog_manager().clone(), - ctx.clone(), - ), - ); - let logical_plan = decoder - .decode(Bytes::from(insert.logical_plan), catalog, false) - .await - .unwrap(); - assert!( - !logical_plan.to_string().contains("Filter"), - "full repair must not include a dirty-window predicate: {logical_plan}" - ); - - // The fixture's source timestamps share one window; this models the - // unfiltered A+B sink repair explicitly rather than asserting source - // rows for distinct windows. - let mut sink = self.sink.lock().unwrap(); - assert_eq!(sink.a, 11); - assert_eq!(sink.b, 20); - sink.b = 21; - let metrics = RecordBatchMetrics { - region_watermarks: vec![ - RegionWatermarkEntry { - region_id: 1, - watermark: Some(11), - }, - RegionWatermarkEntry { - region_id: 2, - watermark: Some(21), - }, - ], - ..Default::default() - }; - Ok(Output::new_with_stream(Box::pin(RepairOutputStream { - schema: aggregate_time_window_sink_schema(), - metrics, - }))) - } -} - -#[tokio::test] -async fn test_persistence_full_repair_retries_unfiltered_snapshot_after_partial_sink_failure() { - let TestTaskParts { - mut task, - query_engine, - .. - } = new_time_window_test_task_with_query( - "SELECT max(number) AS number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window", - ) - .await; - let sink_table = "incremental_failure_repair_sink"; - Arc::get_mut(&mut task.config) - .expect("test task config should be uniquely owned") - .sink_table_name[2] = sink_table.to_string(); - register_twe_sink(&query_engine, sink_table, 9103); - task.set_persistence(Some(Arc::new(TestIncrementalPersistence))) - .await - .unwrap(); - { - let mut state = task.state.write().unwrap(); - // Only A is dirty. B is intentionally absent from the dirty signal; - // the later full repair must nevertheless repair the modeled A+B sink. - state - .dirty_time_windows - .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); - } - - let sink = Arc::new(std::sync::Mutex::new(ModeledSink { a: 10, b: 20 })); - let handler: Arc = - Arc::new(StatefulRepairHandler { - sink: sink.clone(), - query_engine: query_engine.clone(), - }); - let frontend_client = Arc::new(FrontendClient::from_grpc_handler( - Arc::downgrade(&handler), - QueryOptions::default(), - )); - - let first = task - .execute_once_serialized(&query_engine, &frontend_client, Some(2)) - .await; - assert!( - first.is_err(), - "the partial incremental sink write must fail" - ); - assert_eq!(*sink.lock().unwrap(), ModeledSink { a: 11, b: 20 }); - { - let state = task.state.read().unwrap(); - assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot); - assert_eq!( - state.checkpoints(), - &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]) - ); - assert!(state.full_repair_required()); - } - - let second = task - .execute_once_serialized(&query_engine, &frontend_client, Some(2)) - .await - .expect("the unfiltered full repair should succeed"); - assert!(second.is_some()); - assert_eq!(*sink.lock().unwrap(), ModeledSink { a: 11, b: 21 }); - let state = task.state.read().unwrap(); - assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental); - assert_eq!( - state.checkpoints(), - &BTreeMap::from([(1_u64, 11_u64), (2_u64, 21_u64)]) - ); - assert!(!state.full_repair_required()); -} - #[tokio::test] async fn test_unscoped_failure_restores_consumed_dirty_signal() { assert_unscoped_failure_restore(dirty_marker(), DirtyTimeWindows::default(), 1, 0).await; @@ -2876,6 +2761,73 @@ async fn test_scoped_plan_generation_failure_restores_consumed_dirty_windows() { ); } +#[tokio::test] +async fn test_validate_sink_table_schema_with_values_accepts_typed_metadata() { + let sink_table = "metadata_validation_sink"; + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window, number", + ) + .await; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .sink_table_name[2] = sink_table.to_string(); + register_twe_sink_with_metadata(&query_engine, sink_table, 9104); + + let values = BTreeMap::from([( + "metadata".to_string(), + datafusion_common::ScalarValue::Utf8(Some("typed".to_string())), + )]); + assert!( + task.validate_sink_table_schema_with_values(&query_engine, &values) + .await + .is_ok() + ); + assert!(task.state.read().unwrap().dirty_time_windows.is_empty()); +} + +#[tokio::test] +async fn test_validate_sink_table_schema_with_values_rejects_output_mismatch_without_dirty_change() +{ + let sink_table = "metadata_validation_mismatch_sink"; + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window, number", + ) + .await; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .sink_table_name[2] = sink_table.to_string(); + register_number_only_sink(&query_engine, sink_table); + task.state + .write() + .unwrap() + .dirty_time_windows + .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); + + let values = BTreeMap::from([( + "metadata".to_string(), + datafusion_common::ScalarValue::Utf8(Some("typed".to_string())), + )]); + assert!( + task.validate_sink_table_schema_with_values(&query_engine, &values) + .await + .is_err() + ); + let state = task.state.read().unwrap(); + assert_eq!(state.dirty_time_windows.len(), 1); + assert_eq!( + state.dirty_time_windows.window_size(), + std::time::Duration::from_secs(5) + ); +} + #[tokio::test] async fn test_insert_plan_matching_failure_restores_consumed_dirty_marker() { let sink_table = "partial_sink"; diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index fd22e42f22..01beb82714 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -1345,7 +1345,7 @@ impl ColumnMatcherRewriter { return self.modify_project_exprs_with_partial(exprs); } - // Ordinary values are persistence-owned columns, not flow outputs. Remove them from the + // Ordinary values are execution-owned columns, not flow outputs. Remove them from the // effective sink sequence while deciding whether the existing auto-column rules apply. // This keeps those columns from hiding an auto-created update_at column that precedes them. let effective_sink_columns = self @@ -1428,7 +1428,7 @@ impl ColumnMatcherRewriter { &effective_sink_columns, )?; - // Put persistence-owned values back at their physical sink positions only after matching + // Put execution-owned values back at their physical sink positions only after matching // flow expressions against the effective sequence. let mut exprs = exprs; for (idx, column) in self.schema.column_schemas().iter().enumerate() { diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index 94ff463a57..4e5ee26ce5 100644 --- a/src/flow/src/lib.rs +++ b/src/flow/src/lib.rs @@ -43,16 +43,21 @@ mod test_utils; pub use adapter::flownode_impl::FlowDualEngineRef; pub use adapter::{FlowConfig, FlowStreamingEngineRef, StreamingEngine}; -pub use batching_mode::IncrementalMode; -pub use batching_mode::frontend_client::{FrontendClient, GrpcQueryHandlerWithBoxedError}; -pub use batching_mode::persistence::{ - BatchingAttempt, BatchingMetadataColumn, BatchingPersistence, BatchingQueryExecutor, Factory, - FactoryPlugin, PersistenceContext, RestoreOutcome, SinkLayout, +pub use batching_mode::batching_execution::{BatchingExecution, BatchingExecutionFactory}; +pub use batching_mode::frontend_client::{ + FrontendClient, GrpcQueryHandlerWithBoxedError, PeerDesc, }; +pub use batching_mode::task::{ + BatchingTask, DirtyRestore, ExecuteOnceOutcome, PlanInfo, QueryCoverage, TaskArgs, +}; +pub use batching_mode::time_window::{TimeWindowExpr, find_time_window_expr}; +pub use batching_mode::utils::sql_to_df_plan; +pub use batching_mode::{BatchingModeOptions, IncrementalMode}; pub(crate) use engine::{CreateFlowArgs, FlowId, TableName}; pub use error::{Error, Result}; pub use server::{ FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder, FrontendInvoker, }; +pub use transform::register_function_to_query_engine; pub use crate::adapter::FlownodeOptions; diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index fb15a687aa..4a2a2051ea 100644 --- a/src/flow/src/server.rs +++ b/src/flow/src/server.rs @@ -57,7 +57,6 @@ use tonic::{Request, Response, Status}; use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef}; use crate::adapter::{FlowStreamingEngineRef, create_worker}; use crate::batching_mode::engine::BatchingEngine; -use crate::batching_mode::persistence::FactoryPlugin; use crate::error::{ CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err, @@ -66,7 +65,7 @@ use crate::heartbeat::HeartbeatTask; use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS}; use crate::transform::register_function_to_query_engine; use crate::utils::{SizeReportSender, StateReportHandler}; -use crate::{Error, FlownodeOptions, FrontendClient, StreamingEngine}; +use crate::{BatchingExecutionFactory, Error, FlownodeOptions, FrontendClient, StreamingEngine}; pub const FLOW_NODE_SERVER_NAME: &str = "FLOW_NODE_SERVER"; /// wrapping flow node manager to avoid orphan rule with Arc<...> @@ -405,14 +404,14 @@ impl FlownodeBuilder { self.build_manager(query_engine_factory.query_engine()) .await?, ); - let batching = Arc::new(BatchingEngine::new_with_persistence( + let batching = Arc::new(BatchingEngine::new_with_execution( self.frontend_client.clone(), query_engine_factory.query_engine(), self.flow_metadata_manager.clone(), self.table_meta.clone(), self.catalog_manager.clone(), self.opts.flow.batching_mode.clone(), - self.plugins.get::(), + self.plugins.get::>(), )); let dual = Arc::new(FlowDualEngine::new( manager.clone(),