/*
 * 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 anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use blobstore::Blobstore;
use bonsai_hg_mapping::BonsaiHgMapping;
use cacheblob::MemWritesBlobstore;
use context::CoreContext;
use filenodes::Filenodes;
use futures::future::try_join_all;
use metaconfig_types::DerivedDataTypesConfig;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;

use crate::derivable::BonsaiDerivable;
use crate::manager::derive::Rederivation;
use crate::manager::DerivedDataManager;

/// Context for performing derivation.
///
/// This struct is passed to derivation implementations.  They can use it
/// to access repository attributes or request access to dependent
/// derived data types.
#[derive(Clone)]
pub struct DerivationContext {
    manager: DerivedDataManager,
    rederivation: Option<Arc<dyn Rederivation>>,
    blobstore: Arc<dyn Blobstore>,

    /// Write cache layered over the blobstore.  This is the same object
    /// with two views, so we can return a reference to the `Arc<dyn
    /// Blobstore>` version if needed.
    blobstore_write_cache: Option<(
        Arc<dyn Blobstore>,
        Arc<MemWritesBlobstore<Arc<dyn Blobstore>>>,
    )>,
}

impl DerivationContext {
    pub(crate) fn new(
        manager: DerivedDataManager,
        rederivation: Option<Arc<dyn Rederivation>>,
        blobstore: Arc<dyn Blobstore>,
    ) -> Self {
        DerivationContext {
            manager,
            rederivation,
            blobstore,
            blobstore_write_cache: None,
        }
    }

    /// Fetch previously derived data.
    pub async fn fetch_derived<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
    ) -> Result<Option<Derivable>>
    where
        Derivable: BonsaiDerivable,
    {
        if let Some(rederivation) = self.rederivation.as_ref() {
            if rederivation.needs_rederive(Derivable::NAME, csid) == Some(true) {
                return Ok(None);
            }
        }
        let derived = Derivable::fetch(ctx, self, csid).await?;
        Ok(derived)
    }

    /// Fetch a batch of previously derived data.
    pub async fn fetch_derived_batch<Derivable>(
        &self,
        ctx: &CoreContext,
        mut csids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, Derivable>>
    where
        Derivable: BonsaiDerivable,
    {
        if let Some(rederivation) = self.rederivation.as_ref() {
            csids.retain(|csid| rederivation.needs_rederive(Derivable::NAME, *csid) != Some(true));
        }
        let derived = Derivable::fetch_batch(ctx, self, &csids).await?;
        Ok(derived)
    }

    /// Fetch previously derived data that is known to be derived as it is a
    /// dependency of the current changeset.
    pub async fn fetch_dependency<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
    ) -> Result<Derivable>
    where
        Derivable: BonsaiDerivable,
    {
        self.fetch_derived(ctx, csid).await?.ok_or_else(|| {
            anyhow!(
                "dependency '{}' of {} was not already derived",
                Derivable::NAME,
                csid
            )
        })
    }

    /// Fetch a dependency for all parents of a changeset.
    pub async fn fetch_parents<Derivable>(
        &self,
        ctx: &CoreContext,
        bonsai: &BonsaiChangeset,
    ) -> Result<Vec<Derivable>>
    where
        Derivable: BonsaiDerivable,
    {
        self.fetch_unknown_parents(ctx, None, bonsai).await
    }

    /// Fetch a dependency for all parents of a changeset if the dependency is
    /// not already known.
    pub async fn fetch_unknown_parents<Derivable>(
        &self,
        ctx: &CoreContext,
        known: Option<&HashMap<ChangesetId, Derivable>>,
        bonsai: &BonsaiChangeset,
    ) -> Result<Vec<Derivable>>
    where
        Derivable: BonsaiDerivable,
    {
        try_join_all(bonsai.parents().map(|p| async move {
            self.fetch_unknown_dependency(ctx, known, p)
                .await
                .with_context(|| {
                    format!(
                        "could not fetch '{}' for parents of {}",
                        Derivable::NAME,
                        bonsai.get_changeset_id(),
                    )
                })
        }))
        .await
    }

    /// Fetch derived data value for changeset if it is not already known.
    pub async fn fetch_unknown_dependency<Derivable>(
        &self,
        ctx: &CoreContext,
        known: Option<&HashMap<ChangesetId, Derivable>>,
        csid: ChangesetId,
    ) -> Result<Derivable>
    where
        Derivable: BonsaiDerivable,
    {
        if let Some(parent) = known.and_then(|k| k.get(&csid)) {
            return Ok(parent.clone());
        }
        self.fetch_dependency(ctx, csid).await
    }

    /// Cause derivation of a dependency, and fetch the result.
    ///
    /// In the future, this will be removed in favour of making the manager
    /// arrange for dependent derived data to always be derived before
    /// derivation starts.
    pub async fn derive_dependency<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
    ) -> Result<Derivable>
    where
        Derivable: BonsaiDerivable,
    {
        Ok(self
            .manager
            .derive::<Derivable>(ctx, csid, self.rederivation.clone())
            .await?)
    }

    /// The repo id of the repo being derived.
    pub fn repo_id(&self) -> RepositoryId {
        self.manager.repo_id()
    }

    /// The repo name of the repo being derived.
    pub fn repo_name(&self) -> &str {
        self.manager.repo_name()
    }

    /// The blobstore that should be used for storing and retrieving blobs.
    pub fn blobstore(&self) -> &Arc<dyn Blobstore> {
        match &self.blobstore_write_cache {
            Some((blobstore, _)) => blobstore,
            None => &self.blobstore,
        }
    }

    pub fn bonsai_hg_mapping(&self) -> Result<&dyn BonsaiHgMapping> {
        self.manager.bonsai_hg_mapping()
    }

    pub fn filenodes(&self) -> Result<&dyn Filenodes> {
        self.manager.filenodes()
    }

    /// The config that should be used for derivation.
    pub fn config(&self) -> &DerivedDataTypesConfig {
        self.manager.config()
    }

    /// Mapping key prefix for a particular derived data type.
    pub fn mapping_key_prefix<Derivable>(&self) -> &str
    where
        Derivable: BonsaiDerivable,
    {
        self.config()
            .mapping_key_prefixes
            .get(Derivable::NAME)
            .map_or("", String::as_str)
    }

    pub(crate) fn needs_rederive<Derivable>(&self, csid: ChangesetId) -> bool
    where
        Derivable: BonsaiDerivable,
    {
        if let Some(rederivation) = self.rederivation.as_ref() {
            rederivation.needs_rederive(Derivable::NAME, csid) == Some(true)
        } else {
            false
        }
    }

    pub(crate) fn mark_derived<Derivable>(&self, csid: ChangesetId)
    where
        Derivable: BonsaiDerivable,
    {
        if let Some(rederivation) = self.rederivation.as_ref() {
            rederivation.mark_derived(Derivable::NAME, csid);
        }
    }

    /// Enable write batching for this derivation context.
    ///
    /// With write batching enabled, blobstore writes are sent to a write
    /// cache, rather than directly to the blobstore.  They must be flushed by
    /// a call to `flush` to make them persistent.
    pub(crate) fn enable_write_batching(&mut self) {
        if self.blobstore_write_cache.is_none() {
            let blobstore = Arc::new(MemWritesBlobstore::new(self.blobstore.clone()));
            self.blobstore_write_cache = Some((blobstore.clone(), blobstore));
        }
    }

    /// Flush any pending writes for this derivation context.
    pub(crate) async fn flush(&self, ctx: &CoreContext) -> Result<()> {
        if let Some((_, blobstore)) = &self.blobstore_write_cache {
            blobstore.persist(ctx).await?;
        }
        Ok(())
    }
}


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

use std::any::TypeId;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Debug;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use derived_data_service_if::types::DerivedData;
use futures::future::try_join;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use strum::AsRefStr;
use strum::Display;
use strum::EnumString;

use crate::context::DerivationContext;

/// Enum which consolidates all available derived data types
/// It provides access to `const &'static str` representation to
/// use as Name of the derived data type, which is used to
/// identify or name data (for example lease keys) associated with this
/// particular derived data type.
/// It also provides `as_ref()` method for serialization.
/// and implements FromStr trait for deserialization.
#[derive(AsRefStr, EnumString, Display, Debug, Clone, PartialEq, Eq, Hash)]
pub enum DerivableType {
    BlameV2,
    Bssm,
    ChangesetInfo,
    DeletedManifests,
    Fastlog,
    FileNodes,
    Fsnodes,
    HgChangesets,
    GitTree,
    SkeletonManifests,
    Unodes,
}

impl DerivableType {
    const fn name(&self) -> &'static str {
        match self {
            DerivableType::BlameV2 => "blame",
            DerivableType::Bssm => "bssm",
            DerivableType::ChangesetInfo => "changeset_info",
            DerivableType::DeletedManifests => "deleted_manifest",
            DerivableType::Fastlog => "fastlog",
            DerivableType::FileNodes => "filenodes",
            DerivableType::Fsnodes => "fsnodes",
            DerivableType::HgChangesets => "hgchangesets",
            DerivableType::GitTree => "git_trees",
            DerivableType::SkeletonManifests => "skeleton_manifests",
            DerivableType::Unodes => "unodes",
        }
    }
}

/// Defines how derivation occurs.  Each derived data type must implement
/// `BonsaiDerivable` to describe how to derive a new value from its inputs
/// and how to persist mappings from changesets to derived data.
///
/// As a performance enhancement, a derived data type may also implement the
/// `derive_batch` method of this trait to implement a fast path for deriving
/// data from a batch of changesets.  The default implementation derives each
/// changeset in the batch sequentially.
///
/// As a performance enhancement, a derived data type may also implement the
/// `fetch_batch` method of this trait to implement a fast path for fetching
/// previously-derived and persisted values for a batch of changesets.  The
/// default implementation fetches in separate requests.
#[async_trait]
pub trait BonsaiDerivable: Sized + Send + Sync + Clone + Debug + 'static {
    const VARIANT: DerivableType;
    const NAME: &'static str = Self::VARIANT.name();

    /// Types of derived data types on which this derived data type
    /// depends.
    ///
    /// When performing derivation or backfilling, the derived data manager
    /// will ensure all dependencies are either derived or have been derived
    /// before deriving this derived data type.
    ///
    /// Use the `dependencies!` macro to populate this type.
    type Dependencies: DerivationDependencies;

    /// Derive data for a single changeset.
    ///
    /// If the implementation generates any other data (e.g. manifest nodes
    /// when deriving a manifest root), then it must ensure that all of this
    /// data has been written to its backing store before returning.
    async fn derive_single(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsai: BonsaiChangeset,
        parents: Vec<Self>,
    ) -> Result<Self>;

    /// Derive data for a batch of changesets.
    ///
    /// This method may be overridden by BonsaiDerivable implementors if
    /// there's a more efficient way to derive a batch of changesets.
    ///
    /// If a gap size is provided, then implementations may choose to
    /// derive only a subset of commits, as long as the gap between
    /// derived commits does not exceed the gap size.
    ///
    /// Not all implementations support gapped derivation, so some
    /// implementations may still derive all items in the batch.  The default
    /// implementation does not support gapped derivation.
    ///
    /// The batch is provided in topological order (ancestors appear before
    /// descendants).
    ///
    /// This method is called after the parents of the roots of the batch of
    /// changesets have been derived, as well as all the data for any of
    /// the dependent derived data types.  These data can be fetched from
    /// the derivation context.
    async fn derive_batch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsais: Vec<BonsaiChangeset>,
        _gap_size: Option<usize>,
    ) -> Result<HashMap<ChangesetId, Self>> {
        let mut res: HashMap<ChangesetId, Self> = HashMap::new();
        // The default implementation must derive sequentially with no
        // parallelism or concurrency, as dependencies between changesets may
        // cause O(n^2) derivations.
        for bonsai in bonsais {
            let csid = bonsai.get_changeset_id();
            let parents = derivation_ctx
                .fetch_unknown_parents(ctx, Some(&res), &bonsai)
                .await?;
            let derived = Self::derive_single(ctx, derivation_ctx, bonsai, parents).await?;
            res.insert(csid, derived);
        }
        Ok(res)
    }

    /// Store this derived data as the mapped value for a given changeset.
    ///
    /// Once derivation for a particular changeset is complete, this method
    /// is called so that the mapping from the changeset to the derived
    /// data type can be persisted, allowing it to be fetched again in the
    /// future.
    ///
    /// The derived data manager will ensure that any write caches provided to
    /// `derive_single` or `derive_batch` have been flushed before calling
    /// this method, so it is safe to persistently store the mapping
    /// immediately.
    async fn store_mapping(
        self,
        ctx: &CoreContext,
        derivation: &DerivationContext,
        csid: ChangesetId,
    ) -> Result<()>;

    /// Fetch previously derived and persisted data.
    ///
    /// Returns None if the given changeset has not had derived data
    /// persisted previously.
    async fn fetch(
        ctx: &CoreContext,
        derivation: &DerivationContext,
        csid: ChangesetId,
    ) -> Result<Option<Self>>;

    /// Fetch a batch of previously derived data.
    ///
    /// This method may be overridden by BonsaiDerivable implementors if
    /// there's a more efficient way to fetch a batch of commits.
    ///
    /// Returns a hashmap containing only the commits which have been
    /// previously persisted.  Changesets for which derived data has not
    /// been previously persisted are omitted.
    async fn fetch_batch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_ids: &[ChangesetId],
    ) -> Result<HashMap<ChangesetId, Self>> {
        stream::iter(changeset_ids.iter().copied().map(|csid| async move {
            let maybe_derived = Self::fetch(ctx, derivation_ctx, csid).await?;
            Ok(maybe_derived.map(|derived| (csid, derived)))
        }))
        .buffer_unordered(64)
        .try_filter_map(|maybe_value| async move { Ok(maybe_value) })
        .try_collect()
        .await
    }

    fn from_thrift(_data: DerivedData) -> Result<Self>;

    fn into_thrift(_data: Self) -> Result<DerivedData>;
}

#[async_trait]
pub trait DerivationDependencies {
    /// Checks that all dependencies have been derived for this
    /// changeset.
    async fn check_dependencies(
        ctx: &CoreContext,
        derivation: &DerivationContext,
        csid: ChangesetId,
        visited: &mut HashSet<TypeId>,
    ) -> Result<()>;
}

#[async_trait]
impl DerivationDependencies for () {
    async fn check_dependencies(
        _ctx: &CoreContext,
        _derivation: &DerivationContext,
        _csid: ChangesetId,
        _visited: &mut HashSet<TypeId>,
    ) -> Result<()> {
        Ok(())
    }
}

#[async_trait]
impl<Derivable, Rest> DerivationDependencies for (Derivable, Rest)
where
    Derivable: BonsaiDerivable,
    Rest: DerivationDependencies + 'static,
{
    async fn check_dependencies(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        csid: ChangesetId,
        visited: &mut HashSet<TypeId>,
    ) -> Result<()> {
        let type_id = TypeId::of::<Derivable>();
        if visited.insert(type_id) {
            try_join(
                derivation_ctx.fetch_dependency::<Derivable>(ctx, csid),
                Rest::check_dependencies(ctx, derivation_ctx, csid, visited),
            )
            .await?;
            Ok(())
        } else {
            Rest::check_dependencies(ctx, derivation_ctx, csid, visited).await
        }
    }
}

#[macro_export]
macro_rules! dependencies {
    () => { () };
    ($dep:ty) => { ( $dep , () ) };
    ($dep:ty, $( $rest:tt )*) => { ( $dep , dependencies!($( $rest )*) ) };
}


/*
 * 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 mononoke_types::RepositoryId;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum DerivationError {
    #[error("Derivation of {0} is not enabled for repo={2} repoid={1}")]
    Disabled(&'static str, RepositoryId, String),
    #[error(transparent)]
    Error(#[from] Error),
}


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

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

use anyhow::Result;
use cacheblob::LeaseOps;
use context::CoreContext;
use futures::channel::oneshot;
use futures::future::FutureExt;
use slog::warn;

const LEASE_WARNING_THRESHOLD: Duration = Duration::from_secs(60);

#[derive(Clone)]
pub struct DerivedDataLease {
    lease_ops: Arc<dyn LeaseOps>,
}

impl DerivedDataLease {
    pub fn new(lease_ops: Arc<dyn LeaseOps>) -> Self {
        DerivedDataLease { lease_ops }
    }

    pub fn lease_ops(&self) -> &Arc<dyn LeaseOps> {
        &self.lease_ops
    }

    pub async fn try_acquire_in_loop<F, Fut>(
        &self,
        ctx: &CoreContext,
        key: &str,
        mut abort_fn: F,
    ) -> Result<Option<DerivedDataLeaseGuard>>
    where
        F: FnMut() -> Fut,
        Fut: Future<Output = Result<bool>>,
    {
        let mut start = Instant::now();
        let mut total_elapsed = Duration::from_secs(0);
        let mut backoff_ms = 200;
        while !self.lease_ops.try_add_put_lease(key).await? {
            if abort_fn().await? {
                return Ok(None);
            }
            let elapsed = start.elapsed();
            if elapsed > LEASE_WARNING_THRESHOLD {
                total_elapsed += elapsed;
                start = Instant::now();
                warn!(
                    ctx.logger(),
                    "Can not acquire lease {} for more than {:?}", key, total_elapsed
                );
            }
            let sleep = rand::random::<u64>() % backoff_ms;
            tokio::time::sleep(Duration::from_millis(sleep)).await;
            backoff_ms = std::cmp::min(backoff_ms * 2, 1000);
        }
        let (sender, receiver) = oneshot::channel();
        self.lease_ops
            .renew_lease_until(ctx.clone(), key, receiver.map(|_| ()).boxed());
        Ok(Some(DerivedDataLeaseGuard {
            sender: Some(sender),
        }))
    }
}

/// Guard representing an active lease.  We stop renewing the lease when
/// the guard is dropped.
pub struct DerivedDataLeaseGuard {
    sender: Option<oneshot::Sender<()>>,
}

impl Drop for DerivedDataLeaseGuard {
    fn drop(&mut self) {
        if let Some(sender) = self.sender.take() {
            let _ = sender.send(());
        }
    }
}


/*
 * 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 mod context;
pub mod derivable;
pub mod error;
pub mod lease;
pub mod manager;

pub use self::context::DerivationContext;
pub use self::derivable::BonsaiDerivable;
pub use self::derivable::DerivableType;
pub use self::error::DerivationError;
pub use self::lease::DerivedDataLease;
pub use self::manager::derive::BatchDeriveOptions;
pub use self::manager::derive::BatchDeriveStats;
pub use self::manager::derive::Rederivation;
pub use self::manager::DerivedDataManager;


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

use std::sync::Arc;

use anyhow::Context;
use anyhow::Result;
use bonsai_hg_mapping::BonsaiHgMapping;
use cacheblob::LeaseOps;
use changesets::Changesets;
use commit_graph::CommitGraph;
use context::CoreContext;
use derived_data_remote::DerivationClient;
use filenodes::Filenodes;
use metaconfig_types::DerivedDataTypesConfig;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use repo_blobstore::RepoBlobstore;
use scuba_ext::MononokeScubaSampleBuilder;

use crate::lease::DerivedDataLease;

pub mod bubble;
pub mod derive;
pub mod logging;
pub mod util;

/// Manager for derived data.
///
/// The manager is responsible for ordering derivation of data based
/// on the dependencies between derived data types and the changeset
/// graph.
#[derive(Clone)]
pub struct DerivedDataManager {
    inner: Arc<DerivedDataManagerInner>,
}

#[derive(Clone)]
pub struct DerivedDataManagerInner {
    repo_id: RepositoryId,
    repo_name: String,
    changesets: Arc<dyn Changesets>,
    commit_graph: Arc<CommitGraph>,
    bonsai_hg_mapping: Option<Arc<dyn BonsaiHgMapping>>,
    filenodes: Option<Arc<dyn Filenodes>>,
    repo_blobstore: RepoBlobstore,
    lease: DerivedDataLease,
    scuba: MononokeScubaSampleBuilder,
    config_name: String,
    config: DerivedDataTypesConfig,
    /// If a (primary) manager has a secondary manager, that means some of the
    /// changesets should be derived using the primary manager, and some the secondary,
    /// in that order. For example, bubble managers are secondary, as all the data in
    /// the persistent blobstore must be derived BEFORE deriving data in the bubble.
    secondary: Option<SecondaryManagerData>,
    /// If this client is set, then derivation will be done remotely on derived data service
    derivation_service_client: Option<Arc<dyn DerivationClient>>,
}

pub struct DerivationAssignment {
    /// Changesets that should be derived by the primary manager
    pub primary: Vec<ChangesetId>,
    /// Changesets that should be derived by the secondary manager, after the first
    /// part of derivation is done.
    pub secondary: Vec<ChangesetId>,
}

#[async_trait::async_trait]
pub trait DerivationAssigner: Send + Sync {
    /// How to split derivation between primary and secondary managers. If not possible
    /// to split, this function should error.
    async fn assign(&self, ctx: &CoreContext, cs: Vec<ChangesetId>)
    -> Result<DerivationAssignment>;
}

#[derive(Clone)]
pub(crate) struct SecondaryManagerData {
    manager: DerivedDataManager,
    assigner: Arc<dyn DerivationAssigner>,
}

impl DerivedDataManager {
    pub fn new(
        repo_id: RepositoryId,
        repo_name: String,
        changesets: Arc<dyn Changesets>,
        commit_graph: Arc<CommitGraph>,
        bonsai_hg_mapping: Arc<dyn BonsaiHgMapping>,
        filenodes: Arc<dyn Filenodes>,
        repo_blobstore: RepoBlobstore,
        lease: Arc<dyn LeaseOps>,
        scuba: MononokeScubaSampleBuilder,
        config_name: String,
        config: DerivedDataTypesConfig,
        derivation_service_client: Option<Arc<dyn DerivationClient>>,
    ) -> Self {
        let lease = DerivedDataLease::new(lease);
        DerivedDataManager {
            inner: Arc::new(DerivedDataManagerInner {
                repo_id,
                repo_name,
                config_name,
                config,
                changesets,
                commit_graph,
                bonsai_hg_mapping: Some(bonsai_hg_mapping),
                filenodes: Some(filenodes),
                repo_blobstore,
                lease,
                scuba,
                secondary: None,
                derivation_service_client,
            }),
        }
    }

    // For dangerous-override: allow replacement of lease-ops
    pub fn with_replaced_lease(&self, lease: Arc<dyn LeaseOps>) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                lease: DerivedDataLease::new(lease),
                ..self.inner.as_ref().clone()
            }),
        }
    }

    // For dangerous-override: allow replacement of blobstore
    pub fn with_replaced_blobstore(&self, repo_blobstore: RepoBlobstore) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                repo_blobstore,
                ..self.inner.as_ref().clone()
            }),
        }
    }

    // For dangerous-override: allow replacement of changesets
    pub fn with_replaced_changesets(&self, changesets: Arc<dyn Changesets>) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                changesets,
                ..self.inner.as_ref().clone()
            }),
        }
    }

    // For dangerous-override: allow replacement of bonsai-hg-mapping
    pub fn with_replaced_bonsai_hg_mapping(
        &self,
        bonsai_hg_mapping: Arc<dyn BonsaiHgMapping>,
    ) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                bonsai_hg_mapping: Some(bonsai_hg_mapping),
                ..self.inner.as_ref().clone()
            }),
        }
    }

    // For dangerous-override: allow replacement of filenodes
    pub fn with_replaced_filenodes(&self, filenodes: Arc<dyn Filenodes>) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                filenodes: Some(filenodes),
                ..self.inner.as_ref().clone()
            }),
        }
    }

    pub fn with_replaced_config(
        &self,
        config_name: String,
        config: DerivedDataTypesConfig,
    ) -> Self {
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                config_name,
                config,
                ..self.inner.as_ref().clone()
            }),
        }
    }

    pub fn repo_id(&self) -> RepositoryId {
        self.inner.repo_id
    }

    pub fn repo_name(&self) -> &str {
        self.inner.repo_name.as_str()
    }

    pub fn changesets(&self) -> &dyn Changesets {
        self.inner.changesets.as_ref()
    }

    pub fn changesets_arc(&self) -> Arc<dyn Changesets> {
        self.inner.changesets.clone()
    }

    pub fn commit_graph(&self) -> &CommitGraph {
        self.inner.commit_graph.as_ref()
    }

    pub fn commit_graph_arc(&self) -> Arc<CommitGraph> {
        self.inner.commit_graph.clone()
    }

    pub fn repo_blobstore(&self) -> &RepoBlobstore {
        &self.inner.repo_blobstore
    }

    pub fn lease(&self) -> &DerivedDataLease {
        &self.inner.lease
    }

    pub fn scuba(&self) -> &MononokeScubaSampleBuilder {
        &self.inner.scuba
    }

    pub fn config(&self) -> &DerivedDataTypesConfig {
        &self.inner.config
    }

    pub fn config_name(&self) -> String {
        self.inner.config_name.clone()
    }

    pub fn bonsai_hg_mapping(&self) -> Result<&dyn BonsaiHgMapping> {
        self.inner
            .bonsai_hg_mapping
            .as_deref()
            .context("Missing BonsaiHgMapping")
    }

    pub fn filenodes(&self) -> Result<&dyn Filenodes> {
        self.inner.filenodes.as_deref().context("Missing filenodes")
    }

    pub fn derivation_service_client(&self) -> Option<&dyn DerivationClient> {
        self.inner.derivation_service_client.as_deref()
    }
}


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

use std::sync::Arc;

use changesets::ChangesetsArc;
use context::CoreContext;
use ephemeral_blobstore::Bubble;
use ephemeral_blobstore::EphemeralChangesets;
use mononoke_types::ChangesetId;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;

use super::DerivationAssigner;
use super::DerivationAssignment;
use super::DerivedDataManager;
use super::DerivedDataManagerInner;
use super::SecondaryManagerData;

struct BubbleAssigner {
    changesets: Arc<EphemeralChangesets>,
}

#[async_trait::async_trait]
impl DerivationAssigner for BubbleAssigner {
    async fn assign(
        &self,
        _ctx: &CoreContext,
        cs: Vec<ChangesetId>,
    ) -> anyhow::Result<DerivationAssignment> {
        let in_bubble = self.changesets.fetch_gens(&cs).await?;
        let (in_bubble, not_in_bubble) = cs
            .into_iter()
            .partition(|cs_id| in_bubble.contains_key(cs_id));
        Ok(DerivationAssignment {
            primary: not_in_bubble,
            secondary: in_bubble,
        })
    }
}

impl DerivedDataManager {
    pub fn for_bubble(
        self,
        bubble: Bubble,
        // Perhaps this can be fetched from inside the manager in the future
        container: &(impl ChangesetsArc + RepoIdentityRef + RepoBlobstoreRef),
    ) -> Self {
        let changesets = Arc::new(bubble.changesets(container));
        Self {
            inner: Arc::new(DerivedDataManagerInner {
                secondary: Some(SecondaryManagerData {
                    manager: Self {
                        inner: Arc::new(DerivedDataManagerInner {
                            changesets: changesets.clone(),
                            repo_blobstore: bubble
                                .wrap_repo_blobstore(self.inner.repo_blobstore.clone()),
                            filenodes: None,
                            bonsai_hg_mapping: None,
                            ..self.inner.as_ref().clone()
                        }),
                    },
                    assigner: Arc::new(BubbleAssigner { changesets }),
                }),
                ..self.inner.as_ref().clone()
            }),
        }
    }
}


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

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

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_recursion::async_recursion;
use blobstore::Loadable;
use borrowed::borrowed;
use cloned::cloned;
use context::CoreContext;
use derived_data_service_if::DerivationType;
use derived_data_service_if::DeriveRequest;
use derived_data_service_if::DeriveResponse;
use derived_data_service_if::DeriveUnderived;
use derived_data_service_if::DerivedDataType;
use derived_data_service_if::RequestStatus;
use futures::future::try_join;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::join;
use futures::select_biased;
use futures::stream;
use futures::stream::FuturesUnordered;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_stats::TimedFutureExt;
use futures_stats::TimedTryFutureExt;
use mononoke_types::ChangesetId;
use slog::debug;
use topo_sort::TopoSortedDagTraversal;
use tunables::get_duration_from_tunable_or;
use tunables::MononokeTunables;

use super::DerivationAssignment;
use super::DerivedDataManager;
use crate::context::DerivationContext;
use crate::derivable::BonsaiDerivable;
use crate::derivable::DerivationDependencies;
use crate::error::DerivationError;
use crate::manager::util::DiscoveryStats;

#[derive(Clone, Copy)]
pub enum BatchDeriveOptions {
    Parallel { gap_size: Option<usize> },
    Serial,
}

#[derive(Debug)]
pub enum BatchDeriveStats {
    Parallel(Duration),
    Serial(Vec<(ChangesetId, Duration)>),
}

impl BatchDeriveStats {
    fn append(self, other: Self) -> anyhow::Result<Self> {
        use BatchDeriveStats::*;
        Ok(match (self, other) {
            (Parallel(d1), Parallel(d2)) => Parallel(d1 + d2),
            (Serial(mut s1), Serial(mut s2)) => {
                s1.append(&mut s2);
                Serial(s1)
            }
            _ => anyhow::bail!("Incompatible stats"),
        })
    }
}

/// Trait to allow determination of rederivation.
pub trait Rederivation: Send + Sync + 'static {
    /// Determine whether a changeset needs rederivation of
    /// a particular derived data type.
    ///
    /// If this function returns `None`, then it will only be
    /// derived if it isn't already derived.
    fn needs_rederive(&self, derivable_name: &str, csid: ChangesetId) -> Option<bool>;

    /// Marks a changeset as having been derived.  After this
    /// is called, `needs_rederive` should not return `true` for
    /// this changeset.
    fn mark_derived(&self, derivable_name: &str, csid: ChangesetId);
}

impl DerivedDataManager {
    #[async_recursion]
    /// Returns the appropriate manager to derive given changeset, either this
    /// manager, or some secondary manager in the chain.
    async fn get_manager(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> anyhow::Result<&DerivedDataManager> {
        Ok(if let Some(secondary) = &self.inner.secondary {
            if secondary
                .assigner
                .assign(ctx, vec![cs_id])
                .await?
                .secondary
                .is_empty()
            {
                self
            } else {
                secondary.manager.get_manager(ctx, cs_id).await?
            }
        } else {
            self
        })
    }

    pub fn derivation_context(
        &self,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> DerivationContext {
        DerivationContext::new(self.clone(), rederivation, self.repo_blobstore().boxed())
    }

    pub async fn check_derived<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
    ) -> Result<(), DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        if self
            .fetch_derived::<Derivable>(ctx, csid, None)
            .await?
            .is_none()
        {
            return Err(
                anyhow!("expected {} already derived for {}", Derivable::NAME, csid).into(),
            );
        }
        Ok(())
    }

    /// Perform derivation for a single changeset.
    /// Will fail in case data for parents changeset wasn't derived
    async fn perform_single_derivation<Derivable>(
        &self,
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        csid: ChangesetId,
        discovery_stats: &Option<DiscoveryStats>,
    ) -> Result<(ChangesetId, Derivable)>
    where
        Derivable: BonsaiDerivable,
    {
        let mut scuba = ctx.scuba().clone();
        scuba
            .add("changeset_id", csid.to_string())
            .add("derived_data_type", Derivable::NAME);
        scuba
            .clone()
            .log_with_msg("Waiting for derived data to be generated", None);

        debug!(ctx.logger(), "derive {} for {}", Derivable::NAME, csid);
        let lease_key = format!("repo{}.{}.{}", self.repo_id(), Derivable::NAME, csid);

        let ctx = ctx.clone_and_reset();

        let (stats, result) = async {
            let bonsai = csid.load(&ctx, self.repo_blobstore()).map_err(Error::from);
            let guard = async {
                if derivation_ctx.needs_rederive::<Derivable>(csid) {
                    // We are rederiving this changeset, so do not try to take
                    // the lease, as doing so will drop out immediately
                    // because the data is already derived.
                    None
                } else {
                    Some(
                        self.lease()
                            .try_acquire_in_loop(&ctx, &lease_key, || async {
                                Ok(Derivable::fetch(&ctx, derivation_ctx, csid)
                                    .await?
                                    .is_some())
                            })
                            .await,
                    )
                }
            };
            let (bonsai, guard) = join!(bonsai, guard);
            if matches!(guard, Some(Ok(None))) {
                // Something else completed derivation
                let derived = Derivable::fetch(&ctx, derivation_ctx, csid)
                    .await?
                    .ok_or_else(|| {
                        anyhow!("derivation completed elsewhere but data could not be fetched")
                    })?;
                Ok((csid, derived))
            } else {
                // We must perform derivation.  Use the appropriate session
                // class for derivation.
                let ctx = self.set_derivation_session_class(ctx.clone());

                // The derivation process is additionally logged to the derived
                // data scuba table.
                let mut derived_data_scuba = self.derived_data_scuba::<Derivable>(discovery_stats);
                derived_data_scuba.add("changeset", csid.to_string());
                self.log_derivation_start::<Derivable>(&ctx, &mut derived_data_scuba, csid);

                let (derive_stats, derived) = async {
                    let bonsai = bonsai?;
                    let parents = derivation_ctx.fetch_parents(&ctx, &bonsai).await?;
                    Derivable::derive_single(&ctx, derivation_ctx, bonsai, parents).await
                }
                .timed()
                .await;

                self.log_derivation_end::<Derivable>(
                    &ctx,
                    &mut derived_data_scuba,
                    csid,
                    &derive_stats,
                    derived.as_ref().err(),
                );

                let derived = derived?;

                // We may now store the mapping, and flush the blobstore to
                // ensure the mapping is persisted.
                let (persist_stats, persisted) = derived
                    .clone()
                    .store_mapping(&ctx, derivation_ctx, csid)
                    .timed()
                    .await;

                self.log_mapping_insertion(
                    &ctx,
                    &mut derived_data_scuba,
                    Some(&derived),
                    &persist_stats,
                    persisted.as_ref().err(),
                );

                persisted?;

                Ok((csid, derived))
            }
        }
        .timed()
        .await;
        scuba.add_future_stats(&stats);
        if result.is_ok() {
            scuba.log_with_msg("Got derived data", None);
        } else {
            scuba.log_with_msg("Failed to get derived data", None);
        };
        result
    }

    /// Find ancestors of the target changeset that are underived.
    async fn find_underived_inner<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        limit: Option<u64>,
        derivation_ctx: &DerivationContext,
    ) -> Result<HashMap<ChangesetId, Vec<ChangesetId>>>
    where
        Derivable: BonsaiDerivable,
    {
        let underived_commits_parents: HashMap<ChangesetId, Vec<ChangesetId>> =
            bounded_traversal::bounded_traversal_dag_limited(
                    100,
                    csid,
                    move |csid: ChangesetId| {
                        async move {
                            if derivation_ctx
                                .fetch_derived::<Derivable>(ctx, csid)
                                .await?
                                .is_some()
                            {
                                Ok((None, Vec::new()))
                            } else {
                                let parents = self
                                    .changesets()
                                    .get(ctx, csid)
                                    .await?
                                    .ok_or_else(|| anyhow!("changeset not found: {}", csid))?
                                    .parents;
                                Ok((Some((csid, parents.clone())), parents))
                            }
                        }
                        .boxed()
                    },
                    move |out, results: bounded_traversal::Iter<HashMap<ChangesetId, Vec<ChangesetId>>>| {
                        async move {
                            anyhow::Ok(results
                                .chain(std::iter::once(out.into_iter().collect()))
                                .reduce(|mut acc, item| {
                                    acc.extend(item);
                                    acc
                                })
                                .unwrap_or_else(HashMap::new))
                        }
                        .boxed()
                    },
                    limit,
                )
                    .await?
                    // If we visited no nodes, then we want an empty hashmap
                    .unwrap_or_else(HashMap::new);

        // Remove parents that have already been derived.
        let underived_commits_parents = underived_commits_parents
            .iter()
            .map(|(csid, parents)| {
                let parents = parents
                    .iter()
                    .filter(|p| underived_commits_parents.contains_key(p))
                    .cloned()
                    .collect::<Vec<_>>();
                (*csid, parents)
            })
            .collect::<HashMap<_, _>>();

        Ok(underived_commits_parents)
    }

    /// Find which ancestors of `csid` are not yet derived, and necessary for
    /// the derivation of `csid` to complete, and derive them.
    async fn derive_underived<Derivable>(
        &self,
        ctx: &CoreContext,
        derivation_ctx: Arc<DerivationContext>,
        target_csid: ChangesetId,
    ) -> Result<DerivationOutcome<Derivable>, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        let (find_underived_stats, dag_traversal) = async {
            self.find_underived_inner::<Derivable>(ctx, target_csid, None, derivation_ctx.as_ref())
                .await
                .context("Finding underived commits")
        }
        .try_timed()
        .await?;

        let stats = Some(DiscoveryStats {
            find_underived_completion_time: find_underived_stats.completion_time,
            commits_discovered: dag_traversal.len() as u32,
        });
        let mut dag_traversal = TopoSortedDagTraversal::new(dag_traversal);

        let buffer_size = self.max_parallel_derivations();
        let mut derivations = FuturesUnordered::new();
        let mut completed_count = 0;
        let mut target_derived = None;
        while !dag_traversal.is_empty() || !derivations.is_empty() {
            let free = buffer_size.saturating_sub(derivations.len());
            derivations.extend(dag_traversal.drain(free).map(|csid| {
                cloned!(ctx, derivation_ctx);
                let manager = self.clone();
                let stats = stats.clone();
                let derivation = async move {
                    manager
                        .perform_single_derivation(&ctx, &derivation_ctx, csid, &stats)
                        .await
                };
                tokio::spawn(derivation).map_err(Error::from)
            }));
            if let Some(derivation_result) = derivations.try_next().await? {
                let (derived_csid, derived) = derivation_result?;
                if derived_csid == target_csid {
                    target_derived = Some(derived);
                }
                dag_traversal.visited(derived_csid);
                completed_count += 1;
                derivation_ctx.mark_derived::<Derivable>(derived_csid);
            }
        }

        let derived = match target_derived {
            Some(derived) => derived,
            None => {
                // We didn't find the derived data during derivation, as
                // possibly it was already derived, so just try to fetch it.
                derivation_ctx
                    .fetch_derived(ctx, target_csid)
                    .await?
                    .ok_or_else(|| anyhow!("failed to derive target"))?
            }
        };

        Ok(DerivationOutcome {
            derived,
            count: completed_count,
            find_underived_time: find_underived_stats.completion_time,
        })
    }

    /// Count how many ancestors of `csid` are not yet derived.
    pub async fn count_underived<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        limit: Option<u64>,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<u64, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.get_manager(ctx, csid)
            .await?
            .count_underived_impl::<Derivable>(ctx, csid, limit, rederivation)
            .await
    }

    async fn count_underived_impl<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        limit: Option<u64>,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<u64, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.check_enabled::<Derivable>()?;
        let derivation_ctx = self.derivation_context(rederivation);
        let underived = self
            .find_underived_inner::<Derivable>(ctx, csid, limit, &derivation_ctx)
            .await?;
        Ok(underived.len() as u64)
    }

    /// Find which ancestors of `csid` are not yet derived.
    ///
    /// Searches backwards looking for the most recent ancestors which have
    /// been derived, and returns all of their descendants up to the target
    /// changeset.
    ///
    /// Note that gapped derivation may mean that some of the ancestors
    /// of those changesets may also be underived.  These changesets are not
    /// necessary to derive data for the target changeset, and so will
    /// not be included.
    ///
    /// Returns a map of underived changesets to their underived parents,
    /// suitable for input to toposort.
    pub async fn find_underived<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        limit: Option<u64>,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<HashMap<ChangesetId, Vec<ChangesetId>>>
    where
        Derivable: BonsaiDerivable,
    {
        self.get_manager(ctx, csid)
            .await?
            .find_underived_impl::<Derivable>(ctx, csid, limit, rederivation)
            .await
    }

    async fn find_underived_impl<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        limit: Option<u64>,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<HashMap<ChangesetId, Vec<ChangesetId>>>
    where
        Derivable: BonsaiDerivable,
    {
        self.check_enabled::<Derivable>()?;
        let derivation_ctx = self.derivation_context(rederivation);
        self.find_underived_inner::<Derivable>(ctx, csid, limit, &derivation_ctx)
            .await
    }

    /// Derive or retrieve derived data for a changeset.
    pub async fn derive<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<Derivable, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        const RETRY_DELAY_MS: u64 = 100;
        const RETRY_ATTEMPTS_LIMIT: u8 = 10;
        const FALLBACK_TIMEOUT_SECS: u64 = 15;
        if let Some(client) = self.derivation_service_client() {
            let mut attempt = 0;
            let started = Instant::now();
            let fallback_timeout = get_duration_from_tunable_or(
                MononokeTunables::remote_derivation_fallback_timeout_secs,
                Duration::from_secs,
                FALLBACK_TIMEOUT_SECS,
            );
            let retry_delay = get_duration_from_tunable_or(
                MononokeTunables::derivation_request_retry_delay,
                Duration::from_millis,
                RETRY_DELAY_MS,
            );
            let request = DeriveRequest {
                repo_name: self.repo_name().to_string(),
                derived_data_type: DerivedDataType {
                    type_name: Derivable::NAME.to_string(),
                },
                changeset_id: csid.as_ref().to_vec(),
                config_name: self.config_name(),
                derivation_type: DerivationType::derive_underived(DeriveUnderived {}),
            };
            let mut request_state = DerivationState::NotRequested;
            while let Some(true) =
                tunables::tunables().by_repo_enable_remote_derivation(self.repo_name())
            {
                if started.elapsed() >= fallback_timeout {
                    self.derived_data_scuba::<Derivable>(&None)
                        .add("changeset", csid.to_string())
                        .log_with_msg(
                            "Derived data service failed",
                            format!(
                                "Fallback to local derivation after timeout {:?}",
                                fallback_timeout
                            ),
                        );
                    break;
                }

                let res = match request_state {
                    DerivationState::NotRequested => {
                        // return if already derived
                        if let Some(data) =
                            self.fetch_derived(ctx, csid, rederivation.clone()).await?
                        {
                            return Ok(data);
                        }
                        // not yet derived, so request derivation
                        self.derived_data_scuba::<Derivable>(&None)
                            .add("changeset", csid.to_string())
                            .log_with_msg("Requesting remote derivation", None);
                        client.derive_remotely(&request).await
                    }
                    DerivationState::InProgress => client.poll(&request).await,
                };

                match res {
                    Ok(DeriveResponse { data, status }) => match (status, data) {
                        // Derivation was requested, set state InProgress and wait
                        (RequestStatus::IN_PROGRESS, _) => {
                            request_state = DerivationState::InProgress;
                            tokio::time::sleep(retry_delay).await
                        }
                        // Derivation succeeded, return
                        (RequestStatus::SUCCESS, Some(data)) => {
                            self.derived_data_scuba::<Derivable>(&None)
                                .add("changeset", csid.to_string())
                                .log_with_msg("Remote derivation finished", None);
                            return Ok(Derivable::from_thrift(data)?);
                        }
                        // Either data was already derived or wasn't requested
                        // wait before requesting again
                        (RequestStatus::DOES_NOT_EXIST, _) => {
                            request_state = DerivationState::NotRequested;
                            tokio::time::sleep(retry_delay).await
                        }
                        // should not happen, reported success but data wasn't derived
                        (RequestStatus::SUCCESS, None) => {
                            self.derived_data_scuba::<Derivable>(&None)
                                .add("changeset", csid.to_string())
                                .log_with_msg(
                                    "Derived data service failed",
                                    "Request succeeded but derived data is None".to_string(),
                                );
                            break;
                        }
                        (RequestStatus(n), _) => {
                            self.derived_data_scuba::<Derivable>(&None)
                                .add("changeset", csid.to_string())
                                .log_with_msg(
                                    "Derived data service failed",
                                    format!("Response with unknown state: {n}"),
                                );
                            break;
                        }
                    },
                    Err(e) => {
                        if attempt >= RETRY_ATTEMPTS_LIMIT {
                            self.derived_data_scuba::<Derivable>(&None)
                                .add("changeset", csid.to_string())
                                .log_with_msg("Derived data service failed", format!("{:#}", e));
                            break;
                        }
                        attempt += 1;
                        tokio::time::sleep(retry_delay).await;
                    }
                }
            }
        }

        self.derive_locally(ctx, csid, rederivation).await
    }

    pub async fn derive_locally<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<Derivable, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.get_manager(ctx, csid)
            .await?
            .derive_impl::<Derivable>(ctx, csid, rederivation)
            .await
    }

    async fn derive_impl<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<Derivable, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.check_enabled::<Derivable>()?;
        let derivation_ctx = self.derivation_context(rederivation);

        let pc = ctx.clone().fork_perf_counters();

        select_biased! {
            _ = derivation_disabled_watcher(self.repo_name(), Derivable::NAME).fuse() =>
            // Derivation was disabled during the derivation process.
            Err(DerivationError::Disabled(
                Derivable::NAME,
                self.repo_id(),
                self.repo_name().to_string(),
            )),
            (stats, res) = self.derive_underived(ctx, Arc::new(derivation_ctx), csid).timed().fuse() => {
                if self.should_log_slow_derivation(stats.completion_time) {
                    self.log_slow_derivation(ctx, csid, &stats, &pc, &res);
                }
            res.map(|r| r.derived)
            }
        }
    }

    #[async_recursion]
    /// Derive data for exactly a batch of changesets.
    ///
    /// The provided batch of changesets must be in topological
    /// order.
    ///
    /// The difference between "derive_exactly" and "derive", is that for
    /// deriving exactly, the caller must have arranged for the dependencies
    /// and ancestors of the batch to have already been derived.  If
    /// any dependency or ancestor is not already derived, an error
    /// will be returned.
    pub async fn derive_exactly_batch<Derivable>(
        &self,
        ctx: &CoreContext,
        csids: Vec<ChangesetId>,
        batch_options: BatchDeriveOptions,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<BatchDeriveStats, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        let (csids, secondary_derivation) = if let Some(secondary_data) = &self.inner.secondary {
            let DerivationAssignment { primary, secondary } =
                secondary_data.assigner.assign(ctx, csids).await?;
            (primary, {
                cloned!(rederivation);
                async move {
                    secondary_data
                        .manager
                        .derive_exactly_batch::<Derivable>(
                            ctx,
                            secondary,
                            batch_options,
                            rederivation,
                        )
                        .await
                }
                .left_future()
            })
        } else {
            (
                csids,
                future::ready(Ok(match batch_options {
                    BatchDeriveOptions::Serial => BatchDeriveStats::Serial(vec![]),
                    BatchDeriveOptions::Parallel { .. } => {
                        BatchDeriveStats::Parallel(Duration::ZERO)
                    }
                }))
                .right_future(),
            )
        };
        self.check_enabled::<Derivable>()?;
        let mut derivation_ctx = self.derivation_context(rederivation.clone());

        // Enable write batching, so that writes are stored in memory
        // before being flushed.
        derivation_ctx.enable_write_batching();
        let derivation_ctx_ref = &derivation_ctx;

        let mut scuba = ctx.scuba().clone();
        scuba
            .add("stack_size", csids.len())
            .add("derived_data", Derivable::NAME);
        if let (Some(first), Some(last)) = (csids.first(), csids.last()) {
            scuba
                .add("first_csid", first.to_string())
                .add("last_csid", last.to_string());
        }

        // Load all of the bonsais for this batch.
        let bonsais = stream::iter(csids.into_iter().map(|csid| async move {
            let bonsai = csid.load(ctx, derivation_ctx_ref.blobstore()).await?;
            Ok::<_, Error>(bonsai)
        }))
        .buffered(100)
        .try_collect::<Vec<_>>()
        .await?;

        // Dependency checks: check topological order and determine heads
        // and highest ancestors of the batch.
        let mut seen = HashSet::new();
        let mut heads = HashSet::new();
        let mut ancestors = HashSet::new();
        for bonsai in bonsais.iter() {
            let csid = bonsai.get_changeset_id();
            if ancestors.contains(&csid) {
                return Err(anyhow!("batch not in topological order at {}", csid).into());
            }
            for parent in bonsai.parents() {
                if !seen.contains(&parent) {
                    ancestors.insert(parent);
                }
                heads.remove(&parent);
            }
            seen.insert(csid);
            heads.insert(csid);
        }

        // Dependency checks: all ancestors should have this derived
        // data type derived
        let ancestor_checks = async move {
            stream::iter(ancestors)
                .map(|csid| derivation_ctx_ref.fetch_dependency::<Derivable>(ctx, csid))
                .buffered(100)
                .try_for_each(|_| async { Ok(()) })
                .await
                .with_context(|| {
                    format!(
                        "a batch ancestor does not have '{}' derived",
                        Derivable::NAME
                    )
                })
        };

        // Dependency checks: all heads should have their dependent
        // data types derived.
        let dependency_checks = async move {
            stream::iter(heads)
                .map(|csid| async move {
                    Derivable::Dependencies::check_dependencies(
                        ctx,
                        derivation_ctx_ref,
                        csid,
                        &mut HashSet::new(),
                    )
                    .await
                })
                .buffered(100)
                .try_for_each(|_| async { Ok(()) })
                .await
                .context("a batch dependency has not been derived")
        };

        try_join(ancestor_checks, dependency_checks)
            .await
            .context(concat!(
                "derive exactly batch pre-condition not satisfied: ",
                "all ancestors' and dependencies' data must already have been derived",
            ))?;

        let ctx = ctx.clone_and_reset();
        let ctx = self.set_derivation_session_class(ctx.clone());
        borrowed!(ctx);

        let csid_range = if let (Some(first), Some(last)) = (bonsais.first(), bonsais.last()) {
            let first_csid = first.get_changeset_id();
            let last_csid = last.get_changeset_id();
            debug!(
                ctx.logger(),
                "derive exactly {} batch from {} to {}",
                Derivable::NAME,
                first_csid,
                last_csid,
            );
            Some((first_csid, last_csid))
        } else {
            None
        };

        let mut derived_data_scuba = self.derived_data_scuba::<Derivable>(&None);
        derived_data_scuba.add(
            "changesets",
            bonsais
                .iter()
                .map(|bonsai| bonsai.get_changeset_id().to_string())
                .collect::<Vec<_>>(),
        );
        self.log_batch_derivation_start::<Derivable>(ctx, &mut derived_data_scuba, csid_range);
        let (overall_stats, result) = async {
            let derivation_ctx_ref = &derivation_ctx;
            let (batch_stats, derived) = match batch_options {
                BatchDeriveOptions::Parallel { gap_size } => {
                    derived_data_scuba.add("parallel", true);
                    if let Some(gap_size) = gap_size {
                        derived_data_scuba.add("gap_size", gap_size);
                    }
                    let (stats, derived) =
                        Derivable::derive_batch(ctx, derivation_ctx_ref, bonsais, gap_size)
                            .try_timed()
                            .await
                            .with_context(|| {
                                if let Some((first, last)) = csid_range {
                                    format!(
                                        "failed to derive {} batch (start:{}, end:{})",
                                        Derivable::NAME,
                                        first,
                                        last
                                    )
                                } else {
                                    format!("failed to derive empty {} batch", Derivable::NAME)
                                }
                            })?;
                    (BatchDeriveStats::Parallel(stats.completion_time), derived)
                }
                BatchDeriveOptions::Serial => {
                    derived_data_scuba.add("parallel", false);
                    let mut per_commit_stats = Vec::new();
                    let mut per_commit_derived = HashMap::new();
                    for bonsai in bonsais {
                        let csid = bonsai.get_changeset_id();
                        let parents = derivation_ctx_ref
                            .fetch_unknown_parents(ctx, Some(&per_commit_derived), &bonsai)
                            .await?;
                        let (stats, derived) =
                            Derivable::derive_single(ctx, derivation_ctx_ref, bonsai, parents)
                                .try_timed()
                                .await
                                .with_context(|| {
                                    format!("failed to derive {} for {}", Derivable::NAME, csid)
                                })?;
                        per_commit_stats.push((csid, stats.completion_time));
                        per_commit_derived.insert(csid, derived);
                    }
                    (
                        BatchDeriveStats::Serial(per_commit_stats),
                        per_commit_derived,
                    )
                }
            };

            // Flush the blobstore.  If it has been set up to cache writes, these
            // must be flushed before we write the mapping.
            let (stats, _) = derivation_ctx.flush(ctx).try_timed().await?;
            scuba
                .add_future_stats(&stats)
                .log_with_msg("Flushed derived blobs", None);

            let mut derivation_ctx = self.derivation_context(rederivation.clone());
            derivation_ctx.enable_write_batching();
            // Write all mapping values, and flush the blobstore to ensure they
            // are persisted.
            let (persist_stats, persisted) = async {
                let derivation_ctx_ref = &derivation_ctx;
                let csids = stream::iter(derived.into_iter())
                    .map(|(csid, derived)| async move {
                        derived.store_mapping(ctx, derivation_ctx_ref, csid).await?;
                        Ok::<_, Error>(csid)
                    })
                    .buffer_unordered(100)
                    .try_collect::<Vec<_>>()
                    .await?;

                derivation_ctx.flush(ctx).await?;
                if let Some(rederivation) = rederivation {
                    for csid in csids {
                        rederivation.mark_derived(Derivable::NAME, csid);
                    }
                }
                Ok::<_, Error>(())
            }
            .timed()
            .await;

            self.log_mapping_insertion::<Derivable>(
                ctx,
                &mut derived_data_scuba,
                None,
                &persist_stats,
                persisted.as_ref().err(),
            );

            persisted?;

            scuba
                .add_future_stats(&persist_stats)
                .log_with_msg("Flushed mapping", None);

            Ok(batch_stats)
        }
        .timed()
        .await;

        self.log_batch_derivation_end::<Derivable>(
            ctx,
            &mut derived_data_scuba,
            csid_range,
            &overall_stats,
            result.as_ref().err(),
        );

        let batch_stats = result?;

        Ok(batch_stats.append(secondary_derivation.await?)?)
    }

    /// Fetch derived data for a changeset if it has previously been derived.
    pub async fn fetch_derived<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<Option<Derivable>, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.get_manager(ctx, csid)
            .await?
            .fetch_derived_impl::<Derivable>(ctx, csid, rederivation)
            .await
    }

    async fn fetch_derived_impl<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<Option<Derivable>, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        self.check_enabled::<Derivable>()?;
        let derivation_ctx = self.derivation_context(rederivation);
        let derived = derivation_ctx.fetch_derived::<Derivable>(ctx, csid).await?;
        Ok(derived)
    }

    #[async_recursion]
    /// Fetch derived data for a batch of changesets if they have previously
    /// been derived.
    ///
    /// Returns a hashmap from changeset id to the derived data.  Changesets
    /// for which the data has not previously been derived are omitted.
    pub async fn fetch_derived_batch<Derivable>(
        &self,
        ctx: &CoreContext,
        csids: Vec<ChangesetId>,
        rederivation: Option<Arc<dyn Rederivation>>,
    ) -> Result<HashMap<ChangesetId, Derivable>, DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        let (csids, secondary_derivation) = if let Some(secondary_data) = &self.inner.secondary {
            let DerivationAssignment { primary, secondary } =
                secondary_data.assigner.assign(ctx, csids).await?;
            (primary, {
                cloned!(rederivation);
                async move {
                    secondary_data
                        .manager
                        .fetch_derived_batch::<Derivable>(ctx, secondary, rederivation)
                        .await
                }
                .left_future()
            })
        } else {
            (csids, future::ready(Ok(HashMap::new())).right_future())
        };
        self.check_enabled::<Derivable>()?;
        let derivation_ctx = self.derivation_context(rederivation);
        let mut derived = derivation_ctx
            .fetch_derived_batch::<Derivable>(ctx, csids)
            .await?;
        derived.extend(secondary_derivation.await?);
        Ok(derived)
    }
}

pub(super) struct DerivationOutcome<Derivable> {
    /// The derived data.
    pub(super) derived: Derivable,

    /// Number of changesets that were derived.
    pub(super) count: u64,

    /// Time take to find the underived changesets.
    pub(super) find_underived_time: Duration,
}

enum DerivationState {
    NotRequested,
    InProgress,
}

fn emergency_disabled(repo_name: &str, derivable_name: &str) -> bool {
    let disabled_for_repo = tunables::tunables()
        .by_repo_all_derived_data_disabled(repo_name)
        .unwrap_or(false);

    if disabled_for_repo {
        return true;
    }

    let disabled_for_type = tunables::tunables()
        .by_repo_derived_data_types_disabled(repo_name)
        .unwrap_or_else(Vec::new);

    if disabled_for_type
        .iter()
        .any(|ty| ty.as_str() == derivable_name)
    {
        return true;
    }

    // Not disabled
    false
}

async fn derivation_disabled_watcher(repo_name: &str, derivable_name: &'static str) {
    let delay_duration = get_duration_from_tunable_or(
        MononokeTunables::derived_data_disabled_watcher_delay_secs,
        Duration::from_secs,
        10,
    );
    while !emergency_disabled(repo_name, derivable_name) {
        tokio::time::sleep(delay_duration).await;
    }
}


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

use std::time::Duration;

use anyhow::Error;
use anyhow::Result;
use context::CoreContext;
use context::PerfCounters;
use derived_data_constants::*;
use futures_stats::FutureStats;
use mononoke_types::ChangesetId;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::warn;
use time_ext::DurationExt;

use super::derive::DerivationOutcome;
use super::DerivedDataManager;
use crate::derivable::BonsaiDerivable;
use crate::error::DerivationError;

impl DerivedDataManager {
    /// Log the start of derivation to both the request and derived data scuba
    /// tables.
    pub(super) fn log_derivation_start<Derivable>(
        &self,
        ctx: &CoreContext,
        derived_data_scuba: &mut MononokeScubaSampleBuilder,
        csid: ChangesetId,
    ) where
        Derivable: BonsaiDerivable,
    {
        ctx.scuba().clone().log_with_msg(
            DERIVATION_START,
            Some(format!("{} {}", Derivable::NAME, csid)),
        );
        derived_data_scuba.log_with_msg(DERIVATION_START, None);
    }

    /// Log the end of derivation to both the request and derived data scuba
    /// tables.
    pub(super) fn log_derivation_end<Derivable>(
        &self,
        ctx: &CoreContext,
        derived_data_scuba: &mut MononokeScubaSampleBuilder,
        csid: ChangesetId,
        stats: &FutureStats,
        error: Option<&Error>,
    ) where
        Derivable: BonsaiDerivable,
    {
        let (tag, error_str) = match error {
            None => (DERIVATION_END, None),
            Some(error) => (FAILED_DERIVATION, Some(format!("{:#}", error))),
        };

        let mut ctx_scuba = ctx.scuba().clone();
        ctx_scuba.add_future_stats(stats);
        if let Some(error_str) = &error_str {
            ctx_scuba.add("Derive error", error_str.as_str());
        };
        ctx_scuba.log_with_msg(tag, Some(format!("{} {}", Derivable::NAME, csid)));

        ctx.perf_counters().insert_perf_counters(derived_data_scuba);

        derived_data_scuba.add_future_stats(stats);
        derived_data_scuba.log_with_msg(tag, error_str);
    }

    /// Log the start of batch derivation to both the request and derived data
    /// scuba tables.
    pub(super) fn log_batch_derivation_start<Derivable>(
        &self,
        ctx: &CoreContext,
        derived_data_scuba: &mut MononokeScubaSampleBuilder,
        csid_range: Option<(ChangesetId, ChangesetId)>,
    ) where
        Derivable: BonsaiDerivable,
    {
        if let Some((first, last)) = csid_range {
            ctx.scuba().clone().log_with_msg(
                DERIVATION_START_BATCH,
                Some(format!("{} {}-{}", Derivable::NAME, first, last)),
            );
            derived_data_scuba.log_with_msg(DERIVATION_START_BATCH, None);
        }
    }

    /// Log the end of derivation to both the request and derived data scuba
    /// tables.
    pub(super) fn log_batch_derivation_end<Derivable>(
        &self,
        ctx: &CoreContext,
        derived_data_scuba: &mut MononokeScubaSampleBuilder,
        csid_range: Option<(ChangesetId, ChangesetId)>,
        stats: &FutureStats,
        error: Option<&Error>,
    ) where
        Derivable: BonsaiDerivable,
    {
        if let Some((first, last)) = csid_range {
            let (tag, error_str) = match error {
                None => (DERIVATION_END_BATCH, None),
                Some(error) => (FAILED_DERIVATION_BATCH, Some(format!("{:#}", error))),
            };

            let mut ctx_scuba = ctx.scuba().clone();
            ctx_scuba.add_future_stats(stats);
            if let Some(error_str) = &error_str {
                ctx_scuba.add("Derive error", error_str.as_str());
            };
            ctx_scuba.log_with_msg(tag, Some(format!("{} {}-{}", Derivable::NAME, first, last)));

            ctx.perf_counters().insert_perf_counters(derived_data_scuba);

            derived_data_scuba.add_future_stats(stats);
            derived_data_scuba.log_with_msg(tag, error_str);
        }
    }

    /// Log the insertion of a new derived data mapping to the derived data
    /// scuba table.
    pub(super) fn log_mapping_insertion<Derivable>(
        &self,
        ctx: &CoreContext,
        derived_data_scuba: &mut MononokeScubaSampleBuilder,
        value: Option<&Derivable>,
        stats: &FutureStats,
        error: Option<&Error>,
    ) where
        Derivable: BonsaiDerivable,
    {
        let (tag, error_str) = match error {
            None => (INSERTED_MAPPING, None),
            Some(error) => (FAILED_INSERTING_MAPPING, Some(format!("{:#}", error))),
        };

        ctx.perf_counters().insert_perf_counters(derived_data_scuba);

        if let Some(value) = value {
            // Limit how much we log to scuba.
            let value = format!("{:1000?}", value);
            derived_data_scuba.add("mapping_value", value);
        }

        derived_data_scuba
            .add_future_stats(stats)
            .log_with_msg(tag, error_str);
    }

    pub(super) fn should_log_slow_derivation(&self, duration: Duration) -> bool {
        let threshold = tunables::tunables()
            .derived_data_slow_derivation_threshold_secs()
            .unwrap_or_default();
        let threshold = match threshold.try_into() {
            Ok(t) if t > 0 => t,
            _ => return false,
        };
        duration > Duration::from_secs(threshold)
    }

    pub(super) fn log_slow_derivation<Derivable>(
        &self,
        ctx: &CoreContext,
        csid: ChangesetId,
        stats: &FutureStats,
        pc: &PerfCounters,
        result: &Result<DerivationOutcome<Derivable>, DerivationError>,
    ) where
        Derivable: BonsaiDerivable,
    {
        let mut scuba = ctx.scuba().clone();
        pc.insert_perf_counters(&mut scuba);

        scuba.add_future_stats(stats);
        scuba.add("changeset_id", csid.to_string());
        scuba.add("derived_data_type", Derivable::NAME);
        scuba.add("repo", self.repo_name());

        match result {
            Ok(derivation_outcome) => {
                scuba.add("derived", derivation_outcome.count);
                scuba.add(
                    "find_underived_completion_time_ms",
                    derivation_outcome.find_underived_time.as_millis_unchecked(),
                );
                warn!(
                    ctx.logger(),
                    "slow derivation of {} for {}, took {:.2?} (find_underived: {:.2?}), derived {} changesets",
                    Derivable::NAME,
                    csid,
                    stats.completion_time,
                    derivation_outcome.find_underived_time,
                    derivation_outcome.count,
                );
            }
            Err(derivation_error) => {
                warn!(
                    ctx.logger(),
                    "slow derivation of {} for {}, took {:.2?}, failed with Err({:?})",
                    Derivable::NAME,
                    csid,
                    stats.completion_time,
                    derivation_error,
                );
            }
        }

        scuba.log_with_msg(SLOW_DERIVATION, None);
    }
}


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

use std::time::Duration;

use anyhow::Result;
use context::CoreContext;
use context::SessionClass;
use scuba_ext::MononokeScubaSampleBuilder;

use super::DerivedDataManager;
use crate::derivable::BonsaiDerivable;
use crate::error::DerivationError;

#[derive(Clone, Debug)]
pub struct DiscoveryStats {
    pub(crate) find_underived_completion_time: Duration,
    pub(crate) commits_discovered: u32,
}

impl DiscoveryStats {
    fn add_scuba_fields(&self, builder: &mut MononokeScubaSampleBuilder) {
        builder.add(
            "find_underived_completion_time_ms",
            self.find_underived_completion_time.as_millis() as u64,
        );
        builder.add("commits_discovered", self.commits_discovered);
    }
}

impl DerivedDataManager {
    /// Returns the passed-in `CoreContext` with the session class modified to
    /// the one that should be used for derivation.
    pub(super) fn set_derivation_session_class(&self, mut ctx: CoreContext) -> CoreContext {
        if tunables::tunables()
            .by_repo_derived_data_use_background_session_class(self.repo_name())
            .unwrap_or(false)
        {
            ctx.session_mut()
                .override_session_class(SessionClass::BackgroundUnlessTooSlow);
        }
        ctx
    }

    /// Construct a scuba sample builder for logging derivation of this
    /// changeset to the derived data scuba table.
    pub(super) fn derived_data_scuba<Derivable>(
        &self,
        discovery_stats: &Option<DiscoveryStats>,
    ) -> MononokeScubaSampleBuilder
    where
        Derivable: BonsaiDerivable,
    {
        let mut scuba = self.inner.scuba.clone();
        scuba.add("derived_data", Derivable::NAME);
        if let Some(discovery_stats) = discovery_stats {
            discovery_stats.add_scuba_fields(&mut scuba);
        }
        scuba
    }

    pub(super) fn check_enabled<Derivable>(&self) -> Result<(), DerivationError>
    where
        Derivable: BonsaiDerivable,
    {
        if self.config().types.contains(Derivable::NAME) {
            Ok(())
        } else {
            Err(DerivationError::Disabled(
                Derivable::NAME,
                self.repo_id(),
                self.repo_name().to_string(),
            ))
        }
    }

    pub(super) fn max_parallel_derivations(&self) -> usize {
        let buffer_size = tunables::tunables()
            .derived_data_parallel_derivation_buffer()
            .unwrap_or_default();
        if buffer_size > 0 {
            buffer_size
                .try_into()
                .expect("buffer size should convert to usize")
        } else {
            10
        }
    }
}


/*
 * 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::time::Instant;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Error;
use async_trait::async_trait;
use blobrepo_common::changed_files::compute_changed_files;
use blobstore::Blobstore;
use blobstore::Loadable;
use borrowed::borrowed;
use cloned::cloned;
use context::CoreContext;
use derived_data_manager::DerivationError;
use futures::future;
use futures::future::try_join;
use futures::future::try_join_all;
use futures::stream;
use futures::FutureExt;
use futures::TryStreamExt;
use manifest::ManifestChanges;
use manifest::ManifestOps;
use mercurial_types::blobs::ChangesetMetadata;
use mercurial_types::blobs::ContentBlobMeta;
use mercurial_types::blobs::File;
use mercurial_types::blobs::HgBlobChangeset;
use mercurial_types::blobs::HgChangesetContent;
use mercurial_types::blobs::UploadHgFileContents;
use mercurial_types::blobs::UploadHgFileEntry;
use mercurial_types::blobs::UploadHgNodeHash;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestId;
use mercurial_types::HgParents;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::FileChange;
use mononoke_types::FileType;
use mononoke_types::MPath;
use mononoke_types::TrackedFileChange;
use repo_derived_data::RepoDerivedData;
use repo_derived_data::RepoDerivedDataRef;
use stats::prelude::*;

use crate::derive_hg_manifest::derive_hg_manifest;
use crate::derive_hg_manifest::derive_simple_hg_manifest_stack_without_copy_info;
use crate::mapping::HgChangesetDeriveOptions;
use crate::mapping::MappedHgChangesetId;

define_stats! {
    prefix = "mononoke.blobrepo";
    get_hg_from_bonsai_changeset: timeseries(Rate, Sum),
    generate_hg_from_bonsai_changeset: timeseries(Rate, Sum),
    generate_hg_from_bonsai_total_latency_ms: histogram(100, 0, 10_000, Average; P 50; P 75; P 90; P 95; P 99),
    generate_hg_from_bonsai_single_latency_ms: histogram(100, 0, 10_000, Average; P 50; P 75; P 90; P 95; P 99),
    generate_hg_from_bonsai_generated_commit_num: histogram(1, 0, 20, Average; P 50; P 75; P 90; P 95; P 99),
}

async fn can_reuse_filenode(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    parent: HgFileNodeId,
    change: &TrackedFileChange,
) -> Result<Option<HgFileNodeId>, Error> {
    let parent_envelope = parent.load(ctx, blobstore).await?;
    let parent_copyfrom_path = File::extract_copied_from(parent_envelope.metadata())?.map(|t| t.0);
    let parent_content_id = parent_envelope.content_id();

    if parent_content_id == change.content_id()
        && change.copy_from().map(|t| &t.0) == parent_copyfrom_path.as_ref()
    {
        Ok(Some(parent))
    } else {
        Ok(None)
    }
}

pub(crate) async fn store_file_change<'a>(
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    p1: Option<HgFileNodeId>,
    p2: Option<HgFileNodeId>,
    path: &'a MPath,
    change: &'a TrackedFileChange,
    copy_from: Option<(MPath, HgFileNodeId)>,
) -> Result<(FileType, HgFileNodeId), Error> {
    // If we produced a hg change that has copy info, then the Bonsai should have copy info
    // too. However, we could have Bonsai copy info without having copy info in the hg change
    // if we stripped it out to produce a hg changeset for an Octopus merge and the copy info
    // references a step-parent (i.e. neither p1, not p2).
    if copy_from.is_some() {
        assert!(change.copy_from().is_some());
    }

    // Mercurial has complicated logic of finding file parents, especially
    // if a file was also copied/moved.
    // See mercurial/localrepo.py:_filecommit().
    // Mononoke uses simpler rules, which still produce a usable result

    // Simplify parents. This aims to reduce the amount of work done in reuse, and avoids
    // semantically unhelpful duplicate parents.
    let (p1, p2) = match (p1, p2) {
        (Some(p1), None) => (Some(p1), None),
        (None, Some(p2)) => (Some(p2), None),
        (p1, p2) if p1 == p2 => (p1, None),
        (p1, p2) => (p1, p2),
    };

    // we can reuse same HgFileNodeId if we have a filenode that has
    // same content and copyfrom information
    let maybe_filenode = match (p1, p2) {
        (Some(parent), None) | (None, Some(parent)) => {
            can_reuse_filenode(&ctx, &blobstore, parent, change).await?
        }
        (Some(p1), Some(p2)) => {
            let (reuse_p1, reuse_p2) = try_join(
                can_reuse_filenode(&ctx, &blobstore, p1, change),
                can_reuse_filenode(&ctx, &blobstore, p2, change),
            )
            .await?;
            reuse_p1.or(reuse_p2)
        }
        // No filenode to reuse
        (None, None) => None,
    };

    let filenode_id = match maybe_filenode {
        Some(filenode) => filenode,
        None => {
            let p1 = if p1.is_some()
                && p2.is_none()
                && copy_from.is_some()
                && copy_from.as_ref().map(|c| &c.0) != Some(path)
            {
                // Mercurial special-cases the "file copied over existing file" case, and does not
                // put in a parent in that situation - `hg log` then looks down the copyfrom
                // information instead. This is not the best decision, but we should keep it
                // For example:
                // ```
                // echo 1 > 1 && echo 2 > 2 && hg ci -A -m first
                // hg cp 2 1 --force && hg ci -m second
                // # File '1' has both p1 and copy from.
                // ```
                // In this case, Mercurial and Mononoke both drop the p1 information in the filenode,
                // instead relying on the copyfrom for history.
                None
            } else {
                p1
            };

            let upload_entry = UploadHgFileEntry {
                upload_node_id: UploadHgNodeHash::Generate,
                contents: UploadHgFileContents::ContentUploaded(ContentBlobMeta {
                    id: change.content_id(),
                    size: change.size(),
                    copy_from: copy_from.clone(),
                }),
                p1,
                p2,
            };

            upload_entry.upload(ctx, blobstore, Some(path)).await?
        }
    };

    Ok((change.file_type(), filenode_id))
}

async fn resolve_paths(
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    manifest_id: Option<HgManifestId>,
    paths: Vec<MPath>,
) -> Result<HashMap<MPath, HgFileNodeId>, Error> {
    match manifest_id {
        None => Ok(HashMap::new()),
        Some(manifest_id) => {
            let mapping: HashMap<MPath, HgFileNodeId> = manifest_id
                .find_entries(ctx, blobstore, paths)
                .map_ok(|(path, entry)| Some((path?, entry.into_leaf()?.1)))
                .try_filter_map(future::ok)
                .try_collect()
                .await?;
            Ok(mapping)
        }
    }
}

pub async fn get_manifest_from_bonsai(
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    bcs: BonsaiChangeset,
    parent_manifests: Vec<HgManifestId>,
) -> Result<HgManifestId, Error> {
    // NOTE: We ignore further parents beyond p1 and p2 for the purposed of tracking copy info
    // or filenode parents. This is because hg supports just 2 parents at most, so we track
    // copy info & filenode parents relative to the first 2 parents, then ignore other parents.

    let (manifest_p1, manifest_p2) = {
        let mut manifests = parent_manifests.iter();
        (manifests.next().copied(), manifests.next().copied())
    };

    let (p1, p2) = {
        let mut parents = bcs.parents();
        let p1 = parents.next();
        let p2 = parents.next();
        (p1, p2)
    };

    let file_changes = bcs
        .file_changes()
        .map(|(path, fc)| {
            Ok((
                path.clone(),
                match fc {
                    FileChange::Change(tc) => Some(tc.clone()),
                    FileChange::Deletion => None,
                    FileChange::UntrackedChange(_) | FileChange::UntrackedDeletion => {
                        bail!("Can't derive manifest for snapshot")
                    }
                },
            ))
        })
        .collect::<Result<Vec<_>, Error>>()?;

    // paths *modified* by changeset or *copied from parents*
    let mut p1_paths = Vec::new();
    let mut p2_paths = Vec::new();
    for (path, file_change) in file_changes.iter() {
        match file_change {
            Some(file_change) => {
                if let Some((copy_path, bcsid)) = file_change.copy_from() {
                    if Some(bcsid) == p1.as_ref() {
                        p1_paths.push(copy_path.clone());
                    }
                    if Some(bcsid) == p2.as_ref() {
                        p2_paths.push(copy_path.clone());
                    }
                };
                p1_paths.push(path.clone());
                p2_paths.push(path.clone());
            }
            None => {}
        }
    }

    // TODO:
    // `derive_manifest` already provides parents for newly created files, so we
    // can remove **all** lookups to files from here, and only leave lookups for
    // files that were copied (i.e bonsai changes that contain `copy_path`)
    let (p1s, p2s) = try_join(
        resolve_paths(ctx.clone(), blobstore.clone(), manifest_p1, p1_paths),
        resolve_paths(ctx.clone(), blobstore.clone(), manifest_p2, p2_paths),
    )
    .await?;

    let file_changes: Vec<_> = file_changes
        .into_iter()
        .map(|(path, file_change)| Ok::<_, Error>((path, file_change)))
        .collect();
    let changes: Vec<_> = stream::iter(file_changes)
        .map_ok({
            cloned!(ctx, blobstore);
            move |(path, file_change)| match file_change {
                None => future::ok((path, None)).left_future(),
                Some(file_change) => {
                    let copy_from = file_change.copy_from().and_then(|(copy_path, bcsid)| {
                        if Some(bcsid) == p1.as_ref() {
                            p1s.get(copy_path).map(|id| (copy_path.clone(), *id))
                        } else if Some(bcsid) == p2.as_ref() {
                            p2s.get(copy_path).map(|id| (copy_path.clone(), *id))
                        } else {
                            None
                        }
                    });
                    let p1 = p1s.get(&path).cloned();
                    let p2 = p2s.get(&path).cloned();
                    cloned!(ctx, blobstore);
                    let spawned = tokio::spawn(async move {
                        let entry = store_file_change(
                            ctx,
                            blobstore,
                            p1,
                            p2,
                            &path,
                            &file_change,
                            copy_from,
                        )
                        .await?;
                        Ok((path, Some(entry)))
                    });
                    async move { spawned.await? }.boxed().right_future()
                }
            }
        })
        .try_buffer_unordered(100)
        .try_collect()
        .await?;

    let manifest_id = derive_hg_manifest(ctx.clone(), blobstore, parent_manifests, changes).await?;

    Ok(manifest_id)
}

pub(crate) async fn derive_from_parents(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    bonsai: BonsaiChangeset,
    parents: Vec<MappedHgChangesetId>,
    options: &HgChangesetDeriveOptions,
) -> Result<MappedHgChangesetId, Error> {
    let parents = {
        borrowed!(ctx);
        try_join_all(
            parents
                .into_iter()
                .map(|id| async move { id.hg_changeset_id().load(ctx, blobstore).await }),
        )
        .await?
    };

    let parent_manifests = parents.iter().map(|p| p.manifestid()).collect();
    let manifest_id = get_manifest_from_bonsai(
        ctx.clone(),
        blobstore.clone(),
        bonsai.clone(),
        parent_manifests,
    )
    .await?;

    let (hg_cs_id, _) =
        generate_hg_changeset(ctx, blobstore, bonsai, manifest_id, parents, options).await?;
    Ok(MappedHgChangesetId::new(hg_cs_id))
}

pub async fn derive_simple_hg_changeset_stack_without_copy_info(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    bonsais: Vec<BonsaiChangeset>,
    parent: Option<MappedHgChangesetId>,
    options: &HgChangesetDeriveOptions,
) -> Result<HashMap<ChangesetId, MappedHgChangesetId>, Error> {
    let parent = match parent {
        Some(parent) => Some(parent.hg_changeset_id().load(ctx, blobstore).await?),
        None => None,
    };

    let file_changes = bonsais
        .iter()
        .map(|bonsai| {
            let per_commit_file_changes: Result<Vec<_>, Error> = bonsai
                .file_changes()
                .map(|(path, fc)| {
                    use FileChange::*;
                    let tracked_file_change = match fc {
                        Change(tracked_file_change) => Some(tracked_file_change.clone()),
                        Deletion => None,
                        UntrackedChange(_) | UntrackedDeletion => {
                            bail!(
                                "unexpected untracked file change while deriving {}",
                                bonsai.get_changeset_id()
                            );
                        }
                    };
                    Ok((path.clone(), tracked_file_change))
                })
                .collect();
            let per_commit_file_changes = per_commit_file_changes?;

            let mf_changes = ManifestChanges {
                cs_id: bonsai.get_changeset_id(),
                changes: per_commit_file_changes,
            };

            Ok(mf_changes)
        })
        .collect::<Result<Vec<_>, Error>>();
    let file_changes = file_changes?;

    let mf_ids = derive_simple_hg_manifest_stack_without_copy_info(
        ctx.clone(),
        blobstore.clone(),
        file_changes,
        parent.clone().map(|p| p.manifestid()),
    )
    .await?;
    let mut parents = parent.into_iter().collect::<Vec<_>>();

    let mut res = HashMap::with_capacity(bonsais.len());
    for bonsai in bonsais {
        let cs_id = bonsai.get_changeset_id();
        let mf_id = mf_ids.get(&cs_id).ok_or_else(|| {
            anyhow!(
                "not found manifest for {} but should have derived it in this function",
                cs_id
            )
        })?;
        let (hg_changeset_id, hg_cs) =
            generate_hg_changeset(ctx, blobstore, bonsai, *mf_id, parents, options).await?;
        res.insert(cs_id, MappedHgChangesetId::new(hg_changeset_id));
        parents = vec![hg_cs];
    }

    Ok(res)
}

async fn generate_hg_changeset(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    bcs: BonsaiChangeset,
    manifest_id: HgManifestId,
    parents: Vec<HgBlobChangeset>,
    options: &HgChangesetDeriveOptions,
) -> Result<(HgChangesetId, HgBlobChangeset), Error> {
    let start_timestamp = Instant::now();

    // NOTE: We're special-casing the first 2 parents here, since that's all Mercurial
    // supports. Producing the Manifest (in get_manifest_from_bonsai) will consider all
    // parents, but everything else is only presented with the first 2 parents, because that's
    // all Mercurial knows about for now. This lets us produce a meaningful Hg changeset for a
    // Bonsai changeset with > 2 parents (which might be one we imported from Git).
    let mut parents = parents.into_iter();
    let p1 = parents.next();
    let p2 = parents.next();

    let p1_hash = p1.as_ref().map(|p1| p1.get_changeset_id());
    let p2_hash = p2.as_ref().map(|p2| p2.get_changeset_id());

    let mf_p1 = p1.map(|p| p.manifestid());
    let mf_p2 = p2.map(|p| p.manifestid());

    let hg_parents = HgParents::new(
        p1_hash.map(|h| h.into_nodehash()),
        p2_hash.map(|h| h.into_nodehash()),
    );

    // Keep a record of any parents for now (i.e. > 2 parents). We'll store those in extras.
    let step_parents = parents;

    let files = compute_changed_files(
        ctx.clone(),
        blobstore.clone(),
        manifest_id.clone(),
        mf_p1,
        mf_p2,
    )
    .await?;

    let mut metadata = ChangesetMetadata {
        user: bcs.author().to_string(),
        time: *bcs.author_date(),
        extra: bcs
            .hg_extra()
            .map(|(k, v)| (k.as_bytes().to_vec(), v.to_vec()))
            .collect(),
        message: bcs.message().to_string(),
    };
    metadata.record_step_parents(step_parents.map(|blob| blob.get_changeset_id()));
    if options.set_committer_field {
        match (bcs.committer(), bcs.committer_date()) {
            (Some(committer), Some(date)) => {
                // Do not record committer if it's the same as author
                if committer != bcs.author() || date != bcs.author_date() {
                    metadata.record_committer(committer, date)?;
                }
            }
            (None, None) => {}
            _ => {
                bail!("invalid committer/committer date in bonsai changeset");
            }
        };
    }

    let content = HgChangesetContent::new_from_parts(hg_parents, manifest_id, metadata, files);
    let cs = HgBlobChangeset::new(content)?;
    let csid = cs.get_changeset_id();

    cs.save(ctx, blobstore).await?;

    STATS::generate_hg_from_bonsai_single_latency_ms
        .add_value(start_timestamp.elapsed().as_millis() as i64);
    STATS::generate_hg_from_bonsai_generated_commit_num.add_value(1);

    Ok((csid, cs))
}

#[async_trait]
pub trait DeriveHgChangeset {
    async fn derive_hg_changeset(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<HgChangesetId, Error>;
}

pub async fn derive_hg_changeset(
    ctx: &CoreContext,
    derived_data: &RepoDerivedData,
    cs_id: ChangesetId,
) -> Result<HgChangesetId, Error> {
    STATS::get_hg_from_bonsai_changeset.add_value(1);
    let start_timestamp = Instant::now();
    let result = match derived_data.derive::<MappedHgChangesetId>(ctx, cs_id).await {
        Ok(id) => Ok(id.hg_changeset_id()),
        Err(err @ DerivationError::Disabled(..)) => Err(err.into()),
        Err(DerivationError::Error(err)) => Err(err),
    };
    STATS::generate_hg_from_bonsai_total_latency_ms
        .add_value(start_timestamp.elapsed().as_millis() as i64);
    result
}

#[async_trait]
impl<Repo: RepoDerivedDataRef + Send + Sync> DeriveHgChangeset for Repo {
    async fn derive_hg_changeset(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<HgChangesetId, Error> {
        derive_hg_changeset(ctx, self.repo_derived_data(), cs_id).await
    }
}