refactor: remove GithubAppState

Author: vohoanglong0107

src/bin/bors.rs

@@ -5,8 +5,8 @@ use std::time::Duration;
 
 use anyhow::Context;
 use bors::{
-    create_app, create_bors_process, BorsContext, BorsGlobalEvent, CommandParser, GithubAppState,
-    SeaORMClient, ServerState, WebhookSecret,
+    create_app, create_bors_process, create_github_client, BorsContext, BorsGlobalEvent,
+    CommandParser, SeaORMClient, ServerState, WebhookSecret,
 };
 use clap::Parser;
 use sea_orm::Database;
@@ -72,14 +72,14 @@ fn try_main(opts: Opts) -> anyhow::Result<()> {
         .block_on(initialize_db(&opts.db))
         .context("Cannot initialize database")?;
 
-    let state = runtime
-        .block_on(GithubAppState::load(
-            opts.app_id.into(),
-            opts.private_key.into_bytes().into(),
-        ))
-        .context("Cannot load GitHub repository state")?;
-    let ctx = BorsContext::new(CommandParser::new(opts.cmd_prefix), Arc::new(db));
-    let (repository_tx, global_tx, bors_process) = create_bors_process(state, ctx);
+    let client = create_github_client(opts.app_id.into(), opts.private_key.into_bytes().into())?;
+
+    let ctx = runtime.block_on(BorsContext::new(
+        CommandParser::new(opts.cmd_prefix),
+        Arc::new(db),
+        Arc::new(client),
+    ));
+    let (repository_tx, global_tx, bors_process) = create_bors_process(ctx);
 
     let refresh_tx = global_tx.clone();
     let refresh_process = async move {

src/bors/context.rs

@@ -1,13 +1,32 @@
-use crate::{bors::command::CommandParser, database::DbClient};
-use std::sync::Arc;
+use arc_swap::ArcSwap;
 
-pub struct BorsContext {
+use crate::{bors::command::CommandParser, database::DbClient, github::GithubRepoName};
+use std::{collections::HashMap, sync::Arc};
+
+use super::{GlobalClient, RepositoryClient, RepositoryState};
+
+pub struct BorsContext<Client: RepositoryClient> {
     pub parser: CommandParser,
     pub db: Arc<dyn DbClient>,
+    pub global_client: Arc<dyn GlobalClient<Client>>,
+    pub repositories: Arc<ArcSwap<HashMap<GithubRepoName, Arc<RepositoryState<Client>>>>>,
 }
 
-impl BorsContext {
-    pub fn new(parser: CommandParser, db: Arc<dyn DbClient>) -> Self {
-        Self { parser, db }
+impl<Client: RepositoryClient> BorsContext<Client> {
+    pub async fn new(
+        parser: CommandParser,
+        db: Arc<dyn DbClient>,
+        global_client: Arc<dyn GlobalClient<Client>>,
+    ) -> Self {
+        // this unwrap is making me nervous, but if lhe repos loading
+        // fails we might as well restart the bot
+        let repositories = global_client.load_repositories().await.unwrap();
+        let repositories = Arc::new(ArcSwap::new(Arc::new(repositories)));
+        Self {
+            parser,
+            db,
+            global_client,
+            repositories,
+        }
     }
 }

src/bors/handlers/mod.rs

@@ -1,4 +1,6 @@
 use anyhow::Context;
+use arc_swap::ArcSwap;
+use std::collections::HashMap;
 use std::sync::Arc;
 use tracing::Instrument;
 
@@ -11,9 +13,8 @@ use crate::bors::handlers::trybuild::{command_try_build, command_try_cancel, TRY
 use crate::bors::handlers::workflow::{
     handle_check_suite_completed, handle_workflow_completed, handle_workflow_started,
 };
-use crate::bors::{BorsContext, BorsState, Comment, RepositoryClient, RepositoryState};
+use crate::bors::{BorsContext, Comment, RepositoryClient, RepositoryState};
 use crate::database::DbClient;
-use crate::github::GithubRepoName;
 use crate::utils::logging::LogError;
 
 mod help;
@@ -26,11 +27,15 @@ mod workflow;
 /// This function executes a single BORS repository event
 pub async fn handle_bors_repository_event<Client: RepositoryClient>(
     event: BorsRepositoryEvent,
-    state: Arc<dyn BorsState<Client>>,
-    ctx: Arc<BorsContext>,
+    ctx: Arc<BorsContext<Client>>,
 ) -> anyhow::Result<()> {
     let db = Arc::clone(&ctx.db);
-    let Some(repo) = get_repo_state(state, event.repository()) else {
+    let Some(repo) = ctx
+        .repositories
+        .load()
+        .get(event.repository())
+        .map(Arc::clone)
+    else {
         return Err(anyhow::anyhow!(
             "Repository {} not found in the bot state",
             event.repository()
@@ -113,20 +118,27 @@ pub async fn handle_bors_repository_event<Client: RepositoryClient>(
 /// This function executes a single BORS global event
 pub async fn handle_bors_global_event<Client: RepositoryClient>(
     event: BorsGlobalEvent,
-    state: Arc<dyn BorsState<Client>>,
-    ctx: Arc<BorsContext>,
+    ctx: Arc<BorsContext<Client>>,
 ) -> anyhow::Result<()> {
     let db = Arc::clone(&ctx.db);
     match event {
         BorsGlobalEvent::InstallationsChanged => {
-            let span = tracing::info_span!("Repository reload");
-            if let Err(error) = state.reload_repositories().instrument(span.clone()).await {
-                span.log_error(error);
+            let span = tracing::info_span!("Repository reload").entered();
+
+            match ctx.global_client.load_repositories().await {
+                Ok(repos) => {
+                    ctx.repositories.store(Arc::new(repos));
+                }
+                Err(error) => {
+                    span.log_error(error);
+                }
             }
+            span.exit();
         }
         BorsGlobalEvent::Refresh => {
             let span = tracing::info_span!("Refresh");
-            let repos = state.get_all_repos();
+            let repos: Vec<Arc<RepositoryState<Client>>> =
+                ctx.repositories.load().values().cloned().collect();
             futures::future::join_all(repos.into_iter().map(|repo| {
                 let repo = Arc::clone(&repo);
                 async {
@@ -143,23 +155,10 @@ pub async fn handle_bors_global_event<Client: RepositoryClient>(
     Ok(())
 }
 
-fn get_repo_state<Client: RepositoryClient>(
-    state: Arc<dyn BorsState<Client>>,
-    repo: &GithubRepoName,
-) -> Option<Arc<RepositoryState<Client>>> {
-    match state.get_repo_state(repo) {
-        Some(result) => Some(result),
-        None => {
-            tracing::warn!("Repository {} not found", repo);
-            None
-        }
-    }
-}
-
 async fn handle_comment<Client: RepositoryClient>(
     repo: Arc<RepositoryState<Client>>,
     database: Arc<dyn DbClient>,
-    ctx: Arc<BorsContext>,
+    ctx: Arc<BorsContext<Client>>,
     comment: PullRequestComment,
 ) -> anyhow::Result<()> {
     let pr_number = comment.pr_number;

src/bors/mod.rs

@@ -4,6 +4,7 @@ mod context;
 pub mod event;
 mod handlers;
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use arc_swap::ArcSwap;
@@ -70,6 +71,16 @@ pub trait RepositoryClient: Send + Sync {
     fn get_workflow_url(&self, run_id: RunId) -> String;
 }
 
+/// Temporary trait to sastify the test mocking.
+/// TODO: Remove this trait once we move to mock REST API call.
+#[async_trait]
+pub trait GlobalClient<Client: RepositoryClient>: Send + Sync {
+    /// Load state of repositories.
+    async fn load_repositories(
+        &self,
+    ) -> anyhow::Result<HashMap<GithubRepoName, Arc<RepositoryState<Client>>>>;
+}
+
 #[derive(Clone)]
 pub enum CheckSuiteStatus {
     Pending,
@@ -84,20 +95,6 @@ pub struct CheckSuite {
     pub(crate) status: CheckSuiteStatus,
 }
 
-/// Main state holder for the bot.
-/// It is behind a trait to allow easier mocking in tests.
-#[async_trait]
-pub trait BorsState<Client: RepositoryClient>: Send + Sync {
-    /// Get repository and database state for the given repository name.
-    fn get_repo_state(&self, repo: &GithubRepoName) -> Option<Arc<RepositoryState<Client>>>;
-
-    /// Get all repositories.
-    fn get_all_repos(&self) -> Vec<Arc<RepositoryState<Client>>>;
-
-    /// Reload state of repositories due to some external change.
-    async fn reload_repositories(&self) -> anyhow::Result<()>;
-}
-
 /// An access point to a single repository.
 /// Can be used to query permissions for the repository, and also to perform various
 /// actions using the stored client.

src/github/api/client.rs

@@ -1,16 +1,23 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
 use anyhow::Context;
 use axum::async_trait;
 use octocrab::models::{App, Repository, RunId};
 use octocrab::{Error, Octocrab};
 use tracing::log;
 
 use crate::bors::event::PullRequestComment;
-use crate::bors::{CheckSuite, CheckSuiteStatus, Comment, RepositoryClient};
+use crate::bors::{
+    CheckSuite, CheckSuiteStatus, Comment, GlobalClient, RepositoryClient, RepositoryState,
+};
 use crate::config::{RepositoryConfig, CONFIG_FILE_PATH};
 use crate::github::api::base_github_html_url;
 use crate::github::api::operations::{merge_branches, set_branch_to_commit, MergeError};
 use crate::github::{Branch, CommitSha, GithubRepoName, PullRequest, PullRequestNumber};
 
+use super::load_repositories;
+
 /// Provides access to a single app installation (repository) using the GitHub API.
 pub struct GithubRepositoryClient {
     pub app: App,
@@ -263,6 +270,15 @@ impl RepositoryClient for GithubRepositoryClient {
     }
 }
 
+#[async_trait]
+impl GlobalClient<GithubRepositoryClient> for Octocrab {
+    async fn load_repositories(
+        &self,
+    ) -> anyhow::Result<HashMap<GithubRepoName, Arc<RepositoryState<GithubRepositoryClient>>>> {
+        load_repositories(self).await
+    }
+}
+
 fn github_pr_to_pr(pr: octocrab::models::pulls::PullRequest) -> PullRequest {
     PullRequest {
         number: pr.number.into(),

src/github/api/mod.rs

@@ -3,14 +3,13 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use arc_swap::ArcSwap;
-use axum::async_trait;
 use octocrab::models::{App, AppId, InstallationRepositories, Repository};
 use octocrab::Octocrab;
 use secrecy::{ExposeSecret, SecretVec};
 
 use client::GithubRepositoryClient;
 
-use crate::bors::{BorsState, RepositoryClient, RepositoryState};
+use crate::bors::{RepositoryClient, RepositoryState};
 use crate::github::GithubRepoName;
 use crate::permissions::load_permissions;
 
@@ -19,39 +18,24 @@ pub(crate) mod operations;
 
 type GithubRepositoryState = RepositoryState<GithubRepositoryClient>;
 
-type RepositoryMap = HashMap<GithubRepoName, Arc<GithubRepositoryState>>;
+type GithubRepositoryMap = HashMap<GithubRepoName, Arc<GithubRepositoryState>>;
 
 fn base_github_html_url() -> &'static str {
     "https://github.com"
 }
 
-/// Provides access to managed GitHub repositories.
-pub struct GithubAppState {
-    client: Octocrab,
-    repositories: ArcSwap<RepositoryMap>,
-}
+pub fn create_github_client(app_id: AppId, private_key: SecretVec<u8>) -> anyhow::Result<Octocrab> {
+    let key = jsonwebtoken::EncodingKey::from_rsa_pem(private_key.expose_secret().as_ref())
+        .context("Could not encode private key")?;
 
-impl GithubAppState {
-    /// Loads repositories managed by the Bors GitHub app with the given ID.
-    pub async fn load(app_id: AppId, private_key: SecretVec<u8>) -> anyhow::Result<GithubAppState> {
-        let key = jsonwebtoken::EncodingKey::from_rsa_pem(private_key.expose_secret().as_ref())
-            .context("Could not encode private key")?;
-
-        let client = Octocrab::builder()
-            .app(app_id, key)
-            .build()
-            .context("Could not create octocrab builder")?;
-
-        let repositories = load_repositories(&client).await?;
-        Ok(GithubAppState {
-            client,
-            repositories: ArcSwap::new(Arc::new(repositories)),
-        })
-    }
+    Octocrab::builder()
+        .app(app_id, key)
+        .build()
+        .context("Could not create octocrab builder")
 }
 
 /// Loads repositories that are connected to the given GitHub App client.
-pub async fn load_repositories(client: &Octocrab) -> anyhow::Result<RepositoryMap> {
+pub async fn load_repositories(client: &Octocrab) -> anyhow::Result<GithubRepositoryMap> {
     let installations = client
         .apps()
         .installations()
@@ -179,27 +163,3 @@ async fn create_repo_state(
         permissions: ArcSwap::new(Arc::new(permissions)),
     })
 }
-
-#[async_trait]
-impl BorsState<GithubRepositoryClient> for GithubAppState {
-    fn get_repo_state(
-        &self,
-        repo: &GithubRepoName,
-    ) -> Option<Arc<RepositoryState<GithubRepositoryClient>>> {
-        self.repositories
-            .load()
-            .get(repo)
-            .map(|repo| Arc::clone(&repo))
-    }
-
-    fn get_all_repos(&self) -> Vec<Arc<RepositoryState<GithubRepositoryClient>>> {
-        self.repositories.load().values().cloned().collect()
-    }
-
-    /// Re-download information about repositories connected to this GitHub app.
-    async fn reload_repositories(&self) -> anyhow::Result<()> {
-        self.repositories
-            .store(Arc::new(load_repositories(&self.client).await?));
-        Ok(())
-    }
-}

src/github/mod.rs

@@ -11,7 +11,6 @@ pub mod server;
 mod webhook;
 
 pub use api::operations::MergeError;
-pub use api::GithubAppState;
 pub use labels::{LabelModification, LabelTrigger};
 pub use webhook::WebhookSecret;