fix(autoscaling,ec2,kinesis): persist CFN ASG + scheduler/pipes deliveries and guard EC2 restart reconcile (bug-hunt) (#2316)

Author: vieiralucas

crates/fakecloud-autoscaling/src/cfn_provision.rs

@@ -12,9 +12,21 @@
 
 use std::sync::Arc;
 
+use fakecloud_persistence::SnapshotHook;
+
 use crate::AutoScalingService;
 use crate::SharedAutoScalingState;
 
+/// Snapshot hooks the detached CFN capacity reconcile fires after it mutates
+/// shared state, so the launched ASG instances (`autoscaling`) and their
+/// backing EC2 records (`ec2`) survive a restart. `None` fields make the
+/// corresponding persist a no-op (memory mode / unit tests).
+#[derive(Clone, Default)]
+pub struct CfnReconcilePersistHooks {
+    pub autoscaling: Option<SnapshotHook>,
+    pub ec2: Option<SnapshotHook>,
+}
+
 /// Reconcile a CFN-provisioned Auto Scaling Group to its desired capacity by
 /// launching REAL EC2 instances. No-op if the group is gone (e.g. the stack was
 /// deleted before reconciliation ran). Intended to be `tokio::spawn`ed by the
@@ -29,9 +41,21 @@ pub async fn cfn_reconcile_capacity(
     group_name: String,
     account_id: String,
     region: String,
+    persist: CfnReconcilePersistHooks,
 ) {
-    let svc = AutoScalingService::new(asg_state).with_ec2(ec2_state, ec2_runtime);
+    // The EC2 hook is set on the service so `apply_capacity` persists the REAL
+    // EC2 instances it launches; the autoscaling hook is fired below to persist
+    // the group's `instances` list. Both are required because this reconcile
+    // runs DETACHED after `CreateStack`/`UpdateStack` has already serialized the
+    // group with `instances=[]`; without them the launched instances (and their
+    // EC2 records) vanish on restart (bug-hunt restart-dataloss).
+    let svc = AutoScalingService::new(asg_state)
+        .with_ec2(ec2_state, ec2_runtime)
+        .with_ec2_snapshot_hook(persist.ec2);
     svc.reconcile_group(&account_id, &group_name, &region).await;
+    if let Some(hook) = persist.autoscaling {
+        hook().await;
+    }
 }
 
 /// Terminate the REAL EC2 instances launched for a CFN-provisioned Auto Scaling
@@ -54,3 +78,146 @@ pub async fn cfn_terminate_instances(
     svc.cfn_terminate_instances(&account_id, &region, &instance_ids)
         .await;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::state::{AutoScalingAccounts, AutoScalingGroup, AutoScalingSnapshot};
+    use fakecloud_persistence::SnapshotStore;
+    use parking_lot::{Mutex, RwLock};
+
+    /// Shared cell holding the bytes a capturing snapshot store last saved.
+    type Captured = Arc<Mutex<Option<Vec<u8>>>>;
+
+    // A `SnapshotStore` that captures the last-saved bytes in memory so a test
+    // can assert what the persist hook serialized (no tempfile needed).
+    struct CapturingStore(Captured);
+
+    impl SnapshotStore for CapturingStore {
+        fn load(&self) -> std::io::Result<Option<Vec<u8>>> {
+            Ok(self.0.lock().clone())
+        }
+        fn save(&self, bytes: &[u8]) -> std::io::Result<()> {
+            *self.0.lock() = Some(bytes.to_vec());
+            Ok(())
+        }
+    }
+
+    fn capturing() -> (Captured, Arc<dyn SnapshotStore>) {
+        let cell: Captured = Arc::new(Mutex::new(None));
+        (cell.clone(), Arc::new(CapturingStore(cell)))
+    }
+
+    fn group_desired(name: &str, desired: i64) -> AutoScalingGroup {
+        AutoScalingGroup {
+            name: name.to_string(),
+            arn: format!(
+                "arn:aws:autoscaling:us-east-1:123456789012:autoScalingGroup:x:autoScalingGroupName/{name}"
+            ),
+            launch_configuration_name: None,
+            launch_template: None,
+            min_size: 0,
+            max_size: desired,
+            desired_capacity: desired,
+            default_cooldown: 300,
+            availability_zones: vec!["us-east-1a".to_string()],
+            vpc_zone_identifier: None,
+            health_check_type: "EC2".to_string(),
+            health_check_grace_period: 0,
+            target_group_arns: Vec::new(),
+            load_balancer_names: Vec::new(),
+            new_instances_protected_from_scale_in: false,
+            created_time: chrono::Utc::now(),
+            instances: Vec::new(),
+            tags: Vec::new(),
+            status: None,
+            service_linked_role_arn: String::new(),
+        }
+    }
+
+    // bug-hunt restart-dataloss #1/#2: a CFN-provisioned ASG is inserted with
+    // `instances=[]` and reconciled to capacity by a DETACHED task. That task
+    // must persist BOTH the launched ASG instances (autoscaling snapshot) and
+    // their backing EC2 records (EC2 snapshot); otherwise the group reloads
+    // empty and the EC2 containers leak with no owning state on restart.
+    #[tokio::test]
+    async fn cfn_reconcile_persists_asg_and_ec2_instances() {
+        let account = "123456789012";
+        let asg_state: SharedAutoScalingState = Arc::new(RwLock::new(AutoScalingAccounts::new()));
+        asg_state
+            .write()
+            .get_or_create(account)
+            .groups
+            .insert("g".to_string(), group_desired("g", 2));
+
+        let ec2_state: fakecloud_ec2::SharedEc2State = Arc::new(RwLock::new(
+            fakecloud_core::multi_account::MultiAccountState::new(account, "us-east-1", ""),
+        ));
+
+        // Persist hooks over capturing stores, built exactly as the server wires
+        // the autoscaling + EC2 snapshot hooks.
+        let (asg_cell, asg_store) = capturing();
+        let asg_hook = AutoScalingService::new(asg_state.clone())
+            .with_snapshot_store(asg_store)
+            .snapshot_hook();
+        assert!(asg_hook.is_some(), "autoscaling hook must be built");
+
+        let (ec2_cell, ec2_store) = capturing();
+        let ec2_hook = fakecloud_ec2::Ec2Service::with_state(ec2_state.clone())
+            .with_snapshot_store(ec2_store)
+            .snapshot_hook();
+        assert!(ec2_hook.is_some(), "ec2 hook must be built");
+
+        // No EC2 runtime -> metadata-only instances (CI path), still real EC2
+        // records in `ec2_state`.
+        cfn_reconcile_capacity(
+            asg_state.clone(),
+            ec2_state.clone(),
+            None,
+            "g".to_string(),
+            account.to_string(),
+            "us-east-1".to_string(),
+            CfnReconcilePersistHooks {
+                autoscaling: asg_hook,
+                ec2: ec2_hook,
+            },
+        )
+        .await;
+
+        // In-memory: the group reached desired capacity.
+        assert_eq!(
+            asg_state.read().accounts[account].groups["g"]
+                .instances
+                .len(),
+            2,
+            "reconcile must launch 2 instances"
+        );
+
+        // #1: the autoscaling snapshot the detached reconcile fired contains the
+        // launched instances (not the empty list CreateStack serialized).
+        let asg_bytes = asg_cell
+            .lock()
+            .clone()
+            .expect("autoscaling snapshot written");
+        let asg_snap: AutoScalingSnapshot = serde_json::from_slice(&asg_bytes).unwrap();
+        let persisted = asg_snap.accounts.expect("multi-account snapshot");
+        assert_eq!(
+            persisted.accounts[account].groups["g"].instances.len(),
+            2,
+            "launched ASG instances must survive restart"
+        );
+
+        // #2: the EC2 snapshot the reconcile fired contains the backing records,
+        // so EC2 boot-recovery has a row to re-drive (no container leak).
+        let ec2_bytes = ec2_cell.lock().clone().expect("ec2 snapshot written");
+        let ec2_snap: fakecloud_ec2::Ec2Snapshot = serde_json::from_slice(&ec2_bytes).unwrap();
+        let ec2_accounts = ec2_snap.accounts.expect("ec2 multi-account snapshot");
+        let in_memory = ec2_state.read().default_ref().instances.len();
+        assert_eq!(in_memory, 2, "reconcile must create 2 EC2 records");
+        assert_eq!(
+            ec2_accounts.default_ref().instances.len(),
+            in_memory,
+            "ASG-launched EC2 records must be persisted (bug #2)"
+        );
+    }
+}

crates/fakecloud-autoscaling/src/service.rs

@@ -45,6 +45,14 @@ pub struct AutoScalingService {
     /// instances (unit tests).
     ec2_state: Option<fakecloud_ec2::SharedEc2State>,
     ec2_runtime: Option<Arc<fakecloud_ec2::Ec2Runtime>>,
+    /// EC2 whole-state snapshot hook. Reconciling desired capacity launches
+    /// (and terminates) REAL EC2 instances through a bare `Ec2Service` built
+    /// without a snapshot store, so those EC2 records lived only in memory and
+    /// leaked their containers on restart (EC2 boot-recovery had no persisted
+    /// `pending`/`running` row to re-drive). Firing this hook after an EC2
+    /// mutation writes the EC2 state through, the same way a direct
+    /// `RunInstances` API call persists (bug-hunt restart-dataloss).
+    ec2_snapshot_hook: Option<SnapshotHook>,
 }
 
 impl AutoScalingService {
@@ -55,6 +63,7 @@ impl AutoScalingService {
             snapshot_lock: Arc::new(AsyncMutex::new(())),
             ec2_state: None,
             ec2_runtime: None,
+            ec2_snapshot_hook: None,
         }
     }
 
@@ -63,6 +72,16 @@ impl AutoScalingService {
         self
     }
 
+    /// Attach the EC2 snapshot hook so ASG-launched EC2 instances are persisted
+    /// after capacity reconciliation. Without it, ASG-launched instances live
+    /// only in memory and their backing containers leak on restart (there is no
+    /// persisted EC2 row for boot-recovery to re-drive). `None` (memory mode /
+    /// unit tests) makes the persist a no-op.
+    pub fn with_ec2_snapshot_hook(mut self, hook: Option<SnapshotHook>) -> Self {
+        self.ec2_snapshot_hook = hook;
+        self
+    }
+
     /// Attach the EC2 backend so desired-capacity reconciliation launches real
     /// container-backed instances via `RunInstances`.
     pub fn with_ec2(
@@ -995,12 +1014,18 @@ impl AutoScalingService {
 
         let mut launched: Vec<AsgInstance> = Vec::new();
         let mut terminate_ids: Vec<String> = Vec::new();
+        // Whether this reconcile mutated REAL EC2 state (launched or terminated
+        // container-backed instances), so it must be persisted through the EC2
+        // snapshot hook below. No-op when there is no EC2 backend (the ids are
+        // synthesized metadata) or no hook wired.
+        let mut ec2_touched = false;
 
         if current_ids.len() < target {
             let need = target - current_ids.len();
             let mut ids = self
                 .run_ec2_instances(&image_id, &instance_type, subnet.as_deref(), need, req)
                 .await;
+            ec2_touched = self.ec2_state.is_some();
             // No EC2 backend wired (unit tests) or a partial launch: synthesize
             // the remainder so the group still reports its desired capacity.
             while ids.len() < need {
@@ -1024,34 +1049,51 @@ impl AutoScalingService {
             let remove = current_ids.len() - target;
             terminate_ids = current_ids.iter().rev().take(remove).cloned().collect();
             self.terminate_ec2_instances(&terminate_ids, req).await;
+            ec2_touched = self.ec2_state.is_some();
         }
 
-        let mut accounts = self.state.write();
-        let st = accounts.get_or_create(account);
-        let descs: Vec<String> = {
-            let Some(g) = st.groups.get_mut(name) else {
-                return;
-            };
-            let lcn = g.launch_configuration_name.clone();
-            let prot = g.new_instances_protected_from_scale_in;
-            let mut descs = Vec::new();
-            for mut ni in launched {
-                ni.launch_configuration_name = lcn.clone();
-                ni.protected_from_scale_in = prot;
-                descs.push(format!("Launching a new EC2 instance: {}", ni.instance_id));
-                g.instances.push(ni);
+        // Persist the REAL EC2 records this reconcile launched/terminated BEFORE
+        // touching the ASG state below. The ASG-launched instances are driven
+        // through a bare `Ec2Service` built without a snapshot store, so without
+        // firing the EC2 snapshot hook here they live only in memory and leak
+        // their containers on restart (EC2 boot-recovery has no persisted row to
+        // re-drive). Done before the ASG write so a concurrent group delete (the
+        // early `return` below) cannot skip persisting them. No-op when there is
+        // no EC2 backend or no hook wired (bug-hunt restart-dataloss).
+        if ec2_touched {
+            if let Some(hook) = &self.ec2_snapshot_hook {
+                hook().await;
             }
-            if !terminate_ids.is_empty() {
-                g.instances
-                    .retain(|i| !terminate_ids.contains(&i.instance_id));
-                for id in &terminate_ids {
-                    descs.push(format!("Terminating EC2 instance: {id}"));
+        }
+
+        {
+            let mut accounts = self.state.write();
+            let st = accounts.get_or_create(account);
+            let descs: Vec<String> = {
+                let Some(g) = st.groups.get_mut(name) else {
+                    return;
+                };
+                let lcn = g.launch_configuration_name.clone();
+                let prot = g.new_instances_protected_from_scale_in;
+                let mut descs = Vec::new();
+                for mut ni in launched {
+                    ni.launch_configuration_name = lcn.clone();
+                    ni.protected_from_scale_in = prot;
+                    descs.push(format!("Launching a new EC2 instance: {}", ni.instance_id));
+                    g.instances.push(ni);
                 }
+                if !terminate_ids.is_empty() {
+                    g.instances
+                        .retain(|i| !terminate_ids.contains(&i.instance_id));
+                    for id in &terminate_ids {
+                        descs.push(format!("Terminating EC2 instance: {id}"));
+                    }
+                }
+                descs
+            };
+            for d in descs {
+                st.activities.insert(0, activity(name, &d));
             }
-            descs
-        };
-        for d in descs {
-            st.activities.insert(0, activity(name, &d));
         }
     }
 }

crates/fakecloud-cloudformation/src/extras.rs

@@ -1075,7 +1075,8 @@ impl CloudFormationService {
                 // custom-resource Lambda invokes (0.2).
                 {
                     let handles =
-                        crate::service::ContainerBackingHandles::from_provisioner(&provisioner);
+                        crate::service::ContainerBackingHandles::from_provisioner(&provisioner)
+                            .with_snapshot_hooks(&self.snapshot_hooks);
                     handles.spawn_container_intents(std::mem::take(
                         &mut *provisioner.pending_container_spawns.lock(),
                     ));