/*
 * 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::path::Path;
#[cfg(windows)]
use std::path::PathBuf;
use std::process::Command;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use clap::Parser;
use edenfs_commands::is_command_enabled;
#[cfg(fbcode_build)]
use edenfs_telemetry::cli_usage::CliUsageSample;
#[cfg(fbcode_build)]
use edenfs_telemetry::send;
#[cfg(windows)]
use edenfs_utils::strip_unc_prefix;
use fbinit::FacebookInit;
use tracing_subscriber::filter::EnvFilter;

#[cfg(windows)]
const PYTHON_CANDIDATES: &[&str] = &[
    r"c:\tools\fb-python\fb-python39",
    r"c:\tools\fb-python\fb-python38",
    r"c:\Python39",
    r"c:\Python38",
];

#[cfg(windows)]
fn find_python() -> Option<PathBuf> {
    for candidate in PYTHON_CANDIDATES.iter() {
        let candidate = Path::new(candidate);
        let python = candidate.join("python.exe");

        if candidate.exists() && python.exists() {
            tracing::debug!("Found Python runtime at {}", python.display());
            return Some(python);
        }
    }
    None
}

#[cfg(windows)]
fn execute_par(par: PathBuf) -> Result<Command> {
    let python = find_python().ok_or_else(|| {
        anyhow!(
            "Unable to find Python runtime. Paths tried:\n\n - {}",
            PYTHON_CANDIDATES.join("\n - ")
        )
    })?;
    let mut python = Command::new(python);
    python.arg(par);
    Ok(python)
}

fn python_fallback() -> Result<Command> {
    if let Ok(args) = std::env::var("EDENFSCTL_REAL") {
        // We might get a command starting with python.exe here instead of a simple path.
        let mut parts = args.split_ascii_whitespace();
        let binary = parts
            .next()
            .ok_or_else(|| anyhow!("invalid fallback environment variable: {:?}", args))?;
        let mut cmd = Command::new(binary);
        cmd.args(parts);
        return Ok(cmd);
    }

    let binary = std::env::current_exe().context("unable to retrieve path to the executable")?;
    let binary =
        std::fs::canonicalize(binary).context("unable to canonicalize path to the executable")?;
    #[cfg(windows)]
    let binary = strip_unc_prefix(binary);
    let libexec = binary.parent().ok_or_else(|| {
        anyhow!(
            "unable to retrieve parent directory to '{}'",
            binary.display()
        )
    })?;

    let executable = libexec.join(if cfg!(windows) {
        "edenfsctl.real.exe"
    } else {
        "edenfsctl.real"
    });
    tracing::debug!("trying {:?}", executable);
    if executable.exists() {
        return Ok(Command::new(executable));
    }

    // On Windows we are shipping the Python edenfsctl as PAR file that is not executable by itself
    #[cfg(windows)]
    {
        let par = libexec.join("edenfsctl.real.par");
        tracing::debug!("trying {:?}", par);

        if par.exists() {
            return execute_par(par);
        }
    }

    Err(anyhow!("unable to locate fallback binary"))
}

fn fallback(reason: Option<clap::Error>) -> Result<i32> {
    if std::env::var("EDENFS_LOG").is_ok() {
        setup_logging();
    }

    if let Some(reason) = reason {
        tracing::debug!(%reason, "falling back to Python");
    }

    let mut cmd = python_fallback()?;
    // skip arg0
    cmd.args(std::env::args().skip(1));

    // Users have PYTHONHOME and PYTHONPATH variables
    // that break the python version of edenfsctl since it will fail to
    // import modules. So, let's strip the PYTHONHOME and PYTHONPATH variables.
    cmd.env_remove("PYTHONHOME");
    cmd.env_remove("PYTHONPATH");

    tracing::debug!("Falling back to {:?}", cmd);

    // Create a subprocess to run Python edenfsctl
    let status = cmd
        .status()
        .with_context(|| format!("failed to execute: {:?}", cmd))?;
    Ok(status.code().unwrap_or(1))
}

/// Setup tracing logging. If we are in development mode, we use the fancier logger, otherwise a
/// simple logger for production use. Logs will be printined to stderr when `--debug` flag is
/// passed.
fn setup_logging() {
    let subscriber = tracing_subscriber::fmt();
    #[cfg(debug_assertions)]
    let subscriber = subscriber.pretty();
    let subscriber = subscriber.with_env_filter(EnvFilter::from_env("EDENFS_LOG"));

    if let Err(e) = subscriber.try_init() {
        eprintln!(
            "Unable to initialize logger. Logging will be disabled. Cause: {:?}",
            e
        );
    }
}

fn rust_main(cmd: edenfs_commands::MainCommand) -> Result<i32> {
    if cmd.debug {
        setup_logging();
    }
    Ok(cmd.run()?)
}

/// This function takes care of the fallback logic, hijack supported subcommand
/// to Rust implementation and forward the rest to Python.
fn wrapper_main() -> Result<i32> {
    if std::env::var("EDENFSCTL_ONLY_RUST").is_ok() {
        let cmd = edenfs_commands::MainCommand::parse();
        rust_main(cmd)
    } else if std::env::var("EDENFSCTL_SKIP_RUST").is_ok() {
        fallback(None)
    } else {
        match edenfs_commands::MainCommand::try_parse() {
            // The command is defined in Rust, but check whether it's "enabled"
            // for Rust or else fall back to Python.
            Ok(cmd) => {
                if cmd.is_enabled() {
                    rust_main(cmd)
                } else {
                    fallback(None)
                }
            }
            // If the command is defined in Rust, then --help will cause
            // try_parse() to return a DisplayHelp error.  In that case, we
            // should check whether the Rust version of the command is "enabled"
            // to decide whether to print Rust or Python help.
            //
            // If the command isn't defined in Rust then try_parse will fail
            // UnknownArgument (whether or not --help was requested) and we
            // should fall back to Python.
            Err(e) => {
                if (e.kind() == clap::ErrorKind::DisplayHelp
                    || e.kind() == clap::ErrorKind::DisplayHelpOnMissingArgumentOrSubcommand)
                    && should_use_rust_help(std::env::args(), &None).unwrap_or(false)
                {
                    e.exit()
                } else {
                    fallback(Some(e))
                }
            }
        }
    }
}

fn should_use_rust_help<T>(args: T, etc_eden_dir_override: &Option<&Path>) -> Result<bool>
where
    T: Iterator<Item = String>,
{
    // This is gross, but clap v3 doesn't let us make --help a normal bool flag.
    // This means we can't successfully parse a command when --help is
    // requested, so here we manually extract the subcommand name in order to
    // check whether it's enabled for Rust.
    let subcommand_name = args
        .skip(1)
        .find(|a| !a.starts_with('-'))
        .ok_or(anyhow!("missing subcommand"))?;

    Ok(is_command_enabled(
        &subcommand_name,
        &etc_eden_dir_override.map(Path::to_owned),
    ))
}

#[fbinit::main]
fn main(_fb: FacebookInit) -> Result<()> {
    // NOTE: if you are considering passing `FacebookInit` down, you may want to check
    // [`fbinit::expect_init`].
    #[cfg(fbcode_build)]
    let mut sample = CliUsageSample::build(_fb);

    let code = match wrapper_main() {
        Ok(code) => Ok(code),
        Err(e) => {
            #[cfg(fbcode_build)]
            sample.set_exception(&e);
            Err(e)
        }
    };

    #[cfg(fbcode_build)]
    {
        sample.set_exit_code(*code.as_ref().unwrap_or(&1));
        send(sample.builder);
    }

    match code {
        Ok(code) => std::process::exit(code),
        Err(e) => Err(e),
    }
}

#[cfg(test)]
mod tests {
    use std::fs::File;
    use std::io::Write;

    use anyhow::Result;
    use tempfile::TempDir;

    use super::should_use_rust_help;

    macro_rules! args {
        ( $( $x:expr ),* ) => (
            {
                let mut v = Vec::new();
                $(
                    v.push($x.to_string());
                )*
                v.into_iter()
            }
        );
    }

    #[test]
    fn test_should_use_rust_help() -> Result<()> {
        assert!(should_use_rust_help(args!["eden.exe", "minitop"], &None)?);
        {
            let dir = TempDir::new()?;
            assert!(!should_use_rust_help(
                args!["eden.exe", "redirect"],
                &Some(dir.path())
            )?,);
            assert!(!should_use_rust_help(
                args!["eden.exe", "--xyz", "redirect"],
                &Some(dir.path())
            )?,);
        }
        {
            let dir = TempDir::new()?;
            let rollout_path = dir.path().join("edenfsctl_rollout.json");
            let mut rollout_file = File::create(rollout_path)?;
            writeln!(rollout_file, r#"{{"redirect": true}}"#)?;

            assert!(should_use_rust_help(
                args!["eden.exe", "redirect"],
                &Some(dir.path())
            )?,);
        }

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

pub use stack_config_derive::*;

pub mod __private {
    pub use serde_derive::Deserialize;

    pub trait ConfigLayerType {
        type Layer;
    }

    impl<T: ConfigLayerType> ConfigLayerType for Option<T> {
        type Layer = Option<<T as ConfigLayerType>::Layer>;
    }

    pub trait ConfigLayer: Default {
        type Product;

        fn finalize(self) -> Result<Self::Product, String>;
        fn merge(&mut self, other: Self);
    }

    impl<T: ConfigLayer> ConfigLayer for Option<T> {
        type Product = Option<<T as ConfigLayer>::Product>;

        fn finalize(self) -> Result<Self::Product, String> {
            match self {
                Some(inner) => inner.finalize().map(Some),
                None => Ok(None),
            }
        }

        fn merge(&mut self, other: Self) {
            match (self.as_mut(), other) {
                (Some(lhs), Some(rhs)) => lhs.merge(rhs),
                // if we don't have base layer yet, we need to initialize it with the default value.
                (None, Some(rhs)) => {
                    let mut lhs = <T as Default>::default();
                    lhs.merge(rhs);
                    *self = Some(lhs);
                }
                (_, None) => (),
            }
        }
    }
}


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

use serde::Deserialize;
use stack_config::StackConfig;

#[derive(Deserialize, StackConfig, Debug, PartialEq)]
struct ConfigNested {
    field: String,
    remainder: bool,
    another: String,
    #[stack(default)]
    default: bool,
}

#[derive(Deserialize, StackConfig, Debug, PartialEq)]
struct Config {
    #[stack(default)]
    field: String,

    flag: bool,

    #[stack(default = "default_list")]
    list: Vec<String>,

    #[stack(default)]
    opt: Option<String>,

    #[stack(nested)]
    nested: ConfigNested,

    #[stack(nested)]
    nested2: Option<ConfigNested>,

    #[stack(nested)]
    partial: Option<ConfigNested>,
}

fn default_list() -> Vec<String> {
    vec!["default".into(), "list".into()]
}

#[test]
fn basic_test() {
    let mut config_loader = Config::loader();
    config_loader.load(
        toml::from_str(
            r#"
field = "1234"

flag = false

[nested]
field = "1234"
remainder = true
"#,
        )
        .unwrap(),
    );
    config_loader.load(
        toml::from_str(
            r#"
flag = true

unknown = "12345"

[nested]
field = "hello"
another = "world"

[partial]
field = "hello"
another = "world"
remainder = true
"#,
        )
        .unwrap(),
    );

    let result = config_loader.build().unwrap();
    assert_eq!(
        result,
        Config {
            field: "1234".into(),
            flag: true,
            list: default_list(),
            opt: None,
            nested: ConfigNested {
                field: "hello".into(),
                remainder: true,
                another: "world".into(),
                default: false,
            },
            nested2: None,
            partial: Some(ConfigNested {
                field: "hello".into(),
                remainder: true,
                another: "world".into(),
                default: false,
            }),
        }
    );
}


/*
 * 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 darling::FromMeta;
use proc_macro2::TokenStream as TokenStream2;
use proc_macro_error::abort;
use quote::quote;
use syn::Ident;
use syn::Meta;
use syn::Path;

#[derive(Debug, FromMeta)]
struct MetaArgsFlag {
    #[darling(default)]
    default: Option<()>,

    #[darling(default)]
    nested: darling::util::Flag,

    #[darling(default)]
    merge: Option<String>,
}

#[derive(Debug, FromMeta)]
struct MetaArgsFunc {
    default: String,

    #[darling(default)]
    nested: darling::util::Flag,

    #[darling(default)]
    merge: Option<String>,
}

#[derive(Debug)]
pub struct ConfigField {
    pub name: Ident,

    /// Specify where to look for default value
    ///
    /// None - no default
    /// Some(None) - use std::default::Default
    /// Some(path) - use path() as default
    pub default: Option<Option<Path>>,

    pub nested: bool,

    pub merge: Option<Path>,
}

impl ConfigField {
    pub fn new(name: Ident) -> Self {
        Self {
            name,
            default: None,
            nested: false,
            merge: None,
        }
    }

    fn to_path(meta: &Meta, field: &str, st: &str) -> Path {
        match syn::parse_str::<Path>(st) {
            Ok(path) => path,
            Err(e) => abort!(meta, "{} must be a path: {}", field, e),
        }
    }

    pub fn new_with_meta(name: Ident, meta: &Meta) -> Self {
        // We have two struct here because I can't find a way to let darling accept an optional
        // argument. i.e. `stack(default)` and `stack(default = "func")`.
        if let Ok(args) = MetaArgsFlag::from_meta(&meta) {
            let nested = args.nested.is_present();

            if args.default.is_some() && nested {
                abort!(meta, "can't use nested and default at the same time");
            }

            if args.merge.is_some() && nested {
                abort!(meta, "can't use nested and merge at the same time");
            }

            let merge = args.merge.map(|x| Self::to_path(meta, "merge", &x));

            Self {
                name,
                default: args.default.map(|_| None),
                nested,
                merge,
            }
        } else if let Ok(args) = MetaArgsFunc::from_meta(&meta) {
            let nested = args.nested.is_present();

            if nested {
                abort!(meta, "can't use nested and default at the same time");
            }

            let merge = args.merge.map(|x| Self::to_path(meta, "merge", &x));

            Self {
                name,
                default: Some(Some(Self::to_path(meta, "default", &args.default))),
                nested,
                merge,
            }
        } else {
            abort!(meta, "uanble to parse attribute");
        }
    }

    pub fn default(&self) -> TokenStream2 {
        let field = &self.name;
        if let Some(default) = &self.default {
            if let Some(func) = default {
                // Default attribute with user-defined function call
                quote! {
                    #field: Some(#func())
                }
            } else {
                // Default attribute with no parameter, use `Default::default`
                quote! {
                    #field: Some(::std::default::Default::default())
                }
            }
        } else if self.nested {
            quote! {
                #field: std::default::Default::default()
            }
        } else {
            // No default supplied, do nothing
            quote! {
                #field: None
            }
        }
    }

    pub fn finalize(&self) -> TokenStream2 {
        let field = &self.name;
        let field_str = field.to_string();

        if self.nested {
            quote! {
                #field: self.#field.finalize()?
            }
        } else {
            quote! {
                #field: self.#field.ok_or_else(|| format!("field {} is missing", #field_str))?
            }
        }
    }

    pub fn merge(&self) -> TokenStream2 {
        let field = &self.name;

        if self.nested {
            quote! {
                self.#field.merge(other.#field);
            }
        } else {
            if let Some(merge) = &self.merge {
                quote! {
                    if let Some(val) = other.#field {
                        if let Some(current) = self.#field.as_mut() {
                            #merge(current, val);
                        } else {
                            self.#field = Some(val);
                        }
                    }
                }
            } else {
                quote! {
                    if let Some(val) = other.#field {
                        self.#field = Some(val);
                    }
                }
            }
        }
    }
}


/*
 * 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 proc_macro::TokenStream;
use proc_macro_error::abort;
use proc_macro_error::proc_macro_error;
use quote::format_ident;
use quote::quote;
use syn::parse_macro_input;
use syn::parse_quote;
use syn::visit_mut;
use syn::visit_mut::VisitMut;
use syn::Attribute;
use syn::Field;
use syn::Ident;
use syn::ItemStruct;
use syn::Visibility;

mod field;

use crate::field::ConfigField;

fn layer_ident(name: &Ident) -> Ident {
    format_ident!("__StackConfig_{}Layer", name)
}

#[derive(Debug)]
struct StackConfig {
    /// Visibility
    vis: Option<Visibility>,
    /// Struct name
    name: Ident,
    /// List of fields
    fields: Vec<ConfigField>,
}

/// Visitor related impls
impl StackConfig {
    fn take_attr(&self, attrs: &mut Vec<Attribute>) -> Option<Attribute> {
        if let Some(pos) = attrs.iter().position(|attr| {
            if let Ok(meta) = attr.parse_meta() {
                meta.path().is_ident("stack")
            } else {
                false
            }
        }) {
            Some(attrs.remove(pos))
        } else {
            None
        }
    }
}

impl VisitMut for StackConfig {
    fn visit_item_struct_mut(&mut self, item: &mut ItemStruct) {
        self.vis = Some(item.vis.clone());
        item.ident = layer_ident(&item.ident);
        item.attrs.push(parse_quote! { #[derive(Debug)]});
        visit_mut::visit_item_struct_mut(self, item);
    }

    fn visit_field_mut(&mut self, field: &mut Field) {
        let ident = match &field.ident {
            Some(ident) => ident,
            None => abort!(field, "unnamed field is not supported."),
        };

        let config_field = if let Some(attr) = self.take_attr(&mut field.attrs) {
            let meta = match attr.parse_meta() {
                Ok(meta) => meta,
                Err(e) => abort!(attr, "unable to parse attribute: {:?}", e),
            };

            ConfigField::new_with_meta(ident.clone(), &meta)
        } else {
            ConfigField::new(ident.clone())
        };

        if config_field.nested {
            field.attrs.push(parse_quote! { #[serde(default)] });
            let old = &field.ty;
            field.ty = parse_quote! { <#old as ::stack_config::__private::ConfigLayerType>::Layer };
        } else {
            let old = &field.ty;
            field.ty = parse_quote! { Option<#old> };
        }

        self.fields.push(config_field);
        visit_mut::visit_field_mut(self, field);
    }
}

impl StackConfig {
    fn new(name: Ident) -> Self {
        Self {
            vis: None,
            name,
            fields: Vec::new(),
        }
    }

    /// Generates `std::default` implementation
    fn default_impl(&self) -> proc_macro2::TokenStream {
        let layer_ty = layer_ident(&self.name);
        let fields = self.fields.iter().map(|field| field.default());

        quote! {
            impl ::std::default::Default for #layer_ty {
                fn default() -> Self {
                    #layer_ty {
                        #(#fields),*
                    }
                }
            }
        }
    }

    /// Generates conversion into the concrete type
    fn layer_impl(&self) -> proc_macro2::TokenStream {
        let name = &self.name;
        let layer_ty = layer_ident(&self.name);
        let finalizes = self.fields.iter().map(|field| field.finalize());
        let merges = self.fields.iter().map(|field| field.merge());

        quote! {
            impl ::stack_config::__private::ConfigLayer for #layer_ty {
                type Product = #name;

                fn finalize(self) -> Result<#name, String> {
                    Ok(#name {
                        #(#finalizes),*
                    })
                }

                fn merge(&mut self, other: #layer_ty) {
                    #(#merges)*
                }
            }
        }
    }

    // generates `ConfigLayerType` for type calculations
    fn layer_type_impl(&self) -> proc_macro2::TokenStream {
        let name = &self.name;
        let layer_ty = layer_ident(&self.name);

        quote! {
            impl ::stack_config::__private::ConfigLayerType for #name {
                type Layer = #layer_ty;
            }
        }
    }

    fn builder(self) -> proc_macro2::TokenStream {
        let product = &self.name;
        let loader = format_ident!("{}Loader", self.name);
        let layer_ty = layer_ident(&self.name);
        let vis = self.vis.unwrap_or(Visibility::Inherited);

        quote! {
            #vis struct #loader {
                product: #layer_ty,
            }

            impl #loader {
                fn new() -> Self {
                    Self { product: Default::default() }
                }

                pub fn load(&mut self, layer: #layer_ty) {
                    ::stack_config::__private::ConfigLayer::merge(&mut self.product, layer);
                }

                pub fn build(self) -> ::std::result::Result<#product, String> {
                    ::stack_config::__private::ConfigLayer::finalize(self.product)
                }
            }

            impl #product {
                #vis fn loader() -> #loader {
                    #loader::new()
                }
            }
        }
    }
}

#[proc_macro_error]
#[proc_macro_derive(StackConfig, attributes(stack))]
pub fn stack_config(input: TokenStream) -> TokenStream {
    let mut input = parse_macro_input!(input as ItemStruct);

    let mut stack = StackConfig::new(input.ident.clone());
    stack.visit_item_struct_mut(&mut input);

    let layer = stack.layer_impl();
    let default = stack.default_impl();
    let layer_type = stack.layer_type_impl();
    let build = stack.builder();

    let result = quote! {
        #[derive(::stack_config::__private::Deserialize)]
        #[allow(non_camel_case_types)]
        #input

        #default
        #layer_type
        #layer

        #build
    };

    result.into()
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "eden_config crate",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=config_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "eden_config.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "overlay crate",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=serialization__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "overlay.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "eden crate
eden_config config_thrift
fb303_core fb303_core",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "eden.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "eden thrift
streamingeden crate",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=thrift-streaming__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "../streamingeden.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


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

mod trie;

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

use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use commit_graph::ArcCommitGraph;
use context::CoreContext;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use metaconfig_types::AclRegion;
use metaconfig_types::AclRegionConfig;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use trie::PrefixTrieWithRules;

#[derive(Debug, Default, Clone, Eq, PartialEq)]
pub struct AclRegionRuleDescriptor {
    pub name: String,
    pub hipster_acl: String,
}

pub enum AssociatedRulesResult {
    AclRegionsDisabled,
    Rules(Vec<AclRegionRuleDescriptor>),
}

impl AssociatedRulesResult {
    pub fn hipster_acls(&self) -> Vec<&str> {
        match self {
            AssociatedRulesResult::AclRegionsDisabled => Vec::new(),
            AssociatedRulesResult::Rules(rules) => {
                rules.iter().map(|rule| rule.hipster_acl.as_str()).collect()
            }
        }
    }
}

#[async_trait]
#[facet::facet]
pub trait AclRegions: Send + Sync {
    async fn associated_rules(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        path: Option<&MPath>,
    ) -> Result<AssociatedRulesResult>;
}

struct AclRegionsImpl {
    path_rules_index: PrefixTrieWithRules,
    commit_graph: ArcCommitGraph,
}

impl AclRegionsImpl {
    fn new(config: &AclRegionConfig, commit_graph: ArcCommitGraph) -> AclRegionsImpl {
        let mut path_rules_index = PrefixTrieWithRules::new();
        for rule in &config.allow_rules {
            path_rules_index.add_rule(Arc::new(rule.clone()));
        }
        AclRegionsImpl {
            path_rules_index,
            commit_graph,
        }
    }

    async fn is_commit_descendant_of_any(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        candidates: &[ChangesetId],
    ) -> Result<bool> {
        let mut is_descendant_results = stream::iter(candidates)
            .map(|candidate| self.commit_graph.is_ancestor(ctx, *candidate, cs_id))
            .boxed()
            .buffered(10);

        while let Some(is_descendant) = is_descendant_results.try_next().await? {
            if is_descendant {
                return Ok(true);
            }
        }

        Ok(false)
    }

    async fn is_commit_in_region(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        region: &AclRegion,
    ) -> Result<bool> {
        let is_any_root_descendant = region.roots.is_empty()
            || self
                .is_commit_descendant_of_any(ctx, cs_id, &region.roots)
                .await?;
        let is_any_head_descendant = self
            .is_commit_descendant_of_any(ctx, cs_id, &region.heads)
            .await?;
        Ok(is_any_root_descendant && !is_any_head_descendant)
    }
}

#[async_trait]
impl AclRegions for AclRegionsImpl {
    async fn associated_rules(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        path: Option<&MPath>,
    ) -> Result<AssociatedRulesResult> {
        let matched_rules: HashMap<_, _> =
            stream::iter(self.path_rules_index.associated_rules(path))
                .map(|((name, region_index), rule)| async move {
                    let region = &rule.regions.get(region_index.0).ok_or_else(|| {
                        anyhow!(
                            "Incorrect region index {} for rule {}",
                            region_index.0,
                            rule.name
                        )
                    })?;
                    anyhow::Ok(
                        self.is_commit_in_region(ctx, cs_id, region)
                            .await?
                            .then_some((name, rule)),
                    )
                })
                .buffered(10)
                .filter_map(|rule| async { rule.transpose() })
                .try_collect()
                .await?;

        Ok(AssociatedRulesResult::Rules(
            matched_rules
                .into_iter()
                .map(|(name, rule)| AclRegionRuleDescriptor {
                    name,
                    hipster_acl: rule.hipster_acl.clone(),
                })
                .collect(),
        ))
    }
}

struct DisabledAclRegions {}

#[async_trait]
impl AclRegions for DisabledAclRegions {
    async fn associated_rules(
        &self,
        _ctx: &CoreContext,
        _cs_id: ChangesetId,
        _path: Option<&MPath>,
    ) -> Result<AssociatedRulesResult> {
        Ok(AssociatedRulesResult::AclRegionsDisabled)
    }
}

pub fn build_acl_regions(
    config: Option<&AclRegionConfig>,
    commit_graph: ArcCommitGraph,
) -> ArcAclRegions {
    match config {
        Some(config) => Arc::new(AclRegionsImpl::new(config, commit_graph)),
        None => build_disabled_acl_regions(),
    }
}

pub fn build_disabled_acl_regions() -> ArcAclRegions {
    Arc::new(DisabledAclRegions {})
}


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

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

use metaconfig_types::AclRegionRule;
use mononoke_types::MPath;
use mononoke_types::MPathElement;

#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)]
pub struct RegionIndex(pub usize);

/// A trie that stores all paths from all regions for efficient rule matching
/// by path.
#[derive(Default)]
pub struct PrefixTrieWithRules {
    children: HashMap<MPathElement, PrefixTrieWithRules>,
    associated_rules: Vec<(RegionIndex, Arc<AclRegionRule>)>,
}

impl PrefixTrieWithRules {
    pub fn new() -> PrefixTrieWithRules {
        PrefixTrieWithRules {
            children: HashMap::new(),
            associated_rules: vec![],
        }
    }

    /// Add all paths of all regions of the rule and store the rule at the end
    /// of each path.
    pub fn add_rule(&mut self, rule: Arc<AclRegionRule>) {
        for (region_index, region) in rule.regions.iter().enumerate() {
            for path in &region.path_prefixes {
                self.add_rule_on_path(
                    path.iter().flatten(),
                    RegionIndex(region_index),
                    rule.clone(),
                );
            }
        }
    }

    /// Add a given path and store the associated rule at the end.
    fn add_rule_on_path<'a>(
        &mut self,
        path: impl IntoIterator<Item = &'a MPathElement>,
        region_index: RegionIndex,
        rule: Arc<AclRegionRule>,
    ) {
        let mut iter = path.into_iter();
        match iter.next() {
            Some(element) => self
                .children
                .entry(element.clone())
                .or_default()
                .add_rule_on_path(iter, region_index, rule),
            None => self.associated_rules.push((region_index, rule)),
        }
    }

    /// Traverse a given path and collect all rules alongside it. Return them
    /// deduplicated by (rule name, matched region index)
    pub fn associated_rules(
        &self,
        path: Option<&MPath>,
    ) -> HashMap<(String, RegionIndex), Arc<AclRegionRule>> {
        self.associated_rules_inner(path.into_iter().flatten())
            .into_iter()
            .map(|(region_index, rule)| ((rule.name.clone(), region_index), rule))
            .collect()
    }

    fn associated_rules_inner<'a>(
        &self,
        path: impl IntoIterator<Item = &'a MPathElement>,
    ) -> Vec<(RegionIndex, Arc<AclRegionRule>)> {
        let mut iter = path.into_iter();
        let mut rules = match iter.next() {
            None => vec![],
            Some(element) => match self.children.get(element) {
                Some(child) => child.associated_rules_inner(iter),
                None => vec![],
            },
        };
        rules.reserve(self.associated_rules.len());
        rules.extend(self.associated_rules.iter().cloned());
        rules
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use metaconfig_types::AclRegion;
    use pretty_assertions::assert_eq;

    use super::*;

    fn path(raw: &str) -> Option<MPath> {
        MPath::new_opt(raw).unwrap()
    }

    struct TestData {
        path: Option<MPath>,
        expected_regions: HashSet<(String, RegionIndex)>,
    }

    impl TestData {
        fn new(raw_path: &str, raw_regions: &[(&str, usize)]) -> Self {
            let expected_regions = raw_regions
                .iter()
                .map(|(name, index)| (name.to_string(), RegionIndex(*index)))
                .collect::<HashSet<_>>();
            let path = path(raw_path);
            Self {
                path,
                expected_regions,
            }
        }

        fn verify(&self, trie: &PrefixTrieWithRules) {
            let rules = trie.associated_rules(self.path.as_ref());
            assert_eq!(
                rules.into_keys().collect::<HashSet<_>>(),
                self.expected_regions
            );
        }
    }

    #[fbinit::test]
    fn test_prefix_trie() {
        let mut trie = PrefixTrieWithRules::new();
        trie.add_rule(Arc::new(AclRegionRule {
            name: "rule1".to_string(),
            regions: vec![
                AclRegion {
                    roots: vec![],
                    heads: vec![],
                    path_prefixes: vec![path("a"), path("c/d")],
                },
                AclRegion {
                    roots: vec![],
                    heads: vec![],
                    path_prefixes: vec![path("c"), path("c/d")],
                },
            ],
            hipster_acl: "acl1".to_string(),
        }));
        trie.add_rule(Arc::new(AclRegionRule {
            name: "rule2".to_string(),
            regions: vec![AclRegion {
                roots: vec![],
                heads: vec![],
                path_prefixes: vec![path("a/b"), path("b")],
            }],
            hipster_acl: "acl2".to_string(),
        }));
        trie.add_rule(Arc::new(AclRegionRule {
            name: "rule3".to_string(),
            regions: vec![AclRegion {
                roots: vec![],
                heads: vec![],
                path_prefixes: vec![path("")],
            }],
            hipster_acl: "acl3".to_string(),
        }));

        let test_data = vec![
            TestData::new("a", &[("rule1", 0), ("rule3", 0)]),
            TestData::new("aa", &[("rule3", 0)]),
            TestData::new("a/b/c", &[("rule1", 0), ("rule2", 0), ("rule3", 0)]),
            TestData::new("a/bb", &[("rule1", 0), ("rule3", 0)]),
            TestData::new("b/a", &[("rule2", 0), ("rule3", 0)]),
            TestData::new("c/d", &[("rule1", 0), ("rule1", 1), ("rule3", 0)]),
        ];

        for (index, data) in test_data.iter().enumerate() {
            eprintln!("Verifying test data #{}", index);
            data.verify(&trie);
        }
    }
}


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

use acl_regions::AclRegions;
use acl_regions::AclRegionsRef;
use acl_regions::AssociatedRulesResult;
use anyhow::Result;
use bonsai_hg_mapping::BonsaiHgMapping;
use bookmarks::Bookmarks;
use changeset_fetcher::ChangesetFetcher;
use changesets::Changesets;
use context::CoreContext;
use fbinit::FacebookInit;
use filestore::FilestoreConfig;
use metaconfig_types::AclRegion;
use metaconfig_types::AclRegionConfig;
use metaconfig_types::AclRegionRule;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use pretty_assertions::assert_eq;
use repo_blobstore::RepoBlobstore;
use repo_derived_data::RepoDerivedData;
use test_repo_factory::TestRepoFactory;
use tests_utils::drawdag::create_from_dag;

#[facet::container]
#[derive(Clone)]
struct Repo {
    #[facet]
    acl_regions: dyn AclRegions,

    #[facet]
    bonsai_hg_mapping: dyn BonsaiHgMapping,

    #[facet]
    bookmarks: dyn Bookmarks,

    #[facet]
    changesets: dyn Changesets,

    #[facet]
    filestore_config: FilestoreConfig,

    #[facet]
    repo_blobstore: RepoBlobstore,

    #[facet]
    repo_derived_data: RepoDerivedData,

    #[facet]
    changeset_fetcher: dyn ChangesetFetcher,
}

fn path(p: &str) -> Option<MPath> {
    MPath::new_opt(p).unwrap()
}

struct TestData {
    cs_id: ChangesetId,
    path: Option<MPath>,
    expected_names: HashSet<String>,
}

impl TestData {
    fn new(cs_id: ChangesetId, raw_path: &str, raw_names: &[&str]) -> Self {
        let expected_names = raw_names
            .iter()
            .map(ToString::to_string)
            .collect::<HashSet<String>>();
        let path = path(raw_path);
        Self {
            cs_id,
            path,
            expected_names,
        }
    }

    async fn verify(&self, ctx: &CoreContext, acl_regions: &dyn AclRegions) -> Result<()> {
        let res = acl_regions
            .associated_rules(ctx, self.cs_id, self.path.as_ref())
            .await?;
        let rules = match res {
            AssociatedRulesResult::AclRegionsDisabled => {
                panic!("AclRegions should work in test, but they're disabled")
            }
            AssociatedRulesResult::Rules(rules) => rules,
        };
        assert_eq!(
            rules
                .into_iter()
                .map(|descriptor| descriptor.name)
                .collect::<HashSet<String>>(),
            self.expected_names
        );
        Ok(())
    }
}

#[fbinit::test]
async fn test_acl_regions(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let mut factory = TestRepoFactory::new(fb)?;
    let repo: Repo = factory.build().await?;
    let dag = r#"
          I
         /
    A-B-C-D-G-H
       \   /
        E-F
    "#;
    let mapping = create_from_dag(&ctx, &repo, dag).await?;
    let a = mapping["A"];
    let b = mapping["B"];
    let c = mapping["C"];
    let d = mapping["D"];
    let e = mapping["E"];
    let f = mapping["F"];
    let g = mapping["G"];
    let h = mapping["H"];
    let i = mapping["I"];

    let rule1 = AclRegionRule {
        name: "rule1".to_string(),
        regions: vec![AclRegion {
            roots: vec![],
            heads: vec![b],
            path_prefixes: vec![path("")],
        }],
        hipster_acl: "acl1".to_string(),
    };

    let rule2 = AclRegionRule {
        name: "rule2".to_string(),
        regions: vec![
            AclRegion {
                roots: vec![a],
                heads: vec![],
                path_prefixes: vec![path("a")],
            },
            AclRegion {
                roots: vec![h],
                heads: vec![],
                path_prefixes: vec![path("aa")],
            },
        ],
        hipster_acl: "acl2".to_string(),
    };

    let rule3 = AclRegionRule {
        name: "rule3".to_string(),
        regions: vec![
            AclRegion {
                roots: vec![b],
                heads: vec![e, i],
                path_prefixes: vec![path("b")],
            },
            AclRegion {
                roots: vec![e],
                heads: vec![g],
                path_prefixes: vec![path("b/c")],
            },
        ],
        hipster_acl: "acl3".to_string(),
    };

    // We're re-creating the repo with a correct config after we got hashes
    // from the mapping.
    let repo: Repo = factory
        .with_config_override(|config| {
            config.acl_region_config = Some(AclRegionConfig {
                allow_rules: vec![rule1, rule2, rule3],
            });
        })
        .build()
        .await?;

    let same_mapping = create_from_dag(&ctx, &repo, dag).await?;
    assert_eq!(mapping, same_mapping);

    let acl_regions = repo.acl_regions();

    let test_data = vec![
        TestData::new(a, "a", &["rule1", "rule2"]),
        TestData::new(a, "b", &["rule1"]),
        TestData::new(b, "", &[]),
        TestData::new(b, "b", &["rule3"]),
        TestData::new(c, "a", &["rule2"]),
        TestData::new(c, "b", &["rule3"]),
        TestData::new(d, "b", &["rule3"]),
        TestData::new(e, "b", &[]),
        TestData::new(e, "b/c", &["rule3"]),
        TestData::new(f, "b/c", &["rule3"]),
        TestData::new(g, "aa", &[]),
        TestData::new(g, "b", &[]),
        TestData::new(g, "b/c", &[]),
        TestData::new(h, "b", &[]),
        TestData::new(h, "aa", &["rule2"]),
        TestData::new(i, "b", &[]),
    ];

    for (index, data) in test_data.iter().enumerate() {
        eprintln!("Verifying test data #{}", index);
        data.verify(&ctx, acl_regions).await?;
    }

    Ok(())
}


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

use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use blobrepo::BlobRepo;
use bulkops::Direction;
use bulkops::PublicChangesetBulkFetch;
use bulkops::MAX_FETCH_STEP;
use changesets::ChangesetsArc;
use clap::Parser;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use futures::future;
use futures::future::TryFutureExt;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_app::args::RepoArgs;
use mononoke_app::MononokeAppBuilder;
use phases::PhasesArc;
use repo_identity::RepoIdentityRef;
use tokio::runtime::Handle;

/// Benchmark bulkops.
#[derive(Parser)]
struct BenchmarkArgs {
    #[clap(flatten)]
    repo: RepoArgs,

    /// Save results as a baseline under this name, for comparison.
    #[clap(long)]
    save_baseline: Option<String>,

    /// Compare to named baseline instead of last run.
    #[clap(long)]
    use_baseline: Option<String>,

    /// Limit to benchmarks whose name contains this string. Repetition tightens the filter.
    #[clap(long)]
    filter: Vec<String>,
}

pub fn bench_stream<'a, F, S, O, E>(
    c: &'a mut Criterion,
    ctx: &'a CoreContext,
    runtime: &Handle,
    group: String,
    fetcher: &'a PublicChangesetBulkFetch,
    to_stream: F,
) where
    F: Fn(&'a CoreContext, &'a PublicChangesetBulkFetch) -> S,
    S: Stream<Item = Result<O, E>> + 'a,
    E: std::fmt::Debug,
{
    let mut group = c.benchmark_group(group);
    let num_to_load = 100000;
    group.throughput(Throughput::Elements(num_to_load));
    let step = MAX_FETCH_STEP;
    group.bench_with_input(
        BenchmarkId::from_parameter(step),
        &num_to_load,
        |b, &num_to_load| {
            let test = || async {
                let mut loaded: u64 = 0;
                let stream = to_stream(ctx, fetcher).take_while(|_entry| {
                    loaded += 1;
                    future::ready(loaded < num_to_load)
                });
                let _ = stream
                    .try_collect::<Vec<_>>()
                    .await
                    .expect("no stream errors");
            };
            b.iter(|| runtime.block_on(async { test().await }));
        },
    );
    group.finish();
}

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

    let mut criterion = Criterion::default()
        .measurement_time(Duration::from_secs(450))
        .sample_size(10)
        .warm_up_time(Duration::from_secs(60));

    if let Some(baseline) = &args.save_baseline {
        criterion = criterion.save_baseline(baseline.to_string());
    }
    if let Some(baseline) = &args.use_baseline {
        criterion = criterion.retain_baseline(baseline.to_string());
    }

    for filter in args.filter.iter() {
        criterion = criterion.with_filter(filter.to_string())
    }

    let logger = app.logger();
    let ctx = CoreContext::new_with_logger(fb, logger.clone());

    let repo = app.runtime().block_on(async {
        app.open_repo::<BlobRepo>(&args.repo)
            .await
            .context("Failed to open repo")
    })?;
    let fetcher = PublicChangesetBulkFetch::new(repo.changesets_arc(), repo.phases_arc());
    let repo_name = repo.repo_identity().name().to_string();

    // Tests are run from here

    bench_stream(
        &mut criterion,
        &ctx,
        app.runtime(),
        format!(
            "{}{}",
            repo_name, ":PublicChangesetBulkFetch::fetch_best_newest_first_mid"
        ),
        &fetcher,
        |ctx, fetcher| {
            async move {
                let (lower, upper) = fetcher.get_repo_bounds(ctx).await?;
                let mid = (upper - lower) / 2;
                Ok(fetcher.fetch_ids(ctx, Direction::NewestFirst, Some((lower, mid))))
            }
            .try_flatten_stream()
        },
    );

    bench_stream(
        &mut criterion,
        &ctx,
        app.runtime(),
        format!(
            "{}{}",
            repo_name, ":PublicChangesetBulkFetch::fetch_best_oldest_first"
        ),
        &fetcher,
        |ctx, fetcher| fetcher.fetch_ids(ctx, Direction::OldestFirst, None),
    );

    bench_stream(
        &mut criterion,
        &ctx,
        app.runtime(),
        format!(
            "{}{}",
            repo_name, ":PublicChangesetBulkFetch::fetch_entries_oldest_first"
        ),
        &fetcher,
        |ctx, fetcher| fetcher.fetch(ctx, Direction::OldestFirst),
    );

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

//! This benchmark generates a single initial commit that adds 100k files to
//! a single large directory, and then 10 more commits that add, modify, and
//! remove some of those files at random.
//!
//! It then benchmarks deriving one of the derived data types (fsnodes,
//! unodes, skeleton manifest or deleted manifests) for those commits.

use std::collections::BTreeSet;

use anyhow::Result;
use blobrepo::BlobRepo;
use context::CoreContext;
use deleted_manifest::RootDeletedManifestIdCommon;
use deleted_manifest::RootDeletedManifestV2Id;
use derived_data::BonsaiDerived;
use derived_data_manager::BonsaiDerivable as NewBonsaiDerivable;
use fbinit::FacebookInit;
use fsnodes::RootFsnodeId;
use futures_stats::TimedFutureExt;
use mercurial_derivation::MappedHgChangesetId;
use mononoke_types::ChangesetId;
use rand::distributions::Alphanumeric;
use rand::distributions::Uniform;
use rand::thread_rng;
use rand::Rng;
use skeleton_manifest::RootSkeletonManifestId;
use tests_utils::CreateCommitContext;
use unodes::RootUnodeManifestId;

fn gen_filename(rng: &mut impl Rng, len: usize) -> String {
    std::iter::repeat_with(|| rng.sample(Alphanumeric))
        .take(len)
        .map(char::from)
        .collect()
}

async fn make_initial_large_directory(
    ctx: &CoreContext,
    repo: &BlobRepo,
    count: usize,
) -> Result<(ChangesetId, BTreeSet<String>)> {
    let mut filenames = BTreeSet::new();
    let mut rng = thread_rng();
    let len_distr = Uniform::new(5, 50);
    while filenames.len() < count {
        let len = rng.sample(len_distr);
        let filename = gen_filename(&mut rng, len);
        filenames.insert(filename);
    }

    let mut create = CreateCommitContext::new_root(ctx, repo);
    for filename in filenames.iter() {
        create = create.add_file(
            format!("large_directory/{}", filename).as_str(),
            format!("content of {}", filename),
        );
    }
    let csid = create.commit().await?;

    Ok((csid, filenames))
}

async fn modify_large_directory(
    ctx: &CoreContext,
    repo: &BlobRepo,
    filenames: &mut BTreeSet<String>,
    csid: ChangesetId,
    index: usize,
    add_count: usize,
    modify_count: usize,
    delete_count: usize,
) -> Result<ChangesetId> {
    let mut create = CreateCommitContext::new(ctx, repo, vec![csid]);
    let mut rng = thread_rng();
    let len_distr = Uniform::new(5, 50);

    let mut add_filenames = BTreeSet::new();
    while add_filenames.len() < add_count {
        let len = rng.sample(len_distr);
        let filename = gen_filename(&mut rng, len);
        if !filenames.contains(&filename) {
            add_filenames.insert(filename);
        }
    }

    let delete_count = delete_count.min(filenames.len());
    let modify_count = modify_count.min(filenames.len() - delete_count);
    let mut modify_filename_indexes = BTreeSet::new();
    let index_distr = Uniform::new(0, filenames.len());
    while modify_filename_indexes.len() < modify_count {
        let index = rng.sample(index_distr);
        modify_filename_indexes.insert(index);
    }
    let mut delete_filename_indexes = BTreeSet::new();
    while delete_filename_indexes.len() < delete_count {
        let index = rng.sample(index_distr);
        if !modify_filename_indexes.contains(&index) {
            delete_filename_indexes.insert(index);
        }
    }
    let mut modify_filenames = BTreeSet::new();
    let mut delete_filenames = BTreeSet::new();
    for (index, filename) in filenames.iter().enumerate() {
        if modify_filename_indexes.contains(&index) {
            modify_filenames.insert(filename);
        } else if delete_filename_indexes.contains(&index) {
            delete_filenames.insert(filename);
        }
    }

    for filename in add_filenames.iter().chain(modify_filenames) {
        create = create.add_file(
            format!("large_directory/{}", filename).as_str(),
            format!("content {} of {}", index, filename),
        );
    }
    for filename in delete_filenames.iter() {
        create = create.delete_file(format!("large_directory/{}", filename).as_str());
    }

    let csid = create.commit().await?;
    Ok(csid)
}

async fn derive(ctx: &CoreContext, repo: &BlobRepo, data: &str, csid: ChangesetId) -> String {
    match data {
        MappedHgChangesetId::NAME => MappedHgChangesetId::derive(ctx, repo, csid)
            .await
            .unwrap()
            .hg_changeset_id()
            .to_string(),
        RootSkeletonManifestId::NAME => RootSkeletonManifestId::derive(ctx, repo, csid)
            .await
            .unwrap()
            .skeleton_manifest_id()
            .to_string(),
        RootUnodeManifestId::NAME => RootUnodeManifestId::derive(ctx, repo, csid)
            .await
            .unwrap()
            .manifest_unode_id()
            .to_string(),
        RootDeletedManifestV2Id::NAME => RootDeletedManifestV2Id::derive(ctx, repo, csid)
            .await
            .unwrap()
            .id()
            .to_string(),
        RootFsnodeId::NAME => RootFsnodeId::derive(ctx, repo, csid)
            .await
            .unwrap()
            .fsnode_id()
            .to_string(),
        _ => panic!("invalid derived data type: {}", data),
    }
}

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

    let mut args = std::env::args();
    let _ = args.next();
    let data = args.next().unwrap_or_else(|| String::from("fsnodes"));
    println!("Deriving: {}", data);

    let repo: BlobRepo = test_repo_factory::build_empty(ctx.fb).await?;

    let (mut csid, mut filenames) = make_initial_large_directory(&ctx, &repo, 100_000).await?;

    println!("First commit: {}", csid);
    let (stats, derived_id) = derive(&ctx, &repo, &data, csid).timed().await;
    println!("Derived id: {}  stats: {:?}", derived_id, stats);

    let commit_count = 10;

    for commit in 0..commit_count {
        csid =
            modify_large_directory(&ctx, &repo, &mut filenames, csid, commit, 25, 100, 25).await?;
    }

    println!("Last commit: {}", csid);
    let (stats, derived_id) = derive(&ctx, &repo, &data, csid).timed().await;
    println!("Derived id: {}  stats: {:?}", derived_id, stats);

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

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

use std::fmt::Debug;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use anyhow::format_err;
use anyhow::Error;
use blobstore::Blobstore;
use blobstore_factory::make_sql_blobstore_xdb;
use blobstore_factory::BlobstoreOptions;
use blobstore_factory::ReadOnlyStorage;
use bytes::Bytes;
use bytes::BytesMut;
use cacheblob::new_memcache_blobstore;
use cached_config::ConfigStore;
use clap::Parser;
use clap::ValueEnum;
use context::CoreContext;
use fbinit::FacebookInit;
use filestore::FetchKey;
use filestore::FilestoreConfig;
use filestore::StoreRequest;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_stats::FutureStats;
use futures_stats::TimedFutureExt;
use mononoke_app::MononokeAppBuilder;
use mononoke_types::BlobstoreKey;
use mononoke_types::ContentMetadataV2;
use rand::Rng;
use throttledblob::ThrottledBlob;
use tokio::fs::File;
use tokio::io::BufReader;
use tokio_util::codec::BytesCodec;
use tokio_util::codec::FramedRead;

const NAME: &str = "benchmark_filestore";

/// Benchmark the filestore
#[derive(Parser)]
struct BenchmarkArgs {
    #[clap(long, default_value_t = 8192)]
    input_capacity: usize,

    #[clap(long, default_value_t = 1048576)]
    chunk_size: u64,

    #[clap(long, default_value_t = 1)]
    concurrency: usize,

    #[clap(long)]
    memcache: bool,

    #[clap(long)]
    cachelib_size: Option<usize>,

    #[clap(long)]
    delay: Option<u64>,

    #[clap(long)]
    randomize: bool,

    #[clap(long, default_value_t = 2)]
    read_count: usize,

    #[clap(long, required_if_eq("blobstore_type", "manifold"))]
    manifold_bucket: Option<String>,

    #[clap(long, required_if_eq("blobstore_type", "xdb"))]
    shardmap: Option<String>,

    #[clap(long)]
    shard_count: Option<NonZeroUsize>,

    /// Data file to use as input.
    input: PathBuf,

    /// Which type of blobstore to benchmark against.
    #[clap(value_enum)]
    blobstore_type: BlobstoreType,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum BlobstoreType {
    Manifold,
    Memory,
    Xdb,
}

fn log_perf<I, E: Debug>(stats: FutureStats, res: &Result<I, E>, len: u64) {
    match res {
        Ok(_) => {
            let bytes_per_ns = (len as f64) / (stats.completion_time.as_nanos() as f64);
            let mbytes_per_s = bytes_per_ns * (10_u128.pow(9) as f64) / (2_u128.pow(20) as f64);
            let gb_per_s = mbytes_per_s * 8_f64 / 1024_f64;
            eprintln!(
                "Success: {:.2} MB/s ({:.2} Gb/s) ({:?})",
                mbytes_per_s, gb_per_s, stats
            );
        }
        Err(e) => {
            eprintln!("Failure: {:?}", e);
        }
    };
}

async fn read<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    content_metadata: &ContentMetadataV2,
) -> Result<(), Error> {
    let key = FetchKey::Canonical(content_metadata.content_id);
    eprintln!(
        "Fetch start: {:?} ({:?} B)",
        key, content_metadata.total_size
    );

    let stream = filestore::fetch(blobstore, ctx.clone(), &key)
        .await?
        .ok_or_else(|| format_err!("Fetch failed: no stream"))?;

    let (stats, res) = stream.try_for_each(|_| async { Ok(()) }).timed().await;
    log_perf(stats, &res, content_metadata.total_size);

    // ignore errors - all we do is log them in `log_perf`
    match res {
        Ok(_) => Ok(()),
        Err(_) => Ok(()),
    }
}

async fn run_benchmark_filestore<'a>(
    ctx: &'a CoreContext,
    args: &BenchmarkArgs,
    blobstore: Arc<dyn Blobstore>,
) -> Result<(), Error> {
    let config = FilestoreConfig {
        chunk_size: Some(args.chunk_size),
        concurrency: args.concurrency,
    };

    eprintln!("Test with {:?}, writing into {:?}", config, blobstore);

    let file = File::open(&args.input).await?;
    let metadata = file.metadata().await?;

    let data = BufReader::with_capacity(args.input_capacity, file);
    let data = FramedRead::new(data, BytesCodec::new()).map_ok(BytesMut::freeze);
    let len = metadata.len();

    let (len, data) = if args.randomize {
        let bytes = rand::thread_rng().gen::<[u8; 32]>();
        let bytes = Bytes::copy_from_slice(&bytes[..]);
        (
            len + (bytes.len() as u64),
            stream::iter(vec![Ok(bytes)]).chain(data).left_stream(),
        )
    } else {
        (len, data.right_stream())
    };

    eprintln!("Write start: {:?} B", len);

    let req = StoreRequest::new(len);

    let (stats, res) = filestore::store(&blobstore, config, ctx, &req, data.map_err(Error::from))
        .timed()
        .await;
    log_perf(stats, &res, len);

    let metadata = res?;

    match args.delay {
        Some(delay) => {
            tokio_shim::time::sleep(Duration::from_secs(delay)).await;
        }
        None => {}
    }

    eprintln!("Write committed: {:?}", metadata.content_id.blobstore_key());

    for _c in 0..args.read_count {
        read(&blobstore, ctx, &metadata).await?;
    }

    Ok(())
}

#[cfg(fbcode_build)]
const TEST_DATA_TTL: Option<Duration> = Some(Duration::from_secs(3600));

async fn open_blobstore(
    fb: FacebookInit,
    args: &BenchmarkArgs,
    blobstore_options: &BlobstoreOptions,
    readonly_storage: &ReadOnlyStorage,
    config_store: &ConfigStore,
) -> Result<Arc<dyn Blobstore>, Error> {
    let blobstore: Arc<dyn Blobstore> = match args.blobstore_type {
        BlobstoreType::Manifold => {
            #[cfg(fbcode_build)]
            {
                use manifoldblob::ManifoldBlob;
                use prefixblob::PrefixBlobstore;

                let bucket = args
                    .manifold_bucket
                    .as_ref()
                    .expect("Manifold bucket must be set when using manifold blobstore type");
                let put_behaviour = blobstore_options.put_behaviour;
                let manifold = ManifoldBlob::new(
                    fb,
                    bucket,
                    TEST_DATA_TTL,
                    blobstore_options.manifold_options.clone(),
                    put_behaviour,
                )?;
                let blobstore = PrefixBlobstore::new(manifold, format!("{}.", NAME));
                Arc::new(blobstore)
            }
            #[cfg(not(fbcode_build))]
            {
                unimplemented!("Accessing Manifold is not implemented in non fbcode builds");
            }
        }
        BlobstoreType::Memory => Arc::new(memblob::Memblob::default()),
        BlobstoreType::Xdb => {
            let tier_name = args
                .shardmap
                .as_ref()
                .expect("Shardmap must be set when using xdb blobstore type")
                .to_string();
            let blobstore = make_sql_blobstore_xdb(
                fb,
                tier_name,
                args.shard_count,
                blobstore_options,
                *readonly_storage,
                blobstore_options.put_behaviour,
                config_store,
            )
            .await?;
            Arc::new(blobstore)
        }
    };

    let blobstore: Arc<dyn Blobstore> = if args.memcache {
        Arc::new(new_memcache_blobstore(fb, blobstore, NAME, "")?)
    } else {
        blobstore
    };

    let blobstore: Arc<dyn Blobstore> = match args.cachelib_size {
        Some(size) => {
            #[cfg(fbcode_build)]
            {
                cachelib::init_cache(fb, cachelib::LruCacheConfig::new(size))?;

                let presence_pool = cachelib::get_or_create_pool(
                    "presence",
                    cachelib::get_available_space()? / 20,
                )?;
                let blob_pool =
                    cachelib::get_or_create_pool("blobs", cachelib::get_available_space()?)?;

                Arc::new(cacheblob::new_cachelib_blobstore_no_lease(
                    blobstore,
                    Arc::new(blob_pool),
                    Arc::new(presence_pool),
                    blobstore_options.cachelib_options,
                ))
            }
            #[cfg(not(fbcode_build))]
            {
                let _ = size;
                unimplemented!("Using cachelib is not implemented for non fbcode build");
            }
        }
        None => blobstore,
    };

    // ThrottledBlob is a noop if no throttling requested
    let blobstore =
        Arc::new(ThrottledBlob::new(blobstore, blobstore_options.throttle_options).await);

    Ok(blobstore)
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let app = MononokeAppBuilder::new(fb).build::<BenchmarkArgs>()?;
    let args: BenchmarkArgs = app.args()?;

    let runtime = app.runtime();
    let logger = app.logger();
    let config_store = app.config_store();
    let ctx = CoreContext::new_with_logger(fb, logger.clone());

    let blobstore = runtime.block_on(open_blobstore(
        fb,
        &args,
        app.blobstore_options(),
        app.readonly_storage(),
        config_store,
    ))?;

    runtime.block_on(run_benchmark_filestore(&ctx, &args, blobstore))?;

    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 anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use blobstore::Blobstore;
use blobstore_factory::make_blobstore;
use cacheblob::new_memcache_blobstore;
use cacheblob::CachelibBlobstoreOptions;
use clap::Parser;
use context::CoreContext;
use criterion::Criterion;
use environment::Caching;
use mononoke_app::MononokeAppBuilder;
use tokio::runtime::Handle;

mod parallel_puts;
mod single_puts;

mod parallel_different_blob_gets;
mod parallel_same_blob_gets;
mod single_gets;

pub const KB: usize = 1024;
pub const MB: usize = KB * 1024;

/// Benchmark storage config.
#[derive(Parser)]
struct BenchmarkArgs {
    /// The name of the storage config to benchmark.
    #[clap(long)]
    storage_config_name: String,

    /// Save results as a baseline under this name, for comparison.
    #[clap(long)]
    save_baseline: Option<String>,

    /// Compare to named baseline instead of last run.
    #[clap(long)]
    use_baseline: Option<String>,

    /// Limit to benchmarks whose name contains this string. Repetition tightens the filter.
    #[clap(long)]
    filter: Vec<String>,
}

#[fbinit::main]
fn main(fb: fbinit::FacebookInit) -> Result<(), Error> {
    let app = MononokeAppBuilder::new(fb).build::<BenchmarkArgs>()?;
    let args: BenchmarkArgs = app.args()?;

    let mut criterion = Criterion::default()
        .measurement_time(Duration::from_secs(60))
        .sample_size(10)
        .warm_up_time(Duration::from_secs(60));

    if let Some(baseline) = &args.save_baseline {
        criterion = criterion.save_baseline(baseline.to_string());
    }
    if let Some(baseline) = &args.use_baseline {
        criterion = criterion.retain_baseline(baseline.to_string());
    }

    for filter in args.filter.iter() {
        criterion = criterion.with_filter(filter.to_string())
    }

    let caching = app.environment().caching;
    let logger = app.logger();
    let config_store = app.config_store();

    let storage_configs = app.storage_configs();
    let storage_config = storage_configs
        .storage
        .get(&args.storage_config_name)
        .ok_or_else(|| anyhow!("unknown storage config"))?;
    let mysql_options = app.mysql_options();
    let blobstore_options = app.blobstore_options();
    let ctx = CoreContext::new_with_logger(fb, logger.clone());

    let blobstore = || async {
        let blobstore = make_blobstore(
            fb,
            storage_config.blobstore.clone(),
            mysql_options,
            blobstore_factory::ReadOnlyStorage(false),
            blobstore_options,
            logger,
            config_store,
            &blobstore_factory::default_scrub_handler(),
            None,
        )
        .await
        .context("Could not make blobstore")?;
        let blobstore = match caching {
            Caching::Disabled => blobstore,
            Caching::LocalOnly(local_cache_config) => repo_factory::cachelib_blobstore(
                blobstore,
                local_cache_config.blobstore_cache_shards,
                &CachelibBlobstoreOptions::default(),
            )
            .context("repo_factory::cachelib_blobstore failed")?,
            Caching::Enabled(local_cache_config) => {
                let cachelib_blobstore = repo_factory::cachelib_blobstore(
                    blobstore,
                    local_cache_config.blobstore_cache_shards,
                    &CachelibBlobstoreOptions::default(),
                )
                .context("repo_factory::cachelib_blobstore failed")?;
                Arc::new(
                    new_memcache_blobstore(fb, cachelib_blobstore, "benchmark", "")
                        .context("Memcache blobstore issues")?,
                )
            }
        };
        Ok::<_, Error>(blobstore)
    };

    // Cut all the repetition around running a benchmark. Note that a fresh cachee shouldn't be needed,
    // as the warmup will fill it, and the key is randomised
    let mut run_benchmark = {
        let criterion = &mut criterion;
        let runtime = app.runtime();
        move |bench: fn(&mut Criterion, CoreContext, Arc<dyn Blobstore>, &Handle)| -> Result<(), Error> {
            let blobstore = runtime.block_on(blobstore())?;
            bench(criterion, ctx.clone(), blobstore, runtime);
            Ok(())
        }
    };

    // Tests are run from here
    run_benchmark(single_puts::benchmark)?;
    run_benchmark(single_gets::benchmark)?;
    run_benchmark(parallel_same_blob_gets::benchmark)?;
    run_benchmark(parallel_different_blob_gets::benchmark)?;
    run_benchmark(parallel_puts::benchmark)?;

    criterion.final_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::iter::repeat;
use std::sync::Arc;

use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use tokio::runtime::Handle;

use crate::KB;
use crate::MB;

pub fn benchmark(
    c: &mut Criterion,
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    runtime: &Handle,
) {
    let mut group = c.benchmark_group("parallel_different_blob_gets");

    for size in [128, 16 * KB, 512 * KB, 8 * MB] {
        for concurrency in [4, 16, 256] {
            group.throughput(Throughput::Bytes(size as u64 * concurrency as u64));
            group.bench_with_input(
                BenchmarkId::from_parameter(format!("{} x{}", size, concurrency)),
                &size,
                |b, &size| {
                    let keys: Vec<_> = repeat(())
                        .take(concurrency)
                        .map(|()| {
                            let mut block = vec![0; size];
                            thread_rng().fill(&mut block as &mut [u8]);

                            let block = BlobstoreBytes::from_bytes(block);
                            let key = format!("benchmark.{:x}", thread_rng().next_u64());
                            runtime.block_on(async {
                                blobstore
                                    .put(&ctx, key.clone(), block)
                                    .await
                                    .expect("Put failed")
                            });
                            key
                        })
                        .collect();
                    let test = |ctx: CoreContext, blobstore: Arc<dyn Blobstore>| {
                        let keys = keys.clone();
                        async move {
                            let futs: FuturesUnordered<_> = keys
                                .into_iter()
                                .map(|key| {
                                    let ctx = &ctx;
                                    let blobstore = &blobstore;
                                    async move { blobstore.get(ctx, &key).await }
                                })
                                .collect();
                            futs.try_for_each(|_| async move { Ok(()) })
                                .await
                                .expect("Gets failed");
                        }
                    };
                    b.iter(|| {
                        runtime.block_on(async { test(ctx.clone(), Arc::clone(&blobstore)).await })
                    });
                },
            );
        }
    }

    group.finish();
}


/*
 * 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 blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use tokio::runtime::Handle;

use crate::KB;
use crate::MB;

pub fn benchmark(
    c: &mut Criterion,
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    runtime: &Handle,
) {
    let mut group = c.benchmark_group("parallel_puts");

    for size in [128, 16 * KB, 512 * KB, 8 * MB] {
        for concurrency in [4, 16, 256] {
            group.throughput(Throughput::Bytes(size as u64 * concurrency as u64));
            group.bench_with_input(
                BenchmarkId::from_parameter(format!("{} x{}", size, concurrency)),
                &size,
                |b, &size| {
                    let blocks: Vec<_> = std::iter::repeat(())
                        .take(concurrency)
                        .map(|()| {
                            let mut block = vec![0; size];
                            thread_rng().fill(&mut block as &mut [u8]);

                            BlobstoreBytes::from_bytes(block)
                        })
                        .collect();
                    let test = |ctx: CoreContext, blobstore: Arc<dyn Blobstore>| {
                        let blocks = blocks.clone();
                        async move {
                            let futs: FuturesUnordered<_> = blocks
                                .into_iter()
                                .map(|block| {
                                    let key = format!("benchmark.{:x}", thread_rng().next_u64());
                                    blobstore.put(&ctx, key, block)
                                })
                                .collect();

                            futs.try_for_each(|_| async move { Ok(()) })
                                .await
                                .expect("Puts failed");
                        }
                    };
                    b.iter(|| {
                        runtime.block_on(async { test(ctx.clone(), Arc::clone(&blobstore)).await })
                    });
                },
            );
        }
    }
    group.finish();
}


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

use std::iter::repeat;
use std::sync::Arc;

use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use tokio::runtime::Handle;

use crate::KB;
use crate::MB;

pub fn benchmark(
    c: &mut Criterion,
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    runtime: &Handle,
) {
    let mut group = c.benchmark_group("parallel_same_blob_gets");

    for size in [128, 16 * KB, 512 * KB, 8 * MB] {
        for concurrency in [4, 16, 256] {
            group.throughput(Throughput::Bytes(size as u64 * concurrency as u64));
            group.bench_with_input(
                BenchmarkId::from_parameter(format!("{} x{}", size, concurrency)),
                &size,
                |b, &size| {
                    let mut block = vec![0; size];
                    thread_rng().fill(&mut block as &mut [u8]);

                    let block = BlobstoreBytes::from_bytes(block);
                    let key = format!("benchmark.{:x}", thread_rng().next_u64());
                    runtime.block_on(async {
                        blobstore
                            .put(&ctx, key.clone(), block)
                            .await
                            .expect("Put failed")
                    });
                    let keys = repeat(key).take(concurrency);
                    let test = |ctx: CoreContext, blobstore: Arc<dyn Blobstore>| {
                        let keys = keys.clone();
                        async move {
                            let futs: FuturesUnordered<_> = keys
                                .map(|key| {
                                    let ctx = &ctx;
                                    let blobstore = &blobstore;
                                    async move { blobstore.get(ctx, &key).await }
                                })
                                .collect();
                            futs.try_for_each(|_| async move { Ok(()) })
                                .await
                                .expect("Gets failed");
                        }
                    };
                    b.iter(|| {
                        runtime.block_on(async { test(ctx.clone(), Arc::clone(&blobstore)).await })
                    });
                },
            );
        }
    }

    group.finish();
}


/*
 * 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 blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use tokio::runtime::Handle;

use crate::KB;
use crate::MB;

pub fn benchmark(
    c: &mut Criterion,
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    runtime: &Handle,
) {
    let mut group = c.benchmark_group("single_gets");

    for size in [128, 16 * KB, 512 * KB, 8 * MB] {
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            let mut block = vec![0; size];
            thread_rng().fill(&mut block as &mut [u8]);

            let block = BlobstoreBytes::from_bytes(block);
            let key = format!("benchmark.{:x}", thread_rng().next_u64());
            runtime.block_on(async {
                blobstore
                    .put(&ctx, key.clone(), block)
                    .await
                    .expect("Put failed")
            });
            let test = |ctx, blobstore: Arc<dyn Blobstore>| {
                let key = key.clone();
                async move {
                    blobstore.get(&ctx, &key).await.expect("Get failed");
                }
            };
            b.iter(|| runtime.block_on(async { test(ctx.clone(), Arc::clone(&blobstore)).await }));
        });
    }
    group.finish();
}


/*
 * 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 blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use context::CoreContext;
use criterion::BenchmarkId;
use criterion::Criterion;
use criterion::Throughput;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use tokio::runtime::Handle;

use crate::KB;
use crate::MB;

pub fn benchmark(
    c: &mut Criterion,
    ctx: CoreContext,
    blobstore: Arc<dyn Blobstore>,
    runtime: &Handle,
) {
    let mut group = c.benchmark_group("single_puts");

    for size in [128, 16 * KB, 512 * KB, 8 * MB] {
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            let mut block = vec![0; size];
            thread_rng().fill(&mut block as &mut [u8]);

            let block = BlobstoreBytes::from_bytes(block);
            let test = |ctx, blobstore: Arc<dyn Blobstore>| {
                let block = block.clone();
                async move {
                    let key = format!("benchmark.{:x}", thread_rng().next_u64());
                    blobstore.put(&ctx, key, block).await.expect("Put failed");
                }
            };
            b.iter(|| runtime.block_on(async { test(ctx.clone(), Arc::clone(&blobstore)).await }));
        });
    }
    group.finish();
}


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

pub const HIGHEST_IMPORTED_GEN_NUM: &str = "highest-imported-gen-num";


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

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

use anyhow::format_err;
use anyhow::Error;
use ascii::AsciiString;
use blobrepo::BlobRepo;
use blobrepo_hg::BlobRepoHg;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkUpdateReason;
use bookmarks::BookmarksRef;
use cloned::cloned;
use context::CoreContext;
use futures::FutureExt;
use futures::TryFutureExt;
use futures_01_ext::try_boxfuture;
use futures_01_ext::BoxFuture;
use futures_01_ext::FutureExt as _;
use futures_old::prelude::*;
use futures_old::stream;
use mercurial_revlog::RevlogRepo;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use slog::info;
use slog::Logger;

pub fn read_bookmarks(revlogrepo: &RevlogRepo) -> BoxFuture<Vec<(Vec<u8>, HgChangesetId)>, Error> {
    let bookmarks = Arc::new(try_boxfuture!(revlogrepo.get_bookmarks()));

    (*bookmarks)
        .keys()
        .and_then({
            let bookmarks = bookmarks.clone();
            move |key| {
                (*bookmarks).get(&key).and_then(move |cs_id| {
                    cs_id
                        .ok_or_else(|| {
                            format_err!("Bookmark value missing: {}", String::from_utf8_lossy(&key))
                        })
                        .map(move |cs_id| (key, cs_id))
                })
            }
        })
        .collect()
        .boxify()
}

pub type BookmarkKeyTransformer = Box<dyn Fn(BookmarkKey) -> BookmarkKey + Send + Sync + 'static>;

pub fn get_bookmark_prefixer(prefix: AsciiString) -> BookmarkKeyTransformer {
    Box::new(move |bookmark| {
        let mut name = prefix.clone();
        name.push_str(bookmark.as_ascii());
        BookmarkKey::new_ascii(name)
    })
}

pub fn upload_bookmarks(
    ctx: CoreContext,
    logger: &Logger,
    revlogrepo: RevlogRepo,
    blobrepo: BlobRepo,
    stale_bookmarks: Vec<(Vec<u8>, HgChangesetId)>,
    mononoke_bookmarks: Vec<(BookmarkKey, ChangesetId)>,
    bookmark_name_transformer: BookmarkKeyTransformer,
) -> BoxFuture<(), Error> {
    let logger = logger.clone();
    let stale_bookmarks = Arc::new(stale_bookmarks.into_iter().collect::<HashMap<_, _>>());

    read_bookmarks(&revlogrepo)
        .map({
            cloned!(ctx, logger, blobrepo, stale_bookmarks);
            move |bookmarks| {
                stream::futures_unordered(bookmarks.into_iter().map(|(key, cs_id)| {
                    {
                        cloned!(ctx, blobrepo);
                        async move { blobrepo.hg_changeset_exists(ctx, cs_id).await }
                    }
                    .boxed()
                    .compat()
                    .and_then({
                            cloned!(ctx, logger, key, blobrepo, stale_bookmarks);
                            move |exists| {
                                match (exists, stale_bookmarks.get(&key).cloned()) {
                                    (false, Some(stale_cs_id)) => {
                                        info!(
                                            logger,
                                            "current version of bookmark {} couldn't be \
                                            imported, because cs {:?} was not present in blobrepo \
                                            yet; using stale version instead {:?}",
                                            String::from_utf8_lossy(&key),
                                            cs_id,
                                            stale_cs_id,
                                        );

                                        async move {
                                            let exists = blobrepo.hg_changeset_exists(ctx, stale_cs_id).await?;
                                            Ok((key, stale_cs_id, exists))
                                        }
                                        .boxed()
                                        .compat()
                                        .boxify()
                                    }
                                    _ => Ok((key, cs_id, exists)).into_future().boxify(),
                                }
                            }})
                        .and_then({
                            cloned!(ctx, blobrepo, logger);
                            move |(key, cs_id, exists)| async move {
                                if exists {
                                    let bcs_id = blobrepo.bonsai_hg_mapping().get_bonsai_from_hg(&ctx, cs_id)
                                        .await?
                                        .ok_or_else(|| format_err!("failed to resolve hg to bonsai: {}", cs_id))?;
                                    Ok(Some((key, bcs_id)))
                                } else {
                                    info!(
                                        logger,
                                        "did not update bookmark {}, because cs {:?} was not imported yet",
                                        String::from_utf8_lossy(&key),
                                        cs_id,
                                    );
                                    Ok(None)
                                }
                            }
                            .boxed()
                            .compat()
                        })
                }))
            }
        })
        .flatten_stream()
        .filter_map(|key_cs_id| key_cs_id)
        .chunks(100) // send 100 bookmarks in a single transaction
        .and_then({
            let blobrepo = blobrepo;
            let mononoke_bookmarks: HashMap<_, _> = mononoke_bookmarks.into_iter().collect();
            move |vec| {
                let mut transaction = blobrepo.bookmarks().create_transaction(ctx.clone());

                let mut count = 0;
                for (key, value) in vec {
                    let bookmark_name = BookmarkKey::new_ascii(try_boxfuture!(AsciiString::from_ascii(key)));
                    let bookmark_name = bookmark_name_transformer(bookmark_name);
                    if mononoke_bookmarks.get(&bookmark_name) != Some(&value) {
                        count += 1;
                        try_boxfuture!(transaction.force_set(&bookmark_name, value, BookmarkUpdateReason::Blobimport))
                    }
                }

                if count > 0 {
                    transaction.commit()
                        .compat()
                        .and_then(move |ok| {
                            if ok {
                                Ok(count)
                            } else {
                                Err(format_err!("Bookmark transaction failed"))
                            }
                        })
                        .boxify()
                } else {
                    Ok(0).into_future().boxify()
                }
            }
        }).for_each(move |count| {
            info!(logger, "uploaded chunk of {:?} bookmarks", count);
            Ok(())
        }).boxify()
}