/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BTreeMap;
use std::collections::HashSet;

use ::sql::Connection;
use ::sql::Transaction;
use ::sql_ext::mononoke_queries;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::hash::GitSha1;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::SqlConnections;
use stats::prelude::*;

use crate::errors::AddGitMappingErrorKind;
use crate::BonsaiGitMapping;
use crate::BonsaiGitMappingEntry;
use crate::BonsaisOrGitShas;

define_stats! {
    prefix = "mononoke.bonsai_git_mapping";
    gets: timeseries(Rate, Sum),
    gets_master: timeseries(Rate, Sum),
    adds: timeseries(Rate, Sum),
}

pub struct SqlBonsaiGitMapping {
    connections: SqlConnections,
    repo_id: RepositoryId,
}

mononoke_queries! {
    write InsertMapping(values: (
        repo_id: RepositoryId,
        git_sha1: GitSha1,
        bcs_id: ChangesetId,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bonsai_git_mapping (repo_id, git_sha1, bcs_id) VALUES {values}"
    }
    read SelectMappingByBonsai(
        repo_id: RepositoryId,
        >list bcs_id: ChangesetId
    ) -> (GitSha1, ChangesetId) {
        "SELECT git_sha1, bcs_id
         FROM bonsai_git_mapping
         WHERE repo_id = {repo_id}
           AND bcs_id IN {bcs_id}"
    }

    read SelectMappingByGitSha1(
        repo_id: RepositoryId,
        >list git_sha1: GitSha1
    ) -> (GitSha1, ChangesetId) {
        "SELECT git_sha1, bcs_id
         FROM bonsai_git_mapping
         WHERE repo_id = {repo_id}
           AND git_sha1 IN {git_sha1}"
    }

    read SelectGitSha1sByRange(
        repo_id: RepositoryId,
        git_sha1_min: GitSha1,
        git_sha1_max: GitSha1,
        limit: usize
    ) -> (GitSha1) {
        "SELECT git_sha1
         FROM bonsai_git_mapping
         WHERE repo_id = {repo_id}
            AND git_sha1 >= {git_sha1_min} AND git_sha1 <= {git_sha1_max}
            LIMIT {limit}
        "
    }
}

#[async_trait]
impl BonsaiGitMapping for SqlBonsaiGitMapping {
    async fn bulk_add(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGitMappingEntry],
    ) -> Result<(), AddGitMappingErrorKind> {
        let txn = self
            .connections
            .write_connection
            .start_transaction()
            .await?;
        let txn = self
            .bulk_add_git_mapping_in_transaction(ctx, entries, txn)
            .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn bulk_add_git_mapping_in_transaction(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGitMappingEntry],
        transaction: Transaction,
    ) -> Result<Transaction, AddGitMappingErrorKind> {
        STATS::adds.add_value(entries.len().try_into().map_err(Error::from)?);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        let rows: Vec<_> = entries
            .iter()
            .map(|BonsaiGitMappingEntry { git_sha1, bcs_id }| (&self.repo_id, git_sha1, bcs_id))
            .collect();

        let (transaction, res) =
            InsertMapping::query_with_transaction(transaction, &rows[..]).await?;

        let transaction = if res.affected_rows() != rows.len() as u64 {
            // Let's see if there are any conflicting entries in DB.
            let git_shas = entries.iter().map(|x| x.git_sha1).collect::<Vec<_>>();
            let (transaction, git2bonsai_mapping_from_db) =
                SelectMappingByGitSha1::query_with_transaction(
                    transaction,
                    &self.repo_id,
                    &git_shas[..],
                )
                .await?;
            let git2bonsai_mapping_from_db: BTreeMap<_, _> =
                git2bonsai_mapping_from_db.into_iter().collect();

            let bcs_ids = entries.iter().map(|x| x.bcs_id).collect::<Vec<_>>();
            let (transaction, bonsai2git_mapping_from_db) =
                SelectMappingByBonsai::query_with_transaction(
                    transaction,
                    &self.repo_id,
                    &bcs_ids[..],
                )
                .await?;
            let bonsai2git_mapping_from_db: BTreeMap<_, _> = bonsai2git_mapping_from_db
                .into_iter()
                .map(|(a, b)| (b, a))
                .collect();

            for entry in entries.iter() {
                match (
                    git2bonsai_mapping_from_db.get(&entry.git_sha1),
                    bonsai2git_mapping_from_db.get(&entry.bcs_id),
                ) {
                    (Some(bcs_id), _) if bcs_id == &entry.bcs_id => {} // We've tried to insert a duplicate, proceed.
                    (Some(bcs_id), None) => {
                        // Conflict git_sha1 already mapped to a different bcs_id.
                        return Err(AddGitMappingErrorKind::Conflict(
                            Some(BonsaiGitMappingEntry {
                                git_sha1: entry.git_sha1,
                                bcs_id: *bcs_id,
                            }),
                            vec![entry.clone()],
                        ));
                    }
                    (None, Some(git_sha1)) => {
                        // Conflict bcs_id already mapped to a different git_sha1.
                        return Err(AddGitMappingErrorKind::Conflict(
                            Some(BonsaiGitMappingEntry {
                                git_sha1: *git_sha1,
                                bcs_id: entry.bcs_id,
                            }),
                            vec![entry.clone()],
                        ));
                    }
                    _ => {
                        return Err(AddGitMappingErrorKind::Conflict(None, vec![entry.clone()]));
                    }
                }
            }

            transaction
        } else {
            transaction
        };

        Ok(transaction)
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        objects: BonsaisOrGitShas,
    ) -> Result<Vec<BonsaiGitMappingEntry>> {
        STATS::gets.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let mut mappings =
            select_mapping(&self.connections.read_connection, &self.repo_id, &objects).await?;
        let left_to_fetch = filter_fetched_ids(objects, &mappings[..]);

        if !left_to_fetch.is_empty() {
            STATS::gets_master.add_value(1);
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let mut master_mappings = select_mapping(
                &self.connections.read_master_connection,
                &self.repo_id,
                &left_to_fetch,
            )
            .await?;
            mappings.append(&mut master_mappings);
        }
        Ok(mappings)
    }

    /// Return [`GitSha1`] entries in the inclusive range described by `low` and `high`.
    /// Maximum `limit` entries will be returned.
    async fn get_in_range(
        &self,
        ctx: &CoreContext,
        low: GitSha1,
        high: GitSha1,
        limit: usize,
    ) -> Result<Vec<GitSha1>, Error> {
        if low > high {
            return Ok(Vec::new());
        }
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let rows = SelectGitSha1sByRange::query(
            &self.connections.read_connection,
            &self.repo_id,
            &low,
            &high,
            &limit,
        )
        .await?;
        let mut fetched: Vec<GitSha1> = rows.into_iter().map(|row| row.0).collect();
        if fetched.is_empty() {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let rows = SelectGitSha1sByRange::query(
                &self.connections.read_master_connection,
                &self.repo_id,
                &low,
                &high,
                &limit,
            )
            .await?;
            fetched = rows.into_iter().map(|row| row.0).collect();
        }
        Ok(fetched)
    }
}

pub struct SqlBonsaiGitMappingBuilder {
    connections: SqlConnections,
}

impl SqlBonsaiGitMappingBuilder {
    pub fn build(self, repo_id: RepositoryId) -> SqlBonsaiGitMapping {
        let SqlBonsaiGitMappingBuilder { connections } = self;
        SqlBonsaiGitMapping {
            connections,
            repo_id,
        }
    }
}

impl SqlConstruct for SqlBonsaiGitMappingBuilder {
    const LABEL: &'static str = "bonsai_git_mapping";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-bonsai-git-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlBonsaiGitMappingBuilder {}

fn filter_fetched_ids(
    cs: BonsaisOrGitShas,
    mappings: &[BonsaiGitMappingEntry],
) -> BonsaisOrGitShas {
    match cs {
        BonsaisOrGitShas::Bonsai(cs_ids) => {
            let bcs_fetched: HashSet<_> = mappings.iter().map(|m| &m.bcs_id).collect();

            BonsaisOrGitShas::Bonsai(
                cs_ids
                    .iter()
                    .filter_map(|cs| {
                        if !bcs_fetched.contains(cs) {
                            Some(*cs)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
        BonsaisOrGitShas::GitSha1(git_sha1s) => {
            let git_fetched: HashSet<_> = mappings.iter().map(|m| &m.git_sha1).collect();

            BonsaisOrGitShas::GitSha1(
                git_sha1s
                    .iter()
                    .filter_map(|cs| {
                        if !git_fetched.contains(cs) {
                            Some(*cs)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
    }
}

async fn select_mapping(
    connection: &Connection,
    repo_id: &RepositoryId,
    objects: &BonsaisOrGitShas,
) -> Result<Vec<BonsaiGitMappingEntry>> {
    if objects.is_empty() {
        return Ok(vec![]);
    }

    let rows = match objects {
        BonsaisOrGitShas::Bonsai(bcs_ids) => {
            SelectMappingByBonsai::query(connection, repo_id, &bcs_ids[..]).await?
        }
        BonsaisOrGitShas::GitSha1(git_sha1s) => {
            SelectMappingByGitSha1::query(connection, repo_id, &git_sha1s[..]).await?
        }
    };

    Ok(rows
        .into_iter()
        .map(move |(git_sha1, bcs_id)| BonsaiGitMappingEntry { bcs_id, git_sha1 })
        .collect())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use assert_matches::assert_matches;
use bonsai_git_mapping::AddGitMappingErrorKind;
use bonsai_git_mapping::BonsaiGitMapping;
use bonsai_git_mapping::BonsaiGitMappingEntry;
use bonsai_git_mapping::BonsaisOrGitShas;
use bonsai_git_mapping::SqlBonsaiGitMappingBuilder;
use context::CoreContext;
use fbinit::FacebookInit;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::hash::*;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_ext::open_sqlite_in_memory;
use sql_ext::SqlConnections;

#[fbinit::test]
async fn test_add_and_get(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGitMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry = BonsaiGitMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        git_sha1: ONES_GIT_SHA1,
    };

    mapping.bulk_add(&ctx, &[entry.clone()]).await?;

    let result = mapping
        .get(&ctx, BonsaisOrGitShas::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;
    assert_eq!(result, vec![entry.clone()]);
    let result = mapping
        .get_git_sha1_from_bonsai(&ctx, bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, Some(ONES_GIT_SHA1));
    let result = mapping
        .get_bonsai_from_git_sha1(&ctx, ONES_GIT_SHA1)
        .await?;
    assert_eq!(result, Some(bonsai::ONES_CSID));

    Ok(())
}

#[fbinit::test]
async fn test_add_duplicate(fb: FacebookInit) -> Result<(), Error> {
    // Inserting duplicate entries should just be a successful no-op.
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGitMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry = BonsaiGitMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        git_sha1: ONES_GIT_SHA1,
    };

    mapping.bulk_add(&ctx, &[entry.clone()]).await?;
    mapping.bulk_add(&ctx, &[entry.clone()]).await?;

    let result = mapping
        .get_git_sha1_from_bonsai(&ctx, bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, Some(ONES_GIT_SHA1));

    Ok(())
}

#[fbinit::test]
async fn test_add_conflict(fb: FacebookInit) -> Result<(), Error> {
    // Adding conflicting entries should fail. Other entries inserted in the
    // same bulk_add should be inserted.
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGitMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry = BonsaiGitMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        git_sha1: ONES_GIT_SHA1,
    };

    mapping.bulk_add(&ctx, &[entry.clone()]).await?;

    let entries = vec![
        BonsaiGitMappingEntry {
            // This entry could be inserted normally, but it won't because we have a conflicting
            // entry
            bcs_id: bonsai::TWOS_CSID,
            git_sha1: TWOS_GIT_SHA1,
        },
        BonsaiGitMappingEntry {
            // Conflicting entry.
            bcs_id: bonsai::ONES_CSID,
            git_sha1: THREES_GIT_SHA1,
        },
    ];

    let res = mapping.bulk_add(&ctx, &entries).await;
    assert_matches!(
        res,
        Err(AddGitMappingErrorKind::Conflict(
            Some(BonsaiGitMappingEntry {
                bcs_id: bonsai::ONES_CSID,
                git_sha1: ONES_GIT_SHA1,
            }),
            _
        ))
    );

    let result = mapping
        .get_git_sha1_from_bonsai(&ctx, bonsai::TWOS_CSID)
        .await?;
    assert_eq!(result, None);

    let entries = vec![BonsaiGitMappingEntry {
        // Now this entry will be inserted normally
        bcs_id: bonsai::TWOS_CSID,
        git_sha1: TWOS_GIT_SHA1,
    }];

    mapping.bulk_add(&ctx, &entries).await?;
    let result = mapping
        .get_git_sha1_from_bonsai(&ctx, bonsai::TWOS_CSID)
        .await?;
    assert_eq!(result, Some(TWOS_GIT_SHA1));

    Ok(())
}

#[fbinit::test]
async fn test_bulk_add(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGitMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry1 = BonsaiGitMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        git_sha1: ONES_GIT_SHA1,
    };
    let entry2 = BonsaiGitMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        git_sha1: TWOS_GIT_SHA1,
    };
    let entry3 = BonsaiGitMappingEntry {
        bcs_id: bonsai::THREES_CSID,
        git_sha1: THREES_GIT_SHA1,
    };

    mapping
        .bulk_add(&ctx, &[entry1.clone(), entry2.clone(), entry3.clone()])
        .await?;

    Ok(())
}

#[fbinit::test]
async fn test_missing(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGitMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let result = mapping
        .get(&ctx, BonsaisOrGitShas::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;

    assert_eq!(result, vec![]);

    Ok(())
}

#[fbinit::test]
async fn test_add_with_transaction(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let conn = open_sqlite_in_memory()?;
    conn.execute_batch(SqlBonsaiGitMappingBuilder::CREATION_QUERY)?;
    let conn = Connection::with_sqlite(conn);

    let mapping =
        SqlBonsaiGitMappingBuilder::from_sql_connections(SqlConnections::new_single(conn.clone()))
            .build(REPO_ZERO);

    let entry1 = BonsaiGitMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        git_sha1: ONES_GIT_SHA1,
    };
    let entry2 = BonsaiGitMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        git_sha1: TWOS_GIT_SHA1,
    };

    let txn = conn.start_transaction().await?;
    mapping
        .bulk_add_git_mapping_in_transaction(&ctx, &[entry1.clone()], txn)
        .await?
        .commit()
        .await?;

    assert_eq!(
        Some(ONES_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?
    );

    let txn = conn.start_transaction().await?;
    mapping
        .bulk_add_git_mapping_in_transaction(&ctx, &[entry2.clone()], txn)
        .await?
        .commit()
        .await?;

    assert_eq!(
        Some(TWOS_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?
    );

    // Inserting duplicates fails
    let txn = conn.start_transaction().await?;
    let res = async {
        let entry_conflict = BonsaiGitMappingEntry {
            bcs_id: bonsai::TWOS_CSID,
            git_sha1: THREES_GIT_SHA1,
        };
        mapping
            .bulk_add_git_mapping_in_transaction(&ctx, &[entry_conflict], txn)
            .await?
            .commit()
            .await?;
        Result::<_, AddGitMappingErrorKind>::Ok(())
    }
    .await;
    assert_matches!(
        res,
        Err(AddGitMappingErrorKind::Conflict(
            Some(BonsaiGitMappingEntry {
                bcs_id: bonsai::TWOS_CSID,
                git_sha1: TWOS_GIT_SHA1,
            }),
            _
        ))
    );

    assert_eq!(
        Some(TWOS_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?
    );

    Ok(())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use bonsai_globalrev_mapping::add_globalrevs;
use bonsai_globalrev_mapping::AddGlobalrevsErrorKind;
use bonsai_globalrev_mapping::BonsaiGlobalrevMapping;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingEntry;
use bookmarks::BookmarkTransactionError;
use context::CoreContext;
use mononoke_types::globalrev::Globalrev;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use pushrebase_hook::PushrebaseCommitHook;
use pushrebase_hook::PushrebaseHook;
use pushrebase_hook::PushrebaseTransactionHook;
use pushrebase_hook::RebasedChangesets;
use sql::Transaction;

#[cfg(test)]
mod test;

#[derive(Clone)]
pub struct GlobalrevPushrebaseHook {
    ctx: CoreContext,
    mapping: Arc<dyn BonsaiGlobalrevMapping>,
    repository_id: RepositoryId,
    /// If this is a large repo where globalrevs will be backsynced to a small repo
    small_repo_id: Option<RepositoryId>,
}

impl GlobalrevPushrebaseHook {
    pub fn new(
        ctx: CoreContext,
        mapping: Arc<dyn BonsaiGlobalrevMapping>,
        repository_id: RepositoryId,
        small_repo_id: Option<RepositoryId>,
    ) -> Box<dyn PushrebaseHook> {
        Box::new(Self {
            ctx,
            mapping,
            repository_id,
            small_repo_id,
        })
    }
}

#[async_trait]
impl PushrebaseHook for GlobalrevPushrebaseHook {
    async fn in_critical_section(&self) -> Result<Box<dyn PushrebaseCommitHook>, Error> {
        let max = self.mapping.get_max(&self.ctx).await?;

        let next_rev = match (max, self.small_repo_id) {
            (Some(max), _) => Globalrev::new(max.id() + 1),
            // The source-of-truth change just happened, let's get this value from
            // the small repo.
            (None, Some(small_repo_id)) => self
                .mapping
                .get_max_custom_repo(&self.ctx, &small_repo_id)
                .await?
                .context("Small repo didn't have globalrevs")?,
            (None, None) => Globalrev::start_commit(),
        };

        let hook = Box::new(GlobalrevCommitHook {
            repository_id: self.repository_id,
            assignments: HashMap::new(),
            next_rev,
        }) as Box<dyn PushrebaseCommitHook>;

        Ok(hook)
    }
}

struct GlobalrevCommitHook {
    repository_id: RepositoryId,
    assignments: HashMap<ChangesetId, Globalrev>,
    next_rev: Globalrev,
}

#[async_trait]
impl PushrebaseCommitHook for GlobalrevCommitHook {
    fn post_rebase_changeset(
        &mut self,
        bcs_old: ChangesetId,
        bcs_new: &mut BonsaiChangesetMut,
    ) -> Result<(), Error> {
        self.next_rev.set_on_changeset(bcs_new);

        self.assignments.insert(bcs_old, self.next_rev);

        self.next_rev = self.next_rev.increment();

        Ok(())
    }

    async fn into_transaction_hook(
        self: Box<Self>,
        _ctx: &CoreContext,
        rebased: &RebasedChangesets,
    ) -> Result<Box<dyn PushrebaseTransactionHook>, Error> {
        // Let's tie assigned globalrevs to rebased Bonsai changesets:
        let entries = self
            .assignments
            .iter()
            .map(|(cs_id, globalrev)| {
                let replacement_bcs_id = rebased
                    .get(cs_id)
                    .ok_or_else(|| {
                        let e = format!(
                            "Commit was assigned a Globalrev, but is not found in rebased set: {}",
                            cs_id
                        );
                        Error::msg(e)
                    })?
                    .0;

                Ok(BonsaiGlobalrevMappingEntry::new(
                    replacement_bcs_id,
                    *globalrev,
                ))
            })
            .collect::<Result<Vec<_>, Error>>()?;

        // NOTE: This check shouldn't be necessary as long as pushrebase hooks are bug-free, but
        // since they're a new addition, let's be conservative.
        if rebased.len() != self.assignments.len() {
            return Err(anyhow!(
                "Globalrev rebased set ({}) and assignments ({}) have different lengths!",
                rebased.len(),
                self.assignments.len(),
            ));
        }

        Ok(Box::new(GlobalrevTransactionHook {
            repo_id: self.repository_id,
            entries,
        }) as Box<dyn PushrebaseTransactionHook>)
    }
}

struct GlobalrevTransactionHook {
    repo_id: RepositoryId,
    entries: Vec<BonsaiGlobalrevMappingEntry>,
}

#[async_trait]
impl PushrebaseTransactionHook for GlobalrevTransactionHook {
    async fn populate_transaction(
        &self,
        _ctx: &CoreContext,
        txn: Transaction,
    ) -> Result<Transaction, BookmarkTransactionError> {
        let txn = add_globalrevs(txn, self.repo_id, &self.entries[..])
            .await
            .map_err(|e| match e {
                AddGlobalrevsErrorKind::Conflict => BookmarkTransactionError::LogicError,
                e @ AddGlobalrevsErrorKind::InternalError(..) => {
                    BookmarkTransactionError::Other(e.into())
                }
            })?;

        Ok(txn)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::time::Duration;

use anyhow::Error;
use async_trait::async_trait;
use blobrepo::BlobRepo;
use blobstore::Loadable;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingArc;
use bookmarks::BookmarkTransactionError;
use borrowed::borrowed;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::future::try_join_all;
use maplit::hashset;
use mononoke_types::globalrev::Globalrev;
use mononoke_types::globalrev::START_COMMIT_GLOBALREV;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use pushrebase::do_pushrebase_bonsai;
use pushrebase_hook::PushrebaseCommitHook;
use pushrebase_hook::PushrebaseHook;
use pushrebase_hook::PushrebaseTransactionHook;
use pushrebase_hook::RebasedChangesets;
use rand::Rng;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use sql::Transaction;
use test_repo_factory::TestRepoFactory;
use tests_utils::bookmark;
use tests_utils::resolve_cs_id;
use tests_utils::CreateCommitContext;

use crate::GlobalrevPushrebaseHook;

#[fbinit::test]
fn pushrebase_assigns_globalrevs(fb: FacebookInit) -> Result<(), Error> {
    let runtime = tokio::runtime::Runtime::new()?;
    runtime.block_on(pushrebase_assigns_globalrevs_impl(fb))
}

async fn pushrebase_assigns_globalrevs_impl(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo: BlobRepo = TestRepoFactory::new(fb)?
        .with_id(RepositoryId::new(1))
        .build()
        .await?;
    borrowed!(ctx, repo);

    let root = CreateCommitContext::new_root(ctx, repo).commit().await?;

    let cs1 = CreateCommitContext::new(ctx, repo, vec![root])
        .commit()
        .await?;

    let cs2 = CreateCommitContext::new(ctx, repo, vec![root])
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let book = bookmark(ctx, repo, "master").set_to(cs1).await?;

    let hooks = [GlobalrevPushrebaseHook::new(
        ctx.clone(),
        repo.bonsai_globalrev_mapping_arc(),
        repo.repo_identity().id(),
        None,
    )];

    let rebased = do_pushrebase_bonsai(
        ctx,
        repo,
        &Default::default(),
        &book,
        &hashset![cs2.clone()],
        &hooks,
    )
    .await?
    .rebased_changesets;

    let cs2_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs2.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs2"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs3 = CreateCommitContext::new(ctx, repo, vec![root])
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs4 = CreateCommitContext::new(ctx, repo, vec![cs3.get_changeset_id()])
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let rebased = do_pushrebase_bonsai(
        ctx,
        repo,
        &Default::default(),
        &book,
        &hashset![cs3.clone(), cs4.clone()],
        &hooks,
    )
    .await?
    .rebased_changesets;

    let cs3_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs3.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs3"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs4_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs4.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs4"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    assert_eq!(
        Globalrev::new(START_COMMIT_GLOBALREV),
        Globalrev::from_bcs(&cs2_rebased)?
    );
    assert_eq!(
        Globalrev::new(START_COMMIT_GLOBALREV + 1),
        Globalrev::from_bcs(&cs3_rebased)?
    );
    assert_eq!(
        Globalrev::new(START_COMMIT_GLOBALREV + 2),
        Globalrev::from_bcs(&cs4_rebased)?
    );

    Ok(())
}

#[fbinit::test]
fn test_pushrebase_race_assigns_monotonic_globalrevs(fb: FacebookInit) -> Result<(), Error> {
    let runtime = tokio::runtime::Runtime::new()?;
    runtime.block_on(pushrebase_race_assigns_monotonic_globalrevs(fb))
}

async fn pushrebase_race_assigns_monotonic_globalrevs(fb: FacebookInit) -> Result<(), Error> {
    #[derive(Copy, Clone)]
    struct SleepHook;

    #[async_trait]
    impl PushrebaseHook for SleepHook {
        async fn in_critical_section(&self) -> Result<Box<dyn PushrebaseCommitHook>, Error> {
            let us = rand::thread_rng().gen_range(0..100);
            tokio::time::sleep(Duration::from_micros(us)).await;
            Ok(Box::new(*self) as Box<dyn PushrebaseCommitHook>)
        }
    }

    #[async_trait]
    impl PushrebaseCommitHook for SleepHook {
        fn post_rebase_changeset(
            &mut self,
            _bcs_old: ChangesetId,
            _bcs_new: &mut BonsaiChangesetMut,
        ) -> Result<(), Error> {
            Ok(())
        }

        async fn into_transaction_hook(
            self: Box<Self>,
            _ctx: &CoreContext,
            _changesets: &RebasedChangesets,
        ) -> Result<Box<dyn PushrebaseTransactionHook>, Error> {
            Ok(Box::new(*self) as Box<dyn PushrebaseTransactionHook>)
        }
    }

    #[async_trait]
    impl PushrebaseTransactionHook for SleepHook {
        async fn populate_transaction(
            &self,
            _ctx: &CoreContext,
            txn: Transaction,
        ) -> Result<Transaction, BookmarkTransactionError> {
            Ok(txn)
        }
    }

    let ctx = CoreContext::test_mock(fb);
    let repo: BlobRepo = TestRepoFactory::new(fb)?
        .with_id(RepositoryId::new(1))
        .build()
        .await?;
    borrowed!(ctx, repo);

    let root = CreateCommitContext::new_root(ctx, repo).commit().await?;
    let book = bookmark(ctx, repo, "master").set_to(root).await?;

    let hooks = [
        GlobalrevPushrebaseHook::new(
            ctx.clone(),
            repo.bonsai_globalrev_mapping_arc(),
            repo.repo_identity().id(),
            None,
        ),
        Box::new(SleepHook) as Box<dyn PushrebaseHook>,
    ];

    let n: u64 = 10;
    let mut futs = vec![];
    for _ in 0..n {
        let bcs = CreateCommitContext::new(ctx, repo, vec![root])
            .commit()
            .await?
            .load(ctx, repo.repo_blobstore())
            .await?;

        let fut = async {
            let params = Default::default();
            let bcss = hashset![bcs];
            do_pushrebase_bonsai(ctx, repo, &params, &book, &bcss, &hooks).await
        };

        futs.push(fut);
    }

    let results = try_join_all(futs).await?;

    // Check that we had some retries
    assert!(results.iter().fold(0, |sum, e| sum + e.retry_num.0) > 0);

    let mut next_cs_id = resolve_cs_id(ctx, repo, "master").await?;
    let mut next_globalrev = START_COMMIT_GLOBALREV + n - 1;

    // Check history backwards and assert that globalrevs are going down

    loop {
        if next_cs_id == root {
            break;
        }

        let cs = next_cs_id.load(ctx, repo.repo_blobstore()).await?;

        assert_eq!(Globalrev::new(next_globalrev), Globalrev::from_bcs(&cs)?);

        next_cs_id = cs
            .parents()
            .next()
            .ok_or_else(|| Error::msg("no parent found"))?;
        next_globalrev -= 1;
    }

    Ok(())
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "bonsai_globalrev_mapping crate
mononoke_types_thrift mononoke_types_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=bonsai_globalrev_mapping_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "bonsai_globalrev_mapping.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use bonsai_globalrev_mapping_thrift as thrift;
use bytes::Bytes;
use caching_ext::get_or_fill;
use caching_ext::CacheDisposition;
use caching_ext::CacheHandlerFactory;
use caching_ext::CacheTtl;
use caching_ext::CachelibHandler;
use caching_ext::EntityStore;
use caching_ext::KeyedEntityStore;
use caching_ext::McErrorKind;
use caching_ext::McResult;
use caching_ext::MemcacheEntity;
use caching_ext::MemcacheHandler;
use context::CoreContext;
use fbthrift::compact_protocol;
use memcache::KeyGen;
use mononoke_types::ChangesetId;
use mononoke_types::Globalrev;
use mononoke_types::RepositoryId;

use super::BonsaiGlobalrevMapping;
use super::BonsaiGlobalrevMappingEntry;
use super::BonsaisOrGlobalrevs;

#[derive(Abomonation, Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiGlobalrevMappingCacheEntry {
    pub repo_id: RepositoryId,
    pub bcs_id: ChangesetId,
    pub globalrev: Globalrev,
}

impl BonsaiGlobalrevMappingCacheEntry {
    fn into_entry(self, repo_id: RepositoryId) -> Result<BonsaiGlobalrevMappingEntry> {
        if self.repo_id == repo_id {
            Ok(BonsaiGlobalrevMappingEntry {
                bcs_id: self.bcs_id,
                globalrev: self.globalrev,
            })
        } else {
            Err(anyhow!(
                "Cache returned invalid entry: repo {} returned for query to repo {}",
                self.repo_id,
                repo_id
            ))
        }
    }

    fn from_entry(
        entry: BonsaiGlobalrevMappingEntry,
        repo_id: RepositoryId,
    ) -> BonsaiGlobalrevMappingCacheEntry {
        BonsaiGlobalrevMappingCacheEntry {
            repo_id,
            bcs_id: entry.bcs_id,
            globalrev: entry.globalrev,
        }
    }
}

pub struct CachingBonsaiGlobalrevMapping {
    cachelib: CachelibHandler<BonsaiGlobalrevMappingCacheEntry>,
    memcache: MemcacheHandler,
    keygen: KeyGen,
    inner: Arc<dyn BonsaiGlobalrevMapping>,
}

impl CachingBonsaiGlobalrevMapping {
    pub fn new(
        inner: Arc<dyn BonsaiGlobalrevMapping>,
        cache_handler_factory: CacheHandlerFactory,
    ) -> Self {
        Self {
            inner,
            cachelib: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
            keygen: Self::create_key_gen(),
        }
    }

    pub fn new_test(inner: Arc<dyn BonsaiGlobalrevMapping>) -> Self {
        Self::new(inner, CacheHandlerFactory::Mocked)
    }

    fn create_key_gen() -> KeyGen {
        let key_prefix = "scm.mononoke.bonsai_globalrev_mapping";

        KeyGen::new(
            key_prefix,
            thrift::MC_CODEVER as u32,
            thrift::MC_SITEVER as u32,
        )
    }

    pub fn cachelib(&self) -> &CachelibHandler<BonsaiGlobalrevMappingCacheEntry> {
        &self.cachelib
    }
}

#[async_trait]
impl BonsaiGlobalrevMapping for CachingBonsaiGlobalrevMapping {
    fn repo_id(&self) -> RepositoryId {
        self.inner.as_ref().repo_id()
    }

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGlobalrevMappingEntry],
    ) -> Result<(), Error> {
        self.inner.as_ref().bulk_import(ctx, entries).await
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        objects: BonsaisOrGlobalrevs,
    ) -> Result<Vec<BonsaiGlobalrevMappingEntry>, Error> {
        let cache_request = (ctx, self);
        let repo_id = self.repo_id();

        let res = match objects {
            BonsaisOrGlobalrevs::Bonsai(cs_ids) => {
                get_or_fill(&cache_request, cs_ids.into_iter().collect())
                    .await
                    .with_context(|| "Error fetching globalrevs via cache")?
                    .into_values()
                    .map(|val| val.into_entry(repo_id))
                    .collect::<Result<_>>()?
            }
            BonsaisOrGlobalrevs::Globalrev(globalrevs) => {
                get_or_fill(&cache_request, globalrevs.into_iter().collect())
                    .await
                    .with_context(|| "Error fetching bonsais via cache")?
                    .into_values()
                    .map(|val| val.into_entry(repo_id))
                    .collect::<Result<_>>()?
            }
        };

        Ok(res)
    }

    async fn get_closest_globalrev(
        &self,
        ctx: &CoreContext,
        globalrev: Globalrev,
    ) -> Result<Option<Globalrev>, Error> {
        self.inner
            .as_ref()
            .get_closest_globalrev(ctx, globalrev)
            .await
    }

    async fn get_max(&self, ctx: &CoreContext) -> Result<Option<Globalrev>, Error> {
        self.inner.as_ref().get_max(ctx).await
    }

    async fn get_max_custom_repo(
        &self,
        ctx: &CoreContext,
        repo_id: &RepositoryId,
    ) -> Result<Option<Globalrev>, Error> {
        self.inner.as_ref().get_max_custom_repo(ctx, repo_id).await
    }
}

impl MemcacheEntity for BonsaiGlobalrevMappingCacheEntry {
    fn serialize(&self) -> Bytes {
        let entry = thrift::BonsaiGlobalrevMappingEntry {
            repo_id: self.repo_id.id(),
            bcs_id: self.bcs_id.into_thrift(),
            globalrev: self
                .globalrev
                .id()
                .try_into()
                .expect("Globalrevs must fit within a i64"),
        };
        compact_protocol::serialize(&entry)
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        let thrift::BonsaiGlobalrevMappingEntry {
            repo_id,
            bcs_id,
            globalrev,
        } = compact_protocol::deserialize(bytes).map_err(|_| McErrorKind::Deserialization)?;

        let repo_id = RepositoryId::new(repo_id);
        let bcs_id = ChangesetId::from_thrift(bcs_id).map_err(|_| McErrorKind::Deserialization)?;
        let globalrev = Globalrev::new(
            globalrev
                .try_into()
                .map_err(|_| McErrorKind::Deserialization)?,
        );

        Ok(BonsaiGlobalrevMappingCacheEntry {
            repo_id,
            bcs_id,
            globalrev,
        })
    }
}

type CacheRequest<'a> = (&'a CoreContext, &'a CachingBonsaiGlobalrevMapping);

impl EntityStore<BonsaiGlobalrevMappingCacheEntry> for CacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<BonsaiGlobalrevMappingCacheEntry> {
        let (_, mapping) = self;
        &mapping.cachelib
    }

    fn keygen(&self) -> &KeyGen {
        let (_, mapping) = self;
        &mapping.keygen
    }

    fn memcache(&self) -> &MemcacheHandler {
        let (_, mapping) = self;
        &mapping.memcache
    }

    fn cache_determinator(&self, _: &BonsaiGlobalrevMappingCacheEntry) -> CacheDisposition {
        CacheDisposition::Cache(CacheTtl::NoTtl)
    }

    caching_ext::impl_singleton_stats!("bonsai_globalrev_mapping");
}

#[async_trait]
impl KeyedEntityStore<ChangesetId, BonsaiGlobalrevMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &ChangesetId) -> String {
        let (_, mapping) = self;
        format!("{}.bonsai.{}", mapping.repo_id(), key)
    }

    async fn get_from_db(
        &self,
        keys: HashSet<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, BonsaiGlobalrevMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .inner
            .as_ref()
            .get(ctx, BonsaisOrGlobalrevs::Bonsai(keys.into_iter().collect()))
            .await
            .with_context(|| "Error fetching globalrevs from bonsais from SQL")?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| {
                    (
                        e.bcs_id,
                        BonsaiGlobalrevMappingCacheEntry::from_entry(e, repo_id),
                    )
                })
                .collect(),
        )
    }
}

#[async_trait]
impl KeyedEntityStore<Globalrev, BonsaiGlobalrevMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &Globalrev) -> String {
        let (_, mapping) = self;
        format!("{}.globalrev.{}", mapping.repo_id(), key.id())
    }

    async fn get_from_db(
        &self,
        keys: HashSet<Globalrev>,
    ) -> Result<HashMap<Globalrev, BonsaiGlobalrevMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .inner
            .as_ref()
            .get(
                ctx,
                BonsaisOrGlobalrevs::Globalrev(keys.into_iter().collect()),
            )
            .await
            .with_context(|| "Error fetching bonsais from globalrevs from SQL")?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| {
                    (
                        e.globalrev,
                        BonsaiGlobalrevMappingCacheEntry::from_entry(e, repo_id),
                    )
                })
                .collect(),
        )
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

mod caching;
mod sql;

use anyhow::Error;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::Globalrev;
use mononoke_types::RepositoryId;

pub use crate::caching::CachingBonsaiGlobalrevMapping;
pub use crate::sql::add_globalrevs;
pub use crate::sql::bulk_import_globalrevs;
pub use crate::sql::AddGlobalrevsErrorKind;
pub use crate::sql::SqlBonsaiGlobalrevMapping;
pub use crate::sql::SqlBonsaiGlobalrevMappingBuilder;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiGlobalrevMappingEntry {
    pub bcs_id: ChangesetId,
    pub globalrev: Globalrev,
}

impl BonsaiGlobalrevMappingEntry {
    pub fn new(bcs_id: ChangesetId, globalrev: Globalrev) -> Self {
        BonsaiGlobalrevMappingEntry { bcs_id, globalrev }
    }
}

pub enum BonsaisOrGlobalrevs {
    Bonsai(Vec<ChangesetId>),
    Globalrev(Vec<Globalrev>),
}

impl BonsaisOrGlobalrevs {
    pub fn is_empty(&self) -> bool {
        match self {
            BonsaisOrGlobalrevs::Bonsai(v) => v.is_empty(),
            BonsaisOrGlobalrevs::Globalrev(v) => v.is_empty(),
        }
    }
}

impl From<ChangesetId> for BonsaisOrGlobalrevs {
    fn from(cs_id: ChangesetId) -> Self {
        BonsaisOrGlobalrevs::Bonsai(vec![cs_id])
    }
}

impl From<Vec<ChangesetId>> for BonsaisOrGlobalrevs {
    fn from(cs_ids: Vec<ChangesetId>) -> Self {
        BonsaisOrGlobalrevs::Bonsai(cs_ids)
    }
}

impl From<Globalrev> for BonsaisOrGlobalrevs {
    fn from(rev: Globalrev) -> Self {
        BonsaisOrGlobalrevs::Globalrev(vec![rev])
    }
}

impl From<Vec<Globalrev>> for BonsaisOrGlobalrevs {
    fn from(revs: Vec<Globalrev>) -> Self {
        BonsaisOrGlobalrevs::Globalrev(revs)
    }
}

#[facet::facet]
#[async_trait]
pub trait BonsaiGlobalrevMapping: Send + Sync {
    fn repo_id(&self) -> RepositoryId;

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGlobalrevMappingEntry],
    ) -> Result<(), Error>;

    async fn get(
        &self,
        ctx: &CoreContext,
        field: BonsaisOrGlobalrevs,
    ) -> Result<Vec<BonsaiGlobalrevMappingEntry>, Error>;

    async fn get_globalrev_from_bonsai(
        &self,
        ctx: &CoreContext,
        bcs_id: ChangesetId,
    ) -> Result<Option<Globalrev>, Error> {
        let result = self
            .get(ctx, BonsaisOrGlobalrevs::Bonsai(vec![bcs_id]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.globalrev))
    }

    async fn get_bonsai_from_globalrev(
        &self,
        ctx: &CoreContext,
        globalrev: Globalrev,
    ) -> Result<Option<ChangesetId>, Error> {
        let result = self
            .get(ctx, BonsaisOrGlobalrevs::Globalrev(vec![globalrev]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.bcs_id))
    }

    async fn get_closest_globalrev(
        &self,
        ctx: &CoreContext,
        globalrev: Globalrev,
    ) -> Result<Option<Globalrev>, Error>;

    /// Read the most recent Globalrev. This produces the freshest data possible, and is meant to
    /// be used for Globalrev assignment.
    async fn get_max(&self, ctx: &CoreContext) -> Result<Option<Globalrev>, Error>;

    async fn get_max_custom_repo(
        &self,
        ctx: &CoreContext,
        repo_id: &RepositoryId,
    ) -> Result<Option<Globalrev>, Error>;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;

use ::sql::Connection;
use ::sql::Transaction;
use ::sql_ext::mononoke_queries;
use anyhow::Error;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::Globalrev;
use mononoke_types::RepositoryId;
use slog::warn;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::SqlConnections;
use thiserror::Error;

use super::BonsaiGlobalrevMapping;
use super::BonsaiGlobalrevMappingEntry;
use super::BonsaisOrGlobalrevs;

mononoke_queries! {
    write DangerouslyAddGlobalrevs(values: (
        repo_id: RepositoryId,
        bcs_id: ChangesetId,
        globalrev: Globalrev,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bonsai_globalrev_mapping (repo_id, bcs_id, globalrev) VALUES {values}"
    }

    write ReplaceGlobalrevs(values: (
        repo_id: RepositoryId,
        bcs_id: ChangesetId,
        globalrev: Globalrev,
    )) {
        none,
        "REPLACE INTO bonsai_globalrev_mapping (repo_id, bcs_id, globalrev) VALUES {values}"
    }

    read SelectMappingByBonsai(
        repo_id: RepositoryId,
        >list bcs_id: ChangesetId
    ) -> (ChangesetId, Globalrev) {
        "SELECT bcs_id, globalrev
         FROM bonsai_globalrev_mapping
         WHERE repo_id = {repo_id} AND bcs_id in {bcs_id}"
    }

    read SelectMappingByGlobalrev(
        repo_id: RepositoryId,
        >list globalrev: Globalrev
    ) -> (ChangesetId, Globalrev) {
        "SELECT bcs_id, globalrev
         FROM bonsai_globalrev_mapping
         WHERE repo_id = {repo_id} AND globalrev in {globalrev}"
    }

    read SelectMaxEntry(repo_id: RepositoryId) -> (Globalrev,) {
        "
        SELECT globalrev
        FROM bonsai_globalrev_mapping
        WHERE repo_id = {repo_id}
        ORDER BY globalrev DESC
        LIMIT 1
        "
    }

    read SelectClosestGlobalrev(repo_id: RepositoryId, rev: Globalrev) -> (Globalrev,) {
        "
        SELECT globalrev
        FROM bonsai_globalrev_mapping
        WHERE repo_id = {repo_id} AND globalrev <= {rev}
        ORDER BY globalrev DESC
        LIMIT 1
        "
    }
}

pub struct SqlBonsaiGlobalrevMapping {
    connections: SqlConnections,
    repo_id: RepositoryId,
}

#[derive(Clone)]
pub struct SqlBonsaiGlobalrevMappingBuilder {
    connections: SqlConnections,
}

impl SqlConstruct for SqlBonsaiGlobalrevMappingBuilder {
    const LABEL: &'static str = "bonsai_globalrev_mapping";

    const CREATION_QUERY: &'static str =
        include_str!("../schemas/sqlite-bonsai-globalrev-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlBonsaiGlobalrevMappingBuilder {}

impl SqlBonsaiGlobalrevMappingBuilder {
    pub fn build(self, repo_id: RepositoryId) -> SqlBonsaiGlobalrevMapping {
        SqlBonsaiGlobalrevMapping {
            connections: self.connections,
            repo_id,
        }
    }
}

#[async_trait]
impl BonsaiGlobalrevMapping for SqlBonsaiGlobalrevMapping {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGlobalrevMappingEntry],
    ) -> Result<(), Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);
        let repo_id = self.repo_id;

        let entries: Vec<_> = entries
            .iter()
            .map(|entry| (&repo_id, &entry.bcs_id, &entry.globalrev))
            .collect();

        DangerouslyAddGlobalrevs::query(&self.connections.write_connection, &entries[..]).await?;

        Ok(())
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        objects: BonsaisOrGlobalrevs,
    ) -> Result<Vec<BonsaiGlobalrevMappingEntry>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let mut mappings =
            select_mapping(&self.connections.read_connection, self.repo_id, &objects).await?;

        let left_to_fetch = filter_fetched_objects(objects, &mappings[..]);

        if left_to_fetch.is_empty() {
            return Ok(mappings);
        }

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);

        let mut master_mappings = select_mapping(
            &self.connections.read_master_connection,
            self.repo_id,
            &left_to_fetch,
        )
        .await?;
        mappings.append(&mut master_mappings);
        Ok(mappings)
    }

    async fn get_closest_globalrev(
        &self,
        ctx: &CoreContext,
        globalrev: Globalrev,
    ) -> Result<Option<Globalrev>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let row = SelectClosestGlobalrev::query(
            &self.connections.read_connection,
            &self.repo_id,
            &globalrev,
        )
        .await?
        .into_iter()
        .next();

        Ok(row.map(|r| r.0))
    }

    async fn get_max(&self, ctx: &CoreContext) -> Result<Option<Globalrev>, Error> {
        self.get_max_custom_repo(ctx, &self.repo_id).await
    }

    async fn get_max_custom_repo(
        &self,
        ctx: &CoreContext,
        repo_id: &RepositoryId,
    ) -> Result<Option<Globalrev>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);

        let row = SelectMaxEntry::query(&self.connections.read_master_connection, repo_id)
            .await?
            .into_iter()
            .next();

        Ok(row.map(|r| r.0))
    }
}

fn filter_fetched_objects(
    objects: BonsaisOrGlobalrevs,
    mappings: &[BonsaiGlobalrevMappingEntry],
) -> BonsaisOrGlobalrevs {
    match objects {
        BonsaisOrGlobalrevs::Bonsai(cs_ids) => {
            let bcs_fetched: HashSet<_> = mappings.iter().map(|m| &m.bcs_id).collect();

            BonsaisOrGlobalrevs::Bonsai(
                cs_ids
                    .iter()
                    .filter_map(|cs| {
                        if !bcs_fetched.contains(cs) {
                            Some(*cs)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
        BonsaisOrGlobalrevs::Globalrev(globalrevs) => {
            let globalrevs_fetched: HashSet<_> = mappings.iter().map(|m| &m.globalrev).collect();

            BonsaisOrGlobalrevs::Globalrev(
                globalrevs
                    .iter()
                    .filter_map(|globalrev| {
                        if !globalrevs_fetched.contains(globalrev) {
                            Some(*globalrev)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
    }
}

async fn select_mapping(
    connection: &Connection,
    repo_id: RepositoryId,
    objects: &BonsaisOrGlobalrevs,
) -> Result<Vec<BonsaiGlobalrevMappingEntry>, Error> {
    if objects.is_empty() {
        return Ok(vec![]);
    }

    let rows = match objects {
        BonsaisOrGlobalrevs::Bonsai(bcs_ids) => {
            SelectMappingByBonsai::query(connection, &repo_id, &bcs_ids[..]).await?
        }
        BonsaisOrGlobalrevs::Globalrev(globalrevs) => {
            SelectMappingByGlobalrev::query(connection, &repo_id, &globalrevs[..]).await?
        }
    };

    Ok(rows
        .into_iter()
        .map(move |(bcs_id, globalrev)| BonsaiGlobalrevMappingEntry { bcs_id, globalrev })
        .collect())
}

/// This method is for importing Globalrevs in bulk from a set of BonsaiChangesets where you know
/// they are correct. Don't use this to assign new Globalrevs.
pub async fn bulk_import_globalrevs<'a>(
    ctx: &'a CoreContext,
    globalrevs_store: &'a dyn BonsaiGlobalrevMapping,
    changesets: impl IntoIterator<Item = &'a BonsaiChangeset>,
) -> Result<(), Error> {
    let mut entries = vec![];
    for bcs in changesets.into_iter() {
        match Globalrev::from_bcs(bcs) {
            Ok(globalrev) => {
                let entry = BonsaiGlobalrevMappingEntry::new(bcs.get_changeset_id(), globalrev);
                entries.push(entry);
            }
            Err(e) => {
                warn!(
                    ctx.logger(),
                    "Couldn't fetch globalrev from commit: {:?}", e
                );
            }
        }
    }

    globalrevs_store.bulk_import(ctx, &entries).await?;

    Ok(())
}

#[derive(Debug, Error)]
pub enum AddGlobalrevsErrorKind {
    #[error("Conflict detected while inserting Globalrevs")]
    Conflict,

    #[error("Internal error occurred while inserting Globalrevs")]
    InternalError(#[from] Error),
}

// NOTE: For now, this is a top-level function since it doesn't use the connections in the
// SqlBonsaiGlobalrevMapping, but if we were to add more implementations of the
// BonsaiGlobalrevMapping trait, we should probably rethink the design of it, and not actually have
// it contain any connections (instead, they should be passed on by callers).
pub async fn add_globalrevs(
    transaction: Transaction,
    repo_id: RepositoryId,
    entries: impl IntoIterator<Item = &BonsaiGlobalrevMappingEntry>,
) -> Result<Transaction, AddGlobalrevsErrorKind> {
    let rows: Vec<_> = entries
        .into_iter()
        .map(|BonsaiGlobalrevMappingEntry { bcs_id, globalrev }| (&repo_id, bcs_id, globalrev))
        .collect();

    // It'd be really nice if we could rely on the error from an index conflict here, but our SQL
    // crate doesn't allow us to reach into this yet, so for now we check the number of affected
    // rows.

    let (transaction, res) =
        DangerouslyAddGlobalrevs::query_with_transaction(transaction, &rows[..]).await?;

    if res.affected_rows() != rows.len() as u64 {
        return Err(AddGlobalrevsErrorKind::Conflict);
    }

    Ok(transaction)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::Error;
use assert_matches::assert_matches;
use bonsai_globalrev_mapping::add_globalrevs;
use bonsai_globalrev_mapping::AddGlobalrevsErrorKind;
use bonsai_globalrev_mapping::BonsaiGlobalrevMapping;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingEntry;
use bonsai_globalrev_mapping::BonsaisOrGlobalrevs;
use bonsai_globalrev_mapping::CachingBonsaiGlobalrevMapping;
use bonsai_globalrev_mapping::SqlBonsaiGlobalrevMappingBuilder;
use context::CoreContext;
use fbinit::FacebookInit;
use mercurial_types_mocks::globalrev::*;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_ext::open_sqlite_in_memory;
use sql_ext::SqlConnections;

#[fbinit::test]
async fn test_add_and_get(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ZERO,
    };

    mapping.bulk_import(&ctx, &[entry.clone()]).await?;

    let result = mapping
        .get(&ctx, BonsaisOrGlobalrevs::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;
    assert_eq!(result, vec![entry.clone()]);
    let result = mapping
        .get_globalrev_from_bonsai(&ctx, bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, Some(GLOBALREV_ZERO));
    let result = mapping
        .get_bonsai_from_globalrev(&ctx, GLOBALREV_ZERO)
        .await?;
    assert_eq!(result, Some(bonsai::ONES_CSID));

    Ok(())
}

#[fbinit::test]
async fn test_bulk_import(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry1 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ZERO,
    };
    let entry2 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        globalrev: GLOBALREV_ONE,
    };
    let entry3 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::THREES_CSID,
        globalrev: GLOBALREV_TWO,
    };

    mapping
        .bulk_import(&ctx, &[entry1.clone(), entry2.clone(), entry3.clone()])
        .await?;

    Ok(())
}

#[fbinit::test]
async fn test_missing(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let result = mapping
        .get(&ctx, BonsaisOrGlobalrevs::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;

    assert_eq!(result, vec![]);

    Ok(())
}

#[fbinit::test]
async fn test_get_max(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    assert_eq!(None, mapping.get_max(&ctx).await?);

    let e0 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ZERO,
    };
    mapping.bulk_import(&ctx, &[e0.clone()]).await?;
    assert_eq!(Some(GLOBALREV_ZERO), mapping.get_max(&ctx).await?);

    let e1 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        globalrev: GLOBALREV_ONE,
    };
    mapping.bulk_import(&ctx, &[e1.clone()]).await?;
    assert_eq!(Some(GLOBALREV_ONE), mapping.get_max(&ctx).await?);

    Ok(())
}

#[fbinit::test]
async fn test_add_globalrevs(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let conn = open_sqlite_in_memory()?;
    conn.execute_batch(SqlBonsaiGlobalrevMappingBuilder::CREATION_QUERY)?;
    let conn = Connection::with_sqlite(conn);

    let mapping = SqlBonsaiGlobalrevMappingBuilder::from_sql_connections(
        SqlConnections::new_single(conn.clone()),
    )
    .build(REPO_ZERO);

    let e0 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ZERO,
    };

    let e1 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        globalrev: GLOBALREV_ONE,
    };

    let txn = conn.start_transaction().await?;
    let txn = add_globalrevs(txn, REPO_ZERO, &[e0.clone()]).await?;
    txn.commit().await?;

    assert_eq!(
        Some(GLOBALREV_ZERO),
        mapping
            .get_globalrev_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?
    );

    let txn = conn.start_transaction().await?;
    let txn = add_globalrevs(txn, REPO_ZERO, &[e1.clone()]).await?;
    txn.commit().await?;

    assert_eq!(
        Some(GLOBALREV_ONE),
        mapping
            .get_globalrev_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?
    );

    // Inserting duplicates fails

    let txn = conn.start_transaction().await?;
    let res = async move {
        let txn = add_globalrevs(txn, REPO_ZERO, &[e1.clone()]).await?;
        txn.commit().await?;
        Result::<_, AddGlobalrevsErrorKind>::Ok(())
    }
    .await;
    assert_matches!(res, Err(AddGlobalrevsErrorKind::Conflict));

    assert_eq!(
        Some(GLOBALREV_ONE),
        mapping
            .get_globalrev_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?
    );

    Ok(())
}

#[fbinit::test]
async fn test_closest_globalrev(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let e0 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ONE,
    };

    let e1 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        globalrev: GLOBALREV_TWO,
    };

    mapping.bulk_import(&ctx, &[e0, e1]).await?;

    assert_eq!(
        mapping.get_closest_globalrev(&ctx, GLOBALREV_ZERO).await?,
        None
    );

    assert_eq!(
        mapping.get_closest_globalrev(&ctx, GLOBALREV_ONE).await?,
        Some(GLOBALREV_ONE)
    );

    assert_eq!(
        mapping.get_closest_globalrev(&ctx, GLOBALREV_TWO).await?,
        Some(GLOBALREV_TWO)
    );

    assert_eq!(
        mapping.get_closest_globalrev(&ctx, GLOBALREV_THREE).await?,
        Some(GLOBALREV_TWO)
    );

    assert_eq!(
        mapping.get_closest_globalrev(&ctx, GLOBALREV_ZERO).await?,
        None
    );

    Ok(())
}

#[fbinit::test]
async fn test_caching(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping =
        Arc::new(SqlBonsaiGlobalrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO));
    let caching = CachingBonsaiGlobalrevMapping::new_test(mapping.clone());
    let store = caching
        .cachelib()
        .mock_store()
        .expect("new_test gives us a MockStore");

    let e0 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        globalrev: GLOBALREV_ONE,
    };

    let e1 = BonsaiGlobalrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        globalrev: GLOBALREV_TWO,
    };

    mapping.bulk_import(&ctx, &[e0, e1]).await?;

    assert_eq!(
        caching
            .get_globalrev_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?,
        Some(GLOBALREV_ONE)
    );

    assert_eq!(store.stats().gets, 1);
    assert_eq!(store.stats().hits, 0);
    assert_eq!(store.stats().sets, 1);

    assert_eq!(
        caching
            .get_globalrev_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?,
        Some(GLOBALREV_ONE)
    );

    assert_eq!(store.stats().gets, 2);
    assert_eq!(store.stats().hits, 1);
    assert_eq!(store.stats().sets, 1);

    assert_eq!(
        caching
            .get_globalrev_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?,
        Some(GLOBALREV_TWO)
    );
    assert_eq!(store.stats().gets, 3);
    assert_eq!(store.stats().hits, 1);
    assert_eq!(store.stats().sets, 2);

    assert_eq!(
        caching
            .get_bonsai_from_globalrev(&ctx, GLOBALREV_ONE)
            .await?,
        Some(bonsai::ONES_CSID)
    );

    assert_eq!(store.stats().gets, 4);
    assert_eq!(store.stats().hits, 1);
    assert_eq!(store.stats().sets, 3);

    assert_eq!(
        caching
            .get_bonsai_from_globalrev(&ctx, GLOBALREV_ONE)
            .await?,
        Some(bonsai::ONES_CSID)
    );

    assert_eq!(store.stats().gets, 5);
    assert_eq!(store.stats().hits, 2);
    assert_eq!(store.stats().sets, 3);

    Ok(())
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "bonsai_hg_mapping crate
mercurial_thrift mercurial_thrift
mononoke_types_thrift mononoke_types_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=bonsai_hg_mapping_entry_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "bonsai_hg_mapping.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use bonsai_hg_mapping_entry_thrift as thrift;
use bytes::Bytes;
use caching_ext::get_or_fill_chunked;
use caching_ext::CacheDisposition;
use caching_ext::CacheHandlerFactory;
use caching_ext::CacheTtl;
use caching_ext::CachelibHandler;
use caching_ext::EntityStore;
use caching_ext::KeyedEntityStore;
use caching_ext::McErrorKind;
use caching_ext::McResult;
use caching_ext::MemcacheEntity;
use caching_ext::MemcacheHandler;
use context::CoreContext;
use fbthrift::compact_protocol;
use memcache::KeyGen;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use stats::prelude::*;
use tunables::tunables;

use super::BonsaiHgMapping;
use super::BonsaiHgMappingEntry;
use super::BonsaiOrHgChangesetIds;

define_stats! {
    prefix = "mononoke.bonsai_hg_mapping";
    memcache_hit: timeseries("memcache.hit"; Rate, Sum),
    memcache_miss: timeseries("memcache.miss"; Rate, Sum),
    memcache_internal_err: timeseries("memcache.internal_err"; Rate, Sum),
    memcache_deserialize_err: timeseries("memcache.deserialize_err"; Rate, Sum),
}

#[derive(Abomonation, Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiHgMappingCacheEntry {
    pub repo_id: RepositoryId,
    pub hg_cs_id: HgChangesetId,
    pub bcs_id: ChangesetId,
}

impl BonsaiHgMappingCacheEntry {
    fn from_thrift(
        entry: bonsai_hg_mapping_entry_thrift::BonsaiHgMappingCacheEntry,
    ) -> Result<Self> {
        Ok(Self {
            repo_id: RepositoryId::new(entry.repo_id.0),
            hg_cs_id: HgChangesetId::new(HgNodeHash::from_thrift(entry.hg_cs_id)?),
            bcs_id: ChangesetId::from_thrift(entry.bcs_id)?,
        })
    }

    fn into_thrift(self) -> bonsai_hg_mapping_entry_thrift::BonsaiHgMappingCacheEntry {
        bonsai_hg_mapping_entry_thrift::BonsaiHgMappingCacheEntry {
            repo_id: bonsai_hg_mapping_entry_thrift::RepoId(self.repo_id.id()),
            hg_cs_id: self.hg_cs_id.into_nodehash().into_thrift(),
            bcs_id: self.bcs_id.into_thrift(),
        }
    }

    pub fn new(repo_id: RepositoryId, hg_cs_id: HgChangesetId, bcs_id: ChangesetId) -> Self {
        BonsaiHgMappingCacheEntry {
            repo_id,
            hg_cs_id,
            bcs_id,
        }
    }

    pub fn into_entry(self, repo_id: RepositoryId) -> Result<BonsaiHgMappingEntry> {
        if self.repo_id == repo_id {
            Ok(BonsaiHgMappingEntry {
                hg_cs_id: self.hg_cs_id,
                bcs_id: self.bcs_id,
            })
        } else {
            Err(anyhow!(
                "Cache returned invalid entry: repo {} returned for query to repo {}",
                self.repo_id,
                repo_id
            ))
        }
    }

    pub fn from_entry(
        entry: BonsaiHgMappingEntry,
        repo_id: RepositoryId,
    ) -> BonsaiHgMappingCacheEntry {
        BonsaiHgMappingCacheEntry {
            repo_id,
            hg_cs_id: entry.hg_cs_id,
            bcs_id: entry.bcs_id,
        }
    }
}

/// Used for cache key generation
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
enum BonsaiOrHgChangesetId {
    Bonsai(ChangesetId),
    Hg(HgChangesetId),
}

impl From<ChangesetId> for BonsaiOrHgChangesetId {
    fn from(cs_id: ChangesetId) -> Self {
        BonsaiOrHgChangesetId::Bonsai(cs_id)
    }
}

impl From<HgChangesetId> for BonsaiOrHgChangesetId {
    fn from(cs_id: HgChangesetId) -> Self {
        BonsaiOrHgChangesetId::Hg(cs_id)
    }
}

pub struct CachingBonsaiHgMapping {
    mapping: Arc<dyn BonsaiHgMapping>,
    cachelib: CachelibHandler<BonsaiHgMappingCacheEntry>,
    memcache: MemcacheHandler,
    keygen: KeyGen,
}

impl CachingBonsaiHgMapping {
    pub fn new(
        mapping: Arc<dyn BonsaiHgMapping>,
        cache_handler_factory: CacheHandlerFactory,
    ) -> Self {
        Self {
            mapping,
            cachelib: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
            keygen: CachingBonsaiHgMapping::create_key_gen(),
        }
    }

    pub fn new_test(mapping: Arc<dyn BonsaiHgMapping>) -> Self {
        Self::new(mapping, CacheHandlerFactory::Mocked)
    }

    fn create_key_gen() -> KeyGen {
        let key_prefix = "scm.mononoke.bonsai_hg_mapping";

        let sitever = if tunables().bonsai_hg_mapping_sitever().unwrap_or_default() > 0 {
            tunables().bonsai_hg_mapping_sitever().unwrap_or_default() as u32
        } else {
            thrift::MC_SITEVER as u32
        };

        KeyGen::new(key_prefix, thrift::MC_CODEVER as u32, sitever)
    }
}

fn memcache_deserialize(bytes: Bytes) -> McResult<BonsaiHgMappingCacheEntry> {
    let thrift_entry =
        compact_protocol::deserialize(bytes).map_err(|_| McErrorKind::Deserialization);
    thrift_entry.and_then(|entry| {
        BonsaiHgMappingCacheEntry::from_thrift(entry).map_err(|_| McErrorKind::Deserialization)
    })
}

fn memcache_serialize(entry: &BonsaiHgMappingCacheEntry) -> Bytes {
    compact_protocol::serialize(&entry.clone().into_thrift())
}

const CHUNK_SIZE: usize = 1000;
const PARALLEL_CHUNKS: usize = 1;

#[async_trait]
impl BonsaiHgMapping for CachingBonsaiHgMapping {
    fn repo_id(&self) -> RepositoryId {
        self.mapping.repo_id()
    }

    async fn add(&self, ctx: &CoreContext, entry: BonsaiHgMappingEntry) -> Result<bool, Error> {
        self.mapping.add(ctx, entry).await
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        cs: BonsaiOrHgChangesetIds,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error> {
        let cache_request = (ctx, self);
        let repo_id = self.repo_id();

        let cache_entry = match cs {
            BonsaiOrHgChangesetIds::Bonsai(cs_ids) => get_or_fill_chunked(
                &cache_request,
                cs_ids.into_iter().collect(),
                CHUNK_SIZE,
                PARALLEL_CHUNKS,
            )
            .await?
            .into_values()
            .map(|val| val.into_entry(repo_id))
            .collect::<Result<_>>()?,
            BonsaiOrHgChangesetIds::Hg(hg_ids) => get_or_fill_chunked(
                &cache_request,
                hg_ids.into_iter().collect(),
                CHUNK_SIZE,
                PARALLEL_CHUNKS,
            )
            .await?
            .into_values()
            .map(|val| val.into_entry(repo_id))
            .collect::<Result<_>>()?,
        };

        Ok(cache_entry)
    }

    /// Use caching for the ranges of one element, use slower path otherwise.
    async fn get_hg_in_range(
        &self,
        ctx: &CoreContext,
        low: HgChangesetId,
        high: HgChangesetId,
        limit: usize,
    ) -> Result<Vec<HgChangesetId>, Error> {
        if low == high {
            let res = self.get(ctx, low.into()).await?;
            if res.is_empty() {
                return Ok(vec![]);
            } else {
                return Ok(vec![low]);
            }
        }

        self.mapping.get_hg_in_range(ctx, low, high, limit).await
    }
}

fn get_cache_key(repo_id: RepositoryId, cs: &BonsaiOrHgChangesetId) -> String {
    format!("{}.{:?}", repo_id.prefix(), cs)
}

impl MemcacheEntity for BonsaiHgMappingCacheEntry {
    fn serialize(&self) -> Bytes {
        memcache_serialize(self)
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        memcache_deserialize(bytes)
    }
}

type CacheRequest<'a> = (&'a CoreContext, &'a CachingBonsaiHgMapping);

impl EntityStore<BonsaiHgMappingCacheEntry> for CacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<BonsaiHgMappingCacheEntry> {
        let (_, mapping) = self;
        &mapping.cachelib
    }

    fn keygen(&self) -> &KeyGen {
        let (_, mapping) = self;
        &mapping.keygen
    }

    fn memcache(&self) -> &MemcacheHandler {
        let (_, mapping) = self;
        &mapping.memcache
    }

    fn cache_determinator(&self, _: &BonsaiHgMappingCacheEntry) -> CacheDisposition {
        CacheDisposition::Cache(CacheTtl::NoTtl)
    }

    caching_ext::impl_singleton_stats!("bonsai_hg_mapping");
}

#[async_trait]
impl KeyedEntityStore<ChangesetId, BonsaiHgMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &ChangesetId) -> String {
        let (_, mapping) = self;
        get_cache_key(mapping.repo_id(), &BonsaiOrHgChangesetId::Bonsai(*key))
    }

    async fn get_from_db(
        &self,
        keys: HashSet<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, BonsaiHgMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .mapping
            .get(
                ctx,
                BonsaiOrHgChangesetIds::Bonsai(keys.into_iter().collect()),
            )
            .await?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| (e.bcs_id, BonsaiHgMappingCacheEntry::from_entry(e, repo_id)))
                .collect(),
        )
    }
}

#[async_trait]
impl KeyedEntityStore<HgChangesetId, BonsaiHgMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &HgChangesetId) -> String {
        let (_, mapping) = self;
        get_cache_key(mapping.repo_id(), &BonsaiOrHgChangesetId::Hg(*key))
    }

    async fn get_from_db(
        &self,
        keys: HashSet<HgChangesetId>,
    ) -> Result<HashMap<HgChangesetId, BonsaiHgMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .mapping
            .get(ctx, BonsaiOrHgChangesetIds::Hg(keys.into_iter().collect()))
            .await?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| {
                    (
                        e.hg_cs_id,
                        BonsaiHgMappingCacheEntry::from_entry(e, repo_id),
                    )
                })
                .collect(),
        )
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use thiserror::Error;

use super::BonsaiHgMappingEntry;

#[derive(Debug, Eq, Error, PartialEq)]
pub enum ErrorKind {
    #[error("Connection error")]
    ConnectionError,
    #[error("Conflicting entries: stored:{0:?} current:{1:?}")]
    ConflictingEntries(BonsaiHgMappingEntry, BonsaiHgMappingEntry),
    #[error("Conflict detected during insert, but no value was there for: {0:?}")]
    RaceConditionWithDelete(BonsaiHgMappingEntry),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use fbinit::FacebookInit;
use futures::future;
use mercurial_types::HgChangesetId;
use mercurial_types::HgChangesetIdPrefix;
use mercurial_types::HgChangesetIdsResolvedFromPrefix;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use rand::Rng;
use rendezvous::RendezVous;
use rendezvous::RendezVousOptions;
use rendezvous::RendezVousStats;
use rendezvous::TunablesRendezVousController;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use stats::prelude::*;

mod caching;
mod errors;
mod mem_writes_bonsai_hg_mapping;

pub use crate::caching::CachingBonsaiHgMapping;
pub use crate::errors::ErrorKind;
pub use crate::mem_writes_bonsai_hg_mapping::MemWritesBonsaiHgMapping;

define_stats! {
    prefix = "mononoke.bonsai_hg_mapping";
    gets: timeseries(Rate, Sum),
    gets_master: timeseries(Rate, Sum),
    adds: timeseries(Rate, Sum),
    get_many_hg_by_prefix: timeseries(Rate, Sum),
}

#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub struct BonsaiHgMappingEntry {
    pub hg_cs_id: HgChangesetId,
    pub bcs_id: ChangesetId,
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum BonsaiOrHgChangesetIds {
    Bonsai(Vec<ChangesetId>),
    Hg(Vec<HgChangesetId>),
}

impl BonsaiOrHgChangesetIds {
    pub fn is_empty(&self) -> bool {
        match self {
            BonsaiOrHgChangesetIds::Bonsai(v) => v.is_empty(),
            BonsaiOrHgChangesetIds::Hg(v) => v.is_empty(),
        }
    }
}

impl From<ChangesetId> for BonsaiOrHgChangesetIds {
    fn from(cs_id: ChangesetId) -> Self {
        BonsaiOrHgChangesetIds::Bonsai(vec![cs_id])
    }
}

impl From<Vec<ChangesetId>> for BonsaiOrHgChangesetIds {
    fn from(cs_ids: Vec<ChangesetId>) -> Self {
        BonsaiOrHgChangesetIds::Bonsai(cs_ids)
    }
}

impl From<HgChangesetId> for BonsaiOrHgChangesetIds {
    fn from(cs_id: HgChangesetId) -> Self {
        BonsaiOrHgChangesetIds::Hg(vec![cs_id])
    }
}

impl From<Vec<HgChangesetId>> for BonsaiOrHgChangesetIds {
    fn from(cs_ids: Vec<HgChangesetId>) -> Self {
        BonsaiOrHgChangesetIds::Hg(cs_ids)
    }
}

#[facet::facet]
#[async_trait]
pub trait BonsaiHgMapping: Send + Sync {
    fn repo_id(&self) -> RepositoryId;

    async fn add(&self, ctx: &CoreContext, entry: BonsaiHgMappingEntry) -> Result<bool, Error>;

    async fn get(
        &self,
        ctx: &CoreContext,
        cs_id: BonsaiOrHgChangesetIds,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error>;

    async fn get_hg_from_bonsai(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<HgChangesetId>, Error> {
        let result = self.get(ctx, cs_id.into()).await?;
        let hg_cs_id = result.into_iter().next().map(|entry| entry.hg_cs_id);
        Ok(hg_cs_id)
    }

    async fn get_bonsai_from_hg(
        &self,
        ctx: &CoreContext,
        cs_id: HgChangesetId,
    ) -> Result<Option<ChangesetId>, Error> {
        let result = self.get(ctx, cs_id.into()).await?;
        let bcs_id = result.into_iter().next().map(|entry| entry.bcs_id);
        Ok(bcs_id)
    }

    async fn get_many_hg_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: HgChangesetIdPrefix,
        limit: usize,
    ) -> Result<HgChangesetIdsResolvedFromPrefix, Error> {
        let mut fetched_cs = self
            .get_hg_in_range(ctx, cs_prefix.min_cs(), cs_prefix.max_cs(), limit + 1)
            .await?;
        let res = match fetched_cs.len() {
            0 => HgChangesetIdsResolvedFromPrefix::NoMatch,
            1 => HgChangesetIdsResolvedFromPrefix::Single(fetched_cs[0].clone()),
            l if l <= limit => HgChangesetIdsResolvedFromPrefix::Multiple(fetched_cs),
            _ => HgChangesetIdsResolvedFromPrefix::TooMany({
                fetched_cs.pop();
                fetched_cs
            }),
        };
        Ok(res)
    }

    async fn get_hg_in_range(
        &self,
        ctx: &CoreContext,
        low: HgChangesetId,
        high: HgChangesetId,
        limit: usize,
    ) -> Result<Vec<HgChangesetId>, Error>;
}

#[derive(Clone)]
struct RendezVousConnection {
    bonsai: RendezVous<ChangesetId, HgChangesetId>,
    hg: RendezVous<HgChangesetId, ChangesetId>,
    conn: Connection,
}

impl RendezVousConnection {
    fn new(conn: Connection, name: &str, opts: RendezVousOptions) -> Self {
        Self {
            conn,
            bonsai: RendezVous::new(
                TunablesRendezVousController::new(opts),
                Arc::new(RendezVousStats::new(format!(
                    "bonsai_hg_mapping.bonsai.{}",
                    name,
                ))),
            ),
            hg: RendezVous::new(
                TunablesRendezVousController::new(opts),
                Arc::new(RendezVousStats::new(format!(
                    "bonsai_hg_mapping.hg.{}",
                    name,
                ))),
            ),
        }
    }
}

pub struct SqlBonsaiHgMapping {
    write_connection: Connection,
    read_connection: RendezVousConnection,
    read_master_connection: RendezVousConnection,
    repo_id: RepositoryId,
    // Option that forces all `add()` method calls to overwrite values
    // that set in the database. This should be used only when we try to
    // fix broken entries in the db.
    overwrite: bool,
}

mononoke_queries! {
    // Sett almost identical ReplaceMapping below
    write InsertMapping(values: (
        repo_id: RepositoryId,
        hg_cs_id: HgChangesetId,
        bcs_id: ChangesetId,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bonsai_hg_mapping (repo_id, hg_cs_id, bcs_id) VALUES {values}"
    }

    // Sett almost identical InsertMapping above
    write ReplaceMapping(values: (
        repo_id: RepositoryId,
        hg_cs_id: HgChangesetId,
        bcs_id: ChangesetId,
    )) {
        none,
        "REPLACE INTO bonsai_hg_mapping (repo_id, hg_cs_id, bcs_id) VALUES {values}"
    }

    read SelectMappingByBonsai(
        repo_id: RepositoryId,
        tok: i32,
        >list bcs_id: ChangesetId
    ) -> (HgChangesetId, ChangesetId, i32) {
        "SELECT hg_cs_id, bcs_id, {tok}
         FROM bonsai_hg_mapping
         WHERE repo_id = {repo_id}
           AND bcs_id IN {bcs_id}"
    }

    read SelectMappingByHg(
        repo_id: RepositoryId,
        tok: i32,
        >list hg_cs_id: HgChangesetId
    ) -> (HgChangesetId, ChangesetId, i32) {
        "SELECT hg_cs_id, bcs_id, {tok}
         FROM bonsai_hg_mapping
         WHERE repo_id = {repo_id}
           AND hg_cs_id IN {hg_cs_id}"
    }

    read SelectHgChangesetsByRange(repo_id: RepositoryId, hg_cs_min: &[u8], hg_cs_max: &[u8], limit: usize) -> (HgChangesetId) {
        "SELECT hg_cs_id
         FROM bonsai_hg_mapping
         WHERE repo_id = {repo_id}
           AND hg_cs_id >= {hg_cs_min} AND hg_cs_id <= {hg_cs_max}
           LIMIT {limit}
        "
    }
}

#[derive(Clone)]
pub struct SqlBonsaiHgMappingBuilder {
    connections: SqlConnections,
    overwrite: bool,
}

impl SqlConstruct for SqlBonsaiHgMappingBuilder {
    const LABEL: &'static str = "bonsai_hg_mapping";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-bonsai-hg-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            connections,
            overwrite: false,
        }
    }
}

impl SqlBonsaiHgMappingBuilder {
    pub fn with_overwrite(mut self) -> Self {
        self.overwrite = true;
        self
    }

    pub fn build(self, repo_id: RepositoryId, opts: RendezVousOptions) -> SqlBonsaiHgMapping {
        let SqlBonsaiHgMappingBuilder {
            connections,
            overwrite,
        } = self;

        SqlBonsaiHgMapping {
            write_connection: connections.write_connection,
            read_connection: RendezVousConnection::new(connections.read_connection, "reader", opts),
            read_master_connection: RendezVousConnection::new(
                connections.read_master_connection,
                "read_master",
                opts,
            ),
            repo_id,
            overwrite,
        }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlBonsaiHgMappingBuilder {}

impl SqlBonsaiHgMapping {
    async fn verify_consistency(&self, entry: BonsaiHgMappingEntry) -> Result<(), Error> {
        let BonsaiHgMappingEntry { hg_cs_id, bcs_id } = entry.clone();

        let tok: i32 = rand::thread_rng().gen();
        let hg_ids = &[hg_cs_id];
        let by_hg = SelectMappingByHg::query(
            &self.read_master_connection.conn,
            &self.repo_id,
            &tok,
            hg_ids,
        );
        let bcs_ids = &[bcs_id];
        let by_bcs = SelectMappingByBonsai::query(
            &self.read_master_connection.conn,
            &self.repo_id,
            &tok,
            bcs_ids,
        );

        let (by_hg_rows, by_bcs_rows) = future::try_join(by_hg, by_bcs).await?;

        match by_hg_rows
            .into_iter()
            .chain(by_bcs_rows.into_iter())
            .map(|(hg_cs_id, bcs_id, _)| (hg_cs_id, bcs_id))
            .next()
        {
            Some(entry) if entry == (hg_cs_id, bcs_id) => Ok(()),
            Some((hg_cs_id, bcs_id)) => Err(ErrorKind::ConflictingEntries(
                BonsaiHgMappingEntry { hg_cs_id, bcs_id },
                entry,
            )
            .into()),
            None => Err(ErrorKind::RaceConditionWithDelete(entry).into()),
        }
    }
}

#[async_trait]
impl BonsaiHgMapping for SqlBonsaiHgMapping {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn add(&self, ctx: &CoreContext, entry: BonsaiHgMappingEntry) -> Result<bool, Error> {
        STATS::adds.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        let BonsaiHgMappingEntry { hg_cs_id, bcs_id } = entry.clone();

        if self.overwrite {
            let result = ReplaceMapping::query(
                &self.write_connection,
                &[(&self.repo_id, &hg_cs_id, &bcs_id)],
            )
            .await?;
            Ok(result.affected_rows() >= 1)
        } else {
            let result = InsertMapping::query(
                &self.write_connection,
                &[(&self.repo_id, &hg_cs_id, &bcs_id)],
            )
            .await?;
            if result.affected_rows() == 1 {
                Ok(true)
            } else {
                self.verify_consistency(entry).await?;
                Ok(false)
            }
        }
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        ids: BonsaiOrHgChangesetIds,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error> {
        STATS::gets.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let (mut mappings, left_to_fetch) =
            select_mapping(ctx.fb, &self.read_connection, self.repo_id, ids).await?;

        if left_to_fetch.is_empty() {
            return Ok(mappings);
        }

        STATS::gets_master.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);
        let (mut master_mappings, _) = select_mapping(
            ctx.fb,
            &self.read_master_connection,
            self.repo_id,
            left_to_fetch,
        )
        .await?;

        mappings.append(&mut master_mappings);
        Ok(mappings)
    }

    /// Return [`HgChangesetId`] entries in the inclusive range described by `low` and `high`.
    /// Maximum `limit` entries will be returned.
    async fn get_hg_in_range(
        &self,
        ctx: &CoreContext,
        low: HgChangesetId,
        high: HgChangesetId,
        limit: usize,
    ) -> Result<Vec<HgChangesetId>, Error> {
        if low > high {
            return Ok(Vec::new());
        }
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let rows = SelectHgChangesetsByRange::query(
            &self.read_connection.conn,
            &self.repo_id,
            &low.as_bytes(),
            &high.as_bytes(),
            &limit,
        )
        .await?;
        let mut fetched: Vec<HgChangesetId> = rows.into_iter().map(|row| row.0).collect();
        if fetched.is_empty() {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let rows = SelectHgChangesetsByRange::query(
                &self.read_master_connection.conn,
                &self.repo_id,
                &low.as_bytes(),
                &high.as_bytes(),
                &limit,
            )
            .await?;
            fetched = rows.into_iter().map(|row| row.0).collect();
        }
        Ok(fetched)
    }
}

async fn select_mapping(
    fb: FacebookInit,
    connection: &RendezVousConnection,
    repo_id: RepositoryId,
    cs_ids: BonsaiOrHgChangesetIds,
) -> Result<(Vec<BonsaiHgMappingEntry>, BonsaiOrHgChangesetIds), Error> {
    if cs_ids.is_empty() {
        return Ok((vec![], cs_ids));
    }

    let tok: i32 = rand::thread_rng().gen();

    let (found, missing): (Vec<_>, _) = match cs_ids {
        BonsaiOrHgChangesetIds::Bonsai(bcs_ids) => {
            let ret = connection
                .bonsai
                .dispatch(fb, bcs_ids.into_iter().collect(), || {
                    let conn = connection.conn.clone();
                    move |bcs_ids| async move {
                        let bcs_ids = bcs_ids.into_iter().collect::<Vec<_>>();

                        Ok(
                            SelectMappingByBonsai::query(&conn, &repo_id, &tok, &bcs_ids[..])
                                .await?
                                .into_iter()
                                .map(|(hg_cs_id, bcs_id, _)| (bcs_id, hg_cs_id))
                                .collect(),
                        )
                    }
                })
                .await?;

            let mut not_found = vec![];
            let found = ret
                .into_iter()
                .filter_map(|(bcs_id, hg_cs_id)| match hg_cs_id {
                    Some(hg_cs_id) => Some((hg_cs_id, bcs_id)),
                    None => {
                        not_found.push(bcs_id);
                        None
                    }
                })
                .collect();

            (found, BonsaiOrHgChangesetIds::Bonsai(not_found))
        }
        BonsaiOrHgChangesetIds::Hg(hg_cs_ids) => {
            let ret = connection
                .hg
                .dispatch(fb, hg_cs_ids.into_iter().collect(), || {
                    let conn = connection.conn.clone();
                    move |hg_cs_ids| async move {
                        let hg_cs_ids = hg_cs_ids.into_iter().collect::<Vec<_>>();
                        Ok(
                            SelectMappingByHg::query(&conn, &repo_id, &tok, &hg_cs_ids[..])
                                .await?
                                .into_iter()
                                .map(|(hg_cs_id, bcs_id, _)| (hg_cs_id, bcs_id))
                                .collect(),
                        )
                    }
                })
                .await?;

            let mut not_found = vec![];
            let found = ret
                .into_iter()
                .filter_map(|(hg_cs_id, bcs_id)| match bcs_id {
                    Some(bcs_id) => Some((hg_cs_id, bcs_id)),
                    None => {
                        not_found.push(hg_cs_id);
                        None
                    }
                })
                .collect();

            (found, BonsaiOrHgChangesetIds::Hg(not_found))
        }
    };

    Ok((
        found
            .into_iter()
            .map(move |(hg_cs_id, bcs_id)| BonsaiHgMappingEntry { hg_cs_id, bcs_id })
            .collect(),
        missing,
    ))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::hash::Hash;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::anyhow;
use anyhow::Error;
use async_trait::async_trait;
use context::CoreContext;
use lock_ext::LockExt;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;

use crate::BonsaiHgMapping;
use crate::BonsaiHgMappingEntry;
use crate::BonsaiOrHgChangesetIds;

type Cache = (
    HashMap<ChangesetId, HgChangesetId>,
    HashMap<HgChangesetId, ChangesetId>,
);

pub struct MemWritesBonsaiHgMapping {
    inner: Arc<dyn BonsaiHgMapping>,
    cache: Arc<Mutex<Cache>>,
    no_access_to_inner: Arc<AtomicBool>,
    readonly: Arc<AtomicBool>,
    save_noop_writes: Arc<AtomicBool>,
}

impl MemWritesBonsaiHgMapping {
    pub fn new(inner: Arc<dyn BonsaiHgMapping>) -> Self {
        Self {
            inner,
            cache: Default::default(),
            no_access_to_inner: Default::default(),
            readonly: Default::default(),
            save_noop_writes: Default::default(),
        }
    }

    pub fn set_readonly(&self, readonly: bool) {
        self.readonly.store(readonly, Ordering::Relaxed);
    }

    pub fn set_no_access_to_inner(&self, no_access_to_inner: bool) {
        self.no_access_to_inner
            .store(no_access_to_inner, Ordering::Relaxed);
    }

    pub fn set_save_noop_writes(&self, save_noop_writes: bool) {
        self.save_noop_writes
            .store(save_noop_writes, Ordering::Relaxed);
    }
}

impl MemWritesBonsaiHgMapping {
    async fn get_from_cache_and_inner<I, O>(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<I>,
        get_cache: impl Fn(&Cache) -> &HashMap<I, O>,
        make_entry: impl Fn(I, O) -> BonsaiHgMappingEntry,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error>
    where
        Vec<I>: Into<BonsaiOrHgChangesetIds>,
        I: Eq + Hash + Copy,
        O: Copy,
    {
        let mut from_cache = vec![];
        let mut from_inner = vec![];

        for i in cs_ids {
            self.cache.with(|cache| {
                let cache = get_cache(cache);

                match cache.get(&i).copied() {
                    Some(o) => from_cache.push(make_entry(i, o)),
                    None => from_inner.push(i),
                };
            });
        }

        if !self.no_access_to_inner.load(Ordering::Relaxed) {
            let from_inner = self.inner.get(ctx, from_inner.into()).await?;
            from_cache.extend(from_inner);
        }
        Ok(from_cache)
    }
}

#[async_trait]
impl BonsaiHgMapping for MemWritesBonsaiHgMapping {
    fn repo_id(&self) -> RepositoryId {
        self.inner.repo_id()
    }

    async fn add(&self, ctx: &CoreContext, entry: BonsaiHgMappingEntry) -> Result<bool, Error> {
        if self.readonly.load(Ordering::Relaxed) {
            return Err(anyhow!(
                "cannot write to a readonly MemWritesBonsaiHgMapping"
            ));
        }

        let BonsaiHgMappingEntry { hg_cs_id, bcs_id } = entry;

        let entry = self.get_hg_from_bonsai(ctx, bcs_id).await?;
        if entry.is_some() && !self.save_noop_writes.load(Ordering::Relaxed) {
            Ok(false)
        } else {
            self.cache.with(|cache| {
                cache.0.insert(bcs_id, hg_cs_id);
                cache.1.insert(hg_cs_id, bcs_id);
            });
            Ok(true)
        }
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        cs_ids: BonsaiOrHgChangesetIds,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error> {
        match cs_ids {
            BonsaiOrHgChangesetIds::Bonsai(bcs_ids) => {
                self.get_from_cache_and_inner(
                    ctx,
                    bcs_ids,
                    |cache| &cache.0,
                    |bcs_id, hg_cs_id| BonsaiHgMappingEntry { bcs_id, hg_cs_id },
                )
                .await
            }
            BonsaiOrHgChangesetIds::Hg(hg_cs_ids) => {
                self.get_from_cache_and_inner(
                    ctx,
                    hg_cs_ids,
                    |cache| &cache.1,
                    |hg_cs_id, bcs_id| BonsaiHgMappingEntry { bcs_id, hg_cs_id },
                )
                .await
            }
        }
    }

    async fn get_hg_in_range(
        &self,
        _ctx: &CoreContext,
        _low: HgChangesetId,
        _high: HgChangesetId,
        _limit: usize,
    ) -> Result<Vec<HgChangesetId>, Error> {
        unimplemented!("This is not currently implemented in Gitimport")
    }
}