diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index 18db59e3..5fc18cbe 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -12,6 +12,11 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, Mutex}; type SummaryState = Arc; +type PendingOutput = ( + Option, + PrecomputedOutput, + Box, +); struct OperatorAdapter<'a> { binding: &'a BackendExecutableBinding, @@ -68,16 +73,23 @@ struct CommittedState { struct CommitRegistry(Mutex>); impl CommitRegistry { - fn claim_publish(&self, key: &MaterializationCommitKey) -> Result { + fn publish( + &self, + key: &MaterializationCommitKey, + emit: impl FnOnce() -> Result<(), Box>, + ) -> Result<(), Box> { + // Serialize acknowledgement with publication so a concurrent replay + // cannot skip an in-flight write that later fails. let mut commits = self.0.lock().map_err(|_| "commit registry poisoned")?; let committed = commits .get_mut(key) .ok_or_else(|| "maintenance result was not committed".to_string())?; if committed.published { - Ok(false) + Ok(()) } else { + emit()?; committed.published = true; - Ok(true) + Ok(()) } } } @@ -132,9 +144,9 @@ impl MaintenanceDagSink { &self, output: PrecomputedOutput, state: Box, - ) -> Result)>, String> { + ) -> Result, String> { let Some(plan) = self.plans.physical_plan_snapshot() else { - return Ok(vec![(output, state)]); + return Ok(vec![(None, output, state)]); }; let source_definition: asap_types::sds::SummaryDefinitionId = output.policy_fp.into(); let source: SummaryState = Arc::from(state); @@ -220,9 +232,6 @@ impl MaintenanceDagSink { &self.commits, ) .map_err(schedule_error)?; - if !self.commits.claim_publish(&key)? { - continue; - } let target = match installed.binding.node(*sink_node) { Some(BackendNodeBinding::Materialization { summary_definition }) => { asap_types::PolicyFingerprint::from(*summary_definition) @@ -231,13 +240,17 @@ impl MaintenanceDagSink { }; let mut target_output = output.clone(); target_output.policy_fp = target; - derived.push((target_output, value.as_ref().as_ref().clone_boxed_core())); + derived.push(( + Some(key), + target_output, + value.as_ref().as_ref().clone_boxed_core(), + )); } } if matched { Ok(derived) } else { - Ok(vec![(output, (*source).clone_boxed_core())]) + Ok(vec![(None, output, (*source).clone_boxed_core())]) } } } @@ -287,7 +300,25 @@ impl OutputSink for MaintenanceDagSink { .map_err(|e| -> Box { e.into() })?, ); } - self.inner.emit_batch(transformed) + if transformed.iter().all(|(key, _, _)| key.is_none()) { + return self.inner.emit_batch( + transformed + .into_iter() + .map(|(_, output, state)| (output, state)) + .collect(), + ); + } + // The generic sink can partially accept a batch. Acknowledge each + // maintained output independently so retries skip only accepted writes. + for (key, output, state) in transformed { + match key { + Some(key) => self + .commits + .publish(&key, || self.inner.emit_batch(vec![(output, state)]))?, + None => self.inner.emit_batch(vec![(output, state)])?, + } + } + Ok(()) } } @@ -339,6 +370,128 @@ mod tests { Arc::new(accumulator) } + // A failed downstream write must be retried, while accepted outputs remain + // idempotent when the same maintenance lineage is replayed. + #[test] + fn downstream_failure_does_not_acknowledge_maintenance_publication() { + use crate::storage_engines::types::{ + ActivePhysicalPlan, HotReloadActivePhysicalPlan, StreamingConfig, + }; + use control_plane::physical::executable_binding::{ + InstalledPostAsapDag, PostAsapDagDocument, + }; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Default)] + struct FailOnceSink { + attempts: AtomicUsize, + accepted: AtomicUsize, + fail_at: usize, + } + impl OutputSink for FailOnceSink { + fn emit_batch( + &self, + outputs: Vec<(PrecomputedOutput, Box)>, + ) -> Result<(), Box> { + if outputs.is_empty() { + return Ok(()); + } + if self.attempts.fetch_add(1, Ordering::SeqCst) == self.fail_at { + return Err("temporary store failure".into()); + } + self.accepted.fetch_add(outputs.len(), Ordering::SeqCst); + Ok(()) + } + } + + let snapshot: control_plane::physical::compiler::BackendLocalPlanningSnapshot = + serde_json::from_str(include_str!( + "../../../docs/examples/asapquery-planning-snapshot.json" + )) + .unwrap(); + let mut bundle = snapshot.compile().unwrap(); + let mut query = node(2); + query.output_state = planner_types::post_asap::ExecutionDataState::READ_ROWS; + let dag = ExecutableDag { + nodes: vec![node(0), node(1), query], + edges: vec![edge(0, 1), edge(1, 2)], + root: PostAsapNodeId(2), + }; + let binding = BackendExecutableBinding { + nodes: BTreeMap::from([ + ( + PostAsapNodeId(0), + BackendNodeBinding::Materialization { + summary_definition: definition(1), + }, + ), + ( + PostAsapNodeId(1), + BackendNodeBinding::Materialization { + summary_definition: definition(2), + }, + ), + ( + PostAsapNodeId(2), + BackendNodeBinding::Query { + query_node: control_plane::query_plan::QueryNodeId(9), + }, + ), + ]), + query_sink: PostAsapNodeId(2), + query_plan_sink: control_plane::query_plan::QueryNodeId(9), + precompute_sinks: vec![PostAsapNodeId(1)], + }; + bundle.precompute_plan.executable_dags = BTreeMap::from([( + "retry".into(), + InstalledPostAsapDag { + document: PostAsapDagDocument::from_executable("retry".into(), &dag).unwrap(), + binding, + }, + )]); + let active = ActivePhysicalPlan { + envelope: bundle.precompute_plan.envelope.clone(), + summary_catalog: Some(Arc::new(bundle.summary_catalog)), + precompute_plan: bundle.precompute_plan, + transmission_plan: bundle.transmission_plan, + runtime_config: Arc::new(StreamingConfig::new(Default::default())), + query_plan: Arc::new(bundle.query_plan), + storage_routing: Arc::new(Default::default()), + }; + for fail_at in [0, 1] { + let downstream = Arc::new(FailOnceSink { + fail_at, + ..Default::default() + }); + let sink = MaintenanceDagSink::new( + downstream.clone(), + HotReloadStreamingConfig::from_active(HotReloadActivePhysicalPlan::new( + active.clone(), + )), + ); + let batch = || { + (0..2) + .map(|i| { + ( + PrecomputedOutput::new( + i * 10, + (i + 1) * 10, + None, + asap_types::PolicyFingerprint(1), + ), + sum(2.0).clone_boxed_core(), + ) + }) + .collect() + }; + assert!(sink.emit_batch(batch()).is_err()); + sink.emit_batch(batch()).unwrap(); + sink.emit_batch(batch()).unwrap(); + assert_eq!(downstream.accepted.load(Ordering::SeqCst), 2); + assert_eq!(downstream.attempts.load(Ordering::SeqCst), 3); + } + } + #[test] fn shared_summary_node_executes_once_and_summary_over_summary_merges() { // source 0 is shared by both branches; root therefore contains two @@ -396,8 +549,10 @@ mod tests { ) .unwrap(); assert_eq!(result.as_ref().aux_stats().sum, Some(4.0)); - assert!(commits.claim_publish(&key).unwrap()); - assert!(!commits.claim_publish(&key).unwrap()); + commits.publish(&key, || Ok(())).unwrap(); + commits + .publish(&key, || panic!("accepted lineage must not publish twice")) + .unwrap(); } #[test]