diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 876749dea1..1b4efca8ef 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -844,7 +844,13 @@ 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) + Self::apply_query_failure_to_state( + &mut state, + elapsed, + coverage, + reason, + attempt.is_some(), + ) }; if let Some(decision) = decision { Self::record_checkpoint_decision(flow_id, decision); @@ -877,7 +883,7 @@ impl BatchingTask { 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_with_repair( + let decision = Self::apply_query_result_to_state( &mut state, &res, elapsed, diff --git a/src/flow/src/batching_mode/task/ckpt.rs b/src/flow/src/batching_mode/task/ckpt.rs index dc8f34c029..27e867fa4c 100644 --- a/src/flow/src/batching_mode/task/ckpt.rs +++ b/src/flow/src/batching_mode/task/ckpt.rs @@ -58,6 +58,7 @@ impl BatchingTask { elapsed: Duration, coverage: &QueryCoverage, reason: FlowQueryFallbackReason, + persistence_backed: bool, ) -> Option { state.after_query_exec(elapsed, false); let checkpoint_mode = state.checkpoint_mode(); @@ -79,6 +80,9 @@ 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, @@ -87,21 +91,11 @@ impl BatchingTask { } /// Apply checkpoint transitions for a successfully executed query using its - /// terminal watermark proof and declared coverage. - pub(super) fn apply_query_result_to_state( - state: &mut TaskState, - res: &OutputWithMetrics, - elapsed: Duration, - coverage: &QueryCoverage, - ) -> FlowCheckpointDecision { - Self::apply_query_result_to_state_with_repair(state, res, elapsed, coverage, false) - } - - /// Apply checkpoint transitions while retaining whether this attempt + /// terminal watermark proof and declared coverage, while retaining whether this attempt /// started with a persisted full-repair request. That fact is deliberately /// captured by the caller before execution; a repair request raised later /// must not turn an ordinary full snapshot into a completed repair. - pub(super) fn apply_query_result_to_state_with_repair( + pub(super) fn apply_query_result_to_state( state: &mut TaskState, res: &OutputWithMetrics, elapsed: Duration, diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 4e88d9a501..299bc1896f 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,7 +13,10 @@ // 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; @@ -22,8 +25,8 @@ 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::{ @@ -46,6 +49,7 @@ 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; @@ -940,6 +944,7 @@ fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, + false, ); assert_eq!( @@ -971,6 +976,7 @@ 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. @@ -1002,6 +1008,7 @@ fn test_apply_query_result_to_state_rejects_unproved_watermark() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, + false, ); assert_eq!( @@ -1027,6 +1034,7 @@ fn test_apply_query_result_to_state_reports_missing_watermark() { &result, std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, + false, ); assert_eq!( @@ -1057,6 +1065,7 @@ fn test_apply_query_result_to_state_advances_incremental_subset() { &result, std::time::Duration::from_millis(1), &QueryCoverage::IncrementalDelta, + false, ); assert_eq!( @@ -1092,6 +1101,7 @@ 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!( @@ -1151,6 +1161,7 @@ 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!( @@ -1207,6 +1218,7 @@ 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, @@ -1244,6 +1256,7 @@ 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!( @@ -1318,6 +1331,7 @@ 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!( @@ -1361,6 +1375,7 @@ 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, @@ -1407,6 +1422,7 @@ fn test_apply_query_failure_to_state_falls_back_from_incremental() { std::time::Duration::from_millis(1), &QueryCoverage::IncrementalDelta, FlowQueryFallbackReason::IncrementalQueryFailure, + false, ); assert_eq!( @@ -1434,6 +1450,7 @@ fn test_apply_query_failure_to_state_records_full_snapshot_failure() { std::time::Duration::from_millis(1), &QueryCoverage::UnfilteredFull, FlowQueryFallbackReason::QueryFailure, + false, ); assert_eq!( @@ -1525,6 +1542,7 @@ 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, @@ -1587,6 +1605,7 @@ 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!( @@ -2146,6 +2165,7 @@ async fn test_successful_incremental_checkpoint_fallback_consumes_unscoped_dirty &result, std::time::Duration::from_millis(1), &plan_info.coverage, + false, ) }; assert_eq!( @@ -2555,6 +2575,229 @@ 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;