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

use std::collections::BTreeMap;
use std::num::NonZeroU64;

use anyhow::anyhow;
use anyhow::Error;
use blobrepo::save_bonsai_changesets;
use blobrepo::BlobRepo;
use commit_transformation::copy_file_contents;
use context::CoreContext;
use derived_data::BonsaiDerived;
use fsnodes::RootFsnodeId;
use futures::future::try_join;
use futures::TryStreamExt;
use manifest::Entry;
use manifest::ManifestOps;
use mononoke_types::fsnode::FsnodeFile;
use mononoke_types::BonsaiChangeset;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::DateTime;
use mononoke_types::FileChange;
use mononoke_types::MPath;
use regex::Regex;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use slog::debug;
use slog::info;
use sorted_vector_map::SortedVectorMap;

#[derive(Copy, Clone, Debug, Default)]
pub struct Limits {
    pub total_file_num_limit: Option<NonZeroU64>,
    pub total_size_limit: Option<NonZeroU64>,
    pub lfs_threshold: Option<NonZeroU64>,
}

#[derive(Clone, Debug, Default)]
pub struct Options {
    pub maybe_exclude_file_regex: Option<Regex>,
    pub overwrite: bool,
}

pub async fn copy(
    ctx: &CoreContext,
    source_repo: &BlobRepo,
    target_repo: &BlobRepo,
    source_cs_id: ChangesetId,
    target_cs_id: ChangesetId,
    from_to_dirs: Vec<(MPath, MPath)>,
    author: String,
    msg: String,
    limits: Limits,
    options: Options,
    skip_overwrite_warning: Option<&dyn Fn(&CoreContext, &MPath)>,
) -> Result<Vec<ChangesetId>, Error> {
    // These are the file changes that have to be removed first
    let mut remove_file_changes = BTreeMap::new();
    // These are the file changes that have to be copied
    let mut file_changes = BTreeMap::new();
    let mut total_file_size = 0;
    let mut contents_to_upload = vec![];
    let same_repo = source_repo.repo_identity().id() == target_repo.repo_identity().id();

    for (from_dir, to_dir) in from_to_dirs {
        let (from_entries, to_entries) = try_join(
            list_directory(ctx, source_repo, source_cs_id, &from_dir),
            list_directory(ctx, target_repo, target_cs_id, &to_dir),
        )
        .await?;
        let from_entries =
            from_entries.ok_or_else(|| Error::msg("from directory does not exist!"))?;
        let to_entries = to_entries.unwrap_or_else(BTreeMap::new);

        for (from_suffix, fsnode_file) in from_entries {
            if let Some(ref regex) = options.maybe_exclude_file_regex {
                if from_suffix.matches_regex(regex) {
                    continue;
                }
            }

            let from_path = from_dir.join(&from_suffix);
            let to_path = to_dir.join(&from_suffix);

            if let Some(to_fsnode) = to_entries.get(&from_suffix) {
                if to_fsnode == &fsnode_file {
                    continue;
                }

                if options.overwrite {
                    remove_file_changes.insert(to_path.clone(), None);
                } else {
                    if let Some(skip_overwrite_warning) = skip_overwrite_warning {
                        skip_overwrite_warning(ctx, &to_path);
                    }
                    continue;
                }
            }

            debug!(
                ctx.logger(),
                "from {}, to {}, size: {}",
                from_path,
                to_path,
                fsnode_file.size()
            );
            file_changes.insert(to_path, Some((from_path, fsnode_file)));

            if !same_repo {
                contents_to_upload.push(fsnode_file.content_id().clone());
            }

            if let Some(lfs_threshold) = limits.lfs_threshold {
                if fsnode_file.size() < lfs_threshold.get() {
                    total_file_size += fsnode_file.size();
                } else {
                    debug!(
                        ctx.logger(),
                        "size is not accounted because of lfs threshold"
                    );
                }
            } else {
                total_file_size += fsnode_file.size();
            }

            if let Some(limit) = limits.total_file_num_limit {
                if file_changes.len() as u64 >= limit.get() {
                    break;
                }
            }
            if let Some(limit) = limits.total_size_limit {
                if total_file_size > limit.get() {
                    break;
                }
            }
        }
    }

    if !same_repo {
        debug!(
            ctx.logger(),
            "Copying {} files contents from {} to {}",
            contents_to_upload.len(),
            source_repo.repo_identity().name(),
            target_repo.repo_identity().name()
        );
        copy_file_contents(ctx, source_repo, target_repo, contents_to_upload, |_| {}).await?;
    }

    create_changesets(
        ctx,
        target_repo,
        vec![remove_file_changes, file_changes],
        target_cs_id,
        author,
        msg,
        same_repo, /* record_copy_from */
    )
    .await
}

async fn create_changesets(
    ctx: &CoreContext,
    repo: &BlobRepo,
    file_changes: Vec<BTreeMap<MPath, Option<(MPath, FsnodeFile)>>>,
    mut parent: ChangesetId,
    author: String,
    msg: String,
    record_copy_from: bool,
) -> Result<Vec<ChangesetId>, Error> {
    let mut changesets = vec![];
    let mut cs_ids = vec![];
    for path_to_maybe_fsnodes in file_changes {
        if path_to_maybe_fsnodes.is_empty() {
            continue;
        }

        let mut fc = BTreeMap::new();
        for (to_path, maybe_fsnode) in path_to_maybe_fsnodes {
            let file_change = match maybe_fsnode {
                Some((from_path, fsnode_file)) => {
                    let copy_from = if record_copy_from {
                        Some((from_path, parent))
                    } else {
                        None
                    };
                    FileChange::tracked(
                        *fsnode_file.content_id(),
                        *fsnode_file.file_type(),
                        fsnode_file.size(),
                        copy_from,
                    )
                }
                None => FileChange::Deletion,
            };

            fc.insert(to_path, file_change);
        }

        info!(ctx.logger(), "creating csid with {} file changes", fc.len());
        let bcs = create_bonsai_changeset(vec![parent], fc.into(), author.clone(), msg.clone())?;

        let cs_id = bcs.get_changeset_id();
        changesets.push(bcs);
        cs_ids.push(cs_id);
        parent = cs_id;
    }

    save_bonsai_changesets(changesets, ctx.clone(), repo).await?;

    Ok(cs_ids)
}

pub async fn remove_excessive_files(
    ctx: &CoreContext,
    source_repo: &BlobRepo,
    target_repo: &BlobRepo,
    source_cs_id: ChangesetId,
    target_cs_id: ChangesetId,
    from_to_dirs: Vec<(MPath, MPath)>,
    author: String,
    msg: String,
    maybe_total_file_num_limit: Option<NonZeroU64>,
) -> Result<ChangesetId, Error> {
    let mut to_delete = BTreeMap::new();

    for (from_dir, to_dir) in from_to_dirs {
        let (from_entries, to_entries) = try_join(
            list_directory(ctx, source_repo, source_cs_id, &from_dir),
            list_directory(ctx, target_repo, target_cs_id, &to_dir),
        )
        .await?;
        let from_entries =
            from_entries.ok_or_else(|| Error::msg("from directory does not exist!"))?;
        let to_entries = to_entries.unwrap_or_else(BTreeMap::new);

        for to_suffix in to_entries.keys() {
            if !from_entries.contains_key(to_suffix) {
                let to_path = to_dir.join(to_suffix);
                to_delete.insert(to_path, None);
                if let Some(limit) = maybe_total_file_num_limit {
                    if to_delete.len() as u64 >= limit.get() {
                        break;
                    }
                }
            }
        }
    }

    let cs_ids = create_changesets(
        ctx,
        target_repo,
        vec![to_delete],
        target_cs_id,
        author,
        msg,
        false, /* record_copy_from */
    )
    .await?;

    cs_ids
        .last()
        .copied()
        .ok_or_else(|| anyhow!("nothing to remove!"))
}

// Recursively lists all the files under `path` if this is a directory.
// If `path` does not exist then None is returned.
// Note that returned paths are RELATIVE to `path`.
async fn list_directory(
    ctx: &CoreContext,
    repo: &BlobRepo,
    cs_id: ChangesetId,
    path: &MPath,
) -> Result<Option<BTreeMap<MPath, FsnodeFile>>, Error> {
    let root = RootFsnodeId::derive(ctx, repo, cs_id).await?;

    let entries = root
        .fsnode_id()
        .find_entries(
            ctx.clone(),
            repo.repo_blobstore().clone(),
            vec![path.clone()],
        )
        .try_collect::<Vec<_>>()
        .await?;

    let entry = entries.get(0);

    let fsnode_id = match entry {
        Some((_, Entry::Tree(fsnode_id))) => fsnode_id,
        None => {
            return Ok(None);
        }
        Some((_, Entry::Leaf(_))) => {
            return Err(anyhow!(
                "{} is a file, but expected to be a directory",
                path
            ));
        }
    };

    let leaf_entries = fsnode_id
        .list_leaf_entries(ctx.clone(), repo.repo_blobstore().clone())
        .try_collect::<BTreeMap<_, _>>()
        .await?;

    Ok(Some(leaf_entries))
}

fn create_bonsai_changeset(
    parents: Vec<ChangesetId>,
    file_changes: SortedVectorMap<MPath, FileChange>,
    author: String,
    message: String,
) -> Result<BonsaiChangeset, Error> {
    BonsaiChangesetMut {
        parents,
        author,
        author_date: DateTime::now(),
        committer: None,
        committer_date: None,
        message,
        hg_extra: Default::default(),
        git_extra_headers: None,
        git_tree_hash: None,
        file_changes,
        is_snapshot: false,
        git_annotated_tag: None,
    }
    .freeze()
}

#[cfg(test)]
mod test {
    use blobstore::StoreLoadable;
    use changeset_fetcher::ChangesetFetcherRef;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_types::RepositoryId;
    use test_repo_factory::TestRepoFactory;
    use tests_utils::list_working_copy_utf8;
    use tests_utils::CreateCommitContext;

    use super::*;

    #[fbinit::test]
    async fn test_list_directory(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir/a", "a")
            .add_file("dir/b", "b")
            .commit()
            .await?;

        let maybe_dir = list_directory(&ctx, &repo, cs_id, &MPath::new("dir")?).await?;
        let dir = maybe_dir.unwrap();

        assert_eq!(
            dir.keys().collect::<Vec<_>>(),
            vec![&MPath::new("a")?, &MPath::new("b")?]
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_simple(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "a")
            .add_file("dir_from/b", "b")
            .add_file("dir_from/c", "c")
            .add_file("dir_to/a", "dontoverwrite")
            .commit()
            .await?;

        let new_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, new_cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "a".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
                MPath::new("dir_to/a")? => "dontoverwrite".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
                MPath::new("dir_to/c")? => "c".to_string(),
            }
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_multiple(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from_1/a", "a")
            .add_file("dir_from_1/b", "b")
            .add_file("dir_from_1/c", "c")
            .add_file("dir_to_1/a", "dontoverwrite")
            .add_file("dir_from_2/aa", "aa")
            .add_file("dir_from_2/bb", "bb")
            .add_file("dir_from_2/cc", "cc")
            .commit()
            .await?;

        let new_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![
                (MPath::new("dir_from_1")?, MPath::new("dir_to_1")?),
                (MPath::new("dir_from_2")?, MPath::new("dir_to_2")?),
            ],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, new_cs_id,).await?,
            hashmap! {
                MPath::new("dir_from_1/a")? => "a".to_string(),
                MPath::new("dir_from_1/b")? => "b".to_string(),
                MPath::new("dir_from_1/c")? => "c".to_string(),
                MPath::new("dir_to_1/a")? => "dontoverwrite".to_string(),
                MPath::new("dir_to_1/b")? => "b".to_string(),
                MPath::new("dir_to_1/c")? => "c".to_string(),

                MPath::new("dir_from_2/aa")? => "aa".to_string(),
                MPath::new("dir_from_2/bb")? => "bb".to_string(),
                MPath::new("dir_from_2/cc")? => "cc".to_string(),
                MPath::new("dir_to_2/aa")? => "aa".to_string(),
                MPath::new("dir_to_2/bb")? => "bb".to_string(),
                MPath::new("dir_to_2/cc")? => "cc".to_string(),
            }
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_with_limit(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "a")
            .add_file("dir_from/b", "b")
            .add_file("dir_from/c", "c")
            .add_file("dir_to/a", "dontoverwrite")
            .commit()
            .await?;

        let limit = Limits {
            total_file_num_limit: NonZeroU64::new(1),
            total_size_limit: None,
            lfs_threshold: None,
        };
        let first_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            limit.clone(),
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, first_cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "a".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
                MPath::new("dir_to/a")? => "dontoverwrite".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
            }
        );

        let second_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            first_cs_id,
            first_cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            limit,
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, second_cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "a".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
                MPath::new("dir_to/a")? => "dontoverwrite".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
                MPath::new("dir_to/c")? => "c".to_string(),
            }
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_with_excludes(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/BUCK", "buck")
            .add_file("dir_from/b", "b")
            .add_file("dir_from/TARGETS", "targets")
            .add_file("dir_from/subdir/TARGETS", "targets")
            .add_file("dir_from/c.bzl", "bzl")
            .add_file("dir_to/a", "dontoverwrite")
            .commit()
            .await?;

        let cs_id = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options {
                maybe_exclude_file_regex: Some(Regex::new("(BUCK|.*\\.bzl|TARGETS)$")?),
                ..Default::default()
            },
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/BUCK")? => "buck".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/TARGETS")? => "targets".to_string(),
                MPath::new("dir_from/subdir/TARGETS")? => "targets".to_string(),
                MPath::new("dir_from/c.bzl")? => "bzl".to_string(),
                MPath::new("dir_to/a")? => "dontoverwrite".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_with_file_size_limit(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "aaaaaaaaaa")
            .add_file("dir_from/b", "b")
            .add_file("dir_from/c", "c")
            .commit()
            .await?;

        let first_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits {
                total_file_num_limit: None,
                total_size_limit: NonZeroU64::new(5),
                lfs_threshold: None,
            },
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, first_cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "aaaaaaaaaa".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
                MPath::new("dir_to/a")? => "aaaaaaaaaa".to_string(),
            }
        );

        let second_cs_id = copy(
            &ctx,
            &repo,
            &repo,
            first_cs_id,
            first_cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits {
                total_file_num_limit: None,
                total_size_limit: NonZeroU64::new(5),
                lfs_threshold: None,
            },
            Options::default(),
            None,
        )
        .await?
        .last()
        .copied()
        .unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, second_cs_id,).await?,
            hashmap! {
                MPath::new("dir_to/a")? => "aaaaaaaaaa".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
                MPath::new("dir_to/c")? => "c".to_string(),
                MPath::new("dir_from/a")? => "aaaaaaaaaa".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_with_file_size_limit_and_lfs_threshold(
        fb: FacebookInit,
    ) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "aaaaaaaaaa")
            .add_file("dir_from/b", "b")
            .add_file("dir_from/c", "c")
            .commit()
            .await?;

        let cs_ids = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits {
                total_file_num_limit: None,
                total_size_limit: NonZeroU64::new(5),
                lfs_threshold: NonZeroU64::new(2),
            },
            Options::default(),
            None,
        )
        .await?;
        assert_eq!(cs_ids.len(), 1);
        let cs_id = cs_ids.last().copied().unwrap();

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, cs_id,).await?,
            hashmap! {
                MPath::new("dir_to/a")? => "aaaaaaaaaa".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
                MPath::new("dir_to/c")? => "c".to_string(),
                MPath::new("dir_from/a")? => "aaaaaaaaaa".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_from/c")? => "c".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_rsync_with_overwrite(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "aa")
            .add_file("dir_from/b", "b")
            .add_file("dir_to/a", "a")
            .add_file("dir_to/b", "b")
            .commit()
            .await?;

        // No overwrite - nothing should be copied
        let cs_ids = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options {
                overwrite: false,
                ..Default::default()
            },
            None,
        )
        .await?;
        assert!(cs_ids.is_empty());

        // Use overwrite - it should create two commits.
        // First commit removes dir_to/a, second commit copies dir_form/a to dir_to/a
        let cs_ids = copy(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options {
                overwrite: true,
                ..Default::default()
            },
            None,
        )
        .await?;

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, *cs_ids.get(0).unwrap()).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "aa".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
            }
        );

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, *cs_ids.last().unwrap()).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "aa".to_string(),
                MPath::new("dir_from/b")? => "b".to_string(),
                MPath::new("dir_to/a")? => "aa".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
            }
        );

        let copy_bcs = cs_ids
            .last()
            .expect("changeset is expected to exist")
            .load(&ctx, &repo.repo_blobstore().clone())
            .await?;
        let file_changes = copy_bcs.file_changes_map();
        let a_change = match file_changes
            .get(&MPath::new("dir_to/a")?)
            .expect("change to dir_to/a expected to be present in the map")
        {
            FileChange::Change(tc) => tc,
            _ => panic!("change to dir_to/a expected to not be None"),
        };
        // Ensure that there is copy-from inserted when copying within the repo
        assert!(a_change.copy_from().is_some());

        Ok(())
    }

    #[fbinit::test]
    async fn test_delete_excessive_files(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from/a", "a")
            .add_file("dir_to/a", "a")
            .add_file("dir_to/b", "b")
            .add_file("dir_to/c/d", "c/d")
            .commit()
            .await?;

        let cs_id = remove_excessive_files(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            None,
        )
        .await?;

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, cs_id,).await?,
            hashmap! {
                MPath::new("dir_from/a")? => "a".to_string(),
                MPath::new("dir_to/a")? => "a".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_delete_excessive_files_multiple_dirs(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs_id = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("dir_from_1/a", "a")
            .add_file("dir_to_1/a", "a")
            .add_file("dir_to_1/b", "b")
            .add_file("dir_to_1/c/d", "c/d")
            .add_file("dir_from_2/a", "a")
            .add_file("dir_to_2/a", "a")
            .add_file("dir_to_2/b", "b")
            .commit()
            .await?;

        let cs_id = remove_excessive_files(
            &ctx,
            &repo,
            &repo,
            cs_id,
            cs_id,
            vec![
                (MPath::new("dir_from_1")?, MPath::new("dir_to_1")?),
                (MPath::new("dir_from_2")?, MPath::new("dir_to_2")?),
            ],
            "author".to_string(),
            "msg".to_string(),
            None,
        )
        .await?;

        assert_eq!(
            list_working_copy_utf8(&ctx, &repo, cs_id,).await?,
            hashmap! {
                MPath::new("dir_from_1/a")? => "a".to_string(),
                MPath::new("dir_to_1/a")? => "a".to_string(),
                MPath::new("dir_from_2/a")? => "a".to_string(),
                MPath::new("dir_to_2/a")? => "a".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_delete_excessive_files_xrepo(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut factory = TestRepoFactory::new(fb)?;
        let source_repo = factory.with_id(RepositoryId::new(0)).build().await?;
        let target_repo = factory.with_id(RepositoryId::new(1)).build().await?;

        let source_cs_id = CreateCommitContext::new_root(&ctx, &source_repo)
            .add_file("dir_from/a", "a")
            .commit()
            .await?;

        let target_cs_id = CreateCommitContext::new_root(&ctx, &target_repo)
            .add_file("dir_to/a", "a")
            .add_file("dir_to/b", "b")
            .add_file("dir_to/c/d", "c/d")
            .commit()
            .await?;

        let cs_id = remove_excessive_files(
            &ctx,
            &source_repo,
            &target_repo,
            source_cs_id,
            target_cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            None,
        )
        .await?;

        assert_eq!(
            list_working_copy_utf8(&ctx, &target_repo, cs_id).await?,
            hashmap! {
                MPath::new("dir_to/a")? => "a".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_xrepo_rsync_with_overwrite(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut factory = TestRepoFactory::new(fb)?;
        let source_repo = factory.with_id(RepositoryId::new(0)).build().await?;
        let target_repo = factory.with_id(RepositoryId::new(1)).build().await?;

        let source_cs_id = CreateCommitContext::new_root(&ctx, &source_repo)
            .add_file("dir_from/a", "aa")
            .add_file("dir_from/b", "b")
            .add_file("source_random_file", "sr")
            .commit()
            .await?;

        let target_cs_id = CreateCommitContext::new_root(&ctx, &target_repo)
            .add_file("dir_to/a", "different")
            .add_file("target_random_file", "tr")
            .commit()
            .await?;

        let cs_ids = copy(
            &ctx,
            &source_repo,
            &target_repo,
            source_cs_id,
            target_cs_id,
            vec![(MPath::new("dir_from")?, MPath::new("dir_to")?)],
            "author".to_string(),
            "msg".to_string(),
            Limits::default(),
            Options {
                overwrite: true,
                ..Default::default()
            },
            None,
        )
        .await?;

        assert_eq!(
            list_working_copy_utf8(&ctx, &target_repo, *cs_ids.get(0).unwrap()).await?,
            hashmap! {
                MPath::new("target_random_file")? => "tr".to_string(),
            }
        );

        assert_eq!(
            list_working_copy_utf8(&ctx, &target_repo, *cs_ids.get(1).unwrap()).await?,
            hashmap! {
                MPath::new("dir_to/a")? => "aa".to_string(),
                MPath::new("dir_to/b")? => "b".to_string(),
                MPath::new("target_random_file")? => "tr".to_string(),
            }
        );

        assert_eq!(
            target_repo
                .changeset_fetcher()
                .get_parents(&ctx, *cs_ids.get(0).unwrap())
                .await?,
            vec![target_cs_id],
        );

        let copy_bcs = cs_ids
            .get(1)
            .expect("changeset is expected to exist")
            .load(&ctx, &target_repo.repo_blobstore().clone())
            .await?;
        let file_changes = copy_bcs.file_changes_map();
        let a_change = match file_changes
            .get(&MPath::new("dir_to/a")?)
            .expect("change to dir_to/a expected to be present in the map")
        {
            FileChange::Change(tc) => tc,
            _ => panic!("change to dir_to/a expected to not be None"),
        };
        // Ensure that there's no copy-from inserted when copying from another repo
        assert!(a_change.copy_from().is_none());

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

#![deny(warnings, missing_docs, clippy::all, rustdoc::broken_intra_doc_links)]

//! Deduplicate items by replacing them with an index into an array.

use std::borrow::Borrow;
use std::borrow::Cow;
use std::collections::HashMap;
use std::hash::Hash;

/// A `DedupMap` accepts items and assigns unique values to different
/// indices in a deduplicated vector.  Items can be retrieved using
/// these indices, and the map can be frozen into a vector for efficient
/// storage.
pub struct DedupMap<T> {
    items: Vec<T>,
    indexes: HashMap<T, usize>,
}

impl<T> Default for DedupMap<T>
where
    T: Eq + Hash + Clone,
{
    fn default() -> Self {
        Self {
            items: Vec::new(),
            indexes: HashMap::new(),
        }
    }
}

impl<T> DedupMap<T>
where
    T: Eq + Hash + Clone,
{
    /// Create a new, empty `DedupMap`.
    pub fn new() -> Self {
        Self::default()
    }

    /// Insert an item into the `DedupMap`.  If the item has already
    /// been added to the map, then its existing index will be returned.
    /// Otherwise, a new index will be returned.
    pub fn insert<'a, B, N>(&mut self, value: N) -> usize
    where
        N: Into<Cow<'a, B>>,
        B: ToOwned<Owned = T> + Eq + Hash + ?Sized + 'a,
        T: Borrow<B>,
    {
        let value = value.into();
        if let Some(&index) = self.indexes.get(value.as_ref()) {
            index
        } else {
            let index = self.items.len();
            let value = value.into_owned();
            self.indexes.insert(value.clone(), index);
            self.items.push(value);
            index
        }
    }

    /// Look up an item by the index previously returned by `insert`.
    pub fn get(&self, index: usize) -> Option<&T> {
        self.items.get(index)
    }

    /// Freeze the `DedupMap` into a vector of the contained item.
    /// This can be used to look-up the indexes previously returned
    /// by `insert`.
    pub fn into_items(self) -> Vec<T> {
        self.items
    }
}

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

    #[test]
    fn basic() {
        let mut dedup_map: DedupMap<String> = DedupMap::new();

        // Can insert items by reference (they will be cloned if necessary).
        let item1 = dedup_map.insert("hello");
        let item2 = dedup_map.insert("rust");

        // Can also insert owned items.
        let item3 = dedup_map.insert(String::from("test"));

        // Can insert duplicate items and will get the same index.
        let item4 = dedup_map.insert("hello");
        assert_eq!(item1, item4);

        // Can read items from the dedup map.
        assert_eq!(dedup_map.get(item1), Some(&String::from("hello")));
        assert_eq!(dedup_map.get(item2), Some(&String::from("rust")));
        assert_eq!(dedup_map.get(item3), Some(&String::from("test")));

        // Can freeez the dedup map into a vector and read from that.
        let dedup_vec = dedup_map.into_items();
        assert_eq!(dedup_vec.get(item1), Some(&String::from("hello")));
        assert_eq!(dedup_vec.get(item2), Some(&String::from("rust")));
        assert_eq!(dedup_vec.get(item3), Some(&String::from("test")));
    }
}


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

use std::pin::Pin;
use std::time::Duration;
use std::time::Instant;

use futures::future::Future;
use futures::task::Context;
use futures::task::Poll;
use maybe_owned::MaybeOwned;
use pin_project::pin_project;
use slog::Logger;
use slog::Record;

#[pin_project]
pub struct WatchedFuture<R, F> {
    reporter: R,
    #[pin]
    inner: F,
}

impl<R, F> Future for WatchedFuture<R, F>
where
    R: Reporter,
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = self.project();

        let now = Instant::now();
        let ret = this.inner.poll(cx);
        this.reporter.report(now.elapsed());

        ret
    }
}

pub trait Reporter {
    fn report(&self, poll: Duration);
}

pub struct SlogReporter<'a> {
    logger: MaybeOwned<'a, Logger>,
    location: slog::RecordLocation,
    max_poll: Duration,
}

impl Reporter for SlogReporter<'_> {
    fn report(&self, poll: Duration) {
        if poll <= self.max_poll {
            return;
        }

        self.logger.log(&Record::new(
            &slog::RecordStatic {
                location: &self.location,
                level: slog::Level::Warning,
                tag: "futures_watchdog",
            },
            &format_args!("Slow poll() ran for {:?}", poll),
            slog::b!(),
        ));
    }
}

pub trait WatchdogExt: Sized {
    #[track_caller]
    fn watched<'a, L>(self, logger: L) -> WatchedFuture<SlogReporter<'a>, Self>
    where
        L: Into<MaybeOwned<'a, Logger>>;
}

impl<F> WatchdogExt for F
where
    F: Future + Sized,
{
    #[track_caller]
    fn watched<'a, L>(self, logger: L) -> WatchedFuture<SlogReporter<'a>, Self>
    where
        L: Into<MaybeOwned<'a, Logger>>,
    {
        let logger = logger.into();

        let location = std::panic::Location::caller();

        let location = slog::RecordLocation {
            file: location.file(),
            line: location.line(),
            column: location.column(),
            function: "",
            module: "",
        };

        // This is a bit arbitrary but generally a very conservative default.
        let max_poll = Duration::from_millis(500);

        let reporter = SlogReporter {
            logger,
            location,
            max_poll,
        };

        WatchedFuture {
            reporter,
            inner: self,
        }
    }
}


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

/// Break iterable down into chunks, by saturating an accumulator
pub fn chunk_by_accumulation<T, A: Copy>(
    items: impl IntoIterator<Item = T>,
    zero: A,
    add: impl Fn(A, &T) -> A,
    overflows: impl Fn(A) -> bool,
) -> Vec<Vec<T>> {
    let mut res = vec![];
    let mut acc = zero;
    let mut current = vec![];
    for item in items.into_iter() {
        let new_acc = add(acc, &item);
        if current.is_empty() || !overflows(new_acc) {
            current.push(item);
            acc = new_acc;
        } else {
            res.push(current);
            acc = add(zero, &item);
            current = vec![item];
        }
    }

    // current can only be empty
    // if the whole `items` was empty
    if !current.is_empty() {
        res.push(current);
    }

    res
}

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

    #[test]
    fn test_chunk_by_accumulation_simple() {
        let v = vec![1, 2, 3, 1, 1, 1];
        let chunks = chunk_by_accumulation(v, 0, |a, x| a + x, |a| a > 3);
        assert_eq!(chunks, vec![vec![1, 2], vec![3], vec![1, 1, 1]]);
    }

    #[test]
    fn test_chunk_by_accumulation_one_item_overflows() {
        // even though 3 on its own overflows the accumulator,
        // we don't drop it
        let v = vec![1, 2, 3, 1, 1, 1];
        let chunks = chunk_by_accumulation(v, 0, |a, x| a + x, |a| a >= 3);
        assert_eq!(chunks, vec![vec![1], vec![2], vec![3], vec![1, 1], vec![1]]);
    }

    #[test]
    fn test_chunk_by_accumulation_empty() {
        let v = vec![];
        let chunks = chunk_by_accumulation(v, 0, |a, x| a + x, |a| a > 3);
        assert!(chunks.is_empty());
    }

    #[test]
    fn test_chunk_by_accumulation_single_item() {
        let v = vec![5];
        let chunks = chunk_by_accumulation(v, 0, |a, x| a + x, |a| a > 3);
        assert_eq!(chunks, vec![vec![5]]);
    }

    #[test]
    fn test_chunk_by_accumulation_all_fit_in_single_chunk() {
        let v = vec![1, 2, 3, 1, 1, 1];
        let chunks = chunk_by_accumulation(v.clone(), 0, |a, x| a + x, |a| a > 1000);
        assert_eq!(chunks, vec![v]);
    }
}


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

#![feature(try_blocks)]

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use metaconfig_types::LoggingDestination;
use serde::Serialize;

#[async_trait]
pub trait Loggable: Serialize + Send + Sync {
    #[cfg(fbcode_build)]
    async fn log_to_logger(&self, ctx: &CoreContext) -> Result<()>;

    #[cfg(not(fbcode_build))]
    async fn log_to_logger(&self, _ctx: &CoreContext) -> Result<()> {
        Ok(())
    }

    fn log_to_scribe(&self, ctx: &CoreContext, scribe_category: &str) -> Result<()> {
        let json_data = serde_json::to_string(self)?;
        ctx.scribe().offer(scribe_category, &json_data)?;
        Ok(())
    }

    async fn log(&self, ctx: &CoreContext, logging_destination: &LoggingDestination) {
        let res = match logging_destination {
            LoggingDestination::Logger => self.log_to_logger(ctx).await,
            LoggingDestination::Scribe { scribe_category } => {
                self.log_to_scribe(ctx, scribe_category)
            }
        };
        if let Err(err) = res {
            ctx.scuba().clone().log_with_msg(
                "Failed to log",
                Some(format!(
                    "dest: {:?}, type: {}, error: {}",
                    logging_destination,
                    std::any::type_name::<Self>(),
                    err
                )),
            )
        }
    }
}


/*
 * 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::borrow::Borrow;
use std::borrow::Cow;
use std::hash::Hash;

use abomonation_derive::Abomonation;
use mononoke_types::hash;
use mononoke_types::path_bytes_from_mpath;
use mononoke_types::MPath;
use mononoke_types::RepoPath;
use sql::mysql;
use sql::mysql_async::prelude::ConvIr;
use sql::mysql_async::prelude::FromValue;
use sql::mysql_async::FromValueError;
use sql::mysql_async::Value;
#[derive(Abomonation, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[derive(mysql::OptTryFromRowField)]
pub struct PathHashBytes(pub Vec<u8>);

impl PathHashBytes {
    pub fn new(path_bytes: &[u8]) -> Self {
        let hash = {
            let mut hash_content = hash::Context::new("path".as_bytes());
            hash_content.update(path_bytes);
            hash_content.finish()
        };

        PathHashBytes(Vec::from(hash.as_ref()))
    }
}

impl ConvIr<PathHashBytes> for Vec<u8> {
    fn new(v: Value) -> Result<Self, FromValueError> {
        Self::from_value_opt(v)
    }

    fn commit(self) -> PathHashBytes {
        PathHashBytes(self)
    }

    fn rollback(self) -> Value {
        self.into()
    }
}

impl FromValue for PathHashBytes {
    type Intermediate = Vec<u8>;
}

impl From<PathHashBytes> for Value {
    fn from(b: PathHashBytes) -> Self {
        b.0.into()
    }
}

impl std::fmt::Display for PathHashBytes {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        for byte in &self.0 {
            write!(f, "{:02x}", byte)?;
        }
        Ok(())
    }
}

#[derive(Abomonation, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[derive(mysql::OptTryFromRowField)]
pub struct PathBytes(pub Vec<u8>);

impl ConvIr<PathBytes> for Vec<u8> {
    fn new(v: Value) -> Result<Self, FromValueError> {
        Self::from_value_opt(v)
    }

    fn commit(self) -> PathBytes {
        PathBytes(self)
    }

    fn rollback(self) -> Value {
        self.into()
    }
}

impl FromValue for PathBytes {
    type Intermediate = Vec<u8>;
}

impl From<PathBytes> for Value {
    fn from(b: PathBytes) -> Self {
        b.0.into()
    }
}

#[derive(Clone)]
pub struct PathWithHash<'a> {
    pub path: Cow<'a, RepoPath>,
    pub path_bytes: PathBytes,
    pub is_tree: bool,
    pub hash: PathHashBytes,
}

impl<'a> PathWithHash<'a> {
    pub fn from_repo_path(path: &'a RepoPath) -> Self {
        Self::from_repo_path_cow(Cow::Borrowed(path))
    }
    pub fn from_repo_path_cow(path: Cow<'a, RepoPath>) -> Self {
        let (path_bytes, is_tree) = convert_from_repo_path(path.borrow());

        let hash = PathHashBytes::new(&path_bytes);

        Self {
            path,
            path_bytes: PathBytes(path_bytes),
            is_tree,
            hash,
        }
    }

    pub fn shard_number_by_hash(hash: &PathHashBytes, shard_count: usize) -> usize {
        // We don't need crypto strength here - we're just turning a potentially large hash into
        // a shard number.
        let raw_shard_number = hash
            .0
            .iter()
            .fold(0usize, |hash, byte| hash.rotate_left(8) ^ (*byte as usize));

        raw_shard_number % shard_count
    }

    pub fn sql_is_tree(&self) -> &'static i8 {
        if self.is_tree { &1 } else { &0 }
    }
}

fn convert_from_repo_path(path: &RepoPath) -> (Vec<u8>, bool) {
    let is_tree = match path {
        &RepoPath::RootPath | &RepoPath::DirectoryPath(_) => true,
        &RepoPath::FilePath(_) => false,
    };
    let bytes = path_bytes_from_mpath(path.mpath());
    (bytes, is_tree)
}

#[derive(Abomonation, Clone, Debug, Eq, PartialEq)]
pub struct PathHash {
    pub path_bytes: PathBytes,
    pub is_tree: bool,
    pub hash: PathHashBytes,
}

impl PathHash {
    pub fn from_repo_path(path: &RepoPath) -> Self {
        let PathWithHash {
            path_bytes,
            is_tree,
            hash,
            ..
        } = PathWithHash::from_repo_path(path);

        Self {
            path_bytes,
            is_tree,
            hash,
        }
    }

    pub fn from_path_and_is_tree(path: Option<&MPath>, is_tree: bool) -> Self {
        let path_bytes = path_bytes_from_mpath(path);
        let hash = PathHashBytes::new(&path_bytes);

        Self {
            path_bytes: PathBytes(path_bytes),
            is_tree,
            hash,
        }
    }

    pub fn shard_number(&self, shard_count: usize) -> usize {
        PathWithHash::shard_number_by_hash(&self.hash, shard_count)
    }

    pub fn sql_is_tree(&self) -> &'static i8 {
        if self.is_tree { &1 } else { &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.
 */

#![feature(fn_traits)]

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

use anyhow::anyhow;
use anyhow::Result;
use arc_swap::ArcSwap;
use async_trait::async_trait;
use cloned::cloned;
use context::CoreContext;
use futures::select;
use futures::FutureExt;
use futures_ext::future::spawn_controlled;
use futures_ext::future::ControlledHandle;
use rand::Rng;
use slog::warn;
use tokio::sync::Notify;

#[derive(Debug, Clone)]
pub struct Reloader<R> {
    obj: Arc<ArcSwap<R>>,
    _handle: Option<ControlledHandle>,
    notify: Arc<Notify>,
}

impl<R> Reloader<R> {
    pub fn load(&self) -> arc_swap::Guard<Arc<R>> {
        self.obj.load()
    }

    pub fn load_full(&self) -> Arc<R> {
        self.obj.load_full()
    }

    pub fn fixed(r: R) -> Self {
        Self {
            obj: Arc::new(ArcSwap::from_pointee(r)),
            _handle: None,
            notify: Arc::new(Notify::new()),
        }
    }

    // Only used in tests
    #[allow(dead_code)]
    pub async fn wait_for_update(&self) {
        self.notify.notified().await;
    }
}

#[async_trait]
pub trait Loader<R> {
    // Load a new instance of R.
    async fn load(&mut self) -> Result<Option<R>>;

    // Determines if R should be reloaded. This function is called to determine
    // if periodic reload is necessary. It's ignored by forced reload.
    async fn needs_reload(&mut self) -> Result<bool> {
        Ok(true)
    }
}

#[async_trait]
impl<R, O, F> Loader<R> for F
where
    O: Future<Output = Result<Option<R>>> + Send,
    F: FnMut() -> O + Send,
{
    async fn load(&mut self) -> Result<Option<R>> {
        self.call_mut(()).await
    }
}

impl<R: 'static + Send + Sync> Reloader<R> {
    pub async fn reload_periodically<
        I: 'static + FnMut() -> Duration + Send,
        L: 'static + Loader<R> + Send + Sync,
    >(
        ctx: CoreContext,
        interval_getter: I,
        loader: L,
    ) -> Result<Self> {
        Self::reload_periodically_with_force_reload(
            ctx,
            interval_getter,
            loader,
            Arc::new(Notify::new()),
        )
        .await
    }

    pub async fn reload_periodically_with_force_reload<
        I: 'static + FnMut() -> Duration + Send,
        L: 'static + Loader<R> + Send + Sync,
    >(
        ctx: CoreContext,
        mut interval_getter: I,
        mut loader: L,
        force_reload: Arc<Notify>,
    ) -> Result<Self> {
        let obj = Arc::new(ArcSwap::from_pointee(
            loader
                .load()
                .await?
                .ok_or_else(|| anyhow!("Missing initial object"))?,
        ));
        let notify = Arc::new(Notify::new());
        let handle = spawn_controlled({
            cloned!(obj, notify);
            async move {
                loop {
                    let interval = interval_getter();
                    let sleep_fut = tokio::time::sleep(interval).fuse();
                    let force_reload_fut = force_reload.notified().fuse();
                    futures::pin_mut!(sleep_fut, force_reload_fut);
                    let forced_reload = select! {
                        _ = sleep_fut => false,
                        _ = force_reload_fut => true,
                    };
                    let reload_now = if forced_reload {
                        true
                    } else {
                        loader.needs_reload().await.unwrap_or_else(|err| {
                            warn!(
                                ctx.logger(),
                                "Failed to check if reload needed, not reloading: {:?}", err
                            );
                            // I's better to keep a known-good SC than to reload an out-of-date version
                            // and catch up again.
                            //
                            // Since the check for an update failed, keep the current version.
                            false
                        })
                    };
                    if reload_now {
                        match loader.load().await {
                            Ok(Some(new)) => obj.store(Arc::new(new)),
                            // Fetch was successful, but there's nothing to reload
                            Ok(None) => {}
                            Err(err) => {
                                warn!(ctx.logger(), "Failed to reload: {:?}", err)
                            }
                        }
                        notify.notify_waiters();
                    }
                }
            }
        });
        Ok(Self {
            obj,
            _handle: Some(handle),
            notify,
        })
    }

    /// Reload periodically with a fixed interval, but on the first wait, randomly wait
    /// up to 10% more. Useful for preventing reloads at the same time when creating multiple
    /// reloaders at the same time with the same interval.
    pub async fn reload_periodically_with_skew<L: 'static + Loader<R> + Send + Sync>(
        ctx: CoreContext,
        period: Duration,
        loader: L,
    ) -> Result<Self> {
        Self::reload_periodically_with_skew_and_force_reload(
            ctx,
            period,
            loader,
            Arc::new(Notify::new()),
        )
        .await
    }

    pub async fn reload_periodically_with_skew_and_force_reload<
        L: 'static + Loader<R> + Send + Sync,
    >(
        ctx: CoreContext,
        period: Duration,
        loader: L,
        force_reload: Arc<Notify>,
    ) -> Result<Self> {
        let mut first = true;
        Self::reload_periodically_with_force_reload(
            ctx,
            move || {
                if first {
                    first = false;
                    let jitter = rand::thread_rng().gen_range(Duration::from_secs(0)..period / 10);
                    period + jitter
                } else {
                    period
                }
            },
            loader,
            force_reload,
        )
        .await
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::AtomicU32;
    use std::sync::atomic::Ordering::Relaxed;
    use std::time::Duration;

    use fbinit::FacebookInit;

    use super::*;

    #[test]
    fn test_fixed() {
        let l = Reloader::fixed(12);
        assert_eq!(**l.load(), 12);
        assert_eq!(**l.load(), 12);
    }

    #[fbinit::test]
    async fn test_reload(fb: FacebookInit) {
        tokio::time::pause();
        struct NumberLoader {
            cur: u32,
        }
        #[async_trait]
        impl Loader<u32> for NumberLoader {
            async fn load(&mut self) -> Result<Option<u32>> {
                let last = self.cur;
                self.cur += 1;
                Ok(Some(last))
            }
        }
        let loader = NumberLoader { cur: 0 };
        let l = Reloader::reload_periodically(
            CoreContext::test_mock(fb),
            || Duration::from_millis(20),
            loader,
        )
        .await
        .unwrap();

        assert_eq!(**l.load(), 0);
        tokio::time::advance(Duration::from_millis(21)).await;
        l.wait_for_update().await;
        assert_eq!(**l.load(), 1);
        tokio::time::advance(Duration::from_millis(21)).await;
        l.wait_for_update().await;
        assert_eq!(**l.load(), 2);
    }

    #[fbinit::test]
    async fn test_reload_fail_then_succeed(fb: FacebookInit) {
        tokio::time::pause();
        struct NumberLoader {
            cur: AtomicU32,
            failing: AtomicBool,
        }
        #[async_trait]
        impl Loader<u32> for Arc<NumberLoader> {
            async fn load(&mut self) -> Result<Option<u32>> {
                if self.failing.load(Relaxed) {
                    Err(anyhow!("Error"))
                } else {
                    Ok(Some(self.cur.fetch_add(1, Relaxed)))
                }
            }
        }
        let loader = Arc::new(NumberLoader {
            cur: AtomicU32::new(0),
            failing: AtomicBool::new(false),
        });
        let l = Reloader::reload_periodically(
            CoreContext::test_mock(fb),
            || std::time::Duration::from_millis(5),
            loader.clone(),
        )
        .await
        .unwrap();
        loader.failing.store(true, Relaxed);
        assert_eq!(**l.load(), 0);
        tokio::time::advance(std::time::Duration::from_millis(12)).await;
        l.wait_for_update().await;
        assert_eq!(**l.load(), 0);
        loader.failing.store(false, Relaxed);
        tokio::time::advance(std::time::Duration::from_millis(10)).await;
        l.wait_for_update().await;
        assert!(**l.load() > 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 clap::Args;

mod multi_rendez_vous;
mod rendez_vous;
mod rendez_vous_stats;
mod tunables;

#[cfg(test)]
mod test;

pub use multi_rendez_vous::MultiRendezVous;
pub use multi_rendez_vous::MultiRendezVousController;
pub use rendez_vous::RendezVous;
pub use rendez_vous::RendezVousController;
pub use rendez_vous_stats::RendezVousStats;

pub use crate::tunables::ConfigurableRendezVousController;
pub use crate::tunables::TunablesMultiRendezVousController;
pub use crate::tunables::TunablesRendezVousController;

#[derive(Copy, Clone, Debug)]
pub struct RendezVousOptions {
    pub free_connections: usize,
}

impl RendezVousOptions {
    pub fn for_test() -> Self {
        Self {
            free_connections: 0,
        }
    }
}

/// Command line arguments for controlling rendez-vous
#[derive(Args, Debug)]
pub struct RendezVousArgs {
    /// How many concurrent connections to allow before batching kicks in
    #[clap(long, default_value = "5")]
    pub rendezvous_free_connections: usize,
}

impl From<RendezVousArgs> for RendezVousOptions {
    fn from(args: RendezVousArgs) -> Self {
        RendezVousOptions {
            free_connections: args.rendezvous_free_connections,
        }
    }
}

#[cfg(test)]
mod demo {
    use std::collections::HashMap;
    use std::sync::Arc;

    use anyhow::Error;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use maplit::hashset;

    use super::*;

    // NOTE: I'd make this a doctest, but we don't have support for running those at the moment
    #[fbinit::test]
    async fn demo(fb: FacebookInit) -> Result<(), Error> {
        // RendezVousOptions would typically be instantiated from CLI args via cmdlib.
        let opts = RendezVousOptions::for_test();

        let stats = Arc::new(RendezVousStats::new("stats_prefix".into()));

        // Callers sharing a RendezVous instance will be eligible to have their calls batched
        // together.
        let rdv = RendezVous::new(TunablesRendezVousController::new(opts), stats);

        let out = rdv
            .dispatch(fb, hashset! { 1u64 }, || {
                |keys| async move {
                    // Keys will include your own query (`1` in this example), and potentially
                    // other queries batched with yours via the RendezVous instance. Return a
                    // HashMap mapping keys to values as your output. You only need to return a
                    // value for keys that were found.
                    Ok(keys
                        .into_iter()
                        .map(|k| (k, k + 1))
                        .collect::<HashMap<_, _>>())
                }
            })
            .await?;

        // The output from dispatch will include only the keys your requested. If a key is missing,
        // you'll get `None` back as the value.
        assert_eq!(out, hashmap! { 1 => Some(2) });

        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::hash::Hash;
use std::sync::Arc;

use dashmap::DashMap;

use crate::RendezVous;
use crate::RendezVousController;
use crate::RendezVousStats;
use crate::TunablesMultiRendezVousController;

pub trait MultiRendezVousController: Send + Sync + 'static {
    type Controller: RendezVousController;

    fn new_controller(&self) -> Self::Controller;
}

/// A wrapper around RendezVous that can be keyed by a grouping key (G). This is useful when you
/// want multiple RendezVous instances for a set of groups but you don't know the groups ahead of
/// time (e.g. the groups might be repository ids).
pub struct MultiRendezVous<
    G,
    K,
    V,
    C: MultiRendezVousController = TunablesMultiRendezVousController,
> {
    inner: Arc<DashMap<G, RendezVous<K, V, <C as MultiRendezVousController>::Controller>>>,
    multi_controller: C,
    stats: Arc<RendezVousStats>,
}

impl<G, K, V, C> Clone for MultiRendezVous<G, K, V, C>
where
    G: Hash + Eq,
    C: MultiRendezVousController + Clone,
{
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            multi_controller: self.multi_controller.clone(),
            stats: self.stats.clone(),
        }
    }
}

impl<G, K, V, C> MultiRendezVous<G, K, V, C>
where
    G: Hash + Eq,
    C: MultiRendezVousController,
{
    pub fn new(multi_controller: C, stats: RendezVousStats) -> Self {
        Self {
            inner: Arc::new(DashMap::new()),
            multi_controller,
            stats: Arc::new(stats),
        }
    }

    pub fn get(&self, group: G) -> RendezVous<K, V, <C as MultiRendezVousController>::Controller> {
        use dashmap::mapref::entry::Entry;

        let ret = match self.inner.entry(group) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => e
                .insert(RendezVous::new(
                    self.multi_controller.new_controller(),
                    self.stats.clone(),
                ))
                .clone(),
        };

        ret
    }
}


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

use anyhow::Error;
use fbinit::FacebookInit;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::Shared;
use futures_ext::FbTryFutureExt;
use futures_stats::TimedFutureExt;
use shared_error::anyhow::IntoSharedError;
use shared_error::anyhow::SharedError;
use time_ext::DurationExt;
use tokio::sync::Notify;

use crate::RendezVousStats;
use crate::TunablesRendezVousController;

/// The RendezVousController controls the behavior of a RendezVous instance. It notably decides
/// when to wait for a batch to build up and when to kick off queries.
#[async_trait::async_trait]
pub trait RendezVousController: Send + Sync + 'static {
    type RendezVousToken: Sized + Send + Sync + 'static;

    /// Delay sending a batch to give ourselves a chance to accumulate some data. The batch will be
    /// kicked off once this future returns. Note that dispatch might still proceed if we reach the
    /// early_dispatch_threshold, in which case the future returned by wait_for_dispatch will be
    /// dropped. Otherwise, the RendezVousToken that was returned will be dropped once this request
    /// finishes.
    async fn wait_for_dispatch(&self) -> Self::RendezVousToken;

    /// If our number of queued keys exceeds this threshold, then we'll dispatch the query even if
    /// wait_for_dispatch hasn't returned yet.
    fn early_dispatch_threshold(&self) -> usize;
}

struct RendezVousInner<K, V, C> {
    staging: Mutex<
        Option<(
            HashSet<K>,
            Shared<BoxFuture<'static, Result<Arc<HashMap<K, V>>, SharedError>>>,
            Arc<Notify>,
        )>,
    >,
    controller: C,
    stats: Arc<RendezVousStats>,
}

/// RendezVous is a difficultly named library which can be used to batch together and deduplicate queries to a backend.
///
/// It batches together queries that look like `async fn get(keys: HashSet<K>) -> HashMap<K, V>`,
/// where the keys of the response are contained in the keys of the input.
/// That is, querying a bunch of keys from an "async key-value" backend, which must follow these restrictions:
/// - The order of keys in a query doesn't matter.
/// - The set of keys in a query can be merged or split with other sets of queries.
/// - Querying the same key twice in a row should return the same value.
///
/// The RendezVousController trait is used to control how the batching is done, such as number of inflight
/// requests, size of them, etc. It has a reasonable default but can be swapped for extra control.
///
/// Rendezvous was created to help with our SQL query load by:
/// - Reducing the number of queries and connections to SQL
/// - Make it easy to do batching right
/// See D27010317 for more context.
pub struct RendezVous<K, V, C = TunablesRendezVousController> {
    inner: Arc<RendezVousInner<K, V, C>>,
}

impl<K, V, C> Clone for RendezVous<K, V, C> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<K, V, C> RendezVous<K, V, C> {
    pub fn new(controller: C, stats: Arc<RendezVousStats>) -> Self {
        Self {
            inner: Arc::new(RendezVousInner {
                staging: Mutex::new(None),
                controller,
                stats,
            }),
        }
    }
}

impl<K, V, C> RendezVous<K, V, C>
where
    K: Clone + Eq + Hash + Send + Sync + 'static,
    V: Clone + Send + Sync + 'static,
    C: RendezVousController,
{
    /// This function does the necessary batching and deduplication under the hood acts like a normal
    /// `get(k) -> k` function. You should always call this function with the same `f0` value to ensure
    /// consistent behaviour.
    pub fn dispatch<F0, F1, Fut>(
        &self,
        fb: FacebookInit,
        keys: HashSet<K>,
        f0: F0,
    ) -> impl Future<Output = Result<HashMap<K, Option<V>>, Error>>
    where
        F0: FnOnce() -> F1, // Can construct a F1 if we are the first caller here
        F1: FnOnce(HashSet<K>) -> Fut + Send + 'static, // Actually makes the call
        Fut: Future<Output = Result<HashMap<K, V>, Error>> + Send,
    {
        if keys.len() < self.inner.controller.early_dispatch_threshold() {
            self.dispatch_batched(fb, keys, f0).left_future()
        } else {
            self.dispatch_not_batched(fb, keys, f0).right_future()
        }
    }

    fn dispatch_batched<F0, F1, Fut>(
        &self,
        fb: FacebookInit,
        keys: HashSet<K>,
        f0: F0,
    ) -> impl Future<Output = Result<HashMap<K, Option<V>>, Error>>
    where
        F0: FnOnce() -> F1,
        F1: FnOnce(HashSet<K>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<HashMap<K, V>, Error>> + Send,
    {
        let mut deduplicated = 0;

        let mut guard = self.inner.staging.lock().expect("Poisoned lock");

        let fut = match &mut *guard {
            guard @ None => {
                let inner = self.inner.clone();
                let f1 = f0();

                let notify = Arc::new(Notify::new());

                let fut = {
                    let notify = notify.clone();

                    async move {
                        let token = futures::select! {
                            token = inner.controller.wait_for_dispatch().fuse() => Some(token),
                            _ = notify.notified().fuse() => None,
                        };

                        if token.is_none() {
                            inner.stats.dispatch_batch_early.add_value(1);
                        } else {
                            inner.stats.dispatch_batch_scheduled.add_value(1);
                        }

                        let ret = tokio::task::spawn(async move {
                            let (keys, _, _) = inner
                                .staging
                                .lock()
                                .expect("Poisoned lock")
                                .take()
                                .expect("Staging cannot be empty if a task was dispatched");

                            let ret = dispatch_with_stats(fb, f1, keys, &inner.stats).await?;

                            std::mem::drop(token);

                            Result::<_, Error>::Ok(ret)
                        })
                        .flatten_err()
                        .await
                        .shared_error()?;

                        Result::<_, SharedError>::Ok(Arc::new(ret))
                    }
                }
                .boxed()
                .shared();

                *guard = Some((keys.clone(), fut.clone(), notify));

                fut
            }
            Some((ref mut staged_keys, ref fut, ref notify)) => {
                for k in keys.iter().cloned() {
                    if !staged_keys.insert(k) {
                        deduplicated += 1;
                    }
                }

                if staged_keys.len() >= self.inner.controller.early_dispatch_threshold() {
                    notify.notify_one();
                }

                fut.clone()
            }
        };

        std::mem::drop(guard);

        self.inner.stats.keys_deduplicated.add_value(deduplicated);

        async move {
            let shared_ret = fut.await?;
            let ret = keys
                .into_iter()
                .map(|k| {
                    let v = shared_ret.get(&k).cloned();
                    (k, v)
                })
                .collect();
            Ok(ret)
        }
    }

    fn dispatch_not_batched<F0, F1, Fut>(
        &self,
        fb: FacebookInit,
        keys: HashSet<K>,
        f0: F0,
    ) -> impl Future<Output = Result<HashMap<K, Option<V>>, Error>>
    where
        F0: FnOnce() -> F1,
        F1: FnOnce(HashSet<K>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<HashMap<K, V>, Error>> + Send,
    {
        let inner = self.inner.clone();

        async move {
            inner.stats.dispatch_no_batch.add_value(1);

            let mut ret = dispatch_with_stats(fb, f0(), keys.clone(), &inner.stats).await?;

            let ret = keys
                .into_iter()
                .map(|k| {
                    let v = ret.remove(&k);
                    (k, v)
                })
                .collect();

            Ok(ret)
        }
    }
}

async fn dispatch_with_stats<K, V, F1, Fut>(
    fb: FacebookInit,
    f1: F1,
    keys: HashSet<K>,
    rdv_stats: &RendezVousStats,
) -> Result<HashMap<K, V>, Error>
where
    F1: FnOnce(HashSet<K>) -> Fut + Send + 'static,
    Fut: Future<Output = Result<HashMap<K, V>, Error>>,
{
    rdv_stats.keys_dispatched.add_value(keys.len() as i64);

    rdv_stats.inflight.increment_value(fb, 1);
    let (stats, ret) = f1(keys).timed().await;
    rdv_stats.inflight.increment_value(fb, -1); // TODO: This should use a scopeguard...

    rdv_stats
        .fetch_completion_time_ms
        .add_value(stats.completion_time.as_millis_unchecked() as i64);

    ret
}


/*
 * 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 stats::prelude::*;

define_stats_struct! {
    RendezVousStats("mononoke.rdv.{}", label: String),

    dispatch_no_batch: timeseries("dispatch_no_batch"; Sum),
    dispatch_batch_early: timeseries("dispatch_batch_early"; Sum),
    dispatch_batch_scheduled: timeseries("dispatch_batch_scheduled"; Sum),

    keys_dispatched: timeseries("keys_dispatched"; Average, Sum),
    keys_deduplicated: timeseries("keys_deduplicated"; Sum),

    fetch_completion_time_ms: histogram(1, 0, 50, Average; P 50; P 95; P 99),

    inflight: singleton_counter("inflight"),
}


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

use anyhow::Error;
use fbinit::FacebookInit;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use maplit::hashmap;
use maplit::hashset;
use tokio::sync::Notify;

pub use crate::RendezVous;
pub use crate::RendezVousController;
pub use crate::RendezVousStats;

#[derive(Clone)]
struct MockController {
    arrive: Arc<Notify>,
    depart: Arc<Notify>,
    threshold: usize,
}

impl MockController {
    pub fn new(threshold: usize) -> Self {
        Self {
            arrive: Arc::new(Notify::new()),
            depart: Arc::new(Notify::new()),
            threshold,
        }
    }

    /// Blocks until wait_for_dispatch is called, then allows wait_for_dispatch to proceed.
    pub async fn release(&self) {
        self.arrive.notified().await;
        self.depart.notify_one();
    }
}

#[async_trait::async_trait]
impl RendezVousController for MockController {
    type RendezVousToken = ();

    /// Blocks until release is called.
    async fn wait_for_dispatch(&self) -> () {
        self.arrive.notify_one();
        self.depart.notified().await;
    }

    fn early_dispatch_threshold(&self) -> usize {
        self.threshold
    }
}

#[derive(Clone)]
struct MockStore {
    calls: Arc<AtomicUsize>,
}

impl MockStore {
    pub fn new() -> Self {
        Self {
            calls: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn calls(&self) -> usize {
        self.calls.load(Ordering::Relaxed)
    }

    pub fn callback(
        &self,
    ) -> impl FnOnce(HashSet<u64>) -> BoxFuture<'static, Result<HashMap<u64, u64>, Error>> {
        let this = self.clone();
        move |keys| {
            async move {
                this.calls.fetch_add(1, Ordering::Relaxed);
                Ok(keys.into_iter().map(|k| (k, k)).collect())
            }
            .boxed()
        }
    }
}

fn stats() -> Arc<RendezVousStats> {
    Arc::new(RendezVousStats::new("test".into()))
}

#[fbinit::test]
async fn test_batch_wait(fb: FacebookInit) -> Result<(), Error> {
    let store = MockStore::new();
    let controller = MockController::new(usize::MAX);

    let rdv = RendezVous::new(controller.clone(), stats());
    let mut f1 = rdv
        .dispatch(fb, hashset! { 1 }, || store.callback())
        .boxed();
    let mut f2 = rdv
        .dispatch(fb, hashset! { 2 }, || store.callback())
        .boxed();

    // We expect both futures to be blocked since there is nothing to release them.
    assert!(futures::poll!(&mut f1).is_pending());
    assert!(futures::poll!(&mut f2).is_pending());

    // We release everything now.
    controller.release().await;

    assert_eq!(f1.await?, hashmap! { 1 => Some(1) });
    assert_eq!(f2.await?, hashmap! { 2 => Some(2) });

    assert_eq!(store.calls(), 1);

    Ok(())
}

#[fbinit::test]
async fn test_batch_limit(fb: FacebookInit) -> Result<(), Error> {
    let store = MockStore::new();
    let controller = MockController::new(2);

    let rdv = RendezVous::new(controller.clone(), stats());

    let mut f1a = rdv
        .dispatch(fb, hashset! { 1 }, || store.callback())
        .boxed();
    let mut f1b = rdv
        .dispatch(fb, hashset! { 1 }, || store.callback())
        .boxed();

    // We expect both futures to be blocked since we're still just fetching one entry here.
    assert!(futures::poll!(&mut f1a).is_pending());
    assert!(futures::poll!(&mut f1b).is_pending());

    // We release everything now.
    let f2 = rdv
        .dispatch(fb, hashset! { 2 }, || store.callback())
        .boxed();

    assert_eq!(f1a.await?, hashmap! { 1 => Some(1) });
    assert_eq!(f1b.await?, hashmap! { 1 => Some(1) });
    assert_eq!(f2.await?, hashmap! { 2 => Some(2) });

    assert_eq!(store.calls(), 1);

    Ok(())
}

#[fbinit::test]
async fn test_unbatched(fb: FacebookInit) -> Result<(), Error> {
    let store = MockStore::new();
    let controller = MockController::new(2);

    let rdv = RendezVous::new(controller.clone(), stats());

    let mut f1 = rdv
        .dispatch(fb, hashset! { 1 }, || store.callback())
        .boxed();

    // We expect f1 to be blocked because nothing released it.
    assert!(futures::poll!(&mut f1).is_pending());

    // We expect f2 to not join f1 because it exceeds the batch threshold.
    let f2 = rdv
        .dispatch(fb, hashset! { 1, 2 }, || store.callback())
        .boxed();
    assert_eq!(f2.await?, hashmap! { 1 => Some(1), 2 => Some(2) });

    // Therefore, we expect f1 to still be blocked.
    assert!(futures::poll!(&mut f1).is_pending());

    assert_eq!(store.calls(), 1);

    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::sync::Arc;
use std::time::Duration;

use tokio::sync::AcquireError;
use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::Semaphore;

use crate::MultiRendezVousController;
use crate::RendezVousController;
use crate::RendezVousOptions;

#[derive(Copy, Clone)]
pub struct TunablesMultiRendezVousController {
    opts: RendezVousOptions,
}

impl TunablesMultiRendezVousController {
    pub fn new(opts: RendezVousOptions) -> Self {
        Self { opts }
    }
}

impl MultiRendezVousController for TunablesMultiRendezVousController {
    type Controller = TunablesRendezVousController;

    fn new_controller(&self) -> Self::Controller {
        TunablesRendezVousController::new(self.opts)
    }
}

/// This RendezVousController is parameterized in two ways:
///
/// It allows a fixed number of "free" connections. This is how many requests we allow to exist
/// in flight at any point in time. Batching does not kick in until these are exhausted. This cannot be
/// changed after the RendezVousController is initialized, so it is not controlled by a tunable.
///
/// Provided functions control what we do once the free connections are exhausted. The functions are:
/// - max_threshold: number of keys after which we'll dispatch a full-size batch.
/// - max_delay: controls how long we wait before dispatching a small batch.
///
/// The provided functions should run very fast, please make sure they are hard-coded or cached.
///
/// Note that if a batch departs when either of those criteria are met, it will not count against
/// the count of free connections: free connections are just connections not subject to batching,
/// but once batching kicks in there is no limit to how many batches can be in flight concurrently
/// (though unless we receive infinite requests the concurrency will tend to approach the free
/// connection count).

///
pub struct ConfigurableRendezVousController {
    semaphore: Arc<Semaphore>,
    max_delay: fn() -> Duration,
    max_threshold: fn() -> usize,
}

impl ConfigurableRendezVousController {
    pub fn new(
        opts: RendezVousOptions,
        max_delay: fn() -> Duration,
        max_threshold: fn() -> usize,
    ) -> Self {
        Self {
            semaphore: Arc::new(Semaphore::new(opts.free_connections)),
            max_delay,
            max_threshold,
        }
    }
}

#[async_trait::async_trait]
impl RendezVousController for ConfigurableRendezVousController {
    // NOTE: We don't actually care about AcquireError here, since that can only happen when the
    // Semaphore is closed, but we don't close it.
    type RendezVousToken = Option<Result<OwnedSemaphorePermit, AcquireError>>;

    /// Wait for the configured dispatch delay.
    async fn wait_for_dispatch(&self) -> Self::RendezVousToken {
        tokio::time::timeout((self.max_delay)(), self.semaphore.clone().acquire_owned())
            .await
            .ok()
    }

    fn early_dispatch_threshold(&self) -> usize {
        (self.max_threshold)()
    }
}

fn max_delay() -> Duration {
    Duration::from_millis(
        ::tunables::tunables()
            .rendezvous_dispatch_delay_ms()
            .unwrap_or_default()
            .try_into()
            .unwrap_or(0),
    )
}

fn max_threshold() -> usize {
    ::tunables::tunables()
        .rendezvous_dispatch_max_threshold()
        .unwrap_or_default()
        .try_into()
        .unwrap_or(0)
}

/// This uses `ConfigurableRendezVousController`, but uses tunables to fetch the limits,
/// for ease of use and being able to quickly change them. The used tunables are:
/// - rendezvous_dispatch_max_threshold: number of keys after which we'll dispatch a full-size batch.
/// - rendezvous_dispatch_delay_ms: controls how long we wait before dispatching a small batch.
pub struct TunablesRendezVousController {
    inner: ConfigurableRendezVousController,
}

impl TunablesRendezVousController {
    pub fn new(opts: RendezVousOptions) -> Self {
        Self {
            inner: ConfigurableRendezVousController::new(opts, max_delay, max_threshold),
        }
    }
}

#[async_trait::async_trait]
impl RendezVousController for TunablesRendezVousController {
    type RendezVousToken = Option<Result<OwnedSemaphorePermit, AcquireError>>;

    /// Wait for the configured dispatch delay.
    async fn wait_for_dispatch(&self) -> Self::RendezVousToken {
        self.inner.wait_for_dispatch().await
    }

    fn early_dispatch_threshold(&self) -> usize {
        self.inner.early_dispatch_threshold()
    }
}


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

use std::time::Duration;

use futures::Future;
use slog::info;
use slog::Logger;

#[derive(Copy, Clone)]
pub struct RetryAttemptsCount(pub usize);

pub enum RetryLogic {
    /// Multiply by a factor every time
    Exponential { base: Duration, factor: f64 },
    /// Increase every time, and randomly jitter some more
    ExponentialWithJitter {
        base: Duration,
        factor: f64,
        jitter: Duration,
    },
}

impl RetryLogic {
    fn delay(&self, attempt: usize) -> Duration {
        use RetryLogic::*;
        match self {
            Exponential { base, factor } => base.mul_f64(factor.powf(attempt as f64)),
            ExponentialWithJitter {
                base,
                factor,
                jitter,
            } => base.mul_f64(factor.powf(attempt as f64)) + jitter.mul_f64(rand::random::<f64>()),
        }
    }
}

/// Retry a function whenever it fails.
/// See `retry` for more information.
pub async fn retry_always<V, Fut, Func, Error>(
    logger: &Logger,
    func: Func,
    base_delay_ms: u64,
    retry_num: usize,
) -> Result<(V, RetryAttemptsCount), Error>
where
    V: Send + 'static,
    Fut: Future<Output = Result<V, Error>>,
    Func: FnMut(usize) -> Fut + Send,
{
    retry(
        Some(logger),
        func,
        |_| true,
        RetryLogic::Exponential {
            base: Duration::from_millis(base_delay_ms),
            factor: 2.0,
        },
        retry_num,
    )
    .await
}

/// Retry a function.
/// `func` is the function to be retried.
/// `should_retry` tells whether an error should be retried.
/// `retry_num` is the maximum amount of times it will be retried
/// `base_retry_delay_ms` is how much to wait between retries. It does exponential backoffs, doubling this value every time.
pub async fn retry<V, Fut, Func, RetryFunc, Error>(
    logger: Option<&Logger>,
    // Function to be retried.
    mut func: Func,
    // Function that tells whether an error should be retried.
    mut should_retry: RetryFunc,
    retry_logic: RetryLogic,
    retry_num: usize,
) -> Result<(V, RetryAttemptsCount), Error>
where
    V: Send + 'static,
    Fut: Future<Output = Result<V, Error>>,
    Func: FnMut(usize) -> Fut + Send,
    RetryFunc: FnMut(&Error) -> bool + Send,
{
    let mut attempt = 1;
    loop {
        let res = func(attempt).await;
        match res {
            Ok(res) => {
                return Ok((res, RetryAttemptsCount(attempt)));
            }
            Err(err) if attempt < retry_num && should_retry(&err) => {
                if let Some(logger) = logger {
                    info!(
                        logger,
                        "retrying attempt {} of {}...",
                        attempt + 1,
                        retry_num
                    );
                }

                tokio::time::sleep(retry_logic.delay(attempt)).await;
                attempt += 1;
            }
            Err(err) => return Err(err),
        }
    }
}

#[test]
fn test_exponential() {
    let logic = RetryLogic::Exponential {
        base: Duration::from_secs(1),
        factor: 2.0,
    };
    assert_eq!(logic.delay(0), Duration::from_secs(1));
    assert_eq!(logic.delay(1), Duration::from_secs(2));
    assert_eq!(logic.delay(2), Duration::from_secs(4));
    let logic = RetryLogic::Exponential {
        base: Duration::from_secs(8),
        factor: 1.5,
    };
    assert_eq!(logic.delay(1), Duration::from_secs(12));
    assert_eq!(logic.delay(2), Duration::from_secs(18));
}

#[test]
fn test_exponential_jitter() {
    let half_sec = Duration::from_millis(500);
    let one_sec = Duration::from_secs(1);
    let two_sec = Duration::from_secs(2);
    let logic = RetryLogic::ExponentialWithJitter {
        base: one_sec.clone(),
        factor: 2.0,
        jitter: half_sec.clone(),
    };
    let d = logic.delay(0);
    assert!(d >= one_sec && d <= one_sec + half_sec);
    let d = logic.delay(1);
    assert!(d >= two_sec && d <= two_sec + half_sec);
}


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

use std::future::Future;
use std::time::Duration;

use anyhow::Error;
use anyhow::Result;
use futures::future;
use futures::future::Either;
use futures::TryFutureExt;
use slog::error;
use slog::info;
use slog::Logger;
use tokio::signal::unix::signal;
use tokio::signal::unix::SignalKind;
use tokio::time;

/// Run a server future, and wait until a termination signal is received.
///
/// When the termination signal is received, the `quiesce` callback is called.
/// This should perform any steps required to quiesce the server, for example
/// by removing this instance from routing configuration, or asking the load
/// balancer to stop sending requests to this instance.  Requests that do
/// arrive should still be accepted.
///
/// After the configured quiesce timeout, the `shutdown` future is awaited.
/// This should do any additional work to stop accepting connections and wait
/// until all outstanding requests have been handled. The `server` future
/// continues to run while `shutdown` is being awaited.
///
/// Once both `shutdown` and `server` have completed, the process
/// exits. If `shutdown_timeout` is exceeded, the server process is canceled
/// and an error is returned.
pub async fn run_until_terminated<Server, QuiesceFn, ShutdownFut>(
    server: Server,
    logger: &Logger,
    quiesce: QuiesceFn,
    shutdown_grace_period: Duration,
    shutdown: ShutdownFut,
    shutdown_timeout: Duration,
) -> Result<(), Error>
where
    Server: Future<Output = Result<(), Error>> + Send + 'static,
    QuiesceFn: FnOnce(),
    ShutdownFut: Future<Output = ()>,
{
    // We want to prevent Folly's signal handlers overriding our
    // intended action with a termination signal. Mononoke server,
    // in particular, depends on this - otherwise our attempts to
    // catch and handle SIGTERM turn into Folly backtracing and killing us.
    unsafe {
        libc::signal(libc::SIGTERM, libc::SIG_DFL);
    }

    let mut terminate = signal(SignalKind::terminate())?;
    let mut interrupt = signal(SignalKind::interrupt())?;

    let terminate = terminate.recv();
    let interrupt = interrupt.recv();
    futures::pin_mut!(terminate, interrupt);

    // This future becomes ready when we receive a termination signal
    let signalled = future::select(terminate, interrupt);

    // Spawn the server onto its own task
    let server_handle = tokio::task::spawn(server);

    // Now wait for the termination signal, or a server exit.
    let server_result_or_handle = match future::select(server_handle, signalled).await {
        Either::Left((server_result, _)) => {
            let server_result = server_result.map_err(Error::from).and_then(|res| res);
            match server_result.as_ref() {
                Ok(()) => {
                    error!(&logger, "Server has exited! Starting shutdown...");
                }
                Err(e) => {
                    error!(
                        &logger,
                        "Server exited with an error! Starting shutdown... Error: {:?}", e
                    );
                }
            }
            Either::Left(server_result)
        }
        Either::Right((_, server_handle)) => {
            info!(&logger, "Signalled! Starting shutdown...");
            Either::Right(server_handle)
        }
    };

    // Shutting down: wait for the grace period.
    quiesce();
    info!(
        &logger,
        "Waiting {}s before shutting down server",
        shutdown_grace_period.as_secs(),
    );

    time::sleep(shutdown_grace_period).await;

    let shutdown = async move {
        shutdown.await;
        match server_result_or_handle {
            Either::Left(server_result) => server_result,
            Either::Right(server_handle) => server_handle.await?,
        }
    };

    info!(&logger, "Shutting down...");
    time::timeout(shutdown_timeout, shutdown)
        .map_err(|_| Error::msg("Timed out shutting down server"))
        .await?
}


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

use std::collections::HashMap;
use std::hash::Hash;
use std::time::Duration;

use anyhow::Result;
use cachelib::get_cached;
use cachelib::set_cached;
use cachelib::Abomonation;
use cachelib::VolatileLruCachePool;

use crate::mock_store::MockStore;
use crate::CachelibKey;

#[derive(Clone)]
pub enum CachelibHandler<T> {
    Real(VolatileLruCachePool),
    Mock(MockStore<T>),
    Noop,
}

impl<T> From<VolatileLruCachePool> for CachelibHandler<T> {
    fn from(cache: VolatileLruCachePool) -> Self {
        CachelibHandler::Real(cache)
    }
}

impl<T: Abomonation + Clone + Send + 'static> CachelibHandler<T> {
    pub(crate) fn get_multiple_from_cachelib<Key: Eq + Hash>(
        &self,
        keys: Vec<(Key, CachelibKey)>,
    ) -> Result<(HashMap<Key, T>, Vec<(Key, CachelibKey)>)> {
        let mut fetched = HashMap::new();
        let mut left_to_fetch = Vec::new();

        for (key, cache_key) in keys {
            match self.get_cached(&cache_key.0)? {
                Some(value) => {
                    fetched.insert(key, value);
                }
                None => {
                    left_to_fetch.push((key, cache_key));
                }
            }
        }

        Ok((fetched, left_to_fetch))
    }

    pub fn get_cached(&self, key: &String) -> Result<Option<T>> {
        match self {
            CachelibHandler::Real(ref cache) => get_cached(cache, key),
            CachelibHandler::Mock(store) => Ok(store.get(key)),
            CachelibHandler::Noop => Ok(None),
        }
    }

    pub fn set_cached(&self, key: &str, value: &T, ttl: Option<Duration>) -> Result<bool> {
        match self {
            CachelibHandler::Real(ref cache) => set_cached(cache, key, value, ttl),
            CachelibHandler::Mock(store) => {
                store.set(key, value.clone());
                Ok(true)
            }
            CachelibHandler::Noop => Ok(false),
        }
    }

    pub fn create_mock() -> Self {
        CachelibHandler::Mock(MockStore::new())
    }

    pub fn create_noop() -> Self {
        CachelibHandler::Noop
    }

    #[cfg(test)]
    pub(crate) fn gets_count(&self) -> usize {
        use std::sync::atomic::Ordering;
        match self {
            CachelibHandler::Real(_) | CachelibHandler::Noop => unimplemented!(),
            CachelibHandler::Mock(MockStore { ref get_count, .. }) => {
                get_count.load(Ordering::SeqCst)
            }
        }
    }

    pub fn mock_store(&self) -> Option<&MockStore<T>> {
        match self {
            CachelibHandler::Real(_) | CachelibHandler::Noop => None,
            CachelibHandler::Mock(ref mock) => Some(mock),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use quickcheck::quickcheck;
    use quickcheck::TestResult;

    use super::*;

    quickcheck! {
        fn multiple_roundtrip(
            initial_keys: HashMap<String, String>,
            keys_to_query: HashSet<String>
        ) -> TestResult {
            let get_query = keys_to_query.clone().into_iter().map(|key| (key.clone(),  CachelibKey(key))).collect();

            let mock_cachelib = MockStore::new();
            let cachelib_handler = CachelibHandler::Mock(mock_cachelib.clone());

            for (k, v) in initial_keys.iter() {
                let _ = cachelib_handler.set_cached(k, v, None);
            }

            if mock_cachelib.data() != initial_keys {
                return TestResult::error("After fill the content of cache is incorrect");
            }

            let (fetched, left) = cachelib_handler.get_multiple_from_cachelib(get_query).unwrap();

            for (key, cache_key) in &left {
                if key != &cache_key.0 {
                    return TestResult::error("Key and cache key got mixed in left!");
                }
                if initial_keys.get(key).is_some() {
                    return TestResult::error("After get_multiple_from_cachelib left is incorrect");
                }
            }

            for (key, val) in fetched.iter() {
                if initial_keys.get(key) != Some(val) {
                    return TestResult::error("Wrong value returned from get");
                }
            }

            if fetched.len() + left.len() != keys_to_query.len() {
                return TestResult::error("Returned wrong number of items from get");
            }

            let left: HashSet<_> = left.into_iter().map(|(key, _)| key).collect();
            let mut fetched: HashSet<_> = fetched.into_keys().collect();

            if fetched.len() + left.len() != keys_to_query.len() {
                return TestResult::error("Returned wrong number of unique items from get");
            }

            fetched.extend(left);

            if fetched != keys_to_query {
                return TestResult::error("Didn't return all keys from get");
            }

            TestResult::passed()
        }
    }
}


/*
 * 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 abomonation::Abomonation;
use cachelib::VolatileLruCachePool;
use memcache::MemcacheClient;

use crate::CachelibHandler;
use crate::MemcacheHandler;

/// Builder to construct caches, depending on the desired caching mode.
pub enum CacheHandlerFactory {
    /// Caching is via a local cache (cachelib) and a shared cache (memcache).
    Shared {
        /// The cachelib pool to use for local caching.
        cachelib_pool: VolatileLruCachePool,

        /// The memcache client to use for shared caching.
        memcache_client: MemcacheClient,
    },

    /// Caching is via a local cache (cachelib) only.
    Local { cachelib_pool: VolatileLruCachePool },

    /// Caching is mocked for testing purposes, with items cached in an
    /// in-memory store.
    Mocked,

    /// Caching is always a no-op.
    Noop,
}

impl CacheHandlerFactory {
    /// Build cachelib cache handler.
    pub fn cachelib<T>(&self) -> CachelibHandler<T>
    where
        T: Abomonation + Send + Clone + 'static,
    {
        match self {
            Self::Shared { cachelib_pool, .. } | Self::Local { cachelib_pool, .. } => {
                cachelib_pool.clone().into()
            }
            Self::Mocked => CachelibHandler::create_mock(),
            Self::Noop => CachelibHandler::create_noop(),
        }
    }

    /// Build memcache cache handler.
    pub fn memcache(&self) -> MemcacheHandler {
        match self {
            Self::Shared {
                memcache_client, ..
            } => memcache_client.clone().into(),
            Self::Mocked => MemcacheHandler::create_mock(),
            Self::Noop | Self::Local { .. } => MemcacheHandler::create_noop(),
        }
    }
}