add graceful shutdown support

Author: pietroalbini

Cargo.lock

@@ -32,6 +32,6 @@ tracing = "0.1.16"
 tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
 atty = "0.2.14"
 clap = { version = "4.1.8", features = ["derive"] }
-tokio = { version = "1.26.0", features = ["rt-multi-thread", "rt"] }
+tokio = { version = "1.26.0", features = ["rt-multi-thread", "rt", "signal", "macros"] }
 aws-sdk-s3 = "0.26.0"
 aws-config = "0.55.1"

Cargo.toml

@@ -14,7 +14,9 @@ extern crate regex;
 extern crate rust_log_analyzer as rla;
 extern crate serde_json;
 
+use crate::server::QueueItem;
 use clap::Parser;
+use crossbeam::channel::Sender;
 use rla::index::IndexStorage;
 use std::process;
 use std::sync::Arc;
@@ -91,7 +93,10 @@ fn main() {
 
         let (queue_send, queue_recv) = crossbeam::channel::unbounded();
 
-        let service = Arc::new(server::RlaService::new(args.webhook_verify, queue_send)?);
+        let service = Arc::new(server::RlaService::new(
+            args.webhook_verify,
+            queue_send.clone(),
+        )?);
 
         let mut worker = server::Worker::new(
             args.index_file,
@@ -103,7 +108,7 @@ fn main() {
             args.query_builds_from_primary_repo,
         )?;
 
-        thread::spawn(move || {
+        let worker_thread = thread::spawn(move || {
             if let Err(e) = worker.main() {
                 error!("Worker failed, exiting: {}", e);
                 process::exit(1);
@@ -126,9 +131,34 @@ fn main() {
                         }))
                     }
                 }))
+                .with_graceful_shutdown(graceful_shutdown(queue_send))
                 .await
         })?;
 
+        worker_thread.join().expect("worker thread failed");
+
         Ok(())
     });
 }
+
+async fn graceful_shutdown(sender: Sender<QueueItem>) {
+    let ctrl_c = tokio::signal::ctrl_c();
+
+    // ECS uses SIGTERM to signal graceful shutdown must begin.
+    #[cfg(unix)]
+    let mut sigterm =
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
+    #[cfg(unix)]
+    let sigterm = sigterm.recv();
+
+    #[cfg(not(unix))]
+    let sigterm = std::future::pending(); // Never resolves
+
+    tokio::select! {
+        _ = ctrl_c => {}
+        _ = sigterm => {}
+    };
+
+    info!("graceful shutdown signal received");
+    let _ = sender.send(QueueItem::GracefulShutdown);
+}

src/bin/rla-server.rs

@@ -19,6 +19,7 @@ pub enum QueueItem {
         payload: rla::github::PullRequestEvent,
         delivery_id: String,
     },
+    GracefulShutdown,
 }
 
 impl QueueItem {
@@ -27,6 +28,7 @@ impl QueueItem {
             QueueItem::GitHubStatus { delivery_id, .. } => Some(&delivery_id),
             QueueItem::GitHubCheckRun { delivery_id, .. } => Some(&delivery_id),
             QueueItem::GitHubPullRequest { delivery_id, .. } => Some(&delivery_id),
+            QueueItem::GracefulShutdown => None,
         }
     }
 }

src/bin/server/mod.rs

@@ -84,7 +84,8 @@ impl Worker {
             let _enter = span.enter();
 
             match self.process(item, &span) {
-                Ok(()) => (),
+                Ok(ProcessOutcome::Continue) => (),
+                Ok(ProcessOutcome::Exit) => return Ok(()),
                 Err(e) => error!("Processing queue item failed: {}", e),
             }
         }
@@ -97,7 +98,7 @@ impl Worker {
         self.secondary_repos.iter().find(|r| *r == repo).is_some()
     }
 
-    fn process(&mut self, item: QueueItem, span: &tracing::Span) -> rla::Result<()> {
+    fn process(&mut self, item: QueueItem, span: &tracing::Span) -> rla::Result<ProcessOutcome> {
         let (repo, build_id, outcome) = match &item {
             QueueItem::GitHubStatus { payload, .. } => {
                 match self.ci.build_id_from_github_status(&payload) {
@@ -109,7 +110,7 @@ impl Worker {
                             "Ignoring invalid event (ctx: {:?}, url: {:?}).",
                             payload.context, payload.target_url
                         );
-                        return Ok(());
+                        return Ok(ProcessOutcome::Continue);
                     }
                 }
             }
@@ -125,11 +126,20 @@ impl Worker {
                             "Ignoring invalid event (app id: {:?}, url: {:?}).",
                             payload.check_run.app.id, payload.check_run.details_url
                         );
-                        return Ok(());
+                        return Ok(ProcessOutcome::Continue);
                     }
                 }
             }
-            QueueItem::GitHubPullRequest { payload, .. } => return self.process_pr(payload),
+            QueueItem::GitHubPullRequest { payload, .. } => {
+                self.process_pr(payload)?;
+                return Ok(ProcessOutcome::Continue);
+            }
+
+            QueueItem::GracefulShutdown => {
+                info!("persisting the index to disk before shutting down");
+                self.index.save(&self.index_file)?;
+                return Ok(ProcessOutcome::Exit);
+            }
         };
 
         span.record("build_id", &build_id);
@@ -154,7 +164,7 @@ impl Worker {
 
         if !outcome.is_finished() {
             info!("ignoring in-progress build");
-            return Ok(());
+            return Ok(ProcessOutcome::Continue);
         }
 
         // Avoid processing the same build multiple times.
@@ -168,7 +178,7 @@ impl Worker {
             info!("did not learn as it's not an auto build");
         }
 
-        Ok(())
+        Ok(ProcessOutcome::Continue)
     }
 
     fn report_failed(&mut self, build_id: u64, build: &dyn rla::ci::Build) -> rla::Result<()> {
@@ -397,6 +407,11 @@ impl<T: Clone + Eq + Hash> RecentlySeen<T> {
     }
 }
 
+enum ProcessOutcome {
+    Continue,
+    Exit,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;