/*
 * 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 save_mapping_pushrebase_hook;
mod sql_queries;
#[cfg(test)]
mod test;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use pushrebase_hook::PushrebaseHook;
pub use sql_queries::add_pushrebase_mapping;
pub use sql_queries::get_prepushrebase_ids;
pub use sql_queries::SqlPushrebaseMutationMapping;
pub use sql_queries::SqlPushrebaseMutationMappingConnection;

pub struct PushrebaseMutationMappingEntry {
    repo_id: RepositoryId,
    predecessor_bcs_id: ChangesetId,
    successor_bcs_id: ChangesetId,
}

impl PushrebaseMutationMappingEntry {
    fn new(
        repo_id: RepositoryId,
        predecessor_bcs_id: ChangesetId,
        successor_bcs_id: ChangesetId,
    ) -> Self {
        Self {
            repo_id,
            predecessor_bcs_id,
            successor_bcs_id,
        }
    }
}

#[async_trait]
#[facet::facet]
pub trait PushrebaseMutationMapping: Send + Sync {
    fn get_hook(&self) -> Option<Box<dyn PushrebaseHook>>;
    async fn get_prepushrebase_ids(
        &self,
        ctx: &CoreContext,
        successor_bcs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>>;
}


/*
 * 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::Result;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use pushrebase_hook::PushrebaseHook;
use sql::Connection;
use sql::Transaction;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use tunables::tunables;

use crate::save_mapping_pushrebase_hook::SaveMappingPushrebaseHook;
use crate::PushrebaseMutationMapping;
use crate::PushrebaseMutationMappingEntry;

mononoke_queries! {
    read SelectPrepushrebaseIds(
        repo_id: RepositoryId,
        successor_bcs_id: ChangesetId,
    ) -> (ChangesetId,) {
        "SELECT predecessor_bcs_id
        FROM pushrebase_mutation_mapping
        WHERE repo_id = {repo_id} AND successor_bcs_id = {successor_bcs_id}"
    }

    write InsertMappingEntries(values:(
        repo_id: RepositoryId,
        predecessor_bcs_id: ChangesetId,
        successor_bcs_id: ChangesetId,
    )) {
        insert_or_ignore,
       "{insert_or_ignore}
       INTO pushrebase_mutation_mapping
       (repo_id, predecessor_bcs_id, successor_bcs_id)
       VALUES {values}"
    }
}

pub async fn add_pushrebase_mapping(
    transaction: Transaction,
    entries: &[PushrebaseMutationMappingEntry],
) -> Result<Transaction> {
    let entries: Vec<_> = entries
        .iter()
        .map(
            |PushrebaseMutationMappingEntry {
                 repo_id,
                 predecessor_bcs_id,
                 successor_bcs_id,
             }| (repo_id, predecessor_bcs_id, successor_bcs_id),
        )
        .collect();

    let (transaction, _) =
        InsertMappingEntries::query_with_transaction(transaction, &entries).await?;

    Ok(transaction)
}

pub async fn get_prepushrebase_ids(
    connection: &Connection,
    repo_id: RepositoryId,
    successor_bcs_id: ChangesetId,
) -> Result<Vec<ChangesetId>> {
    let rows = SelectPrepushrebaseIds::query(connection, &repo_id, &successor_bcs_id).await?;

    Ok(rows.into_iter().map(|r| r.0).collect())
}

pub struct SqlPushrebaseMutationMapping {
    repo_id: RepositoryId,
    sql_conn: SqlPushrebaseMutationMappingConnection,
}

impl SqlPushrebaseMutationMapping {
    pub fn new(repo_id: RepositoryId, sql_conn: SqlPushrebaseMutationMappingConnection) -> Self {
        Self { repo_id, sql_conn }
    }
}

#[derive(Clone)]
pub struct SqlPushrebaseMutationMappingConnection {
    #[allow(dead_code)]
    write_connection: Connection,
    read_connection: Connection,
    read_master_connection: Connection,
}

impl SqlPushrebaseMutationMappingConnection {
    pub fn with_repo_id(self, repo_id: RepositoryId) -> SqlPushrebaseMutationMapping {
        SqlPushrebaseMutationMapping::new(repo_id, self)
    }

    async fn get_prepushrebase_ids(
        &self,
        ctx: &CoreContext,
        repo_id: RepositoryId,
        successor_bcs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let mut ids =
            get_prepushrebase_ids(&self.read_connection, repo_id, successor_bcs_id).await?;
        if ids.is_empty() {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            ids = get_prepushrebase_ids(&self.read_master_connection, repo_id, successor_bcs_id)
                .await?;
        }
        Ok(ids)
    }
}

impl SqlConstruct for SqlPushrebaseMutationMappingConnection {
    const LABEL: &'static str = "pushrebase_mutation_mapping";

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

    // We don't need the connections because we never use them.
    // But we need SqlConstruct to get our SQL tables created in tests.
    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            write_connection: connections.write_connection,
            read_connection: connections.read_connection,
            read_master_connection: connections.read_master_connection,
        }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlPushrebaseMutationMappingConnection {}

#[async_trait]
impl PushrebaseMutationMapping for SqlPushrebaseMutationMapping {
    fn get_hook(&self) -> Option<Box<dyn PushrebaseHook>> {
        if tunables()
            .disable_save_mapping_pushrebase_hook()
            .unwrap_or_default()
        {
            None
        } else {
            Some(SaveMappingPushrebaseHook::new(self.repo_id))
        }
    }

    async fn get_prepushrebase_ids(
        &self,
        ctx: &CoreContext,
        successor_bcs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        self.sql_conn
            .get_prepushrebase_ids(ctx, self.repo_id, successor_bcs_id)
            .await
    }
}


/*
 * 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::Result;
use fbinit::FacebookInit;
use mononoke_types_mocks::changesetid;
use mononoke_types_mocks::repo;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_ext::open_sqlite_in_memory;

use crate::add_pushrebase_mapping;
use crate::get_prepushrebase_ids;
use crate::PushrebaseMutationMappingEntry;
use crate::SqlPushrebaseMutationMappingConnection;

#[fbinit::test]
async fn test_add_and_get(_fb: FacebookInit) -> Result<()> {
    let conn = open_sqlite_in_memory()?;
    conn.execute_batch(SqlPushrebaseMutationMappingConnection::CREATION_QUERY)?;
    let conn = Connection::with_sqlite(conn);

    let entries = vec![
        PushrebaseMutationMappingEntry::new(
            repo::REPO_ZERO,
            changesetid::ONES_CSID,
            changesetid::TWOS_CSID,
        ),
        PushrebaseMutationMappingEntry::new(
            repo::REPO_ONE,
            changesetid::ONES_CSID,
            changesetid::TWOS_CSID,
        ),
        PushrebaseMutationMappingEntry::new(
            repo::REPO_ONE,
            changesetid::TWOS_CSID,
            changesetid::TWOS_CSID,
        ),
        PushrebaseMutationMappingEntry::new(
            repo::REPO_ONE,
            changesetid::ONES_CSID,
            changesetid::THREES_CSID,
        ),
    ];

    let txn = conn.start_transaction().await?;
    let txn = add_pushrebase_mapping(txn, &entries).await?;
    txn.commit().await?;

    let mut prepushrebase_ids =
        get_prepushrebase_ids(&conn, repo::REPO_ONE, changesetid::TWOS_CSID).await?;
    prepushrebase_ids.sort();

    assert_eq!(
        prepushrebase_ids,
        vec![changesetid::ONES_CSID, changesetid::TWOS_CSID]
    );

    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.
 */

#[cfg(test)]
mod test;

use anyhow::Result;
use async_trait::async_trait;
use bookmarks::BookmarkTransactionError;
use context::CoreContext;
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;

use crate::sql_queries::add_pushrebase_mapping;
use crate::PushrebaseMutationMappingEntry;

pub struct SaveMappingPushrebaseHook {
    repository_id: RepositoryId,
}

impl SaveMappingPushrebaseHook {
    pub fn new(repository_id: RepositoryId) -> Box<dyn PushrebaseHook> {
        Box::new(Self { repository_id })
    }
}

#[async_trait]
impl PushrebaseHook for SaveMappingPushrebaseHook {
    async fn in_critical_section(&self) -> Result<Box<dyn PushrebaseCommitHook>> {
        Ok(Box::new(SaveMappingCommitHook {
            repository_id: self.repository_id,
        }))
    }
}

pub struct SaveMappingCommitHook {
    repository_id: RepositoryId,
}

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

    async fn into_transaction_hook(
        self: Box<Self>,
        _ctx: &CoreContext,
        rebased: &RebasedChangesets,
    ) -> Result<Box<dyn PushrebaseTransactionHook>> {
        let entries = rebased
            .iter()
            .map(|(predecessor_bcs_id, (successor_bcs_id, _))| {
                PushrebaseMutationMappingEntry::new(
                    self.repository_id,
                    *predecessor_bcs_id,
                    *successor_bcs_id,
                )
            })
            .collect();
        Ok(Box::new(SaveMappingTransactionHook { entries }))
    }
}

struct SaveMappingTransactionHook {
    entries: Vec<PushrebaseMutationMappingEntry>,
}

#[async_trait]
impl PushrebaseTransactionHook for SaveMappingTransactionHook {
    async fn populate_transaction(
        &self,
        _ctx: &CoreContext,
        txn: Transaction,
    ) -> Result<Transaction, BookmarkTransactionError> {
        let txn = add_pushrebase_mapping(txn, &self.entries[..]).await?;
        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 anyhow::Result;
use blobrepo::BlobRepo;
use blobstore::Loadable;
use borrowed::borrowed;
use context::CoreContext;
use fbinit::FacebookInit;
use maplit::hashset;
use metaconfig_types::PushrebaseFlags;
use mononoke_types_mocks::repo;
use pushrebase::do_pushrebase_bonsai;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use test_repo_factory::TestRepoFactory;
use tests_utils::bookmark;
use tests_utils::CreateCommitContext;

use super::SaveMappingPushrebaseHook;
use crate::get_prepushrebase_ids;

#[fbinit::test]
async fn pushrebase_saves_mapping(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);

    let mut repo_factory = TestRepoFactory::new(fb)?;
    let repo: BlobRepo = repo_factory.with_id(repo::REPO_ONE).build().await?;

    borrowed!(ctx, repo);

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

    let master = bookmark(ctx, repo, "master").set_to(root).await?;
    let main = bookmark(ctx, repo, "main").set_to(root).await?;

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

    let hooks = [SaveMappingPushrebaseHook::new(repo.repo_identity().id())];

    // Pushrebase the same commit onto different bookmarks that are pointing to
    // the same commit (root).
    do_pushrebase_bonsai(
        ctx,
        repo,
        &PushrebaseFlags {
            rewritedates: false,
            ..Default::default()
        },
        &master,
        &hashset![cs.clone()],
        &hooks,
    )
    .await?;

    let rebased = do_pushrebase_bonsai(
        ctx,
        repo,
        &PushrebaseFlags {
            rewritedates: false,
            ..Default::default()
        },
        &main,
        &hashset![cs.clone()],
        &hooks,
    )
    .await?
    .head;

    let prepushrebase_ids = get_prepushrebase_ids(
        &repo_factory.metadata_db().read_connection,
        repo.repo_identity().id(),
        rebased,
    )
    .await?;

    assert_eq!(
        prepushrebase_ids,
        vec![cs.get_changeset_id(), cs.get_changeset_id()]
    );

    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::io::Error;
use std::io::ErrorKind;
use std::pin::Pin;

use futures::ready;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::io::ReadBuf;

#[pin_project]
pub struct QuietShutdownStream<T> {
    #[pin]
    inner: T,
}

impl<T> QuietShutdownStream<T> {
    pub fn new(inner: T) -> Self {
        Self { inner }
    }
}

impl<T> AsyncRead for QuietShutdownStream<T>
where
    T: AsyncRead,
{
    #[inline]
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<(), Error>> {
        let this = self.project();
        this.inner.poll_read(cx, buf)
    }
}

impl<T> AsyncWrite for QuietShutdownStream<T>
where
    T: AsyncWrite,
{
    #[inline]
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &[u8],
    ) -> Poll<Result<usize, Error>> {
        let this = self.project();
        this.inner.poll_write(cx, buf)
    }

    #[inline]
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
        let this = self.project();
        this.inner.poll_flush(cx)
    }

    #[inline]
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
        // This is useful to wrap a SslStream. See here for why:
        // https://github.com/sfackler/tokio-openssl/issues/27
        let this = self.project();
        let res = ready!(this.inner.poll_shutdown(cx));
        let res = match res {
            Ok(r) => Ok(r),
            Err(e) if e.kind() == ErrorKind::NotConnected => Ok(()),
            Err(e) if e.kind() == ErrorKind::BrokenPipe => Ok(()),
            Err(e) => Err(e),
        };
        Poll::Ready(res)
    }
}


/*
 * 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::str::FromStr;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use serde::de::Deserializer;
use serde::de::Error as _;
use serde::Deserialize;

#[cfg(fbcode_build)]
pub use crate::facebook::get_region_capacity;
#[cfg(not(fbcode_build))]
pub use crate::oss::get_region_capacity;
use crate::LoadShedLimit;
use crate::Metric;
use crate::MononokeRateLimitConfig;
use crate::RateLimit;
use crate::RateLimitBody;
use crate::StaticSlice;
use crate::Target;

impl TryFrom<rate_limiting_config::Target> for Target {
    type Error = Error;

    fn try_from(value: rate_limiting_config::Target) -> Result<Self, Self::Error> {
        match value {
            rate_limiting_config::Target::not_target(t) => {
                Ok(Target::NotTarget(Box::new((*t).try_into()?)))
            }
            rate_limiting_config::Target::and_target(t) => Ok(Target::AndTarget(
                t.into_iter()
                    .map(|t| t.try_into())
                    .collect::<Result<Vec<_>, _>>()?,
            )),
            rate_limiting_config::Target::or_target(t) => Ok(Target::OrTarget(
                t.into_iter()
                    .map(|t| t.try_into())
                    .collect::<Result<Vec<_>, _>>()?,
            )),
            rate_limiting_config::Target::identity(i) => {
                Ok(Target::Identity(FromStr::from_str(&i)?))
            }
            rate_limiting_config::Target::static_slice(s) => {
                let slice_pct = s.slice_pct.try_into()?;
                Ok(Target::StaticSlice(StaticSlice {
                    slice_pct,
                    nonce: s.nonce,
                }))
            }
            _ => Err(anyhow!("Invalid target")),
        }
    }
}

impl TryFrom<rate_limiting_config::RateLimitBody> for RateLimitBody {
    type Error = Error;

    fn try_from(value: rate_limiting_config::RateLimitBody) -> Result<Self, Self::Error> {
        let window: u64 = value.window.try_into().context("Invalid window")?;

        Ok(Self {
            raw_config: value,
            window: Duration::from_secs(window),
        })
    }
}

impl TryFrom<rate_limiting_config::RegionalMetric> for Metric {
    type Error = Error;

    fn try_from(value: rate_limiting_config::RegionalMetric) -> Result<Self, Self::Error> {
        match value {
            rate_limiting_config::RegionalMetric::EgressBytes => Ok(Metric::EgressBytes),
            rate_limiting_config::RegionalMetric::TotalManifests => Ok(Metric::TotalManifests),
            rate_limiting_config::RegionalMetric::GetpackFiles => Ok(Metric::GetpackFiles),
            rate_limiting_config::RegionalMetric::Commits => Ok(Metric::Commits),
            _ => Err(anyhow!("Invalid RegionalMetric")),
        }
    }
}

impl TryFrom<rate_limiting_config::RateLimit> for RateLimit {
    type Error = Error;

    fn try_from(value: rate_limiting_config::RateLimit) -> Result<Self, Self::Error> {
        let body = value
            .limit
            .clone()
            .try_into()
            .context("Invalid limit body")?;

        let target = value
            .target
            .clone()
            .map(Target::try_from)
            .transpose()
            .context("Invalid target")?;

        let metric = value.metric.clone().try_into().context("Invalid metric")?;

        Ok(Self {
            body,
            metric,
            target,
        })
    }
}

impl TryFrom<rate_limiting_config::LoadShedLimit> for LoadShedLimit {
    type Error = Error;

    fn try_from(value: rate_limiting_config::LoadShedLimit) -> Result<Self, Self::Error> {
        let target = value
            .target
            .clone()
            .map(Target::try_from)
            .transpose()
            .context("Invalid target")?;

        Ok(Self {
            raw_config: value,
            target,
        })
    }
}

impl<'de> Deserialize<'de> for LoadShedLimit {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let raw = rate_limiting_config::LoadShedLimit::deserialize(deserializer)?;
        let load_shed_limit = raw
            .try_into()
            .map_err(|e| D::Error::custom(format!("{:?}", e)))?;

        Ok(load_shed_limit)
    }
}

impl<'de> Deserialize<'de> for MononokeRateLimitConfig {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let raw_config = rate_limiting_config::MononokeRateLimits::deserialize(deserializer)?;

        let dc_prefix_capacity = &raw_config.datacenter_prefix_capacity;

        // We scale the limits used for RegionalMetrics according to the amount of capacity in a
        // region. Calculate the fraction of tasks that this region accounts for.
        let region_weight = match get_region_capacity(dc_prefix_capacity) {
            Some(capacity) => {
                capacity as f64 / dc_prefix_capacity.values().map(|c| *c as f64).sum::<f64>()
            }
            None => 1.0,
        };

        let rate_limits = raw_config
            .rate_limits
            .clone()
            .into_iter()
            .map(|r| r.try_into())
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| D::Error::custom(format!("{:?}", e)))?;

        let load_shed_limits = raw_config
            .load_shed_limits
            .clone()
            .into_iter()
            .map(|r| r.try_into())
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| D::Error::custom(format!("{:?}", e)))?;

        let commits_per_author = raw_config
            .commits_per_author
            .clone()
            .try_into()
            .map_err(|e| D::Error::custom(format!("{:?}", e)))?;

        let total_file_changes = raw_config
            .total_file_changes
            .clone()
            .map(|v| v.try_into())
            .transpose()
            .map_err(|e| D::Error::custom(format!("{:?}", e)))?;

        Ok(Self {
            region_weight,
            rate_limits,
            load_shed_limits,
            commits_per_author,
            total_file_changes,
        })
    }
}


/*
 * 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::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Error;
use async_trait::async_trait;
use cached_config::ConfigHandle;
use fbinit::FacebookInit;
use permission_checker::MononokeIdentity;
use permission_checker::MononokeIdentitySet;
use permission_checker::MononokeIdentitySetExt;
use stats::prelude::*;
use thiserror::Error;

#[cfg(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod oss;

#[cfg(fbcode_build)]
pub use facebook::create_rate_limiter;
#[cfg(fbcode_build)]
pub use facebook::get_region_capacity;
#[cfg(not(fbcode_build))]
pub use oss::create_rate_limiter;
#[cfg(not(fbcode_build))]
pub use oss::get_region_capacity;
pub use rate_limiting_config::RateLimitStatus;

pub mod config;

pub type LoadCost = f64;
pub type BoxRateLimiter = Box<dyn RateLimiter + Send + Sync + 'static>;

#[async_trait]
pub trait RateLimiter {
    async fn check_rate_limit(
        &self,
        metric: Metric,
        identities: &MononokeIdentitySet,
    ) -> Result<Result<(), RateLimitReason>, Error>;

    fn check_load_shed(&self, identities: &MononokeIdentitySet) -> Result<(), RateLimitReason>;

    fn bump_load(&self, metric: Metric, load: LoadCost);

    fn category(&self) -> &str;

    fn commits_per_author_limit(&self) -> Option<RateLimitBody>;

    fn total_file_changes_limit(&self) -> Option<RateLimitBody>;
}

define_stats! {
    load_shed_counter: dynamic_singleton_counter("{}", (key: String)),
}

#[derive(Clone)]
pub struct RateLimitEnvironment {
    fb: FacebookInit,
    category: String,
    config: ConfigHandle<MononokeRateLimitConfig>,
}

impl RateLimitEnvironment {
    pub fn new(
        fb: FacebookInit,
        category: String,
        config: ConfigHandle<MononokeRateLimitConfig>,
    ) -> Self {
        Self {
            fb,
            category,
            config,
        }
    }

    pub fn get_rate_limiter(&self) -> BoxRateLimiter {
        let config = self.config.get();

        create_rate_limiter(self.fb, self.category.clone(), config)
    }
}

#[derive(Debug, Clone)]
pub struct RateLimitBody {
    pub raw_config: rate_limiting_config::RateLimitBody,
    pub window: Duration,
}

#[derive(Debug, Clone)]
pub struct MononokeRateLimitConfig {
    pub region_weight: f64,
    pub rate_limits: Vec<RateLimit>,
    pub load_shed_limits: Vec<LoadShedLimit>,
    #[allow(dead_code)]
    commits_per_author: RateLimitBody,
    #[allow(dead_code)]
    total_file_changes: Option<RateLimitBody>,
}

#[derive(Debug, Clone)]
pub struct RateLimit {
    pub body: RateLimitBody,
    #[allow(dead_code)]
    target: Option<Target>,
    #[allow(dead_code)]
    metric: Metric,
}

#[cfg(fbcode_build)]
impl RateLimit {
    fn applies_to_client(&self, identities: &MononokeIdentitySet) -> bool {
        match &self.target {
            // TODO (harveyhunt): Pass identities rather than Some(identities) once LFS server has
            // been updated to require certs.
            Some(t) => t.matches_client(Some(identities)),
            None => true,
        }
    }
}

impl LoadShedLimit {
    // TODO(harveyhunt): Make identities none optional once LFS server enforces that.
    pub fn should_load_shed(
        &self,
        fb: FacebookInit,
        identities: Option<&MononokeIdentitySet>,
    ) -> Result<(), RateLimitReason> {
        let applies_to_client = match &self.target {
            Some(t) => t.matches_client(identities),
            None => true,
        };

        if !applies_to_client {
            return Ok(());
        }

        let metric = self.raw_config.metric.to_string();

        match STATS::load_shed_counter.get_value(fb, (metric.clone(),)) {
            Some(value) if value > self.raw_config.limit => Err(RateLimitReason::LoadShedMetric(
                metric,
                value,
                self.raw_config.limit,
            )),
            _ => Ok(()),
        }
    }
}

#[derive(Debug, Clone)]
pub struct LoadShedLimit {
    pub raw_config: rate_limiting_config::LoadShedLimit,
    target: Option<Target>,
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Metric {
    EgressBytes,
    TotalManifests,
    GetpackFiles,
    Commits,
}

#[must_use]
#[derive(Debug, Error)]
pub enum RateLimitReason {
    #[error("Rate limited by {0:?} over {1:?}")]
    RateLimitedMetric(Metric, Duration),
    #[error("Load shed due to {0} (value: {1}, limit: {2})")]
    LoadShedMetric(String, i64, i64),
}

#[derive(Debug, Clone)]
pub enum Target {
    NotTarget(Box<Target>),
    AndTarget(Vec<Target>),
    OrTarget(Vec<Target>),
    Identity(MononokeIdentity),
    StaticSlice(StaticSlice),
}

#[derive(Debug, Copy, Clone)]
struct SlicePct(u8);

impl TryFrom<i32> for SlicePct {
    type Error = Error;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        if !(0..=100).contains(&value) {
            return Err(anyhow!("Invalid percentage"));
        }

        Ok(Self(value.try_into()?))
    }
}

#[derive(Debug, Clone)]
pub struct StaticSlice {
    slice_pct: SlicePct,
    // This is hashed with a client's hostname to allow us to change
    // which percentage of hosts are in a static slice.
    nonce: String,
}

impl Target {
    pub fn matches_client(&self, identities: Option<&MononokeIdentitySet>) -> bool {
        match self {
            Self::NotTarget(t) => !t.matches_client(identities),
            Self::AndTarget(ts) => ts.iter().all(|t| t.matches_client(identities)),
            Self::OrTarget(ts) => ts.iter().any(|t| t.matches_client(identities)),
            Self::Identity(i) => match identities {
                Some(client_idents) => client_idents.contains(i),
                None => false,
            },
            Self::StaticSlice(s) => in_throttled_slice(identities, s.slice_pct, &s.nonce),
        }
    }
}

fn in_throttled_slice(
    identities: Option<&MononokeIdentitySet>,
    slice_pct: SlicePct,
    nonce: &str,
) -> bool {
    let hostname = if let Some(hostname) = identities.map(|i| i.hostname()) {
        hostname
    } else {
        return false;
    };

    let mut hasher = DefaultHasher::new();
    hostname.hash(&mut hasher);
    nonce.hash(&mut hasher);

    hasher.finish() % 100 < slice_pct.0.into()
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_target_matches() {
        let test_ident = MononokeIdentity::new("USER", "foo");
        let test2_ident = MononokeIdentity::new("USER", "bar");
        let test3_ident = MononokeIdentity::new("USER", "baz");

        let ident_target = Target::Identity(test_ident.clone());
        let ident2_target = Target::Identity(test2_ident);
        let ident3_target = Target::Identity(test3_ident.clone());
        let empty_idents = Some(MononokeIdentitySet::new());

        assert!(!ident_target.matches_client(empty_idents.as_ref()));

        let mut idents = MononokeIdentitySet::new();
        idents.insert(test_ident);
        idents.insert(test3_ident);
        let idents = Some(idents);

        assert!(ident_target.matches_client(idents.as_ref()));

        let and_target = Target::AndTarget(vec![ident_target.clone(), ident3_target]);

        assert!(and_target.matches_client(idents.as_ref()));

        let or_target = Target::OrTarget(vec![ident_target, ident2_target.clone()]);

        assert!(or_target.matches_client(idents.as_ref()));

        let not_target = Target::NotTarget(Box::new(ident2_target));
        assert!(not_target.matches_client(idents.as_ref()));
    }

    #[test]
    fn test_target_in_static_slice() {
        let mut identities = MononokeIdentitySet::new();
        identities.insert(MononokeIdentity::new("MACHINE", "abc123.abc1.facebook.com"));

        assert!(!in_throttled_slice(None, 100.try_into().unwrap(), "abc"));

        assert!(!in_throttled_slice(
            Some(&identities),
            0.try_into().unwrap(),
            "abc"
        ));

        assert!(in_throttled_slice(
            Some(&identities),
            100.try_into().unwrap(),
            "abc"
        ));

        assert!(in_throttled_slice(
            Some(&identities),
            50.try_into().unwrap(),
            "123"
        ));

        // Check that changing the nonce results in a different slice.
        assert!(!in_throttled_slice(
            Some(&identities),
            50.try_into().unwrap(),
            "abc"
        ));
    }
}


/*
 * 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::sync::Arc;

use anyhow::Error;
use async_trait::async_trait;
use fbinit::FacebookInit;
use permission_checker::MononokeIdentitySet;

use crate::BoxRateLimiter;
use crate::LoadCost;
use crate::Metric;
use crate::MononokeRateLimitConfig;
use crate::RateLimitBody;
use crate::RateLimitReason;
use crate::RateLimiter;

pub fn get_region_capacity(_datacenter_capacity: &BTreeMap<String, i32>) -> Option<i32> {
    None
}

pub fn create_rate_limiter(
    _fb: FacebookInit,
    category: String,
    _config: Arc<MononokeRateLimitConfig>,
) -> BoxRateLimiter {
    Box::new(FakeLimiter { category })
}

#[derive(Debug)]
struct FakeLimiter {
    category: String,
}

#[async_trait]
impl RateLimiter for FakeLimiter {
    async fn check_rate_limit(
        &self,
        _metric: Metric,
        _identities: &MononokeIdentitySet,
    ) -> Result<Result<(), RateLimitReason>, Error> {
        Ok(Ok(()))
    }

    fn check_load_shed(&self, _identities: &MononokeIdentitySet) -> Result<(), RateLimitReason> {
        Ok(())
    }

    fn bump_load(&self, _metric: Metric, _load: LoadCost) {}

    fn category(&self) -> &str {
        &self.category
    }

    fn commits_per_author_limit(&self) -> Option<RateLimitBody> {
        None
    }

    fn total_file_changes_limit(&self) -> Option<RateLimitBody> {
        None
    }
}


/*
 * 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::hash::Hasher;
use std::sync::Arc;

use anyhow::Result;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use itertools::Itertools;
use metaconfig_types::RemoteMetadataDatabaseConfig;
use metaconfig_types::ShardableRemoteDatabaseConfig;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql::Connection;
use sql_construct::SqlShardableConstructFromMetadataDatabaseConfig;
use sql_construct::SqlShardedConstruct;
use sql_ext::mononoke_queries;
use sql_ext::SqlShardedConnections;
use twox_hash::XxHash32;
use vec1::Vec1;

pub const MYSQL_INSERT_CHUNK_SIZE: usize = 1000;

mononoke_queries! {
    write InsertBlobKeysForChangesets(values: (repo_id: RepositoryId, cs_id: ChangesetId, blob_key: String)) {
    insert_or_ignore,
    "{insert_or_ignore} INTO bonsai_blob_mapping (repo_id, cs_id, blob_key) VALUES {values}"
    }

    read GetBlobKeysForChangesets(repo_id: RepositoryId, >list cs_ids: ChangesetId) -> (ChangesetId, String) {
        "SELECT cs_id, blob_key FROM bonsai_blob_mapping WHERE repo_id = {repo_id} AND cs_id in {cs_ids}"
    }

    read GetChangesetsForBlobKeys(repo_id: RepositoryId, >list blob_keys: String) -> (ChangesetId, String) {
        "SELECT cs_id, blob_key FROM bonsai_blob_mapping WHERE repo_id = {repo_id} AND blob_key in {blob_keys}"
    }
}

#[facet::facet]
pub struct BonsaiBlobMapping {
    pub sql_bonsai_blob_mapping: SqlBonsaiBlobMapping,
}

impl BonsaiBlobMapping {
    pub async fn get_blob_keys_for_changesets(
        &self,
        repo_id: RepositoryId,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<Vec<(ChangesetId, String)>> {
        self.sql_bonsai_blob_mapping
            .get_blob_keys_for_changesets(repo_id, cs_ids)
            .await
    }

    pub async fn get_changesets_for_blob_keys(
        &self,
        repo_id: RepositoryId,
        blob_keys: Vec<String>,
    ) -> Result<Vec<(ChangesetId, String)>> {
        self.sql_bonsai_blob_mapping
            .get_changesets_for_blob_keys(repo_id, blob_keys)
            .await
    }

    pub async fn insert_blob_keys_for_changesets(
        &self,
        repo_id: RepositoryId,
        mappings: Vec<(ChangesetId, String)>,
    ) -> Result<u64> {
        self.sql_bonsai_blob_mapping
            .insert_blob_keys_for_changesets(repo_id, mappings)
            .await
    }
}

pub struct SqlBonsaiBlobMapping {
    shard_count: usize,
    write_connections: Arc<Vec1<Connection>>,
    read_connections: Arc<Vec1<Connection>>,
}

impl SqlBonsaiBlobMapping {
    pub fn new(
        write_connections: Arc<Vec1<Connection>>,
        read_connections: Arc<Vec1<Connection>>,
    ) -> Self {
        let shard_count = read_connections.len();
        Self {
            shard_count,
            write_connections,
            read_connections,
        }
    }

    pub async fn get_blob_keys_for_changesets(
        &self,
        repo_id: RepositoryId,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<Vec<(ChangesetId, String)>> {
        let mut res = vec![];
        for shard_id in 0..self.shard_count {
            let rows = GetBlobKeysForChangesets::query(
                &self.read_connections[shard_id],
                &repo_id,
                &cs_ids[..],
            )
            .await?;
            res.extend(rows);
        }
        Ok(res)
    }

    pub async fn get_changesets_for_blob_keys(
        &self,
        repo_id: RepositoryId,
        blob_keys: Vec<String>,
    ) -> Result<Vec<(ChangesetId, String)>> {
        let shard_to_blobs = blob_keys
            .into_iter()
            .map(|key| (self.shard(repo_id, &key), key))
            .into_group_map()
            .into_iter()
            .collect::<Vec<_>>();
        Ok(stream::iter(shard_to_blobs)
            .map(|(shard_id, blob_keys)| async move {
                GetChangesetsForBlobKeys::query(
                    &self.read_connections[shard_id],
                    &repo_id,
                    &blob_keys[..],
                )
                .await
            })
            .buffer_unordered(100)
            .try_collect::<Vec<_>>()
            .await?
            .into_iter()
            .flatten()
            .collect())
    }

    pub async fn insert_blob_keys_for_changesets(
        &self,
        repo_id: RepositoryId,
        mappings: Vec<(ChangesetId, String)>,
    ) -> Result<u64> {
        let shard_to_values = mappings
            .into_iter()
            .map(|(cs_id, blob_key)| (self.shard(repo_id, &blob_key), (repo_id, cs_id, blob_key)))
            .into_group_map();
        if shard_to_values.is_empty() {
            return Ok(0);
        }

        stream::iter(shard_to_values)
            .map(|(shard_id, values)| async move {
                let mut affected_rows = 0;
                for chunk in values.chunks(MYSQL_INSERT_CHUNK_SIZE) {
                    // This iter().map() is needed to convert &(_,_,_) to (&_, &_, &_)
                    let chunk: Vec<_> = chunk.iter().map(|(a, b, c)| (a, b, c)).collect();
                    let result = InsertBlobKeysForChangesets::query(
                        &self.write_connections[shard_id],
                        &chunk[..],
                    )
                    .await?;
                    affected_rows += result.affected_rows();
                }
                anyhow::Ok(affected_rows)
            })
            .buffer_unordered(100)
            .try_fold(
                0,
                |acc, affected_rows| async move { Ok(acc + affected_rows) },
            )
            .await
    }

    fn shard(&self, repo_id: RepositoryId, blob_key: &str) -> usize {
        let mut hasher = XxHash32::with_seed(0);
        hasher.write(&repo_id.id().to_ne_bytes());
        hasher.write(blob_key.as_bytes());
        (hasher.finish() % self.shard_count as u64) as usize
    }
}

impl SqlShardedConstruct for SqlBonsaiBlobMapping {
    const LABEL: &'static str = "bonsai_blob_mapping";

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

    fn from_sql_shard_connections(connections: SqlShardedConnections) -> Self {
        let SqlShardedConnections {
            read_connections,
            read_master_connections: _,
            write_connections,
        } = connections;

        let write_connections = Arc::new(write_connections);
        let read_connections = Arc::new(read_connections);
        let shard_count = read_connections.len();

        Self {
            shard_count,
            write_connections,
            read_connections,
        }
    }
}

impl SqlShardableConstructFromMetadataDatabaseConfig for SqlBonsaiBlobMapping {
    fn remote_database_config(
        remote: &RemoteMetadataDatabaseConfig,
    ) -> Option<&ShardableRemoteDatabaseConfig> {
        remote.bonsai_blob_mapping.as_ref()
    }
}

#[cfg(test)]
mod test {
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::THREES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;
    use sql_construct::SqlConstruct;

    use super::*;

    #[tokio::test]
    async fn test_single_write_and_read() -> Result<()> {
        let sql = SqlBonsaiBlobMapping::with_sqlite_in_memory()?;
        let repo_id = RepositoryId::new(1);
        let values = vec!["blob1", "blob2", "blob3"]
            .into_iter()
            .map(|blob| (ONES_CSID, blob.into()))
            .collect::<Vec<_>>();
        let res = sql
            .insert_blob_keys_for_changesets(repo_id, values.clone())
            .await?;

        assert_eq!(res, 3); // we are inserting 3 blobs each mapped to ONES_CSID

        let rows = sql
            .get_changesets_for_blob_keys(repo_id, vec!["blob1".into()])
            .await?;
        assert_eq!(rows, vec![values[0].clone()]);

        let rows = sql
            .get_blob_keys_for_changesets(repo_id, vec![ONES_CSID])
            .await?;

        assert_eq!(rows, values);
        Ok(())
    }

    #[tokio::test]
    async fn test_read_write_multiple_values() -> Result<()> {
        let sql = SqlBonsaiBlobMapping::with_sqlite_in_memory()?;
        let repo_id = RepositoryId::new(1);
        let blobs1 = vec!["blob1", "blob2", "blob3"]
            .into_iter()
            .map(|blob| (ONES_CSID, blob.into()))
            .collect::<Vec<_>>();
        let blobs2 = vec!["blob2", "blob4", "blob5"]
            .into_iter()
            .map(|blob| (TWOS_CSID, blob.into()))
            .collect::<Vec<_>>(); // overlap by blob2 with blobs1
        let blobs3 = vec!["blob3", "blob4", "blob6"]
            .into_iter()
            .map(|blob| (THREES_CSID, blob.into()))
            .collect::<Vec<_>>(); // overlap by blob3 and blob4 with blobs1 and blobs2
        let values = [blobs1.clone(), blobs2.clone(), blobs3.clone()]
            .concat()
            .into_iter()
            .collect::<Vec<_>>();

        let res = sql.insert_blob_keys_for_changesets(repo_id, values).await?;

        assert_eq!(res, 9); // for each of the 3 cs_ids we have 3 blobs to insert so total = 9

        let blob_keys = blobs2.into_iter().map(|(_, s)| s).collect::<Vec<_>>();
        let rows = sql.get_changesets_for_blob_keys(repo_id, blob_keys).await?;

        assert_eq!(
            rows,
            vec![
                (ONES_CSID, "blob2".to_string()),
                (TWOS_CSID, "blob2".to_string()),
                (TWOS_CSID, "blob4".to_string()),
                (TWOS_CSID, "blob5".to_string()),
                (THREES_CSID, "blob4".to_string()),
            ]
        );

        let rows = sql
            .get_blob_keys_for_changesets(repo_id, vec![TWOS_CSID, THREES_CSID])
            .await?;

        assert_eq!(
            rows,
            vec![
                (TWOS_CSID, "blob2".to_string()),
                (TWOS_CSID, "blob4".to_string()),
                (TWOS_CSID, "blob5".to_string()),
                (THREES_CSID, "blob3".to_string()),
                (THREES_CSID, "blob4".to_string()),
                (THREES_CSID, "blob6".to_string()),
            ]
        );

        Ok(())
    }

    #[test]
    fn test_sharding() -> Result<()> {
        let mut sql = SqlBonsaiBlobMapping::with_sqlite_in_memory()?;
        // manually specif multiple shards
        sql.shard_count = 3;
        let repo_id = RepositoryId::new(1);
        let blobs = vec![
            "blob1", "blob2", "blob3", "blob4", "blob5", "blob6", "blob7",
        ];
        let shards: Vec<_> = blobs.iter().map(|blob| sql.shard(repo_id, blob)).collect();
        assert_eq!(shards, vec![1, 2, 1, 2, 0, 2, 0]);

        // verify that blob will consistently hashed to the same shard
        let shard = sql.shard(repo_id, blobs[2]);
        assert_eq!(shard, 1);

        // verify that different repo_id with the same blob will land on a different shard
        let shard = sql.shard(RepositoryId::new(3), blobs[2]);
        assert_eq!(shard, 0);

        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::collections::HashSet;
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use in_memory_commit_graph_storage::InMemoryCommitGraphStorage;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::RepositoryId;
use vec1::Vec1;

#[cfg(test)]
mod tests;

pub struct BufferedCommitGraphStorage {
    in_memory_storage: InMemoryCommitGraphStorage,
    persistent_storage: Arc<dyn CommitGraphStorage>,
    /// The maximum number of changeset edges
    /// allowed to be stored in memory.
    max_in_memory_size: usize,
}

impl BufferedCommitGraphStorage {
    pub fn new(persistent_storage: Arc<dyn CommitGraphStorage>, max_in_memory_size: usize) -> Self {
        Self {
            in_memory_storage: InMemoryCommitGraphStorage::new(persistent_storage.repo_id()),
            persistent_storage,
            max_in_memory_size,
        }
    }

    /// Flushes all changeset edges from the in memory storage
    /// to the persistent storage. Returns the number of added
    /// edges to the persistent storage.
    pub async fn flush(&self, ctx: &CoreContext) -> Result<usize> {
        match Vec1::try_from_vec(self.in_memory_storage.drain()) {
            Ok(many_edges) => self.persistent_storage.add_many(ctx, many_edges).await,
            _ => Ok(0),
        }
    }
}

#[async_trait]
impl CommitGraphStorage for BufferedCommitGraphStorage {
    fn repo_id(&self) -> RepositoryId {
        self.persistent_storage.repo_id()
    }

    async fn add(&self, ctx: &CoreContext, edges: ChangesetEdges) -> Result<bool> {
        if self.in_memory_storage.len() + 1 > self.max_in_memory_size {
            self.flush(ctx).await?;
        }
        self.in_memory_storage.add(ctx, edges).await
    }

    async fn add_many(&self, ctx: &CoreContext, many_edges: Vec1<ChangesetEdges>) -> Result<usize> {
        if self.in_memory_storage.len() + many_edges.len() > self.max_in_memory_size {
            self.flush(ctx).await?;
        }
        self.in_memory_storage.add_many(ctx, many_edges).await
    }

    async fn fetch_edges(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetEdges>> {
        match self.in_memory_storage.fetch_edges(ctx, cs_id).await? {
            Some(edges) => Ok(Some(edges)),
            None => self.persistent_storage.fetch_edges(ctx, cs_id).await,
        }
    }

    async fn fetch_edges_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetEdges> {
        match self.in_memory_storage.fetch_edges(ctx, cs_id).await? {
            Some(edges) => Ok(edges),
            None => {
                self.persistent_storage
                    .fetch_edges_required(ctx, cs_id)
                    .await
            }
        }
    }

    async fn fetch_many_edges(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let mut fetched_edges = self
            .in_memory_storage
            .fetch_many_edges(ctx, cs_ids, prefetch)
            .await?;

        let unfetched_ids = cs_ids
            .iter()
            .filter(|cs_id| !fetched_edges.contains_key(cs_id))
            .copied()
            .collect::<Vec<_>>();

        if !unfetched_ids.is_empty() {
            fetched_edges.extend(
                self.persistent_storage
                    .fetch_many_edges(ctx, unfetched_ids.as_slice(), prefetch)
                    .await?
                    .into_iter(),
            )
        }

        Ok(fetched_edges)
    }

    async fn fetch_many_edges_required(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let mut fetched_edges = self
            .in_memory_storage
            .fetch_many_edges(ctx, cs_ids, prefetch)
            .await?;

        let unfetched_ids = cs_ids
            .iter()
            .filter(|cs_id| !fetched_edges.contains_key(cs_id))
            .copied()
            .collect::<Vec<_>>();

        if !unfetched_ids.is_empty() {
            fetched_edges.extend(
                self.persistent_storage
                    .fetch_many_edges_required(ctx, unfetched_ids.as_slice(), prefetch)
                    .await?
                    .into_iter(),
            )
        }

        Ok(fetched_edges)
    }

    async fn find_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: ChangesetIdPrefix,
        limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        match futures::try_join!(
            self.in_memory_storage.find_by_prefix(ctx, cs_prefix, limit),
            self.persistent_storage
                .find_by_prefix(ctx, cs_prefix, limit)
        )? {
            (in_memory_matches @ ChangesetIdsResolvedFromPrefix::TooMany(_), _) => {
                Ok(in_memory_matches)
            }
            (_, persistent_matches @ ChangesetIdsResolvedFromPrefix::TooMany(_)) => {
                Ok(persistent_matches)
            }
            (in_memory_matches, persistent_matches) => {
                Ok(ChangesetIdsResolvedFromPrefix::from_vec_and_limit(
                    in_memory_matches
                        .to_vec()
                        .into_iter()
                        .chain(persistent_matches.to_vec().into_iter())
                        .collect(),
                    limit,
                ))
            }
        }
    }

    async fn fetch_children(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        Ok(self
            .in_memory_storage
            .fetch_children(ctx, cs_id)
            .await?
            .into_iter()
            .chain(
                self.persistent_storage
                    .fetch_children(ctx, cs_id)
                    .await?
                    .into_iter(),
            )
            .collect::<HashSet<_>>()
            .into_iter()
            .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 std::future::Future;
use std::sync::Arc;

use anyhow::Result;
use commit_graph_testlib::*;
use commit_graph_types::storage::CommitGraphStorage;
use context::CoreContext;
use fbinit::FacebookInit;
use mononoke_types::RepositoryId;
use rendezvous::RendezVousOptions;
use sql_commit_graph_storage::SqlCommitGraphStorageBuilder;
use sql_construct::SqlConstruct;

use crate::BufferedCommitGraphStorage;

async fn run_test<Fut>(
    fb: FacebookInit,
    test_function: impl FnOnce(CoreContext, Arc<dyn CommitGraphStorage>) -> Fut,
) -> Result<()>
where
    Fut: Future<Output = Result<()>>,
{
    let ctx = CoreContext::test_mock(fb);
    let storage = Arc::new(BufferedCommitGraphStorage::new(
        Arc::new(
            SqlCommitGraphStorageBuilder::with_sqlite_in_memory()
                .unwrap()
                .build(RendezVousOptions::for_test(), RepositoryId::new(1)),
        ),
        5,
    ));
    test_function(ctx, storage).await
}

impl_commit_graph_tests!(run_test);


/*
 * 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::ops::Deref;
use std::sync::Arc;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use caching_ext::fill_cache;
use caching_ext::get_or_fill;
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 commit_graph_thrift as thrift;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use fbthrift::compact_protocol;
use maplit::hashset;
use memcache::KeyGen;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::RepositoryId;
use stats::prelude::*;
use tunables::tunables;
use vec1::Vec1;

#[cfg(test)]
mod tests;

define_stats! {
    prefix = "mononoke.cache.commit_graph.prefetch";

    hit: timeseries("hit"; Rate, Sum),
    fetched: timeseries("fetched"; Rate, Sum),
    prefetched: timeseries("prefetched"; Rate, Sum),
}

/// Size of chunk when fetching from the backing store
const CHUNK_SIZE: usize = 1000;

/// Number of chunks to fetch in parallel
const PARALLEL_CHUNKS: usize = 2;

/// Caching Commit Graph Storage
pub struct CachingCommitGraphStorage {
    storage: Arc<dyn CommitGraphStorage>,
    cachelib: CachelibHandler<CachedChangesetEdges>,
    memcache: MemcacheHandler,
    keygen: KeyGen,
    repo_id: RepositoryId,
}

struct CacheRequest<'a> {
    ctx: &'a CoreContext,
    caching_storage: &'a CachingCommitGraphStorage,
    prefetch: Prefetch,
    required: bool,
}

#[derive(Clone, Debug, Abomonation)]
pub struct CachedChangesetEdges {
    /// The cached edges.
    edges: ChangesetEdges,

    /// Whether these edges were originally fetched via a prefetch operation.
    prefetched: bool,
}

impl Deref for CachedChangesetEdges {
    type Target = ChangesetEdges;

    fn deref(&self) -> &ChangesetEdges {
        &self.edges
    }
}

impl CachedChangesetEdges {
    fn fetched(edges: ChangesetEdges) -> Self {
        CachedChangesetEdges {
            edges,
            prefetched: false,
        }
    }

    fn prefetched(edges: ChangesetEdges) -> Self {
        CachedChangesetEdges {
            edges,
            prefetched: true,
        }
    }

    fn take(self) -> ChangesetEdges {
        if self.prefetched {
            STATS::hit.add_value(1);
        }
        self.edges
    }

    fn to_thrift(&self) -> thrift::ChangesetEdges {
        self.edges.to_thrift()
    }

    fn from_thrift(edges: thrift::ChangesetEdges) -> Result<Self> {
        Ok(Self {
            edges: ChangesetEdges::from_thrift(edges)?,
            prefetched: false,
        })
    }
}

impl MemcacheEntity for CachedChangesetEdges {
    fn serialize(&self) -> Bytes {
        compact_protocol::serialize(&self.to_thrift())
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        compact_protocol::deserialize(bytes)
            .and_then(CachedChangesetEdges::from_thrift)
            .map_err(|_| McErrorKind::Deserialization)
    }
}

impl EntityStore<CachedChangesetEdges> for CacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<CachedChangesetEdges> {
        &self.caching_storage.cachelib
    }

    fn keygen(&self) -> &KeyGen {
        &self.caching_storage.keygen
    }

    fn memcache(&self) -> &MemcacheHandler {
        if self.prefetch.is_include() {
            if tunables()
                .disable_commit_graph_memcache_for_prefetch()
                .unwrap_or(false)
            {
                // If asked to prefetch, fetching from memcache is actually
                // slower, so don't perform memcache look-ups.
                &MemcacheHandler::Noop
            } else {
                &self.caching_storage.memcache
            }
        } else {
            &self.caching_storage.memcache
        }
    }

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

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

#[async_trait]
impl KeyedEntityStore<ChangesetId, CachedChangesetEdges> for CacheRequest<'_> {
    fn get_cache_key(&self, cs_id: &ChangesetId) -> String {
        self.caching_storage.cache_key(cs_id)
    }

    async fn get_from_db(
        &self,
        keys: HashSet<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, CachedChangesetEdges>> {
        let cs_ids: Vec<ChangesetId> = keys.iter().copied().collect();
        let entries = if self.required {
            self.caching_storage
                .storage
                .fetch_many_edges_required(self.ctx, &cs_ids, self.prefetch)
                .await?
        } else {
            self.caching_storage
                .storage
                .fetch_many_edges(self.ctx, &cs_ids, self.prefetch)
                .await?
        };
        if self.prefetch.is_include() {
            // We were asked to prefetch. We must separate out the prefetched
            // values from the fetched values as we may only return the
            // fetched values.
            let mut fetched = HashMap::new();
            let mut prefetched = HashMap::new();
            for (cs_id, edges) in entries {
                if keys.contains(&cs_id) {
                    fetched.insert(cs_id, CachedChangesetEdges::fetched(edges));
                } else {
                    prefetched.insert(cs_id, CachedChangesetEdges::prefetched(edges));
                }
            }
            if !prefetched.is_empty() {
                STATS::prefetched.add_value(prefetched.len() as i64);
                fill_cache(self, &prefetched).await;
            }
            STATS::fetched.add_value(fetched.len() as i64);
            Ok(fetched)
        } else {
            Ok(entries
                .into_iter()
                .map(|(cs_id, edges)| (cs_id, CachedChangesetEdges::fetched(edges)))
                .collect())
        }
    }
}

impl CachingCommitGraphStorage {
    fn keygen() -> KeyGen {
        let key_prefix = "scm.mononoke.commitgraph";

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

    fn cache_key(&self, cs_id: &ChangesetId) -> String {
        format!("{}.{}", self.repo_id.prefix(), cs_id)
    }

    pub fn new(
        storage: Arc<dyn CommitGraphStorage>,
        cache_handler_factory: CacheHandlerFactory,
    ) -> Self {
        Self {
            repo_id: storage.repo_id(),
            storage,
            cachelib: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
            keygen: Self::keygen(),
        }
    }

    #[cfg(test)]
    pub fn mocked(storage: Arc<dyn CommitGraphStorage>) -> Self {
        Self::new(storage, CacheHandlerFactory::Mocked)
    }

    fn request<'a>(&'a self, ctx: &'a CoreContext, prefetch: Prefetch) -> CacheRequest<'a> {
        let prefetch = if tunables().disable_commit_graph_prefetch().unwrap_or(false) {
            Prefetch::None
        } else {
            prefetch.include_hint()
        };
        CacheRequest {
            ctx,
            caching_storage: self,
            prefetch,
            required: false,
        }
    }

    fn request_required<'a>(
        &'a self,
        ctx: &'a CoreContext,
        prefetch: Prefetch,
    ) -> CacheRequest<'a> {
        let prefetch = if tunables().disable_commit_graph_prefetch().unwrap_or(false) {
            Prefetch::None
        } else {
            prefetch.include_hint()
        };
        CacheRequest {
            ctx,
            caching_storage: self,
            prefetch,
            required: true,
        }
    }
}

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

    async fn add(&self, ctx: &CoreContext, edges: ChangesetEdges) -> Result<bool> {
        self.storage.add(ctx, edges).await
    }

    async fn add_many(&self, ctx: &CoreContext, many_edges: Vec1<ChangesetEdges>) -> Result<usize> {
        self.storage.add_many(ctx, many_edges).await
    }

    async fn fetch_edges(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetEdges>> {
        let mut found = get_or_fill(&self.request(ctx, Prefetch::None), hashset![cs_id]).await?;
        Ok(found.remove(&cs_id).map(CachedChangesetEdges::take))
    }

    async fn fetch_edges_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetEdges> {
        let mut found =
            get_or_fill(&self.request_required(ctx, Prefetch::None), hashset![cs_id]).await?;
        Ok(found
            .remove(&cs_id)
            .ok_or_else(|| anyhow!("Missing changeset from commit graph storage: {}", cs_id))?
            .take())
    }

    async fn fetch_many_edges(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let cs_ids: HashSet<ChangesetId> = cs_ids.iter().copied().collect();
        let found = get_or_fill_chunked(
            &self.request(ctx, prefetch),
            cs_ids.clone(),
            CHUNK_SIZE,
            PARALLEL_CHUNKS,
        )
        .await?;
        Ok(found
            .into_iter()
            .map(|(cs_id, edges)| (cs_id, edges.take()))
            .collect())
    }

    async fn fetch_many_edges_required(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let cs_ids: HashSet<ChangesetId> = cs_ids.iter().copied().collect();
        let found = get_or_fill_chunked(
            &self.request_required(ctx, prefetch),
            cs_ids.clone(),
            CHUNK_SIZE,
            PARALLEL_CHUNKS,
        )
        .await?;
        Ok(found
            .into_iter()
            .map(|(cs_id, edges)| (cs_id, edges.take()))
            .collect())
    }

    async fn find_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: ChangesetIdPrefix,
        limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        self.storage.find_by_prefix(ctx, cs_prefix, limit).await
    }

    async fn fetch_children(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        self.storage.fetch_children(ctx, cs_id).await
    }
}


/*
 * 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::future::Future;
use std::sync::Arc;

use anyhow::Result;
use commit_graph_testlib::*;
use commit_graph_types::storage::CommitGraphStorage;
use context::CoreContext;
use fbinit::FacebookInit;
use mononoke_types::RepositoryId;
use rendezvous::RendezVousOptions;
use sql_commit_graph_storage::SqlCommitGraphStorageBuilder;
use sql_construct::SqlConstruct;

use crate::CachingCommitGraphStorage;

async fn run_test<Fut>(
    fb: FacebookInit,
    test_function: impl FnOnce(CoreContext, Arc<dyn CommitGraphStorage>) -> Fut,
) -> Result<()>
where
    Fut: Future<Output = Result<()>>,
{
    let ctx = CoreContext::test_mock(fb);
    let storage = Arc::new(CachingCommitGraphStorage::mocked(Arc::new(
        SqlCommitGraphStorageBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(RendezVousOptions::for_test(), RepositoryId::new(1)),
    )));
    test_function(ctx, storage.clone()).await?;
    assert!(storage.cachelib.mock_store().unwrap().stats().hits > 0);
    Ok(())
}

impl_commit_graph_tests!(run_test);


/*
 * 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::pin::Pin;

use anyhow::Result;
use borrowed::borrowed;
use commit_graph_types::frontier::ChangesetFrontier;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use futures::future;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::Future;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::ArcCommitGraph;
use crate::CommitGraph;

/// Builder for a reverse topologically ordered stream of changesets that
/// are ancestors of any set of changesets (heads). This builder allows customizing
/// the stream by:
///
/// - excluding ancestors of a set of changesets (common).
///
/// - excluding changesets that satisify a given property (if this property holds
/// for one changeset then it has to hold for all its parents).
///
/// - including only changesets that satisfy a given property (if this property doesn't
/// hold for one changeset then it mustn't hold for any of its parents).
///
/// - including only changesets that are descendants of any one changeset.
pub struct AncestorsStreamBuilder {
    commit_graph: ArcCommitGraph,
    ctx: CoreContext,
    heads: Vec<ChangesetId>,
    common: Vec<ChangesetId>,
    descendants_of: Option<ChangesetId>,
    property: Box<
        dyn Fn(ChangesetId) -> Pin<Box<dyn Future<Output = Result<bool>> + Send>> + Send + Sync,
    >,
}

impl AncestorsStreamBuilder {
    pub fn new(commit_graph: ArcCommitGraph, ctx: CoreContext, heads: Vec<ChangesetId>) -> Self {
        Self {
            commit_graph,
            ctx,
            heads,
            common: vec![],
            descendants_of: None,
            property: Box::new(|_| Box::pin(future::ready(Ok(true)))),
        }
    }

    pub fn exclude_ancestors_of(mut self, common: Vec<ChangesetId>) -> Self {
        self.common.extend(common);
        self
    }

    pub fn descendants_of(mut self, descendants_of: ChangesetId) -> Self {
        self.descendants_of = Some(descendants_of);
        self
    }

    pub fn with<Property, Out>(mut self, other_property: Property) -> Self
    where
        Property: Fn(ChangesetId) -> Out + Send + Sync + 'static,
        Out: Future<Output = Result<bool>> + Send + 'static,
    {
        self.property = Box::new(move |cs_id| {
            let fut_property = (self.property)(cs_id);
            let fut_other_property = other_property(cs_id);

            Box::pin(async move {
                if !fut_property.await? {
                    Ok(false)
                } else {
                    fut_other_property.await
                }
            })
        });
        self
    }

    pub fn without<Property, Out>(mut self, other_property: Property) -> Self
    where
        Property: Fn(ChangesetId) -> Out + Send + Sync + 'static,
        Out: Future<Output = Result<bool>> + Send + 'static,
    {
        self.property = Box::new(move |cs_id| {
            let fut_property = (self.property)(cs_id);
            let fut_other_property = other_property(cs_id);

            Box::pin(async move {
                if !fut_property.await? {
                    Ok(false)
                } else {
                    Ok(!fut_other_property.await?)
                }
            })
        });
        self
    }

    pub async fn build(self) -> Result<BoxStream<'static, Result<ChangesetId>>> {
        struct AncestorsStreamState {
            commit_graph: ArcCommitGraph,
            ctx: CoreContext,
            heads: ChangesetFrontier,
            common: ChangesetFrontier,
            descendants_of: Option<(ChangesetId, Generation)>,
            property: Box<
                dyn Fn(ChangesetId) -> Pin<Box<dyn Future<Output = Result<bool>> + Send>>
                    + Send
                    + Sync,
            >,
        }

        let heads = match self.descendants_of {
            Some(descendants_of) => {
                stream::iter(self.heads)
                    .map(anyhow::Ok)
                    .try_filter_map(|head| {
                        borrowed!(self.commit_graph: &CommitGraph, self.ctx);
                        async move {
                            match commit_graph.is_ancestor(ctx, descendants_of, head).await? {
                                true => Ok(Some(head)),
                                false => Ok(None),
                            }
                        }
                    })
                    .try_collect()
                    .await?
            }
            None => self.heads,
        };

        let descendants_of = match self.descendants_of {
            Some(descendants_of) => Some((
                descendants_of,
                self.commit_graph
                    .changeset_generation_required(&self.ctx, descendants_of)
                    .await?,
            )),
            None => None,
        };

        let (heads, common) = futures::try_join!(
            self.commit_graph.frontier(&self.ctx, heads),
            self.commit_graph.frontier(&self.ctx, self.common)
        )?;

        Ok(stream::try_unfold(
            Box::new(AncestorsStreamState {
                commit_graph: self.commit_graph,
                ctx: self.ctx,
                heads,
                common,
                descendants_of,
                property: self.property,
            }),
            move |mut state| async move {
                let AncestorsStreamState {
                    commit_graph,
                    ctx,
                    heads,
                    common,
                    descendants_of,
                    property,
                } = &mut *state;

                if let Some((generation, cs_ids)) = heads.pop_last() {
                    commit_graph.lower_frontier(ctx, common, generation).await?;

                    let mut cs_ids_not_excluded = vec![];
                    for cs_id in cs_ids {
                        if !common.highest_generation_contains(cs_id, generation)
                            && property(cs_id).await?
                        {
                            cs_ids_not_excluded.push(cs_id)
                        }
                    }

                    let all_edges = commit_graph
                        .storage
                        .fetch_many_edges_required(
                            ctx,
                            &cs_ids_not_excluded,
                            Prefetch::for_p1_linear_traversal(),
                        )
                        .await?;

                    for (_cs_id, edges) in all_edges.into_iter() {
                        for parent in edges.parents.iter() {
                            if let Some((descendants_of, descendants_of_gen)) = descendants_of {
                                // There is no need to query ancestry if the skip tree parent's generation number
                                // is greater than or equal to the generation number of descendants_of. This is
                                // because the skip tree parent is the common ancestor of all parents, and since
                                // the current changeset is a descendant of descendants_of, all of its parents
                                // will also be descendants of it.
                                if !edges.skip_tree_parent.map_or(false, |skip_tree_parent| {
                                    skip_tree_parent.generation >= *descendants_of_gen
                                }) && !commit_graph
                                    .is_ancestor(ctx, *descendants_of, parent.cs_id)
                                    .await?
                                {
                                    continue;
                                }
                            }
                            heads
                                .entry(parent.generation)
                                .or_default()
                                .insert(parent.cs_id);
                        }
                    }

                    anyhow::Ok(Some((stream::iter(cs_ids_not_excluded).map(Ok), state)))
                } else {
                    Ok(None)
                }
            },
        )
        .try_flatten()
        .boxed())
    }
}


/*
 * 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 anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use buffered_commit_graph_storage::BufferedCommitGraphStorage;
use changeset_fetcher::ArcChangesetFetcher;
use changeset_fetcher::ChangesetFetcher;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::edges::ChangesetParents;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use smallvec::SmallVec;
use smallvec::ToSmallVec;
use vec1::Vec1;

use crate::CommitGraph;

impl CommitGraph {
    /// Same as add but fetches parent edges using the changeset fetcher
    /// if not found in the storage, and recursively tries to add them.
    ///
    /// Changesets should be sorted in topological order.
    pub async fn add_recursive(
        &self,
        ctx: &CoreContext,
        changeset_fetcher: ArcChangesetFetcher,
        changesets: Vec1<(ChangesetId, ChangesetParents)>,
    ) -> Result<usize> {
        let mut edges_map: HashMap<ChangesetId, ChangesetEdges> = Default::default();
        let changesets_set: HashSet<ChangesetId> =
            changesets.iter().map(|(cs_id, _)| cs_id).cloned().collect();
        let mut search_stack: Vec<(ChangesetId, ChangesetParents)> = changesets.into();
        let mut to_add_stack: Vec<(ChangesetId, ChangesetParents)> = Default::default();

        while let Some((cs_id, parents)) = search_stack.pop() {
            // If edges map already has the key there's no need to process it (this may happen if
            // initial vector had duplicates or if we descent into the same parrents via two
            // different paths)
            if edges_map.contains_key(&cs_id) {
                continue;
            }

            to_add_stack.push((cs_id, parents.clone()));

            // We don't need to look up:
            //  * changesets we already have in edges_map
            //  * changesets that are part of changesets set (as they'll be inserted anyway)
            let parents_to_fetch: SmallVec<[ChangesetId; 1]> = parents
                .into_iter()
                .filter(|cs_id| !edges_map.contains_key(cs_id) && !changesets_set.contains(cs_id))
                .collect();

            if !parents_to_fetch.is_empty() {
                edges_map.extend(
                    self.storage
                        .fetch_many_edges(ctx, &parents_to_fetch, Prefetch::None)
                        .await
                        .with_context(|| "during commit_graph::add_recursive (fetch_many_edges)")?
                        .into_iter(),
                );
            }

            for parent in parents_to_fetch {
                if !edges_map.contains_key(&parent) {
                    // If the parents are not present in the commit graph we have to backfilll them
                    // so let's add them to the stack so they can be processed in the next
                    // iteration.
                    search_stack.push((
                        parent,
                        changeset_fetcher
                            .get_parents(ctx, parent)
                            .await
                            .with_context(|| "during commit_graph::add_recursive (get_parents)")?
                            .to_smallvec(),
                    ));
                }
            }
        }

        // We use buffered storage here to be able to do all the writes in parallel.
        // We need to create a new CommitGraph wrapper to work with the buffered storage.
        let buffered_storage =
            Arc::new(BufferedCommitGraphStorage::new(self.storage.clone(), 10000));
        let graph = CommitGraph::new(buffered_storage.clone());
        while let Some((cs_id, parents)) = to_add_stack.pop() {
            let edges = graph.build_edges(ctx, cs_id, parents, &edges_map).await?;
            edges_map.insert(cs_id, edges.clone());
            buffered_storage.add(ctx, edges).await?;
        }
        buffered_storage.flush(ctx).await
    }
}

#[async_trait]
impl ChangesetFetcher for CommitGraph {
    async fn get_generation_number(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Generation> {
        self.changeset_generation_required(ctx, cs_id).await
    }

    async fn get_parents(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<Vec<ChangesetId>> {
        self.changeset_parents_required(ctx, cs_id)
            .await
            .map(SmallVec::into_vec)
    }
}


/*
 * 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 anyhow::anyhow;
use anyhow::Result;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::edges::ChangesetNode;
use commit_graph_types::edges::ChangesetNodeParents;
use commit_graph_types::edges::ChangesetParents;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::CommitGraph;

impl CommitGraph {
    pub(crate) async fn build_edges(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        parents: ChangesetParents,
        edges_map: &HashMap<ChangesetId, ChangesetEdges>,
    ) -> Result<ChangesetEdges> {
        let mut max_parent_gen = 0;
        let mut edge_parents = ChangesetNodeParents::new();
        let mut merge_ancestor = None;

        let mut skip_tree_parent = None;
        let mut first_parent = true;

        let mut p1_linear_depth = 0;

        for parent in &parents {
            let parent_edge = edges_map
                .get(parent)
                .ok_or_else(|| anyhow!("Missing parent: {}", parent))?;
            max_parent_gen = max_parent_gen.max(parent_edge.node.generation.value());
            edge_parents.push(parent_edge.node);
            if parents.len() == 1 {
                merge_ancestor = Some(parent_edge.merge_ancestor.unwrap_or(parent_edge.node));
            }

            // skip_tree_parent is the skip tree lowest common ancestor of all parents
            if first_parent {
                first_parent = false;
                skip_tree_parent = Some(parent_edge.node);

                p1_linear_depth = parent_edge.node.p1_linear_depth + 1;
            } else if let Some(previous_parent) = skip_tree_parent {
                skip_tree_parent = self
                    .lowest_common_ancestor(
                        ctx,
                        previous_parent.cs_id,
                        parent_edge.node.cs_id,
                        |edges| edges.skip_tree_parent,
                        |edges| edges.skip_tree_skew_ancestor,
                        |node| node.skip_tree_depth,
                    )
                    .await?;
            }
        }

        let generation = Generation::new(max_parent_gen + 1);
        let skip_tree_depth = match skip_tree_parent {
            Some(node) => node.skip_tree_depth + 1,
            None => 0,
        };
        let node = ChangesetNode {
            cs_id,
            generation,
            skip_tree_depth,
            p1_linear_depth,
        };

        let p1_parent = edge_parents.first().copied();

        Ok(ChangesetEdges {
            node,
            parents: edge_parents,
            merge_ancestor,
            skip_tree_parent,
            skip_tree_skew_ancestor: self
                .calc_skew_ancestor(
                    ctx,
                    skip_tree_parent,
                    |edges| edges.skip_tree_skew_ancestor,
                    |node| node.skip_tree_depth,
                )
                .await?,
            p1_linear_skew_ancestor: self
                .calc_skew_ancestor(
                    ctx,
                    p1_parent,
                    |edges| edges.p1_linear_skew_ancestor,
                    |node| node.p1_linear_depth,
                )
                .await?,
        })
    }

    /// Calculates the skew binary ancestor of a changeset
    /// given its parent and two closures, one returns the
    /// skew ancestor of a ChangesetEdges and the other
    /// returns the depth of a ChangesetNode.
    pub(crate) async fn calc_skew_ancestor<F, G>(
        &self,
        ctx: &CoreContext,
        parent: Option<ChangesetNode>,
        get_skew_ancestor: F,
        get_depth: G,
    ) -> Result<Option<ChangesetNode>>
    where
        F: Fn(&ChangesetEdges) -> Option<ChangesetNode>,
        G: Fn(ChangesetNode) -> u64,
    {
        // The skew binary ancestor is either the parent of the
        // changeset or the skew binary ancestor of the skew binary
        // ancestor of the parent if it exists, and if the difference
        // in depth between the parent and the first ancestor is the
        // same as the difference between the two ancestors.

        let parent = match parent {
            Some(node) => node,
            None => return Ok(None),
        };

        let parent_edges = self.storage.fetch_edges_required(ctx, parent.cs_id).await?;

        let parent_skew_ancestor = match get_skew_ancestor(&parent_edges) {
            Some(node) => node,
            None => return Ok(Some(parent)),
        };

        let parent_skew_ancestor_edges = self
            .storage
            .fetch_edges_required(ctx, parent_skew_ancestor.cs_id)
            .await?;

        let parent_second_skew_ancestor = match get_skew_ancestor(&parent_skew_ancestor_edges) {
            Some(node) => node,
            None => return Ok(Some(parent)),
        };

        if get_depth(parent) - get_depth(parent_skew_ancestor)
            == get_depth(parent_skew_ancestor) - get_depth(parent_second_skew_ancestor)
        {
            Ok(Some(parent_second_skew_ancestor))
        } else {
            Ok(Some(parent))
        }
    }

    /// Returns the ancestor of a changeset that has depth target_depth,
    /// or None if the changeset's depth is smaller than target_depth.
    async fn level_ancestor<F, G, H>(
        &self,
        ctx: &CoreContext,
        mut cs_id: ChangesetId,
        target_depth: u64,
        get_parent: F,
        get_skew_ancestor: G,
        get_depth: H,
    ) -> Result<Option<ChangesetNode>>
    where
        F: Fn(&ChangesetEdges) -> Option<ChangesetNode>,
        G: Fn(&ChangesetEdges) -> Option<ChangesetNode>,
        H: Fn(ChangesetNode) -> u64,
    {
        loop {
            let node_edges = self.storage.fetch_edges_required(ctx, cs_id).await?;

            if get_depth(node_edges.node) == target_depth {
                return Ok(Some(node_edges.node));
            }

            if get_depth(node_edges.node) < target_depth {
                return Ok(None);
            }

            match (get_skew_ancestor(&node_edges), get_parent(&node_edges)) {
                (Some(skew_ancestor), _) if get_depth(skew_ancestor) >= target_depth => {
                    cs_id = skew_ancestor.cs_id
                }
                (_, Some(parent)) => cs_id = parent.cs_id,
                _ => {
                    return Err(anyhow!(
                        "Changeset has positive depth yet has no parent: {}",
                        cs_id
                    ));
                }
            }
        }
    }

    /// Returns the ancestor of a changeset that has depth target_depth in the skip tree,
    /// or None if the changeset's depth is smaller than target_depth.
    pub async fn skip_tree_level_ancestor(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        target_depth: u64,
    ) -> Result<Option<ChangesetNode>> {
        self.level_ancestor(
            ctx,
            cs_id,
            target_depth,
            |edges| edges.skip_tree_parent,
            |edges| edges.skip_tree_skew_ancestor,
            |node| node.skip_tree_depth,
        )
        .await
    }

    /// Returns the ancestor of a changeset that has depth target_depth in the p1 linear tree,
    /// or None if the changeset's depth is smaller than target_depth.
    pub async fn p1_linear_level_ancestor(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        target_depth: u64,
    ) -> Result<Option<ChangesetNode>> {
        self.level_ancestor(
            ctx,
            cs_id,
            target_depth,
            |edges| edges.parents.first().copied(),
            |edges| edges.p1_linear_skew_ancestor,
            |node| node.p1_linear_depth,
        )
        .await
    }

    /// Returns the lowest common ancestor of two changesets.
    async fn lowest_common_ancestor<F, G, H>(
        &self,
        ctx: &CoreContext,
        cs_id1: ChangesetId,
        cs_id2: ChangesetId,
        get_parent: F,
        get_skew_ancestor: G,
        get_depth: H,
    ) -> Result<Option<ChangesetNode>>
    where
        F: Fn(&ChangesetEdges) -> Option<ChangesetNode> + Copy,
        G: Fn(&ChangesetEdges) -> Option<ChangesetNode> + Copy,
        H: Fn(ChangesetNode) -> u64 + Copy,
    {
        let (edges1, edges2) = futures::try_join!(
            self.storage.fetch_edges_required(ctx, cs_id1),
            self.storage.fetch_edges_required(ctx, cs_id2),
        )?;

        let (mut u, mut v) = (edges1.node, edges2.node);

        if get_depth(u) < get_depth(v) {
            std::mem::swap(&mut u, &mut v);
        }

        // Get ancestor of u that has the same depth
        // as v and change u to it
        u = self
            .level_ancestor(
                ctx,
                u.cs_id,
                get_depth(v),
                get_parent,
                get_skew_ancestor,
                get_depth,
            )
            .await?
            .ok_or_else(|| {
                anyhow!(
                    "Failed to get ancestor of changeset {} that has depth {}",
                    u.cs_id,
                    get_depth(v)
                )
            })?;

        // Now that u and v have the same depth, we check if u
        // and v have different skew binary ancestors, if that
        // is the case we move to those ancestors, otherwise we
        // move to their parents. This way we guarantee ending
        // up in the lowest common ancestor.
        while u.cs_id != v.cs_id {
            let (u_edges, v_edges) = futures::try_join!(
                self.storage.fetch_edges_required(ctx, u.cs_id),
                self.storage.fetch_edges_required(ctx, v.cs_id),
            )?;

            match (
                get_skew_ancestor(&u_edges),
                get_skew_ancestor(&v_edges),
                get_parent(&u_edges),
                get_parent(&v_edges),
            ) {
                (Some(u_skew_ancestor), Some(v_skew_ancestor), _, _)
                    if u_skew_ancestor.cs_id != v_skew_ancestor.cs_id =>
                {
                    u = u_skew_ancestor;
                    v = v_skew_ancestor;
                }
                (_, _, Some(u_parent), Some(v_parent)) => {
                    u = u_parent;
                    v = v_parent;
                }
                _ => return Ok(None),
            }
        }

        Ok(Some(u))
    }

    /// Returns the lowest common ancestor of two changesets in the skip tree.
    pub async fn skip_tree_lowest_common_ancestor(
        &self,
        ctx: &CoreContext,
        cs_id1: ChangesetId,
        cs_id2: ChangesetId,
    ) -> Result<Option<ChangesetNode>> {
        self.lowest_common_ancestor(
            ctx,
            cs_id1,
            cs_id2,
            |edges| edges.skip_tree_parent,
            |edges| edges.skip_tree_skew_ancestor,
            |node| node.skip_tree_depth,
        )
        .await
    }

    /// Returns the lowest common ancestor of two changesets in the p1 linear tree.
    pub async fn p1_linear_lowest_common_ancestor(
        &self,
        ctx: &CoreContext,
        cs_id1: ChangesetId,
        cs_id2: ChangesetId,
    ) -> Result<Option<ChangesetNode>> {
        self.lowest_common_ancestor(
            ctx,
            cs_id1,
            cs_id2,
            |edges| edges.skip_tree_parent,
            |edges| edges.skip_tree_skew_ancestor,
            |node| node.skip_tree_depth,
        )
        .await
    }
}


/*
 * 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::anyhow;
use anyhow::Result;
use borrowed::borrowed;
use commit_graph_types::edges::ChangesetNode;
use commit_graph_types::frontier::ChangesetFrontier;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use futures::future;
use futures::Future;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::CommitGraph;

impl CommitGraph {
    /// Obtain a frontier of changesets from a single changeset id, which must
    /// exist.
    pub(crate) async fn single_frontier(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetFrontier> {
        let generation = self.changeset_generation_required(ctx, cs_id).await?;
        Ok(ChangesetFrontier::new_single(cs_id, generation))
    }

    /// Obtain a frontier of changesets from a list of changeset ids, which
    /// must all exist.
    pub(crate) async fn frontier(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<ChangesetFrontier> {
        let all_edges = self
            .storage
            .fetch_many_edges_required(ctx, &cs_ids, Prefetch::None)
            .await?;

        cs_ids
            .into_iter()
            .map(|cs_id| {
                Ok((
                    cs_id,
                    all_edges
                        .get(&cs_id)
                        .ok_or_else(|| anyhow!("Missing changeset in commit graph: {}", cs_id))?
                        .node
                        .generation,
                ))
            })
            .collect::<Result<_>>()
    }

    /// Pops the highest generation changesets of a frontier, returning any that
    /// satisify a property and lowering the rest of them to either their immediate
    /// parents or their lowest skip tree edge that doesn't satisfy the property.
    /// Repeatedly calling this function and concatenating the output will result
    /// in the frontier of changesets satisifying the property.
    ///
    /// Returns None if the frontier is empty.
    pub(crate) async fn lower_frontier_step<Property, Out>(
        &self,
        ctx: &CoreContext,
        frontier: &mut ChangesetFrontier,
        property: Property,
        prefetch: Prefetch,
    ) -> Result<Option<Vec<ChangesetId>>>
    where
        Property: Fn(ChangesetNode) -> Out + Send + Sync,
        Out: Future<Output = Result<bool>>,
    {
        match frontier.pop_last() {
            None => Ok(None),
            Some((_, cs_ids)) => {
                let cs_ids = cs_ids.into_iter().collect::<Vec<_>>();
                let frontier_edges = self
                    .storage
                    .fetch_many_edges_required(ctx, &cs_ids, prefetch)
                    .await?;

                let mut property_frontier: Vec<_> = Default::default();

                for (_, edges) in frontier_edges {
                    if property(edges.node).await? {
                        property_frontier.push(edges.node.cs_id);
                    } else {
                        match edges
                            .lowest_skip_tree_edge_with(|node| {
                                borrowed!(property);
                                async move { Ok(!property(node).await?) }
                            })
                            .await?
                        {
                            Some(ancestor) => {
                                frontier
                                    .entry(ancestor.generation)
                                    .or_default()
                                    .insert(ancestor.cs_id);
                            }
                            None => {
                                for parent in edges.parents {
                                    frontier
                                        .entry(parent.generation)
                                        .or_default()
                                        .insert(parent.cs_id);
                                }
                            }
                        }
                    }
                }

                Ok(Some(property_frontier))
            }
        }
    }

    /// Lower a frontier so that it contains the highest ancestors of the
    /// frontier that have a generation number less than or equal to
    /// `generation`.
    pub(crate) async fn lower_frontier(
        &self,
        ctx: &CoreContext,
        frontier: &mut ChangesetFrontier,
        target_generation: Generation,
    ) -> Result<()> {
        loop {
            match frontier.last_key_value() {
                None => return Ok(()),
                Some((generation, _)) if *generation <= target_generation => {
                    return Ok(());
                }
                _ => {}
            }

            self.lower_frontier_step(
                ctx,
                frontier,
                move |node| future::ready(Ok(node.generation < target_generation)),
                Prefetch::for_skip_tree_traversal(target_generation),
            )
            .await?;
        }
    }

    /// Lower the highest generation changesets of a frontier
    /// to their immediate parents.
    pub(crate) async fn lower_frontier_highest_generation(
        &self,
        ctx: &CoreContext,
        frontier: &mut ChangesetFrontier,
    ) -> Result<()> {
        if let Some((_, cs_ids)) = frontier.pop_last() {
            let cs_ids = cs_ids.into_iter().collect::<Vec<_>>();
            let frontier_edges = self
                .storage
                .fetch_many_edges_required(ctx, &cs_ids, Prefetch::for_p1_linear_traversal())
                .await?;

            for cs_id in cs_ids {
                let edges = frontier_edges
                    .get(&cs_id)
                    .ok_or_else(|| anyhow!("Missing changeset in commit graph: {}", cs_id))?;

                for parent in edges.parents.iter() {
                    frontier
                        .entry(parent.generation)
                        .or_default()
                        .insert(parent.cs_id);
                }
            }
        }
        Ok(())
    }
}