Take review preferences into account when assigning reviews

Author: Kobzol

src/db/review_prefs.rs

@@ -2,6 +2,7 @@ use crate::db::users::record_username;
 use crate::github::{User, UserId};
 use anyhow::Context;
 use serde::Serialize;
+use std::collections::HashMap;
 
 #[derive(Debug, Serialize)]
 pub struct ReviewPrefs {
@@ -37,6 +38,50 @@ WHERE review_prefs.user_id = $1;";
     Ok(row.map(|r| r.into()))
 }
 
+/// Returns a set of review preferences for all passed usernames.
+/// Usernames are matched regardless of case.
+///
+/// Usernames that are not present in the resulting map have no review preferences configured
+/// in the database.
+pub async fn get_review_prefs_batch<'a>(
+    db: &tokio_postgres::Client,
+    users: &[&'a str],
+) -> anyhow::Result<HashMap<&'a str, ReviewPrefs>> {
+    // We need to make sure that we match users regardless of case, but at the
+    // same time we need to return the originally-cased usernames in the final hashmap.
+    // At the same time, we can't depend on the order of results returned by the DB.
+    // So we need to do some additional bookkeeping here.
+    let lowercase_map: HashMap<String, &str> = users
+        .iter()
+        .map(|name| (name.to_lowercase(), *name))
+        .collect();
+    let lowercase_users: Vec<&str> = lowercase_map.keys().map(|s| s.as_str()).collect();
+
+    // The id/user_id/max_assigned_prs columns have to match the names used in
+    // `From<tokio_postgres::row::Row> for ReviewPrefs`.
+    let query = "
+SELECT lower(u.username) AS username, r.id AS id, r.user_id AS user_id, r.max_assigned_prs AS max_assigned_prs
+FROM review_prefs AS r
+JOIN users AS u ON u.user_id = r.user_id
+WHERE lower(u.username) = ANY($1);";
+
+    Ok(db
+        .query(query, &[&lowercase_users])
+        .await
+        .context("Error retrieving review preferences from usernames")?
+        .into_iter()
+        .map(|row| {
+            // Map back from the lowercase username to the original username.
+            let username_lower: &str = row.get("username");
+            let username = lowercase_map
+                .get(username_lower)
+                .expect("Lowercase username not found");
+            let prefs: ReviewPrefs = row.into();
+            (*username, prefs)
+        })
+        .collect())
+}
+
 /// Updates review preferences of the specified user, or creates them
 /// if they do not exist yet.
 pub async fn upsert_review_prefs(

src/handlers/assign.rs

@@ -20,6 +20,8 @@
 //! `assign.owners` config, it will auto-select an assignee based on the files
 //! the PR modifies.
 
+use crate::db::review_prefs::get_review_prefs_batch;
+use crate::github::UserId;
 use crate::handlers::pr_tracking::ReviewerWorkqueue;
 use crate::{
     config::AssignConfig,
@@ -345,7 +347,9 @@ async fn determine_assignee(
                     e @ FindReviewerError::NoReviewer { .. }
                     | e @ FindReviewerError::ReviewerIsPrAuthor { .. }
                     | e @ FindReviewerError::ReviewerAlreadyAssigned { .. }
-                    | e @ FindReviewerError::ReviewerOnVacation { .. },
+                    | e @ FindReviewerError::ReviewerOnVacation { .. }
+                    | e @ FindReviewerError::DatabaseError(_)
+                    | e @ FindReviewerError::ReviewerAtMaxCapacity { .. },
                 ) => log::trace!(
                     "no reviewer could be determined for PR {}: {e}",
                     event.issue.global_id()
@@ -675,6 +679,10 @@ enum FindReviewerError {
     ReviewerIsPrAuthor { username: String },
     /// Requested reviewer is already assigned to that PR
     ReviewerAlreadyAssigned { username: String },
+    /// Data required for assignment could not be loaded from the DB.
+    DatabaseError(String),
+    /// The reviewer has too many PRs alreayd assigned.
+    ReviewerAtMaxCapacity { username: String },
 }
 
 impl std::error::Error for FindReviewerError {}
@@ -717,6 +725,17 @@ impl fmt::Display for FindReviewerError {
                     REVIEWER_ALREADY_ASSIGNED.replace("{username}", username)
                 )
             }
+            FindReviewerError::DatabaseError(error) => {
+                write!(f, "Database error: {error}")
+            }
+            FindReviewerError::ReviewerAtMaxCapacity { username } => {
+                write!(
+                    f,
+                    r"`{username}` has too many PRs assigned to them.
+
+Please select a different reviewer.",
+                )
+            }
         }
     }
 }
@@ -728,7 +747,7 @@ impl fmt::Display for FindReviewerError {
 /// auto-assign groups, or rust-lang team names. It must have at least one
 /// entry.
 async fn find_reviewer_from_names(
-    _db: &DbClient,
+    db: &DbClient,
     workqueue: Arc<RwLock<ReviewerWorkqueue>>,
     teams: &Teams,
     config: &AssignConfig,
@@ -742,7 +761,10 @@ async fn find_reviewer_from_names(
         }
     }
 
-    let candidates = candidate_reviewers_from_names(workqueue, teams, config, issue, names)?;
+    let candidates =
+        candidate_reviewers_from_names(db, workqueue, teams, config, issue, names).await?;
+    assert!(!candidates.is_empty());
+
     // This uses a relatively primitive random choice algorithm.
     // GitHub's CODEOWNERS supports much more sophisticated options, such as:
     //
@@ -846,20 +868,23 @@ fn expand_teams_and_groups(
 
 /// Returns a list of candidate usernames (from relevant teams) to choose as a reviewer.
 /// If not reviewer is available, returns an error.
-fn candidate_reviewers_from_names<'a>(
+async fn candidate_reviewers_from_names<'a>(
+    db: &DbClient,
     workqueue: Arc<RwLock<ReviewerWorkqueue>>,
     teams: &'a Teams,
     config: &'a AssignConfig,
     issue: &Issue,
     names: &'a [String],
 ) -> Result<HashSet<String>, FindReviewerError> {
+    // Step 1: expand teams and groups into candidate names
     let (expanded, expansion_happened) = expand_teams_and_groups(teams, issue, config, names)?;
     let expanded_count = expanded.len();
 
     // Set of candidate usernames to choose from.
     // We go through each expanded candidate and store either success or an error for them.
     let mut candidates: Vec<Result<String, FindReviewerError>> = Vec::new();
 
+    // Step 2: pre-filter candidates based on checks that we can perform quickly
     for candidate in expanded {
         let name_lower = candidate.to_lowercase();
         let is_pr_author = name_lower == issue.user.login.to_lowercase();
@@ -896,9 +921,50 @@ fn candidate_reviewers_from_names<'a>(
     }
     assert_eq!(candidates.len(), expanded_count);
 
-    let valid_candidates: HashSet<String> = candidates
+    if config.review_prefs.is_some() {
+        // Step 3: gather potential usernames to form a DB query for review preferences
+        let usernames: Vec<String> = candidates
+            .iter()
+            .filter_map(|res| res.as_deref().ok().map(|s| s.to_string()))
+            .collect();
+        let usernames: Vec<&str> = usernames.iter().map(|s| s.as_str()).collect();
+        let review_prefs = get_review_prefs_batch(db, &usernames)
+            .await
+            .context("cannot fetch review preferences")
+            .map_err(|e| FindReviewerError::DatabaseError(e.to_string()))?;
+
+        let workqueue = workqueue.read().await;
+
+        // Step 4: check review preferences
+        candidates = candidates
+            .into_iter()
+            .map(|username| {
+                // Only consider candidates that did not have an earlier error
+                let username = username?;
+
+                // If no review prefs were found, we assume the default unlimited
+                // review capacity.
+                let Some(review_prefs) = review_prefs.get(username.as_str()) else {
+                    return Ok(username);
+                };
+                let Some(capacity) = review_prefs.max_assigned_prs else {
+                    return Ok(username);
+                };
+                let assigned_prs = workqueue.assigned_pr_count(review_prefs.user_id as UserId);
+                // Can we assign one more PR?
+                if (assigned_prs as i32) < capacity {
+                    Ok(username)
+                } else {
+                    Err(FindReviewerError::ReviewerAtMaxCapacity { username })
+                }
+            })
+            .collect();
+    }
+    assert_eq!(candidates.len(), expanded_count);
+
+    let valid_candidates: HashSet<&str> = candidates
         .iter()
-        .filter_map(|res| res.as_ref().ok().cloned())
+        .filter_map(|res| res.as_deref().ok())
         .collect();
 
     if valid_candidates.is_empty() {
@@ -925,6 +991,9 @@ fn candidate_reviewers_from_names<'a>(
             })
         }
     } else {
-        Ok(valid_candidates)
+        Ok(valid_candidates
+            .into_iter()
+            .map(|s| s.to_string())
+            .collect())
     }
 }

src/handlers/pr_tracking.rs

@@ -36,6 +36,13 @@ impl ReviewerWorkqueue {
     pub fn new(reviewers: HashMap<UserId, HashSet<PullRequestNumber>>) -> Self {
         Self { reviewers }
     }
+
+    pub fn assigned_pr_count(&self, user_id: UserId) -> u64 {
+        self.reviewers
+            .get(&user_id)
+            .map(|prs| prs.len() as u64)
+            .unwrap_or(0)
+    }
 }
 
 pub(super) enum ReviewPrefsInput {