Merge pull request #73 from vohoanglong0107/refactor-global-loop

Kobzol · web-flow · commit d9ea15a5f3bc · 2024-04-24T13:52:38.000Z
refactor: separate global event handling to a new loop
diff --git a/src/bin/bors.rs b/src/bin/bors.rs
@@ -5,8 +5,8 @@ use std::time::Duration;
 
 use anyhow::Context;
 use bors::{
-    create_app, create_bors_process, BorsContext, BorsEvent, BorsGlobalEvent, CommandParser,
-    GithubAppState, SeaORMClient, ServerState, WebhookSecret,
+    create_app, create_bors_process, BorsContext, BorsGlobalEvent, CommandParser, GithubAppState,
+    SeaORMClient, ServerState, WebhookSecret,
 };
 use clap::Parser;
 use sea_orm::Database;
@@ -79,19 +79,21 @@ fn try_main(opts: Opts) -> anyhow::Result<()> {
         ))
         .context("Cannot load GitHub repository state")?;
     let ctx = BorsContext::new(CommandParser::new(opts.cmd_prefix), Arc::new(db));
-    let (tx, bors_process) = create_bors_process(state, ctx);
+    let (repository_tx, global_tx, bors_process) = create_bors_process(state, ctx);
 
-    let refresh_tx = tx.clone();
+    let refresh_tx = global_tx.clone();
     let refresh_process = async move {
         loop {
             tokio::time::sleep(PERIODIC_REFRESH).await;
-            refresh_tx
-                .send(BorsEvent::Global(BorsGlobalEvent::Refresh))
-                .await?;
+            refresh_tx.send(BorsGlobalEvent::Refresh).await?;
         }
     };
 
-    let state = ServerState::new(tx, WebhookSecret::new(opts.webhook_secret));
+    let state = ServerState::new(
+        repository_tx,
+        global_tx,
+        WebhookSecret::new(opts.webhook_secret),
+    );
     let server_process = webhook_server(state);
 
     let fut = async move {
diff --git a/src/bors/handlers/mod.rs b/src/bors/handlers/mod.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use tracing::Instrument;
 
 use crate::bors::command::{BorsCommand, CommandParseError};
-use crate::bors::event::{BorsRepositoryEvent, PullRequestComment};
+use crate::bors::event::{BorsGlobalEvent, BorsRepositoryEvent, PullRequestComment};
 use crate::bors::handlers::help::command_help;
 use crate::bors::handlers::ping::command_ping;
 use crate::bors::handlers::refresh::refresh_repository;
@@ -16,32 +16,14 @@ use crate::database::DbClient;
 use crate::github::GithubRepoName;
 use crate::utils::logging::LogError;
 
-use super::event::{BorsEvent, BorsGlobalEvent};
-
 mod help;
 mod labels;
 mod ping;
 mod refresh;
 mod trybuild;
 mod workflow;
 
-/// This function executes a single BORS event, it is the main execution function of the bot.
-pub async fn handle_bors_event<Client: RepositoryClient>(
-    event: BorsEvent,
-    state: Arc<dyn BorsState<Client>>,
-    ctx: Arc<BorsContext>,
-) -> anyhow::Result<()> {
-    match event {
-        BorsEvent::Repository(event) => {
-            handle_bors_repository_event(event, state, ctx).await?;
-        }
-        BorsEvent::Global(event) => {
-            handle_bors_global_event(event, state, ctx).await?;
-        }
-    }
-    Ok(())
-}
-
+/// This function executes a single BORS repository event
 pub async fn handle_bors_repository_event<Client: RepositoryClient>(
     event: BorsRepositoryEvent,
     state: Arc<dyn BorsState<Client>>,
@@ -128,6 +110,7 @@ pub async fn handle_bors_repository_event<Client: RepositoryClient>(
     Ok(())
 }
 
+/// This function executes a single BORS global event
 pub async fn handle_bors_global_event<Client: RepositoryClient>(
     event: BorsGlobalEvent,
     state: Arc<dyn BorsState<Client>>,
diff --git a/src/bors/mod.rs b/src/bors/mod.rs
@@ -17,7 +17,7 @@ use crate::permissions::UserPermissions;
 pub use command::CommandParser;
 pub use comment::Comment;
 pub use context::BorsContext;
-pub use handlers::handle_bors_event;
+pub use handlers::{handle_bors_global_event, handle_bors_repository_event};
 
 /// Provides functionality for working with a remote repository.
 #[async_trait]
diff --git a/src/github/server.rs b/src/github/server.rs
@@ -1,9 +1,10 @@
 use crate::bors::event::BorsEvent;
-use crate::bors::{handle_bors_event, BorsContext, BorsState};
+use crate::bors::{handle_bors_global_event, handle_bors_repository_event, BorsContext};
 use crate::github::api::GithubAppState;
 use crate::github::webhook::GitHubWebhook;
 use crate::github::webhook::WebhookSecret;
 use crate::utils::logging::LogError;
+use crate::{BorsGlobalEvent, BorsRepositoryEvent};
 use axum::extract::State;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
@@ -17,14 +18,20 @@ use tracing::Instrument;
 
 /// Shared server state for all axum handlers.
 pub struct ServerState {
-    webhook_sender: WebhookSender,
+    repository_event_queue: mpsc::Sender<BorsRepositoryEvent>,
+    global_event_queue: mpsc::Sender<BorsGlobalEvent>,
     webhook_secret: WebhookSecret,
 }
 
 impl ServerState {
-    pub fn new(webhook_sender: WebhookSender, webhook_secret: WebhookSecret) -> Self {
+    pub fn new(
+        repository_event_queue: mpsc::Sender<BorsRepositoryEvent>,
+        global_event_queue: mpsc::Sender<BorsGlobalEvent>,
+        webhook_secret: WebhookSecret,
+    ) -> Self {
         Self {
-            webhook_sender,
+            repository_event_queue,
+            global_event_queue,
             webhook_secret,
         }
     }
@@ -48,41 +55,88 @@ pub async fn github_webhook_handler(
     State(state): State<ServerStateRef>,
     GitHubWebhook(event): GitHubWebhook,
 ) -> impl IntoResponse {
-    match state.webhook_sender.send(event).await {
-        Ok(_) => (StatusCode::OK, ""),
-        Err(err) => {
-            tracing::error!("Could not send webhook event: {err:?}");
-            (StatusCode::INTERNAL_SERVER_ERROR, "")
-        }
+    match event {
+        BorsEvent::Global(e) => match state.global_event_queue.send(e).await {
+            Ok(_) => (StatusCode::OK, ""),
+            Err(err) => {
+                tracing::error!("Could not send webhook global event: {err:?}");
+                (StatusCode::INTERNAL_SERVER_ERROR, "")
+            }
+        },
+        BorsEvent::Repository(e) => match state.repository_event_queue.send(e).await {
+            Ok(_) => (StatusCode::OK, ""),
+            Err(err) => {
+                tracing::error!("Could not send webhook repository event: {err:?}");
+                (StatusCode::INTERNAL_SERVER_ERROR, "")
+            }
+        },
     }
 }
 
-type WebhookSender = mpsc::Sender<BorsEvent>;
-
 /// Creates a future with a Bors process that continuously receives webhook events and reacts to
 /// them.
 pub fn create_bors_process(
     state: GithubAppState,
     ctx: BorsContext,
-) -> (WebhookSender, impl Future<Output = ()>) {
-    let (tx, mut rx) = mpsc::channel::<BorsEvent>(1024);
+) -> (
+    mpsc::Sender<BorsRepositoryEvent>,
+    mpsc::Sender<BorsGlobalEvent>,
+    impl Future<Output = ()>,
+) {
+    let (repository_tx, repository_rx) = mpsc::channel::<BorsRepositoryEvent>(1024);
+    let (global_tx, global_rx) = mpsc::channel::<BorsGlobalEvent>(1024);
 
     let service = async move {
-        let state: Arc<dyn BorsState<_>> = Arc::new(state);
+        let state = Arc::new(state);
         let ctx = Arc::new(ctx);
-        while let Some(event) = rx.recv().await {
-            let state = state.clone();
-            let ctx = ctx.clone();
-
-            let span = tracing::info_span!("Event");
-            tracing::debug!("Received event: {event:#?}");
-            if let Err(error) = handle_bors_event(event, state, ctx)
-                .instrument(span.clone())
-                .await
-            {
-                span.log_error(error);
+        tokio::select! {
+            _ = consume_repository_events(state.clone(), ctx.clone(), repository_rx) => {
+                tracing::warn!("Repository event handling process has ended");
+            }
+            _ = consume_global_events(state.clone(), ctx.clone(), global_rx) => {
+                tracing::warn!("Global event handling process has ended");
             }
         }
     };
-    (tx, service)
+    (repository_tx, global_tx, service)
+}
+
+async fn consume_repository_events(
+    state: Arc<GithubAppState>,
+    ctx: Arc<BorsContext>,
+    mut repository_rx: mpsc::Receiver<BorsRepositoryEvent>,
+) {
+    while let Some(event) = repository_rx.recv().await {
+        let state = state.clone();
+        let ctx = ctx.clone();
+
+        let span = tracing::info_span!("RepositoryEvent");
+        tracing::debug!("Received repository event: {event:#?}");
+        if let Err(error) = handle_bors_repository_event(event, state, ctx)
+            .instrument(span.clone())
+            .await
+        {
+            span.log_error(error);
+        }
+    }
+}
+
+async fn consume_global_events(
+    state: Arc<GithubAppState>,
+    ctx: Arc<BorsContext>,
+    mut global_rx: mpsc::Receiver<BorsGlobalEvent>,
+) {
+    while let Some(event) = global_rx.recv().await {
+        let state = state.clone();
+        let ctx = ctx.clone();
+
+        let span = tracing::info_span!("GlobalEvent");
+        tracing::debug!("Received global event: {event:#?}");
+        if let Err(error) = handle_bors_global_event(event, state, ctx)
+            .instrument(span.clone())
+            .await
+        {
+            span.log_error(error);
+        }
+    }
 }
diff --git a/src/github/webhook.rs b/src/github/webhook.rs
@@ -674,8 +674,13 @@ mod tests {
             HeaderValue::from_str(&signature).unwrap(),
         );
 
-        let (tx, _) = mpsc::channel(1024);
-        let server_ref = ServerStateRef::new(ServerState::new(tx, WebhookSecret::new(secret)));
+        let (repository_tx, _) = mpsc::channel(1024);
+        let (global_tx, _) = mpsc::channel(1024);
+        let server_ref = ServerStateRef::new(ServerState::new(
+            repository_tx,
+            global_tx,
+            WebhookSecret::new(secret),
+        ));
         GitHubWebhook::from_request(request, &server_ref).await
     }
 }
diff --git a/src/tests/state.rs b/src/tests/state.rs
@@ -16,7 +16,8 @@ use crate::bors::event::{
     WorkflowStarted,
 };
 use crate::bors::{
-    handle_bors_event, BorsContext, CheckSuite, CommandParser, Comment, RepositoryState,
+    handle_bors_global_event, handle_bors_repository_event, BorsContext, CheckSuite, CommandParser,
+    Comment, RepositoryState,
 };
 use crate::bors::{BorsState, RepositoryClient};
 use crate::config::RepositoryConfig;
@@ -67,16 +68,21 @@ impl TestBorsState {
 
     /// Execute an event.
     pub async fn event(&self, event: BorsEvent) {
-        handle_bors_event(
-            event,
-            Arc::new(self.clone()),
-            Arc::new(BorsContext::new(
-                CommandParser::new("@bors".to_string()),
-                Arc::clone(&self.db) as Arc<dyn DbClient>,
-            )),
-        )
-        .await
-        .unwrap();
+        let state = Arc::new(self.clone());
+        let ctx = Arc::new(BorsContext::new(
+            CommandParser::new("@bors".to_string()),
+            Arc::clone(&self.db) as Arc<dyn DbClient>,
+        ));
+        match event {
+            BorsEvent::Repository(event) => {
+                handle_bors_repository_event(event, state, ctx)
+                    .await
+                    .unwrap();
+            }
+            BorsEvent::Global(event) => {
+                handle_bors_global_event(event, state, ctx).await.unwrap();
+            }
+        }
     }
 
     pub async fn comment<T: Into<PullRequestComment>>(&self, comment: T) {
