Recover failed staging and verify concurrent successful generation cutover

control_plane/src/backend_client.rs

@@ -396,6 +396,33 @@ impl BackendClient {
         }
     }
 
+    pub async fn discard_staged_physical_plan(
+        &self,
+        plan_id: u64,
+        plan_version: u64,
+    ) -> std::result::Result<(), BackendPostError> {
+        let response = self
+            .http
+            .post(format!(
+                "{}/discard",
+                derive_physical_plan_url(&self.endpoint)
+            ))
+            .json(&serde_json::json!({"plan_id": plan_id, "plan_version": plan_version}))
+            .send()
+            .await
+            .map_err(classify_reqwest_error)?;
+        let status = response.status();
+        if status.is_success() {
+            Ok(())
+        } else {
+            Err(classify_http_status(
+                status,
+                response.text().await.unwrap_or_default(),
+                "PhysicalPlan discard POST",
+            ))
+        }
+    }
+
     pub async fn activate_physical_plan(
         &self,
         plan_id: u64,

control_plane/src/main.rs

@@ -675,9 +675,12 @@ async fn handle_compile_and_publish_physical_plan(
         .publish_collector_plans(&bundle.collector_plans, apply_timeout)
         .await
     {
+        let cleanup = backend
+            .discard_staged_physical_plan(bundle.envelope.plan_id, bundle.envelope.plan_version)
+            .await;
         return (
             StatusCode::BAD_GATEWAY,
-            format!("collector physical-plan publication failed: {error}"),
+            format!("collector physical-plan publication failed: {error}; staged backend cleanup: {cleanup:?}"),
         )
             .into_response();
     }

data_plane/src/drivers/query/servers/http.rs

@@ -482,6 +482,10 @@ impl HttpServer {
                 get(handle_get_backend_plan).post(handle_post_backend_plan),
             )
             .route("/api/v1/physical-plan", post(handle_post_physical_plan))
+            .route(
+                "/api/v1/physical-plan/discard",
+                post(handle_discard_physical_plan),
+            )
             .route(
                 "/api/v1/physical-plan/activate",
                 post(handle_activate_physical_plan),
@@ -588,6 +592,10 @@ impl HttpServer {
                 get(handle_get_backend_plan).post(handle_post_backend_plan),
             )
             .route("/api/v1/physical-plan", post(handle_post_physical_plan))
+            .route(
+                "/api/v1/physical-plan/discard",
+                post(handle_discard_physical_plan),
+            )
             .route(
                 "/api/v1/physical-plan/activate",
                 post(handle_activate_physical_plan),
@@ -6076,6 +6084,24 @@ async fn handle_activate_physical_plan(
         .into_response()
 }
 
+async fn handle_discard_physical_plan(
+    State(state): State<AppState>,
+    axum::Json(request): axum::Json<ActivatePhysicalPlanRequest>,
+) -> axum::response::Response {
+    use axum::response::IntoResponse;
+    let Some(lifecycle) = state.physical_plan_lifecycle.as_ref() else {
+        return (
+            StatusCode::SERVICE_UNAVAILABLE,
+            "physical-plan lifecycle is not attached",
+        )
+            .into_response();
+    };
+    match lifecycle.discard_staged(request.plan_id, request.plan_version) {
+        Ok(()) => StatusCode::NO_CONTENT.into_response(),
+        Err(error) => (StatusCode::CONFLICT, error.to_string()).into_response(),
+    }
+}
+
 async fn handle_physical_plan_status(State(state): State<AppState>) -> axum::response::Response {
     use axum::response::IntoResponse;
     let Some(lifecycle) = state.physical_plan_lifecycle.as_ref() else {

data_plane/src/precompute_engine/operators/sum_accumulator.rs

@@ -49,12 +49,15 @@ impl SumAccumulator {
     }
 
     pub fn deserialize_from_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        if buffer.len() < 4 {
-            return Err("Buffer too short for f32".into());
+        match buffer.len() {
+            // Legacy Python scalar sums carry no sample-count evidence.
+            4 => Ok(Self::with_sum(f32::from_le_bytes(buffer.try_into()?) as f64)),
+            // Counted sums use the same fixed layout as the Collector Sum payload.
+            16 => Self::from_sum_bytes(buffer),
+            len => {
+                Err(format!("Invalid persisted Sum payload length: {len} (want 4 or 16)").into())
+            }
         }
-        // Python uses struct.pack("<f", self.sum) which is 4-byte little-endian float
-        let sum = f32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as f64;
-        Ok(Self::with_sum(sum))
     }
 
     /// Decode the fixed Sum payload produced by the first-class Sum
@@ -96,8 +99,15 @@ impl SerializableToSink for SumAccumulator {
     }
 
     fn serialize_to_bytes(&self) -> Vec<u8> {
-        // Match Python's struct.pack("<f", self.sum) - 4-byte little-endian float
-        (self.sum as f32).to_le_bytes().to_vec()
+        match self.observation_count {
+            Some(count) => {
+                let mut bytes = Vec::with_capacity(16);
+                bytes.extend_from_slice(&self.sum.to_le_bytes());
+                bytes.extend_from_slice(&count.to_le_bytes());
+                bytes
+            }
+            None => (self.sum as f32).to_le_bytes().to_vec(),
+        }
     }
 }
 
@@ -272,6 +282,40 @@ mod tests {
         assert_eq!(merged.aux_stats().count, None);
     }
 
+    // Persistence retains known counts, including zero and the full u64 range.
+    #[test]
+    fn counted_sum_binary_round_trip() {
+        for count in [0, 3, u64::MAX] {
+            let acc = SumAccumulator {
+                sum: 1.0000000000001,
+                observation_count: Some(count),
+            };
+            let bytes = acc.serialize_to_bytes();
+            assert_eq!(bytes.len(), 16);
+            let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap();
+            assert_eq!(restored.sum, acc.sum);
+            assert_eq!(restored.observation_count, Some(count));
+        }
+    }
+
+    // Existing scalar-only files remain readable without inventing counts.
+    #[test]
+    fn legacy_binary_sum_has_unknown_count() {
+        let bytes = 42.5f32.to_le_bytes();
+        let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap();
+        assert_eq!(restored.sum, 42.5);
+        assert_eq!(restored.observation_count, None);
+        assert_eq!(restored.serialize_to_bytes(), bytes);
+    }
+
+    // Truncated counted payloads must not silently decode as scalar sums.
+    #[test]
+    fn persisted_sum_rejects_invalid_lengths() {
+        for len in [0, 3, 5, 8, 15, 17] {
+            assert!(SumAccumulator::deserialize_from_bytes(&vec![0; len]).is_err());
+        }
+    }
+
     #[test]
     fn test_sum_accumulator_creation() {
         let acc = SumAccumulator::new();

data_plane/src/storage_engines/sketch_db/index/mod.rs

@@ -3924,6 +3924,70 @@ mod tests {
         drop(p);
     }
 
+    // Raw sample counts must survive production durable storage.
+    #[test]
+    fn raw_count_survives_disk_eviction() {
+        use crate::storage_engines::types::AggregationType;
+        let tmp = tempfile::TempDir::new().unwrap();
+        let idx = Arc::new(SketchStore::new());
+        // Register an ExactAgg(Sum) sid keyed by `zone`.
+        let mut m = meta(8001);
+        m.metric_name = "http_requests_total".into();
+        m.group_by_keys = ["zone".to_string()].into_iter().collect();
+        m.agg_kind = AggKind::ExactAgg {
+            agg_type: AggregationType::Sum,
+            parameters_canonical: String::new(),
+            spatial_filter_canonical: String::new(),
+        };
+        idx.register(m);
+        let p = idx
+            .start_persistence(durable_cfg(tmp.path().to_path_buf()))
+            .unwrap();
+
+        let lv_zone = |v: &str| {
+            let mut x = BTreeMap::new();
+            x.insert("zone".to_string(), v.to_string());
+            x
+        };
+        for i in 0..10u64 {
+            let s = i * 30_000;
+            idx.append_precompute(
+                8001,
+                lv_zone("z0"),
+                (s, s + 30_000),
+                Box::new({
+                    let mut acc = crate::precompute_engine::operators::SumAccumulator::new();
+                    acc.update((i + 1) as f64);
+                    acc.update(10.0);
+                    acc
+                }),
+            );
+        }
+        assert!(
+            wait_until(
+                || idx.approx_memory_bytes() == 0 && idx.list_sealed_epochs_len() == 0,
+                std::time::Duration::from_secs(5),
+            ),
+            "exact-agg windows never fully evicted"
+        );
+        // Query the EVICTED portion [0, 150_000) — must come back from disk.
+        let series = idx.query_exact_agg_range(8001, 0, 150_000);
+        assert!(
+            !series.is_empty(),
+            "exact-agg query returned no result after flush and eviction"
+        );
+        let (_label, samples) = &series[0];
+        assert!(
+            samples.contains_key(&30_000),
+            "evicted exact-agg window missing from disk"
+        );
+        let stats = samples[&30_000].aux_stats();
+        assert_eq!(stats.count, Some(2));
+        assert_eq!(stats.sum, Some(11.0));
+        assert_eq!(stats.sum.unwrap() / stats.count.unwrap() as f64, 5.5);
+        drop(p);
+    }
+
     /// BUG #3: the memory diagnostic + the flusher's memory-pressure
     /// trigger must account for `current_epoch`, not just sealed epochs.
     /// On origin/main `approx_memory_bytes()` sums ONLY sealed epochs, so

data_plane/src/storage_engines/types/hot_reload_config.rs

@@ -309,6 +309,27 @@ impl PhysicalPlanLifecycle {
         Ok(())
     }
 
+    /// Roll back a failed publication without touching active readers or state.
+    pub fn discard_staged(
+        &self,
+        plan_id: u64,
+        plan_version: u64,
+    ) -> Result<(), PhysicalPlanLifecycleError> {
+        let key = (plan_id, plan_version);
+        let mut state = self
+            .state
+            .lock()
+            .expect("physical-plan lifecycle lock poisoned");
+        if state.staged.remove(&key).is_none() {
+            return Err(PhysicalPlanLifecycleError::NotStaged {
+                plan_id,
+                plan_version,
+            });
+        }
+        state.statuses.remove(&key);
+        Ok(())
+    }
+
     pub fn activate(
         &self,
         plan_id: u64,
@@ -981,6 +1002,25 @@ mod tests {
         assert_eq!(lifecycle.statuses()[0].phase, PhysicalPlanPhase::Retired);
     }
 
+    // Failed publication releases only its staging slot; active readers remain valid.
+    #[test]
+    fn discard_staged_allows_retry_and_never_discards_active() {
+        let active = HotReloadActivePhysicalPlan::new(physical_plan(7, 1, 100, None));
+        let held_reader = active.snapshot();
+        let lifecycle = PhysicalPlanLifecycle::new(active.clone());
+        lifecycle
+            .stage(physical_plan(7, 2, 200, None), 150)
+            .unwrap();
+        lifecycle.discard_staged(7, 2).unwrap();
+        lifecycle
+            .stage(physical_plan(7, 2, 300, None), 250)
+            .unwrap();
+        lifecycle.activate(7, 2, 300).unwrap();
+        assert!(lifecycle.discard_staged(7, 2).is_err());
+        assert_eq!(active.snapshot().backend_plan.plan_version, 2);
+        assert_eq!(held_reader.backend_plan.plan_version, 1);
+    }
+
     #[test]
     fn materialization_readiness_is_generation_scoped_and_monotonic() {
         let active = HotReloadActivePhysicalPlan::new(physical_plan(7, 1, 100, None));