Add generation ID to appstate (#34713)

* Add a generation ID to shutdown messages

GitOrigin-RevId: 8830ae1c627e5ec4d809f4faefb38016e212424e

crates/common/src/shutdown.rs

@@ -5,22 +5,26 @@ use std::sync::Arc;
 pub struct ShutdownSignal {
     shutdown_tx: Option<async_broadcast::Sender<ShutdownMessage>>,
     instance_name: String,
+    generation_id: u64,
 }
 
 #[derive(Clone, Debug)]
 pub struct ShutdownMessage {
     pub error: Arc<anyhow::Error>,
     pub instance_name: String,
+    pub generation_id: u64,
 }
 
 impl ShutdownSignal {
     pub fn new(
         shutdown_tx: async_broadcast::Sender<ShutdownMessage>,
         instance_name: String,
+        generation_id: u64,
     ) -> Self {
         Self {
             shutdown_tx: Some(shutdown_tx),
             instance_name,
+            generation_id,
         }
     }
 
@@ -29,6 +33,7 @@ impl ShutdownSignal {
             _ = shutdown_tx.try_broadcast(ShutdownMessage {
                 error: Arc::new(fatal_error),
                 instance_name: self.instance_name.clone(),
+                generation_id: self.generation_id,
             });
         } else {
             // We don't anyone to shutdown signal configured. Just panic.
@@ -41,6 +46,7 @@ impl ShutdownSignal {
         Self {
             shutdown_tx: None,
             instance_name: "".to_owned(),
+            generation_id: 0,
         }
     }
 
@@ -50,6 +56,7 @@ impl ShutdownSignal {
         Self {
             shutdown_tx: Some(sender),
             instance_name: "".to_owned(),
+            generation_id: 0,
         }
     }
 }

crates/local_backend/src/main.rs

@@ -100,7 +100,7 @@ async fn run_server(runtime: ProdRuntime, config: LocalConfig) -> anyhow::Result
 async fn run_server_inner(runtime: ProdRuntime, config: LocalConfig) -> anyhow::Result<()> {
     // Used to receive fatal errors from the database or /preempt endpoint.
     let (preempt_tx, mut preempt_rx) = async_broadcast::broadcast(1);
-    let preempt_signal = ShutdownSignal::new(preempt_tx.clone(), config.name());
+    let preempt_signal = ShutdownSignal::new(preempt_tx.clone(), config.name(), 0);
     // Use to signal to the http service to stop.
     let (shutdown_tx, shutdown_rx) = async_broadcast::broadcast(1);
     let persistence = connect_persistence(

crates/local_backend/src/test_helpers.rs

@@ -72,7 +72,7 @@ pub async fn setup_backend_for_test(runtime: ProdRuntime) -> anyhow::Result<Test
         config.clone(),
         Arc::new(persistence),
         shutdown_rx,
-        ShutdownSignal::new(preempt_tx, config.name()),
+        ShutdownSignal::new(preempt_tx, config.name(), 0),
     )
     .await?;
     let router = router(st.clone());