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

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

use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;
use slog::trace;
use stats::prelude::*;
use tunables::tunables;

use crate::dag::errors;
use crate::dag::errors::programming;
use crate::dag::errors::BackendError;
use crate::dag::errors::DagError;
use crate::dag::id::Group;
use crate::dag::id::Id;
use crate::dag::idmap::IdMapWrite;
use crate::dag::ops::IdConvert;
use crate::dag::ops::PrefixLookup;
use crate::dag::Result;
use crate::dag::VerLink;
use crate::dag::VertexName;
use crate::idmap::ConcurrentMemIdMap;
use crate::idmap::IdMapVersion;
use crate::DagId;
use crate::IdMap;

define_stats! {
    prefix = "mononoke.segmented_changelog.idmap.memory";
    insert_many: timeseries(Sum),
    flush_writes: timeseries(Sum),
}

const DEFAULT_LOG_SAMPLING_RATE: usize = 20000;

/// Type conversion - `VertexName` is used in the `dag` crate to abstract over different ID types
/// such as Mercurial IDs, Bonsai, a theoretical git ID and more.
/// but should always name a Bonsai `ChangesetId` on the server.
///
/// This converts a `VertexName` to a `ChangesetId`
///
/// # Panics
///
/// This conversion will panic if the `VertexName` is not a Bonsai `ChangesetId`
pub fn cs_id_from_vertex_name(name: &VertexName) -> ChangesetId {
    ChangesetId::from_bytes(name).expect("VertexName is not a valid ChangesetId")
}

/// Type conversion - `VertexName` is used in the `dag` crate to abstract over different ID types
/// such as Mercurial IDs, Bonsai, a theoretical git ID and more.
/// but should always name a Bonsai `ChangesetId` on the server.
///
/// This converts a `ChangesetId` to a `VertexName`
pub fn vertex_name_from_cs_id(cs_id: &ChangesetId) -> VertexName {
    VertexName::copy_from(cs_id.blake2().as_ref())
}

#[derive(Clone)]
pub struct IdMapMemWrites {
    /// The actual IdMap
    inner: Arc<dyn IdMap>,
    /// Stores recent writes that haven't yet been persisted to the backing store
    mem: Arc<ConcurrentMemIdMap>,
}

impl IdMapMemWrites {
    pub fn new(inner: Arc<dyn IdMap>) -> Self {
        Self {
            inner,
            mem: Arc::new(ConcurrentMemIdMap::new()),
        }
    }

    pub async fn flush_writes(&self, ctx: &CoreContext) -> anyhow::Result<()> {
        trace!(
            ctx.logger(),
            "flushing {} in-memory IdMap entries to SQL",
            self.mem.len(),
        );
        let writes = self.mem.drain();
        STATS::flush_writes.add_value(
            writes
                .len()
                .try_into()
                .expect("More than an i64 of writes before a flush!"),
        );
        self.inner.insert_many(ctx, writes).await
    }
}

#[async_trait]
impl IdMap for IdMapMemWrites {
    async fn insert_many(
        &self,
        ctx: &CoreContext,
        mappings: Vec<(DagId, ChangesetId)>,
    ) -> anyhow::Result<()> {
        let mappings_size = mappings.len();
        STATS::insert_many.add_value(
            mappings_size
                .try_into()
                .expect("More than an i64 of writes in one go!"),
        );
        let res = self.mem.insert_many(ctx, mappings).await;
        if res.is_ok() {
            let new_size = self.mem.len();
            let old_size = new_size - mappings_size;
            let sampling_rate = tunables()
                .segmented_changelog_idmap_log_sampling_rate()
                .unwrap_or_default();
            let sampling_rate = if sampling_rate <= 0 {
                DEFAULT_LOG_SAMPLING_RATE
            } else {
                sampling_rate as usize
            };
            if new_size / sampling_rate != old_size / sampling_rate {
                trace!(
                    ctx.logger(),
                    "{} entries inserted into in-memory IdMap, new size: {}",
                    mappings_size,
                    new_size,
                );
            }
        }
        res
    }

    async fn find_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> anyhow::Result<HashMap<DagId, ChangesetId>> {
        let mut result = self
            .mem
            .find_many_changeset_ids(ctx, dag_ids.clone())
            .await?;
        let missing: Vec<_> = dag_ids
            .iter()
            .filter(|v| !result.contains_key(v))
            .copied()
            .collect();
        if !missing.is_empty() {
            let inner_result = self.inner.find_many_changeset_ids(ctx, missing).await?;
            result.extend(inner_result);
        }
        Ok(result)
    }

    async fn find_many_dag_ids(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> anyhow::Result<HashMap<ChangesetId, DagId>> {
        let mut result = self.mem.find_many_dag_ids(ctx, cs_ids.clone()).await?;
        let missing: Vec<_> = cs_ids
            .iter()
            .filter(|id| !result.contains_key(id))
            .copied()
            .collect();
        if !missing.is_empty() {
            let inner_result = self.inner.find_many_dag_ids(ctx, missing).await?;
            result.extend(inner_result);
        }
        Ok(result)
    }

    /// Finds the dag ID for given changeset - if possible to do so quickly.
    /// Might return no answers for changesets that have dag ids assigned.
    ///
    /// Should be used by callers that can deal with missing information.
    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> anyhow::Result<HashMap<ChangesetId, DagId>> {
        let mut result = self
            .mem
            .find_many_dag_ids_maybe_stale(ctx, cs_ids.clone())
            .await?;
        let missing: Vec<_> = cs_ids
            .iter()
            .filter(|id| !result.contains_key(id))
            .copied()
            .collect();
        if !missing.is_empty() {
            let inner_result = self
                .inner
                .find_many_dag_ids_maybe_stale(ctx, missing)
                .await?;
            result.extend(inner_result);
        }
        Ok(result)
    }

    async fn get_last_entry(
        &self,
        ctx: &CoreContext,
    ) -> anyhow::Result<Option<(DagId, ChangesetId)>> {
        let mem_last = self.mem.get_last_entry(ctx).await?;
        match mem_last {
            Some(_) => return Ok(mem_last),
            None => self.inner.get_last_entry(ctx).await,
        }
    }

    fn idmap_version(&self) -> Option<IdMapVersion> {
        self.inner.idmap_version()
    }
}

/// The server needs metadata that isn't normally available in the `dag` crate
/// for normal operation.
///
/// This wrapper provides that for any `IdMap`, so that the `dag` crate
/// traits can be used on a server `IdMap`
///
/// # Performance
///
/// This function creates a new `VerLink` of `IdMap` every time for convenience.
///  `dag::NameSet` will conservatively assumes `IdMap` are totally different,
/// and conservatively avoid fast paths when it sees those different `verlink`s.
/// Practically, that means you need to avoid capturing `NameSet` produced
/// outside the `closure` to avoid performance issues. Having all `NameSet`
/// calculations inside the `closure` is fine, even if the final result is
/// passed out to the enclosing scope
#[derive(Clone)]
pub struct IdMapWrapper {
    verlink: VerLink,
    inner: IdMapMemWrites,
    ctx: CoreContext,
}

impl IdMapWrapper {
    /// Create a new wrapper around the server IdMap, using CoreContext
    /// for calling update functions
    pub fn new(ctx: CoreContext, idmap: Arc<dyn IdMap>) -> Self {
        let idmap_memwrites = IdMapMemWrites::new(idmap);
        Self {
            verlink: VerLink::new(),
            inner: idmap_memwrites,
            ctx,
        }
    }

    /// If not called, IdMap changes are discarded when this is dropped
    pub async fn flush_writes(&self) -> anyhow::Result<()> {
        self.inner.flush_writes(&self.ctx).await
    }

    /// Flushes writes and then returns the original IdMap
    pub async fn finish(self) -> anyhow::Result<Arc<dyn IdMap>> {
        self.flush_writes().await?;
        Ok(self.inner.inner)
    }

    /// Access to the inner IdMap
    pub fn as_inner(&self) -> &IdMapMemWrites {
        &self.inner
    }

    /// Get a clone of the original IdMap fed in.
    /// If `flush_writes` has not been called, this will not be updated
    pub fn clone_idmap(&self) -> Arc<dyn IdMap> {
        self.inner.inner.clone()
    }
}

#[async_trait]
impl PrefixLookup for IdMapWrapper {
    async fn vertexes_by_hex_prefix(
        &self,
        _hex_prefix: &[u8],
        _limit: usize,
    ) -> Result<Vec<VertexName>> {
        errors::programming("Server-side IdMap does not support prefix lookup")
    }
}
#[async_trait]
impl IdConvert for IdMapWrapper {
    async fn vertex_id(&self, name: VertexName) -> Result<Id> {
        // NOTE: The server implementation puts all Ids in the "master" group.
        self.vertex_id_with_max_group(&name, Group::MASTER)
            .await?
            .ok_or(DagError::VertexNotFound(name))
    }
    async fn vertex_id_with_max_group(
        &self,
        name: &VertexName,
        _max_group: Group,
    ) -> Result<Option<Id>> {
        // NOTE: The server implementation puts all Ids in the "master" group.
        let cs_id = cs_id_from_vertex_name(name);
        Ok(self
            .inner
            .find_dag_id(&self.ctx, cs_id)
            .await
            .map_err(BackendError::from)?)
    }

    async fn vertex_name(&self, id: Id) -> Result<VertexName> {
        self.inner
            .find_changeset_id(&self.ctx, id)
            .await
            .map_err(BackendError::from)?
            .map(|id| vertex_name_from_cs_id(&id))
            .ok_or(DagError::IdNotFound(id))
    }
    async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool> {
        self.vertex_id_with_max_group(name, Group::MASTER)
            .await
            .map(|d| d.is_some())
    }

    async fn contains_vertex_id_locally(&self, id: &[Id]) -> Result<Vec<bool>> {
        let ids = Vec::from(id);
        let found = self
            .inner
            .find_many_changeset_ids(&self.ctx, ids)
            .await
            .map_err(BackendError::from)?;

        Ok(id.iter().map(|id| found.contains_key(id)).collect())
    }

    async fn contains_vertex_name_locally(&self, name: &[VertexName]) -> Result<Vec<bool>> {
        let cs_ids: Vec<_> = name.iter().map(cs_id_from_vertex_name).collect();
        let found = self
            .inner
            .find_many_dag_ids(&self.ctx, cs_ids)
            .await
            .map_err(BackendError::from)?;

        Ok(name
            .iter()
            .map(|name| found.contains_key(&cs_id_from_vertex_name(name)))
            .collect())
    }

    /// Convert [`Id`]s to [`VertexName`]s in batch.
    async fn vertex_name_batch(&self, id: &[Id]) -> Result<Vec<Result<VertexName>>> {
        let ids = Vec::from(id);
        let found = self
            .inner
            .find_many_changeset_ids(&self.ctx, ids)
            .await
            .map_err(BackendError::from)?;

        Ok(id
            .iter()
            .map(|id| {
                found
                    .get(id)
                    .map(vertex_name_from_cs_id)
                    .ok_or(DagError::IdNotFound(*id))
            })
            .collect())
    }

    /// Convert [`VertexName`]s to [`Id`]s in batch.
    async fn vertex_id_batch(&self, names: &[VertexName]) -> Result<Vec<Result<Id>>> {
        let cs_ids: Vec<_> = names.iter().map(cs_id_from_vertex_name).collect();

        let found = self
            .inner
            .find_many_dag_ids(&self.ctx, cs_ids)
            .await
            .map_err(BackendError::from)?;

        Ok(names
            .iter()
            .map(|name| {
                found
                    .get(&cs_id_from_vertex_name(name))
                    .copied()
                    .ok_or_else(|| DagError::VertexNotFound(name.clone()))
            })
            .collect())
    }

    fn map_id(&self) -> &str {
        "Mononoke segmented changelog"
    }

    fn map_version(&self) -> &VerLink {
        &self.verlink
    }
}

#[async_trait]
impl IdMapWrite for IdMapWrapper {
    async fn insert(&mut self, id: Id, name: &[u8]) -> Result<()> {
        // NB: This is only suitable for tailing right now, as it writes on every call
        // Eventually, this needs to use a batching interface
        let cs_id = ChangesetId::from_bytes(name).map_err(BackendError::from)?;
        Ok(self
            .inner
            .insert(&self.ctx, id, cs_id)
            .await
            .map_err(BackendError::from)?)
    }
    async fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
        let _ = (low, high);
        programming("remove_range() is not implemented server-side")
    }
}


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

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

use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql::Connection;
use sql_ext::mononoke_queries;
use sql_ext::replication::ReplicaLagMonitor;
use sql_ext::replication::WaitForReplicationConfig;
use sql_ext::SqlConnections;
use stats::prelude::*;

use crate::idmap::IdMap;
use crate::types::IdMapVersion;
use crate::DagId;

define_stats! {
    prefix = "mononoke.segmented_changelog.idmap";
    insert: timeseries(Sum),
    find_changeset_id: timeseries(Sum),
    find_dag_id: timeseries(Sum),
    get_last_entry: timeseries(Sum),
    find_by_changeset_id_prefix: timeseries(Sum),
}

const INSERT_MAX: usize = 2_000;

pub struct SqlIdMap {
    connections: SqlConnections,
    replica_lag_monitor: Arc<dyn ReplicaLagMonitor>,
    repo_id: RepositoryId,
    version: IdMapVersion,
}

mononoke_queries! {
    write InsertIdMapEntry(
        values: (repo_id: RepositoryId, version: IdMapVersion, dag_id: u64, cs_id: ChangesetId)
    ) {
        insert_or_ignore,
        "
        {insert_or_ignore} INTO segmented_changelog_idmap (repo_id, version, vertex, cs_id)
        VALUES {values}
        "
    }

    write CopyIdMap(
        values: (repo_id: RepositoryId, new_version: IdMapVersion, source_version: IdMapVersion, copy_limit: u64)
    ) {
        insert_or_ignore,
        "
        {insert_or_ignore} INTO segmented_changelog_idmap_copy_mappings (repo_id, idmap_version, copied_version, copy_limit)
        VALUES {values}
        "
    }

    read GetIdMapCopySource(
        repo_id: RepositoryId,
        version: IdMapVersion,
    ) -> (IdMapVersion, u64) {
        "
        SELECT copied_version, copy_limit
        FROM segmented_changelog_idmap_copy_mappings
        WHERE repo_id = {repo_id} AND idmap_version = {version}
        "
    }

    read CheckCopyMapping(
        repo_id: RepositoryId,
        version: IdMapVersion,
        dag_id: u64,
    ) -> (IdMapVersion) {
        "
        SELECT idmap_version
        FROM segmented_changelog_idmap_copy_mappings
        WHERE repo_id = {repo_id} AND copied_version = {version} AND copy_limit >= {dag_id}
        LIMIT 1
        "
    }

    read SelectManyChangesetIds(
        repo_id: RepositoryId,
        version: IdMapVersion,
        >list dag_ids: u64
    ) -> (u64, ChangesetId) {
        "
        SELECT idmap.vertex as vertex, idmap.cs_id as cs_id
        FROM segmented_changelog_idmap AS idmap
            LEFT JOIN segmented_changelog_idmap_copy_mappings AS copy_mappings
            ON (idmap.repo_id = copy_mappings.repo_id AND idmap.version = copy_mappings.copied_version)
        WHERE
            idmap.repo_id = {repo_id} AND
            (idmap.version = {version} OR copy_mappings.idmap_version = {version}) AND
            (copy_mappings.copy_limit IS NULL OR copy_mappings.copy_limit >= idmap.vertex) AND
            idmap.vertex IN {dag_ids}
        "
    }

    read SelectManyDagIds(
        repo_id: RepositoryId,
        version: IdMapVersion,
        >list cs_ids: ChangesetId
    ) -> (ChangesetId, u64) {
        "
        SELECT idmap.cs_id as cs_id, idmap.vertex as vertex
        FROM segmented_changelog_idmap AS idmap
            LEFT JOIN segmented_changelog_idmap_copy_mappings AS copy_mappings
            ON (idmap.repo_id = copy_mappings.repo_id AND idmap.version = copy_mappings.copied_version)
        WHERE
            idmap.repo_id = {repo_id} AND
            (idmap.version = {version} OR copy_mappings.idmap_version = {version}) AND
            (copy_mappings.copy_limit IS NULL OR copy_mappings.copy_limit >= idmap.vertex) AND
            idmap.cs_id in {cs_ids}
        "
    }

    read SelectLastEntry(repo_id: RepositoryId, version: IdMapVersion) -> (u64, ChangesetId) {
        "
        SELECT idmap.vertex as vertex, idmap.cs_id as cs_id
        FROM segmented_changelog_idmap AS idmap
        WHERE
            idmap.repo_id = {repo_id} AND
            idmap.version = {version} AND
            idmap.vertex = (
                SELECT MAX(inx.vertex)
                FROM segmented_changelog_idmap AS inx
                WHERE
                    inx.repo_id = {repo_id} AND
                    inx.version = {version}
              )
        "
    }

    read SelectHighestCopyLimit(repo_id: RepositoryId, version: IdMapVersion) -> (Option<u64>,) {
        "
        SELECT MAX(copy_mappings.copy_limit)
        FROM segmented_changelog_idmap_copy_mappings AS copy_mappings
        WHERE
            copy_mappings.repo_id = {repo_id} AND
            copy_mappings.idmap_version = {version}
        "

    }

    read SelectHexPrefix(repo_id: RepositoryId, version: IdMapVersion, prefix: &[u8], limit: usize) -> (ChangesetId, u64) {
        "
        SELECT idmap.cs_id as cs_id, idmap.vertex as vertex
        FROM segmented_changelog_idmap AS idmap
            LEFT JOIN segmented_changelog_idmap_copy_mappings AS copy_mappings
            ON (idmap.repo_id = copy_mappings.repo_id AND idmap.version = copy_mappings.copied_version)
        WHERE
            idmap.repo_id = {repo_id} AND
            (idmap.version = {version} OR copy_mappings.idmap_version = {version}) AND
            (copy_mappings.copy_limit IS NULL OR copy_mappings.copy_limit >= idmap.vertex) AND
            HEX(cs_id) LIKE CONCAT({prefix}, '%')
        LIMIT {limit}
        "
    }
}

impl SqlIdMap {
    pub fn new(
        connections: SqlConnections,
        replica_lag_monitor: Arc<dyn ReplicaLagMonitor>,
        repo_id: RepositoryId,
        version: IdMapVersion,
    ) -> Self {
        Self {
            connections,
            replica_lag_monitor,
            repo_id,
            version,
        }
    }

    async fn select_many_changesetids(
        &self,
        conn: &Connection,
        dag_ids: &[u64],
    ) -> Result<HashMap<DagId, ChangesetId>, Error> {
        if dag_ids.is_empty() {
            return Ok(HashMap::new());
        }
        let rows =
            SelectManyChangesetIds::query(conn, &self.repo_id, &self.version, dag_ids).await?;
        Ok::<_, Error>(rows.into_iter().map(|row| (DagId(row.0), row.1)).collect())
    }

    async fn select_many_dag_ids(
        &self,
        conn: &Connection,
        cs_ids: &[ChangesetId],
    ) -> Result<HashMap<ChangesetId, DagId>, Error> {
        if cs_ids.is_empty() {
            return Ok(HashMap::new());
        }
        let rows = SelectManyDagIds::query(conn, &self.repo_id, &self.version, cs_ids).await?;
        Ok(rows.into_iter().map(|row| (row.0, DagId(row.1))).collect())
    }

    pub async fn copy(&self, dag_limit: DagId, new_version: IdMapVersion) -> Result<Self, Error> {
        // Check that this version *has* a DagId at `dag_limit`
        if self
            .select_many_changesetids(&self.connections.write_connection, &[dag_limit.0])
            .await?
            .is_empty()
        {
            return Err(format_err!(
                "repo {} {:?} does not have DagId {}",
                self.repo_id,
                self.version,
                dag_limit.0
            ));
        }
        let new_self = Self::new(
            self.connections.clone(),
            self.replica_lag_monitor.clone(),
            self.repo_id,
            new_version,
        );
        let transaction = self
            .connections
            .write_connection
            .start_transaction()
            .await?;

        // Check copy sources for this version, and create new copy sources for new_version
        let (mut transaction, copy_sources) =
            GetIdMapCopySource::query_with_transaction(transaction, &self.repo_id, &self.version)
                .await?;
        for (version, limit) in copy_sources {
            let limit = dag_limit.0.min(limit);
            let (t, _) = CopyIdMap::query_with_transaction(
                transaction,
                &[(&self.repo_id, &new_version, &version, &limit)],
            )
            .await?;
            transaction = t;
            if limit == dag_limit.0 {
                // Copying existing copy mappings has covered everything. Stop copying here.
                transaction.commit().await?;
                return Ok(new_self);
            }
        }

        // Then, if not everything is covered, add this version as a copy source for new_version
        let (transaction, _) = CopyIdMap::query_with_transaction(
            transaction,
            &[(&self.repo_id, &new_version, &self.version, &dag_limit.0)],
        )
        .await?;
        transaction.commit().await?;
        Ok(new_self)
    }

    async fn check_copy_mapping(&self, conn: &Connection, lowest_id: &DagId) -> Result<()> {
        let copy_source =
            CheckCopyMapping::query(conn, &self.repo_id, &self.version, &lowest_id.0).await?;
        if let Some(copy_source) = copy_source.first() {
            return Err(format_err!(
                "repo {}: dag_ig {} in version {:?} is a copy baseline for version {:?}",
                self.repo_id,
                lowest_id,
                self.version,
                copy_source
            ));
        }
        Ok(())
    }
}

#[async_trait]
impl IdMap for SqlIdMap {
    async fn insert_many(
        &self,
        ctx: &CoreContext,
        mut mappings: Vec<(DagId, ChangesetId)>,
    ) -> Result<()> {
        STATS::insert.add_value(mappings.len() as i64);
        mappings.sort();

        // Sanity check for copy mappings - we want to reject inserts to vertex IDs that are
        // referenced by a copy mapping.
        // This is not perfect (races, replication lag), but should show bugs in tests
        if let Some(min) = mappings.iter().map(|(id, _cs)| id).min() {
            self.check_copy_mapping(&self.connections.read_connection, min)
                .await?;
        }

        // With validation passed, we split the mappings into batches that we write in separate
        // transactions.
        for (i, chunk) in mappings.chunks(INSERT_MAX).enumerate() {
            if i > 0 {
                self.replica_lag_monitor
                    .wait_for_replication(&|| {
                        WaitForReplicationConfig::default().with_logger(ctx.logger())
                    })
                    .await?;
            }
            let mut to_insert = Vec::with_capacity(chunk.len());
            for (dag_id, cs_id) in chunk {
                to_insert.push((&self.repo_id, &self.version, &dag_id.0, cs_id));
            }
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlWrites);
            let mut transaction = self
                .connections
                .write_connection
                .start_transaction()
                .await?;
            let query_result =
                InsertIdMapEntry::query_with_transaction(transaction, &to_insert).await;
            match query_result {
                Err(err) => {
                    // transaction is "lost" to the query
                    return Err(err.context(format_err!(
                        "repo {}: failed inserting IdMap entries",
                        self.repo_id
                    )));
                }
                Ok((t, insert_result)) => {
                    transaction = t;
                    if insert_result.affected_rows() != chunk.len() as u64 {
                        transaction.rollback().await?;
                        return Err(format_err!(
                            "repo {}: failed insert race, total entries {}, batch {}",
                            self.repo_id,
                            mappings.len(),
                            i
                        ));
                    } else {
                        transaction.commit().await?;
                    }
                }
            }
        }
        Ok(())
    }

    async fn find_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> Result<HashMap<DagId, ChangesetId>> {
        STATS::find_changeset_id.add_value(dag_ids.len() as i64);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let to_query: Vec<_> = dag_ids.iter().map(|v| v.0).collect();
        let mut cs_ids = self
            .select_many_changesetids(&self.connections.read_connection, &to_query)
            .await?;
        let not_found_in_replica: Vec<_> = dag_ids
            .iter()
            .filter(|x| !cs_ids.contains_key(x))
            .map(|v| v.0)
            .collect();
        if !not_found_in_replica.is_empty() {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let from_master = self
                .select_many_changesetids(
                    &self.connections.read_master_connection,
                    &not_found_in_replica,
                )
                .await?;
            for (k, v) in from_master {
                cs_ids.insert(k, v);
            }
        }
        Ok(cs_ids)
    }

    async fn find_many_dag_ids(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        STATS::find_dag_id.add_value(cs_ids.len() as i64);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let mut dag_ids = self
            .select_many_dag_ids(&self.connections.read_connection, &cs_ids)
            .await?;
        let not_found_in_replica: Vec<_> = cs_ids
            .iter()
            .filter(|x| !dag_ids.contains_key(x))
            .cloned()
            .collect();
        if !not_found_in_replica.is_empty() {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let from_master = self
                .select_many_dag_ids(
                    &self.connections.read_master_connection,
                    &not_found_in_replica,
                )
                .await?;
            for (k, v) in from_master {
                dag_ids.insert(k, v);
            }
        }
        Ok(dag_ids)
    }

    /// The maybe stale version of the method doesn't reach out beyond
    /// replica servers even if information is missing so it might ommit
    /// newest entries.
    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        STATS::find_dag_id.add_value(cs_ids.len() as i64);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        self.select_many_dag_ids(&self.connections.read_connection, &cs_ids)
            .await
    }

    async fn get_last_entry(&self, ctx: &CoreContext) -> Result<Option<(DagId, ChangesetId)>> {
        STATS::get_last_entry.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);
        // From the update algorithm perspective, it makes most sense to read from master. Because
        // trying to insert a value that was already inserted will fail the whole processing an
        // outdated entry will definitely lead to wasted work.
        let rows = SelectLastEntry::query(
            &self.connections.write_connection,
            &self.repo_id,
            &self.version,
        )
        .await?;
        if rows.is_empty() {
            // This IdMap version doesn't yet have its own entries. Check to see if it's a copy
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            let highest_copy_limit = SelectHighestCopyLimit::query(
                &self.connections.write_connection,
                &self.repo_id,
                &self.version,
            )
            .await?
            .into_iter()
            .next()
            .map(|(r,)| r);

            if let Some(Some(highest_copy_limit)) = highest_copy_limit {
                // It is a copy, so take the last entry from the copy source instead
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::SqlReadsMaster);
                let mut res = self
                    .select_many_changesetids(
                        &self.connections.write_connection,
                        &[highest_copy_limit],
                    )
                    .await?;
                Ok(res.remove_entry(&DagId(highest_copy_limit)))
            } else {
                Ok(None)
            }
        } else {
            Ok(rows.into_iter().next().map(|r| (DagId(r.0), r.1)))
        }
    }

    fn idmap_version(&self) -> Option<IdMapVersion> {
        Some(self.version)
    }
}

#[cfg(test)]
mod tests {
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_types_mocks::changesetid::AS_CSID;
    use mononoke_types_mocks::changesetid::BS_CSID;
    use mononoke_types_mocks::changesetid::FIVES_CSID;
    use mononoke_types_mocks::changesetid::FOURS_CSID;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::THREES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;
    use sql::rusqlite::Connection as SqliteConnection;
    use sql::Connection;
    use sql_construct::SqlConstruct;
    use sql_ext::replication::NoReplicaLagMonitor;

    use super::*;
    use crate::builder::SegmentedChangelogSqlConnections;

    fn new_sql_idmap() -> Result<SqlIdMap> {
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
        Ok(SqlIdMap::new(
            conns.0,
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(0),
            IdMapVersion(0),
        ))
    }

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

        assert_eq!(idmap.get_last_entry(&ctx).await?, None);

        idmap.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap.insert(&ctx, DagId(1), ONES_CSID).await?;
        idmap.insert(&ctx, DagId(2), TWOS_CSID).await?;
        idmap.insert(&ctx, DagId(3), THREES_CSID).await?;

        assert_eq!(
            idmap.get_last_entry(&ctx).await?,
            Some((DagId(3), THREES_CSID))
        );

        Ok(())
    }

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

        assert_eq!(idmap.get_last_entry(&ctx).await?, None);

        idmap.insert_many(&ctx, vec![]).await?;
        assert!(idmap.get_changeset_id(&ctx, DagId(1)).await.is_err());

        idmap.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap
            .insert_many(
                &ctx,
                vec![
                    (DagId(1), ONES_CSID),
                    (DagId(2), TWOS_CSID),
                    (DagId(3), THREES_CSID),
                ],
            )
            .await?;

        assert_eq!(idmap.get_changeset_id(&ctx, DagId(1)).await?, ONES_CSID);
        assert_eq!(idmap.get_changeset_id(&ctx, DagId(3)).await?, THREES_CSID);

        assert!(
            idmap
                .insert_many(
                    &ctx,
                    vec![
                        (DagId(1), ONES_CSID),
                        (DagId(2), TWOS_CSID),
                        (DagId(3), THREES_CSID),
                    ],
                )
                .await
                .is_err()
        );

        idmap
            .insert_many(&ctx, vec![(DagId(4), FOURS_CSID)])
            .await?;
        assert_eq!(idmap.get_changeset_id(&ctx, DagId(4)).await?, FOURS_CSID);

        assert!(
            idmap
                .insert_many(&ctx, vec![(DagId(1), FIVES_CSID)])
                .await
                .is_err()
        );

        Ok(())
    }

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

        let response = idmap
            .find_many_changeset_ids(&ctx, vec![DagId(1), DagId(2), DagId(3), DagId(6)])
            .await?;
        assert!(response.is_empty());

        idmap.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap
            .insert_many(
                &ctx,
                vec![
                    (DagId(1), ONES_CSID),
                    (DagId(2), TWOS_CSID),
                    (DagId(3), THREES_CSID),
                    (DagId(4), FOURS_CSID),
                    (DagId(5), FIVES_CSID),
                ],
            )
            .await?;

        let response = idmap
            .find_many_changeset_ids(&ctx, vec![DagId(1), DagId(2), DagId(3), DagId(6)])
            .await?;
        assert_eq!(
            response,
            hashmap![DagId(1) => ONES_CSID, DagId(2) => TWOS_CSID, DagId(3) => THREES_CSID]
        );

        let response = idmap
            .find_many_changeset_ids(&ctx, vec![DagId(4), DagId(5)])
            .await?;
        assert_eq!(
            response,
            hashmap![DagId(4) => FOURS_CSID, DagId(5) => FIVES_CSID]
        );

        let response = idmap.find_many_changeset_ids(&ctx, vec![DagId(6)]).await?;
        assert!(response.is_empty());

        Ok(())
    }

    #[fbinit::test]
    async fn test_find_many_changeset_ids_leader_fallback(fb: FacebookInit) -> Result<()> {
        fn conn() -> Result<Connection> {
            let con = SqliteConnection::open_in_memory()?;
            con.execute_batch(SegmentedChangelogSqlConnections::CREATION_QUERY)?;
            Ok(Connection::with_sqlite(con))
        }

        let ctx = CoreContext::test_mock(fb);

        let leader = conn()?;
        let replica = conn()?;

        let conns = SqlConnections {
            write_connection: leader.clone(),
            read_connection: replica,
            read_master_connection: leader,
        };

        let idmap = SqlIdMap::new(
            conns,
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(0),
            IdMapVersion(0),
        );

        idmap.insert(&ctx, DagId(0), ONES_CSID).await?;

        let res = idmap.get_changeset_id(&ctx, DagId(0)).await?;
        assert_eq!(res, ONES_CSID);

        let res = idmap.find_many_changeset_ids(&ctx, vec![DagId(0)]).await?;
        assert_eq!(res, hashmap![DagId(0) => ONES_CSID]);

        Ok(())
    }

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

        let response = idmap
            .find_many_dag_ids(&ctx, vec![ONES_CSID, TWOS_CSID, THREES_CSID, FOURS_CSID])
            .await?;
        assert!(response.is_empty());

        idmap.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap
            .insert_many(
                &ctx,
                vec![
                    (DagId(1), ONES_CSID),
                    (DagId(2), TWOS_CSID),
                    (DagId(3), THREES_CSID),
                    (DagId(4), FOURS_CSID),
                ],
            )
            .await?;

        let response = idmap
            .find_many_dag_ids(&ctx, vec![ONES_CSID, TWOS_CSID, THREES_CSID])
            .await?;
        assert_eq!(
            response,
            hashmap![ONES_CSID => DagId(1), TWOS_CSID => DagId(2), THREES_CSID => DagId(3)]
        );

        let response = idmap
            .find_many_dag_ids(&ctx, vec![FOURS_CSID, FIVES_CSID])
            .await?;
        assert_eq!(response, hashmap![FOURS_CSID => DagId(4)]);

        let response = idmap.find_many_dag_ids(&ctx, vec![FIVES_CSID]).await?;
        assert!(response.is_empty());

        Ok(())
    }

    #[fbinit::test]
    async fn test_find_many_dag_maybe_stale_no_leader_fallback(fb: FacebookInit) -> Result<()> {
        fn conn() -> Result<Connection> {
            let con = SqliteConnection::open_in_memory()?;
            con.execute_batch(SegmentedChangelogSqlConnections::CREATION_QUERY)?;
            Ok(Connection::with_sqlite(con))
        }

        let ctx = CoreContext::test_mock(fb);

        let leader = conn()?;
        let replica = conn()?;

        // We're setting those conns separately so we can craft replica content
        // and make it different from leader.
        let write_to_replica_conns = SqlConnections {
            write_connection: replica.clone(),
            read_connection: replica.clone(),
            read_master_connection: replica.clone(),
        };

        let write_to_replica_idmap = SqlIdMap::new(
            write_to_replica_conns,
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(0),
            IdMapVersion(0),
        );

        write_to_replica_idmap
            .insert(&ctx, DagId(0), AS_CSID)
            .await?;
        write_to_replica_idmap
            .insert_many(&ctx, vec![(DagId(1), ONES_CSID), (DagId(2), TWOS_CSID)])
            .await?;

        let conns = SqlConnections {
            write_connection: leader.clone(),
            read_connection: replica,
            read_master_connection: leader,
        };

        let idmap = SqlIdMap::new(
            conns,
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(0),
            IdMapVersion(0),
        );

        idmap.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap
            .insert_many(
                &ctx,
                vec![
                    (DagId(1), ONES_CSID),
                    (DagId(2), TWOS_CSID),
                    (DagId(3), THREES_CSID),
                    (DagId(4), FOURS_CSID),
                ],
            )
            .await?;

        let response = idmap
            .find_many_dag_ids(&ctx, vec![ONES_CSID, TWOS_CSID, THREES_CSID])
            .await?;
        assert_eq!(
            response,
            hashmap![ONES_CSID => DagId(1), TWOS_CSID => DagId(2), THREES_CSID => DagId(3)]
        );

        let response = idmap
            .find_many_dag_ids_maybe_stale(&ctx, vec![ONES_CSID, TWOS_CSID, THREES_CSID])
            .await?;
        assert_eq!(
            response,
            hashmap![ONES_CSID => DagId(1), TWOS_CSID => DagId(2)]
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_many_repo_id_many_versions(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

        let idmap11 = SqlIdMap::new(
            conns.0.clone(),
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(1),
            IdMapVersion(1),
        );
        let idmap12 = SqlIdMap::new(
            conns.0.clone(),
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(1),
            IdMapVersion(2),
        );
        let idmap21 = SqlIdMap::new(
            conns.0.clone(),
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(2),
            IdMapVersion(1),
        );

        idmap11.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap12.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap21.insert(&ctx, DagId(0), AS_CSID).await?;

        idmap11.insert(&ctx, DagId(1), ONES_CSID).await?;
        idmap11.insert(&ctx, DagId(2), TWOS_CSID).await?;
        idmap12.insert(&ctx, DagId(1), TWOS_CSID).await?;
        idmap21.insert(&ctx, DagId(1), FOURS_CSID).await?;
        idmap21.insert(&ctx, DagId(2), ONES_CSID).await?;

        assert_eq!(idmap11.get_changeset_id(&ctx, DagId(1)).await?, ONES_CSID);
        assert_eq!(idmap11.get_changeset_id(&ctx, DagId(2)).await?, TWOS_CSID);
        assert_eq!(idmap12.get_changeset_id(&ctx, DagId(1)).await?, TWOS_CSID);
        assert_eq!(idmap21.get_changeset_id(&ctx, DagId(1)).await?, FOURS_CSID);
        assert_eq!(idmap21.get_changeset_id(&ctx, DagId(2)).await?, ONES_CSID);

        assert_eq!(idmap11.get_dag_id(&ctx, ONES_CSID).await?, DagId(1));
        assert_eq!(idmap11.get_dag_id(&ctx, TWOS_CSID).await?, DagId(2));

        Ok(())
    }

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

        assert!(idmap.find_many_dag_ids(&ctx, vec![]).await?.is_empty());
        assert!(
            idmap
                .find_many_changeset_ids(&ctx, vec![])
                .await?
                .is_empty()
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_copy_map(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

        let idmap1 = SqlIdMap::new(
            conns.0.clone(),
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(1),
            IdMapVersion(1),
        );
        idmap1.insert(&ctx, DagId(0), AS_CSID).await?;
        // This way, if we don't truncate correctly, things will go wrong.
        idmap1.insert(&ctx, DagId(1), THREES_CSID).await?;
        idmap1.insert(&ctx, DagId(2), FOURS_CSID).await?;

        let idmap2 = idmap1.copy(DagId(0), IdMapVersion(2)).await?;
        // Check the copy didn't go too deep
        assert!(idmap2.get_changeset_id(&ctx, DagId(1)).await.is_err());
        assert!(idmap2.get_changeset_id(&ctx, DagId(2)).await.is_err());
        assert!(idmap2.get_dag_id(&ctx, THREES_CSID).await.is_err());
        assert!(idmap2.get_dag_id(&ctx, FOURS_CSID).await.is_err());

        idmap2.insert(&ctx, DagId(1), ONES_CSID).await?;
        idmap2.insert(&ctx, DagId(2), FIVES_CSID).await?;

        let idmap3 = idmap2.copy(DagId(1), IdMapVersion(3)).await?;
        // Check the copy didn't go too deep
        assert!(idmap3.get_changeset_id(&ctx, DagId(2)).await.is_err());
        assert!(idmap3.get_dag_id(&ctx, FOURS_CSID).await.is_err());

        idmap3.insert(&ctx, DagId(2), TWOS_CSID).await?;

        let idmap4 = idmap3.copy(DagId(2), IdMapVersion(4)).await?;
        idmap4.insert(&ctx, DagId(3), THREES_CSID).await?;

        // Check native table still works
        assert_eq!(idmap3.get_changeset_id(&ctx, DagId(2)).await?, TWOS_CSID);
        assert_eq!(idmap3.get_dag_id(&ctx, TWOS_CSID).await?, DagId(2));
        assert_eq!(idmap2.get_changeset_id(&ctx, DagId(1)).await?, ONES_CSID);
        assert_eq!(idmap2.get_dag_id(&ctx, ONES_CSID).await?, DagId(1));
        assert_eq!(idmap1.get_changeset_id(&ctx, DagId(0)).await?, AS_CSID);
        assert_eq!(idmap1.get_dag_id(&ctx, AS_CSID).await?, DagId(0));

        // And both levels of copy
        assert_eq!(idmap3.get_changeset_id(&ctx, DagId(1)).await?, ONES_CSID);
        assert_eq!(idmap3.get_dag_id(&ctx, ONES_CSID).await?, DagId(1));
        assert_eq!(idmap3.get_changeset_id(&ctx, DagId(0)).await?, AS_CSID);
        assert_eq!(idmap3.get_dag_id(&ctx, AS_CSID).await?, DagId(0));
        assert_eq!(idmap2.get_changeset_id(&ctx, DagId(0)).await?, AS_CSID);
        assert_eq!(idmap2.get_dag_id(&ctx, AS_CSID).await?, DagId(0));

        // IdMap 4 is needed to catch cases where the copy doesn't go deep enough
        // Just check that it's fine
        assert_eq!(idmap4.get_changeset_id(&ctx, DagId(0)).await?, AS_CSID);
        assert_eq!(idmap4.get_dag_id(&ctx, AS_CSID).await?, DagId(0));
        assert_eq!(idmap4.get_changeset_id(&ctx, DagId(1)).await?, ONES_CSID);
        assert_eq!(idmap4.get_dag_id(&ctx, ONES_CSID).await?, DagId(1));
        assert_eq!(idmap4.get_changeset_id(&ctx, DagId(2)).await?, TWOS_CSID);
        assert_eq!(idmap4.get_dag_id(&ctx, TWOS_CSID).await?, DagId(2));
        assert_eq!(idmap4.get_changeset_id(&ctx, DagId(3)).await?, THREES_CSID);
        assert_eq!(idmap4.get_dag_id(&ctx, THREES_CSID).await?, DagId(3));

        Ok(())
    }

    #[fbinit::test]
    async fn test_insert_fail_after_copy_map(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

        let idmap1 = SqlIdMap::new(
            conns.0.clone(),
            Arc::new(NoReplicaLagMonitor()),
            RepositoryId::new(1),
            IdMapVersion(1),
        );
        idmap1.insert(&ctx, DagId(0), AS_CSID).await?;
        idmap1.insert(&ctx, DagId(2), BS_CSID).await?;

        let idmap2 = idmap1.copy(DagId(2), IdMapVersion(2)).await?;
        idmap2.insert(&ctx, DagId(3), ONES_CSID).await?;

        assert!(
            idmap1.insert(&ctx, DagId(1), TWOS_CSID).await.is_err(),
            "Inserted into a copy base table"
        );
        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.
 */

//! Segmented Changelog Types

use std::collections::HashMap;

use anyhow::format_err;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
pub use dag;
pub use dag::CloneData;
pub use dag::FirstAncestorConstraint;
pub use dag::FlatSegment;
pub use dag::Group;
pub use dag::Id as DagId;
pub use dag::IdSet as DagIdSet;
pub use dag::InProcessIdDag;
pub use dag::Location;
pub use dag::PreparedFlatSegments;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use thiserror::Error;

#[facet::facet]
#[async_trait]
pub trait SegmentedChangelog: Send + Sync {
    /// Get the identifier of a commit given it's commit graph location.
    ///
    /// The client using segmented changelog will have only a set of identifiers for the commits in
    /// the graph. To retrieve the identifier of an commit that is now known they will provide a
    /// known descendant and the distance from the known commit to the commit we inquire about.
    async fn location_to_changeset_id(
        &self,
        ctx: &CoreContext,
        location: Location<ChangesetId>,
    ) -> Result<ChangesetId> {
        let mut ids = self
            .location_to_many_changeset_ids(ctx, location, 1)
            .await?;
        if ids.len() == 1 {
            if let Some(id) = ids.pop() {
                return Ok(id);
            }
        }
        Err(format_err!(
            "unexpected result from location_to_many_changeset_ids"
        ))
    }

    /// Get identifiers of a continuous set of commit given their commit graph location.
    ///
    /// Similar to `location_to_changeset_id` but instead of returning the ancestor that is
    /// `distance` away from the `known` commit, it returns `count` ancestors following the parents.
    /// It is expected that all but the last ancestor will have a single parent.
    async fn location_to_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        location: Location<ChangesetId>,
        count: u64,
    ) -> Result<Vec<ChangesetId>>;

    /// Get the graph location of a given commit identifier.
    ///
    /// ## Practical use-cases
    ///
    /// A lazy changelog client knows the "shape" of the commit graph where
    /// vertexes in the graph are labeled using numbers. It does not know
    /// all of the commit hashes corresponding to the numbers. When a commit
    /// hash was referred (ex. from user input), the client wants to "locate"
    /// the commit in its local graph, or be confident that the hash does not
    /// exist in its local graph.
    ///
    ///
    /// ## Principle (How it works)
    ///
    /// A lazy changelog client knows certain "anchor" commits (described as
    /// "universally known" in the `dag` crate), and can use those commits
    /// as "anchor points" to locate other commits. For example:
    ///
    /// The client wants to resolve C to its incomplete graph (only A is known).
    /// The server knows the client knows A's location.
    ///
    /// ```plain,ignore
    ///     Client | Server
    ///       A 30 |  A 55
    ///       |    |  |
    ///       ? 29 |  B 54
    ///       |    |  |
    ///       ? 28 |  C 53
    /// ```
    ///
    /// The server might use different integer IDs assigned in the graph so it
    /// cannot return C's server-side integer ID (53) directly. Instead, it
    /// translates `C` into `A~2` (revset notation, 1st parent of 1st parent of
    /// A), sends `A~2` to client, then client resolves `A~2` locally to integer
    /// ID 28.
    ///
    ///
    /// ## Heads
    ///
    /// To understand what "anchor" commits client has, this API requires
    /// `master_heads` from the client, meaning the client's lazy portion of
    /// the graph contains *exactly* `ancestors(master_heads)`. This indicates
    /// important properties like:
    ///
    /// 1. If the server returns `Ok(None)`, then the client can be confident
    ///    that the `cs_id` does not exist in the `ancestors(master_heads)`
    ///    sub-graph.
    /// 2. The server knows that `master_heads` and
    ///    `parents(ancestors(master_heads) & merge())` are "anchor"s known
    ///    by the client. The server should use those commit hashes as the `X`
    ///    part of `X~n` in responses.
    ///
    /// Providing precise heads is important. For example, suppose the client
    /// wants to resolve commit hash `X` with heads `Y`.  Then the server should
    /// return `Ok(None)` despite that the server knows `X`, because `X` is
    /// outside `ancestors(Y)`.
    ///
    /// ```plain,ignore
    ///     Client | Server
    ///            |  X
    ///            |  |
    ///       Y    |  Y
    ///       |    |  |
    ///       ?    |  Z
    /// ```
    ///
    /// Another example, commit `Z` can be represented as `Y~1`, `X~1`, or
    /// `Q~2`.  Client1 tries to resolve `Z` using heads `Y`, then the server
    /// should return `Y~1`, because `X` is unknown to client1. Client2 resolves
    /// `Z` using heads `Q`, then the server should return `Q~2`.
    ///
    /// ```plain,ignore
    ///     Client1 | Client2 | Server
    ///             | Q       | P Q
    ///             | |       | |\|
    ///     Y       | X       | Y X
    ///     |       | |       | |/
    ///     ?       | ?       | Z
    /// ```
    ///
    ///
    /// ## Multiple heads
    ///
    /// In a repo with multiple mainline branches, or when the unique master
    /// bookmark moves backwards, the client might send more than 1 head.
    /// The server's DAG might have more than 1 head too.
    ///
    /// Suppose the client provides two heads `A` and `B`.
    ///
    /// - If the server knows both `A` and `B`, great. Then the server can
    ///   process the request as usual.
    ///
    /// - If the server only knows `A` but does not know `B`, the server should
    ///   never return `Ok(None)`, because it cannot confirm that the `cs_id`
    ///   exists in `B % A` (revset notation) or not. This means the client can
    ///   can never confirm a commit hash does not exist in the graph, and
    ///   probably breaks a bunch of workflows until the client gets rid of the
    ///   troublesome heads (by rebuilding the graph).
    ///   If `cs_id` is an ancestor of `A`, the server can resolve it as if the client provides
    ///   only `A` as the head, or return an error (correct, but provides less optimal
    ///   UX when master moves backwards).
    async fn changeset_id_to_location(
        &self,
        ctx: &CoreContext,
        master_heads: Vec<ChangesetId>,
        cs_id: ChangesetId,
    ) -> Result<Option<Location<ChangesetId>>> {
        let mut ids = self
            .many_changeset_ids_to_locations(ctx, master_heads, vec![cs_id])
            .await?;
        ids.remove(&cs_id).transpose()
    }

    /// Get the graph locations given a set of commit identifier.
    ///
    /// Batch variation of `changeset_id_to_location`. The assumption is that we are dealing with
    /// the same client repository so the `head` parameter stays the same between changesets.
    ///
    /// See `changeset_id_to_location` for corner cases of this method.
    async fn many_changeset_ids_to_locations(
        &self,
        ctx: &CoreContext,
        master_heads: Vec<ChangesetId>,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, Result<Location<ChangesetId>>>>;

    /// Returns data necessary for SegmentedChangelog to be initialized by a client.
    ///
    /// Note that the heads that are sent over in a clone can vary. Strictly speaking the client
    /// only needs one head.
    ///
    /// The HashMap is provided to contain hints for Mercurial clones - it has any pre-cached
    /// bonsai to hg mappings we have loaded as part of clone, and may be empty.
    async fn clone_data(
        &self,
        ctx: &CoreContext,
    ) -> Result<(CloneData<ChangesetId>, HashMap<ChangesetId, HgChangesetId>)>;

    /// Returns data that client can import to perform a lazy fast pull.
    /// `common` specifies heads known by both client and server.
    /// `missing` specifies heads unknown by the client, and are being pulled.
    ///
    /// Note this endpoint does not handle all `pull` use-cases. It requires
    /// `roots(missing % common)` to be fully missing client-side, and
    /// `parents(roots(missing % common))` to exist client-side in its MASTER
    /// group. It's possible that client-side has `parents(roots(missing %
    /// common))` in its NON_MASTER group, for example, the client has some
    /// commits that are merged by the server. At present, that prevents the
    /// client from importing the pull data.
    async fn pull_data(
        &self,
        ctx: &CoreContext,
        common: Vec<ChangesetId>,
        missing: Vec<ChangesetId>,
    ) -> Result<CloneData<ChangesetId>>;

    /// Whether segmented changelog is disabled.
    ///
    /// A quick way to test if the backend supports segmented changelog or not
    /// without doing real work.
    ///
    /// Return true if it is disabled.
    async fn disabled(&self, ctx: &CoreContext) -> Result<bool>;

    /// Test if `ancestor` is an ancestor of `descendant`.
    /// Returns None in case segmented changelog doesn't know about either of those commit.
    async fn is_ancestor(
        &self,
        ctx: &CoreContext,
        ancestor: ChangesetId,
        descendant: ChangesetId,
    ) -> Result<Option<bool>>;

    /// Try update segmented changelog to given heads. No-op by default. Useful
    /// for tests. Returns: `true` if update was successful; `false` if the
    /// implementation doesn't support updates; an error otherwise.
    async fn build_up_to_heads(&self, _ctx: &CoreContext, _heads: &[ChangesetId]) -> Result<bool> {
        Ok(false)
    }
}

#[derive(Debug, Error)]
#[error("server cannot match the clients heads, repo {repo_id}, client_heads: {client_heads:?}")]
pub struct MismatchedHeadsError {
    pub repo_id: RepositoryId,
    pub client_heads: Vec<ChangesetId>,
}

impl MismatchedHeadsError {
    pub fn new(repo_id: RepositoryId, client_heads: Vec<ChangesetId>) -> Self {
        Self {
            repo_id,
            client_heads,
        }
    }
}


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

use std::sync::Arc;
use std::thread;

use chrono::Local;
use context::PerfCounterType;
use context::PerfCounters;

fn main() {
    let ctrs = Arc::new(PerfCounters::default());
    let k = PerfCounterType::BlobGets;
    let k2 = PerfCounterType::BlobPuts;

    let start = Local::now();
    println!("Start: {}", start);

    let n_threads = 10;
    let n_ops = 100000;

    let threads = (0..n_threads).map(|_| {
        thread::spawn({
            let ctrs = ctrs.clone();
            move || {
                for i in 0..n_ops {
                    ctrs.increment_counter(k);
                    ctrs.set_max_counter(k2, i);
                }
            }
        })
    });

    for t in threads {
        t.join().unwrap();
    }

    let done = Local::now();

    assert_eq!(ctrs.get_counter(k), n_threads * n_ops);
    assert_eq!(ctrs.get_counter(k2), n_ops - 1);

    println!("Elapsed: {}ms", (done - start).num_milliseconds());
}


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

use std::sync::Arc;

use fbinit::FacebookInit;
use metadata::Metadata;
use scribe_ext::Scribe;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::Logger;
use slog_glog_fmt::logger_that_can_work_in_tests;

use crate::logging::LoggingContainer;
use crate::logging::SamplingKey;
use crate::perf_counters::PerfCounters;
use crate::perf_counters_stack::PerfCountersStack;
use crate::session::SessionClass;
use crate::session::SessionContainer;

#[derive(Clone)]
pub struct CoreContext {
    pub fb: FacebookInit,
    session: SessionContainer,
    logging: LoggingContainer,
}

impl CoreContext {
    pub fn new(fb: FacebookInit, logging: LoggingContainer, session: SessionContainer) -> Self {
        Self {
            fb,
            logging,
            session,
        }
    }

    pub fn new_with_logger(fb: FacebookInit, logger: Logger) -> Self {
        let session = SessionContainer::new_with_defaults(fb);
        session.new_context(logger, MononokeScubaSampleBuilder::with_discard())
    }

    // Context for bulk processing like scrubbing or bulk backfilling
    pub fn new_for_bulk_processing(fb: FacebookInit, logger: Logger) -> Self {
        let session = SessionContainer::builder(fb)
            .session_class(SessionClass::Background)
            .build();
        session.new_context(logger, MononokeScubaSampleBuilder::with_discard())
    }

    pub fn test_mock(fb: FacebookInit) -> Self {
        let session = SessionContainer::new_with_defaults(fb);

        Self::test_mock_session(session)
    }

    pub fn test_mock_session(session: SessionContainer) -> Self {
        let logger = logger_that_can_work_in_tests().unwrap();
        session.new_context(logger, MononokeScubaSampleBuilder::with_discard())
    }

    /// Create a new CoreContext, with a reset LoggingContainer. This is useful to reset perf
    /// counters. The existing CoreContext is unaffected.
    pub fn clone_and_reset(&self) -> Self {
        self.session
            .new_context(self.logger().clone(), self.scuba().clone())
    }

    pub fn clone_and_sample(&self, sampling_key: SamplingKey) -> Self {
        Self {
            fb: self.fb,
            session: self.session.clone(),
            logging: self.logging.clone_and_sample(sampling_key),
        }
    }

    pub fn clone_with_logger(&self, logger: Logger) -> Self {
        Self {
            fb: self.fb,
            session: self.session.clone(),
            logging: self.logging.clone_with_logger(logger),
        }
    }

    pub fn clone_with_repo_name(&self, repo_name: &str) -> Self {
        Self {
            fb: self.fb,
            session: self.session.clone(),
            logging: self.logging.clone_with_repo_name(repo_name),
        }
    }

    pub fn with_mutated_scuba(
        &self,
        mutator: impl FnOnce(MononokeScubaSampleBuilder) -> MononokeScubaSampleBuilder,
    ) -> Self {
        Self {
            fb: self.fb,
            session: self.session.clone(),
            logging: self.logging.with_mutated_scuba(mutator),
        }
    }

    pub fn logger(&self) -> &Logger {
        self.logging.logger()
    }

    pub fn sampling_key(&self) -> Option<&SamplingKey> {
        self.logging.sampling_key()
    }

    pub fn scuba(&self) -> &MononokeScubaSampleBuilder {
        self.logging.scuba()
    }

    pub fn perf_counters(&self) -> &PerfCountersStack {
        self.logging.perf_counters()
    }

    pub fn metadata(&self) -> &Metadata {
        self.session.metadata()
    }

    pub fn session(&self) -> &SessionContainer {
        &self.session
    }

    pub fn session_mut(&mut self) -> &mut SessionContainer {
        &mut self.session
    }

    pub fn scribe(&self) -> &Scribe {
        self.logging.scribe()
    }

    pub fn fork_perf_counters(&mut self) -> Arc<PerfCounters> {
        self.logging.fork_perf_counters()
    }
}


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

#![feature(never_type)]
#![feature(result_flattening)]

pub use session_id::SessionId;

pub use crate::core::CoreContext;
pub use crate::logging::LoggingContainer;
pub use crate::logging::SamplingKey;
pub use crate::perf_counters::PerfCounterType;
pub use crate::perf_counters::PerfCounters;
pub use crate::session::SessionClass;
pub use crate::session::SessionContainer;
pub use crate::session::SessionContainerBuilder;

mod core;
mod logging;
mod perf_counters;
mod perf_counters_stack;
mod session;


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

use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use fbinit::FacebookInit;
use scribe_ext::Scribe;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::o;
use slog::Logger;

use crate::perf_counters::PerfCounters;
use crate::perf_counters_stack::PerfCountersStack;

/// Used to correlation a high level action on a CoreContext
/// e.g. walk of a repo,  with low level actions using that context
/// e.g. which blobs are read
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct SamplingKey {
    inner_key: u32,
}

static NEXT_SAMPLING_KEY: AtomicU32 = AtomicU32::new(0);

impl SamplingKey {
    pub fn new() -> Self {
        let v = NEXT_SAMPLING_KEY.fetch_add(1, Ordering::Relaxed);
        Self { inner_key: v }
    }

    pub fn inner(&self) -> u32 {
        self.inner_key
    }
}

#[derive(Debug, Clone)]
pub struct LoggingContainer {
    logger: Logger,
    scuba: Arc<MononokeScubaSampleBuilder>,
    perf_counters: PerfCountersStack,
    sampling_key: Option<SamplingKey>,
    scribe: Scribe,
}

impl LoggingContainer {
    pub fn new(fb: FacebookInit, logger: Logger, scuba: MononokeScubaSampleBuilder) -> Self {
        Self {
            logger,
            scuba: Arc::new(scuba),
            perf_counters: Default::default(),
            sampling_key: None,
            scribe: Scribe::new(fb),
        }
    }

    pub fn fork_perf_counters(&mut self) -> Arc<PerfCounters> {
        let (perf_counters, ret) = self.perf_counters.fork();
        self.perf_counters = perf_counters;
        ret
    }

    pub fn clone_and_sample(&self, sampling_key: SamplingKey) -> Self {
        Self {
            logger: self.logger.clone(),
            scuba: self.scuba.clone(),
            perf_counters: self.perf_counters.clone(),
            sampling_key: Some(sampling_key),
            scribe: self.scribe.clone(),
        }
    }

    pub fn clone_with_logger(&self, logger: Logger) -> Self {
        Self {
            logger,
            scuba: self.scuba.clone(),
            perf_counters: self.perf_counters.clone(),
            sampling_key: self.sampling_key.clone(),
            scribe: self.scribe.clone(),
        }
    }

    pub fn clone_with_repo_name(&self, repo_name: &str) -> Self {
        Self {
            logger: self.logger.new(o!("repo" => repo_name.to_string())),
            scuba: self.scuba.clone(),
            perf_counters: self.perf_counters.clone(),
            sampling_key: self.sampling_key.clone(),
            scribe: self.scribe.clone(),
        }
    }

    pub fn with_scribe(&mut self, scribe: Scribe) -> &mut Self {
        self.scribe = scribe;
        self
    }

    pub fn logger(&self) -> &Logger {
        &self.logger
    }

    pub fn scuba(&self) -> &MononokeScubaSampleBuilder {
        &self.scuba
    }

    pub fn perf_counters(&self) -> &PerfCountersStack {
        &self.perf_counters
    }

    pub fn sampling_key(&self) -> Option<&SamplingKey> {
        self.sampling_key.as_ref()
    }

    pub fn scribe(&self) -> &Scribe {
        &self.scribe
    }

    pub fn with_mutated_scuba(
        &self,
        mutator: impl FnOnce(MononokeScubaSampleBuilder) -> MononokeScubaSampleBuilder,
    ) -> Self {
        Self {
            logger: self.logger.clone(),
            scuba: Arc::new(mutator(self.scuba().clone())),
            perf_counters: self.perf_counters.clone(),
            sampling_key: self.sampling_key.clone(),
            scribe: self.scribe.clone(),
        }
    }
}


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

use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;

use scuba_ext::MononokeScubaSampleBuilder;

macro_rules! define_perf_counters {
    (enum $enum_name:ident {
        $($variant:ident),*,
    }) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
        pub enum $enum_name {
            $($variant),*
        }

        impl $enum_name {
            pub fn name(&self) -> &'static str {
                match self {
                    $($enum_name::$variant => stringify!($variant)),*
                }
            }
        }

        pub const PERF_COUNTERS: &[PerfCounterType] = &[
            $($enum_name::$variant),*
        ];

        #[allow(non_snake_case)]
        #[derive(Debug, Default)]
        pub struct PerfCounters {
            $($variant: AtomicI64),*
        }

        impl PerfCounters {
            fn get_counter_atomic(&self, counter: $enum_name) -> &AtomicI64 {
                match counter {
                    $($enum_name::$variant => &self.$variant),*
                }
            }
        }
    };
}

define_perf_counters! {
    enum PerfCounterType {
        BlobGets,
        BlobGetsNotFound,
        BlobGetsAccessWait,
        BlobGetsShardAccessWait,
        BlobGetsMaxLatency,
        BlobGetsNotFoundMaxLatency,
        BlobGetsDeduplicated,
        BlobGetsDeduplicatedLarge,
        BlobGetsTotalSize,
        BlobPresenceChecks,
        BlobPresenceChecksMaxLatency,
        BlobPuts,
        BlobPutsAccessWait,
        BlobPutsShardAccessWait,
        BlobPutsMaxLatency,
        BlobPutsDeduplicated,
        BlobPutsTotalSize,
        BytesSent,
        CachelibHits,
        CachelibMisses,
        EdenapiFiles,
        EdenapiTrees,
        GetbundleFilenodesTotalWeight,
        GetbundleNumCommits,
        GetbundleNumDrafts,
        GetbundleNumFilenodes,
        GetbundleNumManifests,
        GetbundlePartialTraversal,
        GetcommitdataNumCommits,
        GetcommitdataResponseSize,
        GetpackMaxFileSize,
        GetpackNumFiles,
        GetpackNumPossibleLFSFiles,
        GetpackPossibleLFSFilesSumSize,
        GetpackResponseSize,
        FilenodesTooBigHistory,
        GettreepackDesignatedNodes,
        GettreepackNumTreepacks,
        GettreepackResponseSize,
        HgMutationStoreNumAdded,
        HgMutationStoreNumFetched,
        MemcacheHits,
        MemcacheMisses,
        NullLinknode,
        NumKnown,
        NumKnownRequested,
        NumUnknown,
        SegmentedChangelogServerSideOpsHits,
        SegmentedChangelogServerSideOpsFallbacks,
        SqlReadsMaster,
        SqlReadsReplica,
        SqlWrites,
        SumManifoldPollTime,
        UndesiredTreeFetch,
        ManifoldBlobSumDelay,
        ManifoldBlobRetries,
        ManifoldBlobConflicts,
        S3BlobRetries,
        S3BlobSumDelay,
        S3AccessWait,
    }
}

enum PerfCounterTypeUpdateFunc {
    Add,
    Max,
}

impl PerfCounterType {
    /// Whether to log the zero value of the counter
    fn always_log(&self) -> bool {
        use PerfCounterType::*;

        match self {
            BlobGets
            | BlobGetsNotFound
            | BlobGetsMaxLatency
            | BlobGetsNotFoundMaxLatency
            | BlobGetsTotalSize
            | BlobPresenceChecks
            | BlobPresenceChecksMaxLatency
            | BlobPuts
            | BlobPutsMaxLatency
            | BlobPutsTotalSize
            | CachelibHits
            | CachelibMisses
            | GetpackNumPossibleLFSFiles
            | GetpackPossibleLFSFilesSumSize
            | GettreepackDesignatedNodes
            | MemcacheHits
            | MemcacheMisses
            | SqlReadsMaster
            | SqlReadsReplica
            | SqlWrites => true,
            _ => false,
        }
    }

    fn expected_update_func(&self) -> PerfCounterTypeUpdateFunc {
        use PerfCounterType::*;

        match self {
            BlobGets
            | BlobGetsNotFound
            | BlobGetsAccessWait
            | BlobGetsShardAccessWait
            | BlobGetsDeduplicated
            | BlobGetsDeduplicatedLarge
            | BlobGetsTotalSize
            | BlobPresenceChecks
            | BlobPuts
            | BlobPutsAccessWait
            | BlobPutsShardAccessWait
            | BlobPutsDeduplicated
            | BlobPutsTotalSize
            | BytesSent
            | CachelibHits
            | CachelibMisses
            | EdenapiFiles
            | EdenapiTrees
            | GetbundleFilenodesTotalWeight
            | GetbundleNumCommits
            | GetbundleNumDrafts
            | GetbundleNumFilenodes
            | GetbundleNumManifests
            | GetbundlePartialTraversal
            | GetcommitdataNumCommits
            | GetcommitdataResponseSize
            | GetpackNumFiles
            | GetpackNumPossibleLFSFiles
            | GetpackPossibleLFSFilesSumSize
            | GetpackResponseSize
            | FilenodesTooBigHistory
            | GettreepackDesignatedNodes
            | GettreepackNumTreepacks
            | GettreepackResponseSize
            | HgMutationStoreNumAdded
            | HgMutationStoreNumFetched
            | MemcacheHits
            | MemcacheMisses
            | NullLinknode
            | NumKnown
            | NumKnownRequested
            | NumUnknown
            | SegmentedChangelogServerSideOpsHits
            | SegmentedChangelogServerSideOpsFallbacks
            | SqlReadsMaster
            | SqlReadsReplica
            | SqlWrites
            | SumManifoldPollTime
            | UndesiredTreeFetch
            | ManifoldBlobSumDelay
            | ManifoldBlobRetries
            | ManifoldBlobConflicts
            | S3BlobRetries
            | S3BlobSumDelay
            | S3AccessWait => PerfCounterTypeUpdateFunc::Add,
            BlobGetsMaxLatency
            | BlobGetsNotFoundMaxLatency
            | BlobPresenceChecksMaxLatency
            | BlobPutsMaxLatency
            | GetpackMaxFileSize => PerfCounterTypeUpdateFunc::Max,
        }
    }
}

impl PerfCounters {
    pub fn set_counter(&self, counter: PerfCounterType, val: i64) {
        self.get_counter_atomic(counter)
            .store(val, Ordering::Relaxed);
    }

    pub fn increment_counter(&self, counter: PerfCounterType) {
        self.get_counter_atomic(counter)
            .fetch_add(1, Ordering::Relaxed);
    }

    pub fn decrement_counter(&self, counter: PerfCounterType) {
        self.get_counter_atomic(counter)
            .fetch_sub(1, Ordering::Relaxed);
    }

    pub fn add_to_counter(&self, counter: PerfCounterType, val: i64) {
        self.get_counter_atomic(counter)
            .fetch_add(val, Ordering::Relaxed);
    }

    pub fn set_max_counter(&self, counter: PerfCounterType, val: i64) {
        self.get_counter_atomic(counter)
            .fetch_max(val, Ordering::Relaxed);
    }

    pub fn get_counter(&self, counter: PerfCounterType) -> i64 {
        self.get_counter_atomic(counter).load(Ordering::Relaxed)
    }

    pub fn update_with_counters(&self, counters: &PerfCounters) {
        for key in PERF_COUNTERS.iter() {
            match key.expected_update_func() {
                PerfCounterTypeUpdateFunc::Add => {
                    self.add_to_counter(*key, counters.get_counter(*key))
                }
                PerfCounterTypeUpdateFunc::Max => {
                    self.set_max_counter(*key, counters.get_counter(*key))
                }
            };
        }
    }

    pub fn insert_nonzero_perf_counters(&self, builder: &mut MononokeScubaSampleBuilder) {
        for key in PERF_COUNTERS.iter() {
            let value = self.get_counter(*key);
            if value != 0 {
                builder.add(key.name(), value);
            }
        }
    }

    pub fn insert_perf_counters(&self, builder: &mut MononokeScubaSampleBuilder) {
        // NOTE: we log 0 mainly so that we can distinguish
        // nulls i.e. "not logged" and 0 as in "zero calls to blobstore".
        // Logging 0 allows counting avg, p50 etc statistic.
        for key in PERF_COUNTERS.iter() {
            let value = self.get_counter(*key);
            if value != 0 || key.always_log() {
                builder.add(key.name(), value);
            }
        }
    }
}

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

    #[test]
    fn test_perf_counter() {
        // NOTE: This test doesn't try to do anything fancy or test concurrency. It does however
        // check that we pass valid values for Ordering (invalid values panic on atomics).
        let ctrs = PerfCounters::default();
        let k = PerfCounterType::BlobGets;

        ctrs.set_counter(k, 1);
        assert_eq!(ctrs.get_counter(k), 1);

        ctrs.increment_counter(k);
        assert_eq!(ctrs.get_counter(k), 2);

        ctrs.decrement_counter(k);
        assert_eq!(ctrs.get_counter(k), 1);

        ctrs.add_to_counter(k, 1);
        assert_eq!(ctrs.get_counter(k), 2);

        ctrs.set_max_counter(k, 3);
        assert_eq!(ctrs.get_counter(k), 3);

        ctrs.set_max_counter(k, 2);
        assert_eq!(ctrs.get_counter(k), 3);
    }
}


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

use scuba_ext::MononokeScubaSampleBuilder;

use crate::perf_counters::PerfCounterType;
use crate::perf_counters::PerfCounters;

#[derive(Debug, Clone)]
pub struct PerfCountersStack {
    inner: Arc<PerfCountersStackInner>,
}

impl Default for PerfCountersStack {
    fn default() -> Self {
        Self {
            inner: Arc::new(Default::default()),
        }
    }
}

impl PerfCountersStack {
    pub(crate) fn fork(&self) -> (Self, Arc<PerfCounters>) {
        let new = Arc::new(PerfCounters::default());

        let mut inner = PerfCountersStackInner {
            top: self.inner.top.clone(),
            rest: self.inner.rest.clone(),
        };

        inner.rest.push(new.clone());

        (
            Self {
                inner: Arc::new(inner),
            },
            new,
        )
    }
}

/// Helper macro to implement that methods that operate on each of the counters in the stack.
macro_rules! impl_counter_methods {
    ( $( $vis:vis fn $method_name:ident(&self, $( $param_name:ident : $param_type:ty ),*); )* ) => {
        $(
            $vis fn $method_name(
                &self,
                $( $param_name: $param_type ),*
            ) {
                for c in self.iter_counters() {
                    c.$method_name($( $param_name, )*);
                }
            }
        )*
    }
}

impl PerfCountersStack {
    fn iter_counters(&self) -> impl Iterator<Item = &Arc<PerfCounters>> {
        self.inner.rest.iter().chain(iter::once(&self.inner.top))
    }

    impl_counter_methods! {
        pub fn set_counter(&self, counter: PerfCounterType, val: i64);
        pub fn increment_counter(&self, counter: PerfCounterType);
        pub fn decrement_counter(&self, counter: PerfCounterType);
        pub fn add_to_counter(&self, counter: PerfCounterType, val: i64);
        pub fn set_max_counter(&self, counter: PerfCounterType, val: i64);
    }

    pub fn top(&self) -> &PerfCounters {
        &self.inner.top
    }

    pub fn get_counter(&self, counter: PerfCounterType) -> i64 {
        self.inner.top.get_counter(counter)
    }

    pub fn insert_perf_counters(&self, builder: &mut MononokeScubaSampleBuilder) {
        self.inner.top.insert_perf_counters(builder)
    }
}

#[derive(Default, Debug)]
struct PerfCountersStackInner {
    top: Arc<PerfCounters>,
    rest: Vec<Arc<PerfCounters>>,
}

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

    #[test]
    fn test_perf_counter_stack() {
        let k = PerfCounterType::BlobGets;

        let s = PerfCountersStack::default();
        s.increment_counter(k);
        assert_eq!(s.get_counter(k), 1);

        let (s, f1) = s.fork();
        assert_eq!(s.get_counter(k), 1);
        assert_eq!(f1.get_counter(k), 0);

        s.increment_counter(k);
        assert_eq!(s.get_counter(k), 2);
        assert_eq!(f1.get_counter(k), 1);

        let (s, f2) = s.fork();
        assert_eq!(s.get_counter(k), 2);
        assert_eq!(f1.get_counter(k), 1);
        assert_eq!(f2.get_counter(k), 0);

        s.increment_counter(k);
        assert_eq!(s.get_counter(k), 3);
        assert_eq!(f1.get_counter(k), 2);
        assert_eq!(f2.get_counter(k), 1);
    }
}


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

use async_limiter::AsyncLimiter;
use fbinit::FacebookInit;
use governor::Quota;
use governor::RateLimiter;
use metadata::Metadata;
use rate_limiting::BoxRateLimiter;

use super::SessionClass;
use super::SessionContainer;
use super::SessionContainerInner;

pub struct SessionContainerBuilder {
    fb: FacebookInit,
    inner: SessionContainerInner,
    session_class: SessionClass,
}

impl SessionContainerBuilder {
    pub fn build(self) -> SessionContainer {
        SessionContainer {
            fb: self.fb,
            inner: Arc::new(self.inner),
            session_class: self.session_class,
        }
    }

    pub fn new(fb: FacebookInit) -> Self {
        Self {
            fb,
            inner: SessionContainerInner {
                metadata: Arc::new(Metadata::default()),
                rate_limiter: None,
                blobstore_write_limiter: None,
                blobstore_read_limiter: None,
                readonly: false,
            },
            session_class: SessionClass::UserWaiting,
        }
    }

    pub fn metadata(mut self, value: Arc<Metadata>) -> Self {
        self.inner.metadata = value;
        self
    }

    pub fn rate_limiter(mut self, value: impl Into<Option<BoxRateLimiter>>) -> Self {
        self.inner.rate_limiter = value.into();
        self
    }

    pub fn blobstore_read_limiter(mut self, limiter: AsyncLimiter) -> Self {
        self.inner.blobstore_read_limiter = Some(limiter);
        self
    }

    pub async fn blobstore_maybe_read_qps_limiter(mut self, qps: impl TryInto<u32>) -> Self {
        if let Ok(Some(qps)) = qps.try_into().map(NonZeroU32::new) {
            self.inner.blobstore_read_limiter =
                Some(AsyncLimiter::new(RateLimiter::direct(Quota::per_second(qps))).await);
        }
        self
    }

    pub fn blobstore_write_limiter(mut self, limiter: AsyncLimiter) -> Self {
        self.inner.blobstore_write_limiter = Some(limiter);
        self
    }

    pub async fn blobstore_maybe_write_qps_limiter(mut self, qps: impl TryInto<u32>) -> Self {
        if let Ok(Some(qps)) = qps.try_into().map(NonZeroU32::new) {
            self.inner.blobstore_write_limiter =
                Some(AsyncLimiter::new(RateLimiter::direct(Quota::per_second(qps))).await);
        }
        self
    }
    pub fn session_class(mut self, value: SessionClass) -> Self {
        self.session_class = value;
        self
    }

    pub fn readonly(mut self, readonly: bool) -> Self {
        self.inner.readonly = readonly;
        self
    }
}


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

use std::sync::Arc;

use async_limiter::AsyncLimiter;
use fbinit::FacebookInit;
use metadata::Metadata;
use permission_checker::MononokeIdentitySetExt;
use rate_limiting::BoxRateLimiter;
use rate_limiting::LoadCost;
use rate_limiting::Metric;
use rate_limiting::RateLimitReason;
use rate_limiting::RateLimiter;
use scribe_ext::Scribe;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::Logger;

pub use self::builder::SessionContainerBuilder;
use crate::core::CoreContext;
use crate::logging::LoggingContainer;

mod builder;

#[derive(Clone)]
pub struct SessionContainer {
    fb: FacebookInit,
    inner: Arc<SessionContainerInner>,
    session_class: SessionClass,
}

/// Represents the reason this session is running
#[derive(Clone, Copy)]
pub enum SessionClass {
    /// There is someone waiting for this session to complete.
    UserWaiting,
    /// The session is doing background work (e.g. backfilling).
    /// Wherever reasonable, prefer to slow down and wait for work to complete
    /// fully rather than pushing work out to other tasks.
    Background,
    /// Same as Background, but if some work is taking too long to complete
    /// then fallback to normal (i.e. UserWaiting) behavior
    BackgroundUnlessTooSlow,
    /// This session is used by the warm bookmarks cache.
    WarmBookmarksCache,
}

struct SessionContainerInner {
    metadata: Arc<Metadata>,
    rate_limiter: Option<BoxRateLimiter>,
    blobstore_write_limiter: Option<AsyncLimiter>,
    blobstore_read_limiter: Option<AsyncLimiter>,
    // Whether this session is supposed to be readonly, this will cause the right
    // AuthContext to constructed.
    readonly: bool,
}

impl SessionContainer {
    pub fn builder(fb: FacebookInit) -> SessionContainerBuilder {
        SessionContainerBuilder::new(fb)
    }

    pub fn new_with_defaults(fb: FacebookInit) -> Self {
        Self::builder(fb).build()
    }

    pub fn new_context(&self, logger: Logger, scuba: MononokeScubaSampleBuilder) -> CoreContext {
        let logging = LoggingContainer::new(self.fb, logger, scuba);

        CoreContext::new(self.fb, logging, self.clone())
    }

    pub fn new_context_with_scribe(
        &self,
        logger: Logger,
        scuba: MononokeScubaSampleBuilder,
        scribe: Scribe,
    ) -> CoreContext {
        let mut logging = LoggingContainer::new(self.fb, logger, scuba);
        logging.with_scribe(scribe);

        CoreContext::new(self.fb, logging, self.clone())
    }

    pub fn fb(&self) -> FacebookInit {
        self.fb
    }

    pub fn metadata(&self) -> &Metadata {
        &self.inner.metadata
    }

    pub fn rate_limiter(&self) -> Option<&(dyn RateLimiter + Send + Sync)> {
        match self.inner.rate_limiter {
            Some(ref rate_limiter) => Some(&**rate_limiter),
            None => None,
        }
    }

    pub fn bump_load(&self, metric: Metric, load: LoadCost) {
        if let Some(limiter) = self.rate_limiter() {
            limiter.bump_load(metric, load)
        }
    }

    pub fn check_load_shed(&self) -> Result<(), RateLimitReason> {
        match &self.inner.rate_limiter {
            Some(limiter) => limiter.check_load_shed(self.metadata().identities()),
            None => Ok(()),
        }
    }

    pub async fn check_rate_limit(&self, metric: Metric) -> Result<(), RateLimitReason> {
        match &self.inner.rate_limiter {
            Some(limiter) => limiter
                .check_rate_limit(metric, self.metadata().identities())
                .await
                .unwrap_or(Ok(())),
            None => Ok(()),
        }
    }

    pub fn is_quicksand(&self) -> bool {
        self.metadata().identities().is_quicksand()
    }

    pub fn is_readonly(&self) -> bool {
        self.inner.readonly
    }

    pub fn is_hg_sync_job(&self) -> bool {
        self.metadata().identities().is_hg_sync_job()
    }

    pub fn blobstore_read_limiter(&self) -> Option<&AsyncLimiter> {
        self.inner.blobstore_read_limiter.as_ref()
    }

    pub fn blobstore_write_limiter(&self) -> Option<&AsyncLimiter> {
        self.inner.blobstore_write_limiter.as_ref()
    }

    pub fn session_class(&self) -> SessionClass {
        self.session_class
    }

    pub fn override_session_class(&mut self, session_class: SessionClass) {
        self.session_class = session_class;
    }
}


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

#![feature(result_flattening)]

pub mod security;

use std::net::IpAddr;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use clientinfo::ClientInfo;
use permission_checker::MononokeIdentitySet;
use permission_checker::MononokeIdentitySetExt;
use session_id::generate_session_id;
use session_id::SessionId;
use tokio::time::timeout;
use trust_dns_resolver::TokioAsyncResolver;

#[derive(Clone, Debug, Default)]
pub struct Metadata {
    session_id: SessionId,
    identities: MononokeIdentitySet,
    /// If the identities were proxied, this is the true and original
    /// identities from the request.
    original_identities: Option<MononokeIdentitySet>,
    client_debug: bool,
    /// "true" if client connects from untrusted environment.
    /// We're going to apply restrictions in this case, like rejecting pushes
    /// or admin bypass.
    client_untrusted: bool,
    client_ip: Option<IpAddr>,
    client_hostname: Option<String>,
    revproxy_region: Option<String>,
    raw_encoded_cats: Option<String>,
    client_info: Option<ClientInfo>,
}

impl Metadata {
    pub async fn new(
        session_id: Option<&String>,
        identities: MononokeIdentitySet,
        client_debug: bool,
        client_untrusted: bool,
        client_ip: Option<IpAddr>,
    ) -> Self {
        let session_id: SessionId = match session_id {
            Some(id) => SessionId::from_string(id.to_owned()),
            None => generate_session_id(),
        };

        // Hostname of the client is for non-critical use only. We're doing best-effort lookup here:
        // 1) We're extracting it from identities (which requires no remote calls)
        let client_hostname = if let Some(client_hostname) = identities.hostname() {
            Some(client_hostname.to_string())
        }
        // 2) If it's not there we're trying to look it up via reverse dns with timeout of 1s.
        else if let Some(client_ip) = client_ip {
            timeout(Duration::from_secs(1), Metadata::reverse_lookup(client_ip))
                .await
                .map_err(Error::from)
                .flatten()
                .ok()
        } else {
            None
        };

        Self {
            session_id,
            identities,
            original_identities: None,
            client_debug,
            client_untrusted,
            client_ip,
            client_hostname,
            revproxy_region: None,
            raw_encoded_cats: None,
            client_info: None,
        }
    }

    // Reverse lookups an IP to associated hostname. Trailing dots are stripped
    // to remain compatible with historical logging and common usage of reverse
    // hostnames in other logs (even though trailing dot is technically more correct)
    async fn reverse_lookup(client_ip: IpAddr) -> Result<String> {
        // This parses /etc/resolv.conf on each request. Given that this should be in
        // the page cache and the parsing of the text is very minimal, this shouldn't
        // impact performance much. In case this does lead to performance issues we
        // could start caching this, which for now would be preferred to avoid as this
        // might lead to unexpected behavior if the system configuration changes.
        let resolver = TokioAsyncResolver::tokio_from_system_conf()?;
        resolver
            .reverse_lookup(client_ip)
            .await?
            .iter()
            .next()
            .map(|name| name.to_string().trim_end_matches('.').to_string())
            .ok_or_else(|| anyhow!("failed to do reverse lookup"))
    }

    pub fn add_raw_encoded_cats(&mut self, raw_encoded_cats: String) -> &mut Self {
        self.raw_encoded_cats = Some(raw_encoded_cats);
        self
    }

    pub fn add_revproxy_region(&mut self, revproxy_region: String) -> &mut Self {
        self.revproxy_region = Some(revproxy_region);
        self
    }

    pub fn add_client_info(&mut self, client_info: ClientInfo) -> &mut Self {
        self.client_info = Some(client_info);
        self
    }

    pub fn add_original_identities(&mut self, identities: MononokeIdentitySet) -> &mut Self {
        self.original_identities = Some(identities);
        self
    }

    pub fn update_client_untrusted(&mut self, client_untrusted: bool) -> &mut Self {
        // Be conservative: if client was already untrusted, don't allow to make
        // it trusted
        self.client_untrusted |= client_untrusted;
        self
    }

    pub fn session_id(&self) -> &SessionId {
        &self.session_id
    }

    pub fn identities(&self) -> &MononokeIdentitySet {
        &self.identities
    }

    pub fn original_identities(&self) -> Option<&MononokeIdentitySet> {
        self.original_identities.as_ref()
    }

    pub fn raw_encoded_cats(&self) -> &Option<String> {
        &self.raw_encoded_cats
    }

    pub fn set_identities(mut self, identities: MononokeIdentitySet) -> Self {
        self.identities = identities;
        self
    }

    pub fn revproxy_region(&self) -> &Option<String> {
        &self.revproxy_region
    }

    pub fn client_debug(&self) -> bool {
        self.client_debug
    }

    pub fn client_untrusted(&self) -> bool {
        self.client_untrusted
    }

    pub fn client_ip(&self) -> Option<&IpAddr> {
        self.client_ip.as_ref()
    }

    pub fn set_client_ip(mut self, client_ip: Option<IpAddr>) -> Self {
        self.client_ip = client_ip;
        self
    }

    pub fn client_hostname(&self) -> Option<&str> {
        self.client_hostname.as_deref()
    }

    pub fn set_client_hostname(mut self, client_hostname: Option<String>) -> Self {
        self.client_hostname = client_hostname;
        self
    }

    pub fn unix_name(&self) -> Option<&str> {
        for identity in self.identities() {
            if identity.id_type() == "USER" {
                return Some(identity.id_data());
            }
        }

        None
    }

    pub fn sandcastle_alias(&self) -> Option<&str> {
        self.client_info
            .as_ref()
            .and_then(|ci| ci.fb.sandcastle_alias())
    }

    pub fn clientinfo_u64tag(&self) -> Option<u64> {
        self.client_info.as_ref()?.u64token
    }

    pub fn sandcastle_nonce(&self) -> Option<&str> {
        self.client_info
            .as_ref()
            .and_then(|ci| ci.fb.sandcastle_nonce())
    }

    pub fn clientinfo_tw_job(&self) -> Option<&str> {
        self.client_info.as_ref().and_then(|ci| ci.fb.tw_job())
    }

    pub fn clientinfo_tw_task(&self) -> Option<&str> {
        self.client_info.as_ref().and_then(|ci| ci.fb.tw_task())
    }
}


/*
 * 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(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod oss;

#[cfg(fbcode_build)]
pub use facebook::*;
#[cfg(not(fbcode_build))]
pub use oss::*;


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

pub fn is_client_untrusted<GetHeaderFn>(headers: GetHeaderFn) -> Result<bool>
where
    GetHeaderFn: Fn(&str) -> Result<Option<String>>,
{
    Ok(false)
}


/*
 * 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(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod oss;

#[cfg(fbcode_build)]
pub use facebook::Qps;
#[cfg(not(fbcode_build))]
pub use oss::Qps;


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

use anyhow::Error;
use cached_config::ConfigStore;
use fbinit::FacebookInit;
pub struct Qps {}

impl Qps {
    pub fn new(
        _fb: FacebookInit,
        _top_level_tier: String,
        _config_store: &ConfigStore,
    ) -> Result<Qps, Error> {
        Ok(Self {})
    }

    pub fn bump(&self, _region: &str) -> Result<(), Error> {
        Ok(())
    }
}