Merge pull request #71 from rust-lang/pa-s3

Store index in S3, write index less times, and implement graceful shutdown

Author: pietroalbini

Cargo.lock

@@ -32,4 +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

@@ -47,3 +47,10 @@ To initialize a new index file, perform the following steps:
     * You can (temporarily) check the result directory in to the repository to see diffs.
     * Example command: `rla-offline extract-dir --ci actions -i demo.idx -s data/failed -d data/err`
     * *Note: Eventually, the expected results for the test log files will be provided in the repository and used as regression tests.*
+
+### Index file storage
+
+The index file can be stored either in the local filesystem (by providing the
+absolute or relative path to the file) or in S3 (by providing a
+`s3://{bucket}/{key}` URL). The S3 region of the bucket is detected
+automatically at startup.

Readme.md

@@ -1,6 +1,7 @@
 use crate::offline;
 use crate::rla;
 
+use rla::index::IndexStorage;
 use std::fs;
 use std::io::{self, Write};
 use std::path::Path;
@@ -31,7 +32,7 @@ fn load_lines<'a>(ci: &dyn rla::ci::CiPlatform, log: &'a [u8]) -> Vec<Line<'a>>
 
 pub fn dir(
     ci: &dyn rla::ci::CiPlatform,
-    index_file: &Path,
+    index_file: &IndexStorage,
     src_dir: &Path,
     dst_dir: &Path,
 ) -> rla::Result<()> {
@@ -91,7 +92,11 @@ pub fn dir(
     Ok(())
 }
 
-pub fn one(ci: &dyn rla::ci::CiPlatform, index_file: &Path, log_file: &Path) -> rla::Result<()> {
+pub fn one(
+    ci: &dyn rla::ci::CiPlatform,
+    index_file: &IndexStorage,
+    log_file: &Path,
+) -> rla::Result<()> {
     let config = rla::extract::Config::default();
     let index = rla::Index::load(index_file)?;
 

src/bin/offline/extract.rs

@@ -1,14 +1,15 @@
 use crate::offline;
 use crate::rla;
 
-use std::path::{Path, PathBuf};
+use rla::index::IndexStorage;
+use std::path::PathBuf;
 use std::time::Duration;
 use std::time::Instant;
 use walkdir::{self, WalkDir};
 
 pub fn learn(
     ci: &dyn rla::ci::CiPlatform,
-    index_file: &Path,
+    index_file: &IndexStorage,
     inputs: &[PathBuf],
     multiplier: u32,
 ) -> rla::Result<()> {

src/bin/offline/learn.rs

@@ -12,6 +12,7 @@ extern crate rust_log_analyzer as rla;
 extern crate walkdir;
 
 use clap::Parser;
+use rla::index::IndexStorage;
 use std::path::PathBuf;
 
 mod offline;
@@ -53,7 +54,7 @@ enum Cli {
             long = "index-file",
             help = "The index file to read / write. An existing index file is updated."
         )]
-        index_file: PathBuf,
+        index_file: IndexStorage,
         #[arg(
             short = 'm',
             long = "multiplier",
@@ -79,7 +80,7 @@ enum Cli {
             long = "index-file",
             help = "The index file to read / write."
         )]
-        index_file: PathBuf,
+        index_file: IndexStorage,
         #[arg(
             short = 's',
             long = "source",
@@ -106,7 +107,7 @@ enum Cli {
             long = "index-file",
             help = "The index file to read / write."
         )]
-        index_file: PathBuf,
+        index_file: IndexStorage,
         #[arg(help = "The log file to analyze.")]
         log: PathBuf,
     },

src/bin/rla-offline.rs

@@ -14,7 +14,10 @@ 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;
 use std::thread;
@@ -47,7 +50,7 @@ struct Cli {
         long = "index-file",
         help = "The index file to read / write. An existing index file is updated."
     )]
-    index_file: std::path::PathBuf,
+    index_file: IndexStorage,
     #[arg(
         long = "debug-post",
         help = "Post all comments to the given issue instead of the actual PR. Format: \"user/repo#id\""
@@ -90,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,
@@ -102,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);
@@ -125,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

@@ -6,13 +6,29 @@ pub use self::worker::Worker;
 mod service;
 mod worker;
 
-pub struct QueueItem {
-    pub kind: QueueItemKind,
-    pub delivery_id: String,
+pub enum QueueItem {
+    GitHubStatus {
+        payload: rla::github::CommitStatusEvent,
+        delivery_id: String,
+    },
+    GitHubCheckRun {
+        payload: rla::github::CheckRunEvent,
+        delivery_id: String,
+    },
+    GitHubPullRequest {
+        payload: rla::github::PullRequestEvent,
+        delivery_id: String,
+    },
+    GracefulShutdown,
 }
 
-pub enum QueueItemKind {
-    GitHubStatus(rla::github::CommitStatusEvent),
-    GitHubCheckRun(rla::github::CheckRunEvent),
-    GitHubPullRequest(rla::github::PullRequestEvent),
+impl QueueItem {
+    fn delivery_id(&self) -> Option<&str> {
+        match self {
+            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

@@ -1,4 +1,4 @@
-use super::{QueueItem, QueueItemKind};
+use super::QueueItem;
 
 use crate::rla;
 use anyhow::bail;
@@ -76,7 +76,7 @@ impl RlaService {
             return reply(StatusCode::BAD_REQUEST, "Missing delivery ID.\n");
         };
 
-        let item_kind = match event {
+        let item = match event {
             "status" => {
                 let payload = match serde_json::from_slice(body) {
                     Ok(p) => p,
@@ -85,7 +85,10 @@ impl RlaService {
                         return reply(StatusCode::BAD_REQUEST, "Failed to decode payload.\n");
                     }
                 };
-                QueueItemKind::GitHubStatus(payload)
+                QueueItem::GitHubStatus {
+                    payload,
+                    delivery_id,
+                }
             }
             "check_run" => {
                 let payload = match serde_json::from_slice(body) {
@@ -96,10 +99,16 @@ impl RlaService {
                     }
                 };
 
-                QueueItemKind::GitHubCheckRun(payload)
+                QueueItem::GitHubCheckRun {
+                    payload,
+                    delivery_id,
+                }
             }
             "pull_request" => match serde_json::from_slice(body) {
-                Ok(payload) => QueueItemKind::GitHubPullRequest(payload),
+                Ok(payload) => QueueItem::GitHubPullRequest {
+                    payload,
+                    delivery_id,
+                },
                 Err(err) => {
                     error!("Failed to decode 'pull_request' webhook payload: {}", err);
                     return reply(StatusCode::BAD_REQUEST, "Failed to decode payload\n");
@@ -115,10 +124,6 @@ impl RlaService {
             }
         };
 
-        let item = QueueItem {
-            kind: item_kind,
-            delivery_id,
-        };
         match self.queue.send(item) {
             Ok(()) => reply(StatusCode::OK, "Event processed.\n"),
             Err(e) => {