pr_tracking.rs

//! This module updates the PR workqueue of the Rust project contributors
//! Runs after a PR has been assigned or unassigned
//!
//! Purpose:
//!
//! - Adds the PR to the workqueue of one team member (after the PR has been assigned or reopened)
//! - Removes the PR from the workqueue of one team member (after the PR has been unassigned or closed)

use crate::github::{GitHubUser, GitHubUserType, UserId};
use crate::github::{Label, PullRequestNumber};
use crate::{
    config::ReviewPrefsConfig,
    github::{IssuesAction, IssuesEvent},
    handlers::Context,
};
use futures::TryStreamExt;
use octocrab::Octocrab;
use octocrab::models::IssueState;
use octocrab::params::pulls::Sort;
use octocrab::params::{Direction, State};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{RwLock, RwLockWriteGuard};
use tracing as log;

/// Repositories for which we track the reviewer workqueue.
pub fn get_review_tracked_repositories() -> Vec<TrackedRepository> {
    vec![
        TrackedRepository::new("rust-lang", "rust"),
        TrackedRepository::new("rust-lang", "cargo"),
        TrackedRepository::new("rust-lang", "rust-clippy"),
    ]
}

#[derive(Clone, Debug)]
pub struct AssignedPullRequest {
    pub title: String,
}

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct TrackedRepository {
    owner: String,
    name: String,
}

impl TrackedRepository {
    pub fn new(owner: &str, name: &str) -> Self {
        Self {
            owner: owner.to_owned(),
            name: name.to_owned(),
        }
    }
    fn from_full_name(name: &str) -> Option<Self> {
        let (owner, name) = name.split_once('/')?;
        Some(Self::new(owner, name))
    }

    pub fn owner(&self) -> &str {
        &self.owner
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn full_name(&self) -> String {
        format!("{}/{}", self.owner, self.name)
    }
}

/// Maps users to a set of currently assigned open non-draft pull requests in a single repository.
/// We store this map in memory, rather than in the DB, because it can get desynced when webhooks
/// are missed.
/// It is thus reloaded when triagebot starts and also periodically, so it is not needed to store it
/// in the DB.
#[derive(Debug, Default)]
pub struct ReviewerWorkqueue {
    reviewers: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>>,
}

impl ReviewerWorkqueue {
    pub fn new(
        reviewers: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>>,
    ) -> 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)
    }
}

/// Stores per-repository reviewer workqueues.
/// Each workqueue is behind its own `Arc<RwLock<...>>` so repos can be locked independently.
pub struct RepositoryWorkqueueMap {
    repos: HashMap<TrackedRepository, Arc<RwLock<ReviewerWorkqueue>>>,
}

impl RepositoryWorkqueueMap {
    pub fn new(repos: HashMap<TrackedRepository, Arc<RwLock<ReviewerWorkqueue>>>) -> Self {
        Self { repos }
    }

    pub fn get(&self, full_name: &str) -> Option<Arc<RwLock<ReviewerWorkqueue>>> {
        let repo = TrackedRepository::from_full_name(full_name)?;
        self.repos.get(&repo).cloned()
    }

    /// Returns an iterator over all repositories that are being tracked.
    pub fn tracked_repositories(
        &self,
    ) -> impl Iterator<Item = (&TrackedRepository, &Arc<RwLock<ReviewerWorkqueue>>)> {
        self.repos.iter()
    }
}

pub(super) enum ReviewPrefsInput {
    Assigned { assignee: GitHubUser },
    Unassigned { assignee: GitHubUser },
    OtherChange,
}

pub(super) async fn parse_input(
    _ctx: &Context,
    event: &IssuesEvent,
    config: Option<&ReviewPrefsConfig>,
) -> Result<Option<ReviewPrefsInput>, String> {
    // NOTE: this config check MUST exist. Else, the triagebot will emit an error
    // about this feature not being enabled
    if config.is_none() {
        return Ok(None);
    }

    // Execute this handler only if this is a PR ...
    if !event.issue.is_pr() {
        return Ok(None);
    }

    // ... and if the action is an assignment or unassignment with an assignee
    match &event.action {
        IssuesAction::Assigned { assignee } => Ok(Some(ReviewPrefsInput::Assigned {
            assignee: assignee.clone(),
        })),
        IssuesAction::Unassigned { assignee } => Ok(Some(ReviewPrefsInput::Unassigned {
            assignee: assignee.clone(),
        })),
        // We don't need to handle Opened explicitly, because that will trigger the Assigned event
        IssuesAction::Reopened
        | IssuesAction::ReadyForReview
        | IssuesAction::ConvertedToDraft
        | IssuesAction::Closed
        | IssuesAction::Deleted
        | IssuesAction::Transferred
        | IssuesAction::Labeled { .. }
        | IssuesAction::Unlabeled { .. } => Ok(Some(ReviewPrefsInput::OtherChange)),
        _ => Ok(None),
    }
}

pub(super) async fn handle_input(
    ctx: &Context,
    _config: &ReviewPrefsConfig,
    event: &IssuesEvent,
    input: ReviewPrefsInput,
) -> anyhow::Result<()> {
    log::info!("Handling event action {:?} in PR tracking", event.action);

    let pr = &event.issue;
    let pr_number = event.issue.number;

    let repo_name = &event.repository.full_name;
    let Some(workqueue_arc) = ctx.workqueue_map.get(repo_name) else {
        log::debug!("Repository {repo_name} does not have a tracked workqueue, skipping");
        return Ok(());
    };
    let mut workqueue = workqueue_arc.write().await;

    // If the PR doesn't wait for a review, remove it from the workqueue completely.
    // This handles situations such as labels being modified, which make the PR no longer to be
    // in the "waiting for a review" state, or the PR being closed/merged.
    if !waits_for_a_review(&pr.labels, &pr.assignees, &pr.user, pr.is_open(), pr.draft) {
        log::info!(
            "Removing PR {pr_number} from workqueue, because it is not waiting for a review.",
        );
        delete_pr_from_all_queues(&mut workqueue, pr_number);
        return Ok(());
    }

    let assigned_pr = AssignedPullRequest {
        title: pr.title.clone(),
    };

    match input {
        // The PR was assigned to a specific user, and it is waiting for a review.
        ReviewPrefsInput::Assigned { assignee } => {
            log::info!(
                "Adding PR {pr_number} to workqueue of {} because they were assigned.",
                assignee.login
            );

            upsert_pr_into_user_queue(&mut workqueue, assignee.id, pr_number, assigned_pr);
        }
        ReviewPrefsInput::Unassigned { assignee } => {
            log::info!(
                "Removing PR {pr_number} from workqueue of {} because they were unassigned.",
                assignee.login
            );
            delete_pr_from_user_queue(&mut workqueue, assignee.id, pr_number);
        }
        // Some other change has happened (e.g. labels changed or the PR being reopened).
        // Make sure that all assigned users have the PR in their queue.
        // When a PR is opened, it might not yet contain all the information needed to determine
        // whether it waits for a reviewer or not. For example, when you open a PR,
        // triagebot might apply the "S-waiting-on-review" (or similar) label to it, which we
        // currently use to determine whether a PR is truly assigned to someone or not.
        // We thus need to refresh the queue state after every relevant state change that we
        // receive.
        ReviewPrefsInput::OtherChange => {
            for assignee in &event.issue.assignees {
                if upsert_pr_into_user_queue(
                    &mut workqueue,
                    assignee.id,
                    pr_number,
                    assigned_pr.clone(),
                ) {
                    log::info!("Adding PR {pr_number} to workqueue of {}.", assignee.login);
                }
            }
        }
    }

    Ok(())
}

/// Loads the workqueue (mapping of open PRs assigned to users) from GitHub
pub async fn load_workqueue(
    client: &Octocrab,
    repo: &TrackedRepository,
) -> anyhow::Result<ReviewerWorkqueue> {
    tracing::debug!("Loading workqueue for {}/{}", repo.owner, repo.name);
    let prs = retrieve_pull_request_assignments(&repo.owner, &repo.name, client).await?;

    // Aggregate PRs by user
    let aggregated: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>> = prs
        .into_iter()
        .fold(HashMap::new(), |mut acc, (user, pr_number, pr)| {
            let prs = acc.entry(user.id).or_default();
            prs.insert(pr_number, pr);
            acc
        });
    tracing::debug!("PR assignments for `{}`:\n{aggregated:?}", repo.full_name());
    Ok(ReviewerWorkqueue::new(aggregated))
}

/// Retrieve tuples of (user, PR number) where
/// the given user is assigned as a reviewer for that PR
/// and the PR is considered to be "waiting for a review", according to the semantics
/// of the reviewer workqueue.
/// See the [`waits_for_a_review`] function.
pub async fn retrieve_pull_request_assignments(
    owner: &str,
    repository: &str,
    client: &Octocrab,
) -> anyhow::Result<Vec<(GitHubUser, PullRequestNumber, AssignedPullRequest)>> {
    let mut assignments = vec![];

    // We use the REST API to fetch open pull requests, as it is much (~5-10x)
    // faster than using GraphQL here.
    let stream = client
        .pulls(owner, repository)
        .list()
        .state(State::Open)
        .direction(Direction::Ascending)
        .sort(Sort::Created)
        .per_page(100)
        .send()
        .await?
        .into_stream(client);
    let mut stream = std::pin::pin!(stream);
    while let Some(pr) = stream.try_next().await? {
        let labels = pr
            .labels
            .unwrap_or_default()
            .into_iter()
            .map(|l| Label { name: l.name })
            .collect::<Vec<Label>>();
        let assignees = pr
            .assignees
            .as_ref()
            .map(|authors| authors.iter().map(GitHubUser::from).collect::<Vec<_>>())
            .unwrap_or_default();
        let author = pr
            .user
            .as_ref()
            .map(|author| GitHubUser::from(author.as_ref()))
            .unwrap_or_else(|| GitHubUser {
                login: "ghost".to_string(),
                id: 0,
                r#type: GitHubUserType::Bot,
            });
        if waits_for_a_review(
            &labels,
            &assignees,
            &author,
            pr.state == Some(IssueState::Open),
            pr.draft.unwrap_or_default(),
        ) {
            for user in pr.assignees.unwrap_or_default() {
                assignments.push((
                    Into::<GitHubUser>::into(&user),
                    pr.number,
                    AssignedPullRequest {
                        title: pr.title.clone().unwrap_or_default(),
                    },
                ));
            }
        }
    }
    assignments.sort_by(|a, b| a.0.id.cmp(&b.0.id));

    Ok(assignments)
}

/// Get pull request assignments for a team member in a specific repo.
pub async fn get_assigned_prs(
    ctx: &Context,
    repo: &str,
    user_id: UserId,
) -> HashMap<PullRequestNumber, AssignedPullRequest> {
    let Some(workqueue) = ctx.workqueue_map.get(repo) else {
        return Default::default();
    };
    workqueue
        .read()
        .await
        .reviewers
        .get(&user_id)
        .cloned()
        .unwrap_or_default()
}

/// Add a PR to the workqueue of a team member.
/// Updates data of the pull request if it already was in the workqueue.
/// Ensures no accidental PR duplicates.
///
/// Returns true if the PR was actually inserted.
fn upsert_pr_into_user_queue(
    workqueue: &mut RwLockWriteGuard<ReviewerWorkqueue>,
    user_id: UserId,
    pr: PullRequestNumber,
    assigned_pr: AssignedPullRequest,
) -> bool {
    workqueue
        .reviewers
        .entry(user_id)
        .or_default()
        .insert(pr, assigned_pr)
        .is_none()
}

/// Delete a PR from the workqueue of a team member.
fn delete_pr_from_user_queue(
    workqueue: &mut ReviewerWorkqueue,
    user_id: UserId,
    pr: PullRequestNumber,
) {
    if let Some(queue) = workqueue.reviewers.get_mut(&user_id) {
        queue.remove(&pr);
    }
}

/// Delete a PR from the workqueue completely.
fn delete_pr_from_all_queues(workqueue: &mut ReviewerWorkqueue, pr: PullRequestNumber) {
    for queue in workqueue.reviewers.values_mut() {
        queue.retain(|pr_number, _| *pr_number != pr);
    }
}

/// Returns true if the workqueue should assume that this PR is actually waiting for a reviewer.
/// The function receives atomic attributes so that it is compatible both with triagebot's
/// `Issue` struct (used for incremental updates) and octocrab's `PullRequest` struct (used for
/// batch PR loads).
///
/// Note: this functionality is currently hardcoded for rust-lang/rust, other repos might use
/// different labels.
fn waits_for_a_review(
    labels: &[Label],
    assignees: &[GitHubUser],
    author: &GitHubUser,
    is_open: bool,
    is_draft: bool,
) -> bool {
    let is_blocked = labels
        .iter()
        .any(|l| l.name == "S-blocked" || l.name == "S-inactive");
    let is_rollup = labels.iter().any(|l| l.name == "rollup");
    let is_waiting_for_reviewer = labels.iter().any(|l| l.name == "S-waiting-on-review");
    let is_assigned_to_author = assignees.contains(author);
    let has_capacity_tracking_opt_out = labels
        .iter()
        .any(|l| l.name == "S-no-work-capacity-tracking");

    is_open
        && !is_draft
        && !is_blocked
        && !is_rollup
        && is_waiting_for_reviewer
        && !is_assigned_to_author
        && !has_capacity_tracking_opt_out
}

#[cfg(test)]
mod tests {
    use crate::config::Config;
    use crate::github::{GitHubUser, Issue, IssuesAction, IssuesEvent, Repository};
    use crate::github::{Label, PullRequestNumber};
    use crate::handlers::pr_tracking::{
        AssignedPullRequest, handle_input, parse_input, upsert_pr_into_user_queue,
    };
    use crate::tests::github::{default_test_user, issue, pull_request, user};
    use crate::tests::{TestContext, run_db_test};

    #[tokio::test]
    async fn add_pr_to_workqueue_on_assign() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);

            run_handler(
                &ctx,
                IssuesAction::Assigned {
                    assignee: user.clone(),
                },
                pull_request()
                    .number(10)
                    .labels(vec!["S-waiting-on-review"])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[10]).await;

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn ignore_blocked_pr() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);

            run_handler(
                &ctx,
                IssuesAction::Assigned {
                    assignee: user.clone(),
                },
                pull_request()
                    .labels(vec!["S-waiting-on-review", "S-blocked"])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[]).await;

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn remove_pr_from_workqueue_on_unassign() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);
            set_assigned_prs(&ctx, &user, &[10]).await;

            run_handler(
                &ctx,
                IssuesAction::Unassigned {
                    assignee: user.clone(),
                },
                pull_request()
                    .number(10)
                    .labels(vec!["S-waiting-on-review"])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[]).await;

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn add_pr_to_workqueue_on_label() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);

            run_handler(
                &ctx,
                IssuesAction::Assigned {
                    assignee: user.clone(),
                },
                pull_request().number(10).call(),
            )
            .await;
            check_assigned_prs(&ctx, &user, &[]).await;

            run_handler(
                &ctx,
                IssuesAction::Labeled {
                    label: Label {
                        name: "S-waiting-on-review".to_string(),
                    },
                },
                pull_request()
                    .number(10)
                    .labels(vec!["S-waiting-on-review"])
                    .assignees(vec![user.clone()])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[10]).await;

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn remove_pr_from_workqueue_on_pr_closed() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);
            set_assigned_prs(&ctx, &user, &[10]).await;

            run_handler(
                &ctx,
                IssuesAction::Closed,
                pull_request()
                    .number(10)
                    .assignees(vec![user.clone()])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[]).await;

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn add_pr_to_workqueue_on_pr_reopen() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);
            set_assigned_prs(&ctx, &user, &[42]).await;

            run_handler(
                &ctx,
                IssuesAction::Reopened,
                pull_request()
                    .number(10)
                    .labels(vec!["S-waiting-on-review"])
                    .assignees(vec![user.clone()])
                    .call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[10, 42]).await;

            Ok(ctx)
        })
        .await;
    }

    // Make sure that we only consider pull requests, not issues.
    #[tokio::test]
    async fn ignore_issue_assignments() {
        run_db_test(|ctx| async move {
            let user = user("Martin", 2);

            run_handler(
                &ctx,
                IssuesAction::Assigned {
                    assignee: user.clone(),
                },
                issue().number(10).call(),
            )
            .await;

            check_assigned_prs(&ctx, &user, &[]).await;

            Ok(ctx)
        })
        .await;
    }

    const TEST_REPO: &str = "rust-lang-test/triagebot-test";

    async fn check_assigned_prs(
        ctx: &TestContext,
        user: &GitHubUser,
        expected_prs: &[PullRequestNumber],
    ) {
        let workqueue_arc = ctx
            .handler_ctx()
            .workqueue_map
            .get(TEST_REPO)
            .expect("test repo workqueue should exist");
        let mut assigned = workqueue_arc
            .read()
            .await
            .reviewers
            .get(&user.id)
            .cloned()
            .unwrap_or_default()
            .into_keys()
            .collect::<Vec<_>>();
        assigned.sort();
        assert_eq!(assigned, expected_prs);
    }

    async fn set_assigned_prs(ctx: &TestContext, user: &GitHubUser, prs: &[PullRequestNumber]) {
        {
            let workqueue_arc = ctx
                .handler_ctx()
                .workqueue_map
                .get(TEST_REPO)
                .expect("test repo workqueue should exist");
            let mut workqueue = workqueue_arc.write().await;
            for &pr in prs {
                upsert_pr_into_user_queue(
                    &mut workqueue,
                    user.id,
                    pr,
                    AssignedPullRequest {
                        title: format!("PR {pr}"),
                    },
                );
            }
        }
        check_assigned_prs(&ctx, user, prs).await;
    }

    async fn run_handler(ctx: &TestContext, action: IssuesAction, issue: Issue) {
        let handler_ctx = ctx.handler_ctx();
        let config = create_config().pr_tracking;

        let event = IssuesEvent {
            action,
            issue,
            changes: None,
            before: None,
            after: None,
            repository: Repository {
                full_name: "rust-lang-test/triagebot-test".to_string(),
                default_branch: "main".to_string(),
                fork: false,
                parent: None,
            },
            sender: default_test_user(),
        };

        let input = parse_input(&handler_ctx, &event, config.as_ref())
            .await
            .unwrap();
        if let Some(input) = input {
            handle_input(&handler_ctx, &config.unwrap(), &event, input)
                .await
                .unwrap()
        }
    }

    fn create_config() -> Config {
        toml::from_str::<Config>(
            r#"
[pr-tracking]
"#,
        )
        .unwrap()
    }
}