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

use std::io;
use std::io::BufRead;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use blobstore_factory::make_packblob;
use borrowed::borrowed;
use clap::Parser;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::stream;
use futures::stream::TryStreamExt;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use mononoke_app::args::AsRepoArg;
use mononoke_app::args::RepoArgs;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;

mod pack_utils;

#[derive(Parser)]
#[clap(
    about = "Given a set of blob names on stdin, replace them with a packed version that takes less space"
)]
struct MononokePackerArgs {
    #[clap(flatten)]
    repo_args: RepoArgs,

    #[clap(
        long,
        help = "If main blobstore in the storage config is a multiplexed one, use inner blobstore with this id"
    )]
    inner_blobstore_id: Option<u64>,

    #[clap(long, help = "zstd compression level to use")]
    zstd_level: i32,

    #[clap(
        long,
        help = "If true, do not upload the finished pack to the blobstore"
    )]
    dry_run: bool,

    #[clap(
        long,
        default_value_t = 10,
        help = "Maximum number of parallel packs to work on. Default 10"
    )]
    scheduled_max: usize,
}

const PACK_PREFIX: &str = "multiblob-";

fn get_blobconfig(
    mut blob_config: BlobConfig,
    inner_blobstore_id: Option<u64>,
) -> Result<BlobConfig> {
    // If the outer store is a mux, find th requested inner store
    if let Some(inner_blobstore_id) = inner_blobstore_id {
        blob_config = match blob_config {
            BlobConfig::MultiplexedWal { blobstores, .. } => {
                let required_id = BlobstoreId::new(inner_blobstore_id);
                blobstores
                    .into_iter()
                    .find_map(|(blobstore_id, _, blobstore)| {
                        if blobstore_id == required_id {
                            Some(blobstore)
                        } else {
                            None
                        }
                    })
                    .with_context(|| {
                        format!("could not find a blobstore with id {}", inner_blobstore_id)
                    })?
            }
            _ => bail!("inner-blobstore-id can only be supplied for multiplexed blobstores"),
        }
    };

    Ok(blob_config)
}

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

    let args: MononokePackerArgs = app.args()?;
    let inner_id = args.inner_blobstore_id;
    let zstd_level = args.zstd_level;
    let dry_run = args.dry_run;
    let max_parallelism = args.scheduled_max;

    let env = app.environment();
    let logger = app.logger();
    let runtime = app.runtime();
    let config_store = app.config_store();

    let ctx = CoreContext::new_for_bulk_processing(fb, logger.clone());
    let readonly_storage = &env.readonly_storage;
    let blobstore_options = &env.blobstore_options;

    let repo_arg = args.repo_args.as_repo_arg();
    let (_repo_name, repo_config) = app.repo_config(repo_arg)?;
    let blobconfig = repo_config.storage_config.blobstore;
    let repo_prefix = repo_config.repoid.prefix();

    let input_lines: Vec<String> = io::stdin()
        .lock()
        .lines()
        .collect::<Result<_, io::Error>>()?;

    let mut scuba = env.scuba_sample_builder.clone();
    scuba.add_opt("blobstore_id", inner_id);

    runtime.block_on(async move {
        let blobstore = make_packblob(
            fb,
            get_blobconfig(blobconfig, inner_id)?,
            *readonly_storage,
            blobstore_options,
            logger,
            config_store,
        )
        .await?;
        stream::iter(input_lines.split(String::is_empty).map(Result::Ok))
            .try_for_each_concurrent(max_parallelism, |pack_keys| {
                borrowed!(ctx, repo_prefix, blobstore, scuba);
                async move {
                    let pack_keys: Vec<&str> = pack_keys.iter().map(|i| i.as_ref()).collect();
                    pack_utils::repack_keys(
                        ctx,
                        blobstore,
                        PACK_PREFIX,
                        zstd_level,
                        repo_prefix,
                        &pack_keys,
                        dry_run,
                        scuba,
                    )
                    .await
                }
            })
            .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 anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use blobstore::BlobstoreUnlinkOps;
use context::CoreContext;
use futures::future::FutureExt;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use packblob::get_entry_compressed_size;
use packblob::EmptyPack;
use packblob::Pack;
use packblob::PackBlob;
use packblob::SingleCompressed;
use scuba_ext::MononokeScubaSampleBuilder;
use tokio::task::spawn_blocking;

type BlobsWithKeys = Vec<(String, BlobstoreBytes)>;

const BLOBSTORE_KEY: &str = "blobstore_key";
const COMPRESSED_SIZE: &str = "compressed_size";
const PACK_KEY: &str = "pack_key";
const UNCOMPRESSED_SIZE: &str = "uncompressed_size";

// Tries to pack with the first blob from `blobs` as the dictionary for the other blobs
fn try_pack(zstd_level: i32, blobs: Vec<(String, BlobstoreBytes)>) -> Result<Pack> {
    let empty_pack = EmptyPack::new(zstd_level);

    let mut blobs = blobs.into_iter();

    let (dict_key, dict_blob) = blobs.next().ok_or_else(|| anyhow!("No blobs to pack"))?;
    let mut pack = empty_pack.add_base_blob(dict_key.clone(), dict_blob)?;
    for (key, blob) in blobs {
        pack.add_delta_blob(dict_key.clone(), key, blob)?;
    }

    Ok(pack)
}

fn is_better_pack(best_so_far: Option<&Pack>, new: &Pack) -> Result<bool> {
    if let Some(best_so_far) = best_so_far {
        Ok(best_so_far.get_compressed_size()? > new.get_compressed_size()?)
    } else {
        Ok(true)
    }
}

async fn find_best_pack(mut blobs: BlobsWithKeys, zstd_level: i32) -> Result<Option<Pack>> {
    let build_packs = FuturesUnordered::new();
    for _ in 0..blobs.len() {
        build_packs.push({
            let blobs = blobs.clone();
            async { tokio::task::spawn_blocking(move || try_pack(zstd_level, blobs)).await? }
        });
        blobs.rotate_left(1);
    }

    build_packs
        .try_fold(None, |best, new| async move {
            if is_better_pack(best.as_ref(), &new)? {
                Ok(Some(new))
            } else {
                Ok(best)
            }
        })
        .await
}

async fn fetch_blobs<T: BlobstoreUnlinkOps>(
    ctx: &CoreContext,
    blobstore: &PackBlob<T>,
    repo_prefix: &str,
    keys: &[&str],
) -> Result<BlobsWithKeys> {
    let blob_fetches: FuturesUnordered<_> = keys
        .iter()
        .map(|key| async move {
            let blob = blobstore
                .get(ctx, key)
                .await?
                .ok_or_else(|| anyhow!("Blob {} not in store", key))?
                .into_bytes();
            let pack_key = key
                .strip_prefix(repo_prefix)
                .ok_or_else(|| anyhow!("Could not strip {} from {}", repo_prefix, key))?;
            Result::<_>::Ok((pack_key.to_string(), blob))
        })
        .collect();

    blob_fetches.try_collect().await
}

/// Given a list of keys to repack, convert them to a single pack
pub async fn repack_keys<T: BlobstoreUnlinkOps>(
    ctx: &CoreContext,
    blobstore: &PackBlob<T>,
    pack_prefix: &str,
    zstd_level: i32,
    repo_prefix: &str,
    keys: &[&str],
    dry_run: bool,
    scuba: &MononokeScubaSampleBuilder,
) -> Result<()> {
    let blobs = fetch_blobs(ctx, blobstore, repo_prefix, keys).await?;
    let compression_futs: FuturesUnordered<_> = blobs
        .clone()
        .into_iter()
        .map(|(key, blob)| async move {
            let uncompressed_size = blob.as_bytes().len();
            let single = spawn_blocking(move || SingleCompressed::new(zstd_level, blob)).await??;
            Result::<_, Error>::Ok((key, uncompressed_size, single))
        })
        .collect();

    let mut uncompressed_sizes = HashMap::new();
    let single_compressed: Vec<_> = compression_futs.try_collect().await?;
    let pack = if keys.len() > 1 {
        find_best_pack(blobs, zstd_level).await?
    } else {
        None
    };

    let single_compressed_size =
        single_compressed
            .iter()
            .try_fold(0usize, |size, (key, uncompressed_size, item)| {
                uncompressed_sizes.insert(key, *uncompressed_size);
                Ok::<_, Error>(size + item.get_compressed_size()?)
            })?;
    if !dry_run {
        match pack {
            Some(pack) if pack.get_compressed_size()? < single_compressed_size => {
                // gather info for logs
                let logs: Vec<MononokeScubaSampleBuilder> = pack
                    .entries()
                    .iter()
                    .map(|e| {
                        get_entry_compressed_size(e).map(|compressed_size| {
                            let mut scuba = scuba.clone();
                            scuba.add(BLOBSTORE_KEY, format!("{}{}", repo_prefix, e.key));
                            scuba.add_opt(
                                UNCOMPRESSED_SIZE,
                                uncompressed_sizes.get(&e.key).copied(),
                            );
                            scuba.add(COMPRESSED_SIZE, compressed_size);
                            scuba
                        })
                    })
                    .collect::<Result<Vec<MononokeScubaSampleBuilder>>>()?;

                // store
                let pack_key = blobstore
                    .put_packed(ctx, pack, repo_prefix.to_string(), pack_prefix.to_string())
                    .await?;

                // log what we stored
                for mut scuba in logs {
                    scuba.add(PACK_KEY, pack_key.as_str());
                    scuba.log();
                }
            }
            Some(_) | None => {
                let put_futs: FuturesUnordered<_> = single_compressed
                    .into_iter()
                    .map(|(key, uncompressed_size, value)| {
                        let key = format!("{}{}", repo_prefix, key);
                        let mut scuba = scuba.clone();
                        scuba.add(BLOBSTORE_KEY, key.as_str());
                        scuba.add(UNCOMPRESSED_SIZE, uncompressed_size);
                        let compressed_size = value.get_compressed_size();
                        blobstore.put_single(ctx, key, value).map(move |v| {
                            scuba.add(COMPRESSED_SIZE, compressed_size?);
                            scuba.log();
                            v
                        })
                    })
                    .collect();
                put_futs.try_for_each(|_| async { Ok(()) }).await?;
            }
        }
    }
    Ok(())
}


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

mononoke_app::subcommands! {
    mod generation_size;
    mod mark;
}


/*
 * 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::Result;
use clap::Parser;
use cmdlib_logging::ScribeLoggingArgs;
use fbinit::FacebookInit;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;

mod commands;
mod utils;

#[derive(Parser)]
#[clap(
    name = "SQLblob GC",
    about = "Perform garbage collection on a set of SQLblob shards"
)]
pub struct MononokeSQLBlobGCArgs {
    #[clap(flatten)]
    scribe_logging_args: ScribeLoggingArgs,
    /// The name of the storage config to GC. This *must* be an XDB storage config,
    /// or a multiplex containing an XDB (in which case, give the inner blobstore ID, too
    #[clap(long)]
    storage_config_name: String,
    /// If main blobstore in the storage config is a multiplexed one, use inner blobstore with this id
    #[clap(long)]
    inner_blobstore_id: Option<u64>,
    /// Maximum number of parallel keys to GC.  Default 100.
    #[clap(long, default_value_t = 100)]
    scheduled_max: usize,
    /// Metadata shard number to start at (or 0 if not specified)
    #[clap(long, default_value_t = 0)]
    start_shard: usize,
    /// Number of shards to walk (or all shards up to the maximum shard number if not specified
    #[clap(long)]
    shard_count: Option<usize>,
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    MononokeAppBuilder::new(fb)
        .build_with_subcommands::<MononokeSQLBlobGCArgs>(commands::subcommands())?
        .run_basic(async_main)
}

async fn async_main(app: MononokeApp) -> Result<()> {
    commands::dispatch(app).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::ops::Range;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use blobstore_factory::make_sql_blobstore;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use metaconfig_types::ShardableRemoteDatabaseConfig;
use mononoke_app::MononokeApp;
use sqlblob::Sqlblob;

use crate::MononokeSQLBlobGCArgs;

fn remove_wrapper_blobconfigs(mut blob_config: BlobConfig) -> BlobConfig {
    // Pack is a wrapper store - remove it
    while let BlobConfig::Pack { ref blobconfig, .. } = blob_config {
        blob_config = BlobConfig::clone(blobconfig);
    }
    blob_config
}

fn get_blobconfig(
    app: &MononokeApp,
    storage_config_name: String,
    inner_blobstore_id: Option<u64>,
) -> Result<BlobConfig> {
    let blob_config = {
        let storage_configs = app.storage_configs();
        let storage_config = storage_configs
            .storage
            .get(&storage_config_name)
            .context("Requested storage config not found")?;
        storage_config.blobstore.clone()
    };

    let mut blob_config = remove_wrapper_blobconfigs(blob_config);

    // If the outer store is a mux, find th requested inner store
    if let Some(inner_blobstore_id) = inner_blobstore_id {
        blob_config = match blob_config {
            BlobConfig::MultiplexedWal { blobstores, .. } => {
                let required_id = BlobstoreId::new(inner_blobstore_id);
                blobstores
                    .into_iter()
                    .find_map(|(blobstore_id, _, blobstore)| {
                        if blobstore_id == required_id {
                            Some(blobstore)
                        } else {
                            None
                        }
                    })
                    .with_context(|| {
                        format!("could not find a blobstore with id {}", inner_blobstore_id)
                    })?
            }
            _ => bail!("inner-blobstore-id can only be supplied for multiplexed blobstores"),
        }
    };

    Ok(remove_wrapper_blobconfigs(blob_config))
}

fn get_shard_range(
    blobstore_config: &BlobConfig,
    shard_start: usize,
    shard_count: Option<usize>,
) -> Range<usize> {
    let max_shards = match &blobstore_config {
        BlobConfig::Mysql {
            remote: ShardableRemoteDatabaseConfig::Sharded(config),
        } => config.shard_num.get(),
        _ => 1,
    };

    Range {
        start: shard_start,
        end: shard_start + shard_count.unwrap_or(max_shards),
    }
}

async fn get_sqlblob(app: &MononokeApp, blobstore_config: BlobConfig) -> Result<Sqlblob> {
    let config_store = app.config_store();
    let blobstore_options = app.environment().blobstore_options.clone();

    Ok(make_sql_blobstore(
        app.fb,
        blobstore_config,
        blobstore_factory::ReadOnlyStorage(false),
        &blobstore_options,
        config_store,
    )
    .await?
    .into_inner())
}

pub async fn get_sqlblob_and_shard_range(app: &MononokeApp) -> Result<(Sqlblob, Range<usize>)> {
    let common_args: MononokeSQLBlobGCArgs = app.args()?;
    let blobstore_config = get_blobconfig(
        app,
        common_args.storage_config_name,
        common_args.inner_blobstore_id,
    )?;
    let shard_start = common_args.start_shard;
    let shard_count = common_args.shard_count;
    Ok((
        get_sqlblob(app, blobstore_config.clone()).await?,
        get_shard_range(&blobstore_config, shard_start, shard_count),
    ))
}


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

use std::collections::HashMap;

use anyhow::Result;
use bytesize::ByteSize;
use clap::Parser;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::TryFutureExt;
use mononoke_app::MononokeApp;

use crate::utils;
use crate::MononokeSQLBlobGCArgs;

/// measure generation sizes
#[derive(Parser)]
pub struct CommandArgs {}

fn print_sizes(sizes: &HashMap<Option<u64>, u64>) {
    let generations = {
        let mut keys: Vec<_> = sizes.keys().collect();
        keys.sort_unstable();
        keys
    };

    println!("Generation | Size");
    println!("-----------------");

    for generation in generations {
        let size = ByteSize::b(sizes[generation]);
        let generation = match generation {
            None => "NULL".to_string(),
            Some(g) => g.to_string(),
        };
        println!("{:>10} | {}", generation, size.to_string_as(true));
    }
}

pub async fn run(app: MononokeApp, _args: CommandArgs) -> Result<()> {
    let common_args: MononokeSQLBlobGCArgs = app.args()?;

    let max_parallelism: usize = common_args.scheduled_max;

    let (sqlblob, shard_range) = utils::get_sqlblob_and_shard_range(&app).await?;

    let scuba_sample_builder = app.environment().scuba_sample_builder.clone();

    let shard_sizes: Vec<(usize, HashMap<Option<u64>, (u64, u64)>)> = stream::iter(shard_range)
        .map(|shard| {
            sqlblob
                .get_chunk_sizes_by_generation(shard)
                .map_ok(move |gen_to_size| (shard, gen_to_size))
        })
        .buffer_unordered(max_parallelism)
        .try_collect::<Vec<(usize, HashMap<Option<u64>, (u64, u64)>)>>()
        .await?;

    let mut size_summary = HashMap::new();
    for (shard, sizes) in shard_sizes {
        for (generation, (size, chunk_id_count)) in sizes.into_iter() {
            let mut sample = scuba_sample_builder.clone();
            sample.add("shard", shard);
            sample.add_opt("generation", generation);
            sample.add("size", size);
            sample.add("chunk_id_count", chunk_id_count);
            sample.log();
            *size_summary.entry(generation).or_insert(0u64) += size;
        }
    }

    if scuba_sample_builder.is_discard() {
        print_sizes(&size_summary);
    }
    Ok(())
}


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

use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use clap::Parser;
use futures::channel::mpsc;
use futures::sink::SinkExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_app::MononokeApp;
use retry::retry_always;
use slog::info;
use slog::Logger;
use sqlblob::Sqlblob;

use crate::utils;
use crate::MononokeSQLBlobGCArgs;

const BASE_RETRY_DELAY_MS: u64 = 1000;
const RETRIES: usize = 3;

/// mark referenced blobs as not safe to delete
#[derive(Parser)]
pub struct CommandArgs {
    /// Only set generation on blobs that have no generation set yet. Do not do a full mark.
    #[clap(long)]
    initial_generation_only: bool,
    /// Only do the mark; do not set generation on blobs with no generation set yet.
    #[clap(long)]
    skip_initial_generation: bool,
    /// Only set the generation, don't inline small values
    #[clap(long)]
    skip_inline_small_values: bool,
}

async fn handle_one_key(
    key: String,
    store: Arc<Sqlblob>,
    inline_small_values: bool,
    logger: Arc<Logger>,
    mark_generation: u64,
) -> Result<()> {
    retry_always(
        &logger,
        |_| store.set_generation(&key, inline_small_values, mark_generation),
        BASE_RETRY_DELAY_MS,
        RETRIES,
    )
    .await
    .with_context(|| anyhow!("Failed to handle {} after {} retries", &key, RETRIES))?;
    Ok(())
}

async fn handle_initial_generation(store: &Sqlblob, shard: usize, logger: &Logger) -> Result<()> {
    retry_always(
        logger,
        |_| store.set_initial_generation(shard),
        BASE_RETRY_DELAY_MS,
        RETRIES,
    )
    .await
    .with_context(|| {
        anyhow!(
            "Failed to handle initial generation on shard {} after {} retries",
            &shard,
            RETRIES
        )
    })?;
    Ok(())
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<()> {
    let common_args: MononokeSQLBlobGCArgs = app.args()?;

    let logger: Logger = app.logger().clone();

    let max_parallelism: usize = common_args.scheduled_max;

    let (sqlblob, shard_range) = utils::get_sqlblob_and_shard_range(&app).await?;

    if !args.skip_initial_generation {
        info!(logger, "Starting initial generation set");
        let set_initial_generation_futures: Vec<_> = shard_range
            .clone()
            .map(|shard| Ok(handle_initial_generation(&sqlblob, shard, &logger)))
            .collect();
        stream::iter(set_initial_generation_futures.into_iter())
            .try_for_each_concurrent(max_parallelism, |fut| fut)
            .await?;
        info!(logger, "Completed initial generation set");
    }

    if args.initial_generation_only {
        return Ok(());
    }

    let sqlblob = Arc::new(sqlblob);
    let logger = Arc::new(logger);

    let inline_small_values = !args.skip_inline_small_values;

    // Hold mark generation constant for run
    let mark_generation = sqlblob.get_mark_generation();

    info!(logger, "Starting marking generation {}", mark_generation);
    // Set up a task to process each key in parallel in its own task.
    let (key_channel, processor) = {
        let sqlblob = Arc::clone(&sqlblob);
        let logger = Arc::clone(&logger);
        let (tx, rx) = mpsc::channel(10);
        let task = tokio::spawn(async move {
            rx.map(Ok)
                .try_for_each_concurrent(max_parallelism, {
                    |key| {
                        let sqlblob = sqlblob.clone();
                        let logger = logger.clone();
                        async move {
                            tokio::spawn(handle_one_key(
                                key,
                                sqlblob,
                                inline_small_values,
                                logger,
                                mark_generation,
                            ))
                            .await?
                        }
                    }
                })
                .await
        });
        (tx, task)
    };

    // Foreach shard in shard_range
    for shard in shard_range {
        info!(logger, "Starting mark on data keys from shard {}", shard);
        let res = sqlblob
            .get_keys_from_shard(shard)
            .forward(key_channel.clone().sink_err_into())
            .await;
        // Report processing errors ahead of key errors - that way, we don't lose the error if the channel goes away because of an error
        if res.is_err() {
            std::mem::drop(key_channel);
            processor.await??;
            return res;
        }
    }

    // Drop the spare sender so that the processor task can exit
    std::mem::drop(key_channel);

    processor.await??;
    info!(logger, "Completed marking generation {}", mark_generation);
    Ok(())
}


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

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

use anyhow::anyhow;
use anyhow::Error;
use blobstore_factory::make_blobstore;
use blobstore_factory::BlobstoreOptions;
use blobstore_factory::ReadOnlyStorage;
use cacheblob::new_memcache_blobstore;
use cached_config::ConfigStore;
use clap_old::Arg;
use cmdlib::args;
use cmdlib::args::MononokeMatches;
use cmdlib::helpers;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::future;
use futures::future::try_join;
use futures::future::try_join_all;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use metaconfig_types::RepoConfig;
use repo_blobstore::RepoBlobstore;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::error;
use slog::info;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::facebook::MysqlOptions;
use streaming_clone::RevlogStreamingChunks;
use streaming_clone::StreamingClone;
use streaming_clone::StreamingCloneBuilder;
use tokio::time;

const REPO_ARG: &str = "repo";
const REPO_WITH_TAGS_ARG: &str = "repo-with-tags";
const PERIOD_ARG: &str = "warmup-period";

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let app = args::MononokeAppBuilder::new("Utility to keep streaming clone data warm")
        .with_advanced_args_hidden()
        .with_fb303_args()
        .build()
        .about("Utility to keep streaming clone data warm")
        .arg(
            Arg::with_name(REPO_ARG)
                .long(REPO_ARG)
                .takes_value(true)
                .required(false)
                .multiple(true)
                .help("Repository name to warm-up, and empty tag is assumed"),
        )
        .arg(
            Arg::with_name(REPO_WITH_TAGS_ARG)
                .long(REPO_WITH_TAGS_ARG)
                .takes_value(true)
                .required(false)
                .multiple(true)
                .help("Repository name with a list of tags to warmup in format REPO=tag1,tag2."),
        )
        .arg(
            Arg::with_name(PERIOD_ARG)
                .long(PERIOD_ARG)
                .takes_value(true)
                .required(false)
                .default_value("900")
                .help("Period of warmup runs in secods"),
        );
    let (matches, _runtime) = app.get_matches(fb)?;

    let logger = matches.logger();
    let ctx = CoreContext::new_with_logger(fb, logger.clone());
    helpers::block_execute(
        run(ctx, &matches),
        fb,
        &std::env::var("TW_JOB_NAME").unwrap_or_else(|_| "streaming_clone_warmup".to_string()),
        logger,
        &matches,
        cmdlib::monitoring::AliveService,
    )
}

async fn run<'a>(ctx: CoreContext, matches: &'a MononokeMatches<'a>) -> Result<(), Error> {
    let period_secs: u64 = matches
        .value_of(PERIOD_ARG)
        .ok_or_else(|| anyhow!("--{} argument is required", PERIOD_ARG))?
        .parse()?;
    let period = Duration::from_secs(period_secs);

    let mut reponames_with_tags = vec![];
    if let Some(values) = matches.values_of(REPO_ARG) {
        // Assume empty tag
        reponames_with_tags.extend(
            values
                .map(ToString::to_string)
                .map(|reponame| (reponame, None)),
        );
    }

    if let Some(values) = matches.values_of(REPO_WITH_TAGS_ARG) {
        for value in values {
            let (reponame, tags) = split_repo_with_tags(value)?;
            for tag in tags {
                reponames_with_tags.push((reponame.clone(), Some(tag)));
            }
        }
    }

    if reponames_with_tags.is_empty() {
        error!(ctx.logger(), "At least one repo had to be specified");
        return Ok(());
    }

    let config_store = matches.config_store();
    let mysql_options = matches.mysql_options();
    let blobstore_options = matches.blobstore_options();
    let configs = args::load_repo_configs(config_store, matches)?;

    let mut warmers = Vec::new();
    for (reponame, tag) in reponames_with_tags {
        let config = configs
            .repos
            .get(&reponame)
            .ok_or_else(|| anyhow!("unknown repository: {}", reponame))?;
        let warmer = StreamingCloneWarmup::new(
            ctx.clone(),
            reponame,
            tag,
            config,
            mysql_options,
            blobstore_options.clone(),
            config_store,
        )
        .await?;
        warmers.push(warmer);
    }

    let offset_delay = period / warmers.len() as u32;
    let mut tasks = Vec::new();
    for (index, warmer) in warmers.into_iter().enumerate() {
        let ctx = ctx.clone();
        tasks.push(async move {
            // spread fetches over period, to reduce memory consumption
            time::sleep(offset_delay * index as u32).await;
            warmer.warmer_task(ctx.clone(), period).await?;
            Ok::<_, Error>(())
        });
    }
    try_join_all(tasks).await?;
    Ok(())
}

fn split_repo_with_tags(s: &str) -> Result<(String, Vec<String>), Error> {
    if let Some((reponame, tags)) = s.split_once('=') {
        let tags = tags.split(',').map(|s| s.to_string()).collect();

        Ok((reponame.to_string(), tags))
    } else {
        Err(anyhow!("invalid format for repo with tags: {}", s))
    }
}

struct StreamingCloneWarmup {
    streaming_clone: StreamingClone,
    reponame: String,
    tag: Option<String>,
}

impl StreamingCloneWarmup {
    async fn new(
        ctx: CoreContext,
        reponame: String,
        tag: Option<String>,
        config: &RepoConfig,
        mysql_options: &MysqlOptions,
        blobstore_options: BlobstoreOptions,
        config_store: &ConfigStore,
    ) -> Result<Self, Error> {
        // Create blobstore that contains streaming clone chunks, without cachelib
        // layer (we want to hit memcache even if it is available in cachelib), and
        // with memcache layer identical to production setup.
        let blobstore = make_blobstore(
            ctx.fb,
            config.storage_config.blobstore.clone(),
            mysql_options,
            ReadOnlyStorage(true),
            &blobstore_options,
            ctx.logger(),
            config_store,
            &blobstore_factory::default_scrub_handler(),
            None,
        )
        .await?;
        let blobstore = new_memcache_blobstore(ctx.fb, blobstore, "multiplexed", "")?;
        let repo_blobstore = Arc::new(RepoBlobstore::new(
            Arc::new(blobstore),
            None,
            config.repoid,
            MononokeScubaSampleBuilder::with_discard(),
        ));

        // Because we want to use our custom blobstore, we must construct the
        // streaming clone attribute directly.
        let streaming_clone = StreamingCloneBuilder::with_metadata_database_config(
            ctx.fb,
            &config.storage_config.metadata,
            mysql_options,
            true, /*read-only*/
        )?
        .build(config.repoid, repo_blobstore);

        Ok(Self {
            streaming_clone,
            reponame,
            tag,
        })
    }

    /// Periodically fetch streaming clone data
    async fn warmer_task(&self, ctx: CoreContext, period: Duration) -> Result<(), Error> {
        if let Some(ref tag) = self.tag {
            info!(ctx.logger(), "[{}:{}] warmer started", self.reponame, tag);
        } else {
            info!(ctx.logger(), "[{}] warmer started", self.reponame);
        };

        loop {
            let tag = None;
            let start = Instant::now();
            let chunks = self
                .streaming_clone
                .fetch_changelog(ctx.clone(), tag)
                .await?;
            info!(
                ctx.logger(),
                "[{}] index fetched in: {:.2?}",
                self.reponame,
                start.elapsed()
            );

            let size = chunks_warmup(ctx.clone(), chunks).await? as f32;
            let duration = start.elapsed();
            info!(
                ctx.logger(),
                "[{}] fetching complete in: time:{:.2?} speed:{:.1?} b/s size: {}",
                self.reponame,
                duration,
                size / duration.as_secs_f32(),
                size,
            );

            // sleep if needed
            if duration < period {
                let delay = period - duration;
                info!(
                    ctx.logger(),
                    "[{}] sleeping for: {:?}", self.reponame, delay
                );
                time::sleep(delay).await;
            }
        }
    }
}

async fn chunks_warmup(ctx: CoreContext, chunks: RevlogStreamingChunks) -> Result<usize, Error> {
    let RevlogStreamingChunks {
        index_blobs,
        data_blobs,
        index_size: index_size_expected,
        data_size: data_size_expected,
    } = chunks;

    let index = stream::iter(index_blobs.into_iter().map(|f| f.map_ok(|b| b.len())))
        .buffer_unordered(100)
        .try_fold(0usize, |acc, size| future::ok(acc + size));

    let data = stream::iter(data_blobs.into_iter().map(|f| f.map_ok(|b| b.len())))
        .buffer_unordered(100)
        .try_fold(0usize, |acc, size| future::ok(acc + size));

    let (index_size, data_size) = try_join(index, data).await?;
    if index_size_expected != index_size {
        error!(
            ctx.logger(),
            "incorrect index size: expected:{} received:{}", index_size_expected, index_size
        );
    }
    if data_size_expected != data_size {
        error!(
            ctx.logger(),
            "incorrect data size: expected:{} received:{}", data_size_expected, data_size
        );
    }
    Ok(index_size + data_size)
}


/*
 * 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::fs::File;
use std::io::BufRead;
use std::io::BufReader;
use std::str::FromStr;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use backsyncer::backsync_latest;
use backsyncer::format_counter;
use backsyncer::open_backsyncer_dbs;
use backsyncer::BacksyncLimit;
use backsyncer::Repo;
use blobrepo_hg::BlobRepoHg;
use bookmarks::BookmarkUpdateLogRef;
use bookmarks::Freshness;
use clap::Arg;
use clap::SubCommand;
use cloned::cloned;
use cmdlib::args;
use cmdlib::args::MononokeMatches;
use cmdlib::helpers;
use cmdlib_x_repo::create_commit_syncer_from_matches;
use context::CoreContext;
use context::SessionContainer;
use cross_repo_sync::CandidateSelectionHint;
use cross_repo_sync::CommitSyncContext;
use cross_repo_sync::CommitSyncOutcome;
use cross_repo_sync::CommitSyncer;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use executor_lib::ShardedProcessExecutor;
use fbinit::FacebookInit;
use futures::future::FutureExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::try_join;
use live_commit_sync_config::CfgrLiveCommitSyncConfig;
use live_commit_sync_config::LiveCommitSyncConfig;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use once_cell::sync::OnceCell;
use repo_identity::RepoIdentityRef;
use scuba_ext::MononokeScubaSampleBuilder;
use sharding_ext::RepoShard;
use slog::debug;
use slog::error;
use slog::info;
use stats::prelude::*;
use synced_commit_mapping::SqlSyncedCommitMapping;
use synced_commit_mapping::SyncedCommitMapping;
use tokio::runtime::Runtime;
use wireproto_handler::TargetRepoDbs;

const ARG_MODE_BACKSYNC_FOREVER: &str = "backsync-forever";
const ARG_MODE_BACKSYNC_ALL: &str = "backsync-all";
const ARG_MODE_BACKSYNC_COMMITS: &str = "backsync-commits";
const ARG_BATCH_SIZE: &str = "batch-size";
const ARG_INPUT_FILE: &str = "INPUT_FILE";
const SCUBA_TABLE: &str = "mononoke_xrepo_backsync";

define_stats! {
    prefix = "mononoke.backsyncer";
    remaining_entries: dynamic_singleton_counter(
        "{}.{}.remaining_entries",
        (source_repo_name: String, target_repo_name: String)
    ),
    delay_secs: dynamic_singleton_counter(
        "{}.{}.delay_secs",
        (source_repo_name: String, target_repo_name: String)
    ),
}

const SM_SERVICE_SCOPE: &str = "global";
const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;
const APP_NAME: &str = "backsyncer cmd-line tool";

/// Struct representing the Back Syncer BP.
pub struct BacksyncProcess {
    matches: Arc<MononokeMatches<'static>>,
    fb: FacebookInit,
    _runtime: Runtime,
}

impl BacksyncProcess {
    fn new(fb: FacebookInit) -> anyhow::Result<Self> {
        let app = args::MononokeAppBuilder::new(APP_NAME)
            .with_fb303_args()
            .with_source_and_target_repos()
            .with_dynamic_repos()
            .with_scribe_args()
            .build();
        let backsync_forever_subcommand = SubCommand::with_name(ARG_MODE_BACKSYNC_FOREVER)
            .about("Backsyncs all new bookmark moves");

        let sync_loop = SubCommand::with_name(ARG_MODE_BACKSYNC_COMMITS)
            .about("Syncs all commits from the file")
            .arg(
                Arg::with_name(ARG_INPUT_FILE)
                    .takes_value(true)
                    .required(true)
                    .help("list of hg commits to backsync"),
            )
            .arg(
                Arg::with_name(ARG_BATCH_SIZE)
                    .long(ARG_BATCH_SIZE)
                    .takes_value(true)
                    .required(false)
                    .help("how many commits to backsync at once"),
            );

        let backsync_all_subcommand = SubCommand::with_name(ARG_MODE_BACKSYNC_ALL)
            .about("Backsyncs all new bookmark moves once");
        let app = app
            .subcommand(backsync_all_subcommand)
            .subcommand(backsync_forever_subcommand)
            .subcommand(sync_loop);
        let (matches, _runtime) = app.get_matches(fb)?;
        let matches = Arc::new(matches);
        Ok(Self {
            matches,
            fb,
            _runtime,
        })
    }
}

#[async_trait]
impl RepoShardedProcess for BacksyncProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let logger = self.matches.logger();
        // For backsyncer, two repos (i.e. source and target) are required as input
        let source_repo_name = repo.repo_name.clone();
        let target_repo_name = match repo.target_repo_name.clone() {
            Some(repo_name) => repo_name,
            None => {
                let details = format!(
                    "Only source repo name {} provided, target repo name missing in {}",
                    source_repo_name, repo
                );
                error!(logger, "{}", details);
                bail!("{}", details)
            }
        };
        info!(
            logger,
            "Setting up back syncer command from repo {} to repo {}",
            source_repo_name,
            target_repo_name,
        );
        let details = format!(
            "Completed back syncer command setup from repo {} to repo {}",
            source_repo_name, target_repo_name
        );
        let executor = BacksyncProcessExecutor::new(
            self.fb,
            Arc::clone(&self.matches),
            source_repo_name,
            target_repo_name,
        );
        info!(logger, "{}", details);
        Ok(Arc::new(executor))
    }
}

/// Struct representing the execution of the Back Syncer
/// BP over the context of a provided repos.
pub struct BacksyncProcessExecutor {
    fb: FacebookInit,
    matches: Arc<MononokeMatches<'static>>,
    source_repo_name: String,
    target_repo_name: String,
    cancellation_requested: Arc<AtomicBool>,
}

impl BacksyncProcessExecutor {
    fn new(
        fb: FacebookInit,
        matches: Arc<MononokeMatches<'static>>,
        source_repo_name: String,
        target_repo_name: String,
    ) -> Self {
        Self {
            fb,
            matches,
            source_repo_name,
            target_repo_name,
            cancellation_requested: Arc::new(AtomicBool::new(false)),
        }
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for BacksyncProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.matches.logger(),
            "Initiating back syncer command execution for repo pair {}-{}",
            &self.source_repo_name,
            &self.target_repo_name,
        );
        run(
            self.fb,
            Arc::clone(&self.matches),
            self.source_repo_name.clone(),
            self.target_repo_name.clone(),
            Arc::clone(&self.cancellation_requested),
        )
        .await
        .with_context(|| {
            format!(
                "Error during back syncer command execution for repo pair {}-{}",
                &self.source_repo_name, &self.target_repo_name,
            )
        })?;
        info!(
            self.matches.logger(),
            "Finished back syncer command execution for repo pair {}-{}",
            &self.source_repo_name,
            self.target_repo_name
        );
        Ok(())
    }

    async fn stop(&self) -> anyhow::Result<()> {
        info!(
            self.matches.logger(),
            "Terminating back syncer command execution for repo pair {}-{}",
            &self.source_repo_name,
            self.target_repo_name,
        );
        self.cancellation_requested.store(true, Ordering::Relaxed);
        Ok(())
    }
}

fn extract_cs_id_from_sync_outcome(
    source_cs_id: ChangesetId,
    maybe_sync_outcome: Option<CommitSyncOutcome>,
) -> Result<Option<ChangesetId>, Error> {
    use CommitSyncOutcome::*;

    match maybe_sync_outcome {
        Some(RewrittenAs(cs_id, _)) => Ok(Some(cs_id)),
        Some(NotSyncCandidate(_)) => Ok(None),
        Some(EquivalentWorkingCopyAncestor(cs_id, _)) => Ok(Some(cs_id)),
        None => Err(format_err!(
            "sync outcome is not available for {}",
            source_cs_id
        )),
    }
}

async fn derive_target_hg_changesets(
    ctx: &CoreContext,
    maybe_target_cs_id: Option<ChangesetId>,
    commit_syncer: &CommitSyncer<SqlSyncedCommitMapping, Repo>,
) -> Result<(), Error> {
    match maybe_target_cs_id {
        Some(target_cs_id) => {
            let hg_cs_id = commit_syncer
                .get_target_repo()
                .derive_hg_changeset(ctx, target_cs_id)
                .await?;
            info!(
                ctx.logger(),
                "Hg cs id {} derived for {}", hg_cs_id, target_cs_id
            );
            Ok(())
        }
        None => Ok(()),
    }
}

pub async fn backsync_forever<M>(
    ctx: CoreContext,
    commit_syncer: CommitSyncer<M, Repo>,
    target_repo_dbs: Arc<TargetRepoDbs>,
    source_repo_name: String,
    target_repo_name: String,
    live_commit_sync_config: CfgrLiveCommitSyncConfig,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error>
where
    M: SyncedCommitMapping + Clone + 'static,
{
    let target_repo_id = commit_syncer.get_target_repo_id();
    let live_commit_sync_config = Arc::new(live_commit_sync_config);

    loop {
        // Before initiating loop, check if cancellation has been
        // requested. If yes, then exit early.
        if cancellation_requested.load(Ordering::Relaxed) {
            info!(ctx.logger(), "sync stopping due to cancellation request");
            return Ok(());
        }
        // We only care about public pushes because draft pushes are not in the bookmark
        // update log at all.
        let enabled = live_commit_sync_config.push_redirector_enabled_for_public(target_repo_id);

        if enabled {
            let delay = calculate_delay(&ctx, &commit_syncer, &target_repo_dbs).await?;
            log_delay(&ctx, &delay, &source_repo_name, &target_repo_name);
            if delay.remaining_entries == 0 {
                debug!(ctx.logger(), "no entries remained");
                tokio::time::sleep(Duration::new(1, 0)).await;
            } else {
                debug!(ctx.logger(), "backsyncing...");

                backsync_latest(
                    ctx.clone(),
                    commit_syncer.clone(),
                    target_repo_dbs.clone(),
                    BacksyncLimit::NoLimit,
                    Arc::clone(&cancellation_requested),
                    CommitSyncContext::Backsyncer,
                    false,
                )
                .await?
            }
        } else {
            debug!(ctx.logger(), "push redirector is disabled");
            let delay = Delay::no_delay();
            log_delay(&ctx, &delay, &source_repo_name, &target_repo_name);
            tokio::time::sleep(Duration::new(1, 0)).await;
        }
    }
}

struct Delay {
    delay_secs: i64,
    remaining_entries: u64,
}

impl Delay {
    fn no_delay() -> Self {
        Self {
            delay_secs: 0,
            remaining_entries: 0,
        }
    }
}

// Returns logs delay and returns the number of remaining bookmark update log entries
async fn calculate_delay<M>(
    ctx: &CoreContext,
    commit_syncer: &CommitSyncer<M, Repo>,
    target_repo_dbs: &TargetRepoDbs,
) -> Result<Delay, Error>
where
    M: SyncedCommitMapping + Clone + 'static,
{
    let TargetRepoDbs { ref counters, .. } = target_repo_dbs;
    let source_repo_id = commit_syncer.get_source_repo().repo_identity().id();

    let counter_name = format_counter(&source_repo_id);
    let maybe_counter = counters.get_counter(ctx, &counter_name).await?;
    let counter = maybe_counter.ok_or_else(|| format_err!("{} counter not found", counter_name))?;
    let source_repo = commit_syncer.get_source_repo();
    let next_entry = source_repo
        .bookmark_update_log()
        .read_next_bookmark_log_entries(ctx.clone(), counter as u64, 1, Freshness::MostRecent)
        .try_collect::<Vec<_>>();
    let remaining_entries = source_repo
        .bookmark_update_log()
        .count_further_bookmark_log_entries(ctx.clone(), counter as u64, None);

    let (next_entry, remaining_entries) = try_join!(next_entry, remaining_entries)?;
    let delay_secs = next_entry
        .get(0)
        .map_or(0, |entry| entry.timestamp.since_seconds());

    Ok(Delay {
        delay_secs,
        remaining_entries,
    })
}

fn log_delay(ctx: &CoreContext, delay: &Delay, source_repo_name: &str, target_repo_name: &str) {
    STATS::remaining_entries.set_value(
        ctx.fb,
        delay.remaining_entries as i64,
        (source_repo_name.to_owned(), target_repo_name.to_owned()),
    );
    STATS::delay_secs.set_value(
        ctx.fb,
        delay.delay_secs,
        (source_repo_name.to_owned(), target_repo_name.to_owned()),
    );
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let process = BacksyncProcess::new(fb)?;
    match process.matches.value_of("sharded-service-name") {
        Some(service_name) => {
            // The service name needs to be 'static to satisfy SM contract
            static SM_SERVICE_NAME: OnceCell<String> = OnceCell::new();
            let logger = process.matches.logger().clone();
            let matches = Arc::clone(&process.matches);
            let mut executor = ShardedProcessExecutor::new(
                process.fb,
                process.matches.runtime().clone(),
                &logger,
                SM_SERVICE_NAME.get_or_init(|| service_name.to_string()),
                SM_SERVICE_SCOPE,
                SM_CLEANUP_TIMEOUT_SECS,
                Arc::new(process),
                true, // enable shard (repo) level healing
            )?;
            helpers::block_execute(
                executor.block_and_execute(&logger, Arc::new(AtomicBool::new(false))),
                fb,
                &std::env::var("TW_JOB_NAME").unwrap_or_else(|_| APP_NAME.to_string()),
                matches.logger(),
                &matches,
                cmdlib::monitoring::AliveService,
            )
        }
        None => {
            let logger = process.matches.logger().clone();
            let matches = process.matches.clone();
            let config_store = matches.config_store();
            let source_repo_id =
                args::not_shardmanager_compatible::get_source_repo_id(config_store, &matches)?;
            let target_repo_id =
                args::not_shardmanager_compatible::get_target_repo_id(config_store, &matches)?;
            let (source_repo_name, _) =
                args::get_config_by_repoid(config_store, &matches, source_repo_id)?;
            let (target_repo_name, _) =
                args::get_config_by_repoid(config_store, &matches, target_repo_id)?;
            let fut = run(
                fb,
                matches.clone(),
                source_repo_name,
                target_repo_name,
                Arc::new(AtomicBool::new(false)),
            );
            helpers::block_execute(
                fut,
                fb,
                APP_NAME,
                &logger,
                &matches,
                cmdlib::monitoring::AliveService,
            )
        }
    }
}

async fn run(
    fb: FacebookInit,
    matches: Arc<MononokeMatches<'static>>,
    source_repo_name: String,
    target_repo_name: String,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let config_store = matches.config_store();

    let source_repo = args::resolve_repo_by_name(config_store, &matches, &source_repo_name)?;
    let target_repo = args::resolve_repo_by_name(config_store, &matches, &target_repo_name)?;
    let repo_tag = format!("{}=>{}", &source_repo_name, &target_repo_name);
    let session_container = SessionContainer::new_with_defaults(fb);
    let scribe = args::get_scribe(fb, &matches)?;
    let ctx = session_container
        .new_context_with_scribe(
            matches.logger().clone(),
            MononokeScubaSampleBuilder::with_discard(),
            scribe.clone(),
        )
        .clone_with_repo_name(&repo_tag);
    let commit_syncer = create_commit_syncer_from_matches::<Repo>(
        &ctx,
        &matches,
        Some((source_repo.id, target_repo.id)),
    )
    .await?;
    let logger = ctx.logger();

    info!(
        logger,
        "syncing from repoid {:?} into repoid {:?}", source_repo.id, target_repo.id,
    );

    let config_store = matches.config_store();
    let live_commit_sync_config = CfgrLiveCommitSyncConfig::new(logger, config_store)?;

    match matches.subcommand() {
        (ARG_MODE_BACKSYNC_ALL, _) => {
            let scuba_sample = MononokeScubaSampleBuilder::with_discard();
            let ctx =
                session_container.new_context_with_scribe(logger.clone(), scuba_sample, scribe);
            let target_repo_dbs = Arc::new(
                open_backsyncer_dbs(commit_syncer.get_target_repo())
                    .boxed()
                    .await?,
            );

            // TODO(ikostia): why do we use discarding ScubaSample for BACKSYNC_ALL?
            backsync_latest(
                ctx,
                commit_syncer,
                target_repo_dbs,
                BacksyncLimit::NoLimit,
                cancellation_requested,
                CommitSyncContext::Backsyncer,
                false,
            )
            .boxed()
            .await?;
        }
        (ARG_MODE_BACKSYNC_FOREVER, _) => {
            let target_repo_dbs = Arc::new(
                open_backsyncer_dbs(commit_syncer.get_target_repo())
                    .boxed()
                    .await?,
            );

            let mut scuba_sample = MononokeScubaSampleBuilder::new(fb, SCUBA_TABLE)?;
            scuba_sample.add("source_repo", source_repo.id.id());
            scuba_sample.add("source_repo_name", source_repo.name.clone());
            scuba_sample.add("target_repo", target_repo.id.id());
            scuba_sample.add("target_repo_name", target_repo.name.clone());
            scuba_sample.add_common_server_data();

            let ctx =
                session_container.new_context_with_scribe(logger.clone(), scuba_sample, scribe);
            let f = backsync_forever(
                ctx,
                commit_syncer,
                target_repo_dbs,
                source_repo.name,
                target_repo.name,
                live_commit_sync_config,
                cancellation_requested,
            )
            .boxed();
            f.await?;
        }
        (ARG_MODE_BACKSYNC_COMMITS, Some(sub_m)) => {
            let ctx = session_container.new_context_with_scribe(
                logger.clone(),
                MononokeScubaSampleBuilder::with_discard(),
                scribe,
            );
            let inputfile = sub_m
                .value_of(ARG_INPUT_FILE)
                .expect("input file is not set");
            let inputfile = File::open(inputfile)?;
            let file = BufReader::new(&inputfile);
            let batch_size = args::get_usize(matches.as_ref(), ARG_BATCH_SIZE, 100);

            let source_repo = commit_syncer.get_source_repo().clone();

            let mut hg_cs_ids = vec![];
            for line in file.lines() {
                hg_cs_ids.push(HgChangesetId::from_str(&line?)?);
            }
            let total_to_backsync = hg_cs_ids.len();
            info!(ctx.logger(), "backsyncing {} commits", total_to_backsync);

            let ctx = &ctx;
            let commit_syncer = &commit_syncer;

            // Before processing each commit, check if cancellation has
            // been requested and exit if that's the case.
            if cancellation_requested.load(Ordering::Relaxed) {
                info!(ctx.logger(), "sync stopping due to cancellation request");
                return Ok(());
            }
            let f = stream::iter(hg_cs_ids.clone())
                .chunks(batch_size)
                .map(Result::<_, Error>::Ok)
                .and_then({
                    cloned!(ctx);
                    move |chunk| {
                        cloned!(ctx, source_repo);
                        async move { source_repo.get_hg_bonsai_mapping(ctx.clone(), chunk).await }
                    }
                })
                .try_fold(0, move |backsynced_so_far, hg_bonsai_mapping| {
                    hg_bonsai_mapping
                        .into_iter()
                        .map({
                            move |(_, bonsai)| async move {
                                // Backsyncer is always used in the large-to-small direction,
                                // therefore there can be at most one remapped candidate,
                                // so `CandidateSelectionHint::Only` is a safe choice
                                commit_syncer
                                    .sync_commit(
                                        ctx,
                                        bonsai.clone(),
                                        CandidateSelectionHint::Only,
                                        CommitSyncContext::Backsyncer,
                                        false,
                                    )
                                    .await?;

                                let maybe_sync_outcome =
                                    commit_syncer.get_commit_sync_outcome(ctx, bonsai).await?;

                                info!(
                                    ctx.logger(),
                                    "{} backsynced as {:?}", bonsai, maybe_sync_outcome
                                );

                                let maybe_target_cs_id =
                                    extract_cs_id_from_sync_outcome(bonsai, maybe_sync_outcome)?;

                                derive_target_hg_changesets(ctx, maybe_target_cs_id, commit_syncer)
                                    .await
                            }
                        })
                        .collect::<stream::futures_unordered::FuturesUnordered<_>>()
                        .try_fold(backsynced_so_far, {
                            move |backsynced_so_far, _| async move {
                                info!(
                                    ctx.logger(),
                                    "backsynced so far {} out of {}",
                                    backsynced_so_far + 1,
                                    total_to_backsync
                                );
                                Ok::<_, Error>(backsynced_so_far + 1)
                            }
                        })
                });

            f.await?;
        }
        _ => {
            bail!("unknown subcommand");
        }
    }

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

//! Backsyncer
//!
//! Library to sync commits from source repo to target repo by following bookmark update log
//! and doing commit rewrites. The main motivation for backsyncer is to keep "small repo" up to
//! date with "large repo" in a setup where all writes to small repo are redirected to large repo
//! in a push redirector.
//! More details can be found here - <https://fb.quip.com/tZ4yAaA3S4Mc>
//!
//! Target repo tails source repo's bookmark update log and backsync bookmark updates one by one.
//! The latest backsynced log id is stored in mutable_counters table. Backsync consists of the
//! following phases:
//!
//! 1) Given an entry from bookmark update log of a target repo,
//!    find commits to backsync from source repo into a target repo.
//! 2) Rewrite these commits and create rewritten commits in target repo
//! 3) In the same transaction try to update a bookmark in the source repo AND latest backsynced
//!    log id.

use std::collections::HashSet;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Instant;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Error;
use blobstore::Loadable;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingEntry;
use bonsai_hg_mapping::BonsaiHgMapping;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkTransactionError;
use bookmarks::BookmarkUpdateLog;
use bookmarks::BookmarkUpdateLogArc;
use bookmarks::BookmarkUpdateLogEntry;
use bookmarks::BookmarkUpdateLogRef;
use bookmarks::BookmarkUpdateReason;
use bookmarks::Bookmarks;
use bookmarks::BookmarksArc;
use bookmarks::Freshness;
use changeset_fetcher::ChangesetFetcher;
use changeset_fetcher::ChangesetFetcherArc;
use changesets::Changesets;
use cloned::cloned;
use commit_graph::CommitGraph;
use commit_graph::CommitGraphRef;
use context::CoreContext;
use cross_repo_sync::find_toposorted_unsynced_ancestors;
use cross_repo_sync::CandidateSelectionHint;
use cross_repo_sync::CommitSyncContext;
use cross_repo_sync::CommitSyncOutcome;
use cross_repo_sync::CommitSyncer;
use filestore::FilestoreConfig;
use futures::compat::Stream01CompatExt;
use futures::stream;
use futures::FutureExt;
use futures::StreamExt;
use futures::TryStreamExt;
use metaconfig_types::RepoConfig;
use metaconfig_types::RepoConfigRef;
use mononoke_types::ChangesetId;
use mononoke_types::Globalrev;
use mononoke_types::RepositoryId;
use mutable_counters::MutableCounters;
use mutable_counters::MutableCountersArc;
use mutable_counters::SqlMutableCounters;
use phases::Phases;
use repo_blobstore::RepoBlobstore;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedData;
use repo_identity::RepoIdentity;
use repo_identity::RepoIdentityRef;
use repo_update_logger::find_draft_ancestors;
use repo_update_logger::log_new_bonsai_changesets;
use revset::AncestorsNodeStream;
use slog::debug;
use slog::info;
use slog::warn;
use sql::Transaction;
use sql_ext::TransactionResult;
use synced_commit_mapping::SyncedCommitMapping;
use thiserror::Error;
use wireproto_handler::TargetRepoDbs;

#[derive(Clone)]
#[facet::container]
pub struct Repo(
    dyn BonsaiHgMapping,
    dyn Bookmarks,
    dyn BookmarkUpdateLog,
    dyn Changesets,
    dyn ChangesetFetcher,
    FilestoreConfig,
    dyn MutableCounters,
    dyn Phases,
    RepoBlobstore,
    RepoConfig,
    RepoDerivedData,
    RepoIdentity,
    CommitGraph,
);

#[cfg(test)]
mod tests;

#[derive(Debug, Error)]
pub enum BacksyncError {
    #[error("BacksyncError::LogEntryNotFound: {latest_log_id} not found")]
    LogEntryNotFound { latest_log_id: u64 },
    #[error("BacksyncError::Other")]
    Other(#[from] Error),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BacksyncLimit {
    NoLimit,
    Limit(u64),
}

pub async fn backsync_latest<M, R>(
    ctx: CoreContext,
    commit_syncer: CommitSyncer<M, R>,
    target_repo_dbs: Arc<TargetRepoDbs>,
    limit: BacksyncLimit,
    cancellation_requested: Arc<AtomicBool>,
    sync_context: CommitSyncContext,
    disable_lease: bool,
) -> Result<(), Error>
where
    M: SyncedCommitMapping + Clone + 'static,
    R: RepoLike + Send + Sync + Clone + 'static,
{
    // TODO(ikostia): start borrowing `CommitSyncer`, no reason to consume it
    let source_repo_id = commit_syncer.get_source_repo().repo_identity().id();
    let counter_name = format_counter(&source_repo_id);

    let counter = target_repo_dbs
        .counters
        .get_counter(&ctx, &counter_name)
        .await?
        .unwrap_or(0);

    debug!(ctx.logger(), "fetched counter {}", counter);

    let log_entries_limit = match limit {
        BacksyncLimit::Limit(limit) => limit,
        BacksyncLimit::NoLimit => {
            // Set limit extremely high to read all new values
            u64::max_value()
        }
    };
    let next_entries: Vec<_> = commit_syncer
        .get_source_repo()
        .bookmark_update_log()
        .read_next_bookmark_log_entries(
            ctx.clone(),
            counter as u64,
            log_entries_limit,
            Freshness::MostRecent,
        )
        .try_collect()
        .await?;

    // Before syncing entries, check if cancellation has been
    // requested. If yes, then exit early.
    if cancellation_requested.load(Ordering::Relaxed) {
        info!(ctx.logger(), "sync stopping due to cancellation request");
        return Ok(());
    }

    if next_entries.is_empty() {
        debug!(ctx.logger(), "nothing to sync");
        Ok(())
    } else {
        sync_entries(
            ctx,
            &commit_syncer,
            target_repo_dbs,
            next_entries,
            counter,
            cancellation_requested,
            sync_context,
            disable_lease,
        )
        .await
    }
}

async fn sync_entries<M, R>(
    ctx: CoreContext,
    commit_syncer: &CommitSyncer<M, R>,
    target_repo_dbs: Arc<TargetRepoDbs>,
    entries: Vec<BookmarkUpdateLogEntry>,
    mut counter: i64,
    cancellation_requested: Arc<AtomicBool>,
    sync_context: CommitSyncContext,
    disable_lease: bool,
) -> Result<(), Error>
where
    M: SyncedCommitMapping + Clone + 'static,
    R: RepoLike + Send + Sync + Clone + 'static,
{
    for entry in entries {
        // Before processing each entry, check if cancellation has
        // been requested and exit if that's the case.
        if cancellation_requested.load(Ordering::Relaxed) {
            info!(ctx.logger(), "sync stopping due to cancellation request");
            return Ok(());
        }
        let entry_id = entry.id;
        if counter >= entry_id {
            continue;
        }
        debug!(ctx.logger(), "backsyncing {} ...", entry_id);

        let mut scuba_sample = ctx.scuba().clone();
        scuba_sample.add("backsyncer_bookmark_log_entry_id", entry.id);

        let start_instant = Instant::now();

        if let Some(to_cs_id) = entry.to_changeset_id {
            let (_, unsynced_ancestors_versions) =
                find_toposorted_unsynced_ancestors(&ctx, commit_syncer, to_cs_id).await?;

            if !unsynced_ancestors_versions.has_ancestor_with_a_known_outcome() {
                // Not a single ancestor of to_cs_id was ever synced.
                // That means that we can't figure out which commit sync mapping version
                // to use. In that case we just skip this entry and not sync it at all.
                // This seems the safest option (i.e. we won't rewrite a commit with
                // an incorrect version) but it also has a downside that the bookmark that points
                // to this commit is not going to be synced.
                warn!(
                    ctx.logger(),
                    "skipping {}, entry id {}", entry.bookmark_name, entry.id
                );
                scuba_sample.log_with_msg(
                    "Skipping entry because there are no synced ancestors",
                    Some(format!("{}", entry.id)),
                );
                target_repo_dbs
                    .counters
                    .set_counter(
                        &ctx,
                        &format_counter(&commit_syncer.get_source_repo().repo_identity().id()),
                        entry.id,
                        Some(counter),
                    )
                    .await?;
                counter = entry.id;
                continue;
            }

            // Backsyncer is always used in the large-to-small direction,
            // therefore there can be at most one remapped candidate,
            // so `CandidateSelectionHint::Only` is a safe choice
            commit_syncer
                .sync_commit(
                    &ctx,
                    to_cs_id,
                    CandidateSelectionHint::Only,
                    sync_context,
                    disable_lease,
                )
                .await?;
        }

        let new_counter = entry.id;
        let success = backsync_bookmark(
            ctx.clone(),
            commit_syncer,
            target_repo_dbs.clone(),
            Some(counter),
            entry,
        )
        .await?;

        scuba_sample.add(
            "backsync_duration_ms",
            u64::try_from(start_instant.elapsed().as_millis()).unwrap_or(u64::max_value()),
        );
        scuba_sample.add("backsync_previously_done", !success);
        scuba_sample.log_with_msg("Backsyncing", None);

        if success {
            counter = new_counter;
        } else {
            debug!(
                ctx.logger(),
                "failed to backsync {}, most likely another process already synced it ", entry_id
            );
            // Transaction failed, it could be because another process already backsynced it
            // Verify that counter was moved and continue if that's the case

            let source_repo_id = commit_syncer.get_source_repo().repo_identity().id();
            let counter_name = format_counter(&source_repo_id);
            let new_counter = target_repo_dbs
                .counters
                .get_counter(&ctx, &counter_name)
                .await?
                .unwrap_or(0);
            if new_counter <= counter {
                return Err(format_err!(
                    "backsync transaction failed, but the counter didn't move forward. Was {}, became {}",
                    counter,
                    new_counter,
                ));
            } else {
                debug!(
                    ctx.logger(),
                    "verified that another process has already synced {}", entry_id
                );
                counter = new_counter;
            }
        }
    }
    Ok(())
}

/// All "new" commits on this bookmark move. Use with care, creating a bookmark
/// means ALL ancestors are new.
async fn commits_added_by_bookmark_move(
    ctx: &CoreContext,
    repo: &impl RepoLike,
    from_cs_id: Option<ChangesetId>,
    to_cs_id: Option<ChangesetId>,
) -> Result<HashSet<ChangesetId>, Error> {
    match (from_cs_id, to_cs_id) {
        (_, None) => Ok(HashSet::new()),
        (None, Some(to_id)) => {
            AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), to_id)
                .compat()
                .try_collect()
                .await
        }
        (Some(from_id), Some(to_id)) => {
            // If needed, this can be optimised by using range_stream when from_id is
            // an ancestor of from_id
            repo.commit_graph()
                .ancestors_difference_stream(ctx, vec![to_id], vec![from_id])
                .await?
                .try_collect()
                .await
        }
    }
}

async fn backsync_bookmark<M, R>(
    ctx: CoreContext,
    commit_syncer: &CommitSyncer<M, R>,
    target_repo_dbs: Arc<TargetRepoDbs>,
    prev_counter: Option<i64>,
    log_entry: BookmarkUpdateLogEntry,
) -> Result<bool, Error>
where
    M: SyncedCommitMapping + Clone + 'static,
    R: RepoLike + Send + Sync + Clone + 'static,
{
    let target_repo_id = commit_syncer.get_target_repo().repo_identity().id();
    let source_repo_id = commit_syncer.get_source_repo().repo_identity().id();

    debug!(ctx.logger(), "preparing to backsync {:?}", log_entry);

    let new_counter = log_entry.id;
    let bookmark = commit_syncer.get_bookmark_renamer().await?(&log_entry.bookmark_name);
    debug!(ctx.logger(), "bookmark was renamed into {:?}", bookmark);
    let from_cs_id = log_entry.from_changeset_id;
    let to_cs_id = log_entry.to_changeset_id;

    let get_commit_sync_outcome = |maybe_cs_id: Option<ChangesetId>| {
        cloned!(ctx);
        async move {
            match maybe_cs_id {
                Some(cs_id) => {
                    let maybe_outcome = commit_syncer.get_commit_sync_outcome(&ctx, cs_id).await?;
                    match maybe_outcome {
                        Some(outcome) => Ok(Some((outcome, cs_id))),
                        None => Err(format_err!("{} hasn't been backsynced yet", cs_id)),
                    }
                }
                None => Ok(None),
            }
        }
    };

    let get_remapped_cs_id =
        move |maybe_outcome: Option<(CommitSyncOutcome, ChangesetId)>| match maybe_outcome {
            Some((outcome, cs_id)) => {
                use CommitSyncOutcome::*;
                match outcome {
                    NotSyncCandidate(_) => Err(format_err!(
                        "invalid bookmark move: {:?} should not be synced to target repo",
                        cs_id
                    )),
                    RewrittenAs(cs_id, _) | EquivalentWorkingCopyAncestor(cs_id, _) => {
                        Ok(Some(cs_id))
                    }
                }
            }
            None => Ok(None),
        };

    if let Some(bookmark) = bookmark {
        // Fetch sync outcome before transaction to keep transaction as short as possible
        let from_sync_outcome = get_commit_sync_outcome(from_cs_id).await?;
        let to_sync_outcome = get_commit_sync_outcome(to_cs_id).await?;
        debug!(
            ctx.logger(),
            "commit sync outcomes: from_cs: {:?}, to_cs: {:?}", from_sync_outcome, to_sync_outcome
        );

        let from_cs_id = get_remapped_cs_id(from_sync_outcome)?;
        let to_cs_id = get_remapped_cs_id(to_sync_outcome)?;

        if from_cs_id != to_cs_id {
            let target_repo = commit_syncer.get_target_repo();
            // This CANNOT be done after getting the bookmark transaction, because it accesses SQL without a
            // transaction and that causes a deadlock that blocks the syncing.
            let globalrev_entries: Vec<BonsaiGlobalrevMappingEntry> = if target_repo
                .repo_config()
                .pushrebase
                .globalrev_config
                .as_ref()
                .map(|c| &c.publishing_bookmark)
                == Some(&bookmark)
            {
                let all_commits =
                    commits_added_by_bookmark_move(&ctx, target_repo, from_cs_id, to_cs_id).await?;
                let ctx = &ctx;
                let blobstore = target_repo.repo_blobstore();
                stream::iter(all_commits)
                    .map(|bcs_id| async move {
                        let cs = bcs_id.load(ctx, blobstore).await?;
                        // When pushrebasing into the large repo, this commit
                        // should've gotten a globalrev
                        let globalrev = Globalrev::from_bcs(&cs)?;
                        anyhow::Ok(BonsaiGlobalrevMappingEntry { bcs_id, globalrev })
                    })
                    .buffer_unordered(100)
                    .try_collect()
                    .await?
            } else {
                vec![]
            };
            let commits_to_log = async {
                match to_cs_id {
                    Some(to_cs_id) => {
                        let res = find_draft_ancestors(&ctx, target_repo, to_cs_id).await;
                        match res {
                            Ok(bcss) => bcss,
                            Err(err) => {
                                ctx.scuba().clone().log_with_msg(
                                    "Failed to find draft ancestors for logging",
                                    Some(format!("{}", err)),
                                );
                                vec![]
                            }
                        }
                    }
                    None => vec![],
                }
            }
            .await;

            let mut bookmark_txn = target_repo_dbs.bookmarks.create_transaction(ctx.clone());
            debug!(
                ctx.logger(),
                "syncing bookmark {} to {:?}", bookmark, to_cs_id
            );

            match (from_cs_id, to_cs_id) {
                (Some(from), Some(to)) => {
                    debug!(
                        ctx.logger(),
                        "updating bookmark {:?} from {:?} to {:?}", bookmark, from, to
                    );
                    bookmark_txn.update(&bookmark, to, from, BookmarkUpdateReason::Backsyncer)?;
                }
                (Some(from), None) => {
                    debug!(
                        ctx.logger(),
                        "deleting bookmark {:?} with original position {:?}", bookmark, from
                    );
                    bookmark_txn.delete(&bookmark, from, BookmarkUpdateReason::Backsyncer)?;
                }
                (None, Some(to)) => {
                    debug!(
                        ctx.logger(),
                        "creating bookmark {:?} to point to {:?}", bookmark, to
                    );
                    bookmark_txn.create(&bookmark, to, BookmarkUpdateReason::Backsyncer)?;
                }
                (None, None) => {
                    bail!("unexpected bookmark move");
                }
            };

            let txn_hook = Arc::new({
                move |ctx: CoreContext, txn: Transaction| {
                    cloned!(globalrev_entries);
                    async move {
                        // This is an abstraction leak: it only works because the
                        // mutable counters/globalrevs are stored in the same db as the
                        // bookmarks.
                        let txn_result = SqlMutableCounters::set_counter_on_txn(
                            &ctx,
                            target_repo_id,
                            &format_counter(&source_repo_id),
                            new_counter,
                            prev_counter,
                            txn,
                        )
                        .await?;

                        let txn = match txn_result {
                            TransactionResult::Succeeded(txn) => Ok(txn),
                            TransactionResult::Failed => Err(BookmarkTransactionError::LogicError),
                        }?;

                        if !globalrev_entries.is_empty() {
                            bonsai_globalrev_mapping::add_globalrevs(
                                txn,
                                target_repo_id,
                                &globalrev_entries,
                            )
                            .await
                            .map_err(|err| BookmarkTransactionError::Other(err.into()))
                        } else {
                            Ok(txn)
                        }
                    }
                    .boxed()
                }
            });

            let res = bookmark_txn.commit_with_hook(txn_hook).await?;
            log_new_bonsai_changesets(
                &ctx,
                target_repo,
                &bookmark,
                BookmarkKind::Publishing,
                commits_to_log,
            )
            .await;

            return Ok(res);
        } else {
            debug!(
                ctx.logger(),
                "from_cs_id and to_cs_id are the same: {:?}. No sync happening for {:?}",
                from_cs_id,
                bookmark
            );
        }
    } else {
        debug!(ctx.logger(), "Renamed bookmark is None. No sync happening.");
    }

    let updated = target_repo_dbs
        .counters
        .set_counter(
            &ctx,
            &format_counter(&source_repo_id),
            new_counter,
            prev_counter,
        )
        .await?;

    Ok(updated)
}

pub async fn open_backsyncer_dbs(repo: &impl RepoLike) -> Result<TargetRepoDbs, Error> {
    Ok(TargetRepoDbs {
        bookmarks: repo.bookmarks_arc(),
        bookmark_update_log: repo.bookmark_update_log_arc(),
        counters: repo.mutable_counters_arc(),
    })
}

pub fn format_counter(repo_to_backsync_from: &RepositoryId) -> String {
    format!("backsync_from_{}", repo_to_backsync_from.id())
}