From f79ee0bb8840be01693c1344b26cc1134e5d82a1 Mon Sep 17 00:00:00 2001
From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com>
Date: Fri, 28 Feb 2025 10:49:15 -0500
Subject: [PATCH] fix(storcon): loop in chaos injection (#11004)

## Problem

Somehow the previous patch loses the loop in the chaos injector function
so everything will only run once.
https://github.com/neondatabase/neon/pull/10934

## Summary of changes

Add back the loop.

Signed-off-by: Alex Chi Z <chi@neon.tech>
---
 .../src/service/chaos_injector.rs             | 71 ++++++++++---------
 1 file changed, 37 insertions(+), 34 deletions(-)

diff --git a/storage_controller/src/service/chaos_injector.rs b/storage_controller/src/service/chaos_injector.rs
index 2ff68d7037..a0419e0205 100644
--- a/storage_controller/src/service/chaos_injector.rs
+++ b/storage_controller/src/service/chaos_injector.rs
@@ -46,48 +46,51 @@ impl ChaosInjector {
         }
     }
 
+    fn get_cron_interval_sleep_future(&self) -> Option<tokio::time::Sleep> {
+        if let Some(ref chaos_exit_crontab) = self.chaos_exit_crontab {
+            match cron_to_next_duration(chaos_exit_crontab) {
+                Ok(interval_exit) => Some(interval_exit),
+                Err(e) => {
+                    tracing::error!("Error processing the cron schedule: {e}");
+                    None
+                }
+            }
+        } else {
+            None
+        }
+    }
+
     pub async fn run(&mut self, cancel: CancellationToken) {
         let mut interval = tokio::time::interval(self.interval);
-        let cron_interval = {
-            if let Some(ref chaos_exit_crontab) = self.chaos_exit_crontab {
-                match cron_to_next_duration(chaos_exit_crontab) {
-                    Ok(interval_exit) => Some(interval_exit),
-                    Err(e) => {
-                        tracing::error!("Error processing the cron schedule: {e}");
-                        None
-                    }
-                }
-            } else {
-                None
-            }
-        };
+        #[derive(Debug)]
         enum ChaosEvent {
             ShuffleTenant,
             ForceKill,
         }
-        let chaos_type = tokio::select! {
-            _ = interval.tick() => {
-                ChaosEvent::ShuffleTenant
-            }
-            Some(_) = maybe_sleep(cron_interval) => {
-                ChaosEvent::ForceKill
-            }
-            _ = cancel.cancelled() => {
-                tracing::info!("Shutting down");
-                return;
-            }
-        };
-
-        match chaos_type {
-            ChaosEvent::ShuffleTenant => {
-                self.inject_chaos().await;
-            }
-            ChaosEvent::ForceKill => {
-                self.force_kill().await;
+        loop {
+            let cron_interval = self.get_cron_interval_sleep_future();
+            let chaos_type = tokio::select! {
+                _ = interval.tick() => {
+                    ChaosEvent::ShuffleTenant
+                }
+                Some(_) = maybe_sleep(cron_interval) => {
+                    ChaosEvent::ForceKill
+                }
+                _ = cancel.cancelled() => {
+                    tracing::info!("Shutting down");
+                    return;
+                }
+            };
+            tracing::info!("Chaos iteration: {chaos_type:?}...");
+            match chaos_type {
+                ChaosEvent::ShuffleTenant => {
+                    self.inject_chaos().await;
+                }
+                ChaosEvent::ForceKill => {
+                    self.force_kill().await;
+                }
             }
         }
-
-        tracing::info!("Chaos iteration...");
     }
 
     /// If a shard has a secondary and attached location, then re-assign the secondary to be