Fix shutdown race on instance startup (#36382)

GitOrigin-RevId: 278d6249ee14b8cd73eaaa9a701b4592f567a59d

Author: goffrie

Cargo.lock

@@ -1,71 +1,63 @@
 use std::sync::Arc;
 
 use parking_lot::Mutex;
-use tokio::sync::mpsc;
+use tokio::sync::oneshot;
 
 use crate::errors::report_error_sync;
 
-// Used by the database to signal it has encountered a fatal error.
+/// Used by the database to signal it has encountered a fatal error.
 #[derive(Clone)]
 pub struct ShutdownSignal {
-    shutdown_tx: Option<Arc<Mutex<Option<mpsc::UnboundedSender<ShutdownMessage>>>>>,
-    instance_name: String,
-    generation_id: u64,
+    mode: Mode,
 }
 
-#[derive(Debug)]
-pub struct ShutdownMessage {
-    pub error: anyhow::Error,
-    pub instance_name: String,
-    pub generation_id: u64,
+/// Indicates what to do when an error is reported.
+#[derive(Clone)]
+enum Mode {
+    Panic,
+    /// If the `Option` inside the mutex is `Some`, the next fatal error will be
+    /// sent to that sender. Otherwise, signalling will do nothing (under the
+    /// presumption that an earlier error was reported and removed the sender).
+    Notify(Arc<Mutex<Option<oneshot::Sender<anyhow::Error>>>>),
 }
 
 impl ShutdownSignal {
-    pub fn new(
-        shutdown_tx: mpsc::UnboundedSender<ShutdownMessage>,
-        instance_name: String,
-        generation_id: u64,
-    ) -> Self {
+    /// Creates a new ShutdownSignal that sends the first encountered error to
+    /// the provided oneshot sender.
+    pub fn new(shutdown_tx: oneshot::Sender<anyhow::Error>) -> Self {
         Self {
-            shutdown_tx: Some(Arc::new(Mutex::new(Some(shutdown_tx)))),
-            instance_name,
-            generation_id,
+            mode: Mode::Notify(Arc::new(Mutex::new(Some(shutdown_tx)))),
         }
     }
 
+    /// Signals that an instance has encountered a fatal error and needs to be
+    /// shut down.
     pub fn signal(&self, mut fatal_error: anyhow::Error) {
         report_error_sync(&mut fatal_error);
-        if let Some(ref shutdown_tx_mutex) = self.shutdown_tx {
-            let Some(shutdown_tx) = shutdown_tx_mutex.lock().take() else {
-                // A shutdown message has already been sent for this instance. Do nothing.
-                return;
-            };
-            _ = shutdown_tx.send(ShutdownMessage {
-                error: fatal_error,
-                instance_name: self.instance_name.clone(),
-                generation_id: self.generation_id,
-            });
-        } else {
-            // We don't anyone to shutdown signal configured. Just panic.
-            panic!("Shutting down due to fatal error: {}", fatal_error);
+        match &self.mode {
+            Mode::Notify(shutdown_tx_mutex) => {
+                let Some(shutdown_tx) = shutdown_tx_mutex.lock().take() else {
+                    // A shutdown message has already been sent for this instance. Do nothing.
+                    return;
+                };
+                _ = shutdown_tx.send(fatal_error);
+            },
+            Mode::Panic => {
+                // We don't have the shutdown signal configured. Just panic.
+                panic!("Shutting down due to fatal error: {}", fatal_error);
+            },
         }
     }
 
-    // Creates a new ShutdownSignal that panics when signaled.
+    /// Creates a new ShutdownSignal that panics when signaled.
     pub fn panic() -> Self {
-        Self {
-            shutdown_tx: None,
-            instance_name: "".to_owned(),
-            generation_id: 0,
-        }
+        Self { mode: Mode::Panic }
     }
 
     #[cfg(any(test, feature = "testing"))]
     pub fn no_op() -> Self {
         Self {
-            shutdown_tx: Some(Arc::new(Mutex::new(None))),
-            instance_name: "".to_owned(),
-            generation_id: 0,
+            mode: Mode::Notify(Arc::new(Mutex::new(None))),
         }
     }
 }

crates/common/src/shutdown.rs

@@ -1617,9 +1617,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-util"
-version = "0.7.13"
+version = "0.7.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7fcaa8d55a2bdd6b83ace262b016eca0d79ee02818c5c1bcdf0305114081078"
+checksum = "6b9590b93e6fcc1739458317cccd391ad3955e2bde8913edf6f95f9e65a8f034"
 dependencies = [
  "bytes",
  "futures-core",

crates/convex/Cargo.oss.lock

@@ -32,7 +32,7 @@ use tokio::{
     signal::{
         self,
     },
-    sync::mpsc,
+    sync::oneshot,
 };
 
 fn main() -> Result<(), MainError> {
@@ -99,8 +99,8 @@ 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) = mpsc::unbounded_channel();
-    let preempt_signal = ShutdownSignal::new(preempt_tx.clone(), config.name(), 0);
+    let (preempt_tx, preempt_rx) = oneshot::channel();
+    let preempt_signal = ShutdownSignal::new(preempt_tx);
     // Use to signal to the http service to stop.
     let (shutdown_tx, shutdown_rx) = async_broadcast::broadcast(1);
     let persistence = connect_persistence(
@@ -142,9 +142,6 @@ async fn run_server_inner(runtime: ProdRuntime, config: LocalConfig) -> anyhow::
     let serve_future = future::try_join(serve_http_future, proxy_future).fuse();
     futures::pin_mut!(serve_future);
 
-    let preempt_future = async move { preempt_rx.recv().await }.fuse();
-    futures::pin_mut!(preempt_future);
-
     // Start shutdown when we get a manual shutdown signal or with the first
     // ctrl-c.
     let mut force_exit_duration = None;
@@ -153,7 +150,7 @@ async fn run_server_inner(runtime: ProdRuntime, config: LocalConfig) -> anyhow::
             r?;
             panic!("Serve future stopped unexpectedly!")
         },
-        _err = preempt_future => {
+        _err = preempt_rx.fuse() => {
             // If we fail with a fatal error, we want to exit immediately.
             tracing::info!("Received a fatal error. Shutting down immediately");
             force_exit_duration = Some(Duration::from_secs(0));