fix: graceful shutdown of rrelayer (#39)

* fix: graceful shutdown of rrelayer

* fixes

* remove workflow

* remove workflow

Author: joshstevens19

crates/core/src/lib.rs

@@ -3,6 +3,8 @@ pub mod authentication;
 pub mod gas;
 mod logger;
 pub use logger::setup_info_logger;
+mod shutdown;
+pub use shutdown::{enter_critical_operation, is_shutdown_in_progress, request_graceful_shutdown};
 mod middleware;
 pub mod network;
 mod postgres;

crates/core/src/logger.rs

@@ -3,6 +3,7 @@ use std::{
     sync::atomic::{AtomicBool, Ordering},
 };
 
+use crate::is_shutdown_in_progress;
 use once_cell::sync::Lazy;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::{
@@ -27,8 +28,7 @@ impl ShutdownAwareWriter {
 
 impl Write for ShutdownAwareWriter {
     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) {
-            // During shutdown, write directly to stdout
+        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) || is_shutdown_in_progress() {
             let stdout = std::io::stdout();
             let mut handle = stdout.lock();
             handle.write(buf)
@@ -38,7 +38,7 @@ impl Write for ShutdownAwareWriter {
     }
 
     fn flush(&mut self) -> std::io::Result<()> {
-        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) {
+        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) || is_shutdown_in_progress() {
             let stdout = std::io::stdout();
             let mut handle = stdout.lock();
             handle.flush()
@@ -62,7 +62,7 @@ struct CustomTimer;
 
 impl tracing_subscriber::fmt::time::FormatTime for CustomTimer {
     fn format_time(&self, writer: &mut Writer<'_>) -> std::fmt::Result {
-        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) {
+        if SHUTDOWN_IN_PROGRESS.load(Ordering::Relaxed) || is_shutdown_in_progress() {
             let now = chrono::Local::now();
             write!(writer, "{}", now.format("%H:%M:%S"))
         } else {
@@ -93,11 +93,6 @@ pub fn setup_info_logger() {
     setup_logger(LevelFilter::INFO);
 }
 
-#[allow(dead_code)]
-pub fn mark_shutdown_started() {
-    SHUTDOWN_IN_PROGRESS.store(true, Ordering::Relaxed);
-}
-
 #[allow(dead_code)]
 pub struct LoggerGuard;
 

crates/core/src/shutdown.rs

@@ -0,0 +1,228 @@
+use std::{
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use once_cell::sync::Lazy;
+use tokio::sync::{broadcast, Notify};
+use tracing::{info, warn};
+
+static SHUTDOWN_COORDINATOR: Lazy<Arc<ShutdownCoordinator>> =
+    Lazy::new(|| Arc::new(ShutdownCoordinator::new()));
+
+pub struct ShutdownCoordinator {
+    shutdown_sender: broadcast::Sender<()>,
+    operations_complete: Arc<Notify>,
+    active_operations: Arc<AtomicUsize>,
+    shutdown_requested: Arc<AtomicBool>,
+}
+
+impl ShutdownCoordinator {
+    fn new() -> Self {
+        let (shutdown_sender, _) = broadcast::channel(16);
+        Self {
+            shutdown_sender,
+            operations_complete: Arc::new(Notify::new()),
+            active_operations: Arc::new(AtomicUsize::new(0)),
+            shutdown_requested: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    #[cfg(test)]
+    pub fn new_for_test() -> Self {
+        Self::new()
+    }
+
+    pub fn subscribe(&self) -> broadcast::Receiver<()> {
+        self.shutdown_sender.subscribe()
+    }
+
+    pub fn is_shutdown_requested(&self) -> bool {
+        self.shutdown_requested.load(Ordering::Relaxed)
+    }
+
+    pub fn enter_operation(self: &Arc<Self>) -> Option<OperationGuard> {
+        if self.is_shutdown_requested() {
+            return None;
+        }
+        self.active_operations.fetch_add(1, Ordering::SeqCst);
+        Some(OperationGuard { coordinator: Arc::clone(self) })
+    }
+
+    fn exit_operation(&self) {
+        let prev = self.active_operations.fetch_sub(1, Ordering::SeqCst);
+        if prev == 1 {
+            self.operations_complete.notify_waiters();
+        }
+    }
+
+    pub async fn request_shutdown(&self, timeout: Duration) -> bool {
+        info!(
+            "Graceful shutdown requested, waiting for {} active operations to complete",
+            self.active_operations.load(Ordering::Relaxed)
+        );
+
+        self.shutdown_requested.store(true, Ordering::SeqCst);
+
+        let _ = self.shutdown_sender.send(());
+
+        tokio::select! {
+            _ = self.wait_for_operations_complete() => {
+                info!("All operations completed gracefully");
+                true
+            }
+            _ = tokio::time::sleep(timeout) => {
+                warn!("Shutdown timeout reached, {} operations still active",
+                      self.active_operations.load(Ordering::Relaxed));
+                false
+            }
+        }
+    }
+
+    async fn wait_for_operations_complete(&self) {
+        while self.active_operations.load(Ordering::Relaxed) > 0 {
+            self.operations_complete.notified().await;
+        }
+    }
+
+    #[cfg(test)]
+    pub fn active_operations_count(&self) -> usize {
+        self.active_operations.load(Ordering::Relaxed)
+    }
+}
+
+pub struct OperationGuard {
+    coordinator: Arc<ShutdownCoordinator>,
+}
+
+impl Drop for OperationGuard {
+    fn drop(&mut self) {
+        self.coordinator.exit_operation();
+    }
+}
+
+pub fn shutdown_coordinator() -> Arc<ShutdownCoordinator> {
+    Arc::clone(&SHUTDOWN_COORDINATOR)
+}
+
+pub async fn request_graceful_shutdown(timeout: Duration) -> bool {
+    shutdown_coordinator().request_shutdown(timeout).await
+}
+
+pub fn is_shutdown_in_progress() -> bool {
+    shutdown_coordinator().is_shutdown_requested()
+}
+
+pub fn enter_critical_operation() -> Option<OperationGuard> {
+    let coordinator = shutdown_coordinator();
+    coordinator.enter_operation()
+}
+
+pub fn subscribe_to_shutdown() -> broadcast::Receiver<()> {
+    shutdown_coordinator().subscribe()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::{sync::Arc, time::Duration};
+    use tokio::time::sleep;
+
+    #[tokio::test]
+    async fn test_shutdown_coordination() {
+        let coordinator = Arc::new(ShutdownCoordinator::new_for_test());
+
+        let guard1 = coordinator.enter_operation();
+        assert!(guard1.is_some());
+
+        let guard2 = coordinator.enter_operation();
+        assert!(guard2.is_some());
+
+        assert_eq!(coordinator.active_operations_count(), 2);
+
+        let coordinator_clone = coordinator.clone();
+        let shutdown_task = tokio::spawn(async move {
+            coordinator_clone.request_shutdown(Duration::from_millis(100)).await
+        });
+
+        sleep(Duration::from_millis(10)).await;
+
+        let guard3 = coordinator.enter_operation();
+        assert!(guard3.is_none());
+
+        assert_eq!(coordinator.active_operations_count(), 2);
+
+        drop(guard1);
+        assert_eq!(coordinator.active_operations_count(), 1);
+
+        drop(guard2);
+
+        let result = shutdown_task.await.unwrap();
+        assert!(result);
+        assert_eq!(coordinator.active_operations_count(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_shutdown_timeout() {
+        let coordinator = Arc::new(ShutdownCoordinator::new_for_test());
+
+        let _guard = coordinator.enter_operation();
+        assert!(_guard.is_some());
+        assert_eq!(coordinator.active_operations_count(), 1);
+
+        let result = coordinator.request_shutdown(Duration::from_millis(10)).await;
+
+        assert!(!result);
+        assert_eq!(coordinator.active_operations_count(), 1);
+
+        drop(_guard);
+        assert_eq!(coordinator.active_operations_count(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_enter_critical_operation() {
+        let guard = enter_critical_operation();
+
+        if guard.is_some() {
+            drop(guard);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_operation_guard_lifecycle() {
+        let coordinator = Arc::new(ShutdownCoordinator::new_for_test());
+        assert_eq!(coordinator.active_operations_count(), 0);
+
+        {
+            let _guard1 = coordinator.enter_operation();
+            assert_eq!(coordinator.active_operations_count(), 1);
+
+            {
+                let _guard2 = coordinator.enter_operation();
+                assert_eq!(coordinator.active_operations_count(), 2);
+            }
+
+            assert_eq!(coordinator.active_operations_count(), 1);
+        }
+
+        assert_eq!(coordinator.active_operations_count(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_shutdown_signal_propagation() {
+        let coordinator = Arc::new(ShutdownCoordinator::new_for_test());
+        let mut rx = coordinator.subscribe();
+
+        let coordinator_clone = coordinator.clone();
+        tokio::spawn(
+            async move { coordinator_clone.request_shutdown(Duration::from_millis(50)).await },
+        );
+
+        let result = tokio::time::timeout(Duration::from_millis(100), rx.recv()).await;
+        assert!(result.is_ok());
+        assert!(result.unwrap().is_ok());
+    }
+}

crates/core/src/startup.rs

@@ -18,6 +18,7 @@ use crate::{
     schema::apply_schema,
     setup_info_logger,
     shared::cache::Cache,
+    shutdown,
     signing::create_signing_routes,
     transaction::{
         api::create_transactions_routes,
@@ -38,11 +39,15 @@ use axum::{
 };
 use dotenv::dotenv;
 use std::path::Path;
-use std::{net::SocketAddr, sync::Arc, time::Instant};
+use std::{
+    net::SocketAddr,
+    sync::Arc,
+    time::{Duration, Instant},
+};
 use thiserror::Error;
 use tokio::sync::Mutex;
 use tower_http::cors::{AllowOrigin, Any, CorsLayer};
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
 #[derive(Error, Debug)]
 #[allow(clippy::enum_variant_names)]
@@ -258,7 +263,62 @@ async fn start_api(
 
     let listener = tokio::net::TcpListener::bind(&address).await?;
     info!("rrelayer is up on http://{}", address);
-    axum::serve(listener, app).await.map_err(StartApiError::ApiStartupError)?;
+
+    let shutdown_signal = async {
+        let ctrl_c = async {
+            tokio::signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
+        };
+
+        #[cfg(unix)]
+        let terminate = async {
+            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+                .expect("failed to install signal handler")
+                .recv()
+                .await;
+        };
+
+        #[cfg(windows)]
+        let terminate = async {
+            tokio::signal::windows::ctrl_break()
+                .expect("failed to install Ctrl+Break handler")
+                .recv()
+                .await;
+        };
+
+        #[cfg(not(any(unix, windows)))]
+        let terminate = std::future::pending::<()>();
+
+        tokio::select! {
+            _ = ctrl_c => {
+                info!("Received Ctrl+C, initiating graceful shutdown");
+            },
+            _ = terminate => {
+                #[cfg(unix)]
+                info!("Received SIGTERM, initiating graceful shutdown");
+                #[cfg(windows)]
+                info!("Received Ctrl+Break, initiating graceful shutdown");
+                #[cfg(not(any(unix, windows)))]
+                info!("Received terminate signal, initiating graceful shutdown");
+            },
+        }
+    };
+
+    tokio::select! {
+        result = axum::serve(listener, app) => {
+            result.map_err(StartApiError::ApiStartupError)?;
+        }
+        _ = shutdown_signal => {
+            info!("Starting graceful shutdown...");
+
+            let shutdown_successful = shutdown::request_graceful_shutdown(Duration::from_secs(30)).await;
+
+            if shutdown_successful {
+                info!("Graceful shutdown completed successfully");
+            } else {
+                warn!("Some operations did not complete within shutdown timeout");
+            }
+        }
+    }
 
     Ok(())
 }