/*
 * 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 core::fmt::Debug;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Display;
use std::sync::Arc;

use anyhow::bail;
use anyhow::ensure;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use ascii::AsciiString;
use blobrepo::BlobRepo;
use blobrepo_hg::BlobRepoHg;
use blobrepo_hg::ChangesetHandle;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarkKey;
use bytes::Bytes;
use context::CoreContext;
use context::SessionClass;
use futures::compat::Stream01CompatExt;
use futures::future;
use futures::future::try_join_all;
use futures::stream;
use futures::stream::BoxStream;
use futures::try_join;
use futures::Future;
use futures::StreamExt;
use futures::TryStreamExt;
use hooks::HookRejectionInfo;
use lazy_static::lazy_static;
use mercurial_bundles::Bundle2Item;
use mercurial_bundles::PartHeader;
use mercurial_bundles::PartHeaderInner;
use mercurial_bundles::PartHeaderType;
use mercurial_bundles::PartId;
use mercurial_mutation::HgMutationEntry;
use mercurial_revlog::changeset::RevlogChangeset;
use mercurial_types::HgChangesetId;
use metaconfig_types::PushrebaseFlags;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use rate_limiting::RateLimitBody;
use repo_identity::RepoIdentityRef;
use slog::trace;
use topo_sort::sort_topological;
use tunables::tunables;
use wirepack::TreemanifestBundle2Parser;
use wireproto_handler::BackupSourceRepo;

use crate::changegroup::convert_to_revlog_changesets;
use crate::changegroup::convert_to_revlog_filelog;
use crate::changegroup::split_changegroup;
use crate::errors::*;
use crate::hook_running::make_hook_rejection_remapper;
use crate::hook_running::HookRejectionRemapper;
use crate::rate_limits::enforce_file_changes_rate_limits;
use crate::rate_limits::RateLimitedPushKind;
use crate::stats::*;
use crate::upload_blobs::upload_hg_blobs;
use crate::upload_changesets::upload_changeset;
use crate::upload_changesets::Filelogs;
use crate::upload_changesets::Manifests;

#[allow(non_snake_case)]
mod UNBUNDLE_STATS {
    use stats::define_stats;

    define_stats! {
        prefix = "mononoke.unbundle.resolver";
        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),
        resolver_error: dynamic_timeseries("{}.resolver_error", (reponame: String); Rate, Sum),
        total_unbundles: dynamic_timeseries("{}.total_unbundles", (reponame: String); Rate, Sum),
    }

    pub(crate) use self::STATS::*;
}

pub type Changesets = Vec<(HgChangesetId, RevlogChangeset)>;
pub type UploadedBonsais = HashSet<BonsaiChangeset>;
pub type UploadedHgChangesetIds = HashSet<HgChangesetId>;

// This is to match the core hg behavior from https://fburl.com/jf3iyl7y
// Mercurial substitutes the `onto` parameter with this bookmark name when
// the force pushrebase is done, so we need to look for it and make sure we
// do the right thing here too.
lazy_static! {
    static ref DONOTREBASEBOOKMARK: BookmarkKey =
        BookmarkKey::new("__pushrebase_donotrebase__").unwrap();
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum NonFastForwardPolicy {
    Allowed,
    Disallowed,
}

impl From<bool> for NonFastForwardPolicy {
    fn from(allowed: bool) -> Self {
        if allowed {
            Self::Allowed
        } else {
            Self::Disallowed
        }
    }
}

pub struct HgHookRejection {
    pub(crate) hook_name: String,
    pub(crate) hg_cs_id: HgChangesetId,
    pub(crate) reason: HookRejectionInfo,
}

impl HgHookRejection {
    pub fn get_hook_name(&self) -> &str {
        &self.hook_name
    }
}

pub enum BundleResolverError {
    HookError(Vec<HgHookRejection>),
    PushrebaseConflicts(Vec<pushrebase::PushrebaseConflict>),
    Error(Error),
    RateLimitExceeded {
        limit_name: String,
        limit: RateLimitBody,
        entity: String,
        value: f64,
    },
}

impl From<Error> for BundleResolverError {
    fn from(error: Error) -> Self {
        Self::Error(error)
    }
}

impl From<BundleResolverError> for Error {
    fn from(error: BundleResolverError) -> Error {
        // DO NOT CHANGE FORMATTING WITHOUT UPDATING https://fburl.com/diffusion/bs9fys78 first!!
        use BundleResolverError::*;
        match error {
            HookError(hook_outcomes) => {
                let err_msgs: Vec<_> = hook_outcomes
                    .into_iter()
                    .map(|failure| {
                        format!(
                            "{} for {}: {}",
                            failure.hook_name, failure.hg_cs_id, failure.reason.long_description
                        )
                    })
                    .collect();
                format_err!("hooks failed:\n{}", err_msgs.join("\n"))
            }
            PushrebaseConflicts(conflicts) => {
                format_err!("pushrebase failed Conflicts({:?})", conflicts)
            }
            RateLimitExceeded {
                limit_name,
                limit,
                entity,
                value,
            } => format_err!(
                "Rate limit exceeded: {} for {}. \
                 The maximum allowed value is {} over a sliding {}s interval. \
                 If allowed, the value would be {}.",
                limit_name,
                entity,
                limit.raw_config.limit,
                limit.window.as_secs(),
                value,
            ),
            Error(err) => err,
        }
    }
}

pub trait BundleResolverResultExt<T> {
    fn context<C>(self, context: C) -> Result<T, BundleResolverError>
    where
        C: Display + Send + Sync + 'static;
}

impl<T> BundleResolverResultExt<T> for Result<T, BundleResolverError> {
    fn context<C>(self, context: C) -> Result<T, BundleResolverError>
    where
        C: Display + Send + Sync + 'static,
    {
        match self {
            Ok(v) => Ok(v),
            Err(BundleResolverError::Error(err)) => Err(err.context(context).into()),
            Err(e) => Err(e),
        }
    }
}

/// Data, needed to perform post-resolve `Push` action
pub struct PostResolvePush {
    pub changegroup_id: Option<PartId>,
    pub bookmark_pushes: Vec<PlainBookmarkPush<ChangesetId>>,
    pub mutations: Vec<HgMutationEntry>,
    pub maybe_pushvars: Option<HashMap<String, Bytes>>,
    pub non_fast_forward_policy: NonFastForwardPolicy,
    pub uploaded_bonsais: UploadedBonsais,
    pub uploaded_hg_changeset_ids: UploadedHgChangesetIds,
    pub hook_rejection_remapper: Arc<dyn HookRejectionRemapper>,
}

/// Data, needed to perform post-resolve `InfinitePush` action
pub struct PostResolveInfinitePush {
    pub changegroup_id: Option<PartId>,
    pub maybe_bookmark_push: Option<InfiniteBookmarkPush<ChangesetId>>,
    pub mutations: Vec<HgMutationEntry>,
    pub uploaded_bonsais: UploadedBonsais,
    pub uploaded_hg_changeset_ids: UploadedHgChangesetIds,
}

/// Data, needed to perform post-resolve `PushRebase` action
#[derive(Clone)]
pub struct PostResolvePushRebase {
    pub bookmark_push_part_id: Option<PartId>,
    pub bookmark_spec: PushrebaseBookmarkSpec<ChangesetId>,
    pub maybe_pushvars: Option<HashMap<String, Bytes>>,
    pub commonheads: CommonHeads,
    pub uploaded_bonsais: UploadedBonsais,
    pub hook_rejection_remapper: Arc<dyn HookRejectionRemapper>,
}

/// Data, needed to perform post-resolve `BookmarkOnlyPushRebase` action
pub struct PostResolveBookmarkOnlyPushRebase {
    pub bookmark_push: PlainBookmarkPush<ChangesetId>,
    pub maybe_pushvars: Option<HashMap<String, Bytes>>,
    pub non_fast_forward_policy: NonFastForwardPolicy,
    pub hook_rejection_remapper: Arc<dyn HookRejectionRemapper>,
}

/// An action to take after the `unbundle` bundle2 was completely resolved
/// "Completely resolved" here means:
/// - parsed
/// - all received changesets and blobs uploaded
pub enum PostResolveAction {
    Push(PostResolvePush),
    InfinitePush(PostResolveInfinitePush),
    PushRebase(PostResolvePushRebase),
    BookmarkOnlyPushRebase(PostResolveBookmarkOnlyPushRebase),
}

/// The resolve function takes a bundle2, interprets it's content as Changesets, Filelogs and
/// Manifests and uploades all of them to the provided BlobRepo in the correct order.
/// It returns a Future that contains the response that should be send back to the requester.
pub async fn resolve<'a>(
    ctx: &'a CoreContext,
    repo: &'a BlobRepo,
    infinitepush_writes_allowed: bool,
    bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    pure_push_allowed: bool,
    pushrebase_flags: PushrebaseFlags,
    maybe_backup_repo_source: Option<BackupSourceRepo>,
) -> Result<PostResolveAction, BundleResolverError> {
    let result = resolve_impl(
        ctx,
        repo,
        infinitepush_writes_allowed,
        bundle2,
        pure_push_allowed,
        pushrebase_flags,
        maybe_backup_repo_source,
    )
    .await;
    report_unbundle_type(ctx, repo, &result);
    result
}

async fn resolve_impl<'a>(
    ctx: &'a CoreContext,
    repo: &'a BlobRepo,
    infinitepush_writes_allowed: bool,
    bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    pure_push_allowed: bool,
    pushrebase_flags: PushrebaseFlags,
    maybe_backup_repo_source: Option<BackupSourceRepo>,
) -> Result<PostResolveAction, BundleResolverError> {
    let resolver = Bundle2Resolver::new(ctx, repo, infinitepush_writes_allowed, pushrebase_flags);
    let bundle2 = resolver.resolve_stream_params(bundle2).await?;
    let bundle2 = resolver.resolve_replycaps(bundle2).await?;

    let (maybe_commonheads, bundle2) = resolver.maybe_resolve_commonheads(bundle2).await?;
    let (maybe_pushvars, bundle2) = resolver
        .maybe_resolve_pushvars(bundle2)
        .await
        .context("While resolving Pushvars")?;

    let (pushkey_next, bundle2) = resolver.is_next_part_pushkey(bundle2).await?;

    let non_fast_forward_policy = {
        let mut allow_non_fast_forward = false;
        // check the bypass condition
        if let Some(ref pushvars) = maybe_pushvars {
            allow_non_fast_forward = pushvars
                .get("NON_FAST_FORWARD")
                .map(|s| s.to_ascii_lowercase())
                == Some("true".into());
        }
        NonFastForwardPolicy::from(allow_non_fast_forward)
    };

    let post_resolve_action = if let Some(commonheads) = maybe_commonheads {
        if pushkey_next {
            resolve_bookmark_only_pushrebase(
                ctx,
                resolver,
                bundle2,
                maybe_pushvars,
                non_fast_forward_policy,
            )
            .await
            .map_err(BundleResolverError::from)
        } else {
            fn changegroup_always_unacceptable() -> bool {
                false
            }
            resolve_pushrebase(
                ctx,
                commonheads,
                resolver,
                bundle2,
                maybe_pushvars,
                changegroup_always_unacceptable,
                maybe_backup_repo_source,
            )
            .await
        }
    } else {
        resolve_push(
            ctx,
            resolver,
            bundle2,
            maybe_pushvars,
            non_fast_forward_policy,
            move || pure_push_allowed,
            maybe_backup_repo_source,
        )
        .await
        .context("bundle2_resolver error")
        .map_err(BundleResolverError::from)
    };

    match post_resolve_action {
        Err(e) => Err(e),
        Ok(val) => Ok(val),
    }
}

fn report_unbundle_type(
    ctx: &CoreContext,
    repo: &BlobRepo,
    post_resolve_action: &Result<PostResolveAction, BundleResolverError>,
) {
    let repo_name = repo.repo_identity().name().to_string();
    UNBUNDLE_STATS::total_unbundles.add_value(1, (repo_name.clone(),));
    let unbundle_resolved = "Unbundle resolved";
    match post_resolve_action {
        Ok(PostResolveAction::Push(_)) => {
            ctx.scuba()
                .clone()
                .log_with_msg(unbundle_resolved, Some("push".to_owned()));
            UNBUNDLE_STATS::push.add_value(1, (repo_name,))
        }
        Ok(PostResolveAction::PushRebase(_)) => {
            ctx.scuba()
                .clone()
                .log_with_msg(unbundle_resolved, Some("pushrebase".to_owned()));
            UNBUNDLE_STATS::pushrebase.add_value(1, (repo_name,))
        }
        Ok(PostResolveAction::InfinitePush(_)) => {
            ctx.scuba()
                .clone()
                .log_with_msg(unbundle_resolved, Some("infinitepush".to_owned()));
            UNBUNDLE_STATS::infinitepush.add_value(1, (repo_name,))
        }
        Ok(PostResolveAction::BookmarkOnlyPushRebase(_)) => {
            ctx.scuba().clone().log_with_msg(
                unbundle_resolved,
                Some("bookmark_only_pushrebase".to_owned()),
            );
            UNBUNDLE_STATS::bookmark_only_pushrebase.add_value(1, (repo_name,))
        }
        Err(_) => UNBUNDLE_STATS::resolver_error.add_value(1, (repo_name,)),
    }
}

async fn resolve_push<'r>(
    ctx: &'r CoreContext,
    resolver: Bundle2Resolver<'r>,
    bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    maybe_pushvars: Option<HashMap<String, Bytes>>,
    non_fast_forward_policy: NonFastForwardPolicy,
    changegroup_acceptable: impl FnOnce() -> bool + Send + Sync + 'static,
    maybe_backup_repo_source: Option<BackupSourceRepo>,
) -> Result<PostResolveAction, Error> {
    let (cg_push, bundle2) = resolver
        .maybe_resolve_changegroup(bundle2, changegroup_acceptable)
        .await
        .context("While resolving Changegroup")?;
    let (pushkeys, bundle2) = resolver
        .resolve_multiple_parts(bundle2, Bundle2Resolver::maybe_resolve_pushkey)
        .await
        .context("While resolving Pushkey")?;
    let is_infinitepush = cg_push
        .as_ref()
        .and_then(|cg_push| cg_push.infinitepush_payload.as_ref())
        .is_some();
    let infinitepush_bp = cg_push
        .as_ref()
        .and_then(|cg_push| cg_push.infinitepush_payload.as_ref())
        .and_then(|ip_payload| ip_payload.bookmark_push.as_ref());

    let infinitepush_bp = infinitepush_bp.cloned();

    let (mutations, bundle2) = resolver
        .maybe_resolve_infinitepush_mutation(bundle2)
        .await
        .context("While resolving InfinitepushMutation")?;

    let (cg_and_manifests, bundle2) = if let Some(cg_push) = cg_push {
        let (manifests, bundle2) = resolver
            .resolve_b2xtreegroup2(bundle2)
            .await
            .context("While resolving B2xTreegroup2")?;
        (Some((cg_push, manifests)), bundle2)
    } else {
        (None, bundle2)
    };
    // At the moment pushkey part may appear in two places: after Changegroup
    // and after Treegroup.
    let (pushkeys, bundle2) = if pushkeys.is_empty() {
        resolver
            .resolve_multiple_parts(bundle2, Bundle2Resolver::maybe_resolve_pushkey)
            .await
            .context("While resolving Pushkey")?
    } else {
        (pushkeys, bundle2)
    };
    let maybe_hg_bookmark_push = try_collect_all_bookmark_pushes(pushkeys, infinitepush_bp)?;

    let (changegroup_id, uploaded_bonsais, uploaded_hg_changeset_ids) =
        if let Some((cg_push, manifests)) = cg_and_manifests {
            let changegroup_id = Some(cg_push.part_id);
            let (uploaded_bonsais, uploaded_hg_changeset_ids) = resolver
                .upload_changesets(cg_push, manifests, maybe_backup_repo_source)
                .await?;
            // Note: we do not run hooks on pure pushes. This probably has to be changed later.
            (changegroup_id, uploaded_bonsais, uploaded_hg_changeset_ids)
        } else {
            (None, UploadedBonsais::new(), UploadedHgChangesetIds::new())
        };

    let ((), bundle2) = resolver
        .maybe_resolve_infinitepush_bookmarks(bundle2)
        .await
        .context("While resolving B2xInfinitepushBookmarks")?;
    resolver.ensure_stream_finished(bundle2).await?;

    let maybe_bonsai_bookmark_push = match maybe_hg_bookmark_push {
        Some(hg_bookmark_push) => {
            Some(hg_all_bookmark_pushes_to_bonsai(ctx, resolver.repo, hg_bookmark_push).await?)
        }
        None => None,
    };

    if is_infinitepush {
        get_post_resolve_infinitepush(
            changegroup_id,
            maybe_bonsai_bookmark_push,
            mutations,
            uploaded_bonsais,
            uploaded_hg_changeset_ids,
        )
        .map(PostResolveAction::InfinitePush)
    } else {
        let hook_rejection_remapper =
            make_hook_rejection_remapper(ctx, resolver.repo.clone()).into();

        get_post_resolve_push(
            changegroup_id,
            maybe_bonsai_bookmark_push,
            mutations,
            maybe_pushvars,
            non_fast_forward_policy,
            uploaded_bonsais,
            uploaded_hg_changeset_ids,
            hook_rejection_remapper,
        )
        .map(PostResolveAction::Push)
    }
}

fn get_post_resolve_infinitepush(
    changegroup_id: Option<PartId>,
    maybe_bonsai_bookmark_push: Option<AllBookmarkPushes<ChangesetId>>,
    mutations: Vec<HgMutationEntry>,
    uploaded_bonsais: UploadedBonsais,
    uploaded_hg_changeset_ids: UploadedHgChangesetIds,
) -> Result<PostResolveInfinitePush, Error> {
    let maybe_bookmark_push = match maybe_bonsai_bookmark_push {
        Some(AllBookmarkPushes::PlainPushes(_)) => {
            return Err(format_err!(
                "Infinitepush push cannot contain regular bookmarks"
            ));
        }
        Some(AllBookmarkPushes::Inifinitepush(bookmark_push)) => Some(bookmark_push),
        None => None,
    };

    Ok(PostResolveInfinitePush {
        changegroup_id,
        maybe_bookmark_push,
        mutations,
        uploaded_bonsais,
        uploaded_hg_changeset_ids,
    })
}

fn get_post_resolve_push(
    changegroup_id: Option<PartId>,
    maybe_bonsai_bookmark_push: Option<AllBookmarkPushes<ChangesetId>>,
    mutations: Vec<HgMutationEntry>,
    maybe_pushvars: Option<HashMap<String, Bytes>>,
    non_fast_forward_policy: NonFastForwardPolicy,
    uploaded_bonsais: UploadedBonsais,
    uploaded_hg_changeset_ids: UploadedHgChangesetIds,
    hook_rejection_remapper: Arc<dyn HookRejectionRemapper>,
) -> Result<PostResolvePush, Error> {
    let bookmark_pushes = match maybe_bonsai_bookmark_push {
        Some(AllBookmarkPushes::Inifinitepush(_bookmark_push)) => {
            return Err(format_err!(
                "This should actually be impossible: non-infinitepush push with infinitepush bookmarks"
            ));
        }
        Some(AllBookmarkPushes::PlainPushes(bookmark_pushes)) => bookmark_pushes,
        None => vec![],
    };

    Ok(PostResolvePush {
        changegroup_id,
        bookmark_pushes,
        mutations,
        maybe_pushvars,
        non_fast_forward_policy,
        uploaded_bonsais,
        uploaded_hg_changeset_ids,
        hook_rejection_remapper,
    })
}

// Enum used to pass data for normal or forceful pushrebases
// Normal pushrebase is what one would expect: take a (potentially
// stack of) commit(s) and rebase it on top of a given bookmark.
// Force pushrebase is basically a push, which for logging
// and respondin purposes is treated like a pushrebase
#[derive(Clone)]
pub enum PushrebaseBookmarkSpec<T: Copy> {
    NormalPushrebase(BookmarkKey),
    ForcePushrebase(PlainBookmarkPush<T>),
}

impl<T: Copy> PushrebaseBookmarkSpec<T> {
    pub fn get_bookmark_name(&self) -> &BookmarkKey {
        match self {
            PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark) => onto_bookmark,
            PushrebaseBookmarkSpec::ForcePushrebase(plain_push) => &plain_push.name,
        }
    }
}

async fn resolve_pushrebase<'r>(
    ctx: &'r CoreContext,
    commonheads: CommonHeads,
    resolver: Bundle2Resolver<'r>,
    bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    maybe_pushvars: Option<HashMap<String, Bytes>>,
    changegroup_acceptable: impl FnOnce() -> bool + Send + Sync + 'static,
    maybe_backup_repo_source: Option<BackupSourceRepo>,
) -> Result<PostResolveAction, BundleResolverError> {
    let (manifests, bundle2) = resolver
        .resolve_b2xtreegroup2(bundle2)
        .await
        .context("While resolving B2xTreegroup2")?;
    let (maybe_cg_push, bundle2) = resolver
        .maybe_resolve_changegroup(bundle2, changegroup_acceptable)
        .await
        .context("While resolving Changegroup")?;
    let cg_push = maybe_cg_push.ok_or_else(|| Error::msg("Empty pushrebase"))?;
    let onto_bookmark = match cg_push.mparams.get("onto") {
        Some(onto_bookmark) => {
            let v = Vec::from(onto_bookmark.as_ref());
            let onto_bookmark = String::from_utf8(v).map_err(Error::from)?;

            BookmarkKey::new(onto_bookmark)?
        }
        None => return Err(format_err!("onto is not specified").into()),
    };

    let will_rebase = onto_bookmark != *DONOTREBASEBOOKMARK;
    // Mutation information must not be present in public commits
    // See T54101162, S186586
    if !will_rebase {
        for (_, hg_cs) in &cg_push.changesets {
            for key in pushrebase::MUTATION_KEYS {
                if hg_cs.extra.as_ref().contains_key(key.as_bytes()) {
                    return Err(Error::msg("Forced push blocked because it contains mutation metadata.\n\
                                You can remove the metadata from a commit with `hg amend --config mutation.record=false`.\n\
                                For more help, please contact the Source Control team at https://fburl.com/27qnuyl2").into());
                }
            }
        }
    }

    let (uploaded_bonsais, _uploaded_hg_changeset_ids) = resolver
        .upload_changesets(cg_push, manifests, maybe_backup_repo_source)
        .await?;

    let (pushkeys, bundle2) = resolver
        .resolve_multiple_parts(bundle2, Bundle2Resolver::maybe_resolve_pushkey)
        .await
        .context("While resolving Pushkey")?;

    let bookmark_pushes = collect_pushkey_bookmark_pushes(pushkeys);
    if bookmark_pushes.len() > 1 {
        return Err(format_err!("Too many pushkey parts: {:?}", bookmark_pushes).into());
    }

    let (bookmark_push_part_id, bookmark_spec) = match bookmark_pushes.into_iter().next() {
        Some(bk_push) if bk_push.name != onto_bookmark && onto_bookmark != *DONOTREBASEBOOKMARK => {
            return Err(format_err!(
                "allowed only pushes of {} bookmark: {:?}",
                onto_bookmark,
                bk_push
            )
            .into());
        }
        Some(bk_push) if onto_bookmark == *DONOTREBASEBOOKMARK => {
            (
                // This is a force pushrebase scenario. We need to ignore `onto_bookmark`
                // and run normal push (using bk_push), but generate a pushrebase
                // response.
                // See comment next to DONOTREBASEBOOKMARK definition
                Some(bk_push.part_id),
                PushrebaseBookmarkSpec::ForcePushrebase(bk_push),
            )
        }
        Some(bk_push) => (
            Some(bk_push.part_id),
            PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark),
        ),
        None => (
            None,
            PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark),
        ),
    };

    resolver.ensure_stream_finished(bundle2).await?;
    let bookmark_spec =
        hg_pushrebase_bookmark_spec_to_bonsai(ctx, resolver.repo, bookmark_spec).await?;

    let hook_rejection_remapper = make_hook_rejection_remapper(ctx, resolver.repo.clone()).into();

    Ok(PostResolveAction::PushRebase(PostResolvePushRebase {
        bookmark_push_part_id,
        bookmark_spec,
        maybe_pushvars,
        commonheads,
        uploaded_bonsais,
        hook_rejection_remapper,
    }))
}

/// Do the right thing when pushrebase-enabled client only wants to manipulate bookmarks
async fn resolve_bookmark_only_pushrebase<'r>(
    ctx: &'r CoreContext,
    resolver: Bundle2Resolver<'r>,
    bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    maybe_pushvars: Option<HashMap<String, Bytes>>,
    non_fast_forward_policy: NonFastForwardPolicy,
) -> Result<PostResolveAction, Error> {
    // TODO: we probably run hooks even if no changesets are pushed?
    //       however, current run_hooks implementation will no-op such thing

    let (pushkeys, bundle2) = resolver
        .resolve_multiple_parts(bundle2, Bundle2Resolver::maybe_resolve_pushkey)
        .await
        .context("While resolving Pushkey")?;
    let pushkeys_len = pushkeys.len();
    let bookmark_pushes = collect_pushkey_bookmark_pushes(pushkeys);

    // this means we filtered some Phase pushkeys out
    // which is not expected
    if bookmark_pushes.len() != pushkeys_len {
        return Err(Error::msg(
            "Expected bookmark-only push, phases pushkey found",
        ));
    }

    if bookmark_pushes.len() != 1 {
        return Err(format_err!("Too many pushkey parts: {:?}", bookmark_pushes));
    }

    let bookmark_push = bookmark_pushes.into_iter().next().unwrap();
    resolver.ensure_stream_finished(bundle2).await?;
    let bookmark_push = plain_hg_bookmark_push_to_bonsai(ctx, resolver.repo, bookmark_push).await?;
    let hook_rejection_remapper = make_hook_rejection_remapper(ctx, resolver.repo.clone()).into();

    Ok(PostResolveAction::BookmarkOnlyPushRebase(
        PostResolveBookmarkOnlyPushRebase {
            bookmark_push,
            maybe_pushvars,
            non_fast_forward_policy,
            hook_rejection_remapper,
        },
    ))
}

/// Represents all the bookmark pushes that are created
/// by a single unbundle wireproto command. This can
/// be either an exactly one infinitepush, or multiple
/// plain pushes
enum AllBookmarkPushes<T: Copy> {
    PlainPushes(Vec<PlainBookmarkPush<T>>),
    Inifinitepush(InfiniteBookmarkPush<T>),
}

/// Represets a single non-infinitepush bookmark push
/// This can be a result of a normal push or a pushrebase
#[derive(Debug, Clone)]
pub struct PlainBookmarkPush<T: Copy> {
    pub part_id: PartId,
    pub name: BookmarkKey,
    pub old: Option<T>,
    pub new: Option<T>,
}

/// Represents an infinitepush bookmark push
#[derive(Debug, Clone)]
pub struct InfiniteBookmarkPush<T> {
    pub name: BookmarkKey,
    pub create: bool,
    pub force: bool,
    pub old: Option<T>,
    pub new: T,
}

#[derive(Debug, Clone)]
struct InfinitepushPayload {
    /// An Infinitepush bookmark (aka scratch bookmark) that was provided through an Infinitepush
    /// bundle part.
    bookmark_push: Option<InfiniteBookmarkPush<HgChangesetId>>,
}

struct ChangegroupPush {
    part_id: PartId,
    changesets: Changesets,
    filelogs: Filelogs,
    mparams: HashMap<String, Bytes>,
    /// Infinitepush data provided through the Changegroup. If the push was an Infinitepush, this
    /// will be present.
    infinitepush_payload: Option<InfinitepushPayload>,
}

#[derive(Clone)]
pub struct CommonHeads {
    pub heads: Vec<HgChangesetId>,
}

enum Pushkey {
    HgBookmarkPush(PlainBookmarkPush<HgChangesetId>),
    Phases,
}

/// Holds repo and logger for convienience access from it's methods
#[derive(Clone)]
struct Bundle2Resolver<'r> {
    ctx: &'r CoreContext,
    repo: &'r BlobRepo,
    infinitepush_writes_allowed: bool,
    #[allow(dead_code)]
    pushrebase_flags: PushrebaseFlags,
}

impl<'r> Bundle2Resolver<'r> {
    fn new(
        ctx: &'r CoreContext,
        repo: &'r BlobRepo,
        infinitepush_writes_allowed: bool,
        pushrebase_flags: PushrebaseFlags,
    ) -> Self {
        Self {
            ctx,
            repo,
            infinitepush_writes_allowed,
            pushrebase_flags,
        }
    }

    /// Peek at the next `bundle2` item and check if it is a `Pushkey` part
    /// Return unchanged `bundle2`
    async fn is_next_part_pushkey(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<(bool, BoxStream<'static, Result<Bundle2Item<'static>>>)> {
        match bundle2.try_next().await? {
            Some(part) => {
                if let Bundle2Item::Pushkey(header, box_future) = part {
                    Ok((
                        true,
                        stream::once(async { Ok(Bundle2Item::Pushkey(header, box_future)) })
                            .chain(bundle2)
                            .boxed(),
                    ))
                } else {
                    return_with_rest_of_bundle(false, part, bundle2).await
                }
            }
            _ => Ok((false, bundle2)),
        }
    }

    /// Parse the stream header and extract stream params from it
    /// Return the rest of the stream along, excluding the params
    async fn resolve_stream_params(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<BoxStream<'static, Result<Bundle2Item<'static>>>> {
        match bundle2.try_next().await? {
            Some(Bundle2Item::Start(_)) => Ok(bundle2),
            _ => Err(format_err!("Expected Bundle2 Start")),
        }
    }

    /// Parse replycaps and ignore its content
    /// Return the rest of the bundle
    async fn resolve_replycaps(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<BoxStream<'static, Result<Bundle2Item<'static>>>> {
        match bundle2.try_next().await? {
            Some(Bundle2Item::Replycaps(_, part)) => {
                part.await?;
                Ok(bundle2)
            }
            _ => Err(format_err!("Expected Bundle2 Replycaps")),
        }
    }

    // Parse b2x:commonheads
    // This part sent by pushrebase so that server can find out what commits to send back to the
    // client. This part is used as a marker that this push is pushrebase.
    async fn maybe_resolve_commonheads(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<
        (
            Option<CommonHeads>,
            BoxStream<'static, Result<Bundle2Item<'static>>>,
        ),
        Error,
    > {
        match bundle2.try_next().await? {
            Some(Bundle2Item::B2xCommonHeads(_header, heads)) => Ok((
                Some(CommonHeads {
                    heads: heads.try_collect().await?,
                }),
                bundle2,
            )),

            Some(part) => return_with_rest_of_bundle(None, part, bundle2).await,
            _ => Err(format_err!("Unexpected Bundle2 stream end")),
        }
    }

    /// Parse pushvars
    /// It is used to store hook arguments.
    async fn maybe_resolve_pushvars(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<
        (
            Option<HashMap<String, Bytes>>,
            BoxStream<'static, Result<Bundle2Item<'static>>>,
        ),
        Error,
    > {
        let maybe_pushvars = match bundle2.try_next().await? {
            Some(Bundle2Item::Pushvars(header, emptypart)) => {
                let pushvars = header.into_inner().aparams;
                // ignored for now, will be used for hooks
                emptypart.await?;
                Some(pushvars)
            }
            Some(part) => return return_with_rest_of_bundle(None, part, bundle2).await,
            None => None,
        };

        Ok((maybe_pushvars, bundle2))
    }

    /// Parse changegroup.
    /// The ChangegroupId will be used in the last step for preparing response
    /// The Changesets should be parsed as RevlogChangesets and used for uploading changesets
    /// The Filelogs should be scheduled for uploading to BlobRepo and the Future resolving in
    /// their upload should be used for uploading changesets
    /// `pure_push_allowed` argument is responsible for allowing
    /// pure (non-pushrebase and non-infinitepush) pushes
    async fn maybe_resolve_changegroup(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
        changegroup_acceptable: impl FnOnce() -> bool + Send + Sync + 'static,
    ) -> Result<
        (
            Option<ChangegroupPush>,
            BoxStream<'static, Result<Bundle2Item<'static>>>,
        ),
        Error,
    > {
        let infinitepush_writes_allowed = self.infinitepush_writes_allowed;

        let maybe_cg_push: Option<ChangegroupPush> = match bundle2.try_next().await? {
            // XXX: we may be interested in checking that this is a correct changegroup part
            // type
            Some(Bundle2Item::Changegroup(header, parts))
            | Some(Bundle2Item::B2xInfinitepush(header, parts))
            | Some(Bundle2Item::B2xRebase(header, parts)) => {
                if header.part_type() == &PartHeaderType::Changegroup && !changegroup_acceptable() {
                    // Changegroup part type signals that we are in a pure push scenario
                    return Err(format_err!("Pure pushes are disallowed in this repo"));
                }

                let is_infinitepush = header.part_type() == &PartHeaderType::B2xInfinitepush;

                if is_infinitepush && !infinitepush_writes_allowed {
                    bail!(
                        "Infinitepush is not enabled on this server. Contact Source Control @ FB."
                    );
                }

                let push_kind = if is_infinitepush {
                    RateLimitedPushKind::InfinitePush
                } else {
                    RateLimitedPushKind::Public
                };

                let (changesets, filelogs) = split_changegroup(parts);
                let changesets: Vec<(HgChangesetId, RevlogChangeset)> =
                    convert_to_revlog_changesets(changesets)
                        .try_collect()
                        .await?;

                let commit_limit = tunables()
                    .unbundle_limit_num_of_commits_in_push()
                    .unwrap_or_default();
                // Ignore commit limit if hg sync job is pushing. Hg sync job is used
                // to mirror one repository into another, and we can't discard a push
                // even if it's too big
                if commit_limit > 0 && !self.ctx.session().is_hg_sync_job() {
                    let commit_limit: usize = commit_limit.try_into().unwrap();
                    if changesets.len() > commit_limit {
                        bail!(
                            "Trying to push too many commits! Limit is {}, tried to push {}",
                            commit_limit,
                            changesets.len()
                        );
                    }
                }

                let changesets = if is_infinitepush
                    && tunables()
                        .filter_pre_existing_commits_on_infinitepush()
                        .unwrap_or_default()
                {
                    let hg_cs_ids = changesets.iter().map(|(id, _)| *id).collect::<Vec<_>>();

                    let mapping = self
                        .repo
                        .get_hg_bonsai_mapping(self.ctx.clone(), hg_cs_ids)
                        .await
                        .with_context(|| "Failed to query for pre-existing changesets")?;

                    let existing = mapping
                        .into_iter()
                        .map(|(hg_cs_id, _)| hg_cs_id)
                        .collect::<HashSet<_>>();

                    let orig_count = changesets.len();

                    let new_changesets = changesets
                        .into_iter()
                        .filter(|(hg_cs_id, _)| !existing.contains(hg_cs_id))
                        .collect::<Vec<_>>();

                    if new_changesets.len() < orig_count {
                        self.ctx
                            .scuba()
                            .clone()
                            .add("original_changeset_count", orig_count)
                            .add("new_changeset_count", new_changesets.len())
                            .log_with_msg("Filtered out pre-existing changesets", None);
                    }

                    new_changesets
                } else {
                    changesets
                };

                enforce_file_changes_rate_limits(
                    self.ctx,
                    push_kind,
                    changesets.iter().map(|(_, rc)| rc),
                )
                .await?;

                let mut ctx = self.ctx.clone();
                if is_infinitepush
                    && tunables::tunables()
                        .commit_cloud_use_background_session_class()
                        .unwrap_or_default()
                {
                    ctx.session_mut()
                        .override_session_class(SessionClass::BackgroundUnlessTooSlow);
                }

                let filelogs = upload_hg_blobs(
                    &ctx,
                    self.repo,
                    convert_to_revlog_filelog(self.ctx.clone(), self.repo.clone(), filelogs),
                )
                .await
                .context("While uploading File Blobs")?;

                let cg_push =
                    build_changegroup_push(self.ctx, self.repo, header, changesets, filelogs)
                        .await?;

                Some(cg_push)
            }
            Some(part) => return return_with_rest_of_bundle(None, part, bundle2).await,
            _ => return Err(format_err!("Unexpected Bundle2 stream end")),
        };

        Ok((maybe_cg_push, bundle2))
    }

    /// Parses pushkey part if it exists
    /// Returns an error if the pushkey namespace is unknown
    async fn maybe_resolve_pushkey(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<(
        Option<Pushkey>,
        BoxStream<'static, Result<Bundle2Item<'static>>>,
    )> {
        match bundle2.try_next().await? {
            Some(Bundle2Item::Pushkey(header, emptypart)) => {
                let namespace = header
                    .mparams()
                    .get("namespace")
                    .ok_or_else(|| format_err!("pushkey: `namespace` parameter is not set"))?;

                let pushkey = match &namespace[..] {
                    b"phases" => Pushkey::Phases,
                    b"bookmarks" => {
                        let part_id = header.part_id();
                        let mparams = header.mparams();
                        let name = get_ascii_param(mparams, "key")?;
                        let name = BookmarkKey::new_ascii(name);
                        let old = get_optional_changeset_param(mparams, "old")?;
                        let new = get_optional_changeset_param(mparams, "new")?;

                        Pushkey::HgBookmarkPush(PlainBookmarkPush {
                            part_id,
                            name,
                            old,
                            new,
                        })
                    }
                    _ => {
                        return Err(format_err!(
                            "pushkey: unexpected namespace: {:?}",
                            namespace
                        ));
                    }
                };

                emptypart.await?;
                Ok((Some(pushkey), bundle2))
            }
            Some(part) => return_with_rest_of_bundle(None, part, bundle2).await,
            None => Ok((None, bundle2)),
        }
    }

    /// Parse b2xinfinitepushmutation.
    async fn maybe_resolve_infinitepush_mutation(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<
        (
            Vec<HgMutationEntry>,
            BoxStream<'static, Result<Bundle2Item<'static>>>,
        ),
        Error,
    > {
        match bundle2.try_next().await? {
            Some(Bundle2Item::B2xInfinitepushMutation(_, entries)) => {
                let mutations = entries.try_concat().await?;
                Ok((mutations, bundle2))
            }
            Some(part) => return_with_rest_of_bundle(Vec::new(), part, bundle2).await,
            None => Ok((Vec::new(), bundle2)),
        }
    }

    /// Parse b2xtreegroup2.
    /// The Manifests should be scheduled for uploading to BlobRepo and the Future resolving in
    /// their upload as well as their parsed content should be used for uploading changesets.
    async fn resolve_b2xtreegroup2(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<(Manifests, BoxStream<'static, Result<Bundle2Item<'static>>>)> {
        match bundle2.try_next().await? {
            Some(Bundle2Item::B2xTreegroup2(_, parts))
            | Some(Bundle2Item::B2xRebasePack(_, parts)) => {
                let manifests = upload_hg_blobs(
                    self.ctx,
                    self.repo,
                    TreemanifestBundle2Parser::new(parts).compat(),
                )
                .await
                .context("While uploading Manifest Blobs")?;

                Ok((manifests, bundle2))
            }
            _ => Err(format_err!("Expected Bundle2 B2xTreegroup2")),
        }
    }

    /// Parse b2xinfinitepushscratchbookmarks.
    /// This part is ignored, so just parse it and forget it
    async fn maybe_resolve_infinitepush_bookmarks(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<((), BoxStream<'static, Result<Bundle2Item<'static>>>)> {
        match bundle2.try_next().await? {
            Some(Bundle2Item::B2xInfinitepushBookmarks(_, bookmarks)) => {
                bookmarks.try_for_each(|_| future::ok(())).await?;
                Ok(((), bundle2))
            }
            None => Ok(((), bundle2)),
            _ => Err(format_err!(
                "Expected B2xInfinitepushBookmarks or end of the stream"
            )),
        }
    }

    /// Takes parsed Changesets and scheduled for upload Filelogs and Manifests. The content of
    /// Manifests is used to figure out DAG of dependencies between a given Changeset and the
    /// Manifests and Filelogs it adds.
    /// The Changesets are scheduled for uploading and a Future is returned, whose completion means
    /// that the changesets were uploaded
    async fn upload_changesets(
        &self,
        cg_push: ChangegroupPush,
        manifests: Manifests,
        maybe_backup_repo_source: Option<BackupSourceRepo>,
    ) -> Result<(UploadedBonsais, UploadedHgChangesetIds), Error> {
        let changesets = toposort_changesets(cg_push.changesets)?;
        let filelogs = cg_push.filelogs;

        self.ctx
            .scuba()
            .clone()
            .add("changeset_count", changesets.len())
            .add("manifests_count", manifests.len())
            .add("filelogs_count", filelogs.len())
            .log_with_msg("Size of unbundle", None);

        STATS::changesets_count.add_value(changesets.len() as i64);
        STATS::manifests_count.add_value(manifests.len() as i64);
        STATS::filelogs_count.add_value(filelogs.len() as i64);

        let err_context = || {
            let changesets_hashes: Vec<_> = changesets.iter().map(|(hash, _)| *hash).collect();
            ErrorKind::WhileUploadingData(changesets_hashes)
        };

        trace!(self.ctx.logger(), "changesets: {:?}", changesets);
        trace!(self.ctx.logger(), "filelogs: {:?}", filelogs.keys());
        trace!(self.ctx.logger(), "manifests: {:?}", manifests.keys());

        // Each commit gets a future. This future polls futures of parent commits, which poll futures
        // of their parents and so on. However that might cause stackoverflow on very large pushes
        // To avoid it we commit changesets in relatively small chunks.
        let chunk_size = 100;

        let mut bonsais = UploadedBonsais::new();
        let mut hg_cs_ids = UploadedHgChangesetIds::new();
        for chunk in changesets.chunks(chunk_size) {
            let mut uploaded_changesets: HashMap<HgChangesetId, ChangesetHandle> = HashMap::new();
            for (node, revlog_cs) in chunk {
                uploaded_changesets = upload_changeset(
                    self.ctx.clone(),
                    self.repo.clone(),
                    self.ctx.scuba().clone(),
                    *node,
                    revlog_cs,
                    uploaded_changesets,
                    &filelogs,
                    &manifests,
                    maybe_backup_repo_source.clone(),
                )
                .await
                .with_context(err_context)?;
            }

            let uploaded: Vec<(BonsaiChangeset, HgChangesetId)> = stream::iter(uploaded_changesets)
                .map(move |(hg_cs_id, handle): (HgChangesetId, _)| async move {
                    let shared_item_bcs_and_something = handle.get_completed_changeset().await?;

                    let bcs = shared_item_bcs_and_something.0;
                    Result::<_, Error>::Ok((bcs, hg_cs_id))
                })
                .buffered(chunk_size)
                .try_collect()
                .await
                .with_context(err_context)?;

            bonsais.reserve(uploaded.len());
            hg_cs_ids.reserve(uploaded.len());
            for (bcs, hg_cs_id) in uploaded {
                bonsais.insert(bcs);
                hg_cs_ids.insert(hg_cs_id);
            }
        }

        Ok((bonsais, hg_cs_ids))
    }

    /// Ensures that the next item in stream is None
    async fn ensure_stream_finished(
        &self,
        mut bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
    ) -> Result<(), Error> {
        ensure!(
            bundle2.try_next().await?.is_none(),
            "Expected end of Bundle2"
        );
        Ok(())
    }

    /// A method that can use any of the above maybe_resolve_* methods to return
    /// a Vec of (potentailly multiple) Part rather than an Option of Part.
    /// The original use case is to parse multiple pushkey Parts since bundle2 gets
    /// one pushkey part per bookmark.
    async fn resolve_multiple_parts<'a, T, Func, Fut>(
        &'a self,
        bundle2: BoxStream<'static, Result<Bundle2Item<'static>>>,
        mut maybe_resolve: Func,
    ) -> Result<(Vec<T>, BoxStream<'static, Result<Bundle2Item<'static>>>)>
    where
        Fut: Future<Output = Result<(Option<T>, BoxStream<'static, Result<Bundle2Item<'static>>>)>>
            + Sized,
        Func: FnMut(&'a Self, BoxStream<'static, Result<Bundle2Item<'static>>>) -> Fut
            + Send
            + 'static,
        T: Send + 'static,
    {
        let mut result = Vec::new();
        let mut bundle2 = bundle2;
        loop {
            let (maybe_element, rest_of_bundle2) = maybe_resolve(self, bundle2).await?;
            bundle2 = rest_of_bundle2;
            if let Some(element) = maybe_element {
                result.push(element);
            } else {
                break;
            }
        }
        Ok((result, bundle2))
    }
}

fn get_optional_ascii_param(
    params: &HashMap<String, Bytes>,
    param: &str,
) -> Option<Result<AsciiString>> {
    params.get(param).map(|val| {
        AsciiString::from_ascii(val.to_vec())
            .map_err(|err| format_err!("`{}` parameter is not ascii: {}", param, err))
    })
}

fn get_ascii_param(params: &HashMap<String, Bytes>, param: &str) -> Result<AsciiString> {
    get_optional_ascii_param(params, param)
        .unwrap_or_else(|| Err(format_err!("`{}` parameter is not set", param)))
}

fn get_optional_changeset_param(
    params: &HashMap<String, Bytes>,
    param: &str,
) -> Result<Option<HgChangesetId>> {
    let val = get_ascii_param(params, param)?;

    if val.is_empty() {
        Ok(None)
    } else {
        Ok(Some(HgChangesetId::from_ascii_str(&val)?))
    }
}

async fn build_changegroup_push(
    ctx: &CoreContext,
    repo: &BlobRepo,
    part_header: PartHeader,
    changesets: Changesets,
    filelogs: Filelogs,
) -> Result<ChangegroupPush, Error> {
    let PartHeaderInner {
        part_id,
        part_type,
        aparams,
        mparams,
        ..
    } = part_header.into_inner();

    let infinitepush_payload = match part_type {
        PartHeaderType::B2xInfinitepush => {
            let maybe_name_res = get_optional_ascii_param(&aparams, "bookmark")
                .transpose()
                .map(|maybe_name| maybe_name.map(BookmarkKey::new_ascii));

            let bookmark_push = match maybe_name_res {
                Err(e) => return Err(e),
                Ok(maybe_name) => match maybe_name {
                    None => None,
                    Some(name) => {
                        let old = repo.get_bookmark_hg(ctx.clone(), &name).await?;
                        // NOTE: We do not validate that the bookmarknode selected (i.e. the
                        // changeset we should update our bookmark to) is part of the
                        // changegroup being pushed. We do however validate at a later point
                        // that this changeset exists.
                        let new = get_ascii_param(&aparams, "bookmarknode")?;
                        let new = HgChangesetId::from_ascii_str(&new)?;
                        let create = aparams.get("create").is_some();
                        let force = aparams.get("force").is_some();

                        Some(InfiniteBookmarkPush {
                            name,
                            create,
                            force,
                            old,
                            new,
                        })
                    }
                },
            };

            Some(InfinitepushPayload { bookmark_push })
        }
        _ => None,
    };

    Ok(ChangegroupPush {
        part_id,
        changesets,
        filelogs,
        mparams,
        infinitepush_payload,
    })
}

fn collect_pushkey_bookmark_pushes(
    pushkeys: Vec<Pushkey>,
) -> Vec<PlainBookmarkPush<HgChangesetId>> {
    pushkeys
        .into_iter()
        .filter_map(|pushkey| match pushkey {
            Pushkey::Phases => None,
            Pushkey::HgBookmarkPush(bp) => Some(bp),
        })
        .collect()
}

fn try_collect_all_bookmark_pushes(
    pushkeys: Vec<Pushkey>,
    infinitepush_bookmark_push: Option<InfiniteBookmarkPush<HgChangesetId>>,
) -> Result<Option<AllBookmarkPushes<HgChangesetId>>> {
    let bookmark_pushes: Vec<_> = collect_pushkey_bookmark_pushes(pushkeys)
        .into_iter()
        .collect();
    let bookmark_pushes_len = bookmark_pushes.len();
    match (bookmark_pushes_len, infinitepush_bookmark_push) {
        (0, Some(infinitepush_bookmark_push)) => {
            STATS::bookmark_pushkeys_count.add_value(1);
            Ok(Some(AllBookmarkPushes::Inifinitepush(
                infinitepush_bookmark_push,
            )))
        }
        (bookmark_pushes_len, None) if bookmark_pushes_len > 0 => {
            STATS::bookmark_pushkeys_count.add_value(bookmark_pushes_len as i64);
            Ok(Some(AllBookmarkPushes::PlainPushes(bookmark_pushes)))
        }
        // Neither plain, not infinitepush bookmark pushes are present
        (0, None) => Ok(None),
        (_, Some(_)) => Err(format_err!(
            "Same bundle2 can not be used for both plain and infinite push"
        )),
        (_, _) => Err(format_err!("An unreachable pattern. Programmer's error")),
    }
}

/// Helper fn to return some (usually "empty") value and
/// chain together an unused part with the rest of the bundle
async fn return_with_rest_of_bundle<T: Send + 'static>(
    value: T,
    unused_part: Bundle2Item<'static>,
    rest_of_bundle: BoxStream<'static, Result<Bundle2Item<'static>>>,
) -> Result<(T, BoxStream<'static, Result<Bundle2Item<'static>>>)> {
    Ok((
        value,
        stream::once(async { Ok(unused_part) })
            .chain(rest_of_bundle)
            .boxed(),
    ))
}

fn toposort_changesets(
    changesets: Vec<(HgChangesetId, RevlogChangeset)>,
) -> Result<Vec<(HgChangesetId, RevlogChangeset)>> {
    let mut changesets: HashMap<_, _> = changesets.into_iter().collect();

    // Make sure changesets are toposorted
    let cs_id_to_parents: HashMap<_, _> = changesets
        .iter()
        .map(|(cs_id, revlog_cs)| {
            let parents: Vec<_> = revlog_cs
                .parents()
                .into_iter()
                .map(HgChangesetId::new)
                .collect();
            (*cs_id, parents)
        })
        .collect();
    let sorted_css = sort_topological(&cs_id_to_parents)
        .ok_or_else(|| Error::msg("cycle in the pushed changesets!"))?;

    Ok(sorted_css
        .into_iter()
        .filter_map(|cs| changesets.remove(&cs).map(|revlog_cs| (cs, revlog_cs)))
        .collect())
}

async fn bonsai_from_hg_opt(
    ctx: &CoreContext,
    repo: &BlobRepo,
    cs_id: Option<HgChangesetId>,
) -> Result<Option<ChangesetId>, Error> {
    match cs_id {
        None => Ok(None),
        Some(cs_id) => {
            let maybe_bcs_id = repo
                .bonsai_hg_mapping()
                .get_bonsai_from_hg(ctx, cs_id)
                .await?;
            if maybe_bcs_id.is_none() {
                Err(format_err!("No bonsai mapping found for {}", cs_id))
            } else {
                Ok(maybe_bcs_id)
            }
        }
    }
}

async fn plain_hg_bookmark_push_to_bonsai(
    ctx: &CoreContext,
    repo: &BlobRepo,
    bookmark_push: PlainBookmarkPush<HgChangesetId>,
) -> Result<PlainBookmarkPush<ChangesetId>, Error> {
    let PlainBookmarkPush {
        part_id,
        name,
        old,
        new,
    } = bookmark_push;

    let (old, new) = try_join!(
        bonsai_from_hg_opt(ctx, repo, old),
        bonsai_from_hg_opt(ctx, repo, new),
    )?;

    Ok(PlainBookmarkPush {
        part_id,
        name,
        old,
        new,
    })
}

async fn infinite_hg_bookmark_push_to_bonsai(
    ctx: &CoreContext,
    repo: &BlobRepo,
    bookmark_push: InfiniteBookmarkPush<HgChangesetId>,
) -> Result<InfiniteBookmarkPush<ChangesetId>, Error> {
    let InfiniteBookmarkPush {
        name,
        force,
        create,
        old,
        new,
    } = bookmark_push;

    let (old, new) = try_join!(
        bonsai_from_hg_opt(ctx, repo, old),
        repo.bonsai_hg_mapping().get_bonsai_from_hg(ctx, new)
    )?;
    let new = match new {
        Some(new) => new,
        None => bail!("Bonsai Changeset not found"),
    };

    Ok(InfiniteBookmarkPush {
        name,
        force,
        create,
        old,
        new,
    })
}

async fn hg_pushrebase_bookmark_spec_to_bonsai(
    ctx: &CoreContext,
    repo: &BlobRepo,
    bookmark_spec: PushrebaseBookmarkSpec<HgChangesetId>,
) -> Result<PushrebaseBookmarkSpec<ChangesetId>, Error> {
    let pbs = match bookmark_spec {
        PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark) => {
            PushrebaseBookmarkSpec::NormalPushrebase(onto_bookmark)
        }
        PushrebaseBookmarkSpec::ForcePushrebase(plain_push) => {
            PushrebaseBookmarkSpec::ForcePushrebase(
                plain_hg_bookmark_push_to_bonsai(ctx, repo, plain_push).await?,
            )
        }
    };
    Ok(pbs)
}

async fn hg_all_bookmark_pushes_to_bonsai(
    ctx: &CoreContext,
    repo: &BlobRepo,
    all_bookmark_pushes: AllBookmarkPushes<HgChangesetId>,
) -> Result<AllBookmarkPushes<ChangesetId>, Error> {
    let abp = match all_bookmark_pushes {
        AllBookmarkPushes::PlainPushes(plain_pushes) => {
            let r = try_join_all(
                plain_pushes
                    .into_iter()
                    .map(|plain_push| plain_hg_bookmark_push_to_bonsai(ctx, repo, plain_push)),
            )
            .await?;
            AllBookmarkPushes::PlainPushes(r)
        }
        AllBookmarkPushes::Inifinitepush(infinite_bookmark_push) => {
            let r = infinite_hg_bookmark_push_to_bonsai(ctx, repo, infinite_bookmark_push).await?;
            AllBookmarkPushes::Inifinitepush(r)
        }
    };
    Ok(abp)
}


/*
 * 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::io::Cursor;

use anyhow::Context;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobrepo_hg::BlobRepoHg;
use bookmarks::BookmarkKey;
use bytes::Bytes;
use bytes::BytesMut;
use context::CoreContext;
use futures::compat::Future01CompatExt;
use futures::compat::Stream01CompatExt;
use futures::future::try_join;
use futures::TryStreamExt;
use futures_stats::TimedTryFutureExt;
use getbundle_response::create_getbundle_response;
use getbundle_response::PhasesPart;
use getbundle_response::SessionLfsParams;
use mercurial_bundles::create_bundle_stream;
use mercurial_bundles::parts;
use mercurial_bundles::Bundle2EncodeBuilder;
use mercurial_bundles::PartId;
use mercurial_derivation::DeriveHgChangeset;
use metaconfig_types::PushrebaseParams;
use mononoke_types::ChangesetId;

use crate::CommonHeads;

/// Data, needed to generate a `Push` response
pub struct UnbundlePushResponse {
    pub changegroup_id: Option<PartId>,
    pub bookmark_ids: Vec<PartId>,
}

/// Data, needed to generate an `InfinitePush` response
pub struct UnbundleInfinitePushResponse {
    pub changegroup_id: Option<PartId>,
}

/// Data, needed to generate a `PushRebase` response
pub struct UnbundlePushRebaseResponse {
    pub commonheads: CommonHeads,
    pub pushrebased_rev: ChangesetId,
    pub pushrebased_changesets: Vec<pushrebase::PushrebaseChangesetPair>,
    pub onto: BookmarkKey,
    pub bookmark_push_part_id: Option<PartId>,
}

/// Data, needed to generate a bookmark-only `PushRebase` response
pub struct UnbundleBookmarkOnlyPushRebaseResponse {
    pub bookmark_push_part_id: PartId,
}

pub enum UnbundleResponse {
    Push(UnbundlePushResponse),
    InfinitePush(UnbundleInfinitePushResponse),
    PushRebase(UnbundlePushRebaseResponse),
    BookmarkOnlyPushRebase(UnbundleBookmarkOnlyPushRebaseResponse),
}

impl UnbundleResponse {
    fn get_bundle_builder() -> Bundle2EncodeBuilder<Cursor<Vec<u8>>> {
        let writer = Cursor::new(Vec::new());
        let mut bundle = Bundle2EncodeBuilder::new(writer);
        // Mercurial currently hangs while trying to read compressed bundles over the wire:
        // https://bz.mercurial-scm.org/show_bug.cgi?id=5646
        // TODO: possibly enable compression support once this is fixed.
        bundle.set_compressor_type(None);
        bundle
    }

    async fn generate_push_or_infinitepush_response(
        changegroup_id: Option<PartId>,
        bookmark_ids: Vec<PartId>,
    ) -> Result<Bytes> {
        let mut bundle = Self::get_bundle_builder();
        if let Some(changegroup_id) = changegroup_id {
            bundle.add_part(parts::replychangegroup_part(
                parts::ChangegroupApplyResult::Success { heads_num_diff: 0 },
                changegroup_id,
            )?);
        }
        for part_id in bookmark_ids {
            bundle.add_part(parts::replypushkey_part(true, part_id)?);
        }
        let cursor = bundle.build().compat().await?;
        Ok(Bytes::from(cursor.into_inner()))
    }

    async fn generate_push_response_bytes(
        _ctx: &CoreContext,
        data: UnbundlePushResponse,
    ) -> Result<Bytes> {
        let UnbundlePushResponse {
            changegroup_id,
            bookmark_ids,
        } = data;
        Self::generate_push_or_infinitepush_response(changegroup_id, bookmark_ids)
            .await
            .context("While preparing push response")
    }

    async fn generate_inifinitepush_response_bytes(
        _ctx: &CoreContext,
        data: UnbundleInfinitePushResponse,
    ) -> Result<Bytes> {
        let UnbundleInfinitePushResponse { changegroup_id } = data;
        Self::generate_push_or_infinitepush_response(changegroup_id, vec![])
            .await
            .context("While preparing infinitepush response")
    }

    async fn generate_pushrebase_response_bytes(
        ctx: &CoreContext,
        data: UnbundlePushRebaseResponse,
        repo: &BlobRepo,
        pushrebase_params: PushrebaseParams,
        lfs_params: &SessionLfsParams,
    ) -> Result<Bytes> {
        let UnbundlePushRebaseResponse {
            commonheads,
            pushrebased_rev,
            pushrebased_changesets,
            onto,
            bookmark_push_part_id,
        } = data;

        // Send to the client both pushrebased commit and current "onto" bookmark. Normally they
        // should be the same, however they might be different if bookmark
        // suddenly moved before current pushrebase finished.
        let common = commonheads.heads;
        let maybe_onto_head = repo.get_bookmark_hg(ctx.clone(), &onto);
        let pushrebased_hg_rev = repo.derive_hg_changeset(ctx, pushrebased_rev);

        let bookmark_reply_part = match bookmark_push_part_id {
            Some(part_id) => Some(parts::replypushkey_part(true, part_id)?),
            None => None,
        };

        let obsmarkers_part = match pushrebase_params.emit_obsmarkers {
            true => obsolete::pushrebased_changesets_to_obsmarkers_part(
                ctx,
                repo,
                pushrebased_changesets,
            )
            .transpose()?,
            false => None,
        };

        let mut scuba_logger = ctx.scuba().clone();
        let (stats, response_bytes) = async move {
            let (maybe_onto_head, pushrebased_hg_rev) =
                try_join(maybe_onto_head, pushrebased_hg_rev).await?;

            let mut heads = vec![];
            if let Some(onto_head) = maybe_onto_head {
                heads.push(onto_head);
            }
            heads.push(pushrebased_hg_rev);
            let mut cg_part_builder =
                create_getbundle_response(ctx, repo, common, &heads, PhasesPart::Yes, lfs_params)
                    .await?;

            cg_part_builder.extend(bookmark_reply_part.into_iter());
            cg_part_builder.extend(obsmarkers_part.into_iter());
            let compression = None;
            let chunks = create_bundle_stream(cg_part_builder, compression)
                .compat()
                .try_collect::<Vec<_>>()
                .await?;

            let mut total_capacity = 0;
            for c in chunks.iter() {
                total_capacity += c.len();
            }

            // TODO(stash): make push and pushrebase response streamable - T34090105
            let mut res = BytesMut::with_capacity(total_capacity);
            for c in chunks {
                res.extend_from_slice(&c);
            }
            Result::<_>::Ok(res.freeze())
        }
        .try_timed()
        .await
        .context("While preparing pushrebase response")?;

        scuba_logger
            .add_future_stats(&stats)
            .log_with_msg("Pushrebase: prepared the response", None);
        Ok(response_bytes)
    }

    async fn generate_bookmark_only_pushrebase_response_bytes(
        _ctx: &CoreContext,
        data: UnbundleBookmarkOnlyPushRebaseResponse,
    ) -> Result<Bytes> {
        let UnbundleBookmarkOnlyPushRebaseResponse {
            bookmark_push_part_id,
        } = data;

        let mut bundle = Self::get_bundle_builder();
        bundle.add_part(parts::replypushkey_part(true, bookmark_push_part_id)?);
        let cursor = bundle
            .build()
            .compat()
            .await
            .context("While preparing bookmark-only pushrebase response")?;

        Ok(Bytes::from(cursor.into_inner()))
    }

    /// Produce bundle2 response parts for the completed `unbundle` processing
    pub async fn generate_bytes(
        self,
        ctx: &CoreContext,
        repo: &BlobRepo,
        pushrebase_params: PushrebaseParams,
        lfs_params: &SessionLfsParams,
        respondlightly: Option<bool>,
    ) -> Result<Bytes> {
        if let Some(true) = respondlightly {
            let bundle = Self::get_bundle_builder();
            let cursor = bundle.build().compat().await?;
            return Ok(Bytes::from(cursor.into_inner()));
        }
        match self {
            UnbundleResponse::Push(data) => Self::generate_push_response_bytes(ctx, data).await,
            UnbundleResponse::InfinitePush(data) => {
                Self::generate_inifinitepush_response_bytes(ctx, data).await
            }
            UnbundleResponse::PushRebase(data) => {
                Self::generate_pushrebase_response_bytes(
                    ctx,
                    data,
                    repo,
                    pushrebase_params,
                    lfs_params,
                )
                .await
            }
            UnbundleResponse::BookmarkOnlyPushRebase(data) => {
                Self::generate_bookmark_only_pushrebase_response_bytes(ctx, data).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.
 */

pub(crate) use stats::prelude::*;

define_stats! {
    prefix = "mononoke.bundle2_resolver";
    deltacache_dsize: histogram(400, 0, 100_000, Average, Sum, Count; P 50; P 95; P 99),
    deltacache_dsize_large: histogram(400_000, 0, 100_000_000; P 50; P 95; P 99),
    deltacache_fsize: histogram(400, 0, 100_000, Average, Sum, Count; P 50; P 95; P 99),
    deltacache_fsize_large: histogram(400_000, 0, 100_000_000; P 50; P 95; P 99),
    bookmark_pushkeys_count: timeseries(Rate, Average, Sum),
    changesets_count: timeseries(Rate, Average, Sum),
    manifests_count: timeseries(Rate, Average, Sum),
    filelogs_count: timeseries(Rate, Average, Sum),
    content_blobs_count: timeseries(Rate, Average, Sum),
    per_changeset_manifests_count: timeseries(Rate, Average, Sum),
    per_changeset_filelogs_count: timeseries(Rate, Average, Sum),
    per_changeset_content_blobs_count: timeseries(Rate, Average, Sum),
}


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

use std::collections::HashMap;

use anyhow::Result;
use blobrepo::BlobRepo;
use context::CoreContext;
use futures::compat::Future01CompatExt;
use futures::future::BoxFuture;
use futures::FutureExt;
use futures::Stream;
use futures::TryStreamExt;
use futures_ext::future::TryShared;
use futures_ext::FbTryFutureExt;
use mercurial_revlog::manifest::ManifestContent;
use mercurial_types::blobs::UploadHgNodeHash;
use mercurial_types::blobs::UploadHgTreeEntry;
use mercurial_types::HgManifestId;
use mercurial_types::HgNodeHash;
use mercurial_types::HgNodeKey;
use mononoke_types::RepoPath;
use repo_blobstore::RepoBlobstoreArc;
use wirepack::TreemanifestEntry;

/// Represents data that is Mercurial-encoded and can be uploaded to the blobstore.
pub trait UploadableHgBlob {
    type Value: Send + 'static;

    fn upload(self, ctx: &CoreContext, repo: &BlobRepo) -> Result<(HgNodeKey, Self::Value)>;
}

pub(crate) async fn upload_hg_blobs<'a, S, B>(
    ctx: &'a CoreContext,
    repo: &'a BlobRepo,
    blobs: S,
) -> Result<HashMap<HgNodeKey, B::Value>>
where
    S: Stream<Item = Result<B>> + Send + 'a,
    B: UploadableHgBlob + 'a,
{
    blobs
        .try_fold(HashMap::new(), move |mut map, item| async move {
            let (key, value) = item.upload(ctx, repo)?;
            map.insert(key, value);
            Ok(map)
        })
        .await
}

impl UploadableHgBlob for TreemanifestEntry {
    // * Shared is required here because a single tree manifest can be referred to by more than
    //   one changeset, and all of those will want to refer to the corresponding future.
    type Value = (
        ManifestContent,
        Option<HgNodeHash>,
        Option<HgNodeHash>,
        TryShared<BoxFuture<'static, Result<(HgManifestId, RepoPath)>>>,
    );

    fn upload(self, ctx: &CoreContext, repo: &BlobRepo) -> Result<(HgNodeKey, Self::Value)> {
        let node_key = self.node_key;
        let manifest_content = self.manifest_content;
        let p1 = self.p1;
        let p2 = self.p2;
        // The root tree manifest is expected to have the wrong hash in hybrid mode.
        // XXX possibly remove this once hybrid mode is gone
        let upload_node_id = if node_key.path.is_root() {
            UploadHgNodeHash::Supplied(node_key.hash)
        } else {
            UploadHgNodeHash::Checked(node_key.hash)
        };
        let upload = UploadHgTreeEntry {
            upload_node_id,
            contents: self.data,
            p1: self.p1,
            p2: self.p2,
            path: node_key.path.clone(),
        };
        let (_node, value) = upload.upload(ctx.clone(), repo.repo_blobstore_arc())?;
        Ok((
            node_key,
            (
                manifest_content,
                p1,
                p2,
                value.compat().boxed().try_shared(),
            ),
        ))
    }
}


/*
 * 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::ops::AddAssign;

use ::manifest::Entry;
use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobrepo_hg::create_bonsai_changeset_hook;
use blobrepo_hg::ChangesetHandle;
use blobrepo_hg::CreateChangeset;
use context::CoreContext;
use futures::future;
use futures::future::BoxFuture;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::FuturesUnordered;
use futures::FutureExt;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use mercurial_revlog::changeset::RevlogChangeset;
use mercurial_revlog::manifest::Details;
use mercurial_revlog::manifest::ManifestContent;
use mercurial_types::blobs::ChangesetMetadata;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestId;
use mercurial_types::HgNodeHash;
use mercurial_types::HgNodeKey;
use mercurial_types::MPath;
use mercurial_types::RepoPath;
use mercurial_types::NULL_HASH;
use scuba_ext::MononokeScubaSampleBuilder;
use wirepack::TreemanifestEntry;
use wireproto_handler::BackupSourceRepo;

use crate::changegroup::Filelog;
use crate::stats::*;
use crate::upload_blobs::UploadableHgBlob;

pub type Filelogs = HashMap<HgNodeKey, <Filelog as UploadableHgBlob>::Value>;
pub type Manifests = HashMap<HgNodeKey, <TreemanifestEntry as UploadableHgBlob>::Value>;
pub type UploadedChangesets = HashMap<HgChangesetId, ChangesetHandle>;

type HgBlobFuture = BoxFuture<'static, Result<(Entry<HgManifestId, HgFileNodeId>, RepoPath)>>;
type HgBlobStream = BoxStream<'static, Result<(Entry<HgManifestId, HgFileNodeId>, RepoPath)>>;

/// In order to generate the DAG of dependencies between Root Manifest and other Manifests and
/// Filelogs we need to walk that DAG.
/// This represents the manifests and file nodes introduced by a particular changeset.
pub(crate) struct NewBlobs {
    // root_manifest can be None f.e. when commit removes all the content of the repo
    root_manifest: BoxFuture<'static, Result<Option<(HgManifestId, RepoPath)>>>,
    // sub_entries has both submanifest and filenode entries.
    sub_entries: HgBlobStream,
}

struct WalkHelperCounters {
    manifests_count: usize,
    filelogs_count: usize,
    content_blobs_count: usize,
}

impl AddAssign for WalkHelperCounters {
    fn add_assign(&mut self, other: WalkHelperCounters) {
        *self = Self {
            manifests_count: self.manifests_count + other.manifests_count,
            filelogs_count: self.filelogs_count + other.filelogs_count,
            content_blobs_count: self.content_blobs_count + other.content_blobs_count,
        };
    }
}

impl NewBlobs {
    pub(crate) fn new(
        manifest_root_id: HgManifestId,
        manifests: &Manifests,
        filelogs: &Filelogs,
    ) -> Result<Self> {
        if manifest_root_id.into_nodehash() == NULL_HASH {
            // If manifest root id is NULL_HASH then there is no content in this changest
            return Ok(Self {
                root_manifest: future::ok(None).boxed(),
                sub_entries: stream::empty().boxed(),
            });
        }

        let root_key = HgNodeKey {
            path: RepoPath::root(),
            hash: manifest_root_id.clone().into_nodehash(),
        };

        let (entries, root_manifest) = match manifests.get(&root_key) {
            Some((ref manifest_content, ref p1, ref p2, ref manifest_root)) => {
                let (entries, counters) = Self::walk_helper(
                    &RepoPath::root(),
                    manifest_content,
                    get_manifest_parent_content(manifests, RepoPath::root(), p1.clone()),
                    get_manifest_parent_content(manifests, RepoPath::root(), p2.clone()),
                    manifests,
                    filelogs,
                )?;
                STATS::per_changeset_manifests_count.add_value(counters.manifests_count as i64);
                STATS::per_changeset_filelogs_count.add_value(counters.filelogs_count as i64);
                STATS::per_changeset_content_blobs_count
                    .add_value(counters.content_blobs_count as i64);
                let root_manifest = manifest_root
                    .clone()
                    .map_ok(Some)
                    .map_err(Error::from)
                    .boxed();

                (entries, root_manifest)
            }
            None => {
                let entry = (manifest_root_id, RepoPath::RootPath);
                (vec![], future::ok(Some(entry)).boxed())
            }
        };

        let buffer_size = tunables::tunables()
            .repo_client_concurrent_blob_uploads()
            .unwrap_or_default();
        let s = if buffer_size <= 0 {
            entries
                .into_iter()
                .collect::<FuturesUnordered<_>>()
                .left_stream()
        } else {
            stream::iter(entries)
                .buffer_unordered(buffer_size as usize)
                .right_stream()
        };

        Ok(Self {
            root_manifest,
            sub_entries: s
                .map_err(move |err| {
                    err.context(format!(
                        "While walking dependencies of Root Manifest with id {:?}",
                        manifest_root_id
                    ))
                })
                .boxed(),
        })
    }

    fn walk_helper(
        path_taken: &RepoPath,
        manifest_content: &ManifestContent,
        p1: Option<&ManifestContent>,
        p2: Option<&ManifestContent>,
        manifests: &Manifests,
        filelogs: &Filelogs,
    ) -> Result<(Vec<HgBlobFuture>, WalkHelperCounters)> {
        if path_taken.len() > 4096 {
            bail!(
                "Exceeded max manifest path during walking with path: {:?}",
                path_taken
            );
        }

        let mut entries: Vec<HgBlobFuture> = Vec::new();
        let mut counters = WalkHelperCounters {
            manifests_count: 0,
            filelogs_count: 0,
            content_blobs_count: 0,
        };

        for (name, details) in manifest_content.files.iter() {
            if is_entry_present_in_parent(p1, name, details)
                || is_entry_present_in_parent(p2, name, details)
            {
                // If one of the parents contains exactly the same version of entry then either that
                // file or manifest subtree is not new
                continue;
            }

            let nodehash = details.entryid().clone().into_nodehash();
            let next_path = MPath::join_opt(path_taken.mpath(), name);
            let next_path = match next_path {
                Some(path) => path,
                None => bail!("internal error: joined root path with root manifest"),
            };

            if details.is_tree() {
                let key = HgNodeKey {
                    path: RepoPath::DirectoryPath(next_path),
                    hash: nodehash,
                };

                if let Some((manifest_content, p1, p2, blobfuture)) = manifests.get(&key) {
                    counters.manifests_count += 1;
                    entries.push(
                        blobfuture
                            .clone()
                            .map_ok(|(id, path)| (Entry::Tree(id), path))
                            .map_err(Error::from)
                            .boxed(),
                    );
                    let (mut walked_entries, sub_counters) = Self::walk_helper(
                        &key.path,
                        manifest_content,
                        get_manifest_parent_content(manifests, key.path.clone(), p1.clone()),
                        get_manifest_parent_content(manifests, key.path.clone(), p2.clone()),
                        manifests,
                        filelogs,
                    )?;
                    entries.append(&mut walked_entries);
                    counters += sub_counters;
                }
            } else {
                let key = HgNodeKey {
                    path: RepoPath::FilePath(next_path),
                    hash: nodehash,
                };
                if let Some(blobfuture) = filelogs.get(&key) {
                    counters.filelogs_count += 1;
                    counters.content_blobs_count += 1;
                    entries.push(
                        blobfuture
                            .clone()
                            .map_ok(|(id, path)| (Entry::Leaf(id), path))
                            .map_err(Error::from)
                            .boxed(),
                    );
                }
            }
        }

        Ok((entries, counters))
    }
}

fn get_manifest_parent_content(
    manifests: &Manifests,
    path: RepoPath,
    p: Option<HgNodeHash>,
) -> Option<&ManifestContent> {
    p.and_then(|p| manifests.get(&HgNodeKey { path, hash: p }))
        .map(|(content, ..)| content)
}

fn is_entry_present_in_parent(
    p: Option<&ManifestContent>,
    name: &MPath,
    details: &Details,
) -> bool {
    match p.and_then(|p| p.files.get(name)) {
        None => false,
        Some(parent_details) => parent_details == details,
    }
}

/// Retrieves the parent from uploaded changesets, if it is missing then fetches it from BlobRepo
fn get_parent(
    ctx: CoreContext,
    repo: &BlobRepo,
    map: &UploadedChangesets,
    p: Option<HgNodeHash>,
) -> Option<ChangesetHandle> {
    match p {
        None => None,
        Some(p) => match map.get(&HgChangesetId::new(p)) {
            None => Some(ChangesetHandle::ready_cs_handle(
                ctx,
                repo.clone(),
                HgChangesetId::new(p),
            )),
            Some(cs) => Some(cs.clone()),
        },
    }
}

pub async fn upload_changeset(
    ctx: CoreContext,
    repo: BlobRepo,
    scuba_logger: MononokeScubaSampleBuilder,
    node: HgChangesetId,
    revlog_cs: &RevlogChangeset,
    mut uploaded_changesets: UploadedChangesets,
    filelogs: &Filelogs,
    manifests: &Manifests,
    maybe_backup_repo_source: Option<BackupSourceRepo>,
) -> Result<UploadedChangesets, Error> {
    let NewBlobs {
        root_manifest,
        sub_entries,
    } = NewBlobs::new(revlog_cs.manifestid(), manifests, filelogs)?;

    let cs_metadata = ChangesetMetadata {
        user: String::from_utf8(revlog_cs.user().into())?,
        time: revlog_cs.time().clone(),
        extra: revlog_cs.extra().clone(),
        message: String::from_utf8(revlog_cs.message().into())?,
    };

    // DO NOT try to comute p1 and p2 concurrently!
    // It may result in a combinatoral explosion in mergy repos (see D14100259)
    let p1 = get_parent(ctx.clone(), &repo, &uploaded_changesets, revlog_cs.p1);
    let p2 = get_parent(ctx.clone(), &repo, &uploaded_changesets, revlog_cs.p2);

    let create_bonsai_changeset_hook = Some(create_bonsai_changeset_hook(maybe_backup_repo_source));
    let create_changeset = CreateChangeset {
        expected_nodeid: Some(node.into_nodehash()),
        expected_files: Some(Vec::from(revlog_cs.files())),
        p1,
        p2,
        root_manifest,
        sub_entries,
        // XXX pass content blobs to CreateChangeset here
        cs_metadata,
        create_bonsai_changeset_hook,
        upload_to_blobstore_only: false,
    };
    let scheduled_uploading = create_changeset.create(ctx, &repo, scuba_logger);

    uploaded_changesets.insert(node, scheduled_uploading);
    Ok(uploaded_changesets)
}


/*
 * 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::ensure;
use anyhow::Result;
use bytes::Bytes;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use mercurial_bundles::changegroup::CgDeltaChunk;
use mercurial_revlog::changeset::RevlogChangeset;
use mercurial_types::delta;
use mercurial_types::HgBlob;
use mercurial_types::HgBlobNode;
use mercurial_types::HgChangesetId;
use mercurial_types::NULL_HASH;

#[derive(Debug, Eq, PartialEq)]
pub(crate) struct ChangesetDeltaed {
    pub chunk: CgDeltaChunk,
}

pub(crate) fn convert_to_revlog_changesets(
    deltaed: impl Stream<Item = Result<ChangesetDeltaed>>,
) -> impl Stream<Item = Result<(HgChangesetId, RevlogChangeset)>> {
    deltaed.and_then(|ChangesetDeltaed { chunk }| async move {
        ensure!(
            chunk.base == NULL_HASH,
            "Changeset chunk base ({:?}) should be equal to root commit ({:?}), \
                 because it is never deltaed",
            chunk.base,
            NULL_HASH
        );
        ensure!(
            chunk.node == chunk.linknode,
            "Changeset chunk node ({:?}) should be equal to linknode ({:?})",
            chunk.node,
            chunk.linknode
        );

        Ok((
            HgChangesetId::new(chunk.node),
            RevlogChangeset::new(HgBlobNode::new(
                HgBlob::from(Bytes::from(delta::apply(b"", &chunk.delta)?)),
                chunk.p1.into_option(),
                chunk.p2.into_option(),
            ))?,
        ))
    })
}

#[cfg(test)]
mod tests {
    use futures::stream::iter;
    use itertools::equal;
    use mercurial_types::HgNodeHash;

    use super::*;

    enum CheckResult {
        ExpectedOk(bool),
        ExpectedErr(bool),
    }
    use self::CheckResult::*;

    async fn check_null_changeset(
        node: HgNodeHash,
        linknode: HgNodeHash,
        base: HgNodeHash,
        p1: HgNodeHash,
        p2: HgNodeHash,
    ) -> CheckResult {
        let blobnode = HgBlobNode::new(
            RevlogChangeset::new_null()
                .get_node()
                .unwrap()
                .as_blob()
                .clone(),
            p1.into_option(),
            p2.into_option(),
        );

        let delta = delta::Delta::new_fulltext(blobnode.as_blob().as_slice());
        let cs = RevlogChangeset::new(blobnode).unwrap();

        let chunk = CgDeltaChunk {
            node,
            p1,
            p2,
            base,
            linknode,
            delta,
            flags: None,
        };

        let result = convert_to_revlog_changesets(iter(vec![Ok(ChangesetDeltaed { chunk })]))
            .try_collect::<Vec<_>>()
            .await;

        if base == NULL_HASH && node == linknode {
            ExpectedOk(equal(result.unwrap(), vec![(HgChangesetId::new(node), cs)]))
        } else {
            ExpectedErr(result.is_err())
        }
    }

    #[quickcheck_async::tokio]
    async fn null_changeset_random(
        node: HgNodeHash,
        linknode: HgNodeHash,
        base: HgNodeHash,
        p1: HgNodeHash,
        p2: HgNodeHash,
    ) -> bool {
        match check_null_changeset(node, linknode, base, p1, p2).await {
            ExpectedOk(true) | ExpectedErr(true) => true,
            _ => false,
        }
    }

    #[quickcheck_async::tokio]
    async fn null_changeset_correct(node: HgNodeHash, p1: HgNodeHash, p2: HgNodeHash) -> bool {
        match check_null_changeset(node.clone(), node, NULL_HASH, p1, p2).await {
            ExpectedOk(true) => true,
            _ => false,
        }
    }
}