diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index e0ae41ed4b..2bc299f191 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -30,6 +30,32 @@ mod task; mod time_window; pub(crate) mod utils; +/// Reserved internal epoch column name stamped on every emitted sink state row +/// when checkpoint persistence is active. +/// +/// The enterprise state schema/view adds this column to the sink table; OSS +/// plan generation strips it from the schema-matching view and fills it with +/// the current epoch literal, and OSS restart/recovery reads it to validate +/// checkpoint trust. Ordinary flows (whose sinks never contain this column) +/// are byte-for-byte unaffected. +pub const INTERNAL_FLOW_EPOCH_COL_NAME: &str = "__greptime_internal_flow_epoch"; + +/// Sentinel window timestamp (in milliseconds: 9999-12-31T23:59:59.999Z) used +/// to mark the singleton checkpoint row in the sink table's window/time-index +/// column. +/// +/// This value is a private convention between the flow runtime and the +/// internal producer of the sink state schema (the enterprise state schema/ +/// view). It is NOT safe by construction: the flow's own query may bin source +/// timestamps with an arbitrary `date_bin(stride, origin)` lattice, so the +/// internal producer MUST validate at CREATE time that the sentinel cannot +/// collide with any real window for the exact flow `date_bin` (stride and +/// origin). The OSS runtime never auto-creates the sentinel row; it only +/// reads/writes it when the sink schema already contains the reserved epoch +/// column, and ordinary flows (without that column) are byte-for-byte +/// unaffected. +pub const CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS: i64 = 253_402_300_799_999; + /// Incremental read mode for a batching flow, selected only through the /// reserved internal flow option /// [`common_meta::ddl::create_flow::INTERNAL_INCREMENTAL_MODE_KEY`]. diff --git a/src/flow/src/batching_mode/checkpoint.rs b/src/flow/src/batching_mode/checkpoint.rs index 11be951a9a..e9d5329b44 100644 --- a/src/flow/src/batching_mode/checkpoint.rs +++ b/src/flow/src/batching_mode/checkpoint.rs @@ -12,13 +12,81 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; + +use common_time::timestamp::TimeUnit; +use serde::{Deserialize, Serialize}; + +use crate::Error; +use crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS; use crate::batching_mode::state::CheckpointMode; +use crate::error::UnexpectedSnafu; pub(super) const CHECKPOINT_DECISION_ADVANCE: &str = "advance"; pub(super) const CHECKPOINT_DECISION_FALLBACK: &str = "fallback"; pub(super) const CHECKPOINT_DECISION_CONTINUE_REPAIR: &str = "continue_repair"; pub(super) const CHECKPOINT_REASON_NONE: &str = "none"; +/// Version of the private on-disk checkpoint record format. Bump when the +/// serialized shape changes; old versions are rejected on load (backfill). +pub(super) const CHECKPOINT_RECORD_FORMAT_VERSION: u32 = 1; + +/// The private, versioned checkpoint record stored in the sink table's BINARY +/// state column of the singleton sentinel row. +/// +/// `serde_json` is used on purpose: it is an existing flow dependency, and the +/// `BTreeMap` key order makes the encoding deterministic for a given value. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(super) struct CheckpointRecord { + /// Record format version; must equal [`CHECKPOINT_RECORD_FORMAT_VERSION`]. + pub format_version: u32, + /// Epoch of the persisted checkpoint. Rows stamped with a larger epoch are + /// newer than this record and invalidate it (crash between state write and + /// checkpoint write). + pub epoch: u64, + /// Region id -> last consumed watermark sequence map. Must be non-empty to + /// be trusted. + pub checkpoints: BTreeMap, +} + +/// Encode a checkpoint record deterministically. +pub(super) fn encode_checkpoint_record(record: &CheckpointRecord) -> Result, Error> { + serde_json::to_vec(record).map_err(|err| { + UnexpectedSnafu { + reason: err.to_string(), + } + .build() + }) +} + +/// Decode a checkpoint record, rejecting unknown format versions. +pub(super) fn decode_checkpoint_record(bytes: &[u8]) -> Result, Error> { + let record: CheckpointRecord = match serde_json::from_slice(bytes) { + Ok(record) => record, + Err(_) => return Ok(None), + }; + if record.format_version != CHECKPOINT_RECORD_FORMAT_VERSION { + return Ok(None); + } + Ok(Some(record)) +} + +/// Convert the millisecond sentinel window timestamp to the sink window +/// column's native time unit. +/// +/// A nanosecond sentinel at year 9999 would overflow `i64`, so the nanosecond +/// representation is clamped to the largest representable value; every +/// practical source timestamp is far below it. Second/microsecond conversions +/// are exact. +pub(super) fn checkpoint_sentinel_ts_in_unit(unit: TimeUnit) -> i64 { + match unit { + TimeUnit::Second => CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS / 1000, + TimeUnit::Millisecond => CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS, + TimeUnit::Microsecond => CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS * 1000, + TimeUnit::Nanosecond => i64::MAX, + } +} + /// Why the task fell back to full snapshot mode. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum FlowQueryFallbackReason { @@ -48,6 +116,11 @@ pub(super) enum FlowQueryFallbackReason { /// Incremental mode has been permanently disabled for this Flow /// (e.g. because the query shape is not incrementally safe). IncrementalDisabled, + /// The sink state rows were written but the singleton checkpoint row could + /// not be persisted (write failure or ambiguous result). The runtime + /// resets to full snapshot so a later cycle rebuilds and re-persists the + /// checkpoint instead of claiming persistence it does not have. + CheckpointPersistFailure, } impl FlowQueryFallbackReason { @@ -61,6 +134,7 @@ impl FlowQueryFallbackReason { Self::IncrementalQueryFailure => "incremental_query_failure", Self::QueryFailure => "query_failure", Self::IncrementalDisabled => "incremental_disabled", + Self::CheckpointPersistFailure => "checkpoint_persist_failure", } } } @@ -151,3 +225,64 @@ pub(super) fn checkpoint_mode_label(mode: CheckpointMode) -> &'static str { CheckpointMode::Incremental => "incremental", } } + +#[cfg(test)] +mod tests { + use common_time::timestamp::TimeUnit; + + use super::*; + use crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS; + + #[test] + fn test_checkpoint_record_roundtrip_and_version_rejection() { + let record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 42, + checkpoints: BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]), + }; + let encoded = encode_checkpoint_record(&record).unwrap(); + assert_eq!( + record, + decode_checkpoint_record(&encoded) + .unwrap() + .expect("roundtrip") + ); + + // Deterministic encoding: same value -> same bytes. + assert_eq!(encoded, encode_checkpoint_record(&record).unwrap()); + + // Garbage bytes are not a decodable record. + assert!(decode_checkpoint_record(b"not-a-record").unwrap().is_none()); + + // Unknown format versions are rejected (future compat guard). + let mut json = serde_json::to_value(&record).unwrap(); + json["format_version"] = serde_json::json!(2); + let bytes = serde_json::to_vec(&json).unwrap(); + assert!(decode_checkpoint_record(&bytes).unwrap().is_none()); + } + + #[test] + fn test_checkpoint_sentinel_ts_in_unit_conversion() { + let sentinel = CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS; + assert_eq!( + sentinel / 1000, + checkpoint_sentinel_ts_in_unit(TimeUnit::Second) + ); + assert_eq!( + sentinel, + checkpoint_sentinel_ts_in_unit(TimeUnit::Millisecond) + ); + assert_eq!( + sentinel * 1000, + checkpoint_sentinel_ts_in_unit(TimeUnit::Microsecond) + ); + // A year-9999 nanosecond sentinel would overflow i64, so the + // nanosecond representation is clamped to the largest representable + // value. The exact clamped value is a storage detail; the assertion + // pins that it stays representable and far above any real timestamp. + assert_eq!( + i64::MAX, + checkpoint_sentinel_ts_in_unit(TimeUnit::Nanosecond) + ); + } +} diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index 4daf911271..ae9d65ec8a 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -729,6 +729,13 @@ impl BatchingEngine { // CREATE FLOW time instead of surfacing them later in the execution loop. task.check_or_create_sink_table(&engine, &frontend).await?; task.validate_sink_table_schema(&engine).await?; + // Restore the durable checkpoint from the sink table (when persistence + // applies) before the first run, so incremental extensions emit `(C,H]` + // from the restored checkpoint map. Restore scans run through the + // frontend (not the local query engine) so DistTable-backed catalog + // metadata resolves on the frontend. Load errors are warnings: the task + // starts from full snapshot and later cycles rebuild the checkpoint. + task.try_enable_checkpoint_persistence(&frontend).await; let (start_tx, start_rx) = oneshot::channel(); diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index 59c4f0fe10..cf6a20ac96 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -60,6 +60,15 @@ pub struct TaskState { /// Set when the flow's query shape is deterministically incompatible /// with incremental execution (e.g. unsupported aggregate expressions). incremental_disabled: bool, + /// Checkpoint persistence layout, activated only when the batching mode is + /// `SequenceRange` and the sink schema contains the reserved internal epoch + /// column plus a unique BINARY state column. `None` keeps the task + /// byte-for-byte identical to an ordinary flow. + pub(crate) checkpoint_persistence: Option, + /// Epoch of the last durably persisted checkpoint record. Advanced only + /// after the singleton checkpoint row write succeeds; rows stamped with a + /// larger epoch invalidate the durable record on restart. + persisted_epoch: u64, exec_state: ExecState, /// Shutdown receiver pub(crate) shutdown_rx: oneshot::Receiver<()>, @@ -87,6 +96,8 @@ impl TaskState { pending_fenced_repair: None, checkpoints: Default::default(), incremental_disabled: false, + checkpoint_persistence: None, + persisted_epoch: 0, exec_state: ExecState::Idle, shutdown_rx, task_handle: None, @@ -131,6 +142,46 @@ impl TaskState { &self.checkpoints } + /// Returns the resolved checkpoint persistence layout, if activated. + pub fn checkpoint_persistence(&self) -> Option<&CheckpointPersistence> { + self.checkpoint_persistence.as_ref() + } + + /// Epoch of the last durably persisted checkpoint record (0 = none). + pub fn persisted_epoch(&self) -> u64 { + self.persisted_epoch + } + + /// The epoch the current cycle must stamp onto emitted state rows and, on + /// a successful checkpoint write, persist. One past the last durable epoch + /// so rows from an unpersisted cycle always invalidate the older record. + pub fn next_persist_epoch(&self) -> u64 { + self.persisted_epoch.saturating_add(1) + } + + /// Record a successfully persisted checkpoint epoch. Called only after the + /// singleton checkpoint row write succeeds; never before it. + pub fn advance_persisted_epoch(&mut self, epoch: u64) { + self.persisted_epoch = self.persisted_epoch.max(epoch); + } + + /// Activate or deactivate checkpoint persistence for this task. + pub fn set_checkpoint_persistence(&mut self, persistence: Option) { + self.checkpoint_persistence = persistence; + } + + /// Seed the task from a trusted restored checkpoint record: replace the + /// in-memory checkpoint map, pin the durable epoch, and enter Incremental + /// mode (unless incremental is permanently disabled). + pub fn seed_checkpoints_from_record(&mut self, epoch: u64, checkpoints: BTreeMap) { + self.persisted_epoch = epoch; + self.checkpoints = checkpoints; + self.pending_fenced_repair = None; + if !self.incremental_disabled { + self.checkpoint_mode = CheckpointMode::Incremental; + } + } + /// 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> { @@ -918,6 +969,23 @@ pub enum CheckpointMode { Incremental, } +/// Column layout of the sink table required for checkpoint persistence. +/// +/// Resolved once at task creation when the batching mode is `SequenceRange` +/// and the sink schema contains the reserved internal epoch column plus a +/// unique BINARY state column. The window column is the sink time-index column +/// whose sentinel value marks the singleton checkpoint row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CheckpointPersistence { + /// Name of the reserved internal epoch column + /// ([`crate::batching_mode::INTERNAL_FLOW_EPOCH_COL_NAME`]). + pub epoch_col_name: String, + /// Name of the sink BINARY column storing the encoded checkpoint record. + pub state_col_name: String, + /// Name of the sink window/time-index column used for the sentinel row. + pub window_col_name: String, +} + /// Dirty windows that must be repaired under a frozen full-snapshot watermark. /// This is a FullSnapshot sub-state, not a separate checkpoint mode. #[derive(Debug, Clone)] diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 196c9ea42d..779b32cb00 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -20,16 +20,23 @@ use api::v1::{CreateTableExpr, TableName}; use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_query::logical_plan::breakup_insert_plan; +use common_recordbatch::RecordBatches; +use common_recordbatch::util::collect_batches; use common_telemetry::tracing::warn; use common_telemetry::{debug, info}; use common_time::Timestamp; use datafusion::datasource::DefaultTableSource; +use datafusion::functions_aggregate::expr_fn::{count, max}; 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_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; +use datafusion_common::{Column, DFSchema, DFSchemaRef, ScalarValue, TableReference}; +use datafusion_expr::logical_plan::{EmptyRelation, TableScan}; +use datafusion_expr::{DmlStatement, Expr, LogicalPlan, LogicalPlanBuilder, WriteOp, col, lit}; +use datatypes::prelude::ConcreteDataType; use datatypes::schema::Schema; +use datatypes::value::Value; +use datatypes::vectors::Helper; use query::QueryEngineRef; use query::options::FLOW_INCREMENTAL_MODE; use query::query_engine::DefaultSerializer; @@ -43,12 +50,17 @@ use tokio::sync::oneshot::error::TryRecvError; use tokio::sync::{Mutex, oneshot}; use tokio::time::Instant; -use crate::batching_mode::BatchingModeOptions; -use crate::batching_mode::checkpoint::checkpoint_mode_label; +use crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL; +use crate::batching_mode::checkpoint::{ + CHECKPOINT_RECORD_FORMAT_VERSION, CheckpointRecord, FlowCheckpointDecision, + FlowQueryFallbackReason, checkpoint_mode_label, checkpoint_sentinel_ts_in_unit, + decode_checkpoint_record, encode_checkpoint_record, +}; use crate::batching_mode::eval_schedule::{EvalSchedule, select_due_scheduled_times}; use crate::batching_mode::frontend_client::{FrontendClient, PeerDesc}; use crate::batching_mode::state::{ - CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, to_df_literal, + CheckpointMode, CheckpointPersistence, DirtyTimeWindows, FilterExprInfo, TaskState, + to_df_literal, }; use crate::batching_mode::table_creator::{QueryType, create_table_with_expr}; use crate::batching_mode::time_window::TimeWindowExpr; @@ -56,10 +68,11 @@ use crate::batching_mode::utils::{ AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema, get_table_info_df_schema, sql_to_df_plan, }; +use crate::batching_mode::{BatchingModeOptions, INTERNAL_FLOW_EPOCH_COL_NAME, IncrementalMode}; use crate::df_optimizer::apply_df_optimizer; use crate::error::{ - DatafusionSnafu, ExternalSnafu, InvalidQuerySnafu, SubstraitEncodeLogicalPlanSnafu, - UnexpectedSnafu, + DatafusionSnafu, DatatypesSnafu, ExternalSnafu, InvalidQuerySnafu, + SubstraitEncodeLogicalPlanSnafu, UnexpectedSnafu, }; use crate::metrics::{ METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT, METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME, @@ -235,6 +248,67 @@ struct ExecuteOnceOutcome { result: Result, Error>, } +/// True when `data_type` is a plain integer type (signed or unsigned). +fn is_integer_type(data_type: &ConcreteDataType) -> bool { + matches!( + data_type, + ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_) + ) +} + +/// Returns a copy of the sink schema without the reserved internal epoch column +/// and the primary-key indices remapped onto the stripped schema. +/// +/// Plan generation matches flow output against this stripped view; the epoch +/// column is stamped separately by checkpoint persistence +/// ([`BatchingTask::stamp_epoch_into_plan`]). The real query-produced BINARY +/// state column (e.g. the UDDSketch state of an exact EE-like flow) is NOT +/// stripped: it is ordinary flow output and must survive schema matching, +/// incremental rewrite, and the final DML insert. +fn strip_internal_epoch_column( + schema: &Schema, + primary_key_indices: &[usize], +) -> (Schema, Vec) { + let mut column_schemas = Vec::with_capacity(schema.column_schemas().len()); + let mut new_primary_key_indices = Vec::new(); + for (idx, column) in schema.column_schemas().iter().enumerate() { + if column.name == INTERNAL_FLOW_EPOCH_COL_NAME { + continue; + } + let new_idx = column_schemas.len(); + column_schemas.push(column.clone()); + if primary_key_indices.contains(&idx) { + new_primary_key_indices.push(new_idx); + } + } + (Schema::new(column_schemas), new_primary_key_indices) +} + +/// Builds the sentinel window literal in the given timestamp column's native +/// unit. Errors when the column is not a timestamp. +fn checkpoint_sentinel_scalar( + column: &datatypes::schema::ColumnSchema, +) -> Result { + let ts_type = column + .data_type + .as_timestamp() + .with_context(|| UnexpectedSnafu { + reason: format!( + "Expected timestamp column for checkpoint sentinel, found {}", + column.data_type + ), + })?; + let sentinel = checkpoint_sentinel_ts_in_unit(ts_type.unit()); + to_df_literal(Timestamp::new(sentinel, ts_type.unit())) +} + impl BatchingTask { #[allow(clippy::too_many_arguments)] pub fn try_new( @@ -380,13 +454,18 @@ impl BatchingTask { 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(); + // The reserved internal epoch column is not produced by the flow query; + // it is stamped separately when checkpoint persistence is active, so it + // is excluded from schema matching. The real BINARY state column stays. + let (effective_schema, effective_pk_indices) = + strip_internal_epoch_column(&table_meta.schema, &primary_key_indices); gen_plan_with_matching_schema( &self.config.query, query_ctx, engine.clone(), - table_meta.schema.clone(), - &primary_key_indices, + Arc::new(effective_schema), + &effective_pk_indices, merge_mode_last_non_null, ) .await @@ -486,12 +565,19 @@ impl BatchingTask { 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(); + // The reserved internal epoch column is stamped onto every emitted row + // separately (see `stamp_epoch_into_plan`), so plan generation matches + // flow output against the sink schema without it. The real + // query-produced BINARY state column stays in the matched schema. + let (effective_schema, effective_pk_indices) = + strip_internal_epoch_column(&table_meta.schema, &primary_key_indices); + let effective_schema = Arc::new(effective_schema); let new_query = self .gen_query_with_time_window( engine.clone(), - &table.table_info().meta.schema, - &primary_key_indices, + &effective_schema, + &effective_pk_indices, merge_mode_last_non_null, max_window_cnt, ) @@ -629,6 +715,11 @@ impl BatchingTask { } let plan = incremental_plan.unwrap_or_else(|| plan.clone()); + // Stamp the current cycle epoch onto every emitted state row when + // checkpoint persistence is active. The epoch is decided here, once + // per cycle, and reused for the checkpoint row write after success. + let (plan, cycle_epoch) = self.stamp_epoch_into_plan(plan).await?; + let extensions = self .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta()) .await?; @@ -779,9 +870,57 @@ impl BatchingTask { METRIC_FLOW_ROWS .with_label_values(&[format!("{}-out-batching", flow_id).as_str()]) .inc_by(affected_rows as _); + // Checkpoint persistence: apply the single authoritative checkpoint + // transition first, then persist the singleton checkpoint row only when + // the actual decision advanced checkpoints, using the resulting + // `state.checkpoints()` snapshot. The whole cycle runs under + // `execution_lock`, so no other execution interleaves with the write. + // If the write succeeds, advance the durable epoch; if it fails, reset + // to full snapshot and restore the executed plan's consumed dirty work + // so the next cycle re-runs a full repair/backfill and can write a + // replacement checkpoint. Dirty notifications arriving around the + // transition are never erased. let decision = { let mut state = self.state.write().unwrap(); - Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage) + let decision = Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage); + let persist = cycle_epoch + .filter(|_| { + matches!( + decision, + FlowCheckpointDecision::AdvancedFromFullSnapshot { .. } + | FlowCheckpointDecision::AdvancedIncremental { .. } + ) + }) + .map(|epoch| (epoch, state.checkpoints().clone(), state.checkpoint_mode())); + (decision, persist) + }; + let decision = match decision.1 { + Some((epoch, map_to_persist, previous_mode)) => { + match self + .write_checkpoint_row(frontend_client, epoch, &map_to_persist) + .await + { + Ok(()) => { + let mut state = self.state.write().unwrap(); + state.advance_persisted_epoch(epoch); + decision.0 + } + Err(err) => { + warn!( + "Flow {flow_id} failed to persist checkpoint row, falling back to full snapshot and re-scheduling dirty work: {err:?}" + ); + let mut state = self.state.write().unwrap(); + state.mark_full_snapshot(); + drop(state); + self.restore_dirty_windows(dirty_restore); + FlowCheckpointDecision::FallbackToFullSnapshot { + previous_mode, + reason: FlowQueryFallbackReason::CheckpointPersistFailure, + } + } + } + } + None => decision.0, }; Self::record_checkpoint_decision(flow_id, decision); @@ -1509,6 +1648,512 @@ impl BatchingTask { Ok(Some(info)) } + + /// Stamp the current cycle epoch onto every emitted sink state row when + /// checkpoint persistence is active. Returns the stamped plan and the epoch + /// used for this cycle (`None` when persistence is inactive or the plan is + /// not a DML insert). + /// + /// Rows stamped with an epoch newer than the last durable record + /// invalidate that record on restart (crash between state write and + /// checkpoint write), which is exactly the backfill trigger we want. + async fn stamp_epoch_into_plan( + &self, + plan: LogicalPlan, + ) -> Result<(LogicalPlan, Option), Error> { + let persistence = self.state.read().unwrap().checkpoint_persistence().cloned(); + let Some(persistence) = persistence else { + return Ok((plan, None)); + }; + let LogicalPlan::Dml(dml) = &plan else { + return Ok((plan, None)); + }; + let epoch = self.state.read().unwrap().next_persist_epoch(); + let inner = dml.input.as_ref().clone(); + let mut exprs = inner + .schema() + .fields() + .iter() + .map(|field| Expr::Column(Column::new_unqualified(field.name()))) + .collect::>(); + exprs.push(lit(ScalarValue::UInt64(Some(epoch))).alias(&persistence.epoch_col_name)); + let stamped = LogicalPlanBuilder::from(inner) + .project(exprs) + .with_context(|_| DatafusionSnafu { + context: "Failed to stamp flow epoch column onto state rows".to_string(), + })? + .build() + .with_context(|_| DatafusionSnafu { + context: "Failed to build epoch-stamped state row plan".to_string(), + })?; + let stamped = LogicalPlan::Dml(DmlStatement::new( + dml.table_name.clone(), + dml.target.clone(), + dml.op.clone(), + Arc::new(stamped), + )); + Ok((stamped, Some(epoch))) + } + + /// Upsert the singleton checkpoint row through the frontend write + /// machinery: one row whose window/time-index column is the sentinel, epoch + /// column is the cycle epoch, and the BINARY state column holds the encoded + /// versioned checkpoint record. Other sink columns are left to defaults. + async fn write_checkpoint_row( + &self, + frontend_client: &Arc, + epoch: u64, + checkpoints: &BTreeMap, + ) -> Result<(), Error> { + let persistence = self + .state + .read() + .unwrap() + .checkpoint_persistence() + .cloned() + .with_context(|| UnexpectedSnafu { + reason: "checkpoint persistence is not active".to_string(), + })?; + let record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch, + checkpoints: checkpoints.clone(), + }; + let encoded = encode_checkpoint_record(&record)?; + + let (table, _) = get_table_info_df_schema( + self.config.catalog_manager.clone(), + self.config.sink_table_name.clone(), + ) + .await?; + let sink_schema = table.table_info().meta.schema.clone(); + let window_col = sink_schema + .column_schema_by_name(&persistence.window_col_name) + .with_context(|| UnexpectedSnafu { + reason: format!( + "Sink table lost checkpoint window column {}", + persistence.window_col_name + ), + })?; + + let mut exprs = vec![ + lit(checkpoint_sentinel_scalar(window_col)?).alias(&persistence.window_col_name), + lit(ScalarValue::UInt64(Some(epoch))).alias(&persistence.epoch_col_name), + lit(ScalarValue::Binary(Some(encoded))).alias(&persistence.state_col_name), + ]; + // Keep the auto-created update_at column fresh on the checkpoint row + // when the sink has one; any other sink column is filled with its + // default (or NULL) by the insert machinery. + if let Some(update_at) = sink_schema.column_schema_by_name(AUTO_CREATED_UPDATE_AT_TS_COL) + && update_at.data_type.is_timestamp() + { + exprs.push(datafusion::prelude::now().alias(AUTO_CREATED_UPDATE_AT_TS_COL)); + } + let empty = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: true, + schema: Arc::new(DFSchema::empty()), + }); + let row_plan = LogicalPlanBuilder::from(empty) + .project(exprs) + .with_context(|_| DatafusionSnafu { + context: "Failed to build checkpoint row plan".to_string(), + })? + .build() + .with_context(|_| DatafusionSnafu { + context: "Failed to finalize checkpoint row plan".to_string(), + })?; + + let insert_to = TableName { + catalog_name: self.config.sink_table_name[0].clone(), + schema_name: self.config.sink_table_name[1].clone(), + table_name: self.config.sink_table_name[2].clone(), + }; + let req = encode_insert_plan_request(insert_to, &row_plan)?; + let catalog = &self.config.sink_table_name[0]; + let schema = &self.config.sink_table_name[1]; + let mut peer_desc = None; + frontend_client + .query_with_terminal_metrics(catalog, schema, req, &[], &HashMap::new(), &mut peer_desc) + .await?; + debug!( + "Flow {} persisted checkpoint row with epoch {} and {} regions", + self.config.flow_id, + epoch, + checkpoints.len() + ); + Ok(()) + } + + /// Detect whether checkpoint persistence applies to this task and, if so, + /// restore the durable checkpoint from the sink table before the first run. + /// + /// Detection/load errors are warnings: the task starts from full snapshot + /// and later cycles rebuild the checkpoint from scratch. + pub(crate) async fn try_enable_checkpoint_persistence( + &self, + frontend_client: &Arc, + ) { + let persistence = match self.detect_checkpoint_persistence().await { + Ok(persistence) => persistence, + Err(err) => { + warn!( + "Flow {} failed to detect checkpoint persistence, starting from full snapshot: {err:?}", + self.config.flow_id + ); + return; + } + }; + let restored = match &persistence { + Some(persistence) => { + self.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(persistence.clone())); + match self + .read_checkpoint_state(frontend_client, persistence) + .await + { + Ok(restored) => restored, + Err(err) => { + warn!( + "Flow {} failed to load checkpoint state, starting from full snapshot: {err:?}", + self.config.flow_id + ); + None + } + } + } + None => None, + }; + let mut state = self.state.write().unwrap(); + if let Some((epoch, checkpoints)) = restored { + state.seed_checkpoints_from_record(epoch, checkpoints); + info!( + "Flow {} restored incremental checkpoints from sink table at epoch {}", + self.config.flow_id, epoch + ); + } else { + info!( + "Flow {} found no trustworthy checkpoint record; starting from full snapshot", + self.config.flow_id + ); + } + } + + /// Resolve the checkpoint persistence layout. Activated only when the + /// batching mode is `SequenceRange` and the sink schema contains the + /// reserved internal epoch column plus exactly one BINARY state column. + async fn detect_checkpoint_persistence(&self) -> Result, Error> { + if self.config.batch_opts.incremental_mode != IncrementalMode::SequenceRange { + return Ok(None); + } + let (table, _) = get_table_info_df_schema( + self.config.catalog_manager.clone(), + self.config.sink_table_name.clone(), + ) + .await?; + let schema = table.table_info().meta.schema.clone(); + let epoch_col = match schema.column_schema_by_name(INTERNAL_FLOW_EPOCH_COL_NAME) { + Some(column) => column, + None => return Ok(None), + }; + if !is_integer_type(&epoch_col.data_type) { + debug!( + "Flow {} checkpoint epoch column {} has non-integer type {:?}; persistence inactive", + self.config.flow_id, epoch_col.name, epoch_col.data_type + ); + return Ok(None); + } + let state_cols = schema + .column_schemas() + .iter() + .filter(|column| column.data_type == ConcreteDataType::binary_datatype()) + .collect::>(); + if state_cols.len() != 1 { + return Ok(None); + } + let Some(ts_idx) = schema.timestamp_index() else { + return Ok(None); + }; + let window_col = &schema.column_schemas()[ts_idx]; + Ok(Some(CheckpointPersistence { + epoch_col_name: INTERNAL_FLOW_EPOCH_COL_NAME.to_string(), + state_col_name: state_cols[0].name.clone(), + window_col_name: window_col.name.clone(), + })) + } + + /// Read the singleton sentinel checkpoint row and the maximum non-sentinel + /// epoch from the sink table and return a trusted record. + /// + /// The reads are executed through the frontend client (as `Query::LogicalPlan` + /// requests) rather than the local query engine, so the sink table is + /// resolved from DistTable-backed catalog metadata on the frontend. + /// + /// Trust requires exactly one sentinel row, a decodable v1 record with a + /// non-empty checkpoint map, and `max_non_sentinel_epoch <= record.epoch`. + /// NULL epoch rows (pre-upgrade data) make the state untrusted unless + /// there are no state rows at all. + async fn read_checkpoint_state( + &self, + frontend_client: &Arc, + persistence: &CheckpointPersistence, + ) -> Result)>, Error> { + let (table, df_schema) = get_table_info_df_schema( + self.config.catalog_manager.clone(), + self.config.sink_table_name.clone(), + ) + .await?; + let sink_schema = table.table_info().meta.schema.clone(); + let window_col = sink_schema + .column_schema_by_name(&persistence.window_col_name) + .with_context(|| UnexpectedSnafu { + reason: format!( + "Sink table lost checkpoint window column {}", + persistence.window_col_name + ), + })?; + let sentinel = checkpoint_sentinel_scalar(window_col)?; + // The sentinel query projects the window column too, because the + // sentinel filter is applied as a logical Filter node on top of the + // scan (the local test engine's MemTable does not apply scan filters). + let sentinel_plan = sink_scan_plan( + &self.config.sink_table_name, + table.clone(), + &df_schema, + &[ + persistence.window_col_name.clone(), + persistence.epoch_col_name.clone(), + persistence.state_col_name.clone(), + ], + )?; + let sentinel_plan = LogicalPlanBuilder::from(sentinel_plan) + .filter(col(&persistence.window_col_name).eq(lit(sentinel.clone()))) + .with_context(|_| DatafusionSnafu { + context: "Failed to build checkpoint sentinel row scan".to_string(), + })? + .build() + .with_context(|_| DatafusionSnafu { + context: "Failed to finalize checkpoint sentinel row scan".to_string(), + })?; + let sentinel_batches = self + .execute_sink_scan(frontend_client, sentinel_plan) + .await?; + let mut sentinel_states = Vec::new(); + for batch in sentinel_batches.iter() { + let state_vector = + Helper::try_into_vector(batch.column(2)).with_context(|_| DatatypesSnafu { + extra: "Failed to convert sentinel row state column".to_string(), + })?; + for row in 0..batch.num_rows() { + let state_bytes = state_vector + .get_ref(row) + .try_into_binary() + .with_context(|_| DatatypesSnafu { + extra: "Failed to convert sentinel row state".to_string(), + })? + .map(|bytes| bytes.to_vec()); + sentinel_states.push(state_bytes); + } + } + if sentinel_states.len() > 1 { + debug!( + "Flow {} found {} sentinel checkpoint rows, untrusted", + self.config.flow_id, + sentinel_states.len() + ); + return Ok(None); + } + + // Maximum non-sentinel epoch plus row counts to detect NULL epochs. + let non_sentinel_predicate = col(&persistence.window_col_name) + .is_null() + .or(col(&persistence.window_col_name).not_eq(lit(sentinel))); + let epoch_col_name = persistence.epoch_col_name.clone(); + let scan = sink_scan_plan( + &self.config.sink_table_name, + table, + &df_schema, + &[persistence.window_col_name.clone(), epoch_col_name.clone()], + )?; + let agg_plan = LogicalPlanBuilder::from(scan) + .filter(non_sentinel_predicate) + .with_context(|_| DatafusionSnafu { + context: "Failed to build non-sentinel epoch aggregation".to_string(), + })? + .aggregate( + Vec::::new(), + vec![ + count(col(&epoch_col_name)).alias("non_null_epoch_cnt"), + count(lit(1_i64)).alias("total_cnt"), + max(col(&epoch_col_name)).alias("max_epoch"), + ], + ) + .with_context(|_| DatafusionSnafu { + context: "Failed to aggregate non-sentinel epochs".to_string(), + })? + .build() + .with_context(|_| DatafusionSnafu { + context: "Failed to finalize non-sentinel epoch aggregation".to_string(), + })?; + let agg_batches = self.execute_sink_scan(frontend_client, agg_plan).await?; + let (total_cnt, non_null_cnt, max_epoch) = match agg_batches.iter().next() { + Some(batch) => { + let non_null_cnt_vector = + Helper::try_into_vector(batch.column(0)).with_context(|_| DatatypesSnafu { + extra: "Failed to convert non-null epoch count column".to_string(), + })?; + let total_cnt_vector = + Helper::try_into_vector(batch.column(1)).with_context(|_| DatatypesSnafu { + extra: "Failed to convert state row count column".to_string(), + })?; + let max_epoch_vector = + Helper::try_into_vector(batch.column(2)).with_context(|_| DatatypesSnafu { + extra: "Failed to convert max epoch column".to_string(), + })?; + let non_null_cnt = non_null_cnt_vector + .get_ref(0) + .try_into_i64() + .with_context(|_| DatatypesSnafu { + extra: "Failed to convert non-null epoch count".to_string(), + })? + .unwrap_or(0); + let total_cnt = total_cnt_vector + .get_ref(0) + .try_into_i64() + .with_context(|_| DatatypesSnafu { + extra: "Failed to convert state row count".to_string(), + })? + .unwrap_or(0); + let max_epoch = value_as_u64(Value::from(max_epoch_vector.get_ref(0))); + (total_cnt, non_null_cnt, max_epoch) + } + None => (0, 0, None), + }; + + // NULL epoch rows (pre-upgrade state data) are untrusted unless there + // is no state data at all. + if total_cnt > 0 && non_null_cnt != total_cnt { + debug!( + "Flow {} has {} state rows with NULL epochs, checkpoint untrusted", + self.config.flow_id, + total_cnt - non_null_cnt + ); + return Ok(None); + } + + let [state_bytes] = sentinel_states.as_slice() else { + return Ok(None); + }; + let Some(state_bytes) = state_bytes.as_ref() else { + debug!( + "Flow {} sentinel row has NULL state, untrusted", + self.config.flow_id + ); + return Ok(None); + }; + let Some(record) = decode_checkpoint_record(state_bytes)? else { + debug!( + "Flow {} sentinel row is not a decodable v1 record, untrusted", + self.config.flow_id + ); + return Ok(None); + }; + if record.checkpoints.is_empty() { + debug!( + "Flow {} checkpoint record has an empty map, untrusted", + self.config.flow_id + ); + return Ok(None); + } + if let Some(max_epoch) = max_epoch + && max_epoch > record.epoch + { + debug!( + "Flow {} state rows newer than checkpoint record ({} > {}), untrusted", + self.config.flow_id, max_epoch, record.epoch + ); + return Ok(None); + } + Ok(Some((record.epoch, record.checkpoints))) + } + + /// Execute a sink scan plan through the frontend client and collect the + /// returned record batches. + /// + /// The plan is transported as a `Query::LogicalPlan` request (not executed + /// on the local query engine) so restore works with DistTable-backed + /// catalog metadata: the frontend resolves the sink table and executes. + async fn execute_sink_scan( + &self, + frontend_client: &Arc, + plan: LogicalPlan, + ) -> Result { + let message = DFLogicalSubstraitConvertor {} + .encode(&plan, DefaultSerializer) + .context(SubstraitEncodeLogicalPlanSnafu)?; + let req = api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())), + }; + let catalog = &self.config.sink_table_name[0]; + let schema = &self.config.sink_table_name[1]; + let mut peer_desc = None; + let output = frontend_client + .query_with_terminal_metrics(catalog, schema, req, &[], &HashMap::new(), &mut peer_desc) + .await?; + let batches = match output.output.data { + common_query::OutputData::RecordBatches(batches) => batches, + common_query::OutputData::Stream(stream) => collect_batches(stream) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?, + common_query::OutputData::AffectedRows(_) => { + return UnexpectedSnafu { + reason: "Unexpected affected-rows output from sink scan".to_string(), + } + .fail(); + } + }; + Ok(batches) + } +} + +/// Build a table scan over the sink table with the given column projection. +fn sink_scan_plan( + sink_table_name: &[String; 3], + table: table::TableRef, + df_schema: &DFSchema, + projection: &[String], +) -> Result { + let table_ref = TableReference::Full { + catalog: sink_table_name[0].clone().into(), + schema: sink_table_name[1].clone().into(), + table: sink_table_name[2].clone().into(), + }; + let table_provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(table_provider)); + let projection = projection + .iter() + .map(|name| { + df_schema + .index_of_column(&Column::from_name(name.clone())) + .with_context(|_| DatafusionSnafu { + context: format!("Failed to resolve sink column {name} for checkpoint scan"), + }) + }) + .collect::, _>>()?; + let scan = TableScan::try_new(table_ref, table_source, Some(projection), vec![], None) + .with_context(|_| DatafusionSnafu { + context: "Failed to build sink scan for checkpoint".to_string(), + })?; + Ok(LogicalPlan::TableScan(scan)) +} + +/// Extracts a `u64` from an integer-typed value (signed or unsigned). +fn value_as_u64(value: Value) -> Option { + value + .as_u64() + .or_else(|| value.as_i64().map(|value| value as u64)) } #[cfg(test)] diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index dcc9cf9a01..4a71e3f445 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -25,7 +25,8 @@ use common_query::Output; use common_recordbatch::RecordBatch; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; use datatypes::data_type::ConcreteDataType as CDT; -use datatypes::schema::ColumnSchema; +use datatypes::prelude::{MutableVector, ScalarVectorBuilder}; +use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{ TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, }; @@ -36,15 +37,18 @@ use query::options::{ }; use session::context::QueryContext; use snafu::ResultExt; +use substrait::DFLogicalSubstraitConvertor; use table::test_util::MemTable; use super::*; use crate::batching_mode::IncrementalMode; use crate::batching_mode::checkpoint::{ CHECKPOINT_DECISION_ADVANCE, CHECKPOINT_DECISION_FALLBACK, CHECKPOINT_REASON_NONE, - FlowCheckpointDecision, FlowQueryFallbackReason, + CHECKPOINT_RECORD_FORMAT_VERSION, CheckpointRecord, FlowCheckpointDecision, + FlowQueryFallbackReason, decode_checkpoint_record, encode_checkpoint_record, }; -use crate::batching_mode::state::CheckpointMode; +use crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError; +use crate::batching_mode::state::{CheckpointMode, CheckpointPersistence}; use crate::batching_mode::time_window::find_time_window_expr; use crate::test_utils::create_test_query_engine; @@ -2597,3 +2601,1023 @@ async fn test_insert_plan_matching_failure_restores_consumed_dirty_marker() { std::time::Duration::from_secs(5) ); } + +// --------------------------------------------------------------------------- +// Checkpoint persistence: record codec, activation, restore, stamping, and +// checkpoint row writes. +// --------------------------------------------------------------------------- + +/// The exact EE-like sink schema: a real query-produced BINARY state column, a +/// timestamp time-index window column, an auto update_at column, and the +/// reserved internal epoch column added by the enterprise state schema/view. +fn persistence_sink_schema() -> Arc { + Arc::new(Schema::new(vec![ + ColumnSchema::new("state", CDT::binary_datatype(), true), + ColumnSchema::new("window", CDT::timestamp_millisecond_datatype(), false) + .with_time_index(true), + ColumnSchema::new("update_at", CDT::timestamp_millisecond_datatype(), true), + ColumnSchema::new( + crate::batching_mode::INTERNAL_FLOW_EPOCH_COL_NAME, + CDT::uint64_datatype(), + true, + ), + ])) +} + +/// A persistence-sink row: `(window ms, epoch, state)`. +type SinkRow = (Option, Option, Option>); + +/// Builds a record batch for `persistence_sink_schema`. +fn persistence_sink_recordbatch(rows: Vec) -> RecordBatch { + let schema = persistence_sink_schema(); + let mut states = datatypes::vectors::BinaryVectorBuilder::with_capacity(rows.len()); + let mut windows = + datatypes::vectors::TimestampMillisecondVectorBuilder::with_capacity(rows.len()); + let mut epochs = datatypes::vectors::UInt64VectorBuilder::with_capacity(rows.len()); + let mut update_at = + datatypes::vectors::TimestampMillisecondVectorBuilder::with_capacity(rows.len()); + for (window, epoch, state) in rows { + states.push(state.as_deref()); + windows.push(window.map(datatypes::timestamp::TimestampMillisecond::new)); + epochs.push(epoch); + update_at.push(Some(datatypes::timestamp::TimestampMillisecond::new(0))); + } + RecordBatch::new( + schema, + vec![ + states.to_vector(), + windows.to_vector(), + update_at.to_vector(), + epochs.to_vector(), + ], + ) + .unwrap() +} + +fn register_persistence_sink( + query_engine: &QueryEngineRef, + table_name: &str, + rows: Vec, + table_id: u32, +) { + let batch = persistence_sink_recordbatch(rows); + let table = MemTable::table(table_name, batch); + 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 sequence_range_batch_opts() -> Arc { + Arc::new(BatchingModeOptions { + experimental_enable_incremental_read: true, + incremental_mode: IncrementalMode::SequenceRange, + ..Default::default() + }) +} + +async fn new_sequence_range_test_task(sink_table: &str) -> TestTaskParts { + new_test_task_engine_and_plan_with_query_and_opts( + "SELECT number, ts FROM numbers_with_ts", + sink_table, + sequence_range_batch_opts(), + ) + .await +} + +/// Builds a task for an exact EE-like time-window aggregate query (a real +/// BINARY UDDSketch `state` output plus a `date_bin` window) with +/// `SequenceRange` batch options, ready for an EE-like sink registration. +async fn new_ee_sequence_range_task(sink_table: &str, query: &str) -> TestTaskParts { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), query, true) + .await + .unwrap(); + let (column_name, time_window_expr, _, df_schema) = find_time_window_expr( + &plan, + query_engine.engine_state().catalog_manager().clone(), + ctx.clone(), + ) + .await + .unwrap(); + let time_window_expr = time_window_expr.map(|expr| { + TimeWindowExpr::from_expr( + &expr, + &column_name, + &df_schema, + &query_engine.engine_state().session_state(), + ) + .unwrap() + }); + let (_tx, rx) = tokio::sync::oneshot::channel(); + let task = BatchingTask::try_new(TaskArgs { + flow_id: 1, + query, + plan: plan.clone(), + time_window_expr, + expire_after: None, + sink_table_name: [ + "greptime".to_string(), + "public".to_string(), + sink_table.to_string(), + ], + source_table_names: vec![[ + "greptime".to_string(), + "public".to_string(), + "numbers_with_ts".to_string(), + ]], + query_ctx: ctx, + catalog_manager: query_engine.engine_state().catalog_manager().clone(), + shutdown_rx: rx, + batch_opts: sequence_range_batch_opts(), + flow_eval_interval: None, + eval_schedule: None, + }) + .unwrap(); + TestTaskParts { + task, + query_engine, + plan, + } +} + +fn test_persistence() -> CheckpointPersistence { + CheckpointPersistence { + epoch_col_name: crate::batching_mode::INTERNAL_FLOW_EPOCH_COL_NAME.to_string(), + state_col_name: "state".to_string(), + window_col_name: "window".to_string(), + } +} + +#[tokio::test] +async fn test_detect_checkpoint_persistence_requires_sequence_range_and_epoch_column() { + let sink_table = "persistence_detect_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + // Sink with the epoch + BINARY state columns. + register_persistence_sink(&query_engine, sink_table, vec![], 9100); + + let persistence = task + .detect_checkpoint_persistence() + .await + .unwrap() + .expect("sequence range + epoch column should activate persistence"); + assert_eq!(test_persistence(), persistence); + + // MemtableOnly mode never activates persistence. + let sink_table = "persistence_detect_memtable_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_test_task_engine_and_plan_with_query_and_opts( + "SELECT number, ts FROM numbers_with_ts", + sink_table, + incremental_batch_opts(), + ) + .await; + register_persistence_sink(&query_engine, sink_table, vec![], 9101); + assert!( + task.detect_checkpoint_persistence() + .await + .unwrap() + .is_none(), + "MemtableOnly must not activate persistence" + ); + + // A sink without the reserved epoch column never activates persistence. + let sink_table = "persistence_detect_plain_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + register_auto_created_aggregate_sink(&query_engine, sink_table); + assert!( + task.detect_checkpoint_persistence() + .await + .unwrap() + .is_none(), + "sink without epoch column must not activate persistence" + ); +} + +/// Builds a DataFusion session state able to decode sink/source scans: the +/// engine's session state (which carries all flow functions) plus a catalog +/// list that resolves tables through the engine's catalog manager. The engine's +/// bare `session_state()` has an empty catalog list, so scans cannot be +/// resolved from it directly. +fn decode_session_state(query_engine: &QueryEngineRef) -> datafusion::execution::SessionState { + let catalog_list: Arc = + Arc::new(catalog::table_source::dummy_catalog::DummyCatalogList::new( + query_engine.engine_state().catalog_manager().clone(), + )); + datafusion::execution::SessionStateBuilder::new_from_existing( + query_engine.engine_state().session_state(), + ) + .with_catalog_list(catalog_list) + .build() +} + +/// A frontend handler that decodes and executes `Query::LogicalPlan` requests +/// against the test query engine (MemTable catalog) and records every request. +/// Used to prove that checkpoint restore scans travel through the frontend +/// transport instead of the local `QueryEngine::execute` path. +struct CaptureLogicalPlanHandler { + query_engine: QueryEngineRef, + captured: Arc>>, +} + +#[async_trait::async_trait] +impl GrpcQueryHandlerWithBoxedError for CaptureLogicalPlanHandler { + async fn do_query( + &self, + query: api::v1::greptime_request::Request, + ctx: QueryContextRef, + ) -> std::result::Result { + let api::v1::greptime_request::Request::Query(q) = &query else { + return Ok(Output::new_with_affected_rows(0)); + }; + self.captured.lock().unwrap().push(q.clone()); + let Some(api::v1::query_request::Query::LogicalPlan(bytes)) = &q.query else { + return Ok(Output::new_with_affected_rows(0)); + }; + let session_state = decode_session_state(&self.query_engine); + let plan = DFLogicalSubstraitConvertor {} + .decode(bytes::Bytes::from(bytes.clone()), session_state) + .await + .map_err(BoxedError::new)?; + let output = self + .query_engine + .execute(plan, ctx) + .await + .map_err(BoxedError::new)?; + Ok(output) + } +} + +/// Captured `Query::LogicalPlan` requests plus the handler handle. +type RestoreFrontendClient = ( + Arc, + Arc>>, + Arc, +); + +/// Builds a frontend client whose handler executes decoded LogicalPlan +/// requests against `query_engine`, plus the captured request log. The handler +/// `Arc` is returned too so the client's weak handle stays alive for the +/// whole test. +fn restore_frontend_client(query_engine: &QueryEngineRef) -> RestoreFrontendClient { + let captured = Arc::new(std::sync::Mutex::new(Vec::new())); + let handler: Arc = Arc::new(CaptureLogicalPlanHandler { + query_engine: query_engine.clone(), + captured: captured.clone(), + }); + let frontend_client = Arc::new(FrontendClient::from_grpc_handler( + Arc::downgrade(&handler), + QueryOptions::default(), + )); + (frontend_client, captured, handler) +} + +#[tokio::test] +async fn test_restore_via_frontend_sends_logical_plans_and_seeds_incremental() { + let sink_table = "persistence_restore_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + let record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 3, + checkpoints: BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]), + }; + let encoded = encode_checkpoint_record(&record).unwrap(); + register_persistence_sink( + &query_engine, + sink_table, + vec![ + (Some(1_000), Some(3), None), + (Some(2_000), Some(3), None), + ( + Some(crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS), + Some(3), + Some(encoded), + ), + ], + 9102, + ); + + let (frontend_client, captured, _handler) = restore_frontend_client(&query_engine); + task.try_enable_checkpoint_persistence(&frontend_client) + .await; + + // Both restore scans (sentinel row + non-sentinel epoch aggregation) were + // sent as LogicalPlan requests through the frontend client. + let requests = captured.lock().unwrap(); + assert_eq!(2, requests.len(), "expected two restore scan requests"); + for request in requests.iter() { + assert!( + matches!( + request.query, + Some(api::v1::query_request::Query::LogicalPlan(_)) + ), + "restore scans must be transported as Query::LogicalPlan, got {request:?}" + ); + } + drop(requests); + + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental); + assert_eq!( + state.checkpoints(), + &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]) + ); + assert_eq!(state.persisted_epoch(), 3); + assert_eq!(Some(&test_persistence()), state.checkpoint_persistence()); +} + +/// Runs one untrusted-restore scenario through the frontend transport and +/// asserts the task falls back to full snapshot with persistence still armed. +async fn assert_restore_falls_back(table_name: &str, table_id: u32, rows: Vec) { + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(table_name).await; + register_persistence_sink(&query_engine, table_name, rows, table_id); + + let (frontend_client, _captured, _handler) = restore_frontend_client(&query_engine); + task.try_enable_checkpoint_persistence(&frontend_client) + .await; + + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot); + assert_eq!(state.persisted_epoch(), 0); + assert!(state.checkpoints().is_empty()); + // Persistence stays armed so later cycles can persist fresh checkpoints. + assert!(state.checkpoint_persistence().is_some()); +} + +/// Consolidates the missing / corrupt / empty-map / multiple-sentinel / +/// newer-row-epoch / NULL-epoch restore fallbacks. Every case must leave the +/// task in full snapshot with no trusted checkpoint. +#[tokio::test] +async fn test_restore_falls_back_on_untrusted_records() { + let sentinel = Some(crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS); + let record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 3, + checkpoints: BTreeMap::from([(1_u64, 10_u64)]), + }; + let encoded = encode_checkpoint_record(&record).unwrap(); + + // Missing sentinel row. + assert_restore_falls_back( + "persistence_no_sentinel", + 9103, + vec![(Some(1_000), Some(2), None)], + ) + .await; + // Sentinel row with undecodable state bytes. + assert_restore_falls_back( + "persistence_corrupt_record", + 9104, + vec![(sentinel, Some(2), Some(b"garbage".to_vec()))], + ) + .await; + // Sentinel row holding a valid v1 record with an empty checkpoint map. + let empty_record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 2, + checkpoints: BTreeMap::new(), + }; + let empty_encoded = encode_checkpoint_record(&empty_record).unwrap(); + assert_restore_falls_back( + "persistence_empty_map", + 9105, + vec![(sentinel, Some(2), Some(empty_encoded))], + ) + .await; + // State rows stamped with epoch 5 are newer than the record's epoch 3: + // crash between state write and checkpoint write. + assert_restore_falls_back( + "persistence_newer_rows", + 9106, + vec![ + (Some(1_000), Some(5), None), + (sentinel, Some(3), Some(encoded.clone())), + ], + ) + .await; + // Pre-upgrade rows with NULL epochs are untrusted. + assert_restore_falls_back( + "persistence_null_epoch_rows", + 9107, + vec![ + (Some(1_000), None, None), + (sentinel, Some(3), Some(encoded.clone())), + ], + ) + .await; + // Two sentinel rows make the record ambiguous. + assert_restore_falls_back( + "persistence_multi_sentinel", + 9109, + vec![ + (sentinel, Some(3), Some(encoded.clone())), + (sentinel, Some(3), Some(encoded)), + ], + ) + .await; +} + +#[tokio::test] +async fn test_restore_accepts_record_without_state_data() { + // NULL epochs are acceptable when there is no non-sentinel state data. + let sink_table = "persistence_no_state_data"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + let record = CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 7, + checkpoints: BTreeMap::from([(1_u64, 10_u64)]), + }; + let encoded = encode_checkpoint_record(&record).unwrap(); + register_persistence_sink( + &query_engine, + sink_table, + vec![( + Some(crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS), + None, + Some(encoded), + )], + 9108, + ); + let (frontend_client, _captured, _handler) = restore_frontend_client(&query_engine); + task.try_enable_checkpoint_persistence(&frontend_client) + .await; + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental); + assert_eq!(state.checkpoints(), &BTreeMap::from([(1_u64, 10_u64)])); + assert_eq!(state.persisted_epoch(), 7); +} + +#[tokio::test] +async fn test_stamp_epoch_into_plan_is_noop_when_inactive() { + let TestTaskParts { task, plan, .. } = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "missing_sink", + ) + .await; + + let (stamped, epoch) = task.stamp_epoch_into_plan(plan.clone()).await.unwrap(); + assert_eq!(epoch, None); + assert_eq!( + stamped, plan, + "ordinary flow plan must be byte-for-byte unchanged" + ); +} + +#[tokio::test] +async fn test_stamp_epoch_into_plan_appends_epoch_literal() { + let sink_table = "auto_created_aggregate_sink"; + let query = "SELECT max(number) AS number, ts FROM numbers_with_ts GROUP BY ts"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + register_auto_created_aggregate_sink(&query_engine, sink_table); + task.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(test_persistence())); + + let ctx = task.state.read().unwrap().query_ctx.clone(); + let plan = sql_to_df_plan(ctx, query_engine.clone(), query, true) + .await + .unwrap(); + let (sink_table, _) = get_table_info_df_schema( + query_engine.engine_state().catalog_manager().clone(), + [ + "greptime".to_string(), + "public".to_string(), + sink_table.to_string(), + ], + ) + .await + .unwrap(); + let table_provider = Arc::new(DfTableProviderAdapter::new(sink_table)); + let table_source = Arc::new(DefaultTableSource::new(table_provider)); + let dml_plan = LogicalPlan::Dml(DmlStatement::new( + datafusion_common::TableReference::bare("test"), + table_source, + WriteOp::Insert(datafusion_expr::dml::InsertOp::Append), + Arc::new(plan), + )); + + let (stamped, epoch) = task.stamp_epoch_into_plan(dml_plan).await.unwrap(); + assert_eq!(epoch, Some(1), "first cycle stamps epoch 1"); + + let LogicalPlan::Dml(dml) = &stamped else { + panic!("expected DML plan"); + }; + let fields = dml.input.schema().fields(); + assert_eq!( + crate::batching_mode::INTERNAL_FLOW_EPOCH_COL_NAME, + fields.last().unwrap().name(), + "epoch column must be appended as the last output field" + ); + let exprs = dml.input.expressions(); + let last = exprs.last().expect("epoch projection expr"); + assert!( + format!("{last:?}").contains("UInt64(1)"), + "epoch column must be stamped with the current epoch literal, got {last:?}" + ); +} + +#[tokio::test] +async fn test_exact_ee_schema_plan_retains_binary_state_column() { + // Exact EE-like flow: the query itself produces the BINARY UDDSketch + // `state` column plus the `date_bin` window. The sink is + // `[state BINARY, window TIMESTAMP time-index, update_at TIMESTAMP, + // epoch integer]`. The real state column must survive schema matching, + // the incremental rewrite analysis, and the final stamped DML. + let sink_table = "persistence_ee_sink"; + let query = "SELECT uddsketch_state(128, 0.01, CAST(number AS DOUBLE)) AS state, \ + date_bin(INTERVAL '5 second', ts) AS window FROM greptime.public.numbers_with_ts GROUP BY window"; + let TestTaskParts { + task, query_engine, .. + } = new_ee_sequence_range_task(sink_table, query).await; + register_persistence_sink( + &query_engine, + sink_table, + vec![(Some(0), Some(1), None)], + 9110, + ); + task.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(test_persistence())); + task.state + .write() + .unwrap() + .dirty_time_windows + .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); + + // 1. Schema match: plan generation succeeds and the real BINARY `state` + // column is retained (only the epoch column is stripped from matching). + let info = task + .gen_insert_plan_unlocked(&query_engine, None) + .await + .unwrap_or_else(|err| panic!("EE-like schema must not break plan generation, got: {err:?}")) + .expect("dirty windows must produce a plan"); + let LogicalPlan::Dml(dml) = &info.plan else { + panic!("expected DML insert plan, got {:?}", info.plan); + }; + let matched_fields = dml.input.schema().fields(); + let matched_names = matched_fields + .iter() + .map(|field| field.name().clone()) + .collect::>(); + assert_eq!( + vec!["state", "window", "update_at"], + matched_names, + "schema matching must keep the real state column and append update_at; \ + the epoch column is stamped later" + ); + assert_eq!( + CDT::binary_datatype(), + CDT::from_arrow_type(matched_fields[0].data_type()), + "the state column must stay BINARY through schema matching" + ); + + // 2. Incremental rewrite: OSS cannot merge a raw `uddsketch_state` + // aggregate (the enterprise state view adds the merge op), so the + // rewrite honestly reports the EE-like plan as unsafe and disables + // incremental instead of running it as an unfiltered full snapshot. + // The retained BINARY state column must not crash the analyzer. + { + let mut state = task.state.write().unwrap(); + state.advance_checkpoints(HashMap::from([(1_u64, 10_u64)])); + } + let incremental = task + .prepare_plan_for_incremental(&info.plan) + .await + .unwrap_or_else(|err| { + panic!("incremental rewrite must not error on the EE-like plan: {err:?}") + }); + assert!( + incremental.is_none(), + "uddsketch_state is not an OSS-mergeable aggregate; the rewrite must refuse" + ); + assert!( + task.state.read().unwrap().is_incremental_disabled(), + "unsupported EE-like aggregate must permanently disable incremental" + ); + + // 3. Final stamped DML: the epoch literal is appended and the real state + // column keeps its exact name in the insert. + let (stamped, epoch) = task.stamp_epoch_into_plan(info.plan).await.unwrap(); + assert_eq!(Some(1), epoch, "first cycle stamps epoch 1"); + let LogicalPlan::Dml(stamped_dml) = &stamped else { + panic!("expected stamped DML plan"); + }; + let stamped_names = stamped_dml + .input + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + assert_eq!( + vec![ + "state", + "window", + "update_at", + crate::batching_mode::INTERNAL_FLOW_EPOCH_COL_NAME + ], + stamped_names, + "final DML must carry the real state column and the stamped epoch column" + ); +} + +/// A frontend handler that captures the insert request and returns success. +struct CaptureInsertHandler { + captured: Arc>>, +} + +#[async_trait::async_trait] +impl GrpcQueryHandlerWithBoxedError for CaptureInsertHandler { + async fn do_query( + &self, + query: api::v1::greptime_request::Request, + _ctx: QueryContextRef, + ) -> std::result::Result { + if let api::v1::greptime_request::Request::Query(q) = &query { + *self.captured.lock().unwrap() = Some(q.clone()); + } + Ok(Output::new_with_affected_rows(1)) + } +} + +#[tokio::test] +async fn test_write_checkpoint_row_sends_singleton_sentinel_row() { + let sink_table = "persistence_write_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + register_persistence_sink( + &query_engine, + sink_table, + vec![(Some(0), Some(1), None)], + 9111, + ); + task.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(test_persistence())); + + let captured = Arc::new(std::sync::Mutex::new(None)); + let handler: Arc = Arc::new(CaptureInsertHandler { + captured: captured.clone(), + }); + let frontend_client = Arc::new(FrontendClient::from_grpc_handler( + Arc::downgrade(&handler), + QueryOptions::default(), + )); + + let checkpoints = BTreeMap::from([(1_u64, 11_u64)]); + task.write_checkpoint_row(&frontend_client, 7, &checkpoints) + .await + .expect("checkpoint row write should succeed"); + + let captured = captured + .lock() + .unwrap() + .clone() + .expect("handler captured the request"); + let api::v1::query_request::Query::InsertIntoPlan(insert) = + captured.query.expect("insert plan") + else { + panic!("expected InsertIntoPlan"); + }; + assert_eq!( + "persistence_write_sink", + insert.table_name.as_ref().unwrap().table_name + ); + + let session_state = query_engine.engine_state().session_state(); + let plan = DFLogicalSubstraitConvertor {} + .decode(bytes::Bytes::from(insert.logical_plan), session_state) + .await + .unwrap(); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected projection over empty relation, got {plan:?}"); + }; + // window + epoch + state + auto update_at. + assert_eq!(4, projection.expr.len()); + + // window column -> sentinel timestamp + let window_expr = &projection.expr[0]; + let window_sql = format!("{window_expr:?}"); + assert!( + window_sql + .contains(&crate::batching_mode::CHECKPOINT_SENTINEL_WINDOW_TS_MILLIS.to_string()), + "sentinel window literal expected, got {window_sql}" + ); + // epoch column -> 7 + let epoch_expr = &projection.expr[1]; + assert!( + format!("{epoch_expr:?}").contains("UInt64(7)"), + "epoch literal expected, got {epoch_expr:?}" + ); + // state column -> the encoded v1 record bytes must appear verbatim. + let expected = encode_checkpoint_record(&CheckpointRecord { + format_version: CHECKPOINT_RECORD_FORMAT_VERSION, + epoch: 7, + checkpoints, + }) + .unwrap(); + let decoded = decode_checkpoint_record(&expected).unwrap().unwrap(); + assert_eq!(7, decoded.epoch); + assert_eq!(BTreeMap::from([(1_u64, 11_u64)]), decoded.checkpoints); + let state_bytes = match &projection.expr[2] { + Expr::Alias(alias) => match alias.expr.as_ref() { + Expr::Literal(ScalarValue::Binary(Some(bytes)), _) => bytes, + other => panic!("expected binary literal for state column, got {other:?}"), + }, + other => panic!("expected alias for state column, got {other:?}"), + }; + assert_eq!( + &expected, state_bytes, + "checkpoint record bytes must be stored verbatim" + ); +} + +/// A frontend handler that always fails. +struct FailInsertHandler; + +#[async_trait::async_trait] +impl GrpcQueryHandlerWithBoxedError for FailInsertHandler { + async fn do_query( + &self, + _query: api::v1::greptime_request::Request, + _ctx: QueryContextRef, + ) -> std::result::Result { + Err(BoxedError::new(MockError::new(StatusCode::Internal))) + } +} + +#[tokio::test] +async fn test_checkpoint_row_write_failure_is_reported() { + let sink_table = "persistence_write_fail_sink"; + let TestTaskParts { + task, query_engine, .. + } = new_sequence_range_test_task(sink_table).await; + register_persistence_sink( + &query_engine, + sink_table, + vec![(Some(0), Some(1), None)], + 9112, + ); + task.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(test_persistence())); + + let handler: Arc = Arc::new(FailInsertHandler); + let frontend_client = Arc::new(FrontendClient::from_grpc_handler( + Arc::downgrade(&handler), + QueryOptions::default(), + )); + + let err = task + .write_checkpoint_row(&frontend_client, 7, &BTreeMap::from([(1_u64, 11_u64)])) + .await + .unwrap_err(); + assert!(matches!(err, Error::External { .. }), "{err}"); +} + +/// A minimal record-batch stream that reports region watermarks through its +/// terminal metrics handle immediately (it produces no rows). Lets the flow +/// execution path see a full watermark proof without consuming a real query. +struct WatermarkOnlyStream { + schema: datatypes::schema::SchemaRef, + watermarks: Vec<(u64, Option)>, +} + +impl futures::Stream for WatermarkOnlyStream { + type Item = common_recordbatch::error::Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(None) + } + + fn size_hint(&self) -> (usize, Option) { + (0, Some(0)) + } +} + +impl common_recordbatch::RecordBatchStream for WatermarkOnlyStream { + fn name(&self) -> &str { + "WatermarkOnlyStream" + } + + fn schema(&self) -> datatypes::schema::SchemaRef { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + Some(RecordBatchMetrics { + region_watermarks: self + .watermarks + .iter() + .map(|(region_id, watermark)| RegionWatermarkEntry { + region_id: *region_id, + watermark: *watermark, + }) + .collect(), + ..Default::default() + }) + } +} + +/// Walks a decoded logical plan looking for the checkpoint row's +/// `EmptyRelation` root (state-row inserts are scan-based and never contain +/// one). +fn contains_empty_relation(plan: &LogicalPlan) -> bool { + let mut found = false; + let _ = plan.apply(|node| { + if matches!(node, LogicalPlan::EmptyRelation(_)) { + found = true; + } + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }); + found +} + +/// A frontend handler for the checkpoint persist-failure cycle: state-row +/// inserts return terminal watermarks; the first checkpoint-row write fails, +/// later checkpoint writes succeed. +struct PersistFailOnceHandler { + query_engine: QueryEngineRef, + checkpoint_write_calls: std::sync::atomic::AtomicUsize, + checkpoint_write_attempts: Arc>, +} + +#[async_trait::async_trait] +impl GrpcQueryHandlerWithBoxedError for PersistFailOnceHandler { + 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 { + return Ok(Output::new_with_affected_rows(0)); + }; + // Best-effort decode to classify the request. The singleton + // checkpoint-row write is a Projection over EmptyRelation and always + // decodes; the state-row insert carries flow-only UDAFs (e.g. + // `uddsketch_state`) that the test session state cannot resolve from + // substrait anchors, so a decode failure is treated as a state row. + let convertor = DFLogicalSubstraitConvertor {}; + let is_checkpoint_row = match convertor + .decode( + bytes::Bytes::from(insert.logical_plan), + decode_session_state(&self.query_engine), + ) + .await + { + Ok(plan) => contains_empty_relation(&plan), + Err(_) => false, + }; + if is_checkpoint_row { + // Singleton checkpoint row write. + *self.checkpoint_write_attempts.lock().unwrap() += 1; + let call = self + .checkpoint_write_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if call == 0 { + return Err(BoxedError::new(MockError::new(StatusCode::Internal))); + } + return Ok(Output::new_with_affected_rows(1)); + } + // State-row insert: prove full coverage with terminal watermarks. + Ok(Output::new_with_stream(Box::pin(WatermarkOnlyStream { + schema: Arc::new(Schema::new(vec![])), + watermarks: vec![(1, Some(10)), (2, Some(20))], + }))) + } +} + +#[tokio::test] +async fn test_checkpoint_persist_failure_schedules_backfill_and_replacement_checkpoint() { + let sink_table = "persistence_exec_sink"; + let query = "SELECT uddsketch_state(128, 0.01, CAST(number AS DOUBLE)) AS state, \ + date_bin(INTERVAL '5 second', ts) AS window FROM greptime.public.numbers_with_ts GROUP BY window"; + let TestTaskParts { + task, query_engine, .. + } = new_ee_sequence_range_task(sink_table, query).await; + register_persistence_sink( + &query_engine, + sink_table, + vec![(Some(0), Some(1), None)], + 9113, + ); + task.state + .write() + .unwrap() + .set_checkpoint_persistence(Some(test_persistence())); + task.state + .write() + .unwrap() + .dirty_time_windows + .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); + + let checkpoint_write_attempts = Arc::new(std::sync::Mutex::new(0)); + let handler: Arc = Arc::new(PersistFailOnceHandler { + query_engine: query_engine.clone(), + checkpoint_write_calls: std::sync::atomic::AtomicUsize::new(0), + checkpoint_write_attempts: checkpoint_write_attempts.clone(), + }); + let frontend_client = Arc::new(FrontendClient::from_grpc_handler( + Arc::downgrade(&handler), + QueryOptions::default(), + )); + + // Cycle 1: the state insert advances checkpoints, but the checkpoint row + // write fails. The task must fall back to full snapshot, keep the consumed + // dirty work pending, and not advance the durable epoch. + let outcome = task + .execute_once_serialized(&query_engine, &frontend_client, None) + .await; + assert!( + outcome.is_ok(), + "state insert should succeed, got: {outcome:?}" + ); + { + let state = task.state.read().unwrap(); + assert_eq!( + state.checkpoint_mode(), + CheckpointMode::FullSnapshot, + "a failed checkpoint write must reset to full snapshot" + ); + assert_eq!( + state.persisted_epoch(), + 0, + "a failed checkpoint write must not advance the durable epoch" + ); + assert!( + !state.dirty_time_windows.is_empty(), + "the executed plan's consumed dirty work must be restored" + ); + } + + // Cycle 2: the restored dirty work drives a full repair/backfill; the + // replacement checkpoint row write succeeds and the durable epoch advances. + let outcome = task + .execute_once_serialized(&query_engine, &frontend_client, None) + .await; + assert!( + outcome.is_ok(), + "repair cycle should succeed, got: {outcome:?}" + ); + { + let state = task.state.read().unwrap(); + assert_eq!( + state.checkpoint_mode(), + CheckpointMode::Incremental, + "the replacement checkpoint write must restore incremental mode" + ); + assert_eq!( + state.persisted_epoch(), + 1, + "the replacement checkpoint must advance the durable epoch" + ); + assert_eq!( + state.checkpoints(), + &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]) + ); + } + assert_eq!( + 2, + *checkpoint_write_attempts.lock().unwrap(), + "one failed checkpoint write followed by one replacement write" + ); +}