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

use std::sync::Arc;
use std::time::Duration;

use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use blobrepo::AsBlobRepo;
use blobrepo::BlobRepo;
use bonsai_hg_mapping::BonsaiHgMappingArc;
use bookmarks::BookmarkKey;
use bookmarks::BookmarksArc;
use bulkops::PublicChangesetBulkFetch;
use caching_ext::CacheHandlerFactory;
use changeset_fetcher::ChangesetFetcherArc;
use changeset_fetcher::ChangesetFetcherRef;
use changeset_fetcher::PrefetchedChangesetsFetcher;
use changesets::ChangesetEntry;
use changesets::ChangesetsArc;
use changesets::ChangesetsRef;
use context::CoreContext;
use fbinit::FacebookInit;
use fixtures::set_bookmark;
use fixtures::BranchEven;
use fixtures::Linear;
use fixtures::MergeEven;
use fixtures::MergeUneven;
use fixtures::TestRepoFixture;
use fixtures::UnsharedMergeEven;
use futures::compat::Stream01CompatExt;
use futures::future::try_join_all;
use futures::future::FutureExt;
use futures::stream;
use futures::StreamExt;
use maplit::hashmap;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use once_cell::sync::Lazy;
use phases::PhasesArc;
use phases::PhasesRef;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use revset::AncestorsNodeStream;
use sql_construct::SqlConstruct;
use sql_ext::replication::NoReplicaLagMonitor;
use tests_utils::resolve_cs_id;
use tests_utils::CreateCommitContext;
use tunables::override_tunables;
use tunables::with_tunables_async;

use crate::builder::SegmentedChangelogSqlConnections;
use crate::iddag::IdDagSaveStore;
use crate::idmap::CacheHandlers;
use crate::idmap::ConcurrentMemIdMap;
use crate::idmap::IdMap;
use crate::idmap::IdMapFactory;
use crate::idmap::SqlIdMap;
use crate::manager::SegmentedChangelogManager;
use crate::manager::SegmentedChangelogType;
use crate::on_demand::OnDemandUpdateSegmentedChangelog;
use crate::owned::OwnedSegmentedChangelog;
use crate::periodic_reload::PeriodicReloadSegmentedChangelog;
use crate::tailer::SegmentedChangelogTailer;
use crate::types::IdDagVersion;
use crate::types::IdMapVersion;
use crate::types::SegmentedChangelogVersion;
use crate::version_store::SegmentedChangelogVersionStore;
use crate::CloneHints;
use crate::InProcessIdDag;
use crate::Location;
use crate::SeedHead;
use crate::SegmentedChangelog;
use crate::SegmentedChangelogRef;

#[async_trait::async_trait]
trait SegmentedChangelogExt {
    async fn head(&self, ctx: &CoreContext) -> Result<ChangesetId>;
}

#[async_trait::async_trait]
impl<T> SegmentedChangelogExt for T
where
    T: SegmentedChangelog,
{
    async fn head(&self, ctx: &CoreContext) -> Result<ChangesetId> {
        let clone_data = self.clone_data(ctx).await?;
        let cs_id = clone_data
            .0
            .idmap
            .iter()
            .max_by_key(|i| i.0)
            .map(|i| i.1.clone())
            .ok_or_else(|| format_err!("error with clone data, missing idmap entry for head_id"))?;
        Ok(cs_id)
    }
}

static BOOKMARK_NAME: Lazy<BookmarkKey> = Lazy::new(|| BookmarkKey::new("master").unwrap());

async fn new_tailer_for_tailing(
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
) -> Result<SegmentedChangelogTailer> {
    new_tailer(blobrepo, connections, None, None).await
}

async fn new_tailer(
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
    prefetched: Option<Vec<ChangesetEntry>>,
    seed_heads: Option<Vec<ChangesetId>>,
) -> Result<SegmentedChangelogTailer> {
    let prefetched = match prefetched {
        None => stream::empty().boxed(),
        Some(prefetched) => stream::iter(prefetched.into_iter().map(Ok)).boxed(),
    };

    let changeset_fetcher = Arc::new(
        PrefetchedChangesetsFetcher::new(
            blobrepo.repo_identity().id(),
            blobrepo.changesets_arc(),
            prefetched,
        )
        .await?,
    );

    let bulk_fetcher = Arc::new(PublicChangesetBulkFetch::new(
        blobrepo.changesets_arc(),
        blobrepo.phases_arc(),
    ));

    let seed_heads = match seed_heads {
        Some(heads) => heads.into_iter().map(SeedHead::from).collect(),
        None => vec![Some(BOOKMARK_NAME.clone()).into()],
    };

    Ok(SegmentedChangelogTailer::new(
        blobrepo.repo_identity().id(),
        connections.clone(),
        Arc::new(NoReplicaLagMonitor()),
        changeset_fetcher,
        bulk_fetcher,
        blobrepo.bonsai_hg_mapping_arc(),
        Arc::new(blobrepo.repo_blobstore().clone()),
        blobrepo.bookmarks_arc(),
        seed_heads,
        None,
    ))
}

async fn seed_with_prefetched(
    ctx: &CoreContext,
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
    heads: Vec<ChangesetId>,
    prefetched: Option<Vec<ChangesetEntry>>,
) -> Result<()> {
    // Does it make sense to seed up to the master commit instead of head?
    // Something to consider. Depends how it's used by tests.
    blobrepo
        .phases()
        .add_reachable_as_public(ctx, heads.clone())
        .await?;
    let tailer = new_tailer(blobrepo, connections, prefetched, Some(heads)).await?;
    tailer.once(ctx, true).await?;
    Ok(())
}

async fn seed(
    ctx: &CoreContext,
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
    head: ChangesetId,
) -> Result<()> {
    seed_with_prefetched(
        ctx,
        blobrepo,
        connections,
        vec![head],
        None, // no prefetching
    )
    .await
}

async fn load_sc_version(
    ctx: &CoreContext,
    repo_id: RepositoryId,
    connections: &SegmentedChangelogSqlConnections,
) -> Result<SegmentedChangelogVersion> {
    let sc_version_store = SegmentedChangelogVersionStore::new(connections.0.clone(), repo_id);
    let sc_version = sc_version_store
        .get(ctx)
        .await?
        .context("failed to load segmented changelog version")?;
    Ok(sc_version)
}

async fn load_idmap(
    ctx: &CoreContext,
    repo_id: RepositoryId,
    connections: &SegmentedChangelogSqlConnections,
) -> Result<Arc<dyn IdMap>> {
    let sc_version = load_sc_version(ctx, repo_id, connections).await?;
    let idmap = SqlIdMap::new(
        connections.0.clone(),
        Arc::new(NoReplicaLagMonitor()),
        repo_id,
        sc_version.idmap_version,
    );
    Ok(Arc::new(idmap))
}

async fn load_iddag(
    ctx: &CoreContext,
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
) -> Result<InProcessIdDag> {
    let sc_version = load_sc_version(ctx, blobrepo.repo_identity().id(), connections).await?;
    let iddag_save_store = IdDagSaveStore::new(
        blobrepo.repo_identity().id(),
        Arc::new(blobrepo.repo_blobstore().clone()),
    );
    let iddag = iddag_save_store.load(ctx, sc_version.iddag_version).await?;
    Ok(iddag)
}

async fn load_owned(
    ctx: &CoreContext,
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
) -> Result<OwnedSegmentedChangelog> {
    let iddag = load_iddag(ctx, blobrepo, connections).await?;
    let idmap = load_idmap(ctx, blobrepo.repo_identity().id(), connections).await?;
    Ok(OwnedSegmentedChangelog::new(iddag, idmap))
}

async fn get_manager(
    blobrepo: &BlobRepo,
    connections: &SegmentedChangelogSqlConnections,
    seed_heads: Vec<SeedHead>,
    segmented_changelog_type: SegmentedChangelogType,
) -> Result<SegmentedChangelogManager> {
    let repo_id = blobrepo.repo_identity().id();
    let blobstore = Arc::new(blobrepo.repo_blobstore().clone());
    let sc_version_store = SegmentedChangelogVersionStore::new(connections.0.clone(), repo_id);
    let iddag_save_store = IdDagSaveStore::new(repo_id, blobstore.clone());
    let clone_hints = CloneHints::new(connections.0.clone(), repo_id, blobstore);
    let idmap_factory = IdMapFactory::new(
        connections.0.clone(),
        Arc::new(NoReplicaLagMonitor()),
        repo_id,
    );
    let manager = SegmentedChangelogManager::new(
        repo_id,
        sc_version_store,
        iddag_save_store,
        idmap_factory,
        blobrepo.changeset_fetcher_arc(),
        blobrepo.bookmarks_arc(),
        seed_heads,
        segmented_changelog_type,
        Some(clone_hints),
    );
    Ok(manager)
}

async fn validate_build_idmap(
    ctx: CoreContext,
    blobrepo: BlobRepo,
    head: &'static str,
) -> Result<()> {
    let head = resolve_cs_id(&ctx, &blobrepo, head).await?;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    seed(&ctx, &blobrepo, &conns, head).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let mut ancestors =
        AncestorsNodeStream::new(ctx.clone(), &blobrepo.changeset_fetcher_arc(), head).compat();
    while let Some(cs_id) = ancestors.next().await {
        let cs_id = cs_id?;
        let parents = blobrepo
            .changeset_fetcher()
            .get_parents(&ctx, cs_id)
            .await?;
        for parent in parents {
            let parent_dag_id = sc.idmap.get_dag_id(&ctx, parent).await?;
            let dag_id = sc.idmap.get_dag_id(&ctx, cs_id).await?;
            assert!(parent_dag_id < dag_id);
        }
    }
    Ok(())
}

async fn fetch_cs_entry(ctx: &CoreContext, repo: &BlobRepo, cs: &str) -> Result<ChangesetEntry> {
    let cs = resolve_cs_id(ctx, &repo, cs).await?;
    repo.changesets()
        .get(ctx, cs)
        .await?
        .ok_or_else(|| format_err!("can't find changeset entry for {}", cs))
}

#[fbinit::test]
async fn test_iddag_save_store(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    let repo_id = blobrepo.repo_identity().id();

    let known_cs =
        resolve_cs_id(&ctx, &blobrepo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await?;
    seed(&ctx, &blobrepo, &conns, known_cs).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let distance: u64 = 2;
    let answer = sc
        .location_to_changeset_id(&ctx, Location::new(known_cs, distance))
        .await?;
    let expected_cs =
        resolve_cs_id(&ctx, &blobrepo, "3e0e761030db6e479a7fb58b12881883f9f8c63f").await?;
    assert_eq!(answer, expected_cs);

    let iddag_save_store =
        IdDagSaveStore::new(repo_id, Arc::new(blobrepo.repo_blobstore().clone()));
    let iddag_version = iddag_save_store.save(&ctx, &sc.iddag).await?;

    assert!(
        iddag_save_store
            .find(&ctx, IdDagVersion::from_serialized_bytes(b"random"))
            .await?
            .is_none()
    );
    let loaded_id_dag = iddag_save_store.load(&ctx, iddag_version).await?;
    let from_save = OwnedSegmentedChangelog::new(loaded_id_dag, sc.idmap.clone());
    let answer = from_save
        .location_to_changeset_id(&ctx, Location::new(known_cs, distance))
        .await?;
    assert_eq!(answer, expected_cs);

    Ok(())
}

#[fbinit::test]
async fn test_build_idmap(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    validate_build_idmap(
        ctx.clone(),
        Linear::getrepo(fb).await,
        "79a13814c5ce7330173ec04d279bf95ab3f652fb",
    )
    .await?;
    validate_build_idmap(
        ctx.clone(),
        MergeEven::getrepo(fb).await,
        "1f6bc010883e397abeca773192f3370558ee1320",
    )
    .await?;
    validate_build_idmap(
        ctx.clone(),
        MergeUneven::getrepo(fb).await,
        "d35b1875cdd1ed2c687e86f1604b9d7e989450cb",
    )
    .await?;
    Ok(())
}

#[fbinit::test]
async fn test_is_ancestor(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let repo = MergeUneven::get_inner_repo(fb).await;
    let sc = repo.segmented_changelog();
    let blobrepo = repo.as_blob_repo();

    // Graph looks like:
    // a -> b -> c
    //  \-> d -> e
    let a = resolve_cs_id(&ctx, &blobrepo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await?;
    let b = resolve_cs_id(&ctx, &blobrepo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await?;
    let c = resolve_cs_id(&ctx, &blobrepo, "b65231269f651cfe784fd1d97ef02a049a37b8a0").await?;
    let d = resolve_cs_id(&ctx, &blobrepo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await?;
    let e = resolve_cs_id(&ctx, &blobrepo, "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5").await?;

    assert_eq!(sc.is_ancestor(&ctx, a, c).await?, None);
    assert!(sc.build_up_to_heads(&ctx, &[c, e]).await?);
    assert_eq!(sc.is_ancestor(&ctx, a, c).await?, Some(true));
    assert_eq!(sc.is_ancestor(&ctx, b, d).await?, Some(false));
    assert_eq!(sc.is_ancestor(&ctx, b, a).await?, Some(false));
    assert_eq!(sc.is_ancestor(&ctx, d, e).await?, Some(true));
    assert_eq!(sc.is_ancestor(&ctx, a, a).await?, Some(true));

    Ok(())
}

async fn validate_location_to_changeset_ids(
    ctx: CoreContext,
    blobrepo: BlobRepo,
    known: &'static str,
    distance_count: (u64, u64),
    expected: Vec<&'static str>,
) -> Result<()> {
    let (distance, count) = distance_count;
    let known_cs_id = resolve_cs_id(&ctx, &blobrepo, known).await?;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    seed(&ctx, &blobrepo, &conns, known_cs_id).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let answer = sc
        .location_to_many_changeset_ids(&ctx, Location::new(known_cs_id, distance), count)
        .await?;
    let expected_cs_ids = try_join_all(
        expected
            .into_iter()
            .map(|id| resolve_cs_id(&ctx, &blobrepo, id)),
    )
    .await?;
    assert_eq!(answer, expected_cs_ids);

    Ok(())
}

#[fbinit::test]
async fn test_location_to_changeset_ids(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    validate_location_to_changeset_ids(
        ctx.clone(),
        Linear::getrepo(fb).await,
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // master commit, message: modified 10
        (4, 3),
        vec![
            "0ed509bf086fadcb8a8a5384dc3b550729b0fc17", // message: added 7
            "eed3a8c0ec67b6a6fe2eb3543334df3f0b4f202b", // message: added 6
            "cb15ca4a43a59acff5388cea9648c162afde8372", // message: added 5
        ],
    )
    .await?;
    validate_location_to_changeset_ids(
        ctx.clone(),
        MergeUneven::getrepo(fb).await,
        "264f01429683b3dd8042cb3979e8bf37007118bc", // top of merged branch, message: Add 5
        (4, 2),
        vec![
            "795b8133cf375f6d68d27c6c23db24cd5d0cd00f", // message: Add 1
            "4f7f3fd428bec1a48f9314414b063c706d9c1aed", // bottom of branch
        ],
    )
    .await?;
    validate_location_to_changeset_ids(
        ctx.clone(),
        UnsharedMergeEven::getrepo(fb).await,
        "7fe9947f101acb4acf7d945e69f0d6ce76a81113", // master commit
        (1, 1),
        vec!["d592490c4386cdb3373dd93af04d563de199b2fb"], // parent of master, merge commit
    )
    .await?;
    Ok(())
}

#[fbinit::test]
async fn test_location_to_changeset_id_invalid_req(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = UnsharedMergeEven::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    let known_cs_id =
        resolve_cs_id(&ctx, &blobrepo, "7fe9947f101acb4acf7d945e69f0d6ce76a81113").await?;
    seed(&ctx, &blobrepo, &conns, known_cs_id).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;
    assert!(
        sc.location_to_many_changeset_ids(&ctx, Location::new(known_cs_id, 1u64), 2u64)
            .await
            .is_err()
    );
    // TODO(T74320664): Ideally LocationToHash should error when asked to go over merge commit.
    // The parents order is not well defined enough for this not to be ambiguous.
    assert!(
        sc.location_to_many_changeset_ids(&ctx, Location::new(known_cs_id, 2u64), 1u64)
            .await
            .is_ok()
    );
    let second_commit =
        resolve_cs_id(&ctx, &blobrepo, "1700524113b1a3b1806560341009684b4378660b").await?;
    assert!(
        sc.location_to_many_changeset_ids(&ctx, Location::new(second_commit, 1u64), 2u64)
            .await
            .is_err()
    );
    assert!(
        sc.location_to_many_changeset_ids(&ctx, Location::new(second_commit, 2u64), 1u64)
            .await
            .is_err()
    );
    Ok(())
}

async fn validate_changeset_id_to_location(
    ctx: CoreContext,
    blobrepo: BlobRepo,
    server_master: &'static str,
    head: &'static str,
    hg_id: &'static str,
    expected: Option<Location<&'static str>>,
) -> Result<()> {
    let server_master = resolve_cs_id(&ctx, &blobrepo, server_master).await?;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    seed(&ctx, &blobrepo, &conns, server_master).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let head_cs = resolve_cs_id(&ctx, &blobrepo, head).await?;
    let cs_id = resolve_cs_id(&ctx, &blobrepo, hg_id).await?;

    let expected_location = match expected {
        None => None,
        Some(location) => {
            let cs_location = location
                .and_then_descendant(|d| resolve_cs_id(&ctx, &blobrepo, d))
                .await?;
            Some(cs_location)
        }
    };

    let answer = sc
        .changeset_id_to_location(&ctx, vec![head_cs], cs_id)
        .await?;
    assert_eq!(answer, expected_location);

    Ok(())
}

#[fbinit::test]
async fn test_changeset_id_to_location(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    validate_changeset_id_to_location(
        ctx.clone(),
        Linear::getrepo(fb).await,
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // master commit, message: modified 10
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // client head == master_commit
        "0ed509bf086fadcb8a8a5384dc3b550729b0fc17", // message: added 7
        Some(Location::new("79a13814c5ce7330173ec04d279bf95ab3f652fb", 4)),
    )
    .await?;
    validate_changeset_id_to_location(
        ctx.clone(),
        Linear::getrepo(fb).await,
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // master commit, message: modified 10
        "3c15267ebf11807f3d772eb891272b911ec68759", // message: added 9
        "0ed509bf086fadcb8a8a5384dc3b550729b0fc17", // message: added 7
        Some(Location::new("3c15267ebf11807f3d772eb891272b911ec68759", 2)),
    )
    .await?;
    validate_changeset_id_to_location(
        ctx.clone(),
        Linear::getrepo(fb).await,
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // master commit, message: modified 10
        "3c15267ebf11807f3d772eb891272b911ec68759", // message: added 9
        "79a13814c5ce7330173ec04d279bf95ab3f652fb", // our master, client doesn't know
        None,
    )
    .await?;
    validate_changeset_id_to_location(
        ctx.clone(),
        MergeUneven::getrepo(fb).await,
        "d35b1875cdd1ed2c687e86f1604b9d7e989450cb", // master, message: Merge two branches
        "d35b1875cdd1ed2c687e86f1604b9d7e989450cb", // client head == master
        "fc2cef43395ff3a7b28159007f63d6529d2f41ca", // message: Add 4
        Some(Location::new("264f01429683b3dd8042cb3979e8bf37007118bc", 2)), // message: add 5
    )
    .await?;
    validate_changeset_id_to_location(
        ctx.clone(),
        UnsharedMergeEven::getrepo(fb).await,
        "7fe9947f101acb4acf7d945e69f0d6ce76a81113", // master commit
        "33fb49d8a47b29290f5163e30b294339c89505a2", // head merged branch, message: Add 5
        "163adc0d0f5d2eb0695ca123addcb92bab202096", // message: Add 3
        Some(Location::new("33fb49d8a47b29290f5163e30b294339c89505a2", 2)),
    )
    .await?;
    Ok(())
}

#[fbinit::test]
async fn test_changeset_id_to_location_random_hash(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let server_master =
        resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    seed(&ctx, &blobrepo, &conns, server_master).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let head_cs = server_master;
    let random_cs_id = mononoke_types_mocks::changesetid::ONES_CSID;

    let answer = sc
        .changeset_id_to_location(&ctx, vec![head_cs], random_cs_id)
        .await?;
    assert_eq!(answer, None);

    Ok(())
}

#[fbinit::test]
async fn test_changeset_id_to_location_multiple_heads(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = BranchEven::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    // Graph looks like:
    // a -> b -> c
    //  \-> d -> e

    let a = resolve_cs_id(&ctx, &blobrepo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await?;
    let b = resolve_cs_id(&ctx, &blobrepo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await?;
    let c = resolve_cs_id(&ctx, &blobrepo, "b65231269f651cfe784fd1d97ef02a049a37b8a0").await?;
    let d = resolve_cs_id(&ctx, &blobrepo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await?;
    let e = resolve_cs_id(&ctx, &blobrepo, "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5").await?;

    seed(&ctx, &blobrepo, &conns, c).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], a).await?,
        Some(Location::new(c, 2))
    );
    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], d).await?,
        None
    );

    set_bookmark(
        fb,
        &blobrepo,
        "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5",
        BOOKMARK_NAME.clone(),
    )
    .await;
    let tailer = new_tailer_for_tailing(&blobrepo, &conns).await?;
    let _ = tailer.once(&ctx, false).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], b).await?,
        Some(Location::new(c, 1))
    );
    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], d).await?,
        Some(Location::new(e, 1))
    );
    let ambigous = sc
        .changeset_id_to_location(&ctx, vec![c, e], a)
        .await?
        .unwrap();
    assert!(ambigous == Location::new(c, 2) || ambigous == Location::new(e, 2));

    seed(&ctx, &blobrepo, &conns, e).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], a).await?,
        Some(Location::new(e, 2)),
    );
    assert_eq!(
        sc.changeset_id_to_location(&ctx, vec![c, e], b).await?,
        None
    );

    Ok(())
}

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

    {
        // linear
        let inner = Linear::get_inner_repo(fb).await;
        let blobrepo = inner.as_blob_repo();
        let sc = inner.segmented_changelog();

        let known_cs =
            resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
        let distance: u64 = 4;
        let answer = sc
            .location_to_changeset_id(&ctx, Location::new(known_cs, distance))
            .await?;
        let expected_cs =
            resolve_cs_id(&ctx, &blobrepo, "0ed509bf086fadcb8a8a5384dc3b550729b0fc17").await?;
        assert_eq!(answer, expected_cs);
    }
    {
        // merge_uneven
        let inner = MergeUneven::get_inner_repo(fb).await;
        let blobrepo = inner.as_blob_repo();
        let sc = inner.segmented_changelog();

        let known_cs =
            resolve_cs_id(&ctx, &blobrepo, "264f01429683b3dd8042cb3979e8bf37007118bc").await?;
        let distance: u64 = 5;
        let answer = sc
            .location_to_changeset_id(&ctx, Location::new(known_cs, distance))
            .await?;
        let expected_cs =
            resolve_cs_id(&ctx, &blobrepo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await?;
        assert_eq!(answer, expected_cs);
    }

    Ok(())
}

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

    let blobrepo1 = Linear::getrepo_with_id(fb, RepositoryId::new(1)).await;
    let known_cs1 =
        resolve_cs_id(&ctx, &blobrepo1, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    seed(&ctx, &blobrepo1, &conns, known_cs1).await?;
    let sc1 = load_owned(&ctx, &blobrepo1, &conns).await?;

    let blobrepo2 = MergeEven::getrepo_with_id(fb, RepositoryId::new(2)).await;
    let known_cs2 =
        resolve_cs_id(&ctx, &blobrepo2, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await?;
    seed(&ctx, &blobrepo2, &conns, known_cs2).await?;
    let sc2 = load_owned(&ctx, &blobrepo2, &conns).await?;

    let distance: u64 = 4;
    let answer = sc1
        .location_to_changeset_id(&ctx, Location::new(known_cs1, distance))
        .await?;
    let expected_cs_id =
        resolve_cs_id(&ctx, &blobrepo1, "0ed509bf086fadcb8a8a5384dc3b550729b0fc17").await?;
    assert_eq!(answer, expected_cs_id);

    let distance: u64 = 2;
    let answer = sc2
        .location_to_changeset_id(&ctx, Location::new(known_cs2, distance))
        .await?;
    let expected_cs_id =
        resolve_cs_id(&ctx, &blobrepo2, "d7542c9db7f4c77dab4b315edd328edf1514952f").await?;
    assert_eq!(answer, expected_cs_id);

    Ok(())
}

#[fbinit::test]
async fn test_on_demand_update_commit_location_to_changeset_ids(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let inner = Linear::get_inner_repo(fb).await;
    let blobrepo = inner.as_blob_repo();

    // commit modified10 (11)
    let cs11 = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    // commit 10
    let cs10 = resolve_cs_id(&ctx, &blobrepo, "a5ffa77602a066db7d5cfb9fb5823a0895717c5a").await?;
    // commit 6
    let cs6 = resolve_cs_id(&ctx, &blobrepo, "eed3a8c0ec67b6a6fe2eb3543334df3f0b4f202b").await?;

    // commit 5
    let cs5 = resolve_cs_id(&ctx, &blobrepo, "cb15ca4a43a59acff5388cea9648c162afde8372").await?;

    let sc = inner.segmented_changelog();
    let answer = try_join_all(vec![
        sc.location_to_changeset_id(&ctx, Location::new(cs10, 5)),
        sc.location_to_changeset_id(&ctx, Location::new(cs6, 1)),
        sc.location_to_changeset_id(&ctx, Location::new(cs11, 6)),
    ])
    .await?;
    assert_eq!(answer, vec![cs5, cs5, cs5]);

    // Recreate the test repo to reset the segmented changelog
    let inner = Linear::get_inner_repo(fb).await;
    let sc = inner.segmented_changelog();
    let answer = try_join_all(vec![
        sc.changeset_id_to_location(&ctx, vec![cs10], cs5),
        sc.changeset_id_to_location(&ctx, vec![cs6], cs5),
        sc.changeset_id_to_location(&ctx, vec![cs11], cs5),
        sc.changeset_id_to_location(
            &ctx,
            vec![cs11],
            mononoke_types_mocks::changesetid::ONES_CSID,
        ),
    ])
    .await?;
    assert_eq!(
        answer,
        vec![
            Some(Location::new(cs10, 5)),
            Some(Location::new(cs6, 1)),
            Some(Location::new(cs11, 6)),
            None,
        ]
    );

    Ok(())
}

#[fbinit::test]
async fn test_incremental_update_with_desync_iddag(fb: FacebookInit) -> Result<()> {
    // In this test we first build an demand update segmented changelog with an empty idmap and
    // then we reuse the idmap with another on demand update segmented changelog that starts off
    // with an empty iddag.
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    let idmap: Arc<dyn IdMap> = Arc::new(SqlIdMap::new(
        conns.0.clone(),
        Arc::new(NoReplicaLagMonitor()),
        blobrepo.repo_identity().id(),
        IdMapVersion(0),
    ));
    let new_sc = || {
        OnDemandUpdateSegmentedChangelog::new(
            ctx.clone(),
            blobrepo.repo_identity().id(),
            InProcessIdDag::new_in_process(),
            Arc::clone(&idmap),
            blobrepo.changeset_fetcher_arc(),
            blobrepo.bookmarks_arc(),
            vec![Some(BOOKMARK_NAME.clone()).into()],
            None,
        )
    };

    let master_cs =
        resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;

    let initial = new_sc()?;

    let cs7 = resolve_cs_id(&ctx, &blobrepo, "0ed509bf086fadcb8a8a5384dc3b550729b0fc17").await?;
    let distance: u64 = 4;
    let answer = initial
        .location_to_changeset_id(&ctx, Location::new(master_cs, distance))
        .await?;
    assert_eq!(answer, cs7);

    let second = new_sc()?;

    let cs3 = resolve_cs_id(&ctx, &blobrepo, "607314ef579bd2407752361ba1b0c1729d08b281").await?;
    let answer = second
        .location_to_changeset_id(&ctx, Location::new(cs7, distance))
        .await?;
    assert_eq!(answer, cs3);

    let answer = second
        .location_to_changeset_id(&ctx, Location::new(master_cs, distance))
        .await?;
    assert_eq!(answer, cs7);

    Ok(())
}

#[fbinit::test]
async fn test_clone_data(fb: FacebookInit) -> Result<()> {
    // In this test we first build a dag from scratch and then we reuse the idmap in an ondemand
    // segmented changelog that starts off with an empty iddag.
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let head = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    seed(&ctx, &blobrepo, &conns, head).await?;

    let sc = load_owned(&ctx, &blobrepo, &conns).await?;

    let clone_data = sc.clone_data(&ctx).await?;

    let mut new_iddag = InProcessIdDag::new_in_process();
    new_iddag.build_segments_from_prepared_flat_segments(&clone_data.0.flat_segments)?;
    let new_owned = OwnedSegmentedChangelog::new(new_iddag, sc.idmap.clone());

    let distance: u64 = 4;
    let answer = new_owned
        .location_to_changeset_id(&ctx, Location::new(head, distance))
        .await?;
    let expected_cs =
        resolve_cs_id(&ctx, &blobrepo, "0ed509bf086fadcb8a8a5384dc3b550729b0fc17").await?;
    assert_eq!(answer, expected_cs);

    Ok(())
}

#[fbinit::test]
async fn test_caching(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let cache_handlers = CacheHandlers::new(CacheHandlerFactory::Mocked);
    let cs_to_dag_handler = cache_handlers.cs_to_dag.clone();
    let dag_to_cs_handler = cache_handlers.dag_to_cs.clone();

    // It's easier to reason about cache hits and sets when the dag is already built
    let head = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    seed(&ctx, &blobrepo, &conns, head).await?;

    let iddag = load_iddag(&ctx, &blobrepo, &conns).await?;
    let sc_version = load_sc_version(&ctx, blobrepo.repo_identity().id(), &conns).await?;
    let idmap = IdMapFactory::new(
        conns.0.clone(),
        Arc::new(NoReplicaLagMonitor()),
        blobrepo.repo_identity().id(),
    )
    .with_cache_handlers(cache_handlers)
    .for_server(&ctx, sc_version.idmap_version, &iddag)?;
    let sc = OwnedSegmentedChangelog::new(iddag, idmap);

    let distance: u64 = 1;
    let _ = sc
        .location_to_changeset_id(&ctx, Location::new(head, distance))
        .await?;

    let cs_to_dag_stats = || {
        cs_to_dag_handler
            .mock_store()
            .expect("mock handler has mock store")
            .stats()
    };
    let dag_to_cs_stats = || {
        dag_to_cs_handler
            .mock_store()
            .expect("mock handler has mock store")
            .stats()
    };

    assert_eq!(cs_to_dag_stats().gets, 1);
    assert_eq!(cs_to_dag_stats().misses, 1);
    assert_eq!(cs_to_dag_stats().sets, 1);

    assert_eq!(dag_to_cs_stats().gets, 1);
    assert_eq!(dag_to_cs_stats().misses, 1);
    assert_eq!(dag_to_cs_stats().sets, 1);

    let _ = sc
        .location_to_changeset_id(&ctx, Location::new(head, distance))
        .await?;

    assert_eq!(cs_to_dag_stats().gets, 2);
    assert_eq!(cs_to_dag_stats().hits, 1);

    assert_eq!(dag_to_cs_stats().gets, 2);
    assert_eq!(dag_to_cs_stats().hits, 1);

    Ok(())
}

#[fbinit::test]
async fn test_periodic_update(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let bookmark_name = BookmarkKey::new("periodic_update")?;

    let start_hg_id = "607314ef579bd2407752361ba1b0c1729d08b281";
    let start_cs = resolve_cs_id(&ctx, &blobrepo, start_hg_id).await?;
    set_bookmark(fb, &blobrepo, start_hg_id, bookmark_name.clone()).await;

    tokio::time::pause(); // TODO: pause only works with the `current_thread` Runtime.

    let on_demand = OnDemandUpdateSegmentedChangelog::new(
        ctx.clone(),
        blobrepo.repo_identity().id(),
        InProcessIdDag::new_in_process(),
        Arc::new(ConcurrentMemIdMap::new()),
        blobrepo.changeset_fetcher_arc(),
        blobrepo.bookmarks_arc(),
        vec![Some(bookmark_name.clone()).into()],
        None,
    )?;
    let sc =
        Arc::new(on_demand).with_periodic_update_to_master_bookmark(&ctx, Duration::from_secs(5));

    tokio::time::advance(Duration::from_secs(10)).await;
    // wait time is over and the first tick is ready to be scheduled
    sc.wait_for_update().await;

    // We assume that clone_data will not update the graph in any form.
    assert_eq!(sc.head(&ctx).await?, start_cs);

    let new_hg_id = "79a13814c5ce7330173ec04d279bf95ab3f652fb";
    let new_cs = resolve_cs_id(&ctx, &blobrepo, new_hg_id).await?;
    set_bookmark(fb, &blobrepo, new_hg_id, bookmark_name.clone()).await;

    tokio::time::advance(Duration::from_secs(5)).await;
    // second tick is ready to be scheduled
    sc.wait_for_update().await;
    assert_eq!(sc.head(&ctx).await?, new_cs);

    Ok(())
}

#[fbinit::test]
async fn test_seeder_tailer_and_load(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let start_hg_id = "607314ef579bd2407752361ba1b0c1729d08b281"; // commit 4
    let start_cs_id = resolve_cs_id(&ctx, &blobrepo, start_hg_id).await?;

    seed(&ctx, &blobrepo, &conns, start_cs_id).await?;

    let sc = load_owned(&ctx, &blobrepo, &conns).await?;
    assert_eq!(sc.head(&ctx).await?, start_cs_id);

    let tailer = new_tailer_for_tailing(&blobrepo, &conns).await?;
    let _ = tailer.once(&ctx, false).await?;
    let sc = load_owned(&ctx, &blobrepo, &conns).await?;
    let master = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    assert_eq!(sc.head(&ctx).await?, master);

    Ok(())
}

#[fbinit::test]
async fn test_periodic_reload(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Arc::new(Linear::getrepo(fb).await);
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let start_hg_id = "607314ef579bd2407752361ba1b0c1729d08b281"; // commit 4
    let start_cs_id = resolve_cs_id(&ctx, &blobrepo, start_hg_id).await?;

    seed(&ctx, &blobrepo, &conns, start_cs_id).await?;

    tokio::time::pause();
    let manager = get_manager(&blobrepo, &conns, vec![], SegmentedChangelogType::Owned).await?;
    let sc = PeriodicReloadSegmentedChangelog::start_from_manager(
        &ctx,
        Duration::from_secs(5),
        manager,
        blobrepo.repo_identity().name().to_string(),
    )
    .await?;

    assert_eq!(sc.head(&ctx).await?, start_cs_id);

    // Try waiting for segmented changelog update without tailer running.  This
    // should fail as we don't update the SC unless there's an update.
    assert!(
        tokio::time::timeout(Duration::from_secs(15), sc.wait_for_update())
            .await
            .is_err()
    );

    // Start the tailer.
    let tailer = new_tailer_for_tailing(&blobrepo, &conns).await?;
    let _ = tailer.once(&ctx, false).await?;
    // There was an update - wait_for_update should return.
    tokio::time::timeout(Duration::from_secs(15), sc.wait_for_update()).await?;
    let master = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;
    assert_eq!(sc.head(&ctx).await?, master);

    // No updates should happen afterwards because there are no new commits.
    assert!(
        tokio::time::timeout(Duration::from_secs(45), sc.wait_for_update())
            .await
            .is_err()
    );

    // Force the reload should trigger even if there are not updates at all.
    let tunables_override = Arc::new(tunables::MononokeTunables::default());
    tunables_override.update_by_repo_ints(&hashmap! {
        blobrepo.repo_identity().name().to_string() => hashmap! {
            "segmented_changelog_force_reload".to_string() => 2,
        },
    });
    tunables_override.update_ints(&hashmap! {
        "segmented_changelog_force_reload_jitter_secs".to_string() => 5,
    });
    override_tunables(Some(tunables_override.clone()));
    tokio::time::timeout(Duration::from_secs(45), sc.wait_for_update()).await?;
    override_tunables(None);
    Ok(())
}

#[fbinit::test]
async fn test_manager_check_if_indexed(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Arc::new(Linear::getrepo(fb).await);
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;

    let start_cs_id =
        resolve_cs_id(&ctx, &blobrepo, "607314ef579bd2407752361ba1b0c1729d08b281").await?;
    let master = resolve_cs_id(&ctx, &blobrepo, "79a13814c5ce7330173ec04d279bf95ab3f652fb").await?;

    seed(&ctx, &blobrepo, &conns, start_cs_id).await?;

    tokio::time::pause();
    let manager = get_manager(&blobrepo, &conns, vec![], SegmentedChangelogType::Owned).await?;

    // Changeset is not part of segmented changelog before tailer run.
    let version = manager.latest_version(&ctx).await?;
    assert!(
        !manager
            .check_if_changeset_indexed(&ctx, &version, master)
            .await?
    );

    let tailer = new_tailer_for_tailing(&blobrepo, &conns).await?;
    let _ = tailer.once(&ctx, false).await?;
    // Changeset is part of segmented changelog after tailer run.
    let version = manager.latest_version(&ctx).await?;
    assert!(
        manager
            .check_if_changeset_indexed(&ctx, &version, master)
            .await?
    );
    Ok(())
}

#[fbinit::test]
async fn test_mismatched_heads(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let inner = BranchEven::get_inner_repo(fb).await;
    let blobrepo = inner.as_blob_repo();

    let dag = inner.segmented_changelog();
    let h1 = resolve_cs_id(&ctx, &blobrepo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await?;
    let h1_parent =
        resolve_cs_id(&ctx, &blobrepo, "b65231269f651cfe784fd1d97ef02a049a37b8a0").await?;

    assert_eq!(
        dag.changeset_id_to_location(&ctx, vec![h1], h1_parent)
            .await?,
        Some(Location::new(h1, 1))
    );

    // should fail with the small limit for traversing from client heads
    let h2 = resolve_cs_id(&ctx, &blobrepo, "16839021e338500b3cf7c9b871c8a07351697d68").await?;

    let tunables = tunables::MononokeTunables::default();
    tunables.update_ints(&hashmap! {
        "segmented_changelog_client_max_commits_to_traverse".to_string() => 2,
    });
    let f = dag.changeset_id_to_location(&ctx, vec![h1, h2], h1_parent);

    let err = with_tunables_async(tunables, f.boxed())
        .await
        .err()
        .unwrap();
    assert!(err.is::<crate::MismatchedHeadsError>());

    // should succeed as the client head not far from the commits in SC IdMap
    let h2 = resolve_cs_id(&ctx, &blobrepo, "16839021e338500b3cf7c9b871c8a07351697d68").await?;

    let tunables = tunables::MononokeTunables::default();
    tunables.update_ints(&hashmap! {
        "segmented_changelog_client_max_commits_to_traverse".to_string() => 100,
    });
    let f = dag.changeset_id_to_location(&ctx, vec![h1, h2], h1_parent);

    assert_eq!(
        with_tunables_async(tunables, f.boxed()).await?,
        Some(Location::new(h1, 1))
    );

    Ok(())
}
#[fbinit::test]
async fn test_prefetched_seeding(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    let repo_id = blobrepo.repo_identity().id();

    let first_cs_entry =
        fetch_cs_entry(&ctx, &blobrepo, "2d7d4ba9ce0a6ffd222de7785b249ead9c51c536").await?;
    let second_cs_entry =
        fetch_cs_entry(&ctx, &blobrepo, "3e0e761030db6e479a7fb58b12881883f9f8c63f").await?;
    let last_cs =
        resolve_cs_id(&ctx, &blobrepo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await?;

    seed_with_prefetched(
        &ctx,
        &blobrepo,
        &conns,
        vec![last_cs],
        Some(vec![first_cs_entry, second_cs_entry]),
    )
    .await?;

    // Check that new entries were added to the idmap
    let idmap = load_idmap(&ctx, repo_id, &conns).await?;
    assert!(idmap.find_dag_id(&ctx, last_cs).await?.is_some());

    Ok(())
}

#[fbinit::test]
async fn test_seeding_with_included_bonsais(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let blobrepo = Linear::getrepo(fb).await;
    let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
    let repo_id = blobrepo.repo_identity().id();

    let first_cs_id =
        resolve_cs_id(&ctx, &blobrepo, "2d7d4ba9ce0a6ffd222de7785b249ead9c51c536").await?;
    let last_cs =
        resolve_cs_id(&ctx, &blobrepo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await?;

    let cs_id_to_include = CreateCommitContext::new(&ctx, &blobrepo, vec![first_cs_id])
        .add_file("somefile", "file")
        .commit()
        .await?;

    seed_with_prefetched(
        &ctx,
        &blobrepo,
        &conns,
        vec![last_cs, cs_id_to_include],
        None,
    )
    .await?;

    // Check that cs_id_to_include was added to idmap
    let idmap = load_idmap(&ctx, repo_id, &conns).await?;
    let dag_id = idmap.find_dag_id(&ctx, cs_id_to_include).await?.unwrap();
    let iddag = load_iddag(&ctx, &blobrepo, &conns).await?;
    assert!(iddag.find_flat_segment_including_id(dag_id)?.is_some());

    Ok(())
}


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

use std::fmt;

use mononoke_types::hash;
use mononoke_types::hash::Blake2;
use sql::mysql;

#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd)]
#[derive(mysql::OptTryFromRowField)]
pub struct IdMapVersion(pub u64);

impl fmt::Display for IdMapVersion {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl IdMapVersion {
    pub fn bump(&self) -> Self {
        Self(self.0 + 1)
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[derive(mysql::OptTryFromRowField)]
pub struct IdDagVersion(pub Blake2);

impl IdDagVersion {
    pub fn from_serialized_bytes<B: AsRef<[u8]>>(bytes: B) -> Self {
        let mut blake2_builder = hash::Context::new("iddag_version".as_bytes());
        blake2_builder.update(bytes);
        Self(blake2_builder.finish())
    }
}

impl fmt::Display for IdDagVersion {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct SegmentedChangelogVersion {
    pub iddag_version: IdDagVersion,
    pub idmap_version: IdMapVersion,
}

impl SegmentedChangelogVersion {
    pub fn new(iddag_version: IdDagVersion, idmap_version: IdMapVersion) -> Self {
        Self {
            iddag_version,
            idmap_version,
        }
    }
}

impl From<(IdDagVersion, IdMapVersion)> for SegmentedChangelogVersion {
    fn from(t: (IdDagVersion, IdMapVersion)) -> Self {
        SegmentedChangelogVersion::new(t.0, t.1)
    }
}


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

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

use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bookmarks::BookmarkCategory;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkPagination;
use bookmarks::BookmarkPrefix;
use bookmarks::Bookmarks;
use bookmarks::Freshness;
use context::CoreContext;
use futures::future;
use futures::stream;
use futures::stream::TryStreamExt;
use metaconfig_types::SegmentedChangelogConfig;
use metaconfig_types::SegmentedChangelogHeadConfig;
use mononoke_types::ChangesetId;
use slog::info;

use crate::dag::NameDagBuilder;
use crate::dag::VertexListWithOptions;
use crate::dag::VertexName;
use crate::dag::VertexOptions;
use crate::idmap::vertex_name_from_cs_id;
use crate::idmap::IdMap;
use crate::idmap::IdMapWrapper;
use crate::Group;
use crate::InProcessIdDag;

pub type SeedHead = SegmentedChangelogHeadConfig;

#[async_trait::async_trait]
trait IntoVertexList {
    async fn into_vertex_list(
        &self,
        ctx: &CoreContext,
        bookmarks: &dyn Bookmarks,
    ) -> Result<VertexListWithOptions>;
}

#[async_trait::async_trait]
impl IntoVertexList for SeedHead {
    async fn into_vertex_list(
        &self,
        ctx: &CoreContext,
        bookmarks: &dyn Bookmarks,
    ) -> Result<VertexListWithOptions> {
        match self {
            Self::Changeset(id) => Ok(VertexListWithOptions::from(vec![head_with_options(id)])),
            Self::AllPublicBookmarksExcept(exceptions) => {
                all_bookmarks_except_with_options(ctx, exceptions, bookmarks).await
            }
            Self::Bookmark(name) => bookmark_with_options(ctx, name, bookmarks).await,
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum JobType {
    Background,
    Server,
}

pub fn seedheads_from_config(
    ctx: &CoreContext,
    config: &SegmentedChangelogConfig,
    job_type: JobType,
) -> Result<Vec<SeedHead>> {
    let mut heads = config.heads_to_include.clone();

    if job_type == JobType::Background {
        heads.extend(config.extra_heads_to_include_in_background_jobs.clone());
    }

    info!(
        ctx.logger(),
        "Using the following segmented changelog heads: {:?}", heads
    );
    Ok(heads)
}

pub async fn vertexlist_from_seedheads(
    ctx: &CoreContext,
    heads: &[SeedHead],
    bookmarks: &dyn Bookmarks,
) -> Result<VertexListWithOptions> {
    let heads_with_options = stream::iter(heads.iter().map(Result::Ok))
        .try_fold(VertexListWithOptions::default(), {
            move |acc, head| async move {
                Ok::<_, Error>(acc.chain(head.into_vertex_list(ctx, bookmarks).await?))
            }
        })
        .await?;

    Ok(heads_with_options)
}

pub type ServerNameDag = crate::dag::namedag::AbstractNameDag<InProcessIdDag, IdMapWrapper, (), ()>;

/// Convert a server IdDag and IdMap to a NameDag
/// Note: you will need to call NameDag::map().flush_writes
/// to write out updates to the IdMap
pub fn server_namedag(
    ctx: CoreContext,
    iddag: InProcessIdDag,
    idmap: Arc<dyn IdMap>,
) -> Result<ServerNameDag> {
    let idmap = IdMapWrapper::new(ctx, idmap);
    NameDagBuilder::new_with_idmap_dag(idmap, iddag)
        .build()
        .map_err(anyhow::Error::from)
}

fn head_with_options(head: &ChangesetId) -> (VertexName, VertexOptions) {
    let mut options = VertexOptions::default();
    options.reserve_size = 1 << 26;
    options.highest_group = Group::MASTER;
    (vertex_name_from_cs_id(head), options)
}

async fn all_bookmarks_except_with_options(
    ctx: &CoreContext,
    exceptions: &[BookmarkKey],
    bookmarks: &dyn Bookmarks,
) -> Result<VertexListWithOptions> {
    let exceptions: HashSet<_> = exceptions.iter().cloned().collect();
    Ok(bookmarks
        .list(
            ctx.clone(),
            Freshness::MaybeStale,
            &BookmarkPrefix::empty(),
            BookmarkCategory::ALL,
            BookmarkKind::ALL_PUBLISHING,
            &BookmarkPagination::FromStart,
            u64::MAX,
        )
        .try_filter_map(|(bookmark, cs_id)| {
            let res = if exceptions.contains(bookmark.key()) {
                None
            } else {
                Some(cs_id)
            };
            future::ready(Ok(res))
        })
        .map_ok(|cs| head_with_options(&cs))
        .try_collect::<Vec<_>>()
        .await?
        .into())
}

async fn bookmark_with_options(
    ctx: &CoreContext,
    bookmark: &BookmarkKey,
    bookmarks: &dyn Bookmarks,
) -> Result<VertexListWithOptions> {
    let cs = bookmarks
        .get(ctx.clone(), bookmark)
        .await
        .with_context(|| format!("error while fetching changeset for bookmark {}", bookmark))?
        .ok_or_else(move || format_err!("'{}' bookmark could not be found", bookmark))?;
    Ok(VertexListWithOptions::from(vec![head_with_options(&cs)]))
}

#[cfg(test)]
mod tests {
    use blobrepo::BlobRepo;
    use bookmarks::BookmarksRef;
    use fbinit::FacebookInit;
    use fixtures::set_bookmark;
    use fixtures::BranchWide;
    use fixtures::TestRepoFixture;

    use super::*;

    async fn prep_branch_wide_repo(fb: FacebookInit) -> Result<Arc<BlobRepo>> {
        let blobrepo = BranchWide::getrepo(fb).await;
        let second = BookmarkKey::new("second")?;
        set_bookmark(
            fb,
            &blobrepo,
            "04decbb0d1a65789728250ddea2fe8d00248e01c",
            second,
        )
        .await;
        let third = BookmarkKey::new("third")?;
        set_bookmark(
            fb,
            &blobrepo,
            "c27ef5b7f15e9930e5b93b1f32cc2108a2aabe12",
            third,
        )
        .await;
        Ok(Arc::new(blobrepo))
    }

    #[fbinit::test]
    async fn test_bookmark_with_options(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo = prep_branch_wide_repo(fb).await?;
        let second = BookmarkKey::new("second")?;

        let res = bookmark_with_options(&ctx, &second, repo.bookmarks()).await?;
        assert_eq!(
            res.vertexes(),
            vec![VertexName::from_hex(
                b"5ec506306edb84a4a47f901a55cedeec3113eb118bfae119982f45382481e3dc"
            )?,]
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_all_bookmarks_except_with_options(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo = prep_branch_wide_repo(fb).await?;

        let res =
            all_bookmarks_except_with_options(&ctx, vec![].as_slice(), repo.bookmarks()).await?;
        assert_eq!(
            res.vertexes(),
            vec![
                VertexName::from_hex(
                    b"56da5b997e27f2f9020f6ff2d87b321774369e23579bd2c4ce675efad363f4f4"
                )?,
                VertexName::from_hex(
                    b"5ec506306edb84a4a47f901a55cedeec3113eb118bfae119982f45382481e3dc"
                )?,
                VertexName::from_hex(
                    b"7097e8d1e72af16e8135047d8693fb381246be1bc74c1b6c0cb013fc05331fc1"
                )?,
            ]
        );

        let res = all_bookmarks_except_with_options(
            &ctx,
            vec![BookmarkKey::new("second")?].as_slice(),
            repo.bookmarks(),
        )
        .await?;
        assert_eq!(
            res.vertexes(),
            vec![
                VertexName::from_hex(
                    b"56da5b997e27f2f9020f6ff2d87b321774369e23579bd2c4ce675efad363f4f4"
                )?,
                VertexName::from_hex(
                    b"7097e8d1e72af16e8135047d8693fb381246be1bc74c1b6c0cb013fc05331fc1"
                )?,
            ]
        );
        Ok(())
    }
}


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

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::RepositoryId;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use stats::prelude::*;

use crate::logging::log_new_segmented_changelog_version;
use crate::types::IdDagVersion;
use crate::types::IdMapVersion;
use crate::types::SegmentedChangelogVersion;

define_stats! {
    prefix = "mononoke.segmented_changelog.sql_version_store";
    set: timeseries(Sum),
    update: timeseries(Sum),
    get: timeseries(Sum),
}

/// Specifies the versions for the latest SegmentedChangelogVersion. The version contains IdDag and
/// IdMap versions.  The IdDag version can be loaded directly from the blobstore and the IdMap
/// version ties the IdDag back to the bonsai changesets.
pub struct SegmentedChangelogVersionStore {
    connections: SqlConnections,
    repo_id: RepositoryId,
}

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

    pub async fn set(&self, ctx: &CoreContext, version: SegmentedChangelogVersion) -> Result<()> {
        STATS::set.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);
        SetVersion::query(
            &self.connections.write_connection,
            &self.repo_id,
            &version.iddag_version,
            &version.idmap_version,
        )
        .await
        .with_context(|| format!("failed to set segmented changelog version {:?}", version))?;
        log_new_segmented_changelog_version(ctx, self.repo_id, version);
        Ok(())
    }

    pub async fn update(
        &self,
        ctx: &CoreContext,
        version: SegmentedChangelogVersion,
    ) -> Result<()> {
        STATS::update.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);
        let result = UpdateVersion::query(
            &self.connections.write_connection,
            &self.repo_id,
            &version.iddag_version,
            &version.idmap_version,
        )
        .await
        .with_context(|| {
            format!(
                "failed to update segmented changelog version ({:?})",
                version
            )
        })?;
        if result.affected_rows() == 0 {
            bail!(
                "no valid rows to update for segmented changelog version {:?}",
                version
            );
        }
        log_new_segmented_changelog_version(ctx, self.repo_id, version);
        Ok(())
    }

    pub async fn get(&self, ctx: &CoreContext) -> Result<Option<SegmentedChangelogVersion>> {
        STATS::get.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let rows = SelectVersion::query(&self.connections.read_connection, &self.repo_id).await?;
        Ok(rows.into_iter().next().map(|r| r.into()))
    }
}

mononoke_queries! {
    write SetVersion(
        repo_id: RepositoryId,
        iddag_version: IdDagVersion,
        idmap_version: IdMapVersion
    ) {
        none,
        "
        REPLACE INTO segmented_changelog_version (repo_id, iddag_version, idmap_version)
        VALUES ({repo_id}, {iddag_version}, {idmap_version})
        "
    }

    write UpdateVersion(
        repo_id: RepositoryId,
        iddag_version: IdDagVersion,
        idmap_version: IdMapVersion,
    ) {
        none,
        "
        UPDATE segmented_changelog_version
        SET iddag_version = {iddag_version}
        WHERE repo_id = {repo_id} AND idmap_version = {idmap_version}
        "
    }

    read SelectVersion(repo_id: RepositoryId) -> (IdDagVersion, IdMapVersion) {
        "
        SELECT iddag_version, idmap_version
        FROM segmented_changelog_version
        WHERE repo_id = {repo_id}
        "
    }
}

#[cfg(test)]
mod tests {
    use fbinit::FacebookInit;
    use sql_construct::SqlConstruct;

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

    #[fbinit::test]
    async fn test_set_more_than_one_repo(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
        let version_repo1 =
            SegmentedChangelogVersionStore::new(conns.0.clone(), RepositoryId::new(1));
        let version_repo2 =
            SegmentedChangelogVersionStore::new(conns.0.clone(), RepositoryId::new(2));

        assert_eq!(version_repo1.get(&ctx).await?, None);
        assert_eq!(version_repo2.get(&ctx).await?, None);
        let version11 = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"1"),
            IdMapVersion(1),
        );
        let version23 = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"2"),
            IdMapVersion(3),
        );
        version_repo1.set(&ctx, version11).await?;
        assert_eq!(version_repo1.get(&ctx).await?, Some(version11));
        assert_eq!(version_repo2.get(&ctx).await?, None);
        version_repo2.set(&ctx, version23).await?;
        assert_eq!(version_repo1.get(&ctx).await?, Some(version11));
        assert_eq!(version_repo2.get(&ctx).await?, Some(version23));

        Ok(())
    }

    #[fbinit::test]
    async fn test_update(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
        let version_store = SegmentedChangelogVersionStore::new(conns.0, RepositoryId::new(0));

        let vm1 = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"a"),
            IdMapVersion(1),
        );
        let vm1x = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"x"),
            IdMapVersion(1),
        );
        let vm2 = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"b"),
            IdMapVersion(2),
        );
        let vm2y = SegmentedChangelogVersion::new(
            IdDagVersion::from_serialized_bytes(b"y"),
            IdMapVersion(2),
        );

        version_store.set(&ctx, vm1).await?;
        assert_eq!(version_store.get(&ctx).await?, Some(vm1));
        assert!(version_store.update(&ctx, vm2y).await.is_err());
        version_store.update(&ctx, vm1x).await?;
        assert_eq!(version_store.get(&ctx).await?, Some(vm1x));
        version_store.set(&ctx, vm2).await?;
        assert!(version_store.update(&ctx, vm1x).await.is_err());
        assert_eq!(version_store.get(&ctx).await?, Some(vm2));
        version_store.update(&ctx, vm2y).await?;
        assert_eq!(version_store.get(&ctx).await?, Some(vm2y));

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

mod save_store;

pub use self::save_store::IdDagSaveStore;


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

use std::sync::Arc;

use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use mononoke_types::RepositoryId;

use crate::logging::log_new_iddag_version;
use crate::types::IdDagVersion;
use crate::InProcessIdDag;

pub struct IdDagSaveStore {
    repo_id: RepositoryId,
    blobstore: Arc<dyn Blobstore>,
}

impl IdDagSaveStore {
    pub fn new(repo_id: RepositoryId, blobstore: Arc<dyn Blobstore>) -> Self {
        Self { repo_id, blobstore }
    }

    pub async fn find<'a>(
        &'a self,
        ctx: &'a CoreContext,
        iddag_version: IdDagVersion,
    ) -> Result<Option<InProcessIdDag>> {
        let bytes_opt = self
            .blobstore
            .get(ctx, &self.key(iddag_version))
            .await
            .with_context(|| {
                format!(
                    "loading prebuilt segmented changelog iddag version {}",
                    iddag_version.0
                )
            })?;
        let bytes = match bytes_opt {
            None => return Ok(None),
            Some(b) => b,
        };

        let deserialization = move || {
            let iddag: InProcessIdDag = mincode::deserialize(&bytes.into_raw_bytes())?;
            anyhow::Ok(iddag)
        };
        // In tests we can't offload deserialization to spawn blocking because
        // of a tokio bug where in time::pause environment the time
        // auto-advancement won't wait for spawned blocking tasks.
        let iddag = if cfg!(test) {
            deserialization()?
        } else {
            tokio::task::spawn_blocking(deserialization).await??
        };

        Ok(Some(iddag))
    }

    pub async fn load<'a>(
        &'a self,
        ctx: &'a CoreContext,
        iddag_version: IdDagVersion,
    ) -> Result<InProcessIdDag> {
        self.find(ctx, iddag_version).await?.ok_or_else(|| {
            format_err!(
                "Not Found: prebuilt iddag (repo_id: {}, version: {})",
                self.repo_id,
                iddag_version.0,
            )
        })
    }

    pub async fn save<'a>(
        &'a self,
        ctx: &'a CoreContext,
        iddag: &InProcessIdDag,
    ) -> Result<IdDagVersion> {
        let buffer = mincode::serialize(iddag)?;
        let iddag_version = IdDagVersion::from_serialized_bytes(&buffer);
        self.blobstore
            .put(
                ctx,
                self.key(iddag_version),
                BlobstoreBytes::from_bytes(buffer),
            )
            .await
            .context("saving iddag in blobstore")?;
        log_new_iddag_version(ctx, self.repo_id, iddag_version);
        Ok(iddag_version)
    }

    fn key(&self, iddag_version: IdDagVersion) -> String {
        format!("segmented_changelog_iddag.blake2.{}", iddag_version.0)
    }
}


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

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

use abomonation_derive::Abomonation;
use anyhow::Context;
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 futures::future;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use memcache::KeyGen;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;

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

#[derive(Clone)]
pub struct CachedIdMap {
    idmap: Arc<dyn IdMap>,
    cache_handlers: CacheHandlers,
    repo_id: RepositoryId,
    keygen: KeyGen,
}

#[derive(Clone)]
pub struct CacheHandlers {
    pub dag_to_cs: CachelibHandler<ChangesetIdWrapper>,
    pub cs_to_dag: CachelibHandler<DagIdWrapper>,
    pub memcache: MemcacheHandler,
}

impl CacheHandlers {
    pub fn new(cache_handler_factory: CacheHandlerFactory) -> Self {
        Self {
            dag_to_cs: cache_handler_factory.cachelib(),
            cs_to_dag: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
        }
    }
}

impl CachedIdMap {
    pub fn new(
        idmap: Arc<dyn IdMap>,
        cache_handlers: CacheHandlers,
        repo_id: RepositoryId,
        version: IdMapVersion,
    ) -> Self {
        let codever = 0; // bump when logic changes
        let sitever = version.0 as u32;

        let keygen = KeyGen::new("scm.mononoke.segmented_changelog.idmap", codever, sitever);
        Self {
            idmap,
            cache_handlers,
            repo_id,
            keygen,
        }
    }
}

// Number of entries to fetch from DB into cache at a time
const CHUNK_SIZE: usize = 1000;
const PARALLEL_CHUNKS: usize = 10;

#[async_trait]
impl IdMap for CachedIdMap {
    async fn insert_many(
        &self,
        ctx: &CoreContext,
        mappings: Vec<(DagId, ChangesetId)>,
    ) -> Result<()> {
        self.idmap.insert_many(ctx, mappings).await?;
        Ok(())
    }

    async fn find_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> Result<HashMap<DagId, ChangesetId>> {
        let ctx = (ctx, self);
        let res = get_or_fill_chunked(
            &ctx,
            dag_ids.into_iter().collect(),
            CHUNK_SIZE,
            PARALLEL_CHUNKS,
        )
        .await
        .with_context(|| "Error fetching many changeset ids via cache")?
        .into_iter()
        .map(|(k, v)| (k, v.0))
        .collect();
        Ok(res)
    }

    async fn find_many_dag_ids(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        let ctx = (ctx, self, DagIdStaleness::Fresh);
        let res = get_or_fill_chunked(
            &ctx,
            cs_ids.into_iter().collect(),
            CHUNK_SIZE,
            PARALLEL_CHUNKS,
        )
        .await
        .with_context(|| "Error fetching many changeset ids via cache")?
        .into_iter()
        .map(|(k, v)| (k, v.0))
        .collect();
        Ok(res)
    }

    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        let ctx = (ctx, self, DagIdStaleness::MaybeStale);
        let res = get_or_fill_chunked(
            &ctx,
            cs_ids.into_iter().collect(),
            CHUNK_SIZE,
            PARALLEL_CHUNKS,
        )
        .await
        .with_context(|| "Error fetching many changeset ids via cache")?
        .into_iter()
        .map(|(k, v)| (k, v.0))
        .collect();
        Ok(res)
    }

    async fn get_last_entry(&self, ctx: &CoreContext) -> Result<Option<(DagId, ChangesetId)>> {
        self.idmap.get_last_entry(ctx).await
    }

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

type ChangesetIdCacheRequest<'a> = (&'a CoreContext, &'a CachedIdMap);

#[derive(Clone, Copy, Debug, Abomonation)]
pub struct ChangesetIdWrapper(ChangesetId);

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

impl MemcacheEntity for ChangesetIdWrapper {
    fn serialize(&self) -> Bytes {
        Bytes::copy_from_slice(self.0.as_ref())
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        match ChangesetId::from_bytes(&bytes) {
            Ok(cs_id) => Ok(ChangesetIdWrapper(cs_id)),
            Err(_) => Err(McErrorKind::Deserialization),
        }
    }
}

impl EntityStore<ChangesetIdWrapper> for ChangesetIdCacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<ChangesetIdWrapper> {
        let (_, bag) = self;
        &bag.cache_handlers.dag_to_cs
    }

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

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

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

    caching_ext::impl_singleton_stats!("segmented_changelog.idmap.dag2cs");
}

#[async_trait]
impl KeyedEntityStore<DagId, ChangesetIdWrapper> for ChangesetIdCacheRequest<'_> {
    fn get_cache_key(&self, dag_id: &DagId) -> String {
        let (_, bag) = self;
        format!("{}.dag_id.{}", bag.repo_id, dag_id)
    }

    async fn get_from_db(
        &self,
        keys: HashSet<DagId>,
    ) -> Result<HashMap<DagId, ChangesetIdWrapper>> {
        let (ctx, bag) = self;

        let res = bag
            .idmap
            .find_many_changeset_ids(ctx, keys.into_iter().collect())
            .await?
            .into_iter()
            .map(|(k, v)| (k, v.into()))
            .collect();

        Ok(res)
    }
}

enum DagIdStaleness {
    MaybeStale,
    Fresh,
}

type DagIdCacheRequest<'a> = (&'a CoreContext, &'a CachedIdMap, DagIdStaleness);

#[derive(Clone, Copy, Debug, Abomonation)]
pub struct DagIdWrapper(DagId);

impl From<DagId> for DagIdWrapper {
    fn from(dag_id: DagId) -> Self {
        DagIdWrapper(dag_id)
    }
}

impl MemcacheEntity for DagIdWrapper {
    fn serialize(&self) -> Bytes {
        Bytes::copy_from_slice(&self.0.0.to_be_bytes())
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        let arr = bytes
            .as_ref()
            .try_into()
            .map_err(|_| McErrorKind::Deserialization)?;
        Ok(DagIdWrapper(DagId(u64::from_be_bytes(arr))))
    }
}

impl EntityStore<DagIdWrapper> for DagIdCacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<DagIdWrapper> {
        let (_, bag, _) = self;
        &bag.cache_handlers.cs_to_dag
    }

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

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

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

    caching_ext::impl_singleton_stats!("segmented_changelog.idmap.cs2dag");
}

#[async_trait]
impl KeyedEntityStore<ChangesetId, DagIdWrapper> for DagIdCacheRequest<'_> {
    fn get_cache_key(&self, cs_id: &ChangesetId) -> String {
        let (_, bag, _) = self;
        format!("{}.cs.{}", bag.repo_id, cs_id)
    }

    async fn get_from_db(
        &self,
        keys: HashSet<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagIdWrapper>> {
        let (ctx, bag, staleness) = self;

        let futures = keys.into_iter().map(|cs_id| match staleness {
            DagIdStaleness::Fresh => bag
                .idmap
                .find_dag_id(ctx, cs_id)
                .map_ok(move |v| (cs_id, v))
                .left_future(),
            DagIdStaleness::MaybeStale => bag
                .idmap
                .find_dag_id_maybe_stale(ctx, cs_id)
                .map_ok(move |v| (cs_id, v))
                .right_future(),
        });

        let res = future::try_join_all(futures)
            .await?
            .into_iter()
            .filter_map(|(cs_id, opt)| opt.map(move |v| (cs_id, DagIdWrapper(v))))
            .collect();

        Ok(res)
    }
}

#[cfg(test)]
mod tests {
    use fbinit::FacebookInit;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;
    use sql_construct::SqlConstruct;
    use sql_ext::replication::NoReplicaLagMonitor;

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

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

        let cache_handlers = CacheHandlers::new(CacheHandlerFactory::Mocked);
        let conns = SegmentedChangelogSqlConnections::with_sqlite_in_memory()?;
        let new_cached_idmap = |repo_id| {
            let idmap_version = IdMapVersion(0);
            let sql_idmap = SqlIdMap::new(
                conns.0.clone(),
                Arc::new(NoReplicaLagMonitor()),
                repo_id,
                idmap_version,
            );
            CachedIdMap::new(
                Arc::new(sql_idmap),
                cache_handlers.clone(),
                repo_id,
                idmap_version,
            )
        };
        let idmap1 = new_cached_idmap(RepositoryId::new(1));
        let idmap2 = new_cached_idmap(RepositoryId::new(2));

        idmap1.insert(&ctx, DagId(0), ONES_CSID).await?;
        idmap1.insert(&ctx, DagId(1), TWOS_CSID).await?;
        idmap2.insert(&ctx, DagId(0), TWOS_CSID).await?;

        // prime the caches
        assert_eq!(idmap1.get_changeset_id(&ctx, DagId(0)).await?, ONES_CSID);
        assert_eq!(idmap1.get_changeset_id(&ctx, DagId(1)).await?, TWOS_CSID);
        assert_eq!(idmap2.get_changeset_id(&ctx, DagId(0)).await?, TWOS_CSID);

        assert_eq!(idmap1.get_dag_id(&ctx, TWOS_CSID).await?, DagId(1));
        assert_eq!(idmap2.get_dag_id(&ctx, TWOS_CSID).await?, DagId(0));

        // check caches are set
        assert_eq!(
            cache_handlers
                .dag_to_cs
                .mock_store()
                .expect("mock handler has mock store")
                .stats()
                .gets,
            3
        );
        assert_eq!(
            cache_handlers
                .cs_to_dag
                .mock_store()
                .expect("mock handler has mock store")
                .stats()
                .gets,
            2
        );

        // fetch from caches
        assert_eq!(idmap1.get_changeset_id(&ctx, DagId(0)).await?, ONES_CSID);
        assert_eq!(idmap1.get_changeset_id(&ctx, DagId(1)).await?, TWOS_CSID);
        assert_eq!(idmap2.get_changeset_id(&ctx, DagId(0)).await?, TWOS_CSID);

        assert_eq!(idmap1.get_dag_id(&ctx, TWOS_CSID).await?, DagId(1));
        assert_eq!(idmap2.get_dag_id(&ctx, TWOS_CSID).await?, DagId(0));

        Ok(())
    }
}


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

use std::collections::HashMap;

use anyhow::Result;
use context::CoreContext;
use mononoke_types::ChangesetId;
use parking_lot::RwLock;

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

#[derive(Debug)]
pub struct MemIdMap {
    dag_id2cs: HashMap<DagId, ChangesetId>,
    cs2dag_id: HashMap<ChangesetId, DagId>,
    last_entry: Option<(DagId, ChangesetId)>,
}

impl MemIdMap {
    pub fn new() -> Self {
        Self {
            dag_id2cs: HashMap::new(),
            cs2dag_id: HashMap::new(),
            last_entry: None,
        }
    }

    pub fn len(&self) -> usize {
        self.dag_id2cs.len()
    }

    pub fn drain(&mut self) -> impl Iterator<Item = (DagId, ChangesetId)> + '_ {
        self.last_entry = None;
        self.cs2dag_id.clear();
        self.dag_id2cs.drain()
    }

    pub fn insert(&mut self, dag_id: DagId, cs_id: ChangesetId) {
        self.dag_id2cs.insert(dag_id, cs_id);
        self.cs2dag_id.insert(cs_id, dag_id);
        match self.last_entry {
            Some((last_dag_id, _)) if last_dag_id > dag_id => {}
            _ => self.last_entry = Some((dag_id, cs_id)),
        }
    }

    pub fn find_changeset_id(&self, dag_id: DagId) -> Option<ChangesetId> {
        self.dag_id2cs.get(&dag_id).copied()
    }

    pub fn find_dag_id(&self, cs_id: ChangesetId) -> Option<DagId> {
        self.cs2dag_id.get(&cs_id).copied()
    }

    pub fn get_last_entry(&self) -> Result<Option<(DagId, ChangesetId)>> {
        Ok(self.last_entry.clone())
    }
}

#[derive(Debug)]
pub struct ConcurrentMemIdMap {
    inner: RwLock<MemIdMap>,
}

impl ConcurrentMemIdMap {
    pub fn new() -> Self {
        Self {
            inner: RwLock::new(MemIdMap::new()),
        }
    }

    pub fn len(&self) -> usize {
        let inner = self.inner.read();
        inner.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn drain(&self) -> Vec<(DagId, ChangesetId)> {
        let mut inner = self.inner.write();
        inner.drain().collect()
    }
}

#[async_trait::async_trait]
impl IdMap for ConcurrentMemIdMap {
    async fn insert_many(
        &self,
        _ctx: &CoreContext,
        mappings: Vec<(DagId, ChangesetId)>,
    ) -> Result<()> {
        let mut inner = self.inner.write();
        for (dag_id, cs) in mappings {
            inner.insert(dag_id, cs);
        }
        Ok(())
    }

    async fn find_many_changeset_ids(
        &self,
        _ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> Result<HashMap<DagId, ChangesetId>> {
        let inner = self.inner.read();
        let result = dag_ids
            .into_iter()
            .filter_map(|v| inner.find_changeset_id(v).map(|cs| (v, cs)))
            .collect();
        Ok(result)
    }

    async fn find_many_dag_ids(
        &self,
        _ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        let inner = self.inner.read();
        let result = cs_ids
            .into_iter()
            .filter_map(|cs| inner.find_dag_id(cs).map(|v| (cs, v)))
            .collect();
        Ok(result)
    }

    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        self.find_many_dag_ids(ctx, cs_ids).await
    }

    async fn get_last_entry(&self, _ctx: &CoreContext) -> Result<Option<(DagId, ChangesetId)>> {
        let inner = self.inner.read();
        inner.get_last_entry()
    }

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

#[cfg(test)]
mod test {
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_types_mocks::changesetid::AS_CSID;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;

    use super::*;

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

        assert_eq!(
            idmap
                .find_many_dag_ids(&ctx, vec![AS_CSID, ONES_CSID, TWOS_CSID])
                .await?,
            hashmap![]
        );
        assert_eq!(
            idmap
                .find_many_changeset_ids(&ctx, vec![DagId(0), DagId(1), DagId(2)])
                .await?,
            hashmap![]
        );
        assert_eq!(idmap.get_last_entry(&ctx).await?, None);

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

        assert_eq!(
            idmap
                .find_many_dag_ids(&ctx, vec![AS_CSID, ONES_CSID, TWOS_CSID])
                .await?,
            hashmap![AS_CSID => DagId(0), ONES_CSID => DagId(1)]
        );
        assert_eq!(
            idmap
                .find_many_changeset_ids(&ctx, vec![DagId(0), DagId(1), DagId(2)])
                .await?,
            hashmap![DagId(0) => AS_CSID, DagId(1) => ONES_CSID]
        );
        assert_eq!(
            idmap.get_last_entry(&ctx).await?,
            Some((DagId(1), ONES_CSID))
        );

        Ok(())
    }
}


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

mod cache;
mod mem;
mod shared_traits;
mod sql;

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

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql_ext::replication::ReplicaLagMonitor;
use sql_ext::SqlConnections;

pub use self::cache::CacheHandlers;
pub use self::cache::CachedIdMap;
pub use self::mem::ConcurrentMemIdMap;
pub use self::mem::MemIdMap;
pub use self::shared_traits::cs_id_from_vertex_name;
pub use self::shared_traits::vertex_name_from_cs_id;
pub use self::shared_traits::IdMapWrapper;
pub use self::sql::SqlIdMap;
use crate::types::IdMapVersion;
use crate::DagId;
use crate::DagIdSet;
use crate::InProcessIdDag;

#[async_trait]
#[auto_impl::auto_impl(&, Arc)]
pub trait IdMap: Send + Sync {
    async fn insert_many(
        &self,
        ctx: &CoreContext,
        mappings: Vec<(DagId, ChangesetId)>,
    ) -> Result<()>;

    async fn find_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> Result<HashMap<DagId, ChangesetId>>;

    async fn find_many_dag_ids(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>>;

    /// Finds the dag ID for given changeset - if possible to do so quickly.
    /// Might return no answers for changesets that have dag ids assigned.
    ///
    /// Should be used by callers that can deal with missing information.
    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>>;

    async fn get_last_entry(&self, ctx: &CoreContext) -> Result<Option<(DagId, ChangesetId)>>;

    fn idmap_version(&self) -> Option<IdMapVersion>;

    // Default implementations

    async fn insert(&self, ctx: &CoreContext, dag_id: DagId, cs_id: ChangesetId) -> Result<()> {
        self.insert_many(ctx, vec![(dag_id, cs_id)]).await
    }

    async fn find_changeset_id(
        &self,
        ctx: &CoreContext,
        dag_id: DagId,
    ) -> Result<Option<ChangesetId>> {
        Ok(self
            .find_many_changeset_ids(ctx, vec![dag_id])
            .await?
            .remove(&dag_id))
    }

    async fn find_dag_id(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<Option<DagId>> {
        Ok(self
            .find_many_dag_ids(ctx, vec![cs_id])
            .await?
            .remove(&cs_id))
    }

    async fn find_dag_id_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<DagId>> {
        Ok(self
            .find_many_dag_ids_maybe_stale(ctx, vec![cs_id])
            .await?
            .remove(&cs_id))
    }

    async fn get_changeset_id(&self, ctx: &CoreContext, dag_id: DagId) -> Result<ChangesetId> {
        self.find_changeset_id(ctx, dag_id)
            .await?
            .ok_or_else(|| format_err!("Failed to find segmented changelog id {} in IdMap", dag_id))
    }

    async fn get_dag_id(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<DagId> {
        self.find_dag_id(ctx, cs_id)
            .await?
            .ok_or_else(|| format_err!("Failed to find changeset id {} in IdMap", cs_id))
    }
}

/// The idmap works in unison with the iddag. The idmap and the iddag need to be in sync for iddag
/// update operation to be consistent.  When we download an iddag, it has the idmap described by
/// the shared store. The overlay allows us to update the iddag in process by adding an idmap that
/// stores all the changes that we made in process. It's important to note that the shared store is
/// updated constantly by other processes. So we need to be careful not using the shared store
/// for Ids that are not referred by the matching IdDag. To do that, we ask the IdDag about all
/// Ids covered by it, store it in `shared_id_set` and then only use the shared store for Ids in
/// `shared_id_set`. Note that if we were to use the abcence of an entry from the in process store
/// to fetch from the shared store we would likely end up with inconsistent entries from an updated
/// shared store.
pub struct OverlayIdMap {
    /// Source for `Id`s not in `shared_id_set`. Mutable.
    mem: ConcurrentMemIdMap,
    /// Source for `Id`s in `shared_id_set`. Immutable.
    shared: Arc<dyn IdMap>,
    /// `Id`s covered by the `shared` IdMap.
    shared_id_set: DagIdSet,
}

impl OverlayIdMap {
    pub fn new(shared: Arc<dyn IdMap>, shared_id_set: DagIdSet) -> Self {
        let mem = ConcurrentMemIdMap::new();
        Self {
            mem,
            shared,
            shared_id_set,
        }
    }

    pub fn from_iddag_and_idmap(iddag: &InProcessIdDag, shared: Arc<dyn IdMap>) -> Result<Self> {
        let shared_id_set = iddag.all().context("error calculating iddag.all()")?;
        Ok(Self::new(shared, shared_id_set))
    }
}

#[async_trait]
impl IdMap for OverlayIdMap {
    async fn insert_many(
        &self,
        ctx: &CoreContext,
        mappings: Vec<(DagId, ChangesetId)>,
    ) -> Result<()> {
        for (v, _) in mappings.iter() {
            if self.shared_id_set.contains(*v) {
                return Err(format_err!(
                    "overlay idmap asked to insert {} but it is in immutable shared set {:?}",
                    v,
                    &self.shared_id_set
                ));
            }
        }
        self.mem.insert_many(ctx, mappings).await
    }

    async fn find_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        dag_ids: Vec<DagId>,
    ) -> Result<HashMap<DagId, ChangesetId>> {
        let from_mem = dag_ids
            .iter()
            .filter(|&v| !self.shared_id_set.contains(*v))
            .cloned()
            .collect();
        let mut result = self.mem.find_many_changeset_ids(ctx, from_mem).await?;
        let from_shared: Vec<DagId> = dag_ids
            .iter()
            .filter(|&v| self.shared_id_set.contains(*v))
            .cloned()
            .collect();
        if !from_shared.is_empty() {
            let shared_result = self
                .shared
                .find_many_changeset_ids(ctx, from_shared)
                .await?;
            for (v, cs) in shared_result {
                result.insert(v, cs);
            }
        }
        Ok(result)
    }

    async fn find_many_dag_ids(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        let mut result = self.mem.find_many_dag_ids(ctx, cs_ids.clone()).await?;
        for (cs, v) in result.iter() {
            if self.shared_id_set.contains(*v) {
                bail!(
                    "unexpected assignment found in mem idmap: {} for {} but shared_id_set is {:?}",
                    v,
                    cs,
                    &self.shared_id_set
                );
            }
        }
        let to_get_shared = cs_ids
            .into_iter()
            .filter(|cs_id| !result.contains_key(cs_id))
            .collect();
        let from_shared = self.shared.find_many_dag_ids(ctx, to_get_shared).await?;
        for (cs, v) in from_shared {
            if self.shared_id_set.contains(v) {
                result.insert(cs, v);
            }
        }
        Ok(result)
    }

    async fn find_many_dag_ids_maybe_stale(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, DagId>> {
        let mut result = self
            .mem
            .find_many_dag_ids_maybe_stale(ctx, cs_ids.clone())
            .await?;
        for (cs, v) in result.iter() {
            if self.shared_id_set.contains(*v) {
                bail!(
                    "unexpected assignment found in mem idmap: {} for {} but shared_id_set is {:?}",
                    v,
                    cs,
                    &self.shared_id_set
                );
            }
        }
        let to_get_shared = cs_ids
            .into_iter()
            .filter(|cs_id| !result.contains_key(cs_id))
            .collect();
        let from_shared = self
            .shared
            .find_many_dag_ids_maybe_stale(ctx, to_get_shared)
            .await?;
        for (cs, v) in from_shared {
            if self.shared_id_set.contains(v) {
                result.insert(cs, v);
            }
        }
        Ok(result)
    }

    async fn get_last_entry(&self, ctx: &CoreContext) -> Result<Option<(DagId, ChangesetId)>> {
        match self.mem.get_last_entry(ctx).await? {
            Some(x) => Ok(Some(x)),
            None if !self.shared_id_set.is_empty() => {
                let dag_id = self.shared_id_set.max().unwrap();
                let cs_id = self
                    .shared
                    .find_changeset_id(ctx, dag_id)
                    .await?
                    .with_context(|| {
                        format!(
                            "could not find shared entry for dag_id {} (overlay shared_id_set = {:?})",
                            dag_id, &self.shared_id_set
                        )
                    })?;
                Ok(Some((dag_id, cs_id)))
            }
            None => Ok(None),
        }
    }

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

// The builder for the standard IdMap
// Our layers are: SqlIdMap, CachedIdMap, OverlayIdMap
#[derive(Clone)]
pub struct IdMapFactory {
    connections: SqlConnections,
    replica_lag_monitor: Arc<dyn ReplicaLagMonitor>,
    repo_id: RepositoryId,
    cache_handlers: Option<CacheHandlers>,
}

impl IdMapFactory {
    pub fn new(
        connections: SqlConnections,
        replica_lag_monitor: Arc<dyn ReplicaLagMonitor>,
        repo_id: RepositoryId,
    ) -> Self {
        Self {
            connections,
            replica_lag_monitor,
            repo_id,
            cache_handlers: None,
        }
    }

    // Writes go to the SQL table.
    pub fn for_writer(&self, ctx: &CoreContext, version: IdMapVersion) -> Arc<dyn IdMap> {
        let sql_idmap = SqlIdMap::new(
            self.connections.clone(),
            self.replica_lag_monitor.clone(),
            self.repo_id,
            version,
        );
        slog::debug!(
            ctx.logger(),
            "segmented changelog idmap instantiated - version: {}",
            version
        );
        let mut idmap: Arc<dyn IdMap> = Arc::new(sql_idmap);
        if let Some(cache_handlers) = &self.cache_handlers {
            idmap = Arc::new(CachedIdMap::new(
                idmap,
                cache_handlers.clone(),
                self.repo_id,
                version,
            ));
        }
        idmap
    }

    // Servers have an overlay idmap which means that all their updates to the idmap stay confined
    // to the Dag that performed the updates.
    pub fn for_server(
        &self,
        ctx: &CoreContext,
        version: IdMapVersion,
        iddag: &InProcessIdDag,
    ) -> Result<Arc<dyn IdMap>> {
        let overlay = OverlayIdMap::from_iddag_and_idmap(iddag, self.for_writer(ctx, version))?;
        Ok(Arc::new(overlay))
    }

    pub fn with_cache_handlers(mut self, cache_handlers: CacheHandlers) -> Self {
        self.cache_handlers = Some(cache_handlers);
        self
    }
}

#[cfg(test)]
mod tests {
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_types_mocks::changesetid::AS_CSID;
    use mononoke_types_mocks::changesetid::FOURS_CSID;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::THREES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;

    use super::*;

    fn cutoff(n: u64) -> DagIdSet {
        if n == 0 {
            DagIdSet::empty()
        } else {
            DagIdSet::from_spans(vec![DagId(0)..=DagId(n - 1)])
        }
    }

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

        let shared: Arc<dyn IdMap> = Arc::new(ConcurrentMemIdMap::new());

        shared
            .insert_many(&ctx, vec![(DagId(0), AS_CSID), (DagId(1), ONES_CSID)])
            .await?;

        let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(2));

        assert_eq!(
            overlay
                .find_many_changeset_ids(&ctx, vec![DagId(0), DagId(1), DagId(2)])
                .await?,
            hashmap![DagId(0) => AS_CSID, DagId(1) => ONES_CSID]
        );

        overlay
            .insert_many(&ctx, vec![(DagId(2), TWOS_CSID), (DagId(3), THREES_CSID)])
            .await?;
        assert_eq!(
            overlay
                .find_many_changeset_ids(&ctx, vec![DagId(2), DagId(3)])
                .await?,
            hashmap![DagId(2) => TWOS_CSID, DagId(3) => THREES_CSID]
        );

        assert_eq!(
            overlay
                .shared
                .find_many_changeset_ids(&ctx, vec![DagId(2), DagId(3)])
                .await?,
            hashmap![]
        );

        shared
            .insert_many(
                &ctx,
                vec![
                    (DagId(2), THREES_CSID),
                    (DagId(3), TWOS_CSID),
                    (DagId(4), FOURS_CSID),
                ],
            )
            .await?;

        assert_eq!(
            overlay
                .find_many_changeset_ids(&ctx, vec![DagId(3), DagId(4)])
                .await?,
            hashmap![DagId(3) => THREES_CSID]
        );

        Ok(())
    }

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

        let shared: Arc<dyn IdMap> = Arc::new(ConcurrentMemIdMap::new());

        shared
            .insert_many(&ctx, vec![(DagId(0), AS_CSID), (DagId(1), ONES_CSID)])
            .await?;

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(0));
            assert_eq!(overlay.get_last_entry(&ctx).await?, None);
        }

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(1));
            assert_eq!(
                overlay.get_last_entry(&ctx).await?,
                Some((DagId(0), AS_CSID))
            );
        }

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(1));
            overlay.insert(&ctx, DagId(1), THREES_CSID).await?;
            assert_eq!(
                overlay.get_last_entry(&ctx).await?,
                Some((DagId(1), THREES_CSID)),
            );
        }

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(2));
            assert_eq!(
                overlay.get_last_entry(&ctx).await?,
                Some((DagId(1), ONES_CSID)),
            );
        }

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(2));
            overlay.insert(&ctx, DagId(2), TWOS_CSID).await?;
            assert_eq!(
                overlay.get_last_entry(&ctx).await?,
                Some((DagId(2), TWOS_CSID)),
            );
        }

        {
            let overlay = OverlayIdMap::new(Arc::clone(&shared), cutoff(3));
            assert!(overlay.get_last_entry(&ctx).await.is_err());
        }

        Ok(())
    }
}