fix: harden deletion executor lease handling

Co-authored-by: npub122y0pqkertljmedu303rl0aqrj3w8pvu43t6jxm6875lzg6f2pwqegc3xc <5288f082d91aff2de5bc8be23fbfa01ca2e3859cac57a91b7a3fa9f12349505c@buzz.block.builderlab.xyz>
Signed-off-by: npub122y0pqkertljmedu303rl0aqrj3w8pvu43t6jxm6875lzg6f2pwqegc3xc <5288f082d91aff2de5bc8be23fbfa01ca2e3859cac57a91b7a3fa9f12349505c@buzz.block.builderlab.xyz>

crates/buzz-deletion/src/lib.rs

@@ -37,6 +37,14 @@ fn heartbeat_interval() -> Duration {
     HEARTBEAT_INTERVAL
 }
 
+fn worker_health_stale_after() -> u64 {
+    HEARTBEAT_INTERVAL.as_secs().saturating_mul(3)
+}
+
+#[derive(Debug, Clone, thiserror::Error)]
+#[error("deletion execution lease heartbeat failed")]
+struct DeletionLeaseLost;
+
 #[derive(Debug, Clone, thiserror::Error)]
 #[error("{message}")]
 struct ServingWriteLeaseLost {
@@ -86,13 +94,13 @@ impl ServingWriteGuard {
         self.verify().await?;
         let output = tokio::select! {
             biased;
+            output = operation => output,
             _ = self.lost.cancelled() => {
                 return Err(ServingWriteLeaseLost {
                     message: "serving write lease heartbeat was lost".to_string(),
                 }
                 .into())
             }
-            output = operation => output,
         };
         self.verify().await?;
         Ok(output)
@@ -340,7 +348,8 @@ impl WorkerHealth {
             .as_secs();
         !self.draining.load(Ordering::Relaxed)
             && self.dependencies_ready.load(Ordering::Relaxed)
-            && now.saturating_sub(self.last_heartbeat_epoch.load(Ordering::Relaxed)) <= 30
+            && now.saturating_sub(self.last_heartbeat_epoch.load(Ordering::Relaxed))
+                <= worker_health_stale_after()
     }
 }
 
@@ -623,6 +632,24 @@ async fn run_loop(
     }
 }
 
+async fn stop_claim_executor(
+    services: &Services,
+    mode: LoopMode,
+    token: &LeaseToken,
+) -> Result<()> {
+    // A failed draining heartbeat must not prevent the generation-checked release
+    // attempt. `stop_executor` cannot clear a successor's reclaimed lease.
+    let _ = services
+        .store
+        .heartbeat(token, mode.as_str(), DEFAULT_LEASE_DURATION, true)
+        .await;
+    services
+        .store
+        .stop_executor(Some(token), &token.owner)
+        .await?;
+    Ok(())
+}
+
 async fn execute_claim(
     services: &Services,
     mode: LoopMode,
@@ -633,14 +660,7 @@ async fn execute_claim(
     let token = claim.lease.clone();
     loop {
         if shutdown.is_cancelled() {
-            services
-                .store
-                .heartbeat(&token, mode.as_str(), DEFAULT_LEASE_DURATION, true)
-                .await?;
-            services
-                .store
-                .stop_executor(Some(&token), &token.owner)
-                .await?;
+            stop_claim_executor(services, mode, &token).await?;
             let request = services.store.get(token.request_id).await?;
             return Ok(run_output(request));
         }
@@ -653,8 +673,19 @@ async fn execute_claim(
         let stage_result =
             run_stage_with_heartbeat(services, mode, &claim, shutdown, Arc::clone(&health)).await;
         match stage_result {
-            Ok(()) => {}
-            Err(error) => {
+            StageOutcome::Completed => {}
+            StageOutcome::Shutdown => {
+                stop_claim_executor(services, mode, &token).await?;
+                let request = services.store.get(token.request_id).await?;
+                return Ok(run_output(request));
+            }
+            StageOutcome::Failed(error) => {
+                let request = services.store.get(token.request_id).await?;
+                if request.lease_owner.as_deref() != Some(&token.owner)
+                    || request.lease_generation != token.generation
+                {
+                    return Ok(run_output(request));
+                }
                 let message = format!("{error:#}");
                 if is_permanent_error(&error) {
                     services
@@ -699,13 +730,38 @@ async fn dependencies_ready(services: &Services) -> bool {
     redis_ok && storage_ok
 }
 
+enum StageOutcome {
+    Completed,
+    Shutdown,
+    Failed(anyhow::Error),
+}
+
+async fn await_stage<F>(
+    stage: F,
+    shutdown: &CancellationToken,
+    heartbeat_error: &CancellationToken,
+) -> StageOutcome
+where
+    F: std::future::Future<Output = Result<()>>,
+{
+    tokio::select! {
+        biased;
+        _ = shutdown.cancelled() => StageOutcome::Shutdown,
+        _ = heartbeat_error.cancelled() => StageOutcome::Failed(DeletionLeaseLost.into()),
+        result = stage => match result {
+            Ok(()) => StageOutcome::Completed,
+            Err(error) => StageOutcome::Failed(error),
+        },
+    }
+}
+
 async fn run_stage_with_heartbeat(
     services: &Services,
     mode: LoopMode,
     claim: &ClaimedDeletion,
     shutdown: &CancellationToken,
     health: Arc<WorkerHealth>,
-) -> Result<()> {
+) -> StageOutcome {
     let heartbeat_services = services.clone();
     let heartbeat_token = claim.lease.clone();
     let heartbeat_mode = mode.as_str();
@@ -742,22 +798,26 @@ async fn run_stage_with_heartbeat(
         }
     });
 
-    let stage = tokio::select! {
-        _ = shutdown.cancelled() => Err(anyhow::anyhow!("executor shutdown requested")),
-        _ = heartbeat_error.cancelled() => Err(anyhow::anyhow!("deletion lease heartbeat failed")),
-        result = execute_stage(services, claim) => result,
-    };
+    let stage = await_stage(
+        execute_stage(services, claim, &heartbeat_error),
+        shutdown,
+        &heartbeat_error,
+    )
+    .await;
     heartbeat_shutdown.cancel();
     match heartbeat.await {
         Ok(()) => stage,
-        Err(error) => Err(anyhow::anyhow!("deletion heartbeat task failed: {error}")),
+        Err(error) => {
+            StageOutcome::Failed(anyhow::anyhow!("deletion heartbeat task failed: {error}"))
+        }
     }
 }
 
-async fn guarded_external_step<F, Fut>(
+async fn run_guarded_external_step<F, Fut>(
     services: &Services,
     token: &LeaseToken,
     stage: DeletionStage,
+    heartbeat_lost: &CancellationToken,
     operation: F,
 ) -> Result<()>
 where
@@ -765,12 +825,23 @@ where
     Fut: std::future::Future<Output = Result<()>>,
 {
     services.store.verify_execution_token(token, stage).await?;
-    operation().await?;
+    let result = tokio::select! {
+        biased;
+        _ = heartbeat_lost.cancelled() => {
+            return Err(DeletionLeaseLost.into());
+        }
+        result = operation() => result,
+    };
+    result?;
     services.store.verify_execution_token(token, stage).await?;
     Ok(())
 }
 
-async fn execute_stage(services: &Services, claim: &ClaimedDeletion) -> Result<()> {
+async fn execute_stage(
+    services: &Services,
+    claim: &ClaimedDeletion,
+    heartbeat_lost: &CancellationToken,
+) -> Result<()> {
     let request = &claim.request;
     let token = token_with_current_fence(&claim.lease, request);
     match request.stage {
@@ -803,7 +874,20 @@ async fn execute_stage(services: &Services, claim: &ClaimedDeletion) -> Result<(
             services.store.fence(&token).await?;
         }
         DeletionStage::Fenced => {
-            publish_disconnect_community(&services.redis, request.community_id).await?;
+            services
+                .store
+                .verify_execution_token(&token, DeletionStage::Fenced)
+                .await?;
+            let disconnect = tokio::select! {
+                biased;
+                _ = heartbeat_lost.cancelled() => Err(DeletionLeaseLost.into()),
+                result = publish_disconnect_community(&services.redis, request.community_id) => result,
+            };
+            disconnect?;
+            services
+                .store
+                .verify_execution_token(&token, DeletionStage::Fenced)
+                .await?;
             let destructive = match request.destructive_storage_manifest.clone() {
                 Some(value) => serde_json::from_value(value)?,
                 None => {
@@ -858,19 +942,31 @@ async fn execute_stage(services: &Services, claim: &ClaimedDeletion) -> Result<(
                 }
             }
             for key in &storage.tenant_keys {
-                guarded_external_step(services, &token, DeletionStage::Drained, || async {
-                    services.media.delete(key).await?;
-                    Ok(())
-                })
+                run_guarded_external_step(
+                    services,
+                    &token,
+                    DeletionStage::Drained,
+                    heartbeat_lost,
+                    || async {
+                        services.media.delete(key).await?;
+                        Ok(())
+                    },
+                )
                 .await?;
-                guarded_external_step(services, &token, DeletionStage::Drained, || async {
-                    if services.media.head(key).await? {
-                        return Err(transient(format!(
-                            "object binding still exists after delete: {key}"
-                        )));
-                    }
-                    Ok(())
-                })
+                run_guarded_external_step(
+                    services,
+                    &token,
+                    DeletionStage::Drained,
+                    heartbeat_lost,
+                    || async {
+                        if services.media.head(key).await? {
+                            return Err(transient(format!(
+                                "object binding still exists after delete: {key}"
+                            )));
+                        }
+                        Ok(())
+                    },
+                )
                 .await?;
             }
             services
@@ -1154,6 +1250,82 @@ fn print_json(value: &impl Serialize) -> Result<()> {
 mod tests {
     use super::*;
 
+    async fn claimed_test_deletion(prefix: &str) -> Option<(Services, ClaimedDeletion)> {
+        let database_url = std::env::var("BUZZ_TEST_DATABASE_URL")
+            .or_else(|_| std::env::var("DATABASE_URL"))
+            .ok()?;
+        let pool = sqlx::PgPool::connect(&database_url)
+            .await
+            .expect("connect deletion engine test DB");
+        let db = Db::from_pool(pool);
+        db.migrate().await.expect("migrate deletion engine test DB");
+        let store = db.deletion_store();
+        let host = format!("{prefix}-{}.example", Uuid::new_v4().simple());
+        let community = db
+            .ensure_configured_community(&host)
+            .await
+            .expect("create deletion engine test community");
+        let request = store
+            .submit(&host, "test", None)
+            .await
+            .expect("submit deletion request");
+        let inventory = FrozenInventory {
+            schema: store
+                .inventory_schema(community.id)
+                .await
+                .expect("inventory schema"),
+            storage: StorageManifest {
+                version: 1,
+                tenant_keys: Vec::new(),
+                git_pointer_keys: Vec::new(),
+                media_sidecar_keys: Vec::new(),
+                media_upload_keys: Vec::new(),
+                retained_shared_cas_keys: Vec::new(),
+                unknown_keys: Vec::new(),
+                unsupported_version_keys: Vec::new(),
+            },
+        };
+        store
+            .freeze_inventory(request.id, &inventory)
+            .await
+            .expect("freeze deletion inventory");
+        store
+            .approve(request.id, "test", None)
+            .await
+            .expect("approve deletion request");
+        let claim = store
+            .claim_specific(request.id, "test-executor", DEFAULT_LEASE_DURATION)
+            .await
+            .expect("claim deletion request")
+            .expect("runnable deletion request");
+        let services = Services {
+            store,
+            media: Arc::new(
+                MediaStorage::new(&buzz_media::MediaConfig {
+                    s3_endpoint: "http://127.0.0.1:1".to_string(),
+                    s3_access_key: "unused".to_string(),
+                    s3_secret_key: "unused".to_string(),
+                    s3_bucket: "unused".to_string(),
+                    s3_region: "us-east-1".to_string(),
+                    s3_addressing_style: buzz_media::S3AddressingStyle::Path,
+                    max_image_bytes: 1,
+                    max_gif_bytes: 1,
+                    max_video_bytes: 1,
+                    max_file_bytes: 1,
+                    public_base_url: "http://localhost/media".to_string(),
+                    upload_records_enabled: false,
+                    upload_ip_header: None,
+                    upload_port_header: None,
+                })
+                .expect("construct unused media service"),
+            ),
+            redis: deadpool_redis::Config::from_url("redis://127.0.0.1:1")
+                .create_pool(Some(deadpool_redis::Runtime::Tokio1))
+                .expect("construct unused Redis pool"),
+        };
+        Some((services, claim))
+    }
+
     #[test]
     fn permanent_failures_are_typed_not_string_classified() {
         let permanent_error = permanent("catalog drift");
@@ -1170,6 +1342,34 @@ mod tests {
         assert!(!is_permanent_error(&db_transient));
     }
 
+    #[test]
+    fn worker_readiness_requires_dependencies_heartbeat_and_not_draining() {
+        let health = WorkerHealth::default();
+        assert!(!health.ready());
+        health.dependencies_ready.store(true, Ordering::Relaxed);
+        health.mark_heartbeat();
+        assert!(health.ready());
+        health.draining.store(true, Ordering::Relaxed);
+        assert!(!health.ready());
+        health.draining.store(false, Ordering::Relaxed);
+        health.last_heartbeat_epoch.store(1, Ordering::Relaxed);
+        assert!(!health.ready());
+    }
+
+    #[test]
+    fn worker_configuration_requires_every_destructive_dependency() {
+        let variable = format!("BUZZ_DELETION_REQUIRED_TEST_{}", Uuid::new_v4().simple());
+        assert!(required_env(&variable).is_err());
+        std::env::set_var(&variable, "   ");
+        assert!(required_env(&variable).is_err());
+        std::env::set_var(&variable, "configured");
+        assert_eq!(
+            required_env(&variable).expect("configured environment variable"),
+            "configured"
+        );
+        std::env::remove_var(&variable);
+    }
+
     #[test]
     fn redis_absence_requires_terminal_cursor_and_all_pages_empty() {
         assert!(!scan_proves_absence(&[(9, Vec::new())]));
@@ -1226,6 +1426,111 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn guarded_external_step_rejects_preexisting_heartbeat_loss_without_polling_operation() {
+        let Some((services, claim)) = claimed_test_deletion("deletion-heartbeat").await else {
+            return;
+        };
+        let heartbeat_lost = CancellationToken::new();
+        heartbeat_lost.cancel();
+        let polled = Arc::new(AtomicBool::new(false));
+        let operation_polled = Arc::clone(&polled);
+        let result = run_guarded_external_step(
+            &services,
+            &claim.lease,
+            DeletionStage::Approved,
+            &heartbeat_lost,
+            || async move {
+                operation_polled.store(true, Ordering::Relaxed);
+                Ok(())
+            },
+        )
+        .await;
+        assert!(result.is_err(), "heartbeat loss must abort the side effect");
+        assert!(
+            result
+                .expect_err("heartbeat loss error")
+                .downcast_ref::<DeletionLeaseLost>()
+                .is_some(),
+            "heartbeat loss must stay typed"
+        );
+        assert!(
+            !polled.load(Ordering::Relaxed),
+            "a pre-cancelled heartbeat must win before polling the operation"
+        );
+    }
+
+    #[tokio::test]
+    async fn shutdown_during_stage_releases_claim_without_recording_retry() {
+        let Some((services, claim)) = claimed_test_deletion("deletion-shutdown").await else {
+            return;
+        };
+        let request_id = claim.request.id;
+        let retry_count = claim.request.retry_count;
+        let shutdown = CancellationToken::new();
+        let cancel = shutdown.clone();
+        let services_for_run = services.clone();
+        let executor = tokio::spawn(async move {
+            execute_claim(
+                &services_for_run,
+                LoopMode::Worker,
+                claim,
+                &shutdown,
+                Arc::new(WorkerHealth::default()),
+            )
+            .await
+        });
+        tokio::time::sleep(Duration::from_millis(10)).await;
+        cancel.cancel();
+        let output = tokio::time::timeout(Duration::from_secs(2), executor)
+            .await
+            .expect("shutdown must cancel the active stage")
+            .expect("deletion executor task")
+            .expect("graceful deletion executor shutdown");
+        let request = services
+            .store
+            .get(request_id)
+            .await
+            .expect("load deletion request after shutdown");
+        assert_eq!(output.stage, DeletionStage::Approved);
+        assert_eq!(request.stage, DeletionStage::Approved);
+        assert_eq!(request.retry_count, retry_count);
+        assert!(request.last_error.is_none());
+        assert!(request.lease_owner.is_none());
+        assert!(request.lease_until.is_none());
+    }
+
+    #[tokio::test]
+    async fn stage_wait_treats_shutdown_as_control_flow() {
+        let shutdown = CancellationToken::new();
+        shutdown.cancel();
+        let heartbeat_lost = CancellationToken::new();
+        let outcome = await_stage(
+            std::future::pending::<Result<()>>(),
+            &shutdown,
+            &heartbeat_lost,
+        )
+        .await;
+        assert!(matches!(outcome, StageOutcome::Shutdown));
+    }
+
+    #[tokio::test]
+    async fn stage_wait_prioritizes_heartbeat_loss_over_a_ready_operation() {
+        let shutdown = CancellationToken::new();
+        let heartbeat_lost = CancellationToken::new();
+        heartbeat_lost.cancel();
+        let outcome = await_stage(async { Ok(()) }, &shutdown, &heartbeat_lost).await;
+        match outcome {
+            StageOutcome::Failed(error) => assert!(
+                error.downcast_ref::<DeletionLeaseLost>().is_some(),
+                "heartbeat loss must stay typed"
+            ),
+            StageOutcome::Completed | StageOutcome::Shutdown => {
+                panic!("preexisting heartbeat loss must win")
+            }
+        }
+    }
+
     #[tokio::test]
     async fn shutdown_token_can_interrupt_idle_worker_sleep() {
         let token = CancellationToken::new();