/*
 * 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 futures_old::Future;
use futures_old::Poll;
use futures_old::Stream;

pub struct Monitor<T, P> {
    inner: T,
    // We don't actually do anything with this. We just rely on the fact that upon this Monitor
    // being dropped, P will be dropped.
    #[allow(dead_code)]
    payload: P,
}

impl<T, P> Monitor<T, P> {
    pub fn new(inner: T, payload: P) -> Self {
        Self { inner, payload }
    }
}

impl<F: Future, P> Future for Monitor<F, P> {
    type Item = F::Item;
    type Error = F::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.inner.poll()
    }
}

impl<S: Stream, P> Stream for Monitor<S, P> {
    type Item = S::Item;
    type Error = S::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.inner.poll()
    }
}


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

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

use anyhow::Error;
use blobrepo::BlobRepo;
use blobrepo_hg::to_hg_bookmark_stream;
use blobrepo_hg::BlobRepoHg;
use bookmarks::Bookmark;
use bookmarks::BookmarkCategory;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkPagination;
use bookmarks::BookmarkPrefix;
use bookmarks::BookmarksRef;
use bookmarks::Freshness;
use context::CoreContext;
use futures::compat::Future01CompatExt;
use futures::compat::Stream01CompatExt;
use futures::future;
use futures::FutureExt;
use futures::Stream;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use futures_01_ext::StreamExt as OldStreamExt;
use futures_ext::FbFutureExt;
use futures_ext::FbTryFutureExt;
use futures_old::Future;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgChangesetId;
use mononoke_api::Repo;
use tunables::tunables;
use warm_bookmarks_cache::BookmarksCache;

// We'd like to give user a consistent view of thier bookmarks for the duration of the
// whole Mononoke session. SessionBookmarkCache is used for that.
pub struct SessionBookmarkCache<R = Arc<Repo>> {
    cached_publishing_bookmarks_maybe_stale: Arc<Mutex<Option<HashMap<Bookmark, HgChangesetId>>>>,
    repo: R,
}

pub trait BookmarkCacheRepo {
    fn blobrepo(&self) -> &BlobRepo;

    fn repo_client_use_warm_bookmarks_cache(&self) -> bool;

    fn warm_bookmarks_cache(&self) -> &Arc<dyn BookmarksCache + Send + Sync>;
}

impl BookmarkCacheRepo for Arc<Repo> {
    fn blobrepo(&self) -> &BlobRepo {
        Repo::blob_repo(self)
    }

    fn repo_client_use_warm_bookmarks_cache(&self) -> bool {
        Repo::config(self).repo_client_use_warm_bookmarks_cache
    }

    fn warm_bookmarks_cache(&self) -> &Arc<dyn BookmarksCache + Send + Sync> {
        Repo::warm_bookmarks_cache(self)
    }
}

fn bookmarks_timeout() -> Duration {
    let timeout = tunables()
        .repo_client_bookmarks_timeout_secs()
        .unwrap_or_default();
    if timeout > 0 {
        Duration::from_secs(timeout as u64)
    } else {
        Duration::from_secs(3 * 60)
    }
}

impl<R> SessionBookmarkCache<R>
where
    R: BookmarkCacheRepo,
{
    pub fn new(repo: R) -> Self {
        Self {
            cached_publishing_bookmarks_maybe_stale: Arc::new(Mutex::new(None)),
            repo,
        }
    }

    pub fn drop_cache(&self) {
        let _ = self
            .cached_publishing_bookmarks_maybe_stale
            .lock()
            .expect("lock poisoned")
            .take();
    }

    pub async fn get_publishing_bookmarks(
        &self,
        ctx: CoreContext,
    ) -> Result<HashMap<Bookmark, HgChangesetId>, Error> {
        let maybe_cache = {
            self.cached_publishing_bookmarks_maybe_stale
                .lock()
                .expect("lock poisoned")
                .clone()
        };

        match maybe_cache {
            None => {
                self.get_publishing_bookmarks_maybe_stale_updating_cache(ctx)
                    .await
            }
            Some(bookmarks) => Ok(bookmarks),
        }
    }

    // TODO(stash): consider updating from leader db - T70835157?
    pub fn update_publishing_bookmarks_after_push(
        &self,
        ctx: CoreContext,
    ) -> impl Future<Item = (), Error = Error> + '_ {
        let cache = self.cached_publishing_bookmarks_maybe_stale.clone();
        // We just updated the bookmark, so go and fetch them from db to return
        // the newer version
        self.get_publishing_maybe_stale_from_db(ctx)
            .map(move |bookmarks| {
                update_publishing_bookmarks_maybe_stale_cache_raw(cache, bookmarks)
            })
    }

    pub async fn get_bookmarks_by_prefix(
        &self,
        ctx: &CoreContext,
        prefix: &BookmarkPrefix,
        return_max: u64,
    ) -> Result<impl Stream<Item = Result<(BookmarkKey, HgChangesetId), Error>> + '_, Error> {
        let mut kinds = vec![BookmarkKind::Scratch];

        let mut result = HashMap::new();
        if let Some(warm_bookmarks_cache) = self.get_warm_bookmark_cache() {
            let warm_bookmarks = warm_bookmarks_cache
                .list(
                    ctx,
                    prefix,
                    &BookmarkPagination::FromStart,
                    Some(return_max),
                )
                .await?;
            for (book, (cs_id, _)) in warm_bookmarks {
                result.insert(book, cs_id);
            }
        } else {
            // Warm bookmark cache was disabled, so we'll need to fetch publishing
            // bookmarks from db.
            kinds.extend(BookmarkKind::ALL_PUBLISHING);
        }

        let left_to_fetch = return_max.saturating_sub(result.len().try_into().unwrap());
        let new_bookmarks = if left_to_fetch > 0 {
            self.repo
                .blobrepo()
                .bookmarks()
                .list(
                    ctx.clone(),
                    Freshness::MaybeStale,
                    prefix,
                    BookmarkCategory::ALL,
                    &kinds,
                    &BookmarkPagination::FromStart,
                    left_to_fetch,
                )
                .map_ok(|(bookmark, cs_id)| (bookmark.into_key(), cs_id))
                .left_stream()
        } else {
            futures::stream::empty().map(Ok).right_stream()
        };

        Ok(to_hg_bookmark_stream(
            self.repo.blobrepo(),
            ctx,
            futures::stream::iter(result.into_iter())
                .map(Ok)
                .chain(new_bookmarks),
        ))
    }

    // Tries to fetch a bookmark from warm bookmark cache first, but if the bookmark is not found
    // then fallbacks to fetching from db.
    pub async fn get_bookmark(
        &self,
        ctx: CoreContext,
        bookmark: BookmarkKey,
    ) -> Result<Option<HgChangesetId>, Error> {
        if let Some(warm_bookmarks_cache) = self.get_warm_bookmark_cache() {
            if let Some(cs_id) = warm_bookmarks_cache.get(&ctx, &bookmark).await? {
                return self
                    .repo
                    .blobrepo()
                    .derive_hg_changeset(&ctx, cs_id)
                    .map_ok(Some)
                    .await;
            }
        }

        self.repo.blobrepo().get_bookmark_hg(ctx, &bookmark).await
    }

    async fn get_publishing_bookmarks_maybe_stale_updating_cache(
        &self,
        ctx: CoreContext,
    ) -> Result<HashMap<Bookmark, HgChangesetId>, Error> {
        let cache = self.cached_publishing_bookmarks_maybe_stale.clone();
        let bookmarks = self.get_publishing_maybe_stale_raw(ctx).await?;
        update_publishing_bookmarks_maybe_stale_cache_raw(cache, bookmarks.clone());
        Ok(bookmarks)
    }

    async fn get_publishing_maybe_stale_raw(
        &self,
        ctx: CoreContext,
    ) -> Result<HashMap<Bookmark, HgChangesetId>, Error> {
        if let Some(warm_bookmarks_cache) = self.get_warm_bookmark_cache() {
            ctx.scuba()
                .clone()
                .log_with_msg("Fetching bookmarks from Warm bookmarks cache", None);
            let s = futures_old::stream::iter_ok(
                warm_bookmarks_cache
                    .list(
                        &ctx,
                        &BookmarkPrefix::empty(),
                        &BookmarkPagination::FromStart,
                        Some(std::u64::MAX),
                    )
                    .await?
                    .into_iter()
                    .map(|(name, (cs_id, kind))| (Bookmark::new(name, kind), cs_id)),
            )
            .boxify()
            .compat();
            return to_hg_bookmark_stream(self.repo.blobrepo(), &ctx, s)
                .try_collect()
                .await;
        }
        self.get_publishing_maybe_stale_from_db(ctx).compat().await
    }

    fn get_warm_bookmark_cache(&self) -> Option<&Arc<dyn BookmarksCache + Send + Sync>> {
        if self.repo.repo_client_use_warm_bookmarks_cache() {
            if !tunables()
                .disable_repo_client_warm_bookmarks_cache()
                .unwrap_or_default()
            {
                return Some(self.repo.warm_bookmarks_cache());
            }
        }

        None
    }

    fn get_publishing_maybe_stale_from_db(
        &self,
        ctx: CoreContext,
    ) -> impl Future<Item = HashMap<Bookmark, HgChangesetId>, Error = Error> + '_ {
        self.repo
            .blobrepo()
            .get_publishing_bookmarks_maybe_stale_hg(ctx)
            .try_fold(HashMap::new(), |mut map, item| {
                map.insert(item.0, item.1);
                future::ready(Ok(map))
            })
            .timeout(bookmarks_timeout())
            .flatten_err()
            .boxed()
            .compat()
    }
}

fn update_publishing_bookmarks_maybe_stale_cache_raw(
    cache: Arc<Mutex<Option<HashMap<Bookmark, HgChangesetId>>>>,
    bookmarks: HashMap<Bookmark, HgChangesetId>,
) {
    let mut maybe_cache = cache.lock().expect("lock poisoned");
    *maybe_cache = Some(bookmarks);
}

#[cfg(test)]
mod test {
    use bookmarks::BookmarkKey;
    use bookmarks::BookmarkUpdateLogArc;
    use bookmarks::BookmarksArc;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_api_types::InnerRepo;
    use phases::PhasesArc;
    use repo_derived_data::RepoDerivedDataArc;
    use repo_identity::RepoIdentityArc;
    use tests_utils::bookmark;
    use tests_utils::CreateCommitContext;
    use warm_bookmarks_cache::WarmBookmarksCacheBuilder;

    use super::*;

    struct BasicTestRepo {
        repo: BlobRepo,
        wbc: Option<Arc<dyn BookmarksCache + Send + Sync>>,
    }

    impl BookmarkCacheRepo for BasicTestRepo {
        fn blobrepo(&self) -> &BlobRepo {
            &self.repo
        }

        fn repo_client_use_warm_bookmarks_cache(&self) -> bool {
            self.wbc.is_some()
        }

        fn warm_bookmarks_cache(&self) -> &Arc<dyn BookmarksCache + Send + Sync> {
            self.wbc.as_ref().unwrap()
        }
    }

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

        let repo: InnerRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let cs_id = CreateCommitContext::new_root(&ctx, &repo.blob_repo)
            .add_file("1", "1")
            .commit()
            .await?;

        let hg_cs_id = repo.blob_repo.derive_hg_changeset(&ctx, cs_id).await?;
        bookmark(&ctx, &repo.blob_repo, "prefix/scratchbook")
            .create_scratch(cs_id)
            .await?;

        bookmark(&ctx, &repo.blob_repo, "prefix/publishing")
            .create_publishing(cs_id)
            .await?;

        bookmark(&ctx, &repo.blob_repo, "prefix/pulldefault")
            .create_pull_default(cs_id)
            .await?;

        // Let's try without WarmBookmarkCache first
        println!("No warm bookmark cache");
        let session_bookmark_cache = SessionBookmarkCache::new(BasicTestRepo {
            repo: repo.blob_repo.clone(),
            wbc: None,
        });
        validate(&ctx, hg_cs_id, &session_bookmark_cache).await?;

        // Let's try with WarmBookmarkCache next
        println!("With warm bookmark cache");
        let mut builder = WarmBookmarksCacheBuilder::new(
            ctx.clone(),
            repo.bookmarks_arc(),
            repo.bookmark_update_log_arc(),
            repo.repo_identity_arc(),
        );
        builder.add_hg_warmers(&repo.repo_derived_data_arc(), &repo.phases_arc())?;
        let wbc = builder.build().await?;
        let session_bookmark_cache = SessionBookmarkCache::new(BasicTestRepo {
            repo: repo.blob_repo.clone(),
            wbc: Some(Arc::new(wbc)),
        });
        validate(&ctx, hg_cs_id, &session_bookmark_cache).await?;

        Ok(())
    }

    async fn validate(
        ctx: &CoreContext,
        hg_cs_id: HgChangesetId,
        session_bookmark_cache: &SessionBookmarkCache<BasicTestRepo>,
    ) -> Result<(), Error> {
        let maybe_hg_cs_id = session_bookmark_cache
            .get_bookmark(ctx.clone(), BookmarkKey::new("prefix/scratchbook")?)
            .await?;
        assert_eq!(maybe_hg_cs_id, Some(hg_cs_id));

        let res = session_bookmark_cache
            .get_bookmarks_by_prefix(ctx, &BookmarkPrefix::new("prefix")?, 3)
            .await?
            .try_collect::<HashMap<_, _>>()
            .await?;
        assert_eq!(
            hashmap! {
                BookmarkKey::new("prefix/scratchbook")? => hg_cs_id,
                BookmarkKey::new("prefix/publishing")? => hg_cs_id,
                BookmarkKey::new("prefix/pulldefault")? => hg_cs_id,
            },
            res
        );

        let res = session_bookmark_cache
            .get_bookmarks_by_prefix(ctx, &BookmarkPrefix::new("prefix")?, 1)
            .await?
            .try_collect::<HashMap<_, _>>()
            .await?;
        assert_eq!(res.len(), 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.
 */

#![cfg(test)]
use blobstore::Loadable;
use fbinit::FacebookInit;
use fixtures::ManyFilesDirs;
use fixtures::TestRepoFixture;
use futures::compat::Future01CompatExt;
use manifest::Entry;
use manifest::ManifestOps;
use maplit::hashset;
use mercurial_derivation::DeriveHgChangeset;
use metaconfig_types::LfsParams;
use mononoke_api::Repo;
use mononoke_types_mocks::changesetid::ONES_CSID;
use repo_blobstore::RepoBlobstoreRef;
use scuba_ext::MononokeScubaSampleBuilder;
use serde_json::json;
use tests_utils::CreateCommitContext;

use super::*;

#[test]
fn test_parsing_caps_simple() {
    assert_eq!(
        parse_utf8_getbundle_caps(b"cap"),
        Some((String::from("cap"), HashMap::new())),
    );

    let caps = b"bundle2=HG20";

    assert_eq!(
        parse_utf8_getbundle_caps(caps),
        Some((
            String::from("bundle2"),
            hashmap! { "HG20".to_string() => hashset!{} }
        )),
    );

    let caps = b"bundle2=HG20%0Ab2x%253Ainfinitepush%0Ab2x%253Ainfinitepushscratchbookmarks\
        %0Ab2x%253Arebase%0Abookmarks%0Achangegroup%3D01%2C02%0Adigests%3Dmd5%2Csha1%2Csha512%0A\
        error%3Dabort%2Cunsupportedcontent%2Cpushraced%2Cpushkey%0Ahgtagsfnodes%0Alistkeys%0A\
        pushkey%0Aremote-changegroup%3Dhttp%2Chttps%0Aremotefilelog%3DTrue%0Atreemanifest%3DTrue%0Atreeonly%3DTrue";

    assert_eq!(
        parse_utf8_getbundle_caps(caps),
        Some((
            String::from("bundle2"),
            hashmap! {
                "HG20".to_string() => hashset!{},
                "b2x:rebase".to_string() => hashset!{},
                "digests".to_string() => hashset!{"md5".to_string(), "sha512".to_string(), "sha1".to_string()},
                "listkeys".to_string() => hashset!{},
                "remotefilelog".to_string() => hashset!{"True".to_string()},
                "hgtagsfnodes".to_string() => hashset!{},
                "bookmarks".to_string() => hashset!{},
                "b2x:infinitepushscratchbookmarks".to_string() => hashset!{},
                "treeonly".to_string() => hashset!{"True".to_string()},
                "pushkey".to_string() => hashset!{},
                "error".to_string() => hashset!{
                    "pushraced".to_string(),
                    "pushkey".to_string(),
                    "unsupportedcontent".to_string(),
                    "abort".to_string(),
                },
                "b2x:infinitepush".to_string() => hashset!{},
                "changegroup".to_string() => hashset!{"01".to_string(), "02".to_string()},
                "remote-changegroup".to_string() => hashset!{"http".to_string(), "https".to_string()},
                "treemanifest".to_string() => hashset!{"True".to_string()},
            }
        )),
    );
}

#[fbinit::test]
fn get_changed_manifests_stream_test(fb: FacebookInit) -> Result<(), Error> {
    let runtime = tokio::runtime::Runtime::new()?;
    runtime.block_on(get_changed_manifests_stream_test_impl(fb))
}

async fn get_changed_manifests_stream_test_impl(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo = ManyFilesDirs::getrepo(fb).await;

    // Commit that has only dir2 directory
    let root_mf_id = HgChangesetId::from_str("051946ed218061e925fb120dac02634f9ad40ae2")?
        .load(&ctx, &repo.repo_blobstore().clone())
        .await?
        .manifestid();

    let fetched_mfs = fetch_mfs(
        &ctx,
        &repo,
        root_mf_id,
        HgManifestId::new(NULL_HASH),
        None,
        65536,
    )
    .await?;

    let mut res = fetched_mfs
        .into_iter()
        .map(|(_, path)| path)
        .collect::<Vec<_>>();
    res.sort();
    let mut expected = vec![None, Some(MPath::new("dir2")?)];
    expected.sort();
    assert_eq!(res, expected);

    // Now commit that added a few files and directories

    let root_mf_id = HgChangesetId::from_str("d261bc7900818dea7c86935b3fb17a33b2e3a6b4")?
        .load(&ctx, &repo.repo_blobstore().clone())
        .await?
        .manifestid();

    let base_root_mf_id = HgChangesetId::from_str("2f866e7e549760934e31bf0420a873f65100ad63")?
        .load(&ctx, &repo.repo_blobstore().clone())
        .await?
        .manifestid();

    let fetched_mfs = fetch_mfs(&ctx, &repo, root_mf_id, base_root_mf_id, None, 65536).await?;

    let mut res = fetched_mfs
        .into_iter()
        .map(|(_, path)| path)
        .collect::<Vec<_>>();
    res.sort();
    let mut expected = vec![
        None,
        Some(MPath::new("dir1")?),
        Some(MPath::new("dir1/subdir1")?),
        Some(MPath::new("dir1/subdir1/subsubdir1")?),
        Some(MPath::new("dir1/subdir1/subsubdir2")?),
    ];
    expected.sort();
    assert_eq!(res, expected);

    Ok(())
}

#[fbinit::test]
fn get_changed_manifests_stream_test_depth(fb: FacebookInit) -> Result<(), Error> {
    let runtime = tokio::runtime::Runtime::new()?;
    runtime.block_on(get_changed_manifests_stream_test_depth_impl(fb))
}

async fn get_changed_manifests_stream_test_depth_impl(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo = ManyFilesDirs::getrepo(fb).await;

    let root_mf_id = HgChangesetId::from_str("d261bc7900818dea7c86935b3fb17a33b2e3a6b4")?
        .load(&ctx, &repo.repo_blobstore().clone())
        .await?
        .manifestid();

    let base_mf_id = HgManifestId::new(NULL_HASH);
    let fetched_mfs = fetch_mfs(&ctx, &repo, root_mf_id, base_mf_id, None, 65536).await?;

    let paths = fetched_mfs
        .into_iter()
        .map(|(_, path)| path)
        .collect::<Vec<_>>();

    let max_depth = paths
        .iter()
        .map(|path| match path {
            Some(path) => path.num_components(),
            None => 0,
        })
        .max()
        .unwrap();

    for depth in 0..max_depth + 1 {
        println!("depth: {}", depth);
        let fetched_mfs = fetch_mfs(&ctx, &repo, root_mf_id, base_mf_id, None, depth).await?;
        let mut actual = fetched_mfs
            .into_iter()
            .map(|(_, path)| path)
            .collect::<Vec<_>>();
        actual.sort();
        let iter = paths.clone().into_iter();
        // We have a weird hard-coded behaviour for depth=1 that we are preserving for now
        let mut expected: Vec<_> = if depth == 1 {
            let expected: Vec<_> = iter.filter(|path| path.is_none()).collect();
            assert_eq!(expected.len(), 1);
            expected
        } else {
            iter.filter(|path| match path {
                Some(path) => path.num_components() <= depth,
                None => true,
            })
            .collect()
        };
        expected.sort();
        assert_eq!(actual, expected);
    }

    Ok(())
}

#[fbinit::test]
fn get_changed_manifests_stream_test_base_path(fb: FacebookInit) -> Result<(), Error> {
    let runtime = tokio::runtime::Runtime::new()?;
    runtime.block_on(get_changed_manifests_stream_test_base_path_impl(fb))
}

async fn get_changed_manifests_stream_test_base_path_impl(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo = ManyFilesDirs::getrepo(fb).await;

    let root_mf_id = HgChangesetId::from_str("d261bc7900818dea7c86935b3fb17a33b2e3a6b4")?
        .load(&ctx, &repo.repo_blobstore().clone())
        .await?
        .manifestid();

    let base_mf_id = HgManifestId::new(NULL_HASH);
    let fetched_mfs = fetch_mfs(&ctx, &repo, root_mf_id, base_mf_id, None, 65536).await?;

    for (hash, path) in &fetched_mfs {
        println!("base path: {:?}", path);
        let mut actual = fetch_mfs(&ctx, &repo, *hash, base_mf_id, path.clone(), 65536).await?;
        actual.sort();

        let mut expected: Vec<_> = fetched_mfs
            .clone()
            .into_iter()
            .filter(|(_, curpath)| match &path {
                Some(path) => {
                    let elems = MPath::iter_opt(curpath.as_ref());
                    path.is_prefix_of(elems)
                }
                None => true,
            })
            .collect();
        expected.sort();
        assert_eq!(actual, expected);
    }

    Ok(())
}

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

    assert!(!run_and_check_if_lfs(&ctx, LfsParams::default()).await?);

    // Rollout percentage is 100 and threshold is set - enable lfs
    let lfs_params = LfsParams {
        threshold: Some(5),
        rollout_percentage: 100,
        ..Default::default()
    };
    assert!(run_and_check_if_lfs(&ctx, lfs_params).await?);

    // Rollout percentage is 0 - no lfs is enabled
    let lfs_params = LfsParams {
        threshold: Some(5),
        rollout_percentage: 0,
        ..Default::default()
    };
    assert!(!run_and_check_if_lfs(&ctx, lfs_params).await?);

    // Rollout percentage is 100, but threshold is too high
    let lfs_params = LfsParams {
        threshold: Some(500),
        rollout_percentage: 100,
        ..Default::default()
    };
    assert!(!run_and_check_if_lfs(&ctx, lfs_params).await?);
    Ok(())
}

#[fbinit::test]
async fn test_maybe_validate_pushed_bonsais(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;
    let commit = CreateCommitContext::new_root(&ctx, &repo)
        .add_file("largefile", "11111_11111")
        .commit()
        .await?;

    let hg_cs_id = repo.derive_hg_changeset(&ctx, commit).await?;

    // No replay data - ignore
    maybe_validate_pushed_bonsais(&ctx, &repo, &None).await?;

    // Has replay data, but no hgbonsaimapping - ignore
    maybe_validate_pushed_bonsais(&ctx, &repo, &Some("{}".to_string())).await?;

    // Has valid replay data - should succeed
    maybe_validate_pushed_bonsais(
        &ctx,
        &repo,
        &Some(
            json!({
                "hgbonsaimapping": {
                    format!("{}", hg_cs_id): commit,
                }
            })
            .to_string(),
        ),
    )
    .await?;

    // Additional fields doesn't change the result
    maybe_validate_pushed_bonsais(
        &ctx,
        &repo,
        &Some(
            json!({
                "hgbonsaimapping": {
                    format!("{}", hg_cs_id): commit,
                },
                "somefield": "somevalue"
            })
            .to_string(),
        ),
    )
    .await?;

    // Now invalid bonsai - should fail
    assert!(
        maybe_validate_pushed_bonsais(
            &ctx,
            &repo,
            &Some(
                json!({
                    "hgbonsaimapping": {
                        format!("{}", hg_cs_id): ONES_CSID,
                    },
                    "somefield": "somevalue"
                })
                .to_string(),
            ),
        )
        .await
        .is_err()
    );

    // Now invalid hgbonsaimapping field - should fail
    assert!(
        maybe_validate_pushed_bonsais(
            &ctx,
            &repo,
            &Some(
                json!({
                    "hgbonsaimapping": "somevalue"
                })
                .to_string(),
            ),
        )
        .await
        .is_err()
    );
    Ok(())
}

async fn run_and_check_if_lfs(ctx: &CoreContext, lfs_params: LfsParams) -> Result<bool, Error> {
    let repo = Arc::new(
        test_repo_factory::TestRepoFactory::new(ctx.fb)?
            .with_config_override(|config| config.lfs = lfs_params)
            .build::<Repo>()
            .await?,
    );
    let commit = CreateCommitContext::new_root(ctx, &repo)
        .add_file("largefile", "11111_11111")
        .commit()
        .await?;

    let hg_cs_id = repo.derive_hg_changeset(ctx, commit).await?;

    let hg_cs = hg_cs_id.load(ctx, &repo.repo_blobstore().clone()).await?;

    let path = MPath::new("largefile")?;
    let maybe_entry = hg_cs
        .manifestid()
        .find_entry(
            ctx.clone(),
            repo.repo_blobstore().clone(),
            Some(path.clone()),
        )
        .await?
        .unwrap();

    let filenode_id = match maybe_entry {
        Entry::Leaf((_, filenode_id)) => filenode_id,
        Entry::Tree(_) => {
            panic!("should be a leaf");
        }
    };

    let logging = LoggingContainer::new(
        ctx.fb,
        ctx.logger().clone(),
        MononokeScubaSampleBuilder::with_discard(),
    );

    let repo_client = RepoClient::new(
        repo,
        ctx.session().clone(),
        logging,
        None, // No PushRedirectorArgs
        Default::default(),
        None, // No backup repo source
    );

    let bytes = repo_client
        .getpackv2(stream_old::iter_ok(vec![(path.clone(), vec![filenode_id])]).boxify())
        .concat2()
        .compat()
        .await?;

    let lfs_url: &[u8] = b"version https://git-lfs.github.com/spec/v1";

    let found = bytes.windows(lfs_url.len()).any(|w| w == lfs_url);

    Ok(found)
}

async fn fetch_mfs(
    ctx: &CoreContext,
    repo: &BlobRepo,
    root_mf_id: HgManifestId,
    base_root_mf_id: HgManifestId,
    base_path: Option<MPath>,
    depth: usize,
) -> Result<Vec<(HgManifestId, Option<MPath>)>, Error> {
    let fetched_mfs = get_changed_manifests_stream(
        ctx.clone(),
        repo,
        root_mf_id,
        base_root_mf_id,
        base_path,
        depth,
    )
    .collect()
    .compat()
    .await?;

    // Make sure that Manifest ids are present in the repo
    for (hash, _) in &fetched_mfs {
        hash.load(ctx, repo.repo_blobstore()).await?;
    }
    Ok(fetched_mfs)
}

#[test]
fn test_debug_format_directories() {
    assert_eq!(&debug_format_directories(vec![&"foo"]), "foo,");
    assert_eq!(&debug_format_directories(vec![&"foo,bar"]), "foo:obar,");
    assert_eq!(&debug_format_directories(vec![&"foo", &"bar"]), "foo,bar,");
}

#[test]
fn test_parse_git_lookup() -> Result<(), Error> {
    assert!(parse_git_lookup("ololo").is_none());
    assert!(parse_git_lookup("_gitlookup_hg_badhash").is_none());
    assert!(parse_git_lookup("_gitlookup_git_badhash").is_none());
    assert_eq!(
        parse_git_lookup("_gitlookup_hg_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
        Some(GitLookup::HgToGit(HgChangesetId::from_str(
            "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
        )?))
    );

    assert_eq!(
        parse_git_lookup("_gitlookup_git_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
        Some(GitLookup::GitToHg(GitSha1::from_str(
            "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
        )?))
    );

    Ok(())
}


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

use anyhow::Error;
use blobstore::Blobstore;
use bytes::Bytes;
use context::CoreContext;
use context::PerfCounterType;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use mononoke_types::RepositoryId;
use repo_blobstore::ArcRepoBlobstore;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use thiserror::Error;

#[facet::facet]
pub struct StreamingClone {
    connections: SqlConnections,
    repo_id: RepositoryId,
    repo_blobstore: ArcRepoBlobstore,
}

pub struct StreamingCloneBuilder {
    connections: SqlConnections,
}

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("missing blob {0}")]
    MissingStreamingBlob(String),
    #[error("incorrect size {1} (expected {2}) of corrupt blob {0}")]
    CorruptStreamingBlob(String, usize, usize),
}

pub struct RevlogStreamingChunks {
    pub index_size: usize,
    pub data_size: usize,
    pub index_blobs: Vec<BoxFuture<'static, Result<Bytes, Error>>>,
    pub data_blobs: Vec<BoxFuture<'static, Result<Bytes, Error>>>,
}

impl RevlogStreamingChunks {
    pub fn new() -> Self {
        Self {
            data_size: 0,
            index_size: 0,
            data_blobs: Vec::new(),
            index_blobs: Vec::new(),
        }
    }
}

mononoke_queries! {
    read CountChunks(repo_id: RepositoryId, tag: &str) -> (u64) {
        "SELECT count(*)
         FROM streaming_changelog_chunks
         WHERE repo_id = {repo_id} and tag = {tag}"
    }
    read SelectChunks(repo_id: RepositoryId, tag: &str) -> (Vec<u8>, i32, Vec<u8>, i32) {
        "SELECT idx_blob_name, idx_size, data_blob_name, data_size
         FROM streaming_changelog_chunks
         WHERE repo_id = {repo_id} and tag = {tag}
         ORDER BY chunk_num ASC"
    }

    read SelectSizes(repo_id: RepositoryId, tag: &str) -> (Option<u64>, Option<u64>) {
        "SELECT CAST(SUM(idx_size) AS UNSIGNED), CAST(SUM(data_size) AS UNSIGNED)
         FROM streaming_changelog_chunks
         WHERE repo_id = {repo_id} and tag = {tag}"
    }

    write InsertChunks(
        values: (
            repo_id: RepositoryId,
            tag: &str,
            chunk_num: u32,
            idx_blob_name: &str,
            idx_size: u32,
            data_blob_name: &str,
            data_size: u32,
        )
    ) {
        none,
        "INSERT INTO streaming_changelog_chunks \
            (repo_id, tag, chunk_num, idx_blob_name, idx_size, data_blob_name, data_size) \
            VALUES {values}"
    }

    read SelectMaxChunkNum(repo_id: RepositoryId) -> (Option<u32>) {
        "SELECT max(chunk_num)
         FROM streaming_changelog_chunks
         WHERE repo_id = {repo_id}"
    }
}

impl SqlConstruct for StreamingCloneBuilder {
    const LABEL: &'static str = "streaming-chunks";

    const CREATION_QUERY: &'static str = "
        CREATE TABLE IF NOT EXISTS `streaming_changelog_chunks` (
        `repo_id` int(11) NOT NULL,
        `tag` varbinary(100) NOT NULL DEFAULT '',
        `chunk_num` int(11) NOT NULL,
        `idx_blob_name` varbinary(4096) NOT NULL,
        `idx_size` int(11) NOT NULL,
        `data_blob_name` varbinary(4096) NOT NULL,
        `data_size` int(11) NOT NULL,
        PRIMARY KEY (`repo_id`,`tag`,`chunk_num`)
        )
    ";

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for StreamingCloneBuilder {}

impl StreamingCloneBuilder {
    pub fn build(self, repo_id: RepositoryId, repo_blobstore: ArcRepoBlobstore) -> StreamingClone {
        StreamingClone {
            connections: self.connections,
            repo_id,
            repo_blobstore,
        }
    }
}

fn fetch_blob(
    ctx: CoreContext,
    blobstore: impl Blobstore + 'static,
    key: &[u8],
    expected_size: usize,
) -> BoxFuture<'static, Result<Bytes, Error>> {
    let key = String::from_utf8_lossy(key).into_owned();

    async move {
        let data = blobstore.get(&ctx, &key).await?;

        match data {
            None => Err(ErrorKind::MissingStreamingBlob(key).into()),
            Some(data) if data.as_bytes().len() == expected_size => Ok(data.into_raw_bytes()),
            Some(data) => {
                Err(
                    ErrorKind::CorruptStreamingBlob(key, data.as_bytes().len(), expected_size)
                        .into(),
                )
            }
        }
    }
    .boxed()
}

impl StreamingClone {
    pub async fn count_chunks(&self, ctx: &CoreContext, tag: Option<&str>) -> Result<u64, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let tag = tag.unwrap_or("");
        let res =
            CountChunks::query(&self.connections.read_connection, &self.repo_id, &tag).await?;
        Ok(res.first().map_or(0, |x| x.0))
    }

    pub async fn fetch_changelog(
        &self,
        ctx: CoreContext,
        tag: Option<&str>,
    ) -> Result<RevlogStreamingChunks, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let tag = tag.unwrap_or("");
        let rows =
            SelectChunks::query(&self.connections.read_connection, &self.repo_id, &tag).await?;

        let res = rows.into_iter().fold(
            RevlogStreamingChunks::new(),
            move |mut res, (idx_blob_name, idx_size, data_blob_name, data_size)| {
                let data_size = data_size as usize;
                let idx_size = idx_size as usize;
                res.data_size += data_size;
                res.index_size += idx_size;
                res.data_blobs.push(fetch_blob(
                    ctx.clone(),
                    self.repo_blobstore.clone(),
                    &data_blob_name,
                    data_size,
                ));
                res.index_blobs.push(fetch_blob(
                    ctx.clone(),
                    self.repo_blobstore.clone(),
                    &idx_blob_name,
                    idx_size,
                ));
                res
            },
        );

        Ok(res)
    }

    pub async fn insert_chunks(
        &self,
        ctx: &CoreContext,
        tag: Option<&str>,
        chunks: Vec<(u32, &str, u32, &str, u32)>,
    ) -> Result<(), Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        let tag = tag.unwrap_or("");

        let ref_chunks: Vec<_> = chunks
            .iter()
            .map(|row| (&self.repo_id, &tag, &row.0, &row.1, &row.2, &row.3, &row.4))
            .collect();

        InsertChunks::query(&self.connections.write_connection, &ref_chunks[..]).await?;

        Ok(())
    }

    pub async fn select_index_and_data_sizes(
        &self,
        ctx: &CoreContext,
        tag: Option<&str>,
    ) -> Result<Option<(u64, u64)>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let tag = tag.unwrap_or("");

        let res =
            SelectSizes::query(&self.connections.read_connection, &self.repo_id, &tag).await?;
        let (idx, data) = match res.get(0) {
            Some((Some(idx), Some(data))) => (idx, data),
            _ => {
                return Ok(None);
            }
        };

        Ok(Some((*idx, *data)))
    }

    pub async fn select_max_chunk_num(&self, ctx: &CoreContext) -> Result<Option<u32>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let res =
            SelectMaxChunkNum::query(&self.connections.read_connection, &self.repo_id).await?;
        Ok(res.get(0).and_then(|x| x.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 mercurial_types::HgChangesetId;
use thiserror::Error;

#[derive(Debug, Error)]
pub(crate) enum ErrorKind {
    #[error("Error while uploading data for changesets, hashes: {0:?}")]
    WhileUploadingData(Vec<HgChangesetId>),
}


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

use anyhow::Context;
use anyhow::Error;
use blobrepo::BlobRepo;
use bookmarks_movement::BookmarkMovementError;
use context::CoreContext;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use hooks::CrossRepoPushSource;
use hooks::HookManager;
use hooks::HookRejection;
use hooks::PushAuthoredBy;
use mercurial_derivation::DeriveHgChangeset;

use crate::resolver::HgHookRejection;
use crate::resolver::PostResolveAction;
use crate::resolver::PostResolvePushRebase;
use crate::BundleResolverError;

/// A function to remap hook rejections from Bonsai to Hg.
pub(crate) trait HookRejectionRemapper = (Fn(HookRejection) -> BoxFuture<'static, Result<HgHookRejection, Error>>)
    + Send
    + Sync
    + 'static;

pub(crate) fn make_hook_rejection_remapper(
    ctx: &CoreContext,
    repo: BlobRepo,
) -> Box<dyn HookRejectionRemapper> {
    let ctx = ctx.clone();
    Box::new(
        move |HookRejection {
                  hook_name,
                  cs_id,
                  reason,
              }| {
            let ctx = ctx.clone();
            let repo = repo.clone();
            async move {
                let hg_cs_id = repo.derive_hg_changeset(&ctx, cs_id).await?;
                Ok(HgHookRejection {
                    hook_name,
                    hg_cs_id,
                    reason,
                })
            }
            .boxed()
        },
    )
}

pub(crate) async fn map_hook_rejections(
    rejections: Vec<HookRejection>,
    hook_rejection_remapper: &dyn HookRejectionRemapper,
) -> Result<Vec<HgHookRejection>, Error> {
    stream::iter(rejections)
        .map(move |rejection| async move { (*hook_rejection_remapper)(rejection).await })
        .buffered(10)
        .try_collect()
        .await
        .context("Failed to remap hook rejections")
}

pub async fn run_hooks(
    ctx: &CoreContext,
    repo: &BlobRepo,
    hook_manager: &HookManager,
    action: &PostResolveAction,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<(), BundleResolverError> {
    match action {
        // TODO: Need to run hooks on Push, not just Pushrebase
        PostResolveAction::Push(_) => Ok(()),
        PostResolveAction::InfinitePush(_) => Ok(()),
        PostResolveAction::BookmarkOnlyPushRebase(_) => Ok(()),
        PostResolveAction::PushRebase(action) => {
            run_pushrebase_hooks(ctx, repo, hook_manager, action, cross_repo_push_source).await
        }
    }
}

async fn run_pushrebase_hooks(
    ctx: &CoreContext,
    repo: &BlobRepo,
    hook_manager: &HookManager,
    action: &PostResolvePushRebase,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<(), BundleResolverError> {
    match bookmarks_movement::run_hooks(
        ctx,
        hook_manager,
        action.bookmark_spec.get_bookmark_name(),
        action.uploaded_bonsais.iter(),
        action.maybe_pushvars.as_ref(),
        cross_repo_push_source,
        PushAuthoredBy::User,
    )
    .await
    {
        Ok(()) => Ok(()),
        Err(BookmarkMovementError::HookFailure(rejections)) => {
            let hook_rejection_remapper = make_hook_rejection_remapper(ctx, repo.clone());
            let rejections =
                map_hook_rejections(rejections, hook_rejection_remapper.as_ref()).await?;
            Err(BundleResolverError::HookError(rejections))
        }
        Err(e) => Err(Error::from(e).into()),
    }
}


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

#![cfg_attr(test, type_length_limit = "2097152")]
#![feature(trait_alias)]

mod changegroup;
mod errors;
mod hook_running;
mod processing;
mod push_redirector;
mod rate_limits;
mod resolver;
mod response;
mod stats;
mod upload_blobs;
mod upload_changesets;

pub use hook_running::run_hooks;
pub use hooks::CrossRepoPushSource;
pub use processing::run_post_resolve_action;
pub use push_redirector::PushRedirector;
pub use push_redirector::PushRedirectorArgs;
pub use resolver::resolve;
pub use resolver::BundleResolverError;
pub use resolver::BundleResolverResultExt;
pub use resolver::Changesets;
pub use resolver::CommonHeads;
pub use resolver::InfiniteBookmarkPush;
pub use resolver::NonFastForwardPolicy;
pub use resolver::PlainBookmarkPush;
pub use resolver::PostResolveAction;
pub use resolver::PostResolveBookmarkOnlyPushRebase;
pub use resolver::PostResolveInfinitePush;
pub use resolver::PostResolvePush;
pub use resolver::PostResolvePushRebase;
pub use resolver::PushrebaseBookmarkSpec;
pub use resolver::UploadedBonsais;
pub use resolver::UploadedHgChangesetIds;
pub use response::UnbundleBookmarkOnlyPushRebaseResponse;
pub use response::UnbundleInfinitePushResponse;
pub use response::UnbundlePushRebaseResponse;
pub use response::UnbundlePushResponse;
pub use response::UnbundleResponse;
pub use upload_changesets::upload_changeset;


/*
 * 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 anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkUpdateReason;
use bookmarks_movement::BookmarkKindRestrictions;
use bookmarks_movement::BookmarkMovementError;
use bookmarks_movement::BookmarkUpdatePolicy;
use bookmarks_movement::BookmarkUpdateTargets;
use bytes::Bytes;
use context::CoreContext;
use hooks::HookManager;
use mercurial_mutation::HgMutationStoreRef;
use metaconfig_types::Address;
use metaconfig_types::PushrebaseRemoteMode;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use pushrebase::PushrebaseError;
#[cfg(fbcode_build)]
use pushrebase_client::LandServicePushrebaseClient;
use pushrebase_client::LocalPushrebaseClient;
use pushrebase_client::PushrebaseClient;
use repo_authorization::AuthorizationContext;
use repo_identity::RepoIdentityRef;
use repo_update_logger::log_new_commits;
use repo_update_logger::CommitInfo;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::debug;
use stats::prelude::*;
use tunables::tunables;

use crate::hook_running::map_hook_rejections;
use crate::hook_running::HookRejectionRemapper;
use crate::rate_limits::enforce_commit_rate_limits;
use crate::response::UnbundleBookmarkOnlyPushRebaseResponse;
use crate::response::UnbundleInfinitePushResponse;
use crate::response::UnbundlePushRebaseResponse;
use crate::response::UnbundlePushResponse;
use crate::response::UnbundleResponse;
use crate::BundleResolverError;
use crate::BundleResolverResultExt;
use crate::CrossRepoPushSource;
use crate::InfiniteBookmarkPush;
use crate::NonFastForwardPolicy;
use crate::PlainBookmarkPush;
use crate::PostResolveAction;
use crate::PostResolveBookmarkOnlyPushRebase;
use crate::PostResolveInfinitePush;
use crate::PostResolvePush;
use crate::PostResolvePushRebase;
use crate::PushrebaseBookmarkSpec;

define_stats! {
    prefix = "mononoke.unbundle.processed";
    push: dynamic_timeseries("{}.push", (reponame: String); Rate, Sum),
    pushrebase: dynamic_timeseries("{}.pushrebase", (reponame: String); Rate, Sum),
    bookmark_only_pushrebase: dynamic_timeseries("{}.bookmark_only_pushrebase", (reponame: String); Rate, Sum),
    infinitepush: dynamic_timeseries("{}.infinitepush", (reponame: String); Rate, Sum),
}

pub trait Repo = bookmarks_movement::Repo + HgMutationStoreRef;

pub async fn run_post_resolve_action(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    action: PostResolveAction,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<UnbundleResponse, BundleResolverError> {
    enforce_commit_rate_limits(ctx, &action).await?;

    // FIXME: it's used not only in pushrebase, so it worth moving
    // populate_git_mapping outside of PushrebaseParams.
    let unbundle_response = match action {
        PostResolveAction::Push(action) => {
            run_push(ctx, repo, hook_manager, action, cross_repo_push_source)
                .await
                .context("While doing a push")
                .map(UnbundleResponse::Push)?
        }
        PostResolveAction::InfinitePush(action) => {
            run_infinitepush(ctx, repo, hook_manager, action, cross_repo_push_source)
                .await
                .context("While doing an infinitepush")
                .map(UnbundleResponse::InfinitePush)?
        }
        PostResolveAction::PushRebase(action) => {
            run_pushrebase(ctx, repo, hook_manager, action, cross_repo_push_source)
                .await
                .map(UnbundleResponse::PushRebase)?
        }
        PostResolveAction::BookmarkOnlyPushRebase(action) => {
            run_bookmark_only_pushrebase(ctx, repo, hook_manager, action, cross_repo_push_source)
                .await
                .context("While doing a bookmark-only pushrebase")
                .map(UnbundleResponse::BookmarkOnlyPushRebase)?
        }
    };
    report_unbundle_type(repo, &unbundle_response);
    Ok(unbundle_response)
}

fn report_unbundle_type(repo: &impl RepoIdentityRef, unbundle_response: &UnbundleResponse) {
    let repo_name = repo.repo_identity().name().to_string();
    match unbundle_response {
        UnbundleResponse::Push(_) => STATS::push.add_value(1, (repo_name,)),
        UnbundleResponse::PushRebase(_) => STATS::pushrebase.add_value(1, (repo_name,)),
        UnbundleResponse::InfinitePush(_) => STATS::infinitepush.add_value(1, (repo_name,)),
        UnbundleResponse::BookmarkOnlyPushRebase(_) => {
            STATS::bookmark_only_pushrebase.add_value(1, (repo_name,))
        }
    }
}

async fn run_push(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    action: PostResolvePush,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<UnbundlePushResponse, BundleResolverError> {
    debug!(ctx.logger(), "unbundle processing: running push.");
    let PostResolvePush {
        changegroup_id,
        mut bookmark_pushes,
        mutations,
        maybe_pushvars,
        non_fast_forward_policy,
        uploaded_bonsais,
        uploaded_hg_changeset_ids,
        hook_rejection_remapper,
    } = action;

    if tunables()
        .mutation_accept_for_infinitepush()
        .unwrap_or_default()
    {
        repo.hg_mutation_store()
            .add_entries(ctx, uploaded_hg_changeset_ids, mutations)
            .await
            .context("Failed to store mutation data")?;
    }

    if bookmark_pushes.len() > 1 {
        return Err(anyhow!(
            "only push to at most one bookmark is allowed, got {:?}",
            bookmark_pushes
        )
        .into());
    }

    let mut changesets_to_log = vec![];
    let mut new_changesets = HashMap::new();
    for bcs in uploaded_bonsais {
        let changeset_id = bcs.get_changeset_id();
        changesets_to_log.push(CommitInfo::new(&bcs, None));
        new_changesets.insert(changeset_id, bcs);
    }

    let mut bookmark_ids = Vec::new();
    let mut maybe_bookmark = None;
    if let Some(bookmark_push) = bookmark_pushes.pop() {
        bookmark_ids.push(bookmark_push.part_id);

        plain_push_bookmark(
            ctx,
            repo,
            hook_manager,
            &bookmark_push,
            new_changesets,
            non_fast_forward_policy,
            BookmarkUpdateReason::Push,
            maybe_pushvars.as_ref(),
            hook_rejection_remapper.as_ref(),
            cross_repo_push_source,
        )
        .await?;

        maybe_bookmark = Some(bookmark_push.name);
    }

    // Since this is a normal push, any bookmark must be public.
    log_new_commits(
        ctx,
        repo,
        maybe_bookmark
            .as_ref()
            .map(|name| (name, BookmarkKind::Publishing)),
        changesets_to_log,
    )
    .await;

    Ok(UnbundlePushResponse {
        changegroup_id,
        bookmark_ids,
    })
}

async fn run_infinitepush(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    action: PostResolveInfinitePush,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<UnbundleInfinitePushResponse, BundleResolverError> {
    debug!(ctx.logger(), "unbundle processing: running infinitepush");
    let PostResolveInfinitePush {
        changegroup_id,
        maybe_bookmark_push,
        mutations,
        uploaded_bonsais,
        uploaded_hg_changeset_ids,
    } = action;

    if tunables()
        .mutation_accept_for_infinitepush()
        .unwrap_or_default()
    {
        repo.hg_mutation_store()
            .add_entries(ctx, uploaded_hg_changeset_ids, mutations)
            .await
            .context("Failed to store mutation data")?;
    }

    let bookmark = match maybe_bookmark_push {
        Some(bookmark_push) => {
            infinitepush_scratch_bookmark(
                ctx,
                repo,
                hook_manager,
                &bookmark_push,
                cross_repo_push_source,
            )
            .await?;

            Some(bookmark_push.name)
        }
        None => None,
    };

    let changesets_to_log = uploaded_bonsais
        .iter()
        .map(|bcs| CommitInfo::new(bcs, None))
        .collect();
    // Since this is an infinitepush, any bookmark must be a scratch bookmark.
    log_new_commits(
        ctx,
        repo,
        bookmark.as_ref().map(|name| (name, BookmarkKind::Scratch)),
        changesets_to_log,
    )
    .await;

    Ok(UnbundleInfinitePushResponse { changegroup_id })
}

async fn run_pushrebase(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    action: PostResolvePushRebase,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<UnbundlePushRebaseResponse, BundleResolverError> {
    debug!(ctx.logger(), "unbundle processing: running pushrebase.");
    let PostResolvePushRebase {
        bookmark_push_part_id,
        bookmark_spec,
        maybe_pushvars,
        commonheads,
        uploaded_bonsais,
        hook_rejection_remapper,
    } = action;

    let (bookmark, pushrebased_rev, pushrebased_changesets) = match bookmark_spec {
        // There's no `.context()` after `normal_pushrebase`, as it has
        // `Error=BundleResolverError` and doing `.context("bla").from_err()`
        // would turn some useful variant of `BundleResolverError` into generic
        // `BundleResolverError::Error`, which in turn would render incorrectly
        // (see definition of `BundleResolverError`).
        PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark) => {
            let mut changesets_to_log: HashMap<_, _> = uploaded_bonsais
                .iter()
                .map(|bcs| (bcs.get_changeset_id(), CommitInfo::new(bcs, None)))
                .collect();

            let (pushrebased_rev, pushrebased_changesets) = normal_pushrebase(
                ctx,
                repo,
                uploaded_bonsais,
                &onto_bookmark,
                maybe_pushvars.as_ref(),
                hook_manager,
                hook_rejection_remapper.as_ref(),
                cross_repo_push_source,
            )
            .await?;
            // Modify the changeset logs with the newly pushrebased hashes.
            for pair in pushrebased_changesets.iter() {
                let info = changesets_to_log
                    .get_mut(&pair.id_old)
                    .ok_or_else(|| anyhow!("Missing commit info for {}", pair.id_old))?;
                info.update_changeset_id(pair.id_old, pair.id_new)?;
            }
            // Wireprotocol pushrebase is always for public bookmarks
            log_new_commits(
                ctx,
                repo,
                Some((&onto_bookmark, BookmarkKind::Publishing)),
                changesets_to_log.into_values().collect(),
            )
            .await;
            (onto_bookmark, pushrebased_rev, pushrebased_changesets)
        }
        PushrebaseBookmarkSpec::ForcePushrebase(plain_push) => {
            let changesets_to_log = uploaded_bonsais
                .iter()
                .map(|bcs| CommitInfo::new(bcs, None))
                .collect();

            let pushrebased_rev = force_pushrebase(
                ctx,
                repo,
                hook_manager,
                uploaded_bonsais,
                &plain_push,
                maybe_pushvars.as_ref(),
                hook_rejection_remapper.as_ref(),
                cross_repo_push_source,
            )
            .await
            .context("While doing a force pushrebase")?;
            // Wireprotocol pushrebase is always for public bookmarks
            log_new_commits(
                ctx,
                repo,
                Some((&plain_push.name, BookmarkKind::Publishing)),
                changesets_to_log,
            )
            .await;
            // Force pushrebase merely force-moves the bookmark, it does not rebase any commits.
            (plain_push.name, pushrebased_rev, Vec::new())
        }
    };

    repo.phases()
        .add_reachable_as_public(ctx, vec![pushrebased_rev.clone()])
        .await
        .context("While marking pushrebased changeset as public")?;

    Ok(UnbundlePushRebaseResponse {
        commonheads,
        pushrebased_rev,
        pushrebased_changesets,
        onto: bookmark,
        bookmark_push_part_id,
    })
}

async fn run_bookmark_only_pushrebase(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    action: PostResolveBookmarkOnlyPushRebase,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<UnbundleBookmarkOnlyPushRebaseResponse, BundleResolverError> {
    debug!(
        ctx.logger(),
        "unbundle processing: running bookmark-only pushrebase."
    );
    let PostResolveBookmarkOnlyPushRebase {
        bookmark_push,
        maybe_pushvars,
        non_fast_forward_policy,
        hook_rejection_remapper,
    } = action;

    let part_id = bookmark_push.part_id;

    // This is a bookmark-only push, so there are no new changesets.
    let new_changesets = HashMap::new();

    plain_push_bookmark(
        ctx,
        repo,
        hook_manager,
        &bookmark_push,
        new_changesets,
        non_fast_forward_policy,
        BookmarkUpdateReason::Pushrebase,
        maybe_pushvars.as_ref(),
        hook_rejection_remapper.as_ref(),
        cross_repo_push_source,
    )
    .await?;

    Ok(UnbundleBookmarkOnlyPushRebaseResponse {
        bookmark_push_part_id: part_id,
    })
}

async fn convert_bookmark_movement_err(
    err: BookmarkMovementError,
    hook_rejection_remapper: &dyn HookRejectionRemapper,
) -> Result<BundleResolverError> {
    Ok(match err {
        BookmarkMovementError::PushrebaseError(PushrebaseError::Conflicts(conflicts)) => {
            BundleResolverError::PushrebaseConflicts(conflicts)
        }
        BookmarkMovementError::HookFailure(rejections) => {
            let rejections = map_hook_rejections(rejections, hook_rejection_remapper).await?;
            BundleResolverError::HookError(rejections)
        }
        _ => BundleResolverError::Error(err.into()),
    })
}

pub async fn maybe_client_from_address<'a>(
    remote_mode: &'a PushrebaseRemoteMode,
    ctx: &'a CoreContext,
    repo: &'a impl Repo,
) -> Option<Box<dyn PushrebaseClient + 'a>> {
    match remote_mode {
        PushrebaseRemoteMode::RemoteLandService(address)
        | PushrebaseRemoteMode::RemoteLandServiceWithLocalFallback(address) => {
            address_from_land_service(address, ctx, repo).await
        }
        PushrebaseRemoteMode::Local => None,
    }
}

async fn address_from_land_service<'a>(
    address: &'a Address,
    ctx: &'a CoreContext,
    repo: &'a impl Repo,
) -> Option<Box<dyn PushrebaseClient + 'a>> {
    #[cfg(fbcode_build)]
    {
        match address {
            metaconfig_types::Address::Tier(tier) => Some(Box::new(
                LandServicePushrebaseClient::from_tier(ctx, tier.clone(), repo)
                    .await
                    .ok()?,
            )),
            metaconfig_types::Address::HostPort(host_port) => Some(Box::new(
                LandServicePushrebaseClient::from_host_port(ctx, host_port.clone(), repo)
                    .await
                    .ok()?,
            )),
        }
    }
    #[cfg(not(fbcode_build))]
    {
        let _ = (address, ctx, repo);
        unreachable!()
    }
}

async fn normal_pushrebase<'a>(
    ctx: &'a CoreContext,
    repo: &'a impl Repo,
    changesets: HashSet<BonsaiChangeset>,
    bookmark: &'a BookmarkKey,
    maybe_pushvars: Option<&'a HashMap<String, Bytes>>,
    hook_manager: &'a HookManager,
    hook_rejection_remapper: &'a dyn HookRejectionRemapper,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<(ChangesetId, Vec<pushrebase::PushrebaseChangesetPair>), BundleResolverError> {
    let bookmark_restriction = BookmarkKindRestrictions::OnlyPublishing;
    let remote_mode = if tunables().force_local_pushrebase().unwrap_or_default() {
        PushrebaseRemoteMode::Local
    } else {
        repo.repo_config().pushrebase.remote_mode.clone()
    };
    let maybe_fallback_scuba: Option<(MononokeScubaSampleBuilder, BookmarkMovementError)> = {
        let maybe_client: Option<Box<dyn PushrebaseClient>> =
            maybe_client_from_address(&remote_mode, ctx, repo).await;

        if let Some(client) = maybe_client {
            let result = client
                .pushrebase(
                    bookmark,
                    changesets.clone(),
                    maybe_pushvars,
                    cross_repo_push_source,
                    bookmark_restriction,
                    false, // We will log new commits locally
                )
                .await;
            match (result, &remote_mode) {
                (Ok(outcome), _) => {
                    return Ok((outcome.head, outcome.rebased_changesets));
                }
                // No fallback, propagate error
                (Err(err), metaconfig_types::PushrebaseRemoteMode::RemoteLandService(..)) => {
                    return Err(convert_bookmark_movement_err(err, hook_rejection_remapper).await?);
                }
                (Err(err), _) => {
                    slog::warn!(
                        ctx.logger(),
                        "Failed to pushrebase remotely, falling back to local. Error: {}",
                        err
                    );
                    let mut scuba = ctx.scuba().clone();
                    scuba.add("bookmark_name", bookmark.as_str());
                    scuba.add(
                        "changeset_id",
                        changesets
                            .iter()
                            .next()
                            .map(|b| b.get_changeset_id().to_string()),
                    );
                    Some((scuba, err))
                }
            }
        } else {
            None
        }
    };
    let authz = AuthorizationContext::new(ctx);
    let result = LocalPushrebaseClient {
        ctx,
        authz: &authz,
        repo,
        hook_manager,
    }
    .pushrebase(
        bookmark,
        changesets,
        maybe_pushvars,
        cross_repo_push_source,
        bookmark_restriction,
        false, // We log commits to scribe ourselves
    )
    .await;
    if let Some((mut scuba, err)) = maybe_fallback_scuba {
        if result.is_ok() {
            scuba.log_with_msg("failed_remote_pushrebase", err.to_string());
        }
    }

    match result {
        Ok(outcome) => Ok((outcome.head, outcome.rebased_changesets)),
        Err(err) => Err(convert_bookmark_movement_err(err, hook_rejection_remapper).await?),
    }
}

async fn force_pushrebase(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    uploaded_bonsais: HashSet<BonsaiChangeset>,
    bookmark_push: &PlainBookmarkPush<ChangesetId>,
    maybe_pushvars: Option<&HashMap<String, Bytes>>,
    hook_rejection_remapper: &dyn HookRejectionRemapper,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<ChangesetId, BundleResolverError> {
    let new_target = bookmark_push
        .new
        .ok_or_else(|| anyhow!("new changeset is required for force pushrebase"))?;

    let mut new_changesets = HashMap::new();
    for bcs in uploaded_bonsais {
        let cs_id = bcs.get_changeset_id();
        new_changesets.insert(cs_id, bcs);
    }

    plain_push_bookmark(
        ctx,
        repo,
        hook_manager,
        bookmark_push,
        new_changesets,
        NonFastForwardPolicy::Allowed,
        BookmarkUpdateReason::Pushrebase,
        maybe_pushvars,
        hook_rejection_remapper,
        cross_repo_push_source,
    )
    .await?;

    Ok(new_target)
}

async fn plain_push_bookmark(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    bookmark_push: &PlainBookmarkPush<ChangesetId>,
    new_changesets: HashMap<ChangesetId, BonsaiChangeset>,
    non_fast_forward_policy: NonFastForwardPolicy,
    reason: BookmarkUpdateReason,
    maybe_pushvars: Option<&HashMap<String, Bytes>>,
    hook_rejection_remapper: &dyn HookRejectionRemapper,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<(), BundleResolverError> {
    let authz = AuthorizationContext::new(ctx);
    // Override the tunable if we know for sure writes are not allowed
    let only_log_acl_checks = !matches!(
        authz,
        AuthorizationContext::ReadOnlyIdentity | AuthorizationContext::DraftOnlyIdentity,
    ) && tunables()
        .log_only_wireproto_write_acl()
        .unwrap_or_default();
    match (bookmark_push.old, bookmark_push.new) {
        (None, Some(new_target)) => {
            let res =
                bookmarks_movement::CreateBookmarkOp::new(&bookmark_push.name, new_target, reason)
                    .only_if_public()
                    .with_new_changesets(new_changesets)
                    .with_pushvars(maybe_pushvars)
                    .with_push_source(cross_repo_push_source)
                    .only_log_acl_checks(only_log_acl_checks)
                    .run(ctx, &authz, repo, hook_manager)
                    .await;
            match res {
                Ok(()) => {}
                Err(err) => match err {
                    BookmarkMovementError::HookFailure(rejections) => {
                        let rejections =
                            map_hook_rejections(rejections, hook_rejection_remapper).await?;
                        return Err(BundleResolverError::HookError(rejections));
                    }
                    _ => {
                        return Err(BundleResolverError::Error(
                            Error::from(err).context("Failed to create bookmark"),
                        ));
                    }
                },
            }
        }

        (Some(old_target), Some(new_target)) => {
            let res = bookmarks_movement::UpdateBookmarkOp::new(
                &bookmark_push.name,
                BookmarkUpdateTargets {
                    old: old_target,
                    new: new_target,
                },
                if non_fast_forward_policy == NonFastForwardPolicy::Allowed {
                    BookmarkUpdatePolicy::AnyPermittedByConfig
                } else {
                    BookmarkUpdatePolicy::FastForwardOnly
                },
                reason,
            )
            .only_if_public()
            .with_new_changesets(new_changesets)
            .with_pushvars(maybe_pushvars)
            .with_push_source(cross_repo_push_source)
            .only_log_acl_checks(only_log_acl_checks)
            .run(ctx, &authz, repo, hook_manager)
            .await;
            match res {
                Ok(()) => {}
                Err(err) => match err {
                    BookmarkMovementError::HookFailure(rejections) => {
                        let rejections =
                            map_hook_rejections(rejections, hook_rejection_remapper).await?;
                        return Err(BundleResolverError::HookError(rejections));
                    }
                    _ => {
                        return Err(BundleResolverError::Error(Error::from(err).context(
                            if non_fast_forward_policy == NonFastForwardPolicy::Allowed {
                                "Failed to move bookmark"
                            } else {
                                concat!(
                                    "Failed to fast-forward bookmark (set pushvar ",
                                    "NON_FAST_FORWARD=true for a non-fast-forward move)",
                                )
                            },
                        )));
                    }
                },
            }
        }

        (Some(old_target), None) => {
            bookmarks_movement::DeleteBookmarkOp::new(&bookmark_push.name, old_target, reason)
                .only_if_public()
                .with_pushvars(maybe_pushvars)
                .only_log_acl_checks(only_log_acl_checks)
                .run(ctx, &authz, repo)
                .await
                .context("Failed to delete bookmark")?;
        }

        (None, None) => {}
    }
    Ok(())
}

async fn infinitepush_scratch_bookmark(
    ctx: &CoreContext,
    repo: &impl Repo,
    hook_manager: &HookManager,
    bookmark_push: &InfiniteBookmarkPush<ChangesetId>,
    cross_repo_push_source: CrossRepoPushSource,
) -> Result<()> {
    let authz = AuthorizationContext::new(ctx);
    // Override the tunable if we know for sure writes are not allowed
    let only_log_acl_checks = !matches!(
        authz,
        AuthorizationContext::ReadOnlyIdentity | AuthorizationContext::DraftOnlyIdentity,
    ) && tunables()
        .log_only_wireproto_write_acl()
        .unwrap_or_default();
    if bookmark_push.old.is_none() && bookmark_push.create {
        bookmarks_movement::CreateBookmarkOp::new(
            &bookmark_push.name,
            bookmark_push.new,
            BookmarkUpdateReason::Push,
        )
        .only_if_scratch()
        .with_push_source(cross_repo_push_source)
        .only_log_acl_checks(only_log_acl_checks)
        .run(ctx, &authz, repo, hook_manager)
        .await
        .context("Failed to create scratch bookmark")?;
    } else {
        let old_target = bookmark_push.old.ok_or_else(|| {
            anyhow!(
                "Unknown bookmark: {}. Use --create to create one.",
                bookmark_push.name
            )
        })?;
        bookmarks_movement::UpdateBookmarkOp::new(
            &bookmark_push.name,
            BookmarkUpdateTargets {
                old: old_target,
                new: bookmark_push.new,
            },
            if bookmark_push.force {
                BookmarkUpdatePolicy::AnyPermittedByConfig
            } else {
                BookmarkUpdatePolicy::FastForwardOnly
            },
            BookmarkUpdateReason::Push,
        )
        .only_if_scratch()
        .with_push_source(cross_repo_push_source)
        .only_log_acl_checks(only_log_acl_checks)
        .run(ctx, &authz, repo, hook_manager)
        .await
        .context(if bookmark_push.force {
            "Failed to move scratch bookmark"
        } else {
            "Failed to fast-forward scratch bookmark (try --force?)"
        })?;
    }

    Ok(())
}