/*
 * 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 clap::Args;

/// Command line arguments for controlling MySql
// Defaults are derived from `sql_ext::facebook::mysql`
// https://fburl.com/diffusion/n5isd68j, last synced on 17/12/2020
#[derive(Args, Debug)]
pub struct MysqlArgs {
    /// Connect to MySql master only.
    #[clap(long)]
    pub mysql_master_only: bool,

    /// Size of the MySql connection pool
    #[clap(long, default_value = "10000")]
    pub mysql_pool_limit: usize,

    /// MySql connection pool per key limit
    #[clap(long, default_value = "100")]
    pub mysql_pool_per_key_limit: u64,

    /// Number of threads in MySql connection pool (number of real pools)
    #[clap(long, default_value = "10")]
    pub mysql_pool_threads_num: i32,

    /// Mysql connection pool age timeout in millisecs
    #[clap(long, default_value = "60000")]
    pub mysql_pool_age_timeout: u64,

    /// Mysql connection pool idle timeout in millisecs
    #[clap(long, default_value = "4000")]
    pub mysql_pool_idle_timeout: u64,

    /// Size of the MySql connection pool for SqlBlob
    #[clap(long, default_value = "10000", alias = "mysql-sqblob-pool-limit")]
    pub mysql_sqlblob_pool_limit: usize,

    /// MySql connection pool per key limit for SqlBlob
    #[clap(long, default_value = "100", alias = "mysql-sqblob-pool-per-key-limit")]
    pub mysql_sqlblob_pool_per_key_limit: u64,

    /// Number of threads in MySql connection pool (number of real pools) for
    /// SqlBlob
    #[clap(long, default_value = "10", alias = "mysql-sqblob-pool-threads-num")]
    pub mysql_sqlblob_pool_threads_num: i32,

    /// MySql connection pool age timeout in millisecs for SqlBlob
    #[clap(long, default_value = "60000", alias = "mysql-sqblob-pool-age-timeout")]
    pub mysql_sqlblob_pool_age_timeout: u64,

    /// MySql connection pool idle timeout in millisecs for SqlBlob
    #[clap(long, default_value = "4000", alias = "mysql-sqblob-pool-idle-timeout")]
    pub mysql_sqlblob_pool_idle_timeout: u64,

    /// MySql connection open timeout in millisecs
    #[clap(long, default_value = "3000")]
    pub mysql_conn_open_timeout: u64,

    /// Mysql query time limit in millisecs
    #[clap(long, default_value = "10000", alias = "mysql-max-query-time")]
    pub mysql_query_time_limit: u64,
}


/*
 * 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 clap::Args;

/// Commandline args to make the server read-only
#[derive(Args, Debug)]
pub struct ReadonlyArgs {
    /// Makes the server completely readonly by failing all write ACL checks.
    #[clap(long)]
    pub readonly: bool,
}


/*
 * 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::value_parser;
use clap::Arg;
use clap::ArgGroup;
use clap::ArgMatches;
use clap::Args;
use clap::Command;
use clap::Error;
use clap::FromArgMatches;
use mononoke_types::RepositoryId;

/// Command line arguments for specifying a single repo.

// For convenience first we define macro for generating appropiate RepoArgs
// structure that can be used with clap derive structs. Each struct adds args
// for providing repo using id or name.
fn augment_args(
    cmd: Command,
    ident: &'static str,
    required: bool,
    name_arg: &'static str,
    name_arg_short: Option<char>,
    name_help: &'static str,
    id_arg: &'static str,
    id_help: &'static str,
) -> Command {
    cmd.arg(
        Arg::new(id_arg)
            .long(id_arg)
            .value_parser(value_parser!(i32))
            .value_name("REPO_ID")
            .help(id_help),
    )
    .arg({
        let mut arg = Arg::new(name_arg)
            .long(name_arg)
            .value_name("REPO_NAME")
            .help(name_help);
        if let Some(short) = name_arg_short {
            arg = arg.short(short);
        }
        arg
    })
    .group(
        ArgGroup::new(ident)
            .required(required)
            .args([id_arg, name_arg]),
    )
}

#[macro_export]
macro_rules! repo_args {
    ($ident:ident, $name_arg:literal, $maybe_short_name_arg:expr, $name_help:literal, $id_arg:literal, $id_help:literal) => {
        #[derive(Debug)]
        pub struct $ident(RepoArg);

        impl Args for $ident {
            fn augment_args(cmd: Command) -> Command {
                augment_args(
                    cmd,
                    stringify!($ident),
                    true,
                    $name_arg,
                    $maybe_short_name_arg,
                    $name_help,
                    $id_arg,
                    $id_help,
                )
            }
            fn augment_args_for_update(cmd: Command) -> Command {
                Self::augment_args(cmd)
            }
        }

        impl FromArgMatches for $ident {
            fn from_arg_matches(matches: &ArgMatches) -> Result<Self, Error> {
                let repo_id = matches.get_one($id_arg);
                let repo_name: Option<&String> = matches.get_one($name_arg);
                match (repo_id, repo_name) {
                    (Some(repo_id), None) => Ok(Self(RepoArg::Id(RepositoryId::new(*repo_id)))),
                    (None, Some(repo_name)) => Ok(Self(RepoArg::Name(repo_name.clone()))),
                    // This case should never happen - arg grouping in clap will error first.
                    _ => {
                        unreachable!("exactly one of repo-id and repo-name must be specified");
                    }
                }
            }
            fn update_from_arg_matches(&mut self, matches: &ArgMatches) -> Result<(), Error> {
                *self = Self::from_arg_matches(matches)?;
                Ok(())
            }
        }

        impl AsRepoArg for $ident {
            fn as_repo_arg(&self) -> &RepoArg {
                &self.0
            }
        }
    };
}

#[macro_export]
macro_rules! repo_args_optional {
    ($ident:ident, $name_arg:literal, $maybe_short_name_arg:expr, $name_help:literal, $id_arg:literal, $id_help:literal) => {
        #[derive(Debug)]
        pub struct $ident(Option<RepoArg>);

        impl Args for $ident {
            fn augment_args(cmd: Command) -> Command {
                augment_args(
                    cmd,
                    stringify!($ident),
                    false,
                    $name_arg,
                    $maybe_short_name_arg,
                    $name_help,
                    $id_arg,
                    $id_help,
                )
            }

            fn augment_args_for_update(cmd: Command) -> Command {
                Self::augment_args(cmd)
            }
        }

        impl FromArgMatches for $ident {
            fn from_arg_matches(matches: &ArgMatches) -> Result<Self, Error> {
                let repo_id = matches.get_one($id_arg);
                let repo_name: Option<&String> = matches.get_one($name_arg);
                match (repo_id, repo_name) {
                    (Some(repo_id), None) => {
                        Ok(Self(Some(RepoArg::Id(RepositoryId::new(*repo_id)))))
                    }
                    (None, Some(repo_name)) => Ok(Self(Some(RepoArg::Name(repo_name.clone())))),
                    (Some(_), Some(_)) => unreachable!(),
                    (None, None) => Ok(Self(None)),
                }
            }
            fn update_from_arg_matches(&mut self, matches: &ArgMatches) -> Result<(), Error> {
                *self = Self::from_arg_matches(matches)?;
                Ok(())
            }
        }

        impl $ident {
            pub fn as_repo_arg(&self) -> &Option<RepoArg> {
                &self.0
            }
        }
    };
}

repo_args!(
    RepoArgs,
    "repo-name",
    Some('R'),
    "Repository name",
    "repo-id",
    "Numeric repository ID"
);

impl RepoArgs {
    pub fn from_repo_id(repo_id: i32) -> Self {
        Self(RepoArg::Id(RepositoryId::new(repo_id)))
    }
    pub fn from_repo_name(repo_name: String) -> Self {
        Self(RepoArg::Name(repo_name))
    }
}

repo_args_optional!(
    OptRepoArgs,
    "repo-name",
    Some('R'),
    "Repository name",
    "repo-id",
    "Numeric repository ID"
);

/// Command line arguments for specifying multiple repos.
#[derive(Args, Debug)]
#[clap(group(
    ArgGroup::new("multirepos")
        .multiple(true)
        .args(&["repo_id", "repo_name"]),
))]
pub struct MultiRepoArgs {
    /// Numeric repository ID
    #[clap(long)]
    pub repo_id: Vec<i32>,

    /// Repository name
    #[clap(short = 'R', long)]
    pub repo_name: Vec<String>,
}

impl MultiRepoArgs {
    pub fn ids_or_names(&self) -> Result<Vec<RepoArg>> {
        let mut l = vec![];
        for id in &self.repo_id {
            l.push(RepoArg::Id(RepositoryId::new(*id)));
        }
        for name in &self.repo_name {
            l.push(RepoArg::Name(name.to_owned()));
        }

        Ok(l)
    }
}

repo_args!(
    SourceRepoArgs,
    "source-repo-name",
    None,
    "Source repository name",
    "source-repo-id",
    "Numeric source repository ID"
);

repo_args!(
    TargetRepoArgs,
    "target-repo-name",
    None,
    "Target repository name",
    "target-repo-id",
    "Numeric target repository ID"
);

/// Command line arguments for specifying only a source  and a target repos,
/// Necessary for cross-repo operations
/// Only visible if the app was built with a call to `MononokeAppBuilder::with_source_and_target_repos`
#[derive(Args, Debug)]
pub struct SourceAndTargetRepoArgs {
    #[clap(flatten)]
    pub source_repo: SourceRepoArgs,

    #[clap(flatten)]
    pub target_repo: TargetRepoArgs,
}

#[derive(Debug)]
pub enum RepoArg {
    Id(RepositoryId),
    Name(String),
}

pub trait AsRepoArg {
    fn as_repo_arg(&self) -> &RepoArg;
}

impl AsRepoArg for RepoArg {
    fn as_repo_arg(&self) -> &RepoArg {
        self
    }
}


/*
 * 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 clap::ArgGroup;
use clap::Args;
use clap::ValueEnum;

/// Command line arguments for specifying a blobstore, either by
/// repo, or by storage name.
#[derive(Args, Debug)]
#[clap(group(
    ArgGroup::new("repo_blobstore")
        .required(true)
        .args(&["repo_id", "repo_name", "storage_name"]),
))]
pub struct RepoBlobstoreArgs {
    /// Numeric repository ID
    #[clap(long)]
    pub repo_id: Option<i32>,

    /// Repository name
    #[clap(short = 'R', long)]
    pub repo_name: Option<String>,

    /// Storage name
    #[clap(long)]
    pub storage_name: Option<String>,

    /// If the blobstore is multiplexed, use this inner blobstore
    #[clap(long)]
    pub inner_blobstore_id: Option<u64>,

    /// Use memcache to cache access to the blobstore
    #[clap(long, value_enum)]
    pub use_memcache: Option<UseMemcache>,

    /// Don't prepend the repo prefix to the key
    #[clap(long)]
    pub no_prefix: bool,
}

#[derive(Copy, Clone, Debug, ValueEnum, Eq, PartialEq)]
pub enum UseMemcache {
    CacheOnly,
    NoFill,
    FillMc,
}


/*
 * 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::Result;
use clap::Parser;
use environment::MononokeEnvironment;
use regex::Regex;

use crate::AppExtension;

/// Command line argument to filter repositories
#[derive(Parser, Debug)]
pub struct RepoFilterArgs {
    /// Filter repositories using regex
    #[clap(long, value_parser = Regex::new)]
    pub filter_repos: Option<Regex>,
}

pub struct RepoFilterAppExtension;

impl AppExtension for RepoFilterAppExtension {
    type Args = RepoFilterArgs;

    fn environment_hook(&self, args: &Self::Args, env: &mut MononokeEnvironment) -> Result<()> {
        if let Some(filter_repos) = args.filter_repos.clone() {
            if let Some(current_filter_fn) = env.filter_repos.clone() {
                env.filter_repos = Some(Arc::new(move |name: &str| {
                    filter_repos.is_match(name) & current_filter_fn(name)
                }));
            } else {
                env.filter_repos = Some(Arc::new(move |name: &str| filter_repos.is_match(name)));
            }
        }

        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 clap::Args;

/// Command line arguments for controlling the runtime
// Defaults are derived from `sql_ext::facebook::mysql`
// https://fburl.com/diffusion/n5isd68j, last synced on 17/12/2020
#[derive(Args, Debug)]
pub struct RuntimeArgs {
    /// Number of threads to use in the Tokio runtime
    #[clap(long, help_heading = "RUNTIME OPTIONS")]
    pub runtime_threads: Option<usize>,
}


/*
 * 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 clap::Args;

/// Command line arguments for controlling tunables
#[derive(Args, Debug)]
pub struct TunablesArgs {
    /// Tunables dynamic config path in Configerator
    #[clap(long)]
    pub tunables_config: Option<String>,

    /// Tunables static config local path
    #[clap(long, conflicts_with = "tunables_config")]
    pub tunables_local_path: Option<String>,

    /// Use the default values for all tunables (useful for tests)
    #[clap(long, conflicts_with_all = &["tunables_config", "tunables_local_path"])]
    pub disable_tunables: bool,
}

const DEFAULT_TUNABLES_CONFIG: &str = "scm/mononoke/tunables/default";

impl TunablesArgs {
    pub fn tunables_config_or_default(&self) -> String {
        self.tunables_config
            .clone()
            .unwrap_or_else(|| DEFAULT_TUNABLES_CONFIG.to_string())
    }
}


/*
 * 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 environment::MononokeEnvironment;
use environment::WarmBookmarksCacheDerivedData;

use crate::AppExtension;

/// Command line argument for Warm Bookmarks Cache
#[derive(Parser, Debug)]
pub struct WarmBookmarksCacheArgs {
    /// What needs to be warmed
    #[clap(long, value_enum, value_name = "warm derived data")]
    pub enable_wbc_with: Option<WarmBookmarksCacheDerivedData>,
}

pub struct WarmBookmarksCacheExtension;

impl AppExtension for WarmBookmarksCacheExtension {
    type Args = WarmBookmarksCacheArgs;

    fn environment_hook(&self, args: &Self::Args, env: &mut MononokeEnvironment) -> Result<()> {
        if let Some(wbc_config) = args.enable_wbc_with {
            env.warm_bookmarks_cache_derived_data = Some(wbc_config)
        }
        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 clap::Args;

/// Command line arguments for TLS parameters
#[derive(Args, Debug)]
pub struct TLSArgs {
    /// TLS Certificate for auth
    #[clap(long)]
    pub tls_certificate: Option<String>,
    /// TLS private key
    #[clap(long)]
    pub tls_private_key: Option<String>,
    /// TLS CA
    #[clap(long)]
    pub tls_ca: Option<String>,
    /// TLS Ticket Seeds
    #[clap(long)]
    pub tls_ticket_seeds: Option<String>,
}


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

use std::time::Duration;

use anyhow::Result;
use blobstore_factory::ScrubAction;
use blobstore_factory::ScrubOptions;
use blobstore_factory::SrubWriteOnly;
use clap::Args;
use environment::MononokeEnvironment;
use mononoke_app::AppExtension;

#[derive(Args, Debug)]
pub struct ScrubArgs {
    /// Enable ScrubBlobstore with the given action
    ///
    /// Checks for keys missing from the stores.  In ReportOnly mode, this
    /// only logs, otherwise it performs a copy to the missing stores.
    #[clap(long, help_heading = "BLOBSTORE OPTIONS")]
    pub blobstore_scrub_action: Option<ScrubAction>,

    /// Number of seconds grace to give for key to arrive in multiple
    /// blobstores or the healer queue when scrubbing
    #[clap(
        long,
        help_heading = "BLOBSTORE OPTIONS",
        requires = "blobstore_scrub_action"
    )]
    pub blobstore_scrub_grace: Option<u64>,

    /// Number of seconds within which we consider it worth checking the
    /// healer queue
    #[clap(
        long,
        help_heading = "BLOBSTORE OPTIONS",
        requires = "blobstore_scrub_action"
    )]
    pub blobstore_scrub_queue_peek_bound: Option<u64>,

    /// Whether to allow missing values from write-only stores when
    /// scrubbing
    #[clap(
        long,
        help_heading = "BLOBSTORE OPTIONS",
        requires = "blobstore_scrub_action"
    )]
    pub blobstore_scrub_write_only_missing: Option<SrubWriteOnly>,
}

#[derive(Default, Debug)]
pub struct ScrubAppExtension {
    pub action: Option<ScrubAction>,
    pub grace: Option<Duration>,
    pub queue_peek_bound: Option<Duration>,
    pub write_only_missing: Option<SrubWriteOnly>,
}

impl ScrubAppExtension {
    pub fn new() -> Self {
        ScrubAppExtension::default()
    }
}

impl AppExtension for ScrubAppExtension {
    type Args = ScrubArgs;

    fn arg_defaults(&self) -> Vec<(&'static str, String)> {
        let mut defaults = Vec::new();
        if let Some(action) = self.action {
            defaults.push((
                "blobstore_scrub_action",
                <&'static str>::from(action).to_string(),
            ));
        }
        if let Some(grace) = self.grace {
            defaults.push(("blobstore_scrub_grace", grace.as_secs().to_string()));
        }
        if let Some(queue_peek_bound) = self.queue_peek_bound {
            defaults.push((
                "blobstore_scrub_queue_peek_bound",
                queue_peek_bound.as_secs().to_string(),
            ));
        }
        if let Some(write_only_missing) = self.write_only_missing {
            defaults.push((
                "blobstore_scrub_write_only_missing",
                <&'static str>::from(write_only_missing).to_string(),
            ));
        }
        defaults
    }

    fn environment_hook(&self, args: &ScrubArgs, env: &mut MononokeEnvironment) -> Result<()> {
        if let Some(scrub_action) = args.blobstore_scrub_action {
            let mut scrub_options = ScrubOptions {
                scrub_action,
                ..ScrubOptions::default()
            };

            if let Some(scrub_grace) = args.blobstore_scrub_grace {
                scrub_options.scrub_grace = Some(Duration::from_secs(scrub_grace));
            }
            if let Some(action_on_missing) = args.blobstore_scrub_write_only_missing {
                scrub_options.scrub_action_on_missing_write_only = action_on_missing;
            }
            if let Some(queue_peek_bound) = args.blobstore_scrub_queue_peek_bound {
                scrub_options.queue_peek_bound = Duration::from_secs(queue_peek_bound);
            }
            env.blobstore_options.set_scrub_options(scrub_options);
        }
        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::Result;
use clap::Args;
use fbinit::FacebookInit;
use slog::Logger;
use tokio::runtime::Handle;

use crate::RepoShardedProcess;
use crate::ShardedProcessExecutor;

/// Command line arguments for sharded executor
#[derive(Args, Debug)]
pub struct ShardedExecutorArgs {
    /// The name of the ShardManager service corresponding to this service's region.
    /// If this argument isn't provided, the service will operate in non-sharded mode.
    #[clap(long, requires = "sharded_scope_name")]
    pub sharded_service_name: Option<String>,
    /// The scope of the ShardManager service that this service corresponds to.
    #[clap(long, requires = "sharded_service_name")]
    pub sharded_scope_name: Option<String>,
}

impl ShardedExecutorArgs {
    pub fn build_executor(
        self,
        fb: FacebookInit,
        runtime: Handle,
        logger: &Logger,
        process_fn: impl FnOnce() -> Arc<dyn RepoShardedProcess>,
        shard_healing: bool,
        cleanup_timeout_secs: u64,
    ) -> Result<Option<ShardedProcessExecutor>> {
        if let Some((sharded_service_name, sharded_scope_name)) =
            self.sharded_service_name.zip(self.sharded_scope_name)
        {
            let process = process_fn();
            Ok(Some(ShardedProcessExecutor::new(
                fb,
                runtime,
                logger,
                // The service name & scope needs to be 'static to satisfy SM contract
                Box::leak(Box::new(sharded_service_name)),
                Box::leak(Box::new(sharded_scope_name)),
                cleanup_timeout_secs,
                process,
                shard_healing,
            )?))
        } else {
            Ok(None)
        }
    }
}


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

#![feature(hash_drain_filter)]
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use sharding_ext::RepoShard;

pub mod args;
#[cfg(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod oss;

#[cfg(fbcode_build)]
pub use facebook::*;
#[cfg(not(fbcode_build))]
pub use oss::*;

/// Trait outlining the method responsible for performing the initial bootstrapping
/// of the process with the context of the incoming repo. Implementer of this trait should
/// NOT contain any repo-specific state since one instance of RepoShardedProcess
/// caters all repos for the given job.
#[async_trait]
pub trait RepoShardedProcess: Send + Sync {
    /// Method responsible for performing the initial setup of the process in context of
    /// the provided repo. This method should ONLY contain code necessary to build
    /// state (in form of the struct that implements RepoShardedProcessExecutor trait)
    /// that is required to execute the job. The repo (or any related entity)
    /// should be included as part of the RepoShardedProcessExecutor state.
    async fn setup(&self, repo: &RepoShard) -> Result<Arc<dyn RepoShardedProcessExecutor>>;
}

/// Trait outlining the methods to be implemented by Mononoke jobs that require
/// sharding across different repos based on their resource requirements.
/// Implementer of this trait can store repo-specific state if required since there
/// exists a 1-1 mapping between a repo and a RepoShardedProcessExecutor.
#[async_trait]
pub trait RepoShardedProcessExecutor: Send + Sync {
    /// Callback for when the repo-process for the current job is ready to begin execution
    /// with the provided repo. The core process logic should exist within this method.
    /// The process execution can be one-shot, batched or long-running in nature.
    /// Correspondingly, the task associated with this method could return immediately,
    /// or after a while or just keep executing until the executor decides to terminate it.
    /// Depending on the execution mode, the implementer can choose to maintain some
    /// form of state (i.e. signal, oneshot, etc.) that can be used to interrupt the
    /// normal execution of the process when the Sharded Process Manager needs to move
    /// or terminate the process execution for the provided repo.
    async fn execute(&self) -> Result<()>;

    /// Callback for when the executing repo process for the current job is required to
    /// relinquish an existing repo assigned to it earlier. The timeout specified during
    /// creation determines the time for which the executing repo process will wait before
    /// performing a hard-cleaup of the repo associated task. Note that this
    /// method is responsible ONLY for signalling the termination of the execution
    /// of the repo for the given process. Once the control returns from this method,
    /// the executor waits for timeout seconds (specified during executor construction)
    /// before initiating forced termination of the executing process. If there
    /// are any book-keeping activities pending, they should be completed in the
    /// main execution method during this timeout period. The stop() method itself
    /// SHOULD return quickly (i.e. should not be long running).
    async fn stop(&self) -> Result<()>;
}


/*
 * 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::atomic::AtomicBool;
use std::sync::Arc;

use anyhow::Result;
use fbinit::FacebookInit;
use slog::Logger;
use tokio::runtime::Handle;

use crate::RepoShardedProcess;

pub struct ShardedProcessExecutor {}

impl ShardedProcessExecutor {
    pub fn new(
        _fb: FacebookInit,
        _runtime_handle: Handle,
        _logger: &Logger,
        _service_name: &'static str,
        _service_scope: &'static str,
        _timeout_secs: u64,
        _bp_handle: Arc<dyn RepoShardedProcess>,
        _shard_healing: bool,
    ) -> Result<Self> {
        unimplemented!("ShardedProcessExecutor is supported only for fbcode build")
    }

    pub async fn block_and_execute(
        &mut self,
        _logger: &Logger,
        _terminate_process: Arc<AtomicBool>,
    ) -> Result<()> {
        unimplemented!("ShardedProcessExecutor is supported only for fbcode build")
    }

    pub fn execute(&mut self, _logger: &Logger) {
        unimplemented!("ShardedProcessExecutor is supported only for fbcode build")
    }
}


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

#[cfg(test)]
mod test;

use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;

const ENCODED_SLASH: &str = "_SLASH_";
const ENCODED_PLUS: &str = "_PLUS_";
const X_REPO_SEPARATOR: &str = "_TO_";
const CHUNK_SEPARATOR: &str = "_CHUNK_";
const CHUNK_SIZE_SEPARATOR: &str = "_SIZE_";
const CHUNK_PART_SEPARATOR: &str = "_OF_";

/// Struct representing the parsed structure of a Shard assigned by SM
#[derive(Clone, Default, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct RepoShard {
    // The repo corresponding to this shard. If this is a x-repo shard,
    // then repo_name = source_repo_name
    pub repo_name: String,
    // The target repo name for this x-repo shard. It remains None for
    // single repo shards
    pub target_repo_name: Option<String>,
    // The total number of chunks as represented by the chunked shard.
    // i.e. The 16 in ShardId = fbsource_CHUNK_1_OF_16. Remains None in non-chunked
    // shards
    pub total_chunks: Option<usize>,
    // The id of the chunk corresponding to this specific shard.
    // i.e. The 1 in ShardId = fbsource_CHUNK_1_OF_16. Remains None in non-chunked
    // shards
    pub chunk_id: Option<usize>,
    // The size of each chunk as encoded in the chunked shard.
    // i.e. the 1000 in ShardId = fbsource_CHUNK_1_OF_16_SIZE_1000
    pub chunk_size: Option<usize>,
}

impl RepoShard {
    pub fn with_repo_name(repo_name: &str) -> Self {
        Self {
            repo_name: repo_name.to_string(),
            ..Default::default()
        }
    }

    fn with_source_and_target(repo_name: &str, target_repo_name: &str) -> Self {
        Self {
            repo_name: repo_name.to_string(),
            target_repo_name: Some(target_repo_name.to_string()),
            ..Default::default()
        }
    }

    fn with_chunks(repo_name: &str, chunks: &str, target_repo_name: Option<&str>) -> Result<Self> {
        let mut repo_shard = match target_repo_name {
            Some(target_repo_name) => Self::with_source_and_target(repo_name, target_repo_name),
            None => Self::with_repo_name(repo_name),
        };
        let mut chunk_size_split = split_chunk_size(chunks).into_iter();
        if let Some(chunk_parts) = chunk_size_split.next() {
            let mut parts = split_chunk_parts(chunk_parts).into_iter();
            if let (Some(chunk_id), Some(total_chunks)) = (parts.next(), parts.next()) {
                let chunk_id = chunk_id.parse::<usize>().with_context(|| {
                    format_err!(
                        "Failure in creating RepoShard. Invalid chunk_id {}",
                        chunk_id
                    )
                })?;
                let total_chunks = total_chunks.parse::<usize>().with_context(|| {
                    format_err!(
                        "Failure in creating RepoShard. Invalid total_chunks {}",
                        total_chunks
                    )
                })?;
                repo_shard.chunk_id = Some(chunk_id);
                repo_shard.total_chunks = Some(total_chunks);
            } else {
                anyhow::bail!(
                    "Failure in creating RepoShard. Invalid chunk parts format {:?}",
                    parts
                )
            }
        } else {
            anyhow::bail!(
                "Failure in creating RepoShard. Invalid chunk format {:?}",
                chunks
            )
        }
        if let Some(chunk_size) = chunk_size_split.next() {
            let chunk_size = chunk_size.parse::<usize>().with_context(|| {
                format_err!(
                    "Failure in creating RepoShard. Invalid chunk_size {}",
                    chunk_size
                )
            })?;
            repo_shard.chunk_size = Some(chunk_size);
        }
        Ok(repo_shard)
    }

    /// Create the RepoShard based on the full string representation of the ShardID
    pub fn from_shard_id(shard_id: &str) -> Result<Self> {
        let decoded = decode_repo_name(shard_id);
        let mut split = split_repo_names(&decoded).into_iter();

        let repo_shard = match (split.next(), split.next()) {
            (Some(remaining), None) => match get_repo_with_chunks(remaining)? {
                ShardSplit::Repo(repo_name) => RepoShard::with_repo_name(repo_name),
                ShardSplit::RepoWithChunks(repo_name, chunk_parts) => {
                    RepoShard::with_chunks(repo_name, chunk_parts, None)?
                }
            },
            (Some(source_repo_name), Some(remaining)) => match get_repo_with_chunks(remaining)? {
                ShardSplit::Repo(target_repo_name) => {
                    RepoShard::with_source_and_target(source_repo_name, target_repo_name)
                }
                ShardSplit::RepoWithChunks(target_repo_name, chunk_parts) => {
                    RepoShard::with_chunks(source_repo_name, chunk_parts, Some(target_repo_name))?
                }
            },
            _ => anyhow::bail!(
                "Failure in creating RepoShard. Invalid shard id {}",
                shard_id
            ),
        };
        Ok(repo_shard)
    }
}

impl std::fmt::Display for RepoShard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.repo_name)?;
        if let Some(target_repo_name) = self.target_repo_name.as_ref() {
            f.write_str(X_REPO_SEPARATOR)?;
            f.write_str(target_repo_name)?;
        }
        if let (Some(total_chunks), Some(chunk_id)) = (self.total_chunks, self.chunk_id) {
            f.write_fmt(format_args!(
                "{}{}{}{}",
                CHUNK_SEPARATOR, chunk_id, CHUNK_PART_SEPARATOR, total_chunks
            ))?;
        }
        if let Some(chunk_size) = self.chunk_size {
            f.write_fmt(format_args!("{}{}", CHUNK_SIZE_SEPARATOR, chunk_size))?;
        }
        Ok(())
    }
}

enum ShardSplit<'a> {
    Repo(&'a str),
    RepoWithChunks(&'a str, &'a str),
}

fn get_repo_with_chunks<'a>(input: &'a str) -> Result<ShardSplit<'a>> {
    let mut split = split_chunk(input).into_iter();
    let shard_split = match (split.next(), split.next()) {
        (Some(repo), None) => ShardSplit::Repo(repo),
        (Some(repo), Some(chunk_parts)) => ShardSplit::RepoWithChunks(repo, chunk_parts),
        _ => anyhow::bail!(
            "Failure in creating RepoShard. Invalid shard split {}",
            input
        ),
    };
    Ok(shard_split)
}

fn split_chunk(input: &str) -> Vec<&str> {
    input.splitn(2, CHUNK_SEPARATOR).collect()
}

fn split_chunk_size(chunks: &str) -> Vec<&str> {
    chunks.splitn(2, CHUNK_SIZE_SEPARATOR).collect()
}

fn split_chunk_parts(chunk_parts: &str) -> Vec<&str> {
    chunk_parts.splitn(2, CHUNK_PART_SEPARATOR).collect()
}

/// Function responsible for decoding an SM-encoded repo-name.
pub fn decode_repo_name(encoded_repo_name: &str) -> String {
    encoded_repo_name
        .replace(ENCODED_SLASH, "/")
        .replace(ENCODED_PLUS, "+")
}

/// Function responsible for SM-compatible encoding of repo-na
pub fn encode_repo_name(repo_name: &str) -> String {
    repo_name
        .replace('/', ENCODED_SLASH)
        .replace('+', ENCODED_PLUS)
}

/// Function responsible for splitting source and target repo name
/// from combined repo-name string.
fn split_repo_names(combined_repo_names: &str) -> Vec<&str> {
    combined_repo_names.splitn(2, X_REPO_SEPARATOR).collect()
}


/*
 * 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 crate::decode_repo_name;
use crate::encode_repo_name;
use crate::RepoShard;

#[test]
fn basic_create_repo_shard_test() {
    let repo_shard = RepoShard::with_repo_name("repo");
    assert_eq!(repo_shard.repo_name, "repo".to_string());

    let repo_shard = RepoShard::with_source_and_target("source_repo", "target_repo");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(repo_shard.target_repo_name, Some("target_repo".to_string()));
}

#[test]
fn create_repo_shard_with_sizeless_chunks_test() {
    let repo_shard = RepoShard::with_chunks("source_repo", "4_OF_16", Some("target_repo"))
        .expect("Failed in creating RepoShard");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(repo_shard.target_repo_name, Some("target_repo".to_string()));
    assert_eq!(repo_shard.total_chunks, Some(16));
    assert_eq!(repo_shard.chunk_id, Some(4));
    assert_eq!(repo_shard.chunk_size, None);
}

#[test]
fn create_invalid_repo_shard_with_sizeless_chunks_test() {
    RepoShard::with_chunks("source_repo", "4_OF_-16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "_4_OF_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4__OF__16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4-OF-16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4.0_OF_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_of_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_OF_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_OF_", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4 16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_CHUNK_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
}

#[test]
fn create_repo_shard_with_sized_chunks_test() {
    let repo_shard =
        RepoShard::with_chunks("source_repo", "4_OF_16_SIZE_1000", Some("target_repo"))
            .expect("Failed in creating RepoShard");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(repo_shard.target_repo_name, Some("target_repo".to_string()));
    assert_eq!(repo_shard.total_chunks, Some(16));
    assert_eq!(repo_shard.chunk_id, Some(4));
    assert_eq!(repo_shard.chunk_size, Some(1000));
}

#[test]
fn create_invalid_repo_shard_with_sized_chunks_test() {
    RepoShard::with_chunks("source_repo", "4_OF_16_SIZED_1000", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_1000", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16-1000", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_SIZE_16", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_SIZE_", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_SIZE_200.5", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_SIZE_2-6", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::with_chunks("source_repo", "4_OF_16_SIZE_SIZE_200", Some("target_repo"))
        .expect_err("Should have failed in creating RepoShard");
}

#[test]
fn create_basic_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("repo").expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "repo".to_string());
}

#[test]
fn create_x_repo_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("source_repo_TO_target_repo")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(repo_shard.target_repo_name, Some("target_repo".to_string()));
}

#[test]
fn create_invalid_x_repo_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("source_repo_TO__TO_target_repo")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("_TO_target_repo".to_string())
    );

    let repo_shard = RepoShard::from_shard_id("source_repo_TO_target_repo_TO_another_repo")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "source_repo".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("target_repo_TO_another_repo".to_string())
    );

    let repo_shard = RepoShard::from_shard_id("JustARepoWithTOInItsName")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "JustARepoWithTOInItsName".to_string());
    assert_eq!(repo_shard.target_repo_name, None);

    let repo_shard =
        RepoShard::from_shard_id("JustARepoWithTOInItsName_TO_AnotherRepoWithTOInItsName")
            .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "JustARepoWithTOInItsName".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("AnotherRepoWithTOInItsName".to_string())
    );
}

#[test]
fn create_chunked_repo_shard_with_shard_id_test() {
    let repo_shard =
        RepoShard::from_shard_id("repo_CHUNK_2_OF_15").expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "repo".to_string());
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, None);
}

#[test]
fn create_invalid_chunked_repo_shard_with_shard_id_test() {
    RepoShard::from_shard_id("repo_CHUNK__2_OF_15")
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::from_shard_id("repo_CHUNK_2_OFF_15")
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::from_shard_id("repo_CHUNK_2_OF_15_CHUNK")
        .expect_err("Should have failed in creating RepoShard");
}

#[test]
fn create_sized_chunked_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("repo_CHUNK_2_OF_15_SIZE_1000")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "repo".to_string());
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, Some(1000));
}

#[test]
fn create_invalid_sized_chunked_repo_shard_with_shard_id_test() {
    RepoShard::from_shard_id("repo_CHUNK__2_OF_15_SIZE_1000")
        .expect_err("Should have failed in creating RepoShard");
    RepoShard::from_shard_id("repo_CHUNK_2_OF_15_SIZE_1000_")
        .expect_err("Should have failed in creating RepoShard");
    let repo_shard = RepoShard::from_shard_id("repo_SIZE_1000_CHUNK_2_OF_15")
        .expect("Failed while creating RepoShare");
    assert_eq!(repo_shard.repo_name, "repo_SIZE_1000".to_string());
    assert_eq!(repo_shard.target_repo_name, None);
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, None);
    let repo_shard =
        RepoShard::from_shard_id("repo_SIZE_1000").expect("Failed while creating RepoShare");
    assert_eq!(repo_shard.repo_name, "repo_SIZE_1000".to_string());
}

#[test]
fn create_x_repo_chunked_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("repo_TO_another_repo_CHUNK_2_OF_15")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "repo".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("another_repo".to_string())
    );
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, None);
}

#[test]
fn create_x_repo_sized_chunked_repo_shard_with_shard_id_test() {
    let repo_shard = RepoShard::from_shard_id("repo_TO_another_repo_CHUNK_2_OF_15_SIZE_2000")
        .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "repo".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("another_repo".to_string())
    );
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, Some(2000));
}

#[test]
fn encode_decode_repo_name_test() {
    assert_eq!(
        decode_repo_name(&encode_repo_name("whatsapp/server")),
        "whatsapp/server".to_string()
    );
    assert_eq!(
        decode_repo_name(&encode_repo_name("fbsource")),
        "fbsource".to_string()
    );
    assert_eq!(
        decode_repo_name(&encode_repo_name("repo/with/lots/of/backslashes")),
        "repo/with/lots/of/backslashes".to_string()
    );
    assert_eq!(
        decode_repo_name(&encode_repo_name("repo+with+lots+of+pluses")),
        "repo+with+lots+of+pluses".to_string()
    );
    assert_eq!(
        decode_repo_name(&encode_repo_name("repo//with//double_slashes")),
        "repo//with//double_slashes".to_string()
    );
    assert_eq!(
        decode_repo_name(&encode_repo_name("/+/repo/+/")),
        "/+/repo/+/".to_string()
    );
}

#[test]
fn create_full_blown_repo_shard_with_shard_id_test() {
    let repo_shard =
        RepoShard::from_shard_id("whatsapp/server_TO_another+repo_CHUNK_2_OF_15_SIZE_2000")
            .expect("Failed while creating RepoShard");
    assert_eq!(repo_shard.repo_name, "whatsapp/server".to_string());
    assert_eq!(
        repo_shard.target_repo_name,
        Some("another+repo".to_string())
    );
    assert_eq!(repo_shard.chunk_id, Some(2));
    assert_eq!(repo_shard.total_chunks, Some(15));
    assert_eq!(repo_shard.chunk_size, Some(2000));
}


/*
 * 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;
use std::future::Future;
use std::io;
use std::path::Path;
use std::str::FromStr;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobstore::Loadable;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarkKey;
use bookmarks::BookmarksRef;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgChangesetId;
use mercurial_types::HgManifestId;
use mononoke_types::ChangesetId;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use services::Fb303Service;
use slog::error;
use slog::info;
use slog::Logger;
use tokio::io::AsyncBufReadExt;
use tokio::runtime::Runtime;

use crate::args::MononokeMatches;
use crate::monitoring;

pub const ARG_SHUTDOWN_GRACE_PERIOD: &str = "shutdown-grace-period";
pub const ARG_FORCE_SHUTDOWN_PERIOD: &str = "force-shutdown-period";

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum CreateStorage {
    ExistingOnly,
    ExistingOrCreate,
}

pub fn setup_repo_dir<P: AsRef<Path>>(data_dir: P, create: CreateStorage) -> Result<()> {
    let data_dir = data_dir.as_ref();

    if !data_dir.is_dir() {
        bail!("{:?} does not exist or is not a directory", data_dir);
    }

    // Validate directory layout
    #[allow(clippy::single_element_loop)]
    for subdir in &["blobs"] {
        let subdir = data_dir.join(subdir);

        if subdir.exists() && !subdir.is_dir() {
            bail!("{:?} already exists and is not a directory", subdir);
        }

        if !subdir.exists() {
            if CreateStorage::ExistingOnly == create {
                bail!("{:?} not found in ExistingOnly mode", subdir,);
            }
            fs::create_dir(&subdir)
                .with_context(|| format!("failed to create subdirectory {:?}", subdir))?;
        }
    }
    Ok(())
}

/// Resolve changeset id by either bookmark name, hg hash, or changset id hash
pub async fn csid_resolve(
    ctx: &CoreContext,
    container: impl RepoIdentityRef + BonsaiHgMappingRef + BookmarksRef,
    hash_or_bookmark: impl ToString,
) -> Result<ChangesetId, Error> {
    let res = csid_resolve_impl(ctx, &container, hash_or_bookmark).await;
    if let Ok(csid) = &res {
        info!(ctx.logger(), "changeset resolved as: {:?}", csid);
    }
    res
}

/// Read and resolve a list of changeset from file
pub async fn csids_resolve_from_file(
    ctx: &CoreContext,
    container: impl RepoIdentityRef + BonsaiHgMappingRef + BookmarksRef,
    filename: &str,
) -> Result<Vec<ChangesetId>, Error> {
    let file = tokio::fs::File::open(filename).await?;
    let file = tokio::io::BufReader::new(file);
    let mut lines = file.lines();
    let mut csids = vec![];
    while let Some(line) = lines.next_line().await? {
        csids.push(line);
    }

    stream::iter(csids)
        .map(|csid| csid_resolve_impl(ctx, &container, csid))
        .buffered(100)
        .try_collect::<Vec<_>>()
        .await
}

/// Resolve changeset id by either bookmark name, hg hash, or changset id hash
async fn csid_resolve_impl<C>(
    ctx: &CoreContext,
    container: &C,
    hash_or_bookmark: impl ToString,
) -> Result<ChangesetId, Error>
where
    C: BonsaiHgMappingRef + BookmarksRef,
{
    let hash_or_bookmark = hash_or_bookmark.to_string();
    if let Ok(name) = BookmarkKey::new(hash_or_bookmark.clone()) {
        if let Some(cs_id) = container.bookmarks().get(ctx.clone(), &name).await? {
            return Ok(cs_id);
        }
    }
    if let Ok(hg_id) = HgChangesetId::from_str(&hash_or_bookmark) {
        if let Some(cs_id) = container
            .bonsai_hg_mapping()
            .get_bonsai_from_hg(ctx, hg_id)
            .await?
        {
            return Ok(cs_id);
        }
    }
    if let Ok(cs_id) = ChangesetId::from_str(&hash_or_bookmark) {
        return Ok(cs_id);
    }
    bail!(
        "invalid (hash|bookmark) or does not exist in this repository: {}",
        hash_or_bookmark
    )
}

pub async fn get_root_manifest_id(
    ctx: &CoreContext,
    repo: BlobRepo,
    hash_or_bookmark: impl ToString,
) -> Result<HgManifestId, Error> {
    let cs_id = csid_resolve(ctx, &repo, hash_or_bookmark).await?;
    Ok(repo
        .derive_hg_changeset(ctx, cs_id)
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?
        .manifestid())
}

/// Get a tokio `Runtime` with potentially explicitly set number of core threads
pub fn create_runtime(
    log_thread_name_prefix: Option<&str>,
    worker_threads: Option<usize>,
) -> io::Result<Runtime> {
    let mut builder = tokio::runtime::Builder::new_multi_thread();
    builder.enable_all();
    builder.thread_name(log_thread_name_prefix.unwrap_or("tk"));
    if let Some(worker_threads) = worker_threads {
        builder.worker_threads(worker_threads);
    }
    builder.build()
}

/// Executes the future and waits for it to finish.
pub fn block_execute<F, Out, S: Fb303Service + Sync + Send + 'static>(
    future: F,
    fb: FacebookInit,
    app_name: &str,
    logger: &Logger,
    matches: &MononokeMatches,
    service: S,
) -> Result<Out, Error>
where
    F: Future<Output = Result<Out, Error>>,
{
    block_execute_impl(future, fb, app_name, logger, matches, service)
}

/// Executes the future and waits for it to finish.
fn block_execute_impl<F, Out, S: Fb303Service + Sync + Send + 'static>(
    future: F,
    fb: FacebookInit,
    app_name: &str,
    logger: &Logger,
    matches: &MononokeMatches,
    service: S,
) -> Result<Out, Error>
where
    F: Future<Output = Result<Out, Error>>,
{
    monitoring::start_fb303_server(fb, app_name, logger, matches, service)?;

    let result = matches.runtime().block_on(async {
        let stats_agg = stats::schedule_stats_aggregation_preview()
            .map_err(|_| Error::msg("Failed to create stats aggregation worker"))?;
        tokio::task::spawn(stats_agg);

        future.await
    });

    // Log error in glog format (main will log, but not with glog)
    result.map_err(move |e| {
        error!(logger, "Execution error: {:?}", e);
        // Shorten the error that main will print, given that already printed in glog form
        format_err!("Execution failed")
    })
}


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

//! Scaffolding that's generally useful to build CLI tools on top of Mononoke.

#![feature(never_type)]
#![feature(trait_alias)]

pub mod args;
pub mod helpers;
pub mod monitoring;


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

//! Scaffolding for service-level integration and monitoring.

use std::thread;

use anyhow::Error;
use fbinit::FacebookInit;
pub use mononoke_app::fb303::ReadyFlagService;
// Re-eport AliveService for convenience so callers do not have to get the services dependency to
// get AliveService.
pub use services::AliveService;
use services::Fb303Service;
use slog::info;
use slog::Logger;

use crate::args::MononokeMatches;

/// This is a lower-level function that requires you to spawn the stats aggregation future
/// yourself. This is useful if you'd like to be able to drop it in order to cancel it.
pub fn start_fb303_server<S: Fb303Service + Sync + Send + 'static>(
    fb: FacebookInit,
    service_name: &str,
    logger: &Logger,
    matches: &MononokeMatches,
    service: S,
) -> Result<Option<()>, Error> {
    let service_name = service_name.to_string();
    matches
        .value_of("fb303-thrift-port")
        .map(|port| {
            let port = port.parse().map_err(Error::from)?;
            info!(logger, "Initializing fb303 thrift server on port {}", port);

            thread::Builder::new()
                .name("fb303_thrift_service".to_owned())
                .spawn(move || {
                    services::run_service_framework(
                        fb,
                        service_name,
                        port,
                        0, // Disables separate status http server
                        Box::new(service),
                    )
                    .expect("failure while running thrift service framework")
                })
                .map_err(Error::from)?;

            Ok(())
        })
        .transpose()
}