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

use std::sync::Arc;
use std::sync::Once;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use blobrepo::BlobRepo;
use blobrepo_override::DangerousOverride;
use blobstore::Blobstore;
use blobstore::StoreLoadable;
use bonsai_hg_mapping::ArcBonsaiHgMapping;
use bonsai_hg_mapping::MemWritesBonsaiHgMapping;
use borrowed::borrowed;
use cacheblob::MemWritesBlobstore;
use changeset_fetcher::ChangesetFetcherArc;
use clap_old::ArgMatches;
use cmdlib::args;
use cmdlib::args::MononokeMatches;
use context::CoreContext;
use derived_data::BonsaiDerived;
use derived_data_manager::BonsaiDerivable;
use derived_data_utils::derived_data_utils;
use derived_data_utils::DerivedUtils;
use derived_data_utils::DERIVED_DATA_DEPS;
use fsnodes::RootFsnodeId;
use futures::future::try_join;
use futures::future::try_join_all;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use manifest::find_intersection_of_diffs;
use manifest::find_intersection_of_diffs_and_parents;
use manifest::Entry;
use manifest::Manifest;
use mercurial_derivation::MappedHgChangesetId;
use mononoke_types::BlobstoreKey;
use mononoke_types::ChangesetId;
use readonlyblob::ReadOnlyBlobstore;
use repo_blobstore::RepoBlobstoreArc;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataArc;
use repo_identity::RepoIdentityRef;
use skeleton_manifest::RootSkeletonManifestId;
use slog::info;
use slog::warn;
use unodes::RootUnodeManifestId;

use crate::commit_discovery::CommitDiscoveryOptions;
use crate::regenerate;
use crate::ARG_DERIVED_DATA_TYPE;
use crate::ARG_VALIDATE_CHUNK_SIZE;

pub async fn validate(
    ctx: &CoreContext,
    matches: &MononokeMatches<'_>,
    sub_m: &ArgMatches<'_>,
    repo_name: String,
) -> Result<(), Error> {
    if !matches.environment().readonly_storage.0 {
        return Err(anyhow!(
            "validate subcommand should be run only on readonly storage!"
        ));
    }
    let repo: BlobRepo =
        args::open_repo_by_name_unredacted(ctx.fb, ctx.logger(), matches, repo_name).await?;
    let csids = CommitDiscoveryOptions::from_matches(ctx, &repo, sub_m)
        .await?
        .get_commits();

    let derived_data_type = sub_m
        .value_of(ARG_DERIVED_DATA_TYPE)
        .ok_or_else(|| anyhow!("{} is not set", ARG_DERIVED_DATA_TYPE))?;
    info!(
        ctx.logger(),
        "Validating {} on {}...",
        derived_data_type,
        repo.repo_identity().name()
    );
    let opts = regenerate::DeriveOptions::from_matches(sub_m)?;

    let validate_chunk_size = args::get_usize(&sub_m, ARG_VALIDATE_CHUNK_SIZE, 10000);
    let warn_once = Once::new();

    info!(ctx.logger(), "Started validation");
    for chunk in csids.chunks(validate_chunk_size) {
        let chunk = chunk.to_vec();
        info!(
            ctx.logger(),
            "Processing chunk starting from {:?}",
            chunk.get(0)
        );
        let orig_repo = repo.clone();
        let mut memblobstore = None;
        let mut membonsaihgmapping = None;
        let repo = repo
            .dangerous_override(|blobstore| -> Arc<dyn Blobstore> {
                let blobstore = Arc::new(MemWritesBlobstore::new(blobstore));
                memblobstore = Some(blobstore.clone());
                blobstore
            })
            .dangerous_override(|bonsai_hg_mapping| -> ArcBonsaiHgMapping {
                let bonsai_hg_mapping = Arc::new(MemWritesBonsaiHgMapping::new(bonsai_hg_mapping));
                membonsaihgmapping = Some(bonsai_hg_mapping.clone());
                bonsai_hg_mapping
            });
        let memblobstore = memblobstore.unwrap();
        let membonsaihgmapping = membonsaihgmapping.unwrap();
        // By default MemWritesBonsaiHgMapping doesn't save data to cache if it
        // already exists in underlying mapping. This option disables this feature.
        membonsaihgmapping.set_save_noop_writes(true);

        let types = std::iter::once(derived_data_type.to_string())
            .chain(
                DERIVED_DATA_DEPS
                    .get(derived_data_type)
                    .unwrap()
                    .iter()
                    .map(|t| t.to_string()),
            )
            .collect::<Vec<_>>();

        regenerate::regenerate_derived_data(ctx, &repo, chunk.clone(), types, &opts).await?;

        {
            let cache = memblobstore.get_cache().lock().unwrap();
            info!(ctx.logger(), "created {} blobs", cache.len());
        }
        let real_derived_utils = &derived_data_utils(ctx.fb, &orig_repo, derived_data_type)?;

        // Make sure that the generated data was saved in memory blobstore
        membonsaihgmapping.set_no_access_to_inner(true);
        membonsaihgmapping.set_readonly(true);
        memblobstore.set_no_access_to_inner(true);
        let repo = repo.dangerous_override(|blobstore| -> Arc<dyn Blobstore> {
            Arc::new(ReadOnlyBlobstore::new(blobstore))
        });
        let rederived_utils = &derived_data_utils(ctx.fb, &repo, derived_data_type)?;

        borrowed!(ctx, orig_repo, repo, warn_once);
        stream::iter(chunk)
            .map(Ok)
            .try_for_each_concurrent(100, |csid| async move {
                if !rederived_utils.is_derived(ctx, csid).await? {
                    return Err(anyhow!("{} unexpectedly not derived", csid));
                }

                let f1 =
                    real_derived_utils.derive(ctx.clone(), orig_repo.repo_derived_data_arc(), csid);
                let f2 = rederived_utils.derive(ctx.clone(), repo.repo_derived_data_arc(), csid);
                let (real, rederived) = try_join(f1, f2).await?;
                if real != rederived {
                    return Err(anyhow!("mismatch in {}: {} vs {}", csid, real, rederived));
                };

                validate_generated_data(ctx, orig_repo, warn_once, real_derived_utils, csid, repo)
                    .await
                    .with_context(|| format!("failed validating generated data for {}", csid))
            })
            .await?;
        info!(ctx.logger(), "Validation successful!");
    }

    Ok(())
}

async fn validate_generated_data<'a>(
    ctx: &'a CoreContext,
    real_repo: &'a BlobRepo,
    warn_once: &Once,
    real_derived_utils: &'a Arc<dyn DerivedUtils>,
    cs_id: ChangesetId,
    mem_blob_repo: &'a BlobRepo,
) -> Result<(), Error> {
    let mem_blob = mem_blob_repo.repo_blobstore_arc() as Arc<dyn Blobstore>;
    if real_derived_utils.name() == RootFsnodeId::NAME {
        validate_fsnodes(ctx, real_repo, cs_id, &mem_blob).await?;
    } else if real_derived_utils.name() == RootSkeletonManifestId::NAME {
        validate_skeleton_manifests(ctx, real_repo, cs_id, &mem_blob).await?;
    } else if real_derived_utils.name() == RootUnodeManifestId::NAME {
        validate_unodes(ctx, real_repo, cs_id, &mem_blob).await?;
    } else if real_derived_utils.name() == MappedHgChangesetId::NAME {
        validate_hgchangesets(ctx, real_repo, cs_id, &mem_blob).await?;
    } else {
        warn_once.call_once(||
            warn!(
                ctx.logger(),
                "Validating generated blobs is not supported for {}, so no validation of generated blobs was done!",
                real_derived_utils.name()
            )
        );
    }

    Ok(())
}

async fn validate_fsnodes<'a>(
    ctx: &'a CoreContext,
    real_repo: &'a BlobRepo,
    cs_id: ChangesetId,
    mem_blob: &'a Arc<dyn Blobstore>,
) -> Result<(), Error> {
    let real_blobstore = real_repo.repo_blobstore_arc();
    let (fsnode, parents) =
        find_cs_and_parents_derived_data::<RootFsnodeId>(ctx, real_repo, cs_id).await?;
    let fsnode = *fsnode.fsnode_id();
    let parents = parents
        .into_iter()
        .map(|fsnode| *fsnode.fsnode_id())
        .collect::<Vec<_>>();

    validate_new_manifest_entries(
        ctx,
        real_blobstore,
        fsnode,
        parents,
        mem_blob,
        |tree_id| Some(tree_id.blobstore_key()),
        |_| None,
    )
    .await?;

    Ok(())
}

async fn validate_skeleton_manifests<'a>(
    ctx: &'a CoreContext,
    real_repo: &'a BlobRepo,
    cs_id: ChangesetId,
    mem_blob: &'a Arc<dyn Blobstore>,
) -> Result<(), Error> {
    let real_blobstore = real_repo.repo_blobstore_arc();

    let (skeleton_manifest, parents) =
        find_cs_and_parents_derived_data::<RootSkeletonManifestId>(ctx, real_repo, cs_id).await?;
    let skeleton_manifest = *skeleton_manifest.skeleton_manifest_id();
    let parents = parents
        .into_iter()
        .map(|skeleton_manifest| *skeleton_manifest.skeleton_manifest_id())
        .collect::<Vec<_>>();

    validate_new_manifest_entries(
        ctx,
        real_blobstore,
        skeleton_manifest,
        parents,
        mem_blob,
        |tree_id| Some(tree_id.blobstore_key()),
        |_| None,
    )
    .await?;

    Ok(())
}

async fn validate_unodes<'a>(
    ctx: &'a CoreContext,
    real_repo: &'a BlobRepo,
    cs_id: ChangesetId,
    mem_blob: &'a Arc<dyn Blobstore>,
) -> Result<(), Error> {
    let real_blobstore = real_repo.repo_blobstore_arc();
    let (unode, parents) =
        find_cs_and_parents_derived_data::<RootUnodeManifestId>(ctx, real_repo, cs_id).await?;
    let unode = *unode.manifest_unode_id();
    let parents = parents
        .into_iter()
        .map(|unode| *unode.manifest_unode_id())
        .collect::<Vec<_>>();
    validate_new_manifest_entries(
        ctx,
        real_blobstore,
        unode,
        parents,
        mem_blob,
        |tree_id| Some(tree_id.blobstore_key()),
        |leaf_id| Some(leaf_id.blobstore_key()),
    )
    .await?;

    Ok(())
}

async fn validate_hgchangesets<'a>(
    ctx: &'a CoreContext,
    real_repo: &'a BlobRepo,
    cs_id: ChangesetId,
    mem_blob: &'a Arc<dyn Blobstore>,
) -> Result<(), Error> {
    let real_blobstore = real_repo.repo_blobstore_arc();

    let (derived, parents) =
        find_cs_and_parents_derived_data::<MappedHgChangesetId>(ctx, real_repo, cs_id).await?;
    let mem_blob = &mem_blob;
    let manifest = async {
        let hgchangeset = derived
            .hg_changeset_id()
            .load(ctx, real_repo.repo_blobstore())
            .await?;
        check_exists(
            ctx,
            mem_blob,
            hgchangeset.get_changeset_id().blobstore_key(),
        )
        .await?;
        Result::<_, Error>::Ok(hgchangeset.manifestid())
    };
    let parents = try_join_all(parents.into_iter().map(|p| async move {
        let p = p
            .hg_changeset_id()
            .load(ctx, real_repo.repo_blobstore())
            .await?;
        Result::<_, Error>::Ok(p.manifestid())
    }));

    let (manifest, parents) = try_join(manifest, parents).await?;

    let mf_entries =
        find_intersection_of_diffs_and_parents(ctx.clone(), real_blobstore, manifest, parents)
            .try_filter_map(|(_, entry, parent_entries)| async move {
                match entry {
                    Entry::Leaf((ty, filenode_id)) => {
                        for p in parent_entries {
                            if let Entry::Leaf((_ty, parent_filenode_id)) = p {
                                // This is mode-only change, no new blobstore writes were made
                                if parent_filenode_id == filenode_id {
                                    return Ok(None);
                                }
                            }
                        }
                        Ok(Some(Entry::Leaf((ty, filenode_id))))
                    }
                    Entry::Tree(manifest_id) => Ok(Some(Entry::Tree(manifest_id))),
                }
            })
            .try_collect::<Vec<_>>()
            .await?;

    for entry in mf_entries {
        let key = match entry {
            Entry::Tree(mf_id) => mf_id.blobstore_key(),
            Entry::Leaf((_, filenode_id)) => filenode_id.blobstore_key(),
        };

        check_exists(ctx, mem_blob, key).await?;
    }

    Ok(())
}

async fn find_cs_and_parents_derived_data<D: BonsaiDerived>(
    ctx: &CoreContext,
    repo: &BlobRepo,
    cs_id: ChangesetId,
) -> Result<(D, Vec<D>), Error> {
    let parents = repo.changeset_fetcher_arc().get_parents(ctx, cs_id).await?;

    let derived = D::derive(ctx, repo, cs_id).await?;
    let parents = try_join_all(parents.into_iter().map(|p| async move {
        let derived = D::derive(ctx, repo, p).await?;
        Result::<_, Error>::Ok(derived)
    }))
    .await?;

    Ok((derived, parents))
}

async fn validate_new_manifest_entries<TreeId, LeafId>(
    ctx: &CoreContext,
    real_blobstore: Arc<dyn Blobstore>,
    mfid: TreeId,
    parent_mfids: Vec<TreeId>,
    mem_blob: &Arc<dyn Blobstore>,
    tree_blob_key: impl Fn(TreeId) -> Option<String>,
    leaf_blob_key: impl Fn(LeafId) -> Option<String>,
) -> Result<(), Error>
where
    TreeId: StoreLoadable<Arc<dyn Blobstore>> + Clone + Send + Sync + Eq + Unpin + 'static,
    <TreeId as StoreLoadable<Arc<dyn Blobstore>>>::Value:
        Manifest<TreeId = TreeId, LeafId = LeafId> + Send,
    LeafId: Clone + Send + Eq + Unpin + 'static,
{
    let mf_entries = find_intersection_of_diffs(ctx.clone(), real_blobstore, mfid, parent_mfids)
        .map_ok(|(_, entry)| entry)
        .try_collect::<Vec<_>>()
        .await?;

    for entry in mf_entries {
        let maybe_key = match entry {
            Entry::Tree(tree_id) => tree_blob_key(tree_id),
            Entry::Leaf(leaf_id) => leaf_blob_key(leaf_id),
        };

        if let Some(key) = maybe_key {
            check_exists(ctx, mem_blob, key).await?;
        }
    }

    Ok(())
}

async fn check_exists(
    ctx: &CoreContext,
    mem_blob: &Arc<dyn Blobstore>,
    key: String,
) -> Result<(), Error> {
    let maybe_value = mem_blob.get(ctx, &key).await?;

    if maybe_value.is_none() {
        return Err(anyhow!("{} not found", key));
    }

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

//! This dummy crate contains dummy implementation of traits that are being used only in the
//! --dry-run mode to test the healer

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use context::CoreContext;
use metaconfig_types::MultiplexId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::Timestamp;
use slog::info;
use slog::Logger;

#[derive(Debug)]
pub struct DummyBlobstore<B> {
    inner: B,
    logger: Logger,
}

impl<B: std::fmt::Display> std::fmt::Display for DummyBlobstore<B> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "DummyBlobstore<{}>", &self.inner)
    }
}

impl<B: Blobstore> DummyBlobstore<B> {
    pub fn new(inner: B, logger: Logger) -> Self {
        Self { inner, logger }
    }
}

#[async_trait]
impl<B: Blobstore> Blobstore for DummyBlobstore<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.inner.get(ctx, key).await
    }

    async fn put<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        info!(
            self.logger,
            "I would have written blob {} of size {}",
            key,
            value.len()
        );
        Ok(())
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.inner.is_present(ctx, key).await
    }
}

pub struct DummyBlobstoreWal<Q> {
    inner: Q,
}

impl<Q: BlobstoreWal> DummyBlobstoreWal<Q> {
    pub fn new(inner: Q) -> Self {
        Self { inner }
    }
}

#[async_trait]
impl<Q: BlobstoreWal> BlobstoreWal for DummyBlobstoreWal<Q> {
    async fn log_many<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entries: Vec<BlobstoreWalEntry>,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        let entries_str: Vec<_> = entries.iter().map(|e| format!("{:?}", e)).collect();
        info!(
            ctx.logger(),
            "I would have written {}",
            entries_str.join(",\n")
        );
        Ok(entries)
    }

    async fn read<'a>(
        &'a self,
        ctx: &'a CoreContext,
        multiplex_id: &MultiplexId,
        older_than: &Timestamp,
        limit: usize,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        self.inner.read(ctx, multiplex_id, older_than, limit).await
    }

    async fn delete<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entries: &'a [BlobstoreWalEntry],
    ) -> Result<()> {
        let entries: Vec<_> = entries.iter().map(|e| format!("{:?}", e)).collect();
        info!(ctx.logger(), "I would have deleted {}", entries.join(",\n"));
        Ok(())
    }

    async fn delete_by_key(&self, ctx: &CoreContext, entries: &[BlobstoreWalEntry]) -> Result<()> {
        let entries: Vec<_> = entries.iter().map(|e| format!("{:?}", e)).collect();
        info!(
            ctx.logger(),
            "I would have deleted by key {}",
            entries.join(",\n")
        );
        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::iter::Sum;
use std::ops::Add;
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use chrono::Duration as ChronoDuration;
use context::CoreContext;

pub(crate) const MIN_FETCH_FAILURE_DELAY: Duration = Duration::from_millis(1);
pub(crate) const MAX_FETCH_FAILURE_DELAY: Duration = Duration::from_millis(100);
pub(crate) const DEFAULT_BLOB_SIZE_BYTES: u64 = 1024;

#[derive(Clone, Debug)]
pub struct HealResult {
    pub processed_full_batch: bool,
    pub processed_rows: u64,
}

#[async_trait]
pub trait Healer {
    async fn heal(&self, ctx: &CoreContext, minimum_age: ChronoDuration) -> Result<HealResult>;
}

#[derive(Default, Debug, PartialEq)]
pub(crate) struct HealStats {
    pub queue_add: usize,
    pub queue_del: usize,
    pub put_success: usize,
    pub put_failure: usize,
}

impl Add for HealStats {
    type Output = Self;
    fn add(self, other: Self) -> Self {
        Self {
            queue_add: self.queue_add + other.queue_add,
            queue_del: self.queue_del + other.queue_del,
            put_success: self.put_success + other.put_success,
            put_failure: self.put_failure + other.put_failure,
        }
    }
}

impl Sum for HealStats {
    fn sum<I: Iterator<Item = HealStats>>(iter: I) -> HealStats {
        iter.fold(Default::default(), Add::add)
    }
}


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

#![cfg_attr(not(fbcode_build), allow(unused_crate_dependencies))]
#![feature(never_type)]

mod dummy;
mod healer;
mod wal_healer;

use std::collections::HashMap;
use std::num::NonZeroU64;
use std::ops::Bound;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore_factory::make_blobstore;
use blobstore_factory::BlobstoreOptions;
use blobstore_factory::ReadOnlyStorage;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::SqlBlobstoreWal;
use borrowed::borrowed;
use cached_config::ConfigStore;
use chrono::Duration as ChronoDuration;
use clap::Parser;
use context::CoreContext;
use context::SessionContainer;
use dummy::DummyBlobstore;
use dummy::DummyBlobstoreWal;
use fbinit::FacebookInit;
use futures::future;
use futures_03_ext::BufferedParams;
use healer::HealResult;
use healer::Healer;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexedStoreType;
use metaconfig_types::ShardedDatabaseConfig;
use metaconfig_types::StorageConfig;
use mononoke_app::fb303::AliveService;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;
use slog::info;
use slog::o;
use sql_construct::SqlConstructFromShardedDatabaseConfig;
use sql_ext::facebook::MysqlOptions;
use wait_for_replication::WaitForReplication;
use wal_healer::WalHealer;

#[derive(Parser)]
#[clap(about = "Monitors blobstore_sync_queue to heal blobstores with missing data")]
struct MononokeBlobstoreHealerArgs {
    /// set limit for how many queue entries to process
    #[clap(long, default_value_t = 10000)]
    sync_queue_limit: usize,
    /// performs a single healing and prints what would it do without doing it
    #[clap(long)]
    dry_run: bool,
    /// drain the queue without healing. Use with caution.
    #[clap(long)]
    drain_only: bool,
    /// id of storage group to be healed, e.g. manifold_xdb_multiplex
    #[clap(long)]
    storage_id: String,
    /// Log a lot less
    #[clap(long, short = 'q')]
    quiet: bool,
    /// If specified, only perform the given number of iterations
    #[clap(long)]
    iteration_limit: Option<u64>,
    /// Seconds. If specified, override default minimum age to heal of 120 seconds
    #[clap(long, default_value_t = 120)]
    heal_min_age_secs: i64,
    /// How maby blobs to heal concurrently.
    #[clap(long, default_value_t = 100)]
    heal_concurrency: usize,
    /// max combined size of concurrently healed blobs (approximate, will still let individual larger blobs through)
    #[clap(long, default_value_t = 10_000_000_000)]
    heal_max_bytes: u64,
    /// which shards to read from, useful for spawning multiple independent healers
    #[clap(long, default_value = "..")]
    shard_range: ShardRange,
}

#[derive(Clone, Copy)]
struct ShardRange {
    left: Bound<usize>,
    right: Bound<usize>,
}

fn parse_opt(s: &str) -> Result<Option<usize>> {
    if s.is_empty() {
        Ok(None)
    } else {
        Ok(Some(s.parse()?))
    }
}

impl std::str::FromStr for ShardRange {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        Ok(if let Some((left, right)) = s.split_once("..=") {
            let right = Bound::Included(right.parse()?);
            let left = parse_opt(left)?.map_or(Bound::Unbounded, Bound::Included);
            Self { left, right }
        } else if let Some((left, right)) = s.split_once("..") {
            let left = parse_opt(left)?.map_or(Bound::Unbounded, Bound::Included);
            let right = parse_opt(right)?.map_or(Bound::Unbounded, Bound::Excluded);
            Self { left, right }
        } else {
            let num = s.parse()?;
            Self {
                left: Bound::Included(num),
                right: Bound::Included(num),
            }
        })
    }
}

impl ShardRange {
    // We could also implement SliceIndex, but that's unsafe.
    fn into_slice_index(self) -> (Bound<usize>, Bound<usize>) {
        (self.left, self.right)
    }
}

async fn maybe_schedule_healer_for_storage(
    fb: FacebookInit,
    ctx: &CoreContext,
    dry_run: bool,
    drain_only: bool,
    blobstore_sync_queue_limit: usize,
    buffered_params: BufferedParams,
    storage_config: StorageConfig,
    mysql_options: &MysqlOptions,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &BlobstoreOptions,
    iter_limit: Option<u64>,
    heal_min_age: ChronoDuration,
    config_store: &ConfigStore,
    shard_range: ShardRange,
) -> Result<(), Error> {
    let multiplex_healer = match storage_config.clone().blobstore {
        BlobConfig::MultiplexedWal {
            blobstores,
            multiplex_id,
            queue_db,
            inner_blobstores_scuba_table,
            scuba_sample_rate,
            ..
        } => {
            let wal = setup_wal(
                fb,
                mysql_options,
                queue_db,
                readonly_storage,
                dry_run,
                shard_range,
            )?;
            let blobstores = setup_blobstores(
                fb,
                ctx,
                blobstores,
                mysql_options,
                blobstore_options,
                readonly_storage,
                inner_blobstores_scuba_table,
                scuba_sample_rate,
                config_store,
                dry_run,
            )
            .await?;

            let healer: Arc<dyn Healer> = Arc::new(WalHealer::new(
                blobstore_sync_queue_limit,
                buffered_params,
                wal,
                Arc::new(blobstores),
                multiplex_id,
                drain_only,
            ));
            Result::<_, Error>::Ok(healer)
        }
        s => bail!("Storage doesn't use Multiplexed blobstore, got {:?}", s),
    }?;

    let wait_for_replication = WaitForReplication::new(fb, config_store, storage_config, "healer")?;

    schedule_healing(
        ctx,
        multiplex_healer,
        wait_for_replication,
        iter_limit,
        heal_min_age,
    )
    .await
}

fn setup_wal(
    fb: FacebookInit,
    mysql_options: &MysqlOptions,
    queue_db: ShardedDatabaseConfig,
    readonly_storage: ReadOnlyStorage,
    dry_run: bool,
    shard_range: ShardRange,
) -> Result<Arc<dyn BlobstoreWal>> {
    let wal = SqlBlobstoreWal::with_sharded_database_config(
        fb,
        &queue_db,
        mysql_options,
        readonly_storage.0,
    )
    .context("While opening WAL")?;

    let wal: Arc<dyn BlobstoreWal> = if dry_run {
        Arc::new(DummyBlobstoreWal::new(wal))
    } else {
        Arc::new(wal.with_read_range(shard_range.into_slice_index())?)
    };

    Ok(wal)
}

async fn setup_blobstores(
    fb: FacebookInit,
    ctx: &CoreContext,
    blobstore_configs: Vec<(BlobstoreId, MultiplexedStoreType, BlobConfig)>,
    mysql_options: &MysqlOptions,
    blobstore_options: &BlobstoreOptions,
    readonly_storage: ReadOnlyStorage,
    scuba_table: Option<String>,
    scuba_sample_rate: NonZeroU64,
    config_store: &ConfigStore,
    dry_run: bool,
) -> Result<HashMap<BlobstoreId, Arc<dyn Blobstore>>> {
    let blobstores = blobstore_configs.into_iter().map({
        borrowed!(scuba_table);
        move |(id, _, blobconfig)| async move {
            let blobconfig = BlobConfig::Logging {
                blobconfig: Box::new(blobconfig),
                scuba_table: scuba_table.clone(),
                scuba_sample_rate,
            };

            let blobstore = make_blobstore(
                fb,
                blobconfig,
                mysql_options,
                readonly_storage,
                blobstore_options,
                ctx.logger(),
                config_store,
                &blobstore_factory::default_scrub_handler(),
                None,
            )
            .await?;

            let blobstore: Arc<dyn Blobstore> = if dry_run {
                let logger = ctx.logger().new(o!("blobstore" => format!("{:?}", id)));
                Arc::new(DummyBlobstore::new(blobstore, logger))
            } else {
                blobstore
            };

            Result::<_, Error>::Ok((id, blobstore))
        }
    });

    let blobstores = future::try_join_all(blobstores)
        .await?
        .into_iter()
        .collect::<HashMap<_, _>>();
    Ok(blobstores)
}

// Pass None as iter_limit for never ending run
async fn schedule_healing(
    ctx: &CoreContext,
    multiplex_healer: Arc<dyn Healer>,
    wait_for_replication: WaitForReplication,
    iter_limit: Option<u64>,
    heal_min_age: ChronoDuration,
) -> Result<(), Error> {
    let mut count = 0;
    let healing_start_time = Instant::now();
    let mut total_deleted_rows = 0;

    loop {
        let iteration_start_time = Instant::now();
        count += 1;
        if let Some(iter_limit) = iter_limit {
            if count > iter_limit {
                return Ok(());
            }
        }

        wait_for_replication
            .wait_for_replication(ctx.logger())
            .await
            .context("While waiting for replication")?;

        let HealResult {
            processed_full_batch,
            processed_rows,
        } = multiplex_healer
            .heal(ctx, heal_min_age)
            .await
            .context("While healing")?;

        total_deleted_rows += processed_rows;
        let total_elapsed = healing_start_time.elapsed().as_secs_f32();
        let iteration_elapsed = iteration_start_time.elapsed().as_secs_f32();
        info!(
            ctx.logger(),
            "Iteration rows processed: {} rows, {}s; total: {} rows, {}s",
            processed_rows,
            iteration_elapsed,
            total_deleted_rows,
            total_elapsed,
        );

        // if last batch read was not full,  wait at least 1 second, to avoid busy looping as don't
        // want to hammer the database with thousands of reads a second.
        if !processed_full_batch {
            info!(ctx.logger(), "The last batch was not full size, waiting...",);
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    let app = MononokeAppBuilder::new(fb)
        .with_app_extension(Fb303AppExtension {})
        .build::<MononokeBlobstoreHealerArgs>()?;

    app.run_with_monitoring_and_logging(async_main, "blobstore_healer", AliveService)
}

async fn async_main(app: MononokeApp) -> Result<(), Error> {
    let args: MononokeBlobstoreHealerArgs = app.args()?;
    let env = app.environment();

    let storage_id = args.storage_id;
    let logger = app.logger();
    let config_store = app.config_store();
    let mysql_options = &env.mysql_options;
    let readonly_storage = env.readonly_storage;
    let blobstore_options = &env.blobstore_options;
    let storage_configs = app.storage_configs();
    let storage_config = storage_configs
        .storage
        .get(&storage_id)
        .ok_or_else(|| format_err!("Storage id `{}` not found", storage_id))?;
    let blobstore_sync_queue_limit = args.sync_queue_limit;
    let heal_concurrency = args.heal_concurrency;
    let heal_max_bytes = args.heal_max_bytes;
    let dry_run = args.dry_run;
    let drain_only = args.drain_only;

    let iter_limit = args.iteration_limit;
    let healing_min_age = ChronoDuration::seconds(args.heal_min_age_secs);
    let quiet = args.quiet;
    if !quiet {
        info!(logger, "Using storage_config {:?}", storage_config);
    }

    let scuba = env.scuba_sample_builder.clone();

    let ctx = SessionContainer::new_with_defaults(app.fb).new_context(logger.clone(), scuba);
    let buffered_params = BufferedParams {
        weight_limit: heal_max_bytes,
        buffer_size: heal_concurrency,
    };
    let shard_range = args.shard_range;

    maybe_schedule_healer_for_storage(
        app.fb,
        &ctx,
        dry_run,
        drain_only,
        blobstore_sync_queue_limit,
        buffered_params,
        storage_config.clone(),
        mysql_options,
        readonly_storage,
        blobstore_options,
        iter_limit,
        healing_min_age,
        config_store,
        shard_range,
    )
    .await
}


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

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use chrono::Duration as ChronoDuration;
use cloned::cloned;
use context::CoreContext;
use futures::future::join_all;
use futures::stream;
use futures::FutureExt;
use futures::StreamExt;
use futures_03_ext::BufferedParams;
use futures_03_ext::FbStreamExt;
use itertools::Itertools;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::DateTime;
use mononoke_types::Timestamp;
use rand::thread_rng;
use rand::Rng;
use slog::info;
use slog::warn;

use crate::healer::HealResult;
use crate::healer::Healer;
use crate::healer::DEFAULT_BLOB_SIZE_BYTES;
use crate::healer::MAX_FETCH_FAILURE_DELAY;
use crate::healer::MIN_FETCH_FAILURE_DELAY;

#[cfg(test)]
mod tests;

/// How many times to put a blob back in the queue
/// if it couldn't be found.
const MAX_WAL_RETRIES: u32 = 20;

pub struct WalHealer {
    /// The amount of entries healer processes in one go.
    batch_size: usize,
    /// Dynamic batch size, that helps to recover in case of failures.
    current_fetch_size: AtomicUsize,
    /// The buffered params are specified in order to balance our concurrent and parallel
    /// blobs' healing according to their sizes.
    buffered_params: BufferedParams,
    /// Write-ahead log that is used by multiplexed storage to synchronize the data.
    wal: Arc<dyn BlobstoreWal>,
    /// The blob storages healer operates on.
    blobstores: Arc<HashMap<BlobstoreId, Arc<dyn Blobstore>>>,
    /// Multiplex configuration id.
    multiplex_id: MultiplexId,
    /// Drain the queue without healing. Use with caution.
    drain_only: bool,
}

impl WalHealer {
    pub fn new(
        batch_size: usize,
        buffered_params: BufferedParams,
        wal: Arc<dyn BlobstoreWal>,
        blobstores: Arc<HashMap<BlobstoreId, Arc<dyn Blobstore>>>,
        multiplex_id: MultiplexId,
        drain_only: bool,
    ) -> Self {
        Self {
            batch_size,
            current_fetch_size: AtomicUsize::new(batch_size),
            buffered_params,
            wal,
            blobstores,
            multiplex_id,
            drain_only,
        }
    }

    async fn fetch_entries(
        &self,
        ctx: &CoreContext,
        older_than: Timestamp,
    ) -> Result<(usize, Vec<BlobstoreWalEntry>)> {
        let mut fetch_size = self.current_fetch_size.load(Ordering::Relaxed);
        loop {
            match self
                .wal
                .read(ctx, &self.multiplex_id, &older_than, fetch_size)
                .await
            {
                Ok(queue_entries) => {
                    // Success. Update fetch size for next loop
                    let new_fetch_size =
                        if fetch_size == self.current_fetch_size.load(Ordering::Relaxed) {
                            // Fetch size didn't change during the loop, which implies that we succeeded
                            // on the first attempt (since all failures decrease it) - increase it for next
                            // time if it's not yet at the limit. Growth is at least 1 each loop
                            // so that if fetch_size / 10 == 0, we still climb back
                            self.batch_size.min(fetch_size + fetch_size / 10 + 1)
                        } else {
                            fetch_size
                        };

                    self.current_fetch_size
                        .store(new_fetch_size, Ordering::Relaxed);

                    return Ok((fetch_size, queue_entries));
                }
                Err(e) => {
                    // Error, so fall in size fast
                    let new_fetch_size = fetch_size / 2;
                    warn!(
                        ctx.logger(),
                        "Failed to read full batch from th WAL, failing in batch size: old {}, new {}",
                        fetch_size,
                        new_fetch_size
                    );

                    if new_fetch_size < 1 {
                        return Err(e);
                    }

                    fetch_size = new_fetch_size;
                    let delay =
                        thread_rng().gen_range(MIN_FETCH_FAILURE_DELAY..MAX_FETCH_FAILURE_DELAY);
                    tokio::time::sleep(delay).await;
                }
            }
        }
    }

    pub async fn heal_impl(
        &self,
        ctx: &CoreContext,
        minimum_age: ChronoDuration,
    ) -> Result<HealResult> {
        let now = DateTime::now().into_chrono();
        let older_than = DateTime::new(now - minimum_age);
        let (batch_size, queue_entries) = self.fetch_entries(ctx, older_than.into()).await?;

        // all entries in the queue correspond to the different put operations
        let unique_puts = queue_entries.len();
        info!(
            ctx.logger(),
            "Fetched {} distinct put operations", unique_puts
        );

        if unique_puts == 0 {
            info!(ctx.logger(), "All caught up, nothing to do");
            return Ok(HealResult {
                processed_full_batch: false,
                processed_rows: 0,
            });
        }

        if self.drain_only {
            info!(
                ctx.logger(),
                "Draining all entries without healing. [drain-only mode]"
            );
            let deleted_entries =
                cleanup_after_healing(ctx, self.wal.as_ref(), queue_entries).await?;

            return Ok(HealResult {
                processed_full_batch: true,
                processed_rows: deleted_entries,
            });
        }

        let unique_blobstore_keys = queue_entries
            .iter()
            .unique_by(|entry| entry.blobstore_key.clone())
            .count();

        let healing_futures: Vec<(_, u64)> = queue_entries
            .into_iter()
            .sorted_by_key(|entry| entry.blobstore_key.clone())
            .group_by(|entry| entry.blobstore_key.clone())
            .into_iter()
            .map(|(key, entries)| {
                let entries: Vec<_> = entries.into_iter().collect();
                let healing_weight = entries
                    .first()
                    // The "or" never happens. Can be fixed with vec1.
                    .map_or(DEFAULT_BLOB_SIZE_BYTES, |entry| entry.blob_size);

                let fut =
                    heal_blob(ctx, self.blobstores.clone(), key).map(|outcome| (outcome, entries));

                (fut.boxed(), healing_weight)
            })
            .collect();

        info!(
            ctx.logger(),
            "Found {} blobs to be healed... Doing it with weight limit {}, max concurrency: {}",
            healing_futures.len(),
            self.buffered_params.weight_limit,
            self.buffered_params.buffer_size,
        );

        let heal_res: Vec<_> = stream::iter(healing_futures)
            .buffered_weight_limited(self.buffered_params)
            .collect()
            .await;

        let mut healthy_blobs = 0;
        let mut to_enqueue = vec![];
        let mut missing_blobs = vec![];
        let processed_entries: Vec<_> = heal_res
            .into_iter()
            .flat_map(|(outcome, entries)| {
                match outcome {
                    HealBlobOutcome::Healthy => {
                        // blob was healthy and did not require to be healed
                        healthy_blobs += 1;
                    }
                    HealBlobOutcome::Healed => {} // all done!
                    HealBlobOutcome::MissingBlob(key) => {
                        // the blob is missing, we'll requeue the entries in case the blob
                        // was not propagated yet to the blobstores
                        //
                        // TODO: log missing blobs to scuba
                        let retries = entries.first().map_or(0, |e| e.retry_count);
                        warn!(
                            ctx.logger(),
                            "Missing blob detected: key {} ({} retries so far)", key, retries
                        );
                        to_enqueue.push(
                            entries
                                .iter()
                                .filter(|e| e.retry_count < MAX_WAL_RETRIES)
                                .cloned()
                                .update(BlobstoreWalEntry::increment_retry)
                                .collect(),
                        );
                        missing_blobs.push(key);
                    }
                    HealBlobOutcome::MissingBlobstores(key, blobstores) => {
                        info!(
                            ctx.logger(),
                            "Couldn't heal blob {} in these blobstores: {:?}", key, blobstores
                        );
                        to_enqueue.push(entries.clone());
                    }
                }
                entries
            })
            .collect();

        let blobs_not_healed = to_enqueue.len();
        let to_enqueue: Vec<_> = to_enqueue.into_iter().flatten().collect();

        info!(
            ctx.logger(),
            "For {} processed entries and {} blobstore keys: healthy blobs {}, healed blobs {}, failed to heal {}, missing blobs {}",
            processed_entries.len(),
            unique_blobstore_keys,
            healthy_blobs,
            unique_blobstore_keys - blobs_not_healed - healthy_blobs,
            blobs_not_healed,
            missing_blobs.len()
        );

        enqueue_entries(ctx, self.wal.as_ref(), to_enqueue).await?;
        let deleted_entries =
            cleanup_after_healing(ctx, self.wal.as_ref(), processed_entries).await?;

        Ok(HealResult {
            processed_full_batch: unique_puts == batch_size,
            processed_rows: deleted_entries,
        })
    }
}

#[async_trait]
impl Healer for WalHealer {
    async fn heal(&self, ctx: &CoreContext, minimum_age: ChronoDuration) -> Result<HealResult> {
        self.heal_impl(ctx, minimum_age).await
    }
}

struct HealingBlob {
    blob_get_data: Option<BlobstoreGetData>,
    missing_blobstores: HashMap<BlobstoreId, Arc<dyn Blobstore>>,
    failing_blobstores: Vec<BlobstoreId>,
}

enum HealBlobOutcome {
    // The blob was found in all of the blobstores and did not require healing
    Healthy,
    // The blobstores missing the blob were successfully healed
    Healed,
    // The blob was not found in any of the blobstores
    MissingBlob(String),
    // The blobstores that don't have the blob and could not be healed. These
    // are mainly the blobstores that failed on healing write.
    MissingBlobstores(String, HashSet<BlobstoreId>),
}

async fn heal_blob(
    ctx: &CoreContext,
    blobstores: Arc<HashMap<BlobstoreId, Arc<dyn Blobstore>>>,
    key: String,
) -> HealBlobOutcome {
    let HealingBlob {
        blob_get_data,
        missing_blobstores,
        failing_blobstores,
    } = fetch_blob(ctx, blobstores, &key).await;

    let blob = if let Some(get_data) = blob_get_data {
        // we successfully found the blob in at least one of the blobstores
        if missing_blobstores.is_empty() && failing_blobstores.is_empty() {
            // the blob is well and healthy across all the blobstores
            return HealBlobOutcome::Healthy;
        }

        get_data.into_bytes()
    } else {
        // We couldn't find the blob in the stores: it's either missing,
        // or the blobstores with the blob are failing
        if failing_blobstores.is_empty() {
            // no blobstore failed, the blob is missing
            return HealBlobOutcome::MissingBlob(key);
        }
        return HealBlobOutcome::MissingBlobstores(key, missing_blobstores.into_keys().collect());
    };

    let missing = try_heal(ctx, key.clone(), blob, missing_blobstores).await;
    if missing.is_empty() {
        return HealBlobOutcome::Healed;
    }
    HealBlobOutcome::MissingBlobstores(key, missing)
}

async fn fetch_blob(
    ctx: &CoreContext,
    blobstores: Arc<HashMap<BlobstoreId, Arc<dyn Blobstore>>>,
    key: &str,
) -> HealingBlob {
    let gets = join_all(blobstores.iter().map(|(bid, blobstore)| async move {
        let result = blobstore.get(ctx, key).await;
        (bid, result)
    }))
    .await;

    let mut blob = None;
    let mut failing_blobstores = vec![];
    let missing_storages_ids: HashSet<_> = gets
        .into_iter()
        .filter_map(|(bid, result)| match result {
            Ok(None) => Some(bid),
            Ok(Some(get_data)) => {
                if blob.is_none() {
                    blob = Some(get_data);
                }
                None
            }
            Err(_) => {
                failing_blobstores.push(bid.clone());
                Some(bid)
            }
        })
        .collect();
    let missing_blobstores = blobstores
        .iter()
        .filter_map(|(bid, blobstore)| {
            if missing_storages_ids.contains(bid) {
                Some((bid.clone(), blobstore.clone()))
            } else {
                None
            }
        })
        .collect();

    HealingBlob {
        blob_get_data: blob,
        missing_blobstores,
        failing_blobstores,
    }
}

async fn try_heal(
    ctx: &CoreContext,
    key: String,
    blob: BlobstoreBytes,
    missing_blobstores: HashMap<BlobstoreId, Arc<dyn Blobstore>>,
) -> HashSet<BlobstoreId> {
    let puts = missing_blobstores.into_iter().map(|(bid, blobstore)| {
        cloned!(key, blob);
        async move {
            let put_result = blobstore.put(ctx, key, blob).await;
            (bid, put_result)
        }
    });
    let results = join_all(puts).await;

    results
        .into_iter()
        .filter_map(
            |(bid, put_result)| {
                if put_result.is_ok() { None } else { Some(bid) }
            },
        )
        .collect()
}

async fn enqueue_entries(
    ctx: &CoreContext,
    wal: &dyn BlobstoreWal,
    entries: Vec<BlobstoreWalEntry>,
) -> Result<()> {
    if entries.is_empty() {
        return Ok(());
    }
    info!(
        ctx.logger(),
        "Requeueing {} queue entries for another healing attempt",
        entries.len()
    );
    let new_entries = entries
        .into_iter()
        .map(|entry| {
            let BlobstoreWalEntry {
                blobstore_key,
                multiplex_id,
                blob_size,
                ..
            } = entry;

            BlobstoreWalEntry::new(blobstore_key, multiplex_id, Timestamp::now(), blob_size)
        })
        .collect();

    wal.log_many(ctx, new_entries).await?;
    Ok(())
}

/// Removed healed entries from the queue.
async fn cleanup_after_healing(
    ctx: &CoreContext,
    wal: &dyn BlobstoreWal,
    entries: Vec<BlobstoreWalEntry>,
) -> Result<u64> {
    let n = entries.len() as u64;
    info!(ctx.logger(), "Deleting {} actioned queue entries", n);
    wal.delete(ctx, &entries).await?;
    Ok(n)
}


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

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

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use blobstore::BlobstoreGetData;
use blobstore_sync_queue::SqlBlobstoreWal;
use bytes::Bytes;
use context::CoreContext;
use fbinit::FacebookInit;
use futures_03_ext::BufferedParams;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use sql_construct::SqlConstruct;

use super::*;
use crate::wal_healer::WalHealer;

#[derive(Clone, Debug, Default)]
struct GoodBlob {
    inner: Arc<Mutex<HashMap<String, BlobstoreBytes>>>,
}

impl std::fmt::Display for GoodBlob {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "GoodBlob")
    }
}

#[async_trait]
impl Blobstore for GoodBlob {
    async fn put<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        let mut inner = self.inner.lock().expect("lock poison");
        inner.insert(key, value);
        Ok(())
    }

    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let inner = self.inner.lock().expect("lock poison");
        let bytes = inner.get(key).map(|bytes| bytes.clone().into());
        Ok(bytes)
    }
}

#[derive(Clone, Debug)]
struct FailingBlob;

impl std::fmt::Display for FailingBlob {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "FailingBlob")
    }
}

#[async_trait]
impl Blobstore for FailingBlob {
    async fn put<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _key: String,
        _value: BlobstoreBytes,
    ) -> Result<()> {
        anyhow::bail!("Failed put!");
    }

    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        anyhow::bail!("Failed get!");
    }
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(FailingBlob);
    let bs2 = Arc::new(FailingBlob);
    let bs3 = Arc::new(FailingBlob);
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();

    // all blobatores fail, so it doesn't matter whether we actually write the blob or not

    // the queue will have an entry for the previous write
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 12);
    wal.log_many(&ctx, vec![entry]).await?;

    let expected = vec![key.clone()];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue have the entry, because the blob couldn't be healed
    let expected = vec![key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3 = Arc::new(GoodBlob::default());
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();
    let value = make_value("value");

    // make sure blob is available in each of the blobstores
    bs1.put(&ctx, key.clone(), value.clone()).await?;
    bs2.put(&ctx, key.clone(), value.clone()).await?;
    bs3.put(&ctx, key.clone(), value.clone()).await?;

    // the queue will have an entry for the previous write
    // it can even have multiple entries, if the blob was written twice
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry1 = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 13);
    let entry2 = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 14);
    wal.log_many(&ctx, vec![entry1, entry2]).await?;

    let expected = vec![key.clone(), key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue doesn't have any entries, because the blob is healthy
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), vec![]).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3 = Arc::new(GoodBlob::default());
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();
    let value = make_value("value");

    // make sure blob is available at least in one of the blobstores
    bs1.put(&ctx, key.clone(), value.clone()).await?;

    // the queue will have an entry for the previous write
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 15);
    wal.log_many(&ctx, vec![entry]).await?;

    let expected = vec![key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue is empty, the blob was healed
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), vec![]).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3: Arc<dyn Blobstore> = Arc::new(FailingBlob);
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();
    let value = make_value("value");

    // make sure blob is available at least in one of the blobstores
    bs1.put(&ctx, key.clone(), value.clone()).await?;

    // the queue will have an entry for the previous write
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 16);
    wal.log_many(&ctx, vec![entry]).await?;

    let expected = vec![key.clone()];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue still has the entry, because the blob couldn't be healed
    // in the failing blobstore
    let expected = vec![key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3: Arc<dyn Blobstore> = Arc::new(FailingBlob);
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();

    // we assume the blob is only available in the failing blobstore

    // the queue will have an entry for the previous write
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 17);
    wal.log_many(&ctx, vec![entry]).await?;

    let expected = vec![key.clone()];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue still has the entry, because the blob couldn't be healed
    let expected = vec![key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3 = Arc::new(GoodBlob::default());
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let mid1 = MultiplexId::new(1);
    let mid2 = MultiplexId::new(2);

    let ts = Timestamp::now();
    let key1 = "key1".to_string();
    let value1 = make_value("value1");
    let key2 = "key2".to_string();
    let value2 = make_value("value2");
    let key3 = "key3".to_string();
    let value3 = make_value("value3");

    // first blob will be available in 1 blobstore
    bs1.put(&ctx, key1.clone(), value1.clone()).await?;
    // second blob will be available in 2 blobstore
    bs2.put(&ctx, key2.clone(), value2.clone()).await?;
    // third blob will be available in 3 blobstore
    bs3.put(&ctx, key3.clone(), value3.clone()).await?;

    // the queue will have an entry for the previous writes and some other
    // entries from different multiplex configuration
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);

    let entry1 = BlobstoreWalEntry::new(key1.clone(), mid1, ts, 18);
    let entry2 = BlobstoreWalEntry::new(key2.clone(), mid1, ts, 19);
    let entry3 = BlobstoreWalEntry::new(key3.clone(), mid2, ts, 20);

    wal.log_many(&ctx, vec![entry1, entry2, entry3]).await?;

    let expected = vec![key1, key2];
    validate_queue(&ctx, wal.clone(), mid1, Timestamp::now(), expected).await?;
    let expected = vec![key3.clone()];
    validate_queue(&ctx, wal.clone(), mid2, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, mid1, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue has only entries from different multiplex configuration
    validate_queue(&ctx, wal.clone(), mid1, Timestamp::now(), vec![]).await?;
    let expected = vec![key3.clone()];
    validate_queue(&ctx, wal.clone(), mid2, Timestamp::now(), expected).await?;

    // also check that the third blob wasn't healed
    assert!(bs1.get(&ctx, &key3).await?.is_none());
    assert!(bs2.get(&ctx, &key3).await?.is_none());

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2 = Arc::new(GoodBlob::default());
    let bs3 = Arc::new(GoodBlob::default());
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
            (BlobstoreId::new(3), bs3.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let ts = Timestamp::now();
    let key = "key".to_string();

    // the queue will have an entry for the previous write
    // it can even have multiple entries, if the blob was written twice
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, ts, 21);
    wal.log_many(&ctx, vec![entry]).await?;

    let expected = vec![key.clone()];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue has the entry, because the blob is completely
    // missing (all blobstore reads succeeded, but couldn't find the blob) and we are
    // unable to heal it now, but maybe it wasn't yet written to the storages
    let expected = vec![key];
    validate_queue(&ctx, wal.clone(), multiplex_id, Timestamp::now(), expected).await?;

    Ok(())
}

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

    let bs1: Arc<dyn Blobstore> = Arc::new(GoodBlob::default());
    let bs2: Arc<dyn Blobstore> = Arc::new(FailingBlob);
    let blobstores: Arc<HashMap<_, _>> = Arc::new(
        vec![
            (BlobstoreId::new(1), bs1.clone()),
            (BlobstoreId::new(2), bs2.clone()),
        ]
        .into_iter()
        .collect(),
    );

    // set up some variables
    let multiplex_id = MultiplexId::new(1);

    let original_ts = Timestamp::now();
    let key = "key".to_string();
    let value = make_value("value");

    bs1.put(&ctx, key.clone(), value.clone()).await?;

    // the queue will have an entry for the previous write
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);
    let entry = BlobstoreWalEntry::new(key.clone(), multiplex_id, original_ts, 22);
    wal.log_many(&ctx, vec![entry]).await?;

    let original_entries = wal
        .read(&ctx, &multiplex_id, &Timestamp::now(), 100)
        .await?;
    assert!(original_entries.len() == 1);
    let original_entry = &original_entries[0];
    assert_eq!(original_ts, original_entry.timestamp);

    let buf_params = BufferedParams {
        weight_limit: 1000,
        buffer_size: 100,
    };
    let healer = WalHealer::new(10, buf_params, wal.clone(), blobstores, multiplex_id, false);

    let age = ChronoDuration::seconds(0);
    healer.heal(&ctx, age).await?;

    // check that the queue entry has a new timestamp
    let new_entries = wal
        .read(&ctx, &multiplex_id, &Timestamp::now(), 100)
        .await?;
    assert!(new_entries.len() == 1);
    let new_entry = &new_entries[0];
    assert!(original_entry.timestamp < new_entry.timestamp);

    Ok(())
}

async fn validate_queue<'a>(
    ctx: &'a CoreContext,
    wal: Arc<dyn BlobstoreWal>,
    multiplex_id: MultiplexId,
    older_than: Timestamp,
    mut expected: Vec<String>,
) -> Result<()> {
    let mut entries: Vec<_> = wal
        .read(ctx, &multiplex_id, &older_than, 100)
        .await?
        .into_iter()
        .map(|e| e.blobstore_key)
        .collect();

    assert_eq!(entries.len(), expected.len());

    entries.sort();
    expected.sort();

    for (a, b) in entries.iter().zip(expected.iter()) {
        assert_eq!(a, b);
    }

    Ok(())
}

fn make_value(value: &str) -> BlobstoreBytes {
    BlobstoreBytes::from_bytes(Bytes::copy_from_slice(value.as_bytes()))
}


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

use anyhow::Context;
use anyhow::Result;
use mercurial_types::HgChangesetId;
use mononoke_types::DateTime;
use serde_derive::Deserialize;
use toml::value;

/// Configuration for the bonsai verify tool.
#[derive(Clone, Debug, Default)]
pub struct Config {
    /// Which changesets are known to be broken and therefore should skip verification.
    pub ignores: Vec<HgChangesetId>,

    /// Old versions of treemanifest had a bug with merges: *some* p2 nodes would be set to null,
    /// while others would have the correct p2 manifest. This affects manifest hashes, but not
    /// correctness.
    ///
    /// If this is set and the commit time is before this, on a root manifest mismatch with a merge
    /// commit the verifier will check to see if there are any file changes. If there are, it will
    /// return "valid but different hash".
    ///
    /// Using the commit time is a somewhat unfortunate proxy for when that commit was converted
    /// to treemanifest, but it's the best we have.
    ///
    /// To treat all merges as potentially broken in this way, set this to a time in the future.
    pub broken_merges_before: Option<DateTime>,
}

impl Config {
    fn new(config_serde: ConfigSerde) -> Result<Self> {
        let broken_merges_before = match config_serde.broken_merges_before {
            Some(dt) => Some(
                DateTime::from_rfc3339(&dt.to_string())
                    .context("error while parsing broken_merges_before")?,
            ),
            None => None,
        };

        Ok(Self {
            ignores: config_serde.ignores,
            broken_merges_before,
        })
    }
}

#[derive(Clone, Debug, Deserialize)]
struct ConfigSerde {
    ignores: Vec<HgChangesetId>,
    broken_merges_before: Option<value::Datetime>,
}

pub fn get_config(config_file: Option<String>) -> Result<Config> {
    let config_file = match config_file {
        Some(config_file) => config_file,
        None => return Ok(Config::default()),
    };

    let config_toml = std::fs::read_to_string(config_file)?;
    let config_serde: ConfigSerde =
        toml::from_str(&config_toml).context("error while reading config")?;
    Config::new(config_serde)
}


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

mod config;
use std::collections::HashSet;
use std::io::Write;
use std::process;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Instant;

use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobrepo_utils::BonsaiMFVerify;
use blobrepo_utils::BonsaiMFVerifyResult;
use blobstore::Loadable;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use changeset_fetcher::ChangesetFetcherArc;
use clap::Parser;
use clap::Subcommand;
use cloned::cloned;
use context::CoreContext;
use failure_ext::DisplayChain;
use fbinit::FacebookInit;
use futures::compat::Future01CompatExt;
use futures::compat::Stream01CompatExt;
use futures::future;
use futures::future::TryFutureExt;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_old::future::Either;
use futures_old::future::{self as old_future};
use futures_old::Future;
use futures_old::Stream;
use lock_ext::LockExt;
use mercurial_derivation::get_manifest_from_bonsai;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgChangesetId;
use mononoke_app::args::RepoArgs;
use mononoke_app::MononokeAppBuilder;
use repo_blobstore::RepoBlobstoreArc;
use repo_blobstore::RepoBlobstoreRef;
use revset::AncestorsNodeStream;
use slog::debug;
use slog::error;
use slog::info;
use slog::warn;
use slog::Logger;

#[derive(Parser)]
struct CommandArgs {
    #[clap(flatten)]
    repo: RepoArgs,

    #[clap(subcommand)]
    subcmd: BonsaiSubCommand,
}

#[derive(Subcommand)]
enum BonsaiSubCommand {
    RoundTrip(RoundTrip),
    HgManifest(HgManifest),
}

/// Verify that bonsai changesets roundtrip correctly
#[derive(Parser)]
struct RoundTrip {
    /// How many changesets to follow before stopping
    #[clap(long, default_value_t = 1024)]
    limit: usize,

    /// Print list of changed entries between manifests
    #[clap(long)]
    changes: bool,

    /// Debug mode
    #[clap(long)]
    debug: bool,

    /// Configuration file, see source code for spec
    #[clap(long, value_name = "TOML")]
    config: Option<String>,

    /// Changesets from which to start traversing
    #[clap(required = true)]
    start_points: Vec<HgChangesetId>,
}

/// Verify generation of various things
#[derive(Parser)]
struct HgManifest {
    /// Starting point of traversal
    hg_changeset_id: HgChangesetId,

    /// Count of changeset to traverse
    count: usize,
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    let app = MononokeAppBuilder::new(fb).build::<CommandArgs>()?;
    let args: CommandArgs = app.args()?;
    let runtime = app.runtime();
    let repo = runtime.block_on(app.open_repo(&args.repo))?;
    let logger = app.logger();
    let ctx = CoreContext::new_with_logger(fb, logger.clone());

    match args.subcmd {
        BonsaiSubCommand::RoundTrip(args) => {
            subcommand_round_trip(ctx, logger.clone(), runtime, repo, args)
        }
        BonsaiSubCommand::HgManifest(args) => {
            subcommmand_hg_manifest_verify(&ctx, runtime, &repo, args)
        }
    }
}

fn subcommand_round_trip(
    ctx: CoreContext,
    logger: Logger,
    runtime: &tokio::runtime::Handle,
    repo: BlobRepo,
    args: RoundTrip,
) -> Result<()> {
    let config = config::get_config(args.config).expect("getting configuration failed");
    let start_points = args.start_points;
    let follow_limit = args.limit;
    let print_changes = args.changes;
    let debug_bonsai_diff = args.debug && args.changes;

    let valid = Arc::new(AtomicUsize::new(0));
    let invalid = Arc::new(AtomicUsize::new(0));
    let errors = Arc::new(AtomicUsize::new(0));
    let ignored = Arc::new(AtomicUsize::new(0));
    // The number of changesets at the end should be really small, so an unbounded buffer doesn't
    // matter much.
    let (end_sender, end_receiver) = ::std::sync::mpsc::channel();

    let verify_fut = old_future::lazy(|| {
        let logger = logger.clone();
        let valid = valid.clone();
        let invalid = invalid.clone();
        let errors = errors.clone();
        let ignored = ignored.clone();
        let bonsai_verify = BonsaiMFVerify {
            ctx: ctx.clone(),
            logger: logger.clone(),
            repo,
            follow_limit,
            ignores: config.ignores.into_iter().collect(),
            broken_merges_before: config.broken_merges_before,
            debug_bonsai_diff,
        };
        bonsai_verify
            .verify(start_points)
            .and_then({
                cloned!(ctx, logger);
                move |(result, meta)| {
                    let logger =
                        logger.new(slog::o!["changeset_id" => format!("{}", meta.changeset_id)]);

                    if !result.is_ignored() {
                        let followed = follow_limit - meta.follow_remaining;
                        if followed % 10000 == 0 {
                            info!(
                                logger,
                                "Followed {} changesets, {} remaining",
                                followed,
                                meta.follow_remaining,
                            );
                        }
                        if meta.follow_remaining == 0 {
                            end_sender
                                .send(meta.changeset_id)
                                .expect("end_receiver is still alive");
                        }
                    }

                    match &result {
                        BonsaiMFVerifyResult::Valid { .. } => {
                            debug!(logger, "VALID");
                            valid.fetch_add(1, Ordering::Relaxed);
                            Either::A(old_future::ok(()))
                        }
                        BonsaiMFVerifyResult::ValidDifferentId(difference) => {
                            debug!(
                                logger,
                                "VALID but with a different hash: \
                                expected manifest ID: {}, roundtrip ID: {}",
                                difference.expected_mf_id,
                                difference.roundtrip_mf_id,
                            );
                            valid.fetch_add(1, Ordering::Relaxed);
                            Either::A(old_future::ok(()))
                        }
                        BonsaiMFVerifyResult::Invalid(difference) => {
                            warn!(logger, "INVALID");
                            info!(
                                logger, "manifest hash differs";
                                "expected manifest ID" => difference.expected_mf_id,
                                "roundtrip ID" => difference.roundtrip_mf_id,
                            );
                            invalid.fetch_add(1, Ordering::Relaxed);
                            if print_changes {
                                let logger = logger.clone();
                                let diff_fut = difference
                                    .changes(ctx.clone())
                                    .map(move |changed_entry| {
                                        info!(logger, "Change: {:?}", changed_entry,);
                                    })
                                    .collect()
                                    .map(|_| ());
                                Either::B(diff_fut)
                            } else {
                                Either::A(old_future::ok(()))
                            }
                        }
                        BonsaiMFVerifyResult::Ignored(..) => {
                            ignored.fetch_add(1, Ordering::Relaxed);
                            Either::A(old_future::ok(()))
                        }
                    }
                }
            })
            .then(move |res| {
                // collect() below will stop after the first error, but we care about all errors.
                // So report them now and keep returning Ok.
                if let Err(err) = &res {
                    error!(logger, "ERROR: {}", DisplayChain::from(err));
                    errors.fetch_add(1, Ordering::Relaxed);
                }
                Ok::<_, ()>(())
            })
            // collect to turn the stream into a future that will finish when the stream is done
            .collect()
    });

    let _ = runtime.block_on(verify_fut.compat());

    let end_points: Vec<_> = end_receiver.into_iter().collect();
    process::exit(summarize(
        logger, end_points, valid, invalid, errors, ignored,
    ));
}

fn summarize(
    logger: Logger,
    end_points: Vec<HgChangesetId>,
    valid: Arc<AtomicUsize>,
    invalid: Arc<AtomicUsize>,
    errors: Arc<AtomicUsize>,
    ignored: Arc<AtomicUsize>,
) -> i32 {
    let end_points: Vec<_> = end_points
        .iter()
        .map(|changeset_id| format!("{}", changeset_id))
        .collect();
    let valid = valid.load(Ordering::Acquire);
    let invalid = invalid.load(Ordering::Acquire);
    let errors = errors.load(Ordering::Acquire);
    let ignored = ignored.load(Ordering::Acquire);
    let total = valid + invalid + errors;
    let percent_valid = 100.0 * (valid as f64) / (total as f64);

    let logger = logger.new(slog::o!["summary" => ""]);

    info!(
        logger,
        "{:.2}% valid", percent_valid;
        "ignored" => ignored,
        "errors" => errors,
        "valid" => valid,
        "total" => total,
    );

    if !end_points.is_empty() {
        info!(
            logger,
            "To resume verification, run with arguments: {}",
            end_points.join(" "),
        );
    }

    // Return the appropriate exit code for this process.
    if errors > 0 {
        2
    } else if invalid > 0 {
        1
    } else {
        0
    }
}

fn subcommmand_hg_manifest_verify(
    ctx: &CoreContext,
    runtime: &tokio::runtime::Handle,
    repo: &BlobRepo,
    args: HgManifest,
) -> Result<()> {
    let total = &AtomicUsize::new(0);
    let total_millis = &AtomicU64::new(0);
    let bad = &Mutex::new(HashSet::new());

    let run = async move {
        let count = args.count;
        let hg_csid = args.hg_changeset_id;
        let csid = repo
            .bonsai_hg_mapping()
            .get_bonsai_from_hg(ctx, hg_csid)
            .await?
            .ok_or_else(|| format_err!("failed to fetch bonsai changeset"))?;

        AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), csid)
            .compat()
            .take(count)
            .map(|res| async move {
                match res {
                    Ok(csid) => {
                        let cs_id = repo.derive_hg_changeset(ctx, csid).await?;
                        let bonsai_fut = csid.load(ctx, repo.repo_blobstore()).map_err(Error::from);

                        let parents_fut = async move {
                            let blob_cs = cs_id.load(ctx, repo.repo_blobstore()).await?;
                            let expected = blob_cs.manifestid();

                            let hg_csid = blob_cs.get_changeset_id();

                            let mut parent_hashes = vec![];
                            parent_hashes.extend(blob_cs.p1());
                            parent_hashes.extend(blob_cs.p2());
                            parent_hashes.extend(blob_cs.step_parents()?);

                            let parents = future::try_join_all(
                                parent_hashes
                                    .into_iter()
                                    .map(|p| {
                                        cloned!(ctx, repo);
                                        let cs_id = HgChangesetId::new(p);
                                        async move {
                                            cs_id
                                                .load(&ctx, repo.repo_blobstore())
                                                .map_ok(|cs| cs.manifestid())
                                                .map_err(Error::from)
                                                .await
                                        }
                                    })
                                    .collect::<Vec<_>>(),
                            )
                            .await?;

                            Ok((parents, hg_csid, expected))
                        };

                        let (bonsai, (parents, hg_csid, expected)) =
                            future::try_join(bonsai_fut, parents_fut).await?;

                        let start = Instant::now();

                        get_manifest_from_bonsai(
                            ctx.clone(),
                            repo.repo_blobstore_arc(),
                            bonsai.clone(),
                            parents,
                        )
                        .map_ok(move |result| {
                            if result != expected {
                                println!(
                                    "\x1b[KBAD hg_cisd:{} result:{} expected:{}\x1b[m",
                                    hg_csid, result, expected,
                                );
                                bad.with(|bad| bad.insert((csid, hg_csid, result, expected)));
                            }

                            let all = total.fetch_add(1, Ordering::SeqCst) + 1;
                            total_millis
                                .fetch_add(start.elapsed().as_millis() as u64, Ordering::SeqCst);
                            print!(
                                "\x1b[K {} total:{} bad:{} mean_time:{:.2} ms \r",
                                hg_csid,
                                all,
                                bad.with(|bad| bad.len()),
                                total_millis.load(Ordering::SeqCst) as f32 / all as f32,
                            );
                            std::io::stdout().flush().expect("flush on stdout failed");
                        })
                        .await
                    }
                    Err(e) => Err(e),
                }
            })
            .buffer_unordered(100)
            .try_for_each(|_| async { Ok(()) })
            .map_ok(move |_| {
                let bad = bad.with(|bad| std::mem::take(bad));
                if bad.is_empty() {
                    println!()
                } else {
                    println!("\n BAD: {:#?}", bad)
                }
            })
            .await
    };

    runtime.block_on(run)
}


/*
 * 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::str::FromStr;

use anyhow::Result;
use check_git_wc::check_git_wc;
use clap::Parser;
use context::CoreContext;
use fbinit::FacebookInit;
use git2::Repository;
use git2::RepositoryOpenFlags;
use mononoke_app::args::RepoArgs;
use mononoke_app::fb303::AliveService;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;
use mononoke_types::ChangesetId;
use repo_blobstore::RepoBlobstore;
use repo_derived_data::RepoDerivedData;

#[facet::container]
struct HgRepo {
    #[facet]
    repo_blobstore: RepoBlobstore,

    #[facet]
    repo_derived_data: RepoDerivedData,
}

/// Validate the working copy in Git
#[derive(Parser)]
#[clap(about = "Tool for checking that a working copy will match a git checkout.")]
struct CheckGitWCArgs {
    /// The repo against which the command needs to be executed
    #[clap(flatten)]
    repo: RepoArgs,
    /// Bonsai changeset whose working copy should be verified
    #[clap(long, value_name = "BONSAI")]
    csid: String,
    /// Path to the git repo to compare to
    #[clap(long, value_name = "PATH")]
    git_repo_path: String,
    /// The git commit to compare to
    #[clap(long, value_name = "HASH")]
    git_commit: String,
    /// Enable git-lfs pointer parsing
    #[clap(long)]
    git_lfs: bool,
    /// Maximum number of directories to check in parallel. Default 100
    #[clap(long, default_value_t = 100)]
    scheduled_max: usize,
}

async fn run_check_git_wc(app: MononokeApp) -> Result<()> {
    let logger = app.logger();
    let ctx = CoreContext::new_with_logger(app.fb, logger.clone());

    let args: CheckGitWCArgs = app.args()?;
    let cs = ChangesetId::from_str(&args.csid)?;

    let git_repo = Repository::open_ext(
        args.git_repo_path,
        RepositoryOpenFlags::NO_SEARCH | RepositoryOpenFlags::BARE | RepositoryOpenFlags::NO_DOTGIT,
        std::iter::empty::<std::ffi::OsString>(),
    )?;

    let hg_repo: HgRepo = app.open_repo(&args.repo).await?;

    check_git_wc(
        &ctx,
        &hg_repo,
        cs,
        git_repo,
        args.git_commit,
        args.git_lfs,
        args.scheduled_max,
    )
    .await
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    let app = MononokeAppBuilder::new(fb)
        .with_app_extension(Fb303AppExtension {})
        .build::<CheckGitWCArgs>()?;
    app.run_with_monitoring_and_logging(run_check_git_wc, "check_git_wc", AliveService)
}


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

//! Tool to copy a list of blobs from one repo to another.
//! The tool is quite dumb - it just copies the blobs in any order so e.g. aliases might be copied
//! after the file content was copied. So it should be used only in the urgent situations
//! like e.g. repo corruption. List of keys are passed in a file, and this list of keys
//! can be generated by any tool e.g. walker.

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use blobrepo::BlobRepo;
use blobstore::Blobstore;
use blobstore::PutBehaviour;
use clap::Arg;
use cmdlib::args;
use cmdlib::args::get_config_by_repoid;
use cmdlib::args::load_common_config;
use cmdlib::args::MononokeMatches;
use cmdlib::helpers::setup_repo_dir;
use cmdlib::helpers::CreateStorage;
use context::CoreContext;
use context::SessionClass;
use fbinit::FacebookInit;
use futures::future;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use mononoke_types::RepositoryId;
use repo_blobstore::RepoBlobstoreRef;
use repo_factory::RepoFactory;
use slog::debug;
use slog::info;
use slog::warn;
use slog::Logger;
use thiserror::Error;
use tokio::fs::File;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio_stream::wrappers::LinesStream;

const ARG_CONCURRENCY: &str = "concurrency";
const ARG_ERROR_KEYS: &str = "error-keys-output";
const ARG_IGNORE_ERRORS: &str = "ignore-errors";
const ARG_IN_FILE: &str = "input-file";
const ARG_MISSING_KEYS: &str = "missing-keys-output";
const ARG_SUCCESSFUL_KEYS: &str = "success-keys-output";
const ARG_STRIP_SOURCE_REPO_PREFIX: &str = "strip-source-repo-prefix";
const ARG_TARGET_INNER_BLOBSTORE_ID: &str = "target-inner-blobstore-id";
const ARG_SOURCE_INNER_BLOBSTORE_ID: &str = "source-inner-blobstore-id";

struct OutputFiles {
    error_file: File,
    missing_file: File,
    successful_file: File,
}

impl OutputFiles {
    pub async fn new(matches: &MononokeMatches<'_>) -> Result<Self, Error> {
        let error_file = Self::open_file(matches, ARG_ERROR_KEYS);
        let missing_file = Self::open_file(matches, ARG_MISSING_KEYS);
        let successful_file = Self::open_file(matches, ARG_SUCCESSFUL_KEYS);

        let (error_file, missing_file, successful_file) =
            future::try_join3(error_file, missing_file, successful_file).await?;

        Ok(Self {
            error_file,
            missing_file,
            successful_file,
        })
    }

    pub async fn record_copy_result(
        &mut self,
        key: &str,
        res: Result<(), CopyError>,
    ) -> Result<(), Error> {
        let file = match res {
            Ok(()) => &mut self.successful_file,
            Err(CopyError::NotFound) => &mut self.missing_file,
            Err(CopyError::Error(_)) => &mut self.error_file,
        };

        file.write_all(key.as_bytes()).await?;
        file.write_all(b"\n").await?;

        res.with_context(|| format!("failed to copy {}", key))?;

        Ok(())
    }

    async fn open_file(matches: &MononokeMatches<'_>, arg: &str) -> Result<File, Error> {
        let filename = matches
            .value_of(arg)
            .ok_or_else(|| anyhow!("{} is not set", arg))?;
        let file = File::create(filename).await?;

        Ok(file)
    }
}

#[derive(Error, Debug)]
enum CopyError {
    #[error("Not found")]
    NotFound,
    #[error(transparent)]
    Error(#[from] Error),
}

async fn run<'a>(fb: FacebookInit, matches: &'a MononokeMatches<'a>) -> Result<(), Error> {
    let logger = matches.logger();
    let mut ctx = CoreContext::new_with_logger(fb, logger.clone());
    // Background session class tells multiplexed blobstore to wait
    // for all blobstores to finish.
    // TODO(stash): T89677672 at the moment Background session class writes to
    // blobstore sync queue if blobstore has failed. This might not be what we
    // want here.
    ctx.session_mut()
        .override_session_class(SessionClass::Background);

    let source_repo_id =
        args::not_shardmanager_compatible::get_source_repo_id(matches.config_store(), matches)?;
    let source_inner_blobstore_id = args::get_u64_opt(matches, ARG_SOURCE_INNER_BLOBSTORE_ID);

    let source_repo = open_repo(
        ctx.logger(),
        source_repo_id,
        source_inner_blobstore_id,
        matches,
    );

    let target_repo_id =
        args::not_shardmanager_compatible::get_target_repo_id(matches.config_store(), matches)?;
    let target_inner_blobstore_id = args::get_u64_opt(matches, ARG_TARGET_INNER_BLOBSTORE_ID);

    let target_repo = open_repo(
        ctx.logger(),
        target_repo_id,
        target_inner_blobstore_id,
        matches,
    );

    let (source_repo, target_repo) = future::try_join(source_repo, target_repo).await?;

    let mut keys = vec![];
    let source_repo_prefix = source_repo_id.prefix();

    let strip_source_repo_prefix = matches.is_present(ARG_STRIP_SOURCE_REPO_PREFIX);
    let inputfile = matches
        .value_of(ARG_IN_FILE)
        .ok_or_else(|| anyhow!("{} not set", ARG_IN_FILE))?;
    let mut inputfile = File::open(inputfile).await?;
    let file = BufReader::new(&mut inputfile);
    let mut lines = LinesStream::new(file.lines());
    while let Some(line) = lines.try_next().await? {
        if strip_source_repo_prefix {
            match line.strip_prefix(&source_repo_prefix) {
                Some(key) => {
                    keys.push(key.to_string());
                }
                None => {
                    return Err(anyhow!(
                        "key {} doesn't start with prefix {}",
                        line,
                        source_repo_prefix
                    ));
                }
            }
        } else {
            keys.push(line);
        }
    }

    let concurrency = args::get_usize(matches, ARG_CONCURRENCY, 100);
    let ignore_errors = matches.is_present(ARG_IGNORE_ERRORS);

    info!(ctx.logger(), "{} keys to copy", keys.len());
    let log_step = std::cmp::max(1, keys.len() / 10);

    let mut s = stream::iter(keys)
        .map(|key| async {
            let copy_key = key.clone();
            let res = async {
                let source_blobstore = source_repo.repo_blobstore().clone();
                let target_blobstore = target_repo.repo_blobstore().clone();
                let maybe_value = source_blobstore.get(&ctx, &key).await?;
                let value = maybe_value.ok_or(CopyError::NotFound)?;
                debug!(ctx.logger(), "copying {}", key);
                target_blobstore.put(&ctx, key, value.into_bytes()).await?;
                Result::<_, CopyError>::Ok(())
            }
            .await;

            (copy_key, res)
        })
        .buffered(concurrency);

    let mut copied = 0;
    let mut processed = 0;
    let mut output_files = OutputFiles::new(matches).await?;
    while let Some((key, res)) = s.next().await {
        let res = output_files.record_copy_result(&key, res).await;
        match res {
            Ok(()) => {
                copied += 1;
            }
            Err(err) => {
                if ignore_errors {
                    warn!(ctx.logger(), "key: {} {:#}", key, err);
                } else {
                    return Err(err);
                }
            }
        };
        processed += 1;
        if processed % log_step == 0 {
            info!(ctx.logger(), "{} keys processed", processed);
        }
    }

    info!(ctx.logger(), "{} keys were copied", copied);
    Ok(())
}

async fn open_repo<'a>(
    logger: &Logger,
    repo_id: RepositoryId,
    maybe_inner_blobstore_id: Option<u64>,
    matches: &'a MononokeMatches<'a>,
) -> Result<BlobRepo, Error> {
    let config_store = matches.config_store();
    let common_config = load_common_config(config_store, matches)?;
    let (reponame, mut config) = get_config_by_repoid(config_store, matches, repo_id)?;
    if let Some(inner_id) = maybe_inner_blobstore_id {
        override_blobconfig(&mut config.storage_config.blobstore, inner_id)?;
    }
    info!(logger, "using repo \"{}\" repoid {:?}", reponame, repo_id);
    match &config.storage_config.blobstore {
        BlobConfig::Files { path } | BlobConfig::Sqlite { path } => {
            setup_repo_dir(path, CreateStorage::ExistingOnly)?;
        }
        _ => {}
    };

    let repo_factory = RepoFactory::new(matches.environment().clone());

    let repo = repo_factory.build(reponame, config, common_config).await?;

    Ok(repo)
}

fn override_blobconfig(blob_config: &mut BlobConfig, inner_blobstore_id: u64) -> Result<(), Error> {
    match blob_config {
        BlobConfig::MultiplexedWal { ref blobstores, .. } => {
            let sought_id = BlobstoreId::new(inner_blobstore_id);
            let inner_blob_config = blobstores
                .iter()
                .find_map(|(blobstore_id, _, blobstore)| {
                    if blobstore_id == &sought_id {
                        Some(blobstore.clone())
                    } else {
                        None
                    }
                })
                .ok_or_else(|| {
                    anyhow!("could not find a blobstore with id {}", inner_blobstore_id)
                })?;
            *blob_config = inner_blob_config;
        }
        _ => {
            return Err(anyhow!(
                "inner-blobstore-id supplied but blobstore is not multiplexed"
            ));
        }
    }
    Ok(())
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let (matches, runtime) = args::MononokeAppBuilder::new(
        "Tool to copy a list of blobs from one blobstore to another",
    )
    .with_advanced_args_hidden()
    .with_source_and_target_repos()
    .with_special_put_behaviour(PutBehaviour::Overwrite)
    .build()
    .arg(
        Arg::with_name(ARG_IN_FILE)
            .long(ARG_IN_FILE)
            .required(true)
            .takes_value(true)
            .help("input filename with a list of keys"),
    )
    .arg(
        Arg::with_name(ARG_CONCURRENCY)
            .long(ARG_CONCURRENCY)
            .required(false)
            .takes_value(true)
            .help("How many blobs to copy at once"),
    )
    .arg(
        Arg::with_name(ARG_IGNORE_ERRORS)
            .long(ARG_IGNORE_ERRORS)
            .required(false)
            .takes_value(false)
            .help("Don't terminate if failed to process a key"),
    )
    .arg(
        Arg::with_name(ARG_STRIP_SOURCE_REPO_PREFIX)
            .long(ARG_STRIP_SOURCE_REPO_PREFIX)
            .required(false)
            .takes_value(false)
            .help(
                "If a key starts with 'repoXXXX' prefix \
                      (where XXXX matches source repository) then strip this \
                      prefix before copying",
            ),
    )
    .arg(
        Arg::with_name(ARG_SUCCESSFUL_KEYS)
            .long(ARG_SUCCESSFUL_KEYS)
            .takes_value(true)
            .required(true)
            .help("A file to write successfully copied keys to"),
    )
    .arg(
        Arg::with_name(ARG_MISSING_KEYS)
            .long(ARG_MISSING_KEYS)
            .takes_value(true)
            .required(true)
            .help("A file to write missing keys to"),
    )
    .arg(
        Arg::with_name(ARG_ERROR_KEYS)
            .long(ARG_ERROR_KEYS)
            .takes_value(true)
            .required(true)
            .help("A file to write error fetching keys to"),
    )
    .arg(
        Arg::with_name(ARG_TARGET_INNER_BLOBSTORE_ID)
            .long(ARG_TARGET_INNER_BLOBSTORE_ID)
            .takes_value(true)
            .required(false)
            .requires(ARG_SOURCE_INNER_BLOBSTORE_ID)
            .help("In case of multiplexed blobstore this will be target id of inner blobstore"),
    )
    .arg(
        Arg::with_name(ARG_SOURCE_INNER_BLOBSTORE_ID)
            .long(ARG_SOURCE_INNER_BLOBSTORE_ID)
            .takes_value(true)
            .required(false)
            .requires(ARG_TARGET_INNER_BLOBSTORE_ID)
            .help("In case of multiplexed blobstore this will be source id of inner blobstore"),
    )
    .get_matches(fb)?;

    runtime.block_on(run(fb, &matches))
}