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

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

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

use crate::HgMutationEntry;
use crate::HgMutationStore;

const DEFAULT_TTL_SECS: u64 = 3600;

/// Struct representing the cache entry for
/// (repo_id, cs_id) -> Vec<HgMutationEntry> mapping
#[derive(Abomonation, Clone, Debug, Eq, Hash, PartialEq)]
pub struct HgMutationCacheEntry {
    /// The mutation entries that are part of the cache record
    pub mutation_entries: Vec<HgMutationEntry>,
    /// The ID of the repo corresponding to the mutation entries
    pub repo_id: RepositoryId,
    /// The ID of the changeset corresponding to the mutation
    /// entries
    pub changeset_id: HgChangesetId,
}

impl HgMutationCacheEntry {
    fn from_thrift(entry: thrift::HgMutationCacheEntry) -> Result<Self> {
        Ok(Self {
            repo_id: RepositoryId::new(entry.repo_id.0),
            changeset_id: HgChangesetId::new(HgNodeHash::from_thrift(entry.changeset_id)?),
            mutation_entries: entry
                .mutation_entries
                .into_iter()
                .map(HgMutationEntry::from_thrift)
                .collect::<Result<_>>()?,
        })
    }

    fn into_thrift(self) -> thrift::HgMutationCacheEntry {
        thrift::HgMutationCacheEntry {
            repo_id: thrift::RepoId(self.repo_id.id()),
            changeset_id: self.changeset_id.into_nodehash().into_thrift(),
            mutation_entries: self
                .mutation_entries
                .into_iter()
                .map(HgMutationEntry::into_thrift)
                .collect(),
        }
    }

    pub fn into_entries(
        self,
        repo_id: RepositoryId,
        changeset_id: HgChangesetId,
    ) -> Result<(HgChangesetId, Vec<HgMutationEntry>)> {
        if self.repo_id == repo_id && self.changeset_id == changeset_id {
            Ok((self.changeset_id, self.mutation_entries))
        } else {
            Err(anyhow!(
                "Cache returned invalid entry: repo {} & changeset {} returned for query to repo {} and changeset {}",
                self.repo_id,
                self.changeset_id,
                repo_id,
                changeset_id,
            ))
        }
    }

    pub fn from_entries(
        mutation_entries: Vec<HgMutationEntry>,
        repo_id: RepositoryId,
        changeset_id: HgChangesetId,
    ) -> HgMutationCacheEntry {
        HgMutationCacheEntry {
            repo_id,
            changeset_id,
            mutation_entries,
        }
    }
}

#[derive(Clone)]
pub struct CachedHgMutationStore {
    inner_store: Arc<dyn HgMutationStore>,
    cachelib: CachelibHandler<HgMutationCacheEntry>,
    memcache: MemcacheHandler,
    keygen: KeyGen,
}

impl CachedHgMutationStore {
    pub fn new(
        inner_store: Arc<dyn HgMutationStore>,
        cache_handler_factory: CacheHandlerFactory,
    ) -> Self {
        Self {
            inner_store,
            cachelib: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
            keygen: CachedHgMutationStore::create_key_gen(),
        }
    }

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

    fn create_key_gen() -> KeyGen {
        let key_prefix = "scm.mononoke.hg_mutation_store";
        let sitever = if tunables().hg_mutation_store_sitever().unwrap_or_default() > 0 {
            tunables().hg_mutation_store_sitever().unwrap_or_default() as u32
        } else {
            thrift::MC_SITEVER as u32
        };
        KeyGen::new(key_prefix, thrift::MC_CODEVER as u32, sitever)
    }
}

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

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

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

#[async_trait]
impl HgMutationStore for CachedHgMutationStore {
    fn repo_id(&self) -> RepositoryId {
        self.inner_store.repo_id()
    }

    async fn add_entries(
        &self,
        ctx: &CoreContext,
        new_changeset_ids: HashSet<HgChangesetId>,
        entries: Vec<HgMutationEntry>,
    ) -> Result<()> {
        self.inner_store
            .add_entries(ctx, new_changeset_ids, entries)
            .await
    }

    async fn all_predecessors_by_changeset(
        &self,
        ctx: &CoreContext,
        changeset_ids: HashSet<HgChangesetId>,
    ) -> Result<HashMap<HgChangesetId, Vec<HgMutationEntry>>> {
        let cache_request = (ctx, self);
        let repo_id = self.repo_id();

        let mutation_entries_by_changeset =
            get_or_fill_chunked(&cache_request, changeset_ids, CHUNK_SIZE, PARALLEL_CHUNKS)
                .await?
                .into_iter()
                .map(|(cs_id, val)| val.into_entries(repo_id, cs_id))
                .collect::<Result<_>>()?;

        Ok(mutation_entries_by_changeset)
    }
}

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

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

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

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

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

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

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

    fn cache_determinator(&self, _: &HgMutationCacheEntry) -> CacheDisposition {
        let ttl = if tunables()
            .hg_mutation_store_caching_ttl_secs()
            .unwrap_or_default()
            > 0
        {
            tunables()
                .hg_mutation_store_caching_ttl_secs()
                .unwrap_or_default() as u64
        } else {
            DEFAULT_TTL_SECS
        };
        CacheDisposition::Cache(CacheTtl::Ttl(Duration::from_secs(ttl)))
    }

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

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

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

        let res = store
            .inner_store
            .all_predecessors_by_changeset(ctx, keys)
            .await?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|(cs_id, entries)| {
                    (
                        cs_id,
                        HgMutationCacheEntry::from_entries(entries, repo_id, cs_id),
                    )
                })
                .collect(),
        )
    }
}


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

use std::collections::hash_map;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::Read;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use edenapi_types::Extra;
use edenapi_types::HgMutationEntryContent;
use hg_mutation_entry_thrift as thrift;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use types::mutation::MutationEntry;
use types::HgId;

use crate::grouper::Grouper;

/// Record of a Mercurial mutation operation (e.g. amend or rebase).
#[derive(Abomonation, Clone, Debug, Hash, Eq, PartialEq)]
pub struct HgMutationEntry {
    /// The commit that resulted from the mutation operation.
    successor: HgChangesetId,
    /// The commits that were mutated to create the successor.
    ///
    /// There may be multiple predecessors, e.g. if the commits were folded.
    predecessors: Vec<HgChangesetId>,
    /// Other commits that were created by the mutation operation splitting the predecessors.
    ///
    /// Where a commit is split into two or more commits, the successor will be the final commit,
    /// and this list will contain the other commits.
    split: Vec<HgChangesetId>,
    /// The name of the operation.
    op: String,
    /// The user who performed the mutation operation.  This may differ from the commit author.
    user: String,
    /// The timestamp of the mutation operation.  This may differ from the commit time.
    timestamp: i64,
    /// The timezone offset of the mutation operation.  This may differ from the commit time.
    timezone: i32,
    /// Extra information about this mutation operation.
    extra: Vec<(String, String)>,
}

impl HgMutationEntry {
    pub fn new(
        successor: HgChangesetId,
        predecessors: Vec<HgChangesetId>,
        split: Vec<HgChangesetId>,
        op: String,
        user: String,
        timestamp: i64,
        timezone: i32,
        extra: Vec<(String, String)>,
    ) -> Self {
        Self {
            successor,
            predecessors,
            split,
            op,
            user,
            timestamp,
            timezone,
            extra,
        }
    }

    pub fn deserialize(r: &mut dyn Read) -> Result<Self> {
        HgMutationEntry::try_from(MutationEntry::deserialize(r)?)
    }

    pub fn successor(&self) -> &HgChangesetId {
        &self.successor
    }

    pub fn predecessors(&self) -> &[HgChangesetId] {
        self.predecessors.as_slice()
    }

    pub fn split(&self) -> &[HgChangesetId] {
        self.split.as_slice()
    }

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

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

    pub fn timestamp(&self) -> i64 {
        self.timestamp
    }

    pub fn timezone(&self) -> i32 {
        self.timezone
    }

    pub fn extra(&self) -> &[(String, String)] {
        self.extra.as_slice()
    }

    /// Add the next predecessor to the entry.
    pub(crate) fn add_predecessor(&mut self, index: u64, pred: HgChangesetId) -> Result<()> {
        // This method is used when progressively loading entries from the
        // database. Each predecessor is received in a separate row, and we may
        // receive each predecessor multiple times.  They should always be
        // received in order, so only extend the list of predecessors if the
        // index of the new predecessor matches the expected index of the next
        // predecessor.
        let expected_index = self.predecessors.len() as u64;
        if index > expected_index {
            // We have received a predecessor past the end of the current
            // predecessor list.  This probably means the predecessor table is
            // missing a row.
            return Err(anyhow!(
                "Unexpected out-of-order predecessor {}, expected index {}",
                pred,
                expected_index
            ));
        }
        if index == expected_index {
            self.predecessors.push(pred);
        }
        Ok(())
    }

    /// Add the next split to the entry.
    pub(crate) fn add_split(&mut self, index: u64, split: HgChangesetId) -> Result<()> {
        // This method is used when progressively loading entries from the
        // database. Each split successor is received in a separate row, and we
        // may receive each split successor multiple times.  They should always
        // be received in order, so only extend the list of split successors if
        // the index of the new split successor matches the expected index of
        // the next split successor.
        let expected_index = self.split.len() as u64;
        if index > expected_index {
            // We have received a split successor past the end of the current
            // split successor list.  This probably means the split table is
            // missing a row.
            return Err(anyhow!(
                "Unexpected out-of-order split successor {}, expected index {}",
                split,
                expected_index
            ));
        }
        if index == expected_index {
            self.split.push(split);
        }
        Ok(())
    }

    pub(crate) fn from_thrift(entry: thrift::HgMutationEntry) -> Result<Self> {
        let preds = entry
            .predecessors
            .into_iter()
            .map(HgChangesetId::from_thrift)
            .collect::<Result<Vec<HgChangesetId>>>()?;
        let split = entry
            .split
            .into_iter()
            .map(HgChangesetId::from_thrift)
            .collect::<Result<Vec<HgChangesetId>>>()?;
        let extra = entry
            .extra
            .into_iter()
            .map(|e| (e.key, e.value))
            .collect::<Vec<(String, String)>>();
        Ok(HgMutationEntry::new(
            HgChangesetId::from_thrift(entry.successor)?,
            preds,
            split,
            entry.op,
            entry.user,
            entry.timestamp,
            entry.timezone,
            extra,
        ))
    }

    pub(crate) fn into_thrift(self) -> thrift::HgMutationEntry {
        thrift::HgMutationEntry {
            successor: HgChangesetId::into_thrift(self.successor),
            predecessors: self
                .predecessors
                .into_iter()
                .map(HgChangesetId::into_thrift)
                .collect(),
            split: self
                .split
                .into_iter()
                .map(HgChangesetId::into_thrift)
                .collect(),
            op: self.op,
            user: self.user,
            timestamp: self.timestamp,
            timezone: self.timezone,
            extra: self
                .extra
                .into_iter()
                .map(|(key, value)| thrift::ExtraProperty { key, value })
                .collect(),
        }
    }
}

// Conversion from client mutation entry
impl TryFrom<MutationEntry> for HgMutationEntry {
    type Error = Error;

    fn try_from(entry: MutationEntry) -> Result<HgMutationEntry> {
        let entry = HgMutationEntry {
            successor: HgChangesetId::new(HgNodeHash::from(entry.succ)),
            predecessors: entry
                .preds
                .into_iter()
                .map(HgNodeHash::from)
                .map(HgChangesetId::new)
                .collect(),
            split: entry
                .split
                .into_iter()
                .map(HgNodeHash::from)
                .map(HgChangesetId::new)
                .collect(),
            op: entry.op,
            user: entry.user,
            timestamp: entry.time,
            timezone: entry.tz,
            extra: entry
                .extra
                .into_iter()
                .map(|(key, value)| -> Result<(String, String), Error> {
                    Ok((
                        String::from_utf8(key.into())?,
                        String::from_utf8(value.into())?,
                    ))
                })
                .collect::<Result<_>>()?,
        };
        Ok(entry)
    }
}

impl From<HgMutationEntry> for MutationEntry {
    fn from(m: HgMutationEntry) -> MutationEntry {
        MutationEntry {
            succ: m.successor.into_nodehash().into(),
            preds: m
                .predecessors
                .into_iter()
                .map(HgChangesetId::into_nodehash)
                .map(HgId::from)
                .collect(),
            split: m
                .split
                .into_iter()
                .map(HgChangesetId::into_nodehash)
                .map(HgId::from)
                .collect(),
            op: m.op,
            user: m.user,
            time: m.timestamp,
            tz: m.timezone,
            extra: m
                .extra
                .into_iter()
                .map(|(key, value)| {
                    (
                        key.into_bytes().into_boxed_slice(),
                        value.into_bytes().into_boxed_slice(),
                    )
                })
                .collect(),
        }
    }
}

impl TryFrom<HgMutationEntryContent> for HgMutationEntry {
    type Error = Error;
    fn try_from(mutation: HgMutationEntryContent) -> Result<Self> {
        let successor = HgChangesetId::new(HgNodeHash::from(mutation.successor));
        let predecessors = mutation
            .predecessors
            .into_iter()
            .map(HgNodeHash::from)
            .map(HgChangesetId::new)
            .collect::<Vec<_>>();
        let split = mutation
            .split
            .into_iter()
            .map(HgNodeHash::from)
            .map(HgChangesetId::new)
            .collect::<Vec<_>>();
        let op = mutation.op;
        let user = std::str::from_utf8(&mutation.user)?.to_string();
        let timestamp = mutation.time;
        let timezone = mutation.tz;
        let exta = mutation
            .extras
            .into_iter()
            .map(|extra| {
                Ok((
                    std::str::from_utf8(&extra.key)?.to_string(),
                    std::str::from_utf8(&extra.value)?.to_string(),
                ))
            })
            .collect::<Result<_, Error>>()?;

        Ok(HgMutationEntry::new(
            successor,
            predecessors,
            split,
            op,
            user,
            timestamp,
            timezone,
            exta,
        ))
    }
}
impl From<HgMutationEntry> for HgMutationEntryContent {
    fn from(mutation: HgMutationEntry) -> Self {
        let successor = mutation.successor.into();
        let predecessors = mutation
            .predecessors
            .into_iter()
            .map(Into::into)
            .collect::<Vec<_>>();
        let split = mutation
            .split
            .into_iter()
            .map(Into::into)
            .collect::<Vec<_>>();
        let op = mutation.op;
        let user = mutation.user.into_bytes();
        let time = mutation.timestamp;
        let tz = mutation.timezone;
        let extras = mutation
            .extra
            .into_iter()
            .map(|(key, value)| Extra {
                key: key.into_bytes(),
                value: value.into_bytes(),
            })
            .collect();

        Self {
            successor,
            predecessors,
            split,
            op,
            user,
            time,
            tz,
            extras,
        }
    }
}

pub(crate) struct HgMutationEntrySet {
    // The loaded entries, indexed by successor.
    pub(crate) entries: HashMap<HgChangesetId, HgMutationEntry>,

    // The known primordial changeset ID for any changeset ID.
    pub(crate) changeset_primordials: HashMap<HgChangesetId, HgChangesetId>,
}

/// Result of a request to add new entries to an entry set.
pub(crate) struct HgMutationEntrySetAdded {
    /// The keys for entries that were successfully added.
    pub(crate) added: Vec<HgChangesetId>,

    /// The changeset IDs of changesets that could not be added because their
    /// primordial changesets are not known.
    pub(crate) missing_primordials: Vec<HgChangesetId>,

    /// The remaining entries that were not added.
    pub(crate) remaining_entries: HashMap<HgChangesetId, HgMutationEntry>,
}

impl HgMutationEntrySet {
    pub(crate) fn new() -> Self {
        Self {
            entries: HashMap::new(),
            changeset_primordials: HashMap::new(),
        }
    }

    /// Add new entries for the given ids to the entry set.
    ///
    /// New entries associated with `new_ids` are moved from `new_entries` into
    /// the entry set, provided that their primordial IDs are already known.
    ///
    /// Returns the changeset IDs for the entries that were added and the
    /// changeset IDs which could not be added because their primordial IDs are
    /// not known.
    pub(crate) fn add_entries<'a>(
        &mut self,
        mut new_entries: HashMap<HgChangesetId, HgMutationEntry>,
        new_ids: impl IntoIterator<Item = &'a HgChangesetId>,
    ) -> Result<HgMutationEntrySetAdded> {
        let mut added = Vec::new();
        let mut missing_primordials = Vec::new();

        for changeset_id in new_ids {
            // Check if we already have an entry for this changeset.
            if self.entries.contains_key(changeset_id) {
                continue;
            }

            // This changeset does not have a stored entry for which it is the
            // successor.
            match new_entries.entry(*changeset_id) {
                hash_map::Entry::Vacant(_) => {
                    // This changeset is a new primordial changeset.
                    added.push(*changeset_id);
                }
                hash_map::Entry::Occupied(entry) => {
                    // This changeset has a new entry to store.  See if all its
                    // predecessors' primordials are known.
                    for predecessor_id in entry.get().predecessors().iter() {
                        if !self.changeset_primordials.contains_key(predecessor_id) {
                            missing_primordials.push(*predecessor_id);
                        }
                    }
                    // The first predecessor's primordial should be
                    // propagated to this changeset.
                    let predecessor_id =
                        entry.get().predecessors().iter().next().ok_or_else(|| {
                            anyhow!(
                                "Mutation entry for {} has no predecessors",
                                entry.get().successor()
                            )
                        })?;
                    match self.changeset_primordials.get(predecessor_id) {
                        Some(&primordial_id) => {
                            // The entry's first predecessor's primordial is known.
                            // Move the new entry over to the entry set, and
                            // copy the first predecessor's primordial.
                            self.entries.insert(*changeset_id, entry.remove());
                            self.changeset_primordials
                                .insert(*changeset_id, primordial_id);
                            added.push(*changeset_id);
                        }
                        None => {
                            // The entry's first predecessor's primordial is not
                            // known yet.  We need to include this changeset in
                            // the primordial search as well, so that it will be
                            // filled in once it is found. Add this changeset to
                            // the missing set.
                            missing_primordials.push(*changeset_id)
                        }
                    }
                }
            }
        }

        Ok(HgMutationEntrySetAdded {
            added,
            missing_primordials,
            remaining_entries: new_entries,
        })
    }

    /// Add new entries for the given ids and their predecessors to the entry
    /// set.
    ///
    /// New entries associated with `new_ids` are moved from `new_entries` into
    /// the entry set.
    ///
    /// For these new entries, the primordial changeset IDs are filled in by
    /// searching all predecessors for either a known primordial, or a new
    /// primordial.
    ///
    /// If the set of entries contains cycles, then it may not be possible to
    /// determine an appropriate primordial commit.  In which case, the
    /// entries that form a cycle will be ignored.
    ///
    /// Returns all changeset IDs that were added.
    pub(crate) fn add_entries_and_find_primordials<'a>(
        &mut self,
        mut new_entries: HashMap<HgChangesetId, HgMutationEntry>,
        new_ids: impl IntoIterator<Item = &'a HgChangesetId>,
    ) -> Result<Vec<HgChangesetId>> {
        // The changesets IDs that were added
        let mut added = Vec::new();

        // We will allocate primordials by seeking back to the first commit with
        // a known primordial, or the primordial commit itself.
        //
        // Commits that are yet to be processed are candidates.
        let mut candidates: Vec<_> = new_ids.into_iter().copied().collect();

        // Commits that are queued to be processed or have been
        // processed (to break cycles).
        let mut seen: HashSet<_> = candidates.iter().copied().collect();

        // A Grouper to group commits together into primordial groups.
        let mut grouper = Grouper::new();

        // Look at each candidate.  If it is primoridial, or we know
        // what its primordial should be, then we are done.  Otherwise,
        // expand it to its predecessors, which are the new candidates.
        while let Some(candidate) = candidates.pop() {
            if let Some(primordial_id) = self.changeset_primordials.get(&candidate) {
                // We have reached a changeset with a known primordial
                grouper.set_primordial(candidate, *primordial_id);
            } else if let Some(entry) = new_entries.get(&candidate) {
                // This is not the primordial commit, we must look at its
                // predecessors.
                let predecessors = entry.predecessors();
                if let Some(first) = predecessors.first() {
                    // Merge this candidate's group with the group of its
                    // first predecessor: they will have the same primordial.
                    grouper.merge(candidate, *first);
                } else {
                    return Err(anyhow!(
                        "Mutation entry for {} has no predecessors",
                        entry.successor()
                    ));
                }
                for &predecessor in predecessors.iter() {
                    if seen.insert(predecessor) {
                        candidates.push(predecessor);
                    }
                }
            } else {
                // We have reached a new primordial changeset.
                grouper.set_primordial(candidate, candidate);
                added.push(candidate);
            };
        }

        // Apply calculated primordials to their groups, and work out
        // which entries should be moved into the store.
        let mut move_entries = Vec::new();
        for (primordial, members) in grouper.groups() {
            if let Some(primordial) = primordial {
                // Apply this primordial changeset to all of the members
                // of this group.
                for changeset_id in members {
                    if self
                        .changeset_primordials
                        .insert(changeset_id, primordial)
                        .is_none()
                    {
                        move_entries.push(changeset_id);
                    }
                }
            }
        }

        // Move valid entries into the store.
        for changeset_id in move_entries.into_iter() {
            if let Some(new_entry) = new_entries.remove(&changeset_id) {
                if new_entry
                    .predecessors()
                    .iter()
                    .all(|predecessor| self.changeset_primordials.contains_key(predecessor))
                {
                    // We have found primordials for all predecessors of this
                    // entry, so we can add it.
                    self.entries.insert(changeset_id, new_entry);
                    added.push(changeset_id);
                }
            }
        }

        Ok(added)
    }

    /// Extracts all entries for predecessors of the given changeset ids.
    pub(crate) fn into_all_predecessors_by_changeset(
        self,
        successor_ids: HashSet<HgChangesetId>,
    ) -> HashMap<HgChangesetId, Vec<HgMutationEntry>> {
        let mut mutation_history = HashMap::new();
        for successor_id in successor_ids {
            let mut entries = Vec::new();
            let mut processed = HashSet::new();
            let mut changeset_ids = vec![successor_id];
            while let Some(changeset_id) = changeset_ids.pop() {
                // See if we have an entry for this changeset_id.
                if let Some(entry) = self.entries.get(&changeset_id).cloned() {
                    // Add all of this entry's predecessors to the queue of
                    // additional changesets we will need to process.  Push
                    // predecessors in reverse order so that we process them in
                    // forwards order.
                    for predecessor_id in entry.predecessors().iter().rev() {
                        // Only enqueue the predecessor if:
                        // 1. There is an entry for it
                        if self.entries.contains_key(predecessor_id)
                            // 2. AND we haven't already processed this predecessor
                            && !processed.contains(predecessor_id)
                            // 3. AND its not the same as the successor
                            && successor_id != *predecessor_id
                        {
                            changeset_ids.push(predecessor_id.clone());
                        }
                    }
                    entries.push(entry);
                }
                processed.insert(changeset_id);
            }
            mutation_history.insert(successor_id, entries);
        }
        mutation_history
    }
}


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

//! Utility for grouping commits with their primordials.

use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;

/// A group of commits, and their assigned primordial commit (if known).
#[derive(Debug)]
pub(crate) struct Group<T> {
    /// The assigned primordial commit.
    primordial: Option<T>,

    /// The commits in the group.
    members: HashSet<T>,
}

/// A utility struct to aid grouping of commits and assignment of primordials.
///
/// This struct collects groups of commits together that will be assigned to
/// the same primordial commit.
///
/// Groups can be merged once it is known that they will end up with the same
/// primordial assignment.
#[derive(Debug)]
pub(crate) struct Grouper<T> {
    /// The groups of commits.
    groups: Vec<Group<T>>,

    /// A map from member commit to the group they are currently part of
    /// (by index into the `groups` vector).
    member_groups: HashMap<T, usize>,
}

impl<T: Copy + Hash + Eq> Group<T> {
    fn new() -> Group<T> {
        Group {
            primordial: None,
            members: HashSet::new(),
        }
    }

    fn with_member(mut self, member: T) -> Group<T> {
        self.members.insert(member);
        self
    }

    fn with_primordial(mut self, primordial: T) -> Group<T> {
        self.primordial = Some(primordial);
        self
    }
}

impl<T: Copy + Hash + Eq> Grouper<T> {
    /// Create a new grouper to group sets of commits and assign primordials.
    pub(crate) fn new() -> Grouper<T> {
        Grouper {
            groups: Vec::new(),
            member_groups: HashMap::new(),
        }
    }

    /// Set the primordial commit for the group containing `member`.  If
    /// no such group exists then a group with that member will be created
    /// and its primordial assigned.
    pub(crate) fn set_primordial(&mut self, member: T, primordial: T) {
        match self.member_groups.get(&member) {
            Some(&group) => self.groups[group].primordial = Some(primordial),
            None => {
                self.member_groups.insert(member, self.groups.len());
                self.groups
                    .push(Group::new().with_member(member).with_primordial(primordial));
            }
        }
    }

    /// Merge the group for `other` into the group for `member`.  If
    /// either don't exist, then they are added to the other's group.
    /// If neither exist, a new group is created.
    pub(crate) fn merge(&mut self, member: T, other: T) {
        match (
            self.member_groups.get(&member),
            self.member_groups.get(&other),
        ) {
            (Some(&group1), Some(&group2)) => {
                if group1 != group2 {
                    if self.groups[group1].members.len() > self.groups[group2].members.len() {
                        self.move_group(group1, group2);
                    } else {
                        self.move_group(group2, group1);
                    }
                }
            }
            (Some(&group), None) => {
                self.groups[group].members.insert(other);
                self.member_groups.insert(other, group);
            }
            (None, Some(&group)) => {
                self.groups[group].members.insert(member);
                self.member_groups.insert(member, group);
            }
            (None, None) => {
                self.member_groups.insert(member, self.groups.len());
                self.member_groups.insert(other, self.groups.len());
                self.groups
                    .push(Group::new().with_member(member).with_member(other));
            }
        }
    }

    /// Move all members of `src` to `dest`.
    fn move_group(&mut self, dest: usize, src: usize) {
        let members = std::mem::take(&mut self.groups[src].members);
        for member in members {
            self.member_groups.insert(member, dest);
            self.groups[dest].members.insert(member);
        }
        if let Some(primordial) = self.groups[src].primordial.take() {
            self.groups[dest].primordial = Some(primordial);
        }
    }

    /// Iterate over all groups, returning their primordial (if known) and members.
    pub(crate) fn groups(self) -> impl Iterator<Item = (Option<T>, HashSet<T>)> {
        self.groups
            .into_iter()
            .map(|group| (group.primordial, group.members))
    }
}


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

//! A store for Mercurial mutation information.
//!
//! This allows mutation information (i.e., which commits where amended or rebased) to be exchanged
//! between Mercurial clients via Mononoke.
//!
//! See <https://fburl.com/m2y3nr5c> for more details.

use std::collections::HashMap;
use std::collections::HashSet;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mercurial_types::HgChangesetId;
use mononoke_types::RepositoryId;

mod builder;
mod caching;
mod entry;
mod grouper;
mod store;

pub use crate::builder::SqlHgMutationStoreBuilder;
pub use crate::caching::CachedHgMutationStore;
pub use crate::entry::HgMutationEntry;
pub use crate::store::SqlHgMutationStore;

#[facet::facet]
#[async_trait]
pub trait HgMutationStore: Send + Sync {
    /// Add new entries to the mutation store.
    ///
    /// Adds mutation information for `new_changeset_ids` using the given `entries`.
    async fn add_entries(
        &self,
        ctx: &CoreContext,
        new_changeset_ids: HashSet<HgChangesetId>,
        entries: Vec<HgMutationEntry>,
    ) -> Result<()>;

    /// Get all predecessor information for the given changeset ids.
    ///
    /// Returns all entries that describe the mutation history of the commits.
    async fn all_predecessors(
        &self,
        ctx: &CoreContext,
        changeset_ids: HashSet<HgChangesetId>,
    ) -> Result<Vec<HgMutationEntry>> {
        let entries_by_changeset = self
            .all_predecessors_by_changeset(ctx, changeset_ids)
            .await?;
        Ok(entries_by_changeset
            .into_iter()
            .flat_map(|(_, entries)| entries)
            // Collect into a hashset since the preds for different
            // successors might overlap due to fold and split.
            .collect::<HashSet<_>>()
            .into_iter()
            .collect())
    }

    /// Get all predecessor information for the given changeset id, keyed by
    /// the successor changeset id.
    ///
    /// Returns all entries that describe the mutation history of the commits.
    /// keyed by the successor changeset ids.
    async fn all_predecessors_by_changeset(
        &self,
        ctx: &CoreContext,
        changeset_ids: HashSet<HgChangesetId>,
    ) -> Result<HashMap<HgChangesetId, Vec<HgMutationEntry>>>;

    /// Get the repository for which the mutation history is being added
    /// and retrieved.
    fn repo_id(&self) -> RepositoryId;
}


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

use std::collections::hash_map;
use std::collections::HashMap;
use std::collections::HashSet;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use futures::future;
use futures::future::FutureExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use itertools::Itertools;
use mercurial_types::HgChangesetId;
use mononoke_types::RepositoryId;
use slog::debug;
use sql::Connection;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;

use crate::entry::HgMutationEntry;
use crate::entry::HgMutationEntrySet;
use crate::entry::HgMutationEntrySetAdded;
use crate::HgMutationStore;

/// To avoid overloading the database with too many changesets in a single
/// select, we chunk selects to this size.
const SELECT_CHUNK_SIZE: usize = 100;

pub struct SqlHgMutationStore {
    repo_id: RepositoryId,
    connections: SqlConnections,
    mutation_chain_limit: usize,
}

/// Convenience alias for the type of each row returned from the entries queries.
type EntryRow = (
    HgChangesetId, // Successor
    HgChangesetId, // Primordial
    u64,           // Predecessor Count
    u64,           // Predecessor Index
    HgChangesetId, // Predecessor
    HgChangesetId, // Predecessor's Primordial
    u64,           // Split Count
    String,        // Op
    String,        // User
    i64,           // Timestamp
    i32,           // Time Zone
    String,        // Extras
);

impl SqlHgMutationStore {
    pub fn new(
        repo_id: RepositoryId,
        connections: SqlConnections,
        mutation_chain_limit: usize,
    ) -> Self {
        Self {
            repo_id,
            connections,
            mutation_chain_limit,
        }
    }

    /// Store new mutation entries in the mutation store.
    ///
    /// Stores the mutation entries for the changesets in `changeset_ids` into
    /// the store.
    ///
    /// Mutation entries in `entry_set` for which these changesets are the
    /// successor will be written into the store.  Other changesets (which are
    /// expected to be primordial changesets) will just be recorded as having
    /// been added to the store.
    async fn store<'a>(
        &self,
        ctx: &CoreContext,
        entry_set: &HgMutationEntrySet,
        changeset_ids: impl IntoIterator<Item = &'a HgChangesetId>,
    ) -> Result<()> {
        let txn = self
            .connections
            .write_connection
            .start_transaction()
            .await?;

        let mut db_csets = Vec::new();
        let mut db_entries = Vec::new();
        let mut db_preds = Vec::new();
        let mut db_splits = Vec::new();
        for changeset_id in changeset_ids.into_iter() {
            db_csets.push((&self.repo_id, changeset_id));
            let entry = match entry_set.entries.get(changeset_id) {
                Some(entry) => entry,
                None => continue,
            };
            let primordial = entry_set
                .changeset_primordials
                .get(entry.successor())
                .ok_or_else(|| {
                    anyhow!(
                        "failed to determine primordial commit for successor {}",
                        entry.successor()
                    )
                })?;
            db_entries.push((
                &self.repo_id,
                entry.successor(),
                primordial,
                entry.predecessors().len() as u64,
                entry.split().len() as u64,
                entry.op(),
                entry.user(),
                entry.timestamp(),
                entry.timezone(),
                serde_json::to_string(entry.extra())?,
            ));
            for (index, predecessor) in entry.predecessors().iter().enumerate() {
                let primordial = entry_set
                    .changeset_primordials
                    .get(predecessor)
                    .ok_or_else(|| {
                        anyhow!(
                            "failed to determine primordial commit for predecessor {}",
                            predecessor
                        )
                    })?;
                db_preds.push((
                    &self.repo_id,
                    entry.successor(),
                    index as u64,
                    predecessor,
                    primordial,
                ));
            }
            for (index, split_successor) in entry.split().iter().enumerate() {
                db_splits.push((
                    &self.repo_id,
                    entry.successor(),
                    index as u64,
                    split_successor,
                ));
            }
        }

        // Sort the entries so that they are inserted in a stable
        // order.  This prevents deadlocks where one transaction is
        // adding commits "A, B, C" and gains the lock on A while
        // another transaction is adding "C, D, A" and gains the lock
        // on C.
        db_csets.sort_unstable();
        db_entries.sort_unstable();
        db_preds.sort_unstable();
        db_splits.sort_unstable();

        // Convert the numbers to references.
        let ref_db_entries: Vec<_> = db_entries
            .iter()
            .map(
                |&(
                    repo_id,
                    successor,
                    primordial,
                    ref extra_pred_count,
                    ref split_count,
                    op,
                    user,
                    ref time,
                    ref tz,
                    ref extra,
                )| {
                    (
                        repo_id,
                        successor,
                        primordial,
                        extra_pred_count,
                        split_count,
                        op,
                        user,
                        time,
                        tz,
                        extra.as_str(),
                    )
                },
            )
            .collect();
        let ref_db_preds: Vec<_> = db_preds
            .iter()
            .map(
                |&(repo_id, successor, ref index, predecessor, primordial)| {
                    (repo_id, successor, index, predecessor, primordial)
                },
            )
            .collect();
        let ref_db_splits: Vec<_> = db_splits
            .iter()
            .map(|&(repo_id, successor, ref index, split_successor)| {
                (repo_id, successor, index, split_successor)
            })
            .collect();

        let (txn, _) = AddChangesets::query_with_transaction(txn, db_csets.as_slice()).await?;
        let (txn, _) = AddEntries::query_with_transaction(txn, ref_db_entries.as_slice()).await?;
        let (txn, _) = AddPreds::query_with_transaction(txn, ref_db_preds.as_slice()).await?;
        let (txn, _) = AddSplits::query_with_transaction(txn, ref_db_splits.as_slice()).await?;

        txn.commit().await?;

        debug!(
            ctx.logger(),
            "Mutation store added {} entries for {} changesets",
            db_entries.len(),
            db_csets.len()
        );
        ctx.perf_counters().add_to_counter(
            PerfCounterType::HgMutationStoreNumAdded,
            db_entries.len() as i64,
        );

        Ok(())
    }

    /// Get an appropriate read connection for querying mutation information
    /// about some changesets.
    ///
    /// Checks if mutation information for all of the given changesets have been
    /// replicated to the read-only replica.  If so, returns a read connection
    /// to the replica.  If not all mutation information has been replicated, or
    /// if no mutation information for the changesets exists, returns a read
    /// connection to the master.
    ///
    /// Note that this only guarantees that mutation information for this
    /// commit and all of its predecessors are present.  For checking the
    /// successors of a commit, the replica may still lag behind.
    async fn read_connection_for_changesets(
        &self,
        changeset_ids: &HashSet<HgChangesetId>,
    ) -> Result<&Connection> {
        // Check if the replica is up-to-date with respect to all changesets we
        // are interested in.
        let changeset_ids: Vec<_> = changeset_ids.iter().collect();
        if changeset_ids.is_empty() {
            // There are no interesting changesets, so just use the replica.
            return Ok(&self.connections.read_connection);
        }
        let count = CountChangesets::query(
            &self.connections.read_connection,
            &self.repo_id,
            changeset_ids.as_slice(),
        )
        .await?;
        if let Some((count,)) = count.into_iter().next() {
            if count as usize == changeset_ids.len() {
                // The replica knows all of the changesets, so use it.
                return Ok(&self.connections.read_connection);
            }
        }
        // The replica doesn't know all of the changesets, use the connection to
        // the master.
        Ok(&self.connections.read_master_connection)
    }

    /// Collect entries from the database into an entry set.
    async fn collect_entries<I>(
        &self,
        connection: &Connection,
        entry_set: &mut HgMutationEntrySet,
        rows: I,
    ) -> Result<()>
    where
        I: IntoIterator<Item = EntryRow>,
    {
        let mut to_fetch_split = Vec::new();
        for (
            successor,
            primordial,
            pred_count,
            p_seq,
            p_predecessor,
            p_primordial,
            split_count,
            op,
            user,
            timestamp,
            tz,
            extra,
        ) in rows
        {
            if !entry_set.changeset_primordials.contains_key(&successor) {
                entry_set
                    .changeset_primordials
                    .insert(successor.clone(), primordial.clone());
            }
            if !entry_set.changeset_primordials.contains_key(&p_predecessor) {
                entry_set
                    .changeset_primordials
                    .insert(p_predecessor.clone(), p_primordial.clone());
            }
            let entry = match entry_set.entries.entry(successor) {
                hash_map::Entry::Occupied(entry) => entry.into_mut(),
                hash_map::Entry::Vacant(entry) => {
                    let successor = entry.key().clone();
                    if split_count > 0 {
                        to_fetch_split.push(successor.clone());
                    }
                    entry.insert(HgMutationEntry::new(
                        successor,
                        Vec::with_capacity(pred_count as usize),
                        Vec::with_capacity(split_count as usize),
                        op,
                        user,
                        timestamp,
                        tz,
                        serde_json::from_str(&extra)?,
                    ))
                }
            };
            entry
                .add_predecessor(p_seq, p_predecessor)
                .with_context(|| {
                    format!("Failed to collect predecessor {} for {}", p_seq, successor)
                })?;
        }
        if !to_fetch_split.is_empty() {
            let rows = SelectSplitsBySuccessor::query(
                connection,
                &self.repo_id,
                to_fetch_split.as_slice(),
            )
            .await?;
            for (successor, seq, split_successor) in rows {
                if let Some(entry) = entry_set.entries.get_mut(&successor) {
                    entry.add_split(seq, split_successor).with_context(|| {
                        format!(
                            "Failed to collect split successor {} for {}",
                            seq, successor
                        )
                    })?;
                }
            }
        }
        Ok(())
    }

    /// Fetch all mutation information where the given changesets are the
    /// successor and add it to the entry set.
    async fn fetch_by_successor(
        &self,
        connection: &Connection,
        entry_set: &mut HgMutationEntrySet,
        changesets: &HashSet<HgChangesetId>,
    ) -> Result<()> {
        let chunks = changesets
            .iter()
            .cloned()
            .filter(|hg_cs_id| !entry_set.entries.contains_key(hg_cs_id))
            .chunks(SELECT_CHUNK_SIZE)
            .into_iter()
            .map(|chunk| chunk.collect::<Vec<_>>())
            .collect::<Vec<_>>();

        let chunk_rows = stream::iter(chunks.into_iter().map(move |chunk| async move {
            SelectBySuccessor::query(connection, &self.repo_id, chunk.as_slice())
                .await
                .with_context(|| format!("Error fetching successors: {:?}", chunk))
        }))
        .buffered(10)
        .try_collect::<Vec<_>>()
        .await?;

        self.collect_entries(connection, entry_set, chunk_rows.into_iter().flatten())
            .await?;
        Ok(())
    }

    /// Fetch all predecessor entries for the entries in the entry set.
    async fn fetch_all_predecessors(
        &self,
        connection: &Connection,
        entry_set: &mut HgMutationEntrySet,
    ) -> Result<()> {
        let mut fetched_primordials = HashSet::new();
        loop {
            let mut to_fetch = HashSet::new();
            for primordial in entry_set.changeset_primordials.values() {
                if !fetched_primordials.contains(primordial) {
                    to_fetch.insert(primordial.clone());
                }
            }
            if to_fetch.is_empty() {
                break;
            }
            fetched_primordials.extend(to_fetch.iter().copied());

            let obtained_rows = stream::iter(to_fetch.into_iter().map(|changeset| async move {
                let (mut primordial_rows, successor_rows) = future::try_join(
                    SelectByPrimordial::query(
                        connection,
                        &self.repo_id,
                        &self.mutation_chain_limit,
                        &[changeset],
                    )
                    .map(|r| {
                        r.with_context(|| format!("Error fetching primordials: {:?}", &changeset))
                    }),
                    SelectBySuccessor::query(connection, &self.repo_id, &[changeset]).map(|r| {
                        r.with_context(|| format!("Error fetching successors: {:?}", &changeset))
                    }),
                )
                .await?;
                // If we reach the limit and ended up cutting some predecessors of a mutation
                // then remove the mutation all together.
                if primordial_rows.len() == self.mutation_chain_limit {
                    if let Some(last) = primordial_rows.iter().last() {
                        let (pred_count, seq) = (last.2, last.3);
                        if seq + 1 != pred_count {
                            // we don't have the whole of the last entry, remove it
                            let last_count = primordial_rows
                                .iter()
                                .rev()
                                .take_while(|row| row.0 == last.0) //0 -> successor
                                .count();
                            primordial_rows.truncate(primordial_rows.len() - last_count)
                        }
                    }
                }
                Ok::<_, Error>(
                    primordial_rows
                        .into_iter()
                        .chain(successor_rows.into_iter())
                        .collect::<Vec<_>>(),
                )
            }))
            .buffered(100)
            .try_collect::<Vec<_>>()
            .await?;

            self.collect_entries(connection, entry_set, obtained_rows.into_iter().flatten())
                .await?;
        }
        Ok(())
    }
}

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

    async fn add_entries(
        &self,
        ctx: &CoreContext,
        new_changeset_ids: HashSet<HgChangesetId>,
        entries: Vec<HgMutationEntry>,
    ) -> Result<()> {
        if new_changeset_ids.is_empty() {
            // Nothing to do.
            return Ok(());
        }

        let mut entry_set = HgMutationEntrySet::new();
        let entries: HashMap<_, _> = entries
            .into_iter()
            .filter(|entry| {
                // Filter out entries that are self-referential (these are
                // revive obsmarkers incorrectly converted to mutation entries).
                let successor = entry.successor();
                !entry
                    .predecessors()
                    .iter()
                    .any(|predecessor| predecessor == successor)
            })
            .map(|entry| (*entry.successor(), entry))
            .collect();

        // First fetch all the new changesets and their immediate predecessors.
        // Most of the time this will be enough, as we can learn the primordial
        // changeset ID from the immediate predecessor.
        let mut changeset_ids = HashSet::with_capacity(new_changeset_ids.len() * 2);
        for changeset_id in &new_changeset_ids {
            changeset_ids.insert(changeset_id.clone());
            if let Some(entry) = entries.get(changeset_id) {
                changeset_ids.extend(entry.predecessors().iter().cloned());
            }
        }
        self.fetch_by_successor(
            &self.connections.read_master_connection,
            &mut entry_set,
            &changeset_ids,
        )
        .await?;

        // Attempt to add these new entries to the entry set, and find out which
        // changesets have missing primordial IDs.
        let HgMutationEntrySetAdded {
            mut added,
            missing_primordials,
            remaining_entries,
        } = entry_set.add_entries(entries, &new_changeset_ids)?;

        debug!(
            ctx.logger(),
            "Mutation store fast path found primordial changesets for {} changesets ({} remaining)",
            added.len(),
            missing_primordials.len(),
        );

        // If any entries were missing their primoridal IDs, we will need to
        // fetch them and all of their predecessors in order to find which
        // predecessor entries are missing and what their primordial commits
        // are.
        if !missing_primordials.is_empty() {
            // Find all predecessors of the entries that are missing in the
            // store by iterating over the entries provided by the client,
            // starting with the missing ones, and fetching all of their
            // predecessors.
            let mut predecessor_stack = missing_primordials.clone();
            let mut predecessor_ids = HashSet::new();
            while let Some(missing_id) = predecessor_stack.pop() {
                if let Some(entry) = remaining_entries.get(&missing_id) {
                    for predecessor_id in entry.predecessors() {
                        if predecessor_ids.insert(*predecessor_id) {
                            predecessor_stack.push(*predecessor_id);
                        }
                    }
                }
            }

            debug!(
                ctx.logger(),
                "Mutation store slow path fetching {} additional entries",
                predecessor_ids.len(),
            );

            self.fetch_by_successor(
                &self.connections.read_master_connection,
                &mut entry_set,
                &predecessor_ids,
            )
            .await?;

            // Add these new entries, finding their primordials in the process.
            added.extend(
                entry_set
                    .add_entries_and_find_primordials(remaining_entries, &missing_primordials)?,
            );
        }

        // Store the new entries.
        self.store(ctx, &entry_set, &added).await?;
        Ok(())
    }

    /// Get all predecessor information for the given changeset id, keyed by
    /// the successor changeset id.
    ///
    /// Returns all entries that describe the mutation history of the commits.
    /// keyed by the successor changeset ids.
    async fn all_predecessors_by_changeset(
        &self,
        ctx: &CoreContext,
        changeset_ids: HashSet<HgChangesetId>,
    ) -> Result<HashMap<HgChangesetId, Vec<HgMutationEntry>>> {
        if changeset_ids.is_empty() {
            // Nothing to fetch
            return Ok(HashMap::new());
        }

        let mut entry_set = HgMutationEntrySet::new();
        let connection = self.read_connection_for_changesets(&changeset_ids).await?;
        self.fetch_by_successor(connection, &mut entry_set, &changeset_ids)
            .await?;
        self.fetch_all_predecessors(connection, &mut entry_set)
            .await?;
        let changeset_count = changeset_ids.len();
        let entries = entry_set.into_all_predecessors_by_changeset(changeset_ids);
        debug!(
            ctx.logger(),
            "Mutation store fetched {} entries for {} changesets",
            entries.len(),
            changeset_count,
        );
        ctx.perf_counters().add_to_counter(
            PerfCounterType::HgMutationStoreNumFetched,
            entries.len() as i64,
        );
        Ok(entries)
    }
}

mononoke_queries! {
    write AddChangesets(
        values: (
            repo_id: RepositoryId,
            changeset_id: HgChangesetId,
        )
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO hg_mutation_changesets (
            repo_id,
            changeset_id
        ) VALUES {values}"
    }

    write AddEntries(
        values: (
            repo_id: RepositoryId,
            successor: HgChangesetId,
            primordial: HgChangesetId,
            pred_count: u64,
            split_count: u64,
            op: str,
            user: str,
            timestamp: i64,
            tz: i32,
            extra: str,
        )
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO hg_mutation_info (
            repo_id,
            successor,
            primordial,
            pred_count,
            split_count,
            op,
            user,
            timestamp,
            tz,
            extra
        ) VALUES {values}"
    }

    write AddPreds(
        values: (
            repo_id: RepositoryId,
            successor: HgChangesetId,
            seq: u64,
            predecessor: HgChangesetId,
            primordial: HgChangesetId
        )
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO hg_mutation_preds (
            repo_id,
            successor,
            seq,
            predecessor,
            primordial
        ) VALUES {values}"
    }

    write AddSplits(
        values: (
            repo_id: RepositoryId,
            successor: HgChangesetId,
            seq: u64,
            split_successor: HgChangesetId,
        )
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO hg_mutation_splits (
            repo_id,
            successor,
            seq,
            split_successor
        ) VALUES {values}"
    }

    read CountChangesets(repo_id: RepositoryId, >list cs_id: &HgChangesetId) -> (i64) {
        "SELECT
            COUNT(*)
        FROM
            hg_mutation_changesets
        WHERE repo_id = {repo_id} AND changeset_id IN {cs_id}"
    }

    read SelectBySuccessor(repo_id: RepositoryId, >list cs_id: HgChangesetId) -> (
        HgChangesetId,
        HgChangesetId,
        u64,
        u64,
        HgChangesetId,
        HgChangesetId,
        u64,
        String,
        String,
        i64,
        i32,
        String
    ) {
        "SELECT
            m.successor, m.primordial,
            m.pred_count, p.seq, p.predecessor, p.primordial,
            m.split_count,
            m.op, m.user, m.timestamp, m.tz, m.extra
        FROM
            hg_mutation_info m LEFT JOIN hg_mutation_preds p
            ON m.repo_id = p.repo_id AND m.successor = p.successor
        WHERE m.repo_id = {repo_id} AND m.successor IN {cs_id}
        ORDER BY m.successor, p.seq ASC"
    }

    read SelectByPrimordial(repo_id: RepositoryId, mut_lim: usize, >list cs_id: HgChangesetId) -> (
        HgChangesetId,
        HgChangesetId,
        u64,
        u64,
        HgChangesetId,
        HgChangesetId,
        u64,
        String,
        String,
        i64,
        i32,
        String
    ) {
        "SELECT
            m.successor, m.primordial,
            m.pred_count, p.seq, p.predecessor, p.primordial,
            m.split_count,
            m.op, m.user, m.timestamp, m.tz, m.extra
        FROM
            hg_mutation_info m LEFT JOIN hg_mutation_preds p
            ON m.repo_id = p.repo_id AND m.successor = p.successor
        WHERE m.repo_id = {repo_id} AND m.primordial IN {cs_id}
        ORDER BY m.id DESC, p.seq ASC
        LIMIT {mut_lim}"
    }

    read SelectSplitsBySuccessor(repo_id: RepositoryId, >list cs_id: HgChangesetId) -> (
        HgChangesetId,
        u64,
        HgChangesetId,
    ) {
        "SELECT successor, seq, split_successor
        FROM hg_mutation_splits
        WHERE repo_id = {repo_id} AND successor IN {cs_id}
        ORDER BY successor, seq ASC"
    }
}


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

//! Basic tests.

use std::collections::HashMap;
use std::collections::HashSet;

use anyhow::Result;
use context::CoreContext;
use fbinit::FacebookInit;
use maplit::hashmap;
use maplit::hashset;
use mercurial_mutation::HgMutationEntry;
use mercurial_mutation::HgMutationStore;
use mercurial_mutation::SqlHgMutationStoreBuilder;
use mercurial_types_mocks::nodehash::make_hg_cs_id;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql_construct::SqlConstruct;

use crate::util::check_entries;

pub(crate) fn create_entries() -> HashMap<usize, HgMutationEntry> {
    // Generate the mutation graph:
    //
    //
    //                   3  --.                                   .--> 7
    //                         \                                 /
    //   1  --(amend)->  2  --(fold)-> 4  --(rebase)->  5  --(split)-> 6
    hashmap! {
        2 => HgMutationEntry::new(
            make_hg_cs_id(2),
            vec![make_hg_cs_id(1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        4 => HgMutationEntry::new(
            make_hg_cs_id(4),
            vec![make_hg_cs_id(2), make_hg_cs_id(3)],
            vec![],
            String::from("fold"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        5 => HgMutationEntry::new(
            make_hg_cs_id(5),
            vec![make_hg_cs_id(4)],
            vec![],
            String::from("rebase"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        6 => HgMutationEntry::new(
            make_hg_cs_id(6),
            vec![make_hg_cs_id(5)],
            vec![make_hg_cs_id(7)],
            String::from("split"),
            String::from("testuser"),
            0,
            0,
            vec![(String::from("testextra"), String::from("testextravalue"))],
        ),
    }
}

#[fbinit::test]
async fn add_entries_and_fetch_predecessors(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let store = SqlHgMutationStoreBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_repo_id(REPO_ZERO);

    // Add the initial set of entries.
    let mut entries = create_entries();
    store
        .add_entries(
            &ctx,
            hashset![make_hg_cs_id(6), make_hg_cs_id(7)],
            entries.values().cloned().collect(),
        )
        .await?;

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(6)],
        &entries,
        &[2, 4, 5, 6],
    )
    .await?;
    check_entries(&store, &ctx, hashset![make_hg_cs_id(4)], &entries, &[2, 4]).await?;
    check_entries(&store, &ctx, hashset![make_hg_cs_id(1)], &entries, &[]).await?;
    check_entries(&store, &ctx, hashset![make_hg_cs_id(7)], &entries, &[]).await?;

    // Add one new entry.
    let new_entry = HgMutationEntry::new(
        make_hg_cs_id(8),
        vec![make_hg_cs_id(6)],
        vec![],
        String::from("amend"),
        String::from("testuser"),
        0,
        0,
        vec![],
    );
    store
        .add_entries(&ctx, hashset![make_hg_cs_id(8)], vec![new_entry.clone()])
        .await?;
    entries.insert(8, new_entry);

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(8)],
        &entries,
        &[2, 4, 5, 6, 8],
    )
    .await?;

    // Add an entry with a gap.
    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(9),
            vec![make_hg_cs_id(6)],
            vec![],
            String::from("rebase"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(10),
            vec![make_hg_cs_id(9)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(10)], new_entries.clone())
        .await?;
    entries.extend((9..).zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(10)],
        &entries,
        &[2, 4, 5, 6, 9, 10],
    )
    .await?;

    // Add a more complex fold with some gaps and a new primordial.  This is
    // somewhat artificial as this kind of entry is unlikely to happen in normal
    // use, however it will exercise some edge cases.
    //
    // Current graph is:
    //
    //       3 -.               .-> 7
    //           \             /
    //   1 --> 2 --> 4 -->  5 --> 6 --> 8
    //                             \
    //                              '-> 9 --> 10
    //
    // We will add:
    //
    //     10  ------------> 13
    //      7  ----------/
    //      8  --> 11 --/
    //             12 -'

    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(11),
            vec![make_hg_cs_id(8)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(13),
            vec![
                make_hg_cs_id(10),
                make_hg_cs_id(7),
                make_hg_cs_id(11),
                make_hg_cs_id(12),
            ],
            vec![],
            String::from("combine"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(13)], new_entries.clone())
        .await?;
    entries.extend(vec![11, 13].into_iter().zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(13)],
        &entries,
        &[2, 4, 5, 6, 8, 9, 10, 11, 13],
    )
    .await?;

    // Extend history backwards.  The original client was missing some earlier
    // data.  A new client sends commits with additional history for commit 1.
    //
    //   14 --> 15 --> 1 --> 16 --> 17
    //
    // Note that the fast-path addition of mutation data won't process this,
    // but that's a reasonable trade-off.  We add two additional successors to
    // ensure we use the slow path.
    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(15),
            vec![make_hg_cs_id(14)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(1),
            vec![make_hg_cs_id(15)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(16),
            vec![make_hg_cs_id(1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(17),
            vec![make_hg_cs_id(16)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(17)], new_entries.clone())
        .await?;
    entries.extend(vec![15, 1, 16, 17].into_iter().zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(17)],
        &entries,
        &[1, 15, 16, 17],
    )
    .await?;
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(4)],
        &entries,
        &[1, 2, 4, 15],
    )
    .await?;

    Ok(())
}

#[fbinit::test]
async fn check_mutations_are_cut_when_reaching_limit(fb: FacebookInit) -> Result<()> {
    const TEST_MUTATION_LIMIT: usize = 10;

    let ctx = CoreContext::test_mock(fb);
    let store = SqlHgMutationStoreBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_mutation_limit(TEST_MUTATION_LIMIT)
        .with_repo_id(REPO_ZERO);

    // Add a lot of entries

    let mut entries = hashmap! {};

    let mut new_entries = Vec::with_capacity(20);

    let mut amend_count: u64 = 20;

    for index in 1..amend_count {
        new_entries.push(HgMutationEntry::new(
            make_hg_cs_id(index),
            vec![make_hg_cs_id(index - 1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ));
    }

    store
        .add_entries(
            &ctx,
            (1..20).map(make_hg_cs_id).collect::<HashSet<_>>(),
            new_entries.clone(),
        )
        .await?;

    entries.extend((1..20).zip(new_entries));

    // First we want to make sure that we are not fetching the entire history
    let fetched_entries = store
        .all_predecessors(&ctx, hashset![make_hg_cs_id(19)])
        .await?;
    assert_eq!(fetched_entries.len(), TEST_MUTATION_LIMIT);

    // Now we want to make sure that the mutations we are collecting, correspond to the latest amends
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(19)],
        &entries,
        &(10..20).collect::<Vec<_>>(),
    )
    .await?;

    // What we want to do here is make multiple folds with 3 amends each.
    // This way, we reach the maximum number of changes (10*) BUT made specially
    // for the case where the limit cuts a fold in half. The expected behaviour is
    // for the half that wasn't cut by the limit to be removed.
    //
    // *In this case the maximum number is 10 because we are overriding it
    // in the `.with_mutation_limit(TEST_MUTATION_LIMIT)` call
    //
    // Add a lot of entries
    //
    //                             30-.
    //                                 \ <--- Here it is exceding the limit
    //                               29-.
    //                                   \
    //                                 28-.
    //                                     \
    //   19 --> 20 --> 21 --> 22 --> ... --> 27
    new_entries = Vec::with_capacity(11);

    amend_count = 7;
    // We add 7 amend operations and then create a fold that
    // will exceed the limit to check that it is actually removing
    // the mutation.
    for index in 0..amend_count {
        new_entries.push(HgMutationEntry::new(
            make_hg_cs_id(20 + index),
            vec![make_hg_cs_id(20 + index - 1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ));
    }

    new_entries.push(HgMutationEntry::new(
        make_hg_cs_id(20 + amend_count),
        vec![
            make_hg_cs_id(20 + amend_count - 1),
            make_hg_cs_id(20 + amend_count + 1),
            make_hg_cs_id(20 + amend_count + 2),
            make_hg_cs_id(20 + amend_count + 3),
        ],
        vec![],
        String::from("combine"),
        String::from("testuser"),
        0,
        0,
        vec![],
    ));

    store
        .add_entries(
            &ctx,
            (20..32).map(make_hg_cs_id).collect::<HashSet<_>>(),
            new_entries.clone(),
        )
        .await?;

    entries.extend((20..32).zip(new_entries));

    // First we want to make sure that we are not fetching the entire history
    let fetched_entries = store
        .all_predecessors(&ctx, hashset![make_hg_cs_id(20 + amend_count)])
        .await?;

    // The last fold should be erased because it exceeds the maximum (10 in this case) and it is then
    // cut in half. So only amends stay (7)
    assert_ne!(fetched_entries.len(), 10);
    assert_eq!(fetched_entries.len(), 7);

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

//! Caching tests.

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use fbinit::FacebookInit;
use maplit::hashmap;
use maplit::hashset;
use mercurial_mutation::CachedHgMutationStore;
use mercurial_mutation::HgMutationEntry;
use mercurial_mutation::HgMutationStore;
use mercurial_mutation::SqlHgMutationStoreBuilder;
use mercurial_types::HgChangesetId;
use mercurial_types_mocks::nodehash::make_hg_cs_id;
use mononoke_types::RepositoryId;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql_construct::SqlConstruct;

use crate::basic::create_entries;
use crate::util::check_entries;

struct CountedHgMutationStore {
    inner_store: Arc<dyn HgMutationStore>,
    add_entries: Arc<AtomicUsize>,
    all_predecessors: Arc<AtomicUsize>,
}

impl CountedHgMutationStore {
    fn new(
        inner_store: Arc<dyn HgMutationStore>,
        add_entries: Arc<AtomicUsize>,
        all_predecessors: Arc<AtomicUsize>,
    ) -> Self {
        Self {
            inner_store,
            add_entries,
            all_predecessors,
        }
    }
}

#[async_trait]
impl HgMutationStore for CountedHgMutationStore {
    fn repo_id(&self) -> RepositoryId {
        self.inner_store.repo_id()
    }

    async fn add_entries(
        &self,
        ctx: &CoreContext,
        new_changeset_ids: HashSet<HgChangesetId>,
        entries: Vec<HgMutationEntry>,
    ) -> Result<()> {
        self.add_entries.fetch_add(1, Ordering::Relaxed);
        self.inner_store
            .add_entries(ctx, new_changeset_ids, entries)
            .await
    }

    async fn all_predecessors_by_changeset(
        &self,
        ctx: &CoreContext,
        changeset_ids: HashSet<HgChangesetId>,
    ) -> Result<HashMap<HgChangesetId, Vec<HgMutationEntry>>> {
        self.all_predecessors.fetch_add(1, Ordering::Relaxed);
        self.inner_store
            .all_predecessors_by_changeset(ctx, changeset_ids)
            .await
    }
}

#[fbinit::test]
async fn add_entries_and_fetch_predecessors_with_caching(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let sql_store = SqlHgMutationStoreBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_repo_id(REPO_ZERO);

    let add_entries = Arc::new(AtomicUsize::new(0));
    let all_predecessors = Arc::new(AtomicUsize::new(0));
    let counted_store = CountedHgMutationStore::new(
        Arc::new(sql_store),
        add_entries.clone(),
        all_predecessors.clone(),
    );
    let store = CachedHgMutationStore::new_test(Arc::new(counted_store));

    // Add the initial set of entries.
    let mut entries = create_entries();
    store
        .add_entries(
            &ctx,
            hashset![make_hg_cs_id(6), make_hg_cs_id(7)],
            entries.values().cloned().collect(),
        )
        .await?;
    assert_eq!(add_entries.load(Ordering::Relaxed), 1);
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(6)],
        &entries,
        &[2, 4, 5, 6],
    )
    .await?;
    // First time query, cache miss. Counters should go up.
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 1);

    // Querying for different changeset, counters should again go up.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(4)], &entries, &[2, 4]).await?;
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 2);

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(4), make_hg_cs_id(6)],
        &entries,
        &[2, 4, 5, 6],
    )
    .await?;
    // Querying for known changesets, counters should remain the same.
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 2);

    // Querying for different changeset, counters should again go up.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(1)], &entries, &[]).await?;
    check_entries(&store, &ctx, hashset![make_hg_cs_id(7)], &entries, &[]).await?;
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 4);

    // Add one new entry.
    let new_entry = HgMutationEntry::new(
        make_hg_cs_id(8),
        vec![make_hg_cs_id(6)],
        vec![],
        String::from("amend"),
        String::from("testuser"),
        0,
        0,
        vec![],
    );
    // Entries always bypass the cache, so the counters should always go up
    store
        .add_entries(&ctx, hashset![make_hg_cs_id(8)], vec![new_entry.clone()])
        .await?;
    assert_eq!(add_entries.load(Ordering::Relaxed), 2);
    entries.insert(8, new_entry);

    // Adding a new entry shouldn't change the result for an existing changeset.
    // The value should still be returned from cache and the counters should
    // remain the same.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(7)], &entries, &[]).await?;
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 4);

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(8)],
        &entries,
        &[2, 4, 5, 6, 8],
    )
    .await?;

    // Add an entry with a gap.
    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(9),
            vec![make_hg_cs_id(6)],
            vec![],
            String::from("rebase"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(10),
            vec![make_hg_cs_id(9)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    // Entries always bypass the cache, so the counters should always go up
    store
        .add_entries(&ctx, hashset![make_hg_cs_id(10)], new_entries.clone())
        .await?;
    assert_eq!(add_entries.load(Ordering::Relaxed), 3);
    entries.extend((9..).zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(10)],
        &entries,
        &[2, 4, 5, 6, 9, 10],
    )
    .await?;

    // Add a more complex fold with some gaps and a new primordial.  This is
    // somewhat artificial as this kind of entry is unlikely to happen in normal
    // use, however it will exercise some edge cases.
    //
    // Current graph is:
    //
    //       3 -.               .-> 7
    //           \             /
    //   1 --> 2 --> 4 -->  5 --> 6 --> 8
    //                             \
    //                              '-> 9 --> 10
    //
    // We will add:
    //
    //     10  ------------> 13
    //      7  ----------/
    //      8  --> 11 --/
    //             12 -'

    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(11),
            vec![make_hg_cs_id(8)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(13),
            vec![
                make_hg_cs_id(10),
                make_hg_cs_id(7),
                make_hg_cs_id(11),
                make_hg_cs_id(12),
            ],
            vec![],
            String::from("combine"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(13)], new_entries.clone())
        .await?;
    entries.extend(vec![11, 13].into_iter().zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(13)],
        &entries,
        &[2, 4, 5, 6, 8, 9, 10, 11, 13],
    )
    .await?;

    // Extend history backwards.  The original client was missing some earlier
    // data.  A new client sends commits with additional history for commit 1.
    //
    //   14 --> 15 --> 1 --> 16 --> 17
    //
    // Note that the fast-path addition of mutation data won't process this,
    // but that's a reasonable trade-off.  We add two additional successors to
    // ensure we use the slow path.
    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(15),
            vec![make_hg_cs_id(14)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(1),
            vec![make_hg_cs_id(15)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(16),
            vec![make_hg_cs_id(1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(17),
            vec![make_hg_cs_id(16)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(17)], new_entries.clone())
        .await?;
    entries.extend(vec![15, 1, 16, 17].into_iter().zip(new_entries));

    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(17)],
        &entries,
        &[1, 15, 16, 17],
    )
    .await?;
    // The mutation history for 4 was fetched earlier, and inspite of change
    // in the primordial the history would still be served through cache.
    // Thus, instead of getting [1, 2, 4, 15] as the history we get [2, 4].
    // The new values would take place once the cache expires.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(4)], &entries, &[2, 4]).await?;

    Ok(())
}

#[fbinit::test]
async fn check_mutations_are_cut_when_reaching_limit_with_caching(fb: FacebookInit) -> Result<()> {
    const TEST_MUTATION_LIMIT: usize = 10;
    let ctx = CoreContext::test_mock(fb);
    let sql_store = SqlHgMutationStoreBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_mutation_limit(TEST_MUTATION_LIMIT)
        .with_repo_id(REPO_ZERO);

    let add_entries = Arc::new(AtomicUsize::new(0));
    let all_predecessors = Arc::new(AtomicUsize::new(0));
    let counted_store = CountedHgMutationStore::new(
        Arc::new(sql_store),
        add_entries.clone(),
        all_predecessors.clone(),
    );
    let store = CachedHgMutationStore::new_test(Arc::new(counted_store));

    // Add a lot of entries

    let mut entries = hashmap! {};

    let mut new_entries = Vec::with_capacity(20);

    let mut amend_count: u64 = 20;

    for index in 1..amend_count {
        new_entries.push(HgMutationEntry::new(
            make_hg_cs_id(index),
            vec![make_hg_cs_id(index - 1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ));
    }

    // New entry added, cache bypassed, counter updated.
    store
        .add_entries(
            &ctx,
            (1..20).map(make_hg_cs_id).collect::<HashSet<_>>(),
            new_entries.clone(),
        )
        .await?;
    assert_eq!(add_entries.load(Ordering::Relaxed), 1);
    entries.extend((1..20).zip(new_entries));

    // First we want to make sure that we are not fetching the entire history
    let fetched_entries = store
        .all_predecessors(&ctx, hashset![make_hg_cs_id(19)])
        .await?;
    assert_eq!(fetched_entries.len(), TEST_MUTATION_LIMIT);
    // First time fetch, no entry in cache, counter should go up.
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 1);

    // Now we want to make sure that the mutations we are collecting, correspond to the latest amends
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(19)],
        &entries,
        &(10..20).collect::<Vec<_>>(),
    )
    .await?;
    // Now the entry is in cache, counter shouldn't move
    assert_eq!(all_predecessors.load(Ordering::Relaxed), 1);

    // What we want to do here is make multiple folds with 3 amends each.
    // This way, we reach the maximum number of changes (10*) BUT made specially
    // for the case where the limit cuts a fold in half. The expected behaviour is
    // for the half that wasn't cut by the limit to be removed.
    //
    // *In this case the maximum number is 10 because we are overriding it
    // in the `.with_mutation_limit(TEST_MUTATION_LIMIT)` call
    //
    // Add a lot of entries
    //
    //                             30-.
    //                                 \ <--- Here it is exceding the limit
    //                               29-.
    //                                   \
    //                                 28-.
    //                                     \
    //   19 --> 20 --> 21 --> 22 --> ... --> 27
    new_entries = Vec::with_capacity(11);

    amend_count = 7;
    // We add 7 amend operations and then create a fold that
    // will exceed the limit to check that it is actually removing
    // the mutation.
    for index in 0..amend_count {
        new_entries.push(HgMutationEntry::new(
            make_hg_cs_id(20 + index),
            vec![make_hg_cs_id(20 + index - 1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ));
    }

    new_entries.push(HgMutationEntry::new(
        make_hg_cs_id(20 + amend_count),
        vec![
            make_hg_cs_id(20 + amend_count - 1),
            make_hg_cs_id(20 + amend_count + 1),
            make_hg_cs_id(20 + amend_count + 2),
            make_hg_cs_id(20 + amend_count + 3),
        ],
        vec![],
        String::from("combine"),
        String::from("testuser"),
        0,
        0,
        vec![],
    ));

    store
        .add_entries(
            &ctx,
            (20..32).map(make_hg_cs_id).collect::<HashSet<_>>(),
            new_entries.clone(),
        )
        .await?;

    entries.extend((20..32).zip(new_entries));

    // First we want to make sure that we are not fetching the entire history
    let fetched_entries = store
        .all_predecessors(&ctx, hashset![make_hg_cs_id(20 + amend_count)])
        .await?;

    // The last fold should be erased because it exceeds the maximum (10 in this case) and it is then
    // cut in half. So only amends stay (7)
    assert_ne!(fetched_entries.len(), 10);
    assert_eq!(fetched_entries.len(), 7);

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

//! Tests involving cycles.
//!
//! Ordinarily, cycles in mutation data is invalid, however they can occur
//! with synthetic mutation data backfilled from obsmarkers.  The mutation
//! store deals with cycles in two ways:
//!
//! * Self-referential entries (where the successor is one of the
//!   predecessors) have been converted from revive obsmarkers and are
//!   dropped.
//!
//! * For larger cycles, the mutation store is mostly unaffected, except for
//!   one case: primordial determination.  If the mutation store cannot
//!   determine a primordial commit because of cycles, those entries that
//!   create the cycle will be dropped.
//!
//! If a cycle is formed within the store itself (via multiple additions that
//! eventually form a cycle), this is not a problem for the store.  Any
//! request for a commit within that cycle will return the full cycle, and
//! client-side cycle-breaking will resolve the issue.

use std::collections::HashMap;

use anyhow::Result;
use context::CoreContext;
use fbinit::FacebookInit;
use maplit::hashmap;
use maplit::hashset;
use mercurial_mutation::HgMutationEntry;
use mercurial_mutation::HgMutationStore;
use mercurial_mutation::SqlHgMutationStoreBuilder;
use mercurial_types_mocks::nodehash::make_hg_cs_id;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql_construct::SqlConstruct;

use crate::util::check_entries;
use crate::util::get_entries;

fn create_entries() -> HashMap<usize, HgMutationEntry> {
    // Generate the mutation graph:
    //
    // Self-cycle:
    //
    //   1  --(amend)-> 2
    //  / \
    //  \_/ (revive)
    //
    // Simple cycle:
    //
    //   3  --(amend)--> 4 --(rebase)--> 5 --(undo)--> 3
    //
    // Cycle through second predecessor:
    //
    //   6  --(amend)--> 7 --(fold)--> 8 --(unfold)--> 9
    //                       /
    //                   9 -'
    //
    // Second predecessor is in a separate cycle:
    //
    //                  10 --(amend)--> 11 --(fold)--> 14 --(rebase)--> 15
    //                                       /
    //                    12 --(amend)-->  13 --(undo)--> 12
    //
    hashmap! {
        1 => HgMutationEntry::new(
            make_hg_cs_id(1),
            vec![make_hg_cs_id(1)],
            vec![],
            String::from("revive"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        2 => HgMutationEntry::new(
            make_hg_cs_id(2),
            vec![make_hg_cs_id(1)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        3 => HgMutationEntry::new(
            make_hg_cs_id(3),
            vec![make_hg_cs_id(5)],
            vec![],
            String::from("undo"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        4 => HgMutationEntry::new(
            make_hg_cs_id(4),
            vec![make_hg_cs_id(3)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        5 => HgMutationEntry::new(
            make_hg_cs_id(5),
            vec![make_hg_cs_id(4)],
            vec![],
            String::from("rebase"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        7 => HgMutationEntry::new(
            make_hg_cs_id(7),
            vec![make_hg_cs_id(6)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        8 => HgMutationEntry::new(
            make_hg_cs_id(8),
            vec![make_hg_cs_id(7), make_hg_cs_id(9)],
            vec![],
            String::from("fold"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        9 => HgMutationEntry::new(
            make_hg_cs_id(9),
            vec![make_hg_cs_id(8)],
            vec![],
            String::from("unfold"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        11 => HgMutationEntry::new(
            make_hg_cs_id(11),
            vec![make_hg_cs_id(10)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        12 => HgMutationEntry::new(
            make_hg_cs_id(12),
            vec![make_hg_cs_id(13)],
            vec![],
            String::from("undo"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        13 => HgMutationEntry::new(
            make_hg_cs_id(13),
            vec![make_hg_cs_id(12)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        14 => HgMutationEntry::new(
            make_hg_cs_id(14),
            vec![make_hg_cs_id(11), make_hg_cs_id(13)],
            vec![],
            String::from("fold"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        15 => HgMutationEntry::new(
            make_hg_cs_id(15),
            vec![make_hg_cs_id(14)],
            vec![],
            String::from("rebase"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    }
}

#[fbinit::test]
async fn add_entries_and_fetch_predecessors(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let store = SqlHgMutationStoreBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_repo_id(REPO_ZERO);

    // Add the initial set of entries.
    let mut entries = create_entries();
    store
        .add_entries(
            &ctx,
            hashset![
                make_hg_cs_id(2),
                make_hg_cs_id(5),
                make_hg_cs_id(9),
                make_hg_cs_id(15)
            ],
            entries.values().cloned().collect(),
        )
        .await?;

    // The entry for 1 was ignored as it was self-referential.  The entry for 2
    // was stored.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(2)], &entries, &[2]).await?;

    // The entries for 3, 4 and 5 were all dropped due to the loop.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(5)], &entries, &[]).await?;

    // The entries for 7, 8 and 9 were stored ok as the loop did not affect the
    // primordial search.
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(9)],
        &entries,
        &[7, 8, 9],
    )
    .await?;

    // The history for 15 was cut at the fold as it has a cyclic predecessor.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(15)], &entries, &[15]).await?;

    // A different client pushes without the cycle.
    store
        .add_entries(
            &ctx,
            hashset![make_hg_cs_id(14)],
            get_entries(&entries, &[11, 13, 14]),
        )
        .await?;

    // The history for 15 now goes back to 11.
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(15)],
        &entries,
        &[11, 13, 14, 15],
    )
    .await?;

    // Another client sends entries that create a cycle with existing entries.
    let new_entries = vec![
        HgMutationEntry::new(
            make_hg_cs_id(16),
            vec![make_hg_cs_id(2)],
            vec![],
            String::from("amend"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
        HgMutationEntry::new(
            make_hg_cs_id(1),
            vec![make_hg_cs_id(16)],
            vec![],
            String::from("undo"),
            String::from("testuser"),
            0,
            0,
            vec![],
        ),
    ];

    store
        .add_entries(&ctx, hashset![make_hg_cs_id(1)], new_entries.clone())
        .await?;
    entries.extend(vec![16, 1].into_iter().zip(new_entries));

    // The replacement entry for 1 was ignored.
    check_entries(&store, &ctx, hashset![make_hg_cs_id(2)], &entries, &[2]).await?;

    // But the entry for 16 was accepted.
    check_entries(
        &store,
        &ctx,
        hashset![make_hg_cs_id(16)],
        &entries,
        &[2, 16],
    )
    .await?;

    Ok(())
}


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

//! Tests for the hg mutation store.

mod basic;
mod caching;
mod cycles;
mod util;