From d910183cc46983074d8dfbb46166568ef819cb77 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 23 Sep 2026 09:21:40 +0000 Subject: [PATCH] feat(flow): freeze recovery windows and retention bounds for incremental flows (#9312) * feat(flow): discover recovery windows from exact sequence reads Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(flow): freeze explicit recovery windows and retention bounds Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): validate retained source windows before recovery Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): qualify recovery source tables before serialization Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/flow/Cargo.toml | 1 + src/flow/src/batching_mode/state.rs | 139 +++- src/flow/src/batching_mode/task.rs | 437 ++++++++++++- src/flow/src/batching_mode/task/test.rs | 820 +++++++++++++++++++++++- 5 files changed, 1376 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6c7916d8194..3ab47005352 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5468,6 +5468,7 @@ dependencies = [ "operator", "pretty_assertions", "prometheus 0.14.0", + "prost 0.14.1", "query", "rand 0.9.4", "serde", diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index d0ac5eee9ac..fea7b1eeda2 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -65,6 +65,7 @@ catalog = { workspace = true, features = ["testing"] } common-catalog.workspace = true common-meta = { workspace = true, features = ["testing"] } pretty_assertions.workspace = true +prost.workspace = true query.workspace = true session.workspace = true table.workspace = true diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index b9710a25a0d..55307dbcdec 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -208,6 +208,25 @@ impl TaskState { self.pending_fenced_repair.as_ref() } + /// Start repairing explicit bounded windows under a frozen high `H`. + /// + /// The caller supplies the complete repair scope. Live dirty windows are + /// left unchanged so signals received after `H` remain separate. + pub fn start_fenced_repair_windows( + &mut self, + high: BTreeMap, + windows: Vec<(Timestamp, Timestamp)>, + ) { + let mut pending_windows = self.dirty_time_windows.clone(); + pending_windows.clean(); + pending_windows.add_windows(windows); + self.pending_fenced_repair = Some(FencedRepair { + high, + pending_windows, + }); + self.checkpoint_mode = CheckpointMode::FullSnapshot; + } + /// Finish the fenced repair and promote the frozen high watermark to the /// checkpoint map. Incremental-disabled flows stay in FullSnapshot mode. pub fn finish_fenced_repair(&mut self) -> Option> { @@ -259,22 +278,18 @@ impl TaskState { task_ctx: Option<&BatchingTask>, ) -> Result, Error> { if let Some(repair) = self.pending_fenced_repair.as_mut() { - let expr = repair.pending_windows.gen_filter_exprs( + // Fenced windows are an explicit frozen repair scope. They must not + // be pruned by the moving live-data expiration boundary, and an + // empty repair must remain active until its high watermark is + // explicitly finished. + return repair.pending_windows.gen_filter_exprs( col_name, - expire_lower_bound, + None, window_size, window_cnt, flow_id, task_ctx, - )?; - if expr.is_some() || !repair.pending_windows.is_empty() { - return Ok(expr); - } - - // All pending repair windows may have expired during merge. Clear - // the empty repair so this call can fall back to live dirty windows - // instead of routing future executions to an empty queue forever. - self.pending_fenced_repair = None; + ); } self.dirty_time_windows.gen_filter_exprs( @@ -1393,6 +1408,108 @@ mod test { assert_eq!(state.dirty_time_windows.len(), 2); } + #[test] + fn test_explicit_fenced_repair_windows_keep_live_windows_separate() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows( + high.clone(), + vec![(Timestamp::new_second(10), Timestamp::new_second(15))], + ); + + assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 1); + + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(10), Timestamp::new_second(15))] + ); + state + .dirty_time_windows + .add_window(Timestamp::new_second(20), Some(Timestamp::new_second(25))); + state.restore_scoped_windows(&filter); + + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 2); + } + + #[test] + fn test_explicit_fenced_repair_keeps_empty_high_and_old_windows() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows(high.clone(), Vec::new()); + + assert!( + state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .is_none() + ); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!(state.dirty_time_windows.len(), 1); + assert_eq!(state.finish_fenced_repair(), Some(high)); + + state.start_fenced_repair_windows( + BTreeMap::from([(1, 11)]), + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))], + ); + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))] + ); + } + #[test] fn test_disable_incremental_persists_full_snapshot_mode() { let query_ctx = QueryContext::arc(); diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 346d39d1d46..40020f639c1 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -18,21 +18,29 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use api::v1::{CreateTableExpr, TableName}; use catalog::CatalogManagerRef; +use catalog::kvbackend::KvBackendCatalogManager; use client::OutputWithMetrics; use common_error::ext::BoxedError; +use common_meta::key::schema_name::SchemaNameKey; +use common_query::OutputData; use common_query::logical_plan::breakup_insert_plan; use common_telemetry::tracing::warn; use common_telemetry::{debug, info}; -use common_time::Timestamp; +use common_time::{TimeToLive, Timestamp}; 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, ScalarValue, TableReference}; -use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; +use datafusion_expr::{DmlStatement, LogicalPlan, Projection, WriteOp, col, lit}; use datatypes::schema::Schema; +use datatypes::vectors::Helper; +use futures::TryStreamExt; use query::QueryEngineRef; -use query::options::FLOW_INCREMENTAL_MODE; +use query::options::{ + FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE, + FLOW_RETURN_REGION_SEQ, +}; use query::query_engine::DefaultSerializer; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; @@ -55,7 +63,7 @@ use crate::batching_mode::state::{ 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, + AddFilterRewriter, ColumnMatcherRewriter, analyze_incremental_aggregate_plan, 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; @@ -192,6 +200,56 @@ fn encode_insert_plan_request( }) } +fn recovery_aggregate_input(plan: &LogicalPlan) -> Result { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + _ => plan, + }; + let LogicalPlan::Aggregate(aggregate) = plan else { + return UnexpectedSnafu { + reason: "Recovery timestamp projection did not find an aggregate".to_string(), + } + .fail(); + }; + Ok(aggregate.input.as_ref().clone()) +} + +fn capture_recovery_batch_windows( + batch: &common_recordbatch::RecordBatch, + time_window_expr: &TimeWindowExpr, + windows: &mut BTreeSet<(Timestamp, Timestamp)>, +) -> Result<(), Error> { + if batch.num_columns() != 1 { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned {} columns instead of one", + batch.num_columns() + ), + } + .fail(); + } + let values = Helper::try_into_vector(batch.column(0).clone()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + for index in 0..values.len() { + let timestamp = values.get(index).as_timestamp().context(UnexpectedSnafu { + reason: "Recovery timestamp projection returned a null or non-timestamp value" + .to_string(), + })?; + let (start, end) = time_window_expr.eval(timestamp)?; + let window = ( + start.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no lower bound".to_string(), + })?, + end.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no upper bound".to_string(), + })?, + ); + windows.insert(window); + } + Ok(()) +} + fn format_insert_target_columns(plan: &LogicalPlan) -> String { plan.schema() .fields() @@ -483,6 +541,377 @@ impl BatchingTask { Arc::new(query_ctx.fork()) } + /// Returns the retention lower bound aligned to this task's time window. + pub fn recovery_retention_lower_bound(&self) -> Result, Error> { + let Some(expire_after) = self.config.expire_after else { + return Ok(None); + }; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after}", + self.config.flow_id + ), + } + .build() + })?; + let now = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to read recovery wall clock: {err}"), + } + .build() + })? + .as_secs() as i64, + ); + let lower = now + .sub_duration(Duration::from_secs(expire_after)) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: "Recovery expiry requires a time-window expression".to_string(), + })? + .eval(lower)? + .0 + .context(UnexpectedSnafu { + reason: "Recovery expiry time-window expression returned no lower bound" + .to_string(), + }) + .map(Some) + } + + /// Verifies that recovery can still read all source data required by its retained scope. + pub async fn validate_recovery_retention( + &self, + retention_lower: Option, + windows: &[(Timestamp, Timestamp)], + ) -> Result<(), Error> { + for name in &self.config.source_table_names { + let table = self + .config + .catalog_manager + .table(&name[0], &name[1], &name[2], None) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} is unavailable for recovery retention validation", + self.config.flow_id, + name.join(".") + ), + })?; + let ttl = if let Some(ttl) = table.table_info().meta.options.ttl { + ttl + } else { + let manager = self + .config + .catalog_manager + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} cannot resolve inherited TTL for source table {} during recovery", + self.config.flow_id, + name.join(".") + ), + })?; + manager + .table_metadata_manager_ref() + .schema_manager() + .get(SchemaNameKey::new(&name[0], &name[1])) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} schema {}.{} is unavailable for recovery retention validation", + self.config.flow_id, name[0], name[1] + ), + })? + .ttl + .map(Into::into) + .unwrap_or(TimeToLive::Forever) + }; + + match ttl { + TimeToLive::Forever => {} + TimeToLive::Instant => { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has instant TTL and cannot be recovered", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + TimeToLive::Duration(ttl) => { + let expire_after = self.config.expire_after.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has TTL {ttl:?}, but recovery has no expire_after", + self.config.flow_id, + name.join(".") + ), + })?; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after} during recovery retention validation", + self.config.flow_id + ), + } + .build() + })?; + if ttl <= Duration::from_secs(expire_after) { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} TTL {ttl:?} must exceed expire_after {expire_after}s for recovery", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + let mut oldest = retention_lower.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has finite TTL {ttl:?}, but recovery has no retained lower bound", + self.config.flow_id, + name.join(".") + ), + })?; + for (start, _) in windows { + oldest = oldest.min(*start); + } + let cutoff = Timestamp::current_millis() + .sub_duration(ttl) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if oldest <= cutoff { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery requires source table {} data from {oldest:?}, older than TTL {ttl:?} cutoff {cutoff:?}", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + } + } + } + Ok(()) + } + + /// Discovers the source time windows touched by the exact sequence range `(C, H]`. + /// + /// This compatibility wrapper captures the retention bound before discovery. + pub async fn capture_recovery_windows( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + let retention_lower = self.recovery_retention_lower_bound()?; + self.capture_recovery_windows_since(engine, frontend_client, lower, retention_lower) + .await + } + + /// Discovers windows under `lower` using a caller-frozen retention bound. + /// + /// The caller owns the execution guard and freezes this bound with its recovery scope. + /// This read-only method leaves task state untouched. Terminal proof must cover every + /// region in `lower`; subset or pruned proofs are rejected. + pub async fn capture_recovery_windows_since( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + retention_lower: Option, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + if lower.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires nonempty lower sequence bounds", + self.config.flow_id + ), + } + .fail(); + } + if !self.sequence_range_capable().await? { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires sequence_range-capable sources", + self.config.flow_id + ), + } + .fail(); + } + + let query_ctx = self.query_context_snapshot(); + let plan = sql_to_df_plan(query_ctx, engine.clone(), &self.config.query, false).await?; + let input = recovery_aggregate_input(&plan)?; + let Some(analysis) = analyze_incremental_aggregate_plan(&plan)? else { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a supported aggregate input", + self.config.flow_id + ), + } + .fail(); + }; + if !analysis.unsupported_exprs.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture has unsupported aggregate expressions: {:?}", + self.config.flow_id, analysis.unsupported_exprs + ), + } + .fail(); + } + let time_window_expr = self + .config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a time-window expression", + self.config.flow_id + ), + })?; + let input = if let Some(retention_lower) = retention_lower { + let mut add_filter = AddFilterRewriter::new( + col(&time_window_expr.column_name).gt_eq(lit(to_df_literal(retention_lower)?)), + ); + input + .rewrite(&mut add_filter) + .with_context(|_| DatafusionSnafu { + context: "Failed to apply recovery expire_after filter".to_string(), + })? + .data + } else { + input + }; + let timestamp_plan = LogicalPlan::Projection( + Projection::try_new(vec![col(&time_window_expr.column_name)], Arc::new(input)) + .context(DatafusionSnafu { + context: "Failed to project recovery source timestamps".to_string(), + })?, + ); + let catalog = &self.config.sink_table_name[0]; + let schema = &self.config.sink_table_name[1]; + let timestamp_plan = timestamp_plan + .clone() + .transform_down_with_subqueries(|p| { + if let LogicalPlan::TableScan(mut table_scan) = p { + let resolved = table_scan.table_name.resolve(catalog, schema); + table_scan.table_name = resolved.into(); + Ok(Transformed::yes(LogicalPlan::TableScan(table_scan))) + } else { + Ok(Transformed::no(p)) + } + }) + .with_context(|_| DatafusionSnafu { + context: format!( + "Failed to fix table ref in recovery timestamp plan, plan={:?}", + timestamp_plan + ), + })? + .data; + let message = DFLogicalSubstraitConvertor {} + .encode(×tamp_plan, DefaultSerializer) + .context(SubstraitEncodeLogicalPlanSnafu)?; + let request = api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())), + }; + let lower_json = serde_json::to_string(lower).map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to serialize recovery lower sequence bounds: {err}"), + } + .build() + })?; + let extensions = [ + (FLOW_RETURN_REGION_SEQ, "true"), + (FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE), + (FLOW_INCREMENTAL_AFTER_SEQS, lower_json.as_str()), + ]; + let mut peer_desc = None; + let result = frontend_client + .query_with_terminal_metrics( + catalog, + schema, + request, + &extensions, + &HashMap::new(), + &mut peer_desc, + ) + .await?; + let mut aligned_windows = BTreeSet::new(); + let metrics = result.metrics.clone(); + match result.output.data { + OutputData::AffectedRows(_) => { + return UnexpectedSnafu { + reason: "Recovery timestamp projection unexpectedly returned affected rows" + .to_string(), + } + .fail(); + } + OutputData::RecordBatches(batches) => { + for batch in batches.iter() { + capture_recovery_batch_windows(batch, time_window_expr, &mut aligned_windows)?; + } + } + OutputData::Stream(mut stream) => { + while let Some(batch) = stream + .try_next() + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + { + capture_recovery_batch_windows(&batch, time_window_expr, &mut aligned_windows)?; + } + } + } + if !metrics.is_ready() { + return UnexpectedSnafu { + reason: "Recovery timestamp projection ended without terminal metrics".to_string(), + } + .fail(); + } + let participating = metrics.participating_regions().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no participating-region proof".to_string(), + })?; + let high = metrics.region_watermark_map().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no terminal watermark proof".to_string(), + })?; + if participating.len() != lower.len() + || high.len() != lower.len() + || participating + .iter() + .any(|region| !lower.contains_key(region)) + || lower + .iter() + .any(|(region, low)| high.get(region).is_none_or(|watermark| watermark < low)) + || high.keys().any(|region| !participating.contains(region)) + { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned incomplete or regressing terminal proof: lower={lower:?}, participating={participating:?}, high={high:?}" + ), + } + .fail(); + } + Ok(( + high.into_iter().collect(), + aligned_windows.into_iter().collect(), + )) + } + /// Validates that the sink table schema can accept this flow's ordinary output. pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> { self.validate_sink_table_schema_with_values(engine, &BTreeMap::new()) diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 7add59664a5..29a15600f15 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::task::Poll; +use std::task::{Context, Poll}; +use std::time::Duration; use catalog::memory::MemoryCatalogManager; use catalog::{DeregisterTableRequest, RegisterTableRequest}; @@ -27,18 +29,21 @@ use common_error::status_code::StatusCode; use common_query::Output; use common_recordbatch::RecordBatch; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; +use common_time::{TimeToLive, Timestamp}; use datatypes::data_type::ConcreteDataType as CDT; -use datatypes::schema::ColumnSchema; +use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{ TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, }; use pretty_assertions::assert_eq; +use prost::Message; use query::options::{ FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY, FLOW_SCHEDULED_TIME_MILLIS, FLOW_SINK_TABLE_ID, QueryOptions, }; use session::context::QueryContext; use snafu::ResultExt; +use substrait::substrait_proto_df::proto::Plan; use table::Table; use table::metadata::FilterPushDownType; use table::test_util::MemTable; @@ -815,6 +820,47 @@ fn register_auto_created_aggregate_sink(query_engine: &QueryEngineRef, table_nam memory_catalog.register_table_sync(request).unwrap(); } +async fn configure_source_ttl(query_engine: &QueryEngineRef, ttl: Option) { + let catalog_manager = query_engine.engine_state().catalog_manager(); + let source = catalog_manager + .table( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "numbers_with_ts", + None, + ) + .await + .unwrap() + .unwrap(); + let mut info = (*source.table_info()).clone(); + info.meta.options.ttl = ttl; + let source = Arc::new(Table::new( + Arc::new(info), + FilterPushDownType::Unsupported, + source.data_source(), + )); + let memory_catalog = catalog_manager + .as_any() + .downcast_ref::() + .unwrap(); + memory_catalog + .deregister_table_sync(DeregisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + }) + .unwrap(); + memory_catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + table_id: source.table_info().table_id(), + table: source, + }) + .unwrap(); +} + async fn configure_source_capability( query_engine: &QueryEngineRef, engine: &str, @@ -871,6 +917,127 @@ async fn configure_source_capability( .unwrap(); } +#[tokio::test] +async fn test_validate_recovery_retention_explicit_source_ttls() { + struct Case { + name: &'static str, + ttl: Option, + expire_after: Option, + retention_age: Option, + window_age: Option, + should_pass: bool, + } + + let cases = [ + Case { + name: "finite TTL exceeding expire_after is admissible", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: true, + }, + Case { + name: "TTL equal to expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "TTL shorter than expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(59 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "aligned window older than TTL is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: Some(Duration::from_secs(3 * 60 * 60)), + should_pass: false, + }, + Case { + name: "old retention cutoff is rejected even without windows", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(3 * 60 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "explicit forever is admissible", + ttl: Some(TimeToLive::Forever), + expire_after: None, + retention_age: None, + window_age: None, + should_pass: true, + }, + Case { + name: "instant TTL is rejected", + ttl: Some(TimeToLive::Instant), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + ]; + + for case in cases { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "sink", + ) + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = case.expire_after; + configure_source_ttl(&query_engine, case.ttl).await; + + let now = Timestamp::current_millis(); + let retention_lower = case.retention_age.map(|age| now.sub_duration(age).unwrap()); + let windows = case + .window_age + .map(|age| { + let start = now.sub_duration(age).unwrap(); + vec![(start, now)] + }) + .unwrap_or_default(); + assert_eq!( + task.validate_recovery_retention(retention_lower, &windows) + .await + .is_ok(), + case.should_pass, + "{}", + case.name + ); + } +} + +#[tokio::test] +async fn test_validate_recovery_retention_rejects_unknown_inherited_ttl() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink") + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = Some(60 * 60); + configure_source_ttl(&query_engine, None).await; + + assert!( + task.validate_recovery_retention(Some(Timestamp::current_millis()), &[],) + .await + .is_err() + ); +} + fn dirty_marker() -> DirtyTimeWindows { let mut dirty = DirtyTimeWindows::default(); dirty.set_dirty(); @@ -1018,6 +1185,313 @@ impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError } } +struct RecoveryCaptureHandler { + output: std::sync::Mutex>, + lower: String, + expire_after: Option, + expiry_lower: std::sync::Mutex>, +} + +struct RecoveryMetricsStream { + schema: Arc, + batches: Vec, + metrics: Option, + fail: bool, +} + +impl futures::Stream for RecoveryMetricsStream { + type Item = common_recordbatch::error::Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + if self.fail { + self.fail = false; + return Poll::Ready(Some(Err(common_recordbatch::error::Error::External { + source: BoxedError::new(MockError::new(StatusCode::Unexpected)), + location: snafu::Location::new(file!(), line!(), column!()), + }))); + } + Poll::Ready(self.batches.pop().map(Ok)) + } +} + +impl common_recordbatch::RecordBatchStream for RecoveryMetricsStream { + fn name(&self) -> &str { + "RecoveryMetricsStream" + } + + fn schema(&self) -> Arc { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + self.batches + .is_empty() + .then(|| self.metrics.clone()) + .flatten() + } +} + +fn field_selection(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::field_reference::ReferenceType; + use substrait::substrait_proto_df::proto::expression::reference_segment::ReferenceType as Segment; + + let RexType::Selection(field) = expr.rex_type.as_ref()? else { + return None; + }; + let ReferenceType::DirectReference(reference) = field.reference_type.as_ref()? else { + return None; + }; + let Segment::StructField(field) = reference.reference_type.as_ref()? else { + return None; + }; + Some(field.field) +} + +fn timestamp_literal(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + + let RexType::Literal(literal) = expr.rex_type.as_ref()? else { + return None; + }; + let LiteralType::PrecisionTimestamp(timestamp) = literal.literal_type.as_ref()? else { + return None; + }; + match timestamp.precision { + 0 => Some(Timestamp::new_second(timestamp.value)), + 3 => Some(Timestamp::new_millisecond(timestamp.value)), + 6 => Some(Timestamp::new_microsecond(timestamp.value)), + 9 => Some(Timestamp::new_nanosecond(timestamp.value)), + _ => None, + } +} + +fn assert_recovery_timestamp_plan(plan: &Plan, expire_after: Option) -> Option { + use substrait::substrait_proto_df::proto::plan_rel::RelType as PlanRelType; + use substrait::substrait_proto_df::proto::rel::RelType; + + let function_names = plan + .extensions + .iter() + .filter_map(|extension| match extension.mapping_type.as_ref()? { + substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::MappingType::ExtensionFunction(function) => { + Some((function.function_anchor, function.name.as_str())) + } + _ => None, + }) + .collect::>(); + let [root] = plan.relations.as_slice() else { + panic!("recovery capture must encode exactly one relation"); + }; + let Some(PlanRelType::Root(root)) = root.rel_type.as_ref() else { + panic!("recovery capture must encode a root relation"); + }; + assert_eq!(root.names, ["ts"]); + let Some(RelType::Project(project)) = root.input.as_ref().and_then(|rel| rel.rel_type.as_ref()) + else { + panic!("recovery capture must project the raw timestamp"); + }; + assert_eq!(project.expressions.len(), 1); + assert_eq!(field_selection(&project.expressions[0]), Some(1)); + + fn visit( + rel: &substrait::substrait_proto_df::proto::Rel, + function_names: &HashMap, + scans: &mut Vec>, + has_number_filter: &mut bool, + expiry_lower: &mut Option, + ) { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + use substrait::substrait_proto_df::proto::function_argument::ArgType; + use substrait::substrait_proto_df::proto::read_rel::ReadType; + use substrait::substrait_proto_df::proto::rel::RelType; + + match rel.rel_type.as_ref().expect("recovery relation type") { + RelType::Aggregate(_) | RelType::Write(_) | RelType::Ddl(_) | RelType::Update(_) => { + panic!("recovery capture must not aggregate or write") + } + RelType::Read(read) => match read.read_type.as_ref() { + Some(ReadType::NamedTable(table)) => scans.push(table.names.clone()), + _ => panic!("recovery capture must read a named source table"), + }, + RelType::Filter(filter) => { + let condition = filter + .condition + .as_ref() + .expect("recovery filter condition"); + let RexType::ScalarFunction(function) = condition + .rex_type + .as_ref() + .expect("recovery filter expression") + else { + panic!("recovery filter must be a scalar function"); + }; + let args = function + .arguments + .iter() + .map(|arg| match arg.arg_type.as_ref() { + Some(ArgType::Value(expr)) => expr, + _ => panic!("recovery filter must use value arguments"), + }) + .collect::>(); + match function_names.get(&function.function_reference) { + Some(&"equal") + if args.len() == 2 + && field_selection(args[0]) == Some(0) + && matches!( + args[1].rex_type.as_ref(), + Some(RexType::Literal(literal)) + if matches!(literal.literal_type, Some(LiteralType::I64(42))) + ) => + { + *has_number_filter = true + } + Some(&"gte") if args.len() == 2 && field_selection(args[0]) == Some(1) => { + *expiry_lower = timestamp_literal(args[1]); + } + _ => panic!("unexpected recovery filter"), + } + visit( + filter.input.as_ref().expect("recovery filter input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ); + } + RelType::Project(project) => visit( + project.input.as_ref().expect("recovery project input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ), + _ => panic!("unexpected recovery relation"), + } + } + + let mut scans = Vec::new(); + let mut has_number_filter = false; + let mut expiry_lower = None; + visit( + project.input.as_ref().expect("recovery projection input"), + &function_names, + &mut scans, + &mut has_number_filter, + &mut expiry_lower, + ); + assert_eq!( + scans, + vec![vec![ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + "numbers_with_ts".to_string(), + ]] + ); + assert!(has_number_filter, "recovery capture lost WHERE number = 42"); + assert_eq!(expiry_lower.is_some(), expire_after.is_some()); + expiry_lower +} + +#[async_trait::async_trait] +impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError + for RecoveryCaptureHandler +{ + async fn do_query( + &self, + query: api::v1::greptime_request::Request, + ctx: QueryContextRef, + ) -> std::result::Result { + assert_eq!(ctx.extension(FLOW_INCREMENTAL_MODE), Some("sequence_range")); + assert_eq!( + ctx.extension(FLOW_INCREMENTAL_AFTER_SEQS), + Some(self.lower.as_str()) + ); + assert_eq!( + ctx.extension(query::options::FLOW_RETURN_REGION_SEQ), + Some("true") + ); + let api::v1::greptime_request::Request::Query(request) = query else { + panic!("recovery capture must issue a query request"); + }; + let Some(api::v1::query_request::Query::LogicalPlan(plan)) = request.query else { + panic!("recovery capture must issue a logical plan read"); + }; + let plan = Plan::decode(plan.as_slice()).unwrap(); + *self.expiry_lower.lock().unwrap() = + assert_recovery_timestamp_plan(&plan, self.expire_after); + Ok(self.output.lock().unwrap().take().unwrap()) + } +} + +fn recovery_timestamp_schema() -> Arc { + Arc::new(Schema::new(vec![ + ColumnSchema::new("ts", CDT::timestamp_millisecond_datatype(), true).with_time_index(true), + ])) +} + +fn recovery_stream_output( + batches: Vec, + watermarks: Option)>>, + fail: bool, +) -> Output { + Output::new_with_stream(Box::pin(RecoveryMetricsStream { + schema: recovery_timestamp_schema(), + batches, + metrics: watermarks.map(|watermarks| RecordBatchMetrics { + region_watermarks: watermarks + .into_iter() + .map(|(region_id, watermark)| RegionWatermarkEntry { + region_id, + watermark, + }) + .collect(), + ..Default::default() + }), + fail, + })) +} + +fn recovery_timestamp_batch(values: Vec>) -> RecordBatch { + RecordBatch::new( + recovery_timestamp_schema(), + vec![Arc::new(TimestampMillisecondVector::from(values)) as VectorRef], + ) + .unwrap() +} + +fn seed_recovery_state( + task: &BatchingTask, + lower: &BTreeMap, +) -> (BTreeMap, String) { + let mut state = task.state.write().unwrap(); + state.advance_checkpoints(lower.iter().map(|(region, seq)| (*region, *seq)).collect()); + state.dirty_time_windows.add_window( + Timestamp::new_millisecond(20_000), + Some(Timestamp::new_millisecond(25_000)), + ); + ( + state.checkpoints().clone(), + format!("{:?}", state.dirty_time_windows), + ) +} + +fn assert_recovery_state_unchanged( + task: &BatchingTask, + checkpoints: &BTreeMap, + dirty: &str, +) { + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), checkpoints); + assert_eq!(format!("{:?}", state.dirty_time_windows), dirty); +} + struct CaptureScheduledNowHandler { expected_extension: String, captured_sql: Arc>>, @@ -1370,6 +1844,331 @@ fn output_with_region_watermarks( result } +#[tokio::test] +async fn test_capture_recovery_windows_decodes_raw_timestamp_projection_with_where_and_expiry() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(src.number) AS output_value, date_bin(INTERVAL '5 second', src.ts) AS output_window FROM numbers_with_ts AS src WHERE src.number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let expire_after = 5; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .expire_after = Some(expire_after); + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(expire_after), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &lower); + let before = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + let (high, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + let after = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + let expiry_lower = handler.expiry_lower.lock().unwrap().unwrap(); + let time_window_expr = task.config.time_window_expr.as_ref().unwrap(); + // The helper and samples may straddle a window boundary. + let aligned_expiry = |now: Timestamp| { + time_window_expr + .eval( + now.sub_duration(Duration::from_secs(expire_after as u64)) + .unwrap(), + ) + .unwrap() + .0 + .unwrap() + }; + assert!(aligned_expiry(before) <= expiry_lower && expiry_lower <= aligned_expiry(after)); + assert_eq!( + time_window_expr.eval(expiry_lower).unwrap().0, + Some(expiry_lower) + ); + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); +} + +#[tokio::test] +async fn test_capture_recovery_windows_since_uses_supplied_retention_lower() { + let TestTaskParts { + task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let retention_lower = Timestamp::new_second(-10); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(1), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + + let (high, windows) = task + .capture_recovery_windows_since(&query_engine, &frontend, &lower, Some(retention_lower)) + .await + .unwrap(); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + assert_eq!(*handler.expiry_lower.lock().unwrap(), Some(retention_lower)); +} + +#[tokio::test] +async fn test_capture_recovery_windows_terminal_proof_cases_leave_state_unchanged() { + struct Case { + name: &'static str, + lower: BTreeMap, + watermarks: Option)>>, + succeeds: bool, + } + + let cases = vec![ + Case { + name: "empty_stream_complete_proof_h_equals_c", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: true, + }, + Case { + name: "missing_metrics", + lower: BTreeMap::from([(1, 10)]), + watermarks: None, + succeeds: false, + }, + Case { + name: "missing_c_region", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: false, + }, + Case { + name: "unknown_none_watermark", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10)), (2, None)]), + succeeds: false, + }, + Case { + name: "unexpected_region", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10)), (2, Some(20))]), + succeeds: false, + }, + Case { + name: "regressing_h", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(9))]), + succeeds: false, + }, + ]; + + for case in cases { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + case.watermarks, + false, + ))), + lower: serde_json::to_string(&case.lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&handler_dyn), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &case.lower); + + let result = task + .capture_recovery_windows(&query_engine, &frontend, &case.lower) + .await; + assert_eq!(result.is_ok(), case.succeeds, "{}: {result:?}", case.name); + if case.succeeds { + assert_eq!( + result.unwrap(), + (case.lower.clone(), vec![]), + "{}", + case.name + ); + } + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); + } +} + +#[tokio::test] +async fn test_capture_recovery_windows_streams_sorted_unique_windows() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![ + recovery_timestamp_batch(vec![Some(10_000), Some(-1), Some(1_000)]), + recovery_timestamp_batch(vec![Some(11_000), Some(1_000), Some(-5_000)]), + ], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + let (_, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + assert_eq!( + windows, + vec![ + ( + Timestamp::new_millisecond(-5_000), + Timestamp::new_millisecond(0) + ), + ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(5_000) + ), + ( + Timestamp::new_millisecond(10_000), + Timestamp::new_millisecond(15_000) + ), + ] + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_stream_error() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + true, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_null_timestamp() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![recovery_timestamp_batch(vec![None])], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + #[test] fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() { let query_ctx = QueryContext::arc(); @@ -2327,7 +3126,7 @@ async fn test_full_snapshot_scoped_plan_marks_checkpoint_advance_safe_only_after } #[tokio::test] -async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { +async fn test_expired_fenced_repair_uses_frozen_scope() { let TestTaskParts { mut task, query_engine, @@ -2365,11 +3164,18 @@ async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { ) .await .unwrap() - .expect("expired empty repair should fall back to live dirty"); + .unwrap(); - assert!(matches!(plan.coverage, QueryCoverage::ScopedBaseRepair)); - assert!(plan.coverage.snapshot_seqs().is_empty()); - assert!(task.state.read().unwrap().pending_fenced_repair().is_none()); + // A fenced repair is a frozen explicit scope: expiry cannot discard its + // old window or make this query consume the live post-fence signal. + assert!(matches!( + plan.coverage, + QueryCoverage::FencedRepairChunk { .. } + )); + let state = task.state.read().unwrap(); + assert!(state.pending_fenced_repair().is_some()); + assert!(state.fenced_repair_pending_is_empty()); + assert_eq!(state.dirty_time_windows.len(), 1); } #[tokio::test]