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

use anyhow::bail;
use anyhow::Error;
use cloned::cloned;
use fbinit::FacebookInit;
use metaconfig_types::LocalDatabaseConfig;
use metaconfig_types::MetadataDatabaseConfig;
use metaconfig_types::RemoteMetadataDatabaseConfig;
use metaconfig_types::ShardableRemoteDatabaseConfig;
use sql::sqlite::SqliteMultithreaded;
use sql::sqlite::SqliteQueryType;
use sql::Connection;
use sql::SqlConnections;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_construct::SqlShardableConstructFromMetadataDatabaseConfig;
use sql_ext::facebook::MysqlOptions;
use sql_ext::open_existing_sqlite_path;
use sql_ext::open_sqlite_path;

use crate::ReadOnlyStorage;

/// Create instances of SQL database managers for repository metadata based on metadata database
/// config.
#[derive(Clone)]
pub struct MetadataSqlFactory {
    fb: FacebookInit,
    readonly: ReadOnlyStorage,
    connections_factory: MetadataSqlConnectionsFactory,
}

#[derive(Clone)]
pub enum MetadataSqlConnectionsFactory {
    Local {
        connections: SqlConnections,
        schema_connection: SqliteMultithreaded,
    },
    Remote {
        remote_config: RemoteMetadataDatabaseConfig,
        mysql_options: MysqlOptions,
    },
}

#[derive(Clone)]
pub struct SqlTierInfo {
    pub tier_name: String,
    /// Returns None for unsharded, Some(number of shards) otherwise.
    /// NB does not tell you if shards are 0 or 1 based, just the overall number
    pub shard_num: Option<usize>,
}

impl MetadataSqlFactory {
    pub async fn new(
        fb: FacebookInit,
        dbconfig: MetadataDatabaseConfig,
        mysql_options: MysqlOptions,
        readonly: ReadOnlyStorage,
    ) -> Result<Self, Error> {
        let connections_factory = match dbconfig {
            MetadataDatabaseConfig::Local(LocalDatabaseConfig { path }) => {
                let path = path.join("sqlite_dbs");
                let schema_connection =
                    SqliteMultithreaded::new(open_sqlite_path(path.clone(), false)?);
                let read_connection =
                    Connection::with_sqlite(open_existing_sqlite_path(path, true)?);
                let connections = SqlConnections {
                    write_connection: if readonly.0 {
                        read_connection.clone()
                    } else {
                        schema_connection.clone().into()
                    },
                    read_master_connection: read_connection.clone(),
                    read_connection,
                };
                MetadataSqlConnectionsFactory::Local {
                    connections,
                    schema_connection,
                }
            }
            MetadataDatabaseConfig::Remote(remote_config) => {
                MetadataSqlConnectionsFactory::Remote {
                    remote_config,
                    mysql_options,
                }
            }
        };
        Ok(Self {
            fb,
            readonly,
            connections_factory,
        })
    }

    pub async fn open<T: SqlConstructFromMetadataDatabaseConfig>(&self) -> Result<T, Error> {
        match &self.connections_factory {
            MetadataSqlConnectionsFactory::Local {
                connections,
                schema_connection,
            } => {
                schema_connection
                    .acquire_sqlite_connection(SqliteQueryType::SchemaChange)
                    .await?
                    .execute_batch(T::CREATION_QUERY)?;
                Ok(T::from_sql_connections(connections.clone()))
            }
            MetadataSqlConnectionsFactory::Remote {
                remote_config,
                mysql_options,
            } => T::with_remote_metadata_database_config(
                self.fb,
                remote_config,
                mysql_options,
                self.readonly.0,
            ),
        }
    }

    pub async fn open_shardable<T: SqlShardableConstructFromMetadataDatabaseConfig>(
        &self,
    ) -> Result<T, Error> {
        match &self.connections_factory {
            MetadataSqlConnectionsFactory::Local {
                connections,
                schema_connection,
            } => {
                schema_connection
                    .acquire_sqlite_connection(SqliteQueryType::SchemaChange)
                    .await?
                    .execute_batch(<T as SqlConstruct>::CREATION_QUERY)?;
                Ok(T::from_sql_connections(connections.clone()))
            }
            MetadataSqlConnectionsFactory::Remote {
                remote_config,
                mysql_options,
            } => {
                // Large numbers of shards means this can take some time.
                // Spawn it as a blocking task.
                tokio::task::spawn_blocking({
                    cloned!(self.fb, remote_config, mysql_options, self.readonly);
                    move || {
                        T::with_remote_metadata_database_config(
                            fb,
                            &remote_config,
                            &mysql_options,
                            readonly.0,
                        )
                    }
                })
                .await?
            }
        }
    }

    pub fn tier_info_shardable<T: SqlShardableConstructFromMetadataDatabaseConfig>(
        &self,
    ) -> Result<SqlTierInfo, Error> {
        Ok(match &self.connections_factory {
            MetadataSqlConnectionsFactory::Local { .. } => SqlTierInfo {
                tier_name: "sqlite".to_string(),
                shard_num: None,
            },
            MetadataSqlConnectionsFactory::Remote { remote_config, .. } => {
                match T::remote_database_config(remote_config) {
                    Some(ShardableRemoteDatabaseConfig::Unsharded(config)) => SqlTierInfo {
                        tier_name: config.db_address.clone(),
                        shard_num: None,
                    },
                    Some(ShardableRemoteDatabaseConfig::Sharded(config)) => SqlTierInfo {
                        tier_name: config.shard_map.clone(),
                        shard_num: Some(config.shard_num.get()),
                    },
                    None => bail!(
                        "missing tier name in configuration for {}",
                        <T as SqlConstruct>::LABEL
                    ),
                }
            }
        })
    }
}


/*
 * 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 std::fs::create_dir_all;
use std::ops::RangeBounds;
use std::path::Path;
use std::path::PathBuf;
use std::time::SystemTime;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreEnumerationData;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreKeyParam;
use blobstore::BlobstoreKeySource;
use blobstore::BlobstoreMetadata;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use percent_encoding::percent_encode;
use percent_encoding::AsciiSet;
use percent_encoding::CONTROLS;
use tempfile::NamedTempFile;
use tempfile::PersistError;
use tokio::fs::hard_link;
use tokio::fs::remove_file;
use tokio::fs::File;
use tokio::io;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use walkdir::WalkDir;

const PREFIX: &str = "blob";
const PREFIX_HYPHEN: &str = "blob-";
// https://url.spec.whatwg.org/#fragment-percent-encode-set
const FRAGMENT: &AsciiSet = &CONTROLS.add(b' ').add(b'"').add(b'<').add(b'>').add(b'`');
// https://url.spec.whatwg.org/#path-percent-encode-set
const PATH: &AsciiSet = &FRAGMENT.add(b'#').add(b'?').add(b'{').add(b'}');

#[derive(Debug, Clone)]
pub struct Fileblob {
    base: PathBuf,
    put_behaviour: PutBehaviour,
}

impl Fileblob {
    pub fn open<P: AsRef<Path>>(base: P, put_behaviour: PutBehaviour) -> Result<Self> {
        let base = base.as_ref();

        if !base.is_dir() {
            bail!("Base {:?} doesn't exist or is not directory", base);
        }

        Ok(Self {
            base: base.to_owned(),
            put_behaviour,
        })
    }

    pub fn create<P: AsRef<Path>>(base: P, put_behaviour: PutBehaviour) -> Result<Self> {
        let base = base.as_ref();
        create_dir_all(base)?;
        Self::open(base, put_behaviour)
    }

    fn path(&self, key: &str) -> PathBuf {
        let key = percent_encode(key.as_bytes(), PATH);
        self.base.join(format!("{}-{}", PREFIX, key))
    }

    /// Stripping the prepended prefix (if its exists) before returning
    /// keys back to the caller. Safe to call with or without the prefix.
    fn strip_file_prefix<'a>(&self, key: &'a str) -> &'a str {
        match key.strip_prefix(PREFIX_HYPHEN) {
            // Remove the prefix if present and return the remaining
            // slice.
            Some(key_without_prefix) => key_without_prefix,
            // If not present, return the key as-is.
            None => key,
        }
    }
}

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

async fn ctime(file: &File) -> Option<i64> {
    let meta = file.metadata().await.ok()?;
    let ctime = meta.modified().ok()?;
    let ctime_dur = ctime.duration_since(SystemTime::UNIX_EPOCH).ok()?;
    i64::try_from(ctime_dur.as_secs()).ok()
}

#[async_trait]
impl BlobstorePutOps for Fileblob {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        let p = self.path(&key);
        // block_in_place on tempfile would be ideal here, but it interacts
        // badly with tokio_compat
        let tempfile = NamedTempFile::new_in(&self.base)?;
        let new_file = tempfile.as_file().try_clone()?;
        let mut tokio_file = File::from_std(new_file);
        tokio_file.write_all(value.as_bytes().as_ref()).await?;
        tokio_file.flush().await?;
        tokio_file.sync_all().await?;
        let status = match put_behaviour {
            PutBehaviour::Overwrite => {
                tempfile.persist(&p)?;
                OverwriteStatus::NotChecked
            }
            PutBehaviour::IfAbsent | PutBehaviour::OverwriteAndLog => {
                let temp_path = tempfile.path().to_owned();
                match tempfile.persist_noclobber(&p) {
                    Ok(_) => OverwriteStatus::New,
                    // Key already existed
                    Err(PersistError { file: f, error: e })
                        if f.path() == temp_path
                            && e.kind() == std::io::ErrorKind::AlreadyExists =>
                    {
                        if put_behaviour.should_overwrite() {
                            f.persist(&p)?;
                            OverwriteStatus::Overwrote
                        } else {
                            OverwriteStatus::Prevented
                        }
                    }
                    Err(e) => return Err(e.into()),
                }
            }
        };

        Ok(status)
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.put_explicit(ctx, key, value, self.put_behaviour).await
    }
}

#[async_trait]
impl Blobstore for Fileblob {
    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let p = self.path(key);

        let ret = match File::open(&p).await {
            Err(ref r) if r.kind() == io::ErrorKind::NotFound => None,
            Err(e) => return Err(e.into()),
            Ok(mut f) => {
                let mut v = Vec::new();
                f.read_to_end(&mut v).await?;

                Some(BlobstoreGetData::new(
                    BlobstoreMetadata::new(ctime(&f).await, None),
                    BlobstoreBytes::from_bytes(v),
                ))
            }
        };
        Ok(ret)
    }

    async fn is_present<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let p = self.path(key);

        let present = match File::open(&p).await {
            Err(ref e) if e.kind() == io::ErrorKind::NotFound => false,
            Err(e) => return Err(e.into()),
            Ok(_) => true,
        };
        Ok(if present {
            BlobstoreIsPresent::Present
        } else {
            BlobstoreIsPresent::Absent
        })
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        BlobstorePutOps::put_with_status(self, ctx, key, value).await?;
        Ok(())
    }

    // This uses hardlink semantics as the production blobstores also have hardlink like semantics
    // (i.e. you can't discover a canonical link source when loading by the target)
    async fn copy<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        // from std::fs::hard_link: The dst path will be a link pointing to the src path
        let src_path = self.path(old_key);
        let dst_path = self.path(&new_key);
        // hard_link will fail if dst_path exists. Race it in a task of its own
        Ok(tokio::task::spawn(async move {
            let _ = remove_file(&dst_path).await;
            hard_link(src_path, dst_path).await
        })
        .await??)
    }
}

#[async_trait]
impl BlobstoreUnlinkOps for Fileblob {
    async fn unlink<'a>(&'a self, _ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        let path = self.path(key);
        Ok(remove_file(path).await?)
    }
}

#[async_trait]
impl BlobstoreKeySource for Fileblob {
    async fn enumerate<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData> {
        match range {
            BlobstoreKeyParam::Start(ref range) => {
                let mut enum_data = BlobstoreEnumerationData {
                    keys: HashSet::new(),
                    next_token: None,
                };
                WalkDir::new(&self.base)
                    .into_iter()
                    .filter_map(|v| v.ok())
                    .for_each(|entry| {
                        // Need the filename not the directory, since the directory
                        // structure is not exposed to the caller.
                        let entry = entry.file_name().to_str();
                        if let Some(data) = entry {
                            let key = self.strip_file_prefix(data).to_string();
                            if range.contains(&key) {
                                enum_data.keys.insert(key);
                            }
                        }
                    });
                Ok(enum_data)
            }
            _ => Err(format_err!("Fileblob does not support token, only ranges")),
        }
    }
}

#[cfg(test)]
mod test {
    use fbinit::FacebookInit;

    use super::*;

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

        let blob = Fileblob {
            base: PathBuf::from("/mononoke/fileblob/test/path/should/not/exist"),
            put_behaviour: PutBehaviour::IfAbsent,
        };

        let ret = blob
            .put(&ctx, "key".into(), BlobstoreBytes::from_bytes("value"))
            .await;

        assert!(ret.is_err());

        Ok(())
    }
}


// @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"),
        "memcache_lock 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=memcache-lock-thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

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


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


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

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_stats::record_get_stats;
use blobstore_stats::record_is_present_stats;
use blobstore_stats::record_put_stats;
use blobstore_stats::OperationType;
use context::CoreContext;
use context::PerfCounterType;
use futures_stats::TimedFutureExt;
use mononoke_types::BlobstoreBytes;
use scuba_ext::MononokeScubaSampleBuilder;

#[derive(Debug)]
pub struct LogBlob<B> {
    inner: B,
    scuba: MononokeScubaSampleBuilder,
    scuba_sample_rate: NonZeroU64,
}

impl<B: std::fmt::Debug> LogBlob<B> {
    pub fn new(
        inner: B,
        mut scuba: MononokeScubaSampleBuilder,
        scuba_sample_rate: NonZeroU64,
    ) -> Self {
        scuba.add_common_server_data();
        Self {
            inner,
            scuba,
            scuba_sample_rate,
        }
    }
}

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

#[async_trait]
impl<B: Blobstore + BlobstorePutOps> Blobstore for LogBlob<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let mut ctx = ctx.clone();
        let mut scuba = self.scuba.clone();
        scuba.sampled(self.scuba_sample_rate);

        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobGets);

        let pc = ctx.fork_perf_counters();

        let get = self.inner.get(&ctx, key);
        let (stats, result) = get.timed().await;
        record_get_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            key,
            ctx.metadata().session_id().as_str(),
            OperationType::Get,
            None,
            &self.inner,
        );

        match result {
            Ok(Some(ref data)) => {
                ctx.perf_counters().add_to_counter(
                    PerfCounterType::BlobGetsTotalSize,
                    data.len().try_into().unwrap_or(0),
                );
            }
            _ => {}
        }

        result
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let mut ctx = ctx.clone();
        let mut scuba = self.scuba.clone();
        scuba.sampled(self.scuba_sample_rate);

        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobPresenceChecks);

        let pc = ctx.fork_perf_counters();

        let is_present = self.inner.is_present(&ctx, key);
        let (stats, result) = is_present.timed().await;
        record_is_present_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            key,
            ctx.metadata().session_id().as_str(),
            None,
            &self.inner,
        );

        result
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        BlobstorePutOps::put_with_status(self, ctx, key, value).await?;
        Ok(())
    }
}

impl<B: BlobstorePutOps> LogBlob<B> {
    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
    ) -> Result<OverwriteStatus> {
        let mut ctx = ctx.clone();
        let mut scuba = self.scuba.clone();
        let size = value.len();

        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobPuts);

        let pc = ctx.fork_perf_counters();

        let put = if let Some(put_behaviour) = put_behaviour {
            self.inner
                .put_explicit(&ctx, key.clone(), value, put_behaviour)
        } else {
            self.inner.put_with_status(&ctx, key.clone(), value)
        };
        let (stats, result) = put.timed().await;
        record_put_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            &key,
            ctx.metadata().session_id().as_str(),
            size,
            None,
            &self.inner,
            None,
        );

        if result.is_ok() {
            ctx.perf_counters().add_to_counter(
                PerfCounterType::BlobPutsTotalSize,
                size.try_into().unwrap_or(0),
            );
        }

        result
    }
}

#[async_trait]
impl<B: BlobstorePutOps> BlobstorePutOps for LogBlob<B> {
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        self.put_impl(ctx, key, value, Some(put_behaviour)).await
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.put_impl(ctx, key, value, None).await
    }
}


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

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

use anyhow::format_err;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreEnumerationData;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreKeyParam;
use blobstore::BlobstoreKeySource;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore::DEFAULT_PUT_BEHAVIOUR;
use context::CoreContext;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use mononoke_types::BlobstoreBytes;

// Implements hardlink-style links
#[derive(Default, Debug)]
struct MemState {
    next_id: usize,
    data: HashMap<usize, BlobstoreBytes>,
    links: BTreeMap<String, usize>,
}

impl MemState {
    fn put(
        &mut self,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> OverwriteStatus {
        match put_behaviour {
            PutBehaviour::Overwrite => {
                let id = self.next_id;
                self.data.insert(id, value);
                self.links.insert(key, id);
                self.next_id += 1;
                OverwriteStatus::NotChecked
            }
            PutBehaviour::IfAbsent | PutBehaviour::OverwriteAndLog => {
                if self.links.contains_key(&key) {
                    if put_behaviour.should_overwrite() {
                        self.put(key, value, PutBehaviour::Overwrite);
                        OverwriteStatus::Overwrote
                    } else {
                        OverwriteStatus::Prevented
                    }
                } else {
                    self.put(key, value, PutBehaviour::Overwrite);
                    OverwriteStatus::New
                }
            }
        }
    }

    fn link(&mut self, existing_key: &str, link_key: String) -> Result<()> {
        if let Some(existing_id) = self.links.get(existing_key) {
            let existing_id = *existing_id;
            self.links.insert(link_key, existing_id);
            return Ok(());
        }
        panic!("Unknown existing_key {} {:?}", existing_key, self.links)
    }

    fn get(&self, key: &str) -> Option<&BlobstoreBytes> {
        if let Some(id) = self.links.get(key) {
            self.data.get(id)
        } else {
            None
        }
    }

    fn unlink(&mut self, key: &str) -> Option<()> {
        self.links.remove(key).map(|_| ())
    }
}

/// In-memory "blob store"
///
/// Pure in-memory implementation for testing.
#[derive(Clone)]
pub struct Memblob {
    state: Arc<Mutex<MemState>>,
    put_behaviour: PutBehaviour,
}

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

impl Memblob {
    pub fn new(put_behaviour: PutBehaviour) -> Self {
        Self {
            state: Arc::new(Mutex::new(MemState::default())),
            put_behaviour,
        }
    }

    pub fn unlink(&self, key: String) -> BoxFuture<'static, Result<Option<()>>> {
        let state = self.state.clone();

        async move {
            let mut inner = state.lock().expect("lock poison");
            Ok(inner.unlink(&key))
        }
        .boxed()
    }
}

impl Default for Memblob {
    fn default() -> Self {
        Self::new(DEFAULT_PUT_BEHAVIOUR)
    }
}

#[async_trait]
impl BlobstorePutOps for Memblob {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        let state = self.state.clone();

        let mut inner = state.lock().expect("lock poison");
        Ok(inner.put(key, value, put_behaviour))
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.put_explicit(ctx, key, value, self.put_behaviour).await
    }
}

#[async_trait]
impl Blobstore for Memblob {
    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let state = self.state.clone();

        let inner = state.lock().expect("lock poison");
        Ok(inner.get(key).map(|bytes| bytes.clone().into()))
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        BlobstorePutOps::put_with_status(self, ctx, key, value).await?;
        Ok(())
    }

    async fn copy<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        let state = self.state.clone();

        let mut inner = state.lock().expect("lock poison");
        inner.link(old_key, new_key)
    }
}

#[async_trait]
impl BlobstoreUnlinkOps for Memblob {
    async fn unlink<'a>(&'a self, _ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        let state = self.state.clone();
        let mut inner = state.lock().expect("lock poison");
        if inner.unlink(key).is_some() {
            Ok(())
        } else {
            Err(format_err!("Unknown key {} to Memblob::unlink()", key))
        }
    }
}

#[async_trait]
impl BlobstoreKeySource for Memblob {
    async fn enumerate<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData> {
        match range {
            BlobstoreKeyParam::Start(range) => {
                let state = self.state.lock().expect("lock poison");
                Ok(BlobstoreEnumerationData {
                    keys: state.links.range(range).map(|(k, _)| k.clone()).collect(),
                    next_token: None,
                })
            }
            BlobstoreKeyParam::Continuation(_) => {
                Err(format_err!("Continuation not supported for memblob"))
            }
        }
    }
}

impl fmt::Debug for Memblob {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Memblob")
            .field("state", &self.state.lock().expect("lock poisoned"))
            .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::collections::HashMap;
use std::collections::HashSet;
use std::future::Future;
use std::hash::Hasher;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Error;
use anyhow::Result;
use blobstore::BlobstoreGetData;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_stats::record_put_stats;
use context::CoreContext;
use futures::future;
use futures_stats::TimedFutureExt;
use itertools::Either;
use metaconfig_types::BlobstoreId;
use mononoke_types::BlobstoreBytes;
use scuba_ext::MononokeScubaSampleBuilder;
use thiserror::Error;
use tokio::time::timeout;
use twox_hash::XxHash;

use crate::scrub::SrubWriteOnly;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(600);

type BlobstoresWithEntry = Vec<HashSet<BlobstoreId>>;
type BlobstoresReturnedNone = HashSet<BlobstoreId>;
type BlobstoresReturnedError = HashMap<BlobstoreId, Error>;

#[derive(Error, Debug, Clone)]
pub enum ErrorKind {
    #[error("Some blobstores failed, and other returned None: {main_errors:?}")]
    SomeFailedOthersNone {
        main_errors: Arc<BlobstoresReturnedError>,
        write_only_errors: Arc<BlobstoresReturnedError>,
    },
    #[error("All blobstores failed: {main_errors:?}")]
    AllFailed {
        main_errors: Arc<BlobstoresReturnedError>,
        write_only_errors: Arc<BlobstoresReturnedError>,
    },
    // Errors below this point are from ScrubBlobstore only. If they include an
    // Option<BlobstoreBytes>, this implies that this error is recoverable
    #[error(
        "Different blobstores have different values for this item: {0:?} are grouped by content, {1:?} do not have"
    )]
    ValueMismatch(Arc<BlobstoresWithEntry>, Arc<BlobstoresReturnedNone>),
    #[error("Some blobstores missing this item: {missing_main:?}")]
    SomeMissingItem {
        missing_main: Arc<BlobstoresReturnedNone>,
        missing_write_only: Arc<BlobstoresReturnedNone>,
        value: BlobstoreGetData,
    },
    #[error("Multiple failures on put: {0:?}")]
    MultiplePutFailures(Arc<BlobstoresReturnedError>),
}

fn blobstores_failed_error(
    main_blobstore_ids: impl Iterator<Item = BlobstoreId>,
    main_errors: HashMap<BlobstoreId, Error>,
    write_only_errors: HashMap<BlobstoreId, Error>,
) -> ErrorKind {
    let main_errored_ids: HashSet<BlobstoreId> = main_errors.keys().copied().collect();
    let all_main_ids: HashSet<BlobstoreId> = main_blobstore_ids.collect();
    if main_errored_ids == all_main_ids {
        // The write only stores that returned None might not have been fully populated
        ErrorKind::AllFailed {
            main_errors: Arc::new(main_errors),
            write_only_errors: Arc::new(write_only_errors),
        }
    } else {
        ErrorKind::SomeFailedOthersNone {
            main_errors: Arc::new(main_errors),
            write_only_errors: Arc::new(write_only_errors),
        }
    }
}

type GetResult = (BlobstoreId, Result<Option<BlobstoreGetData>, Error>);

/// Get normal and write only results based on SrubWriteOnly
/// mode of getting results, which might optimise for less access
/// to main blobstores or less access to write only blobstores
pub async fn scrub_get_results<MF, WF>(
    get_main_results: impl FnOnce() -> MF,
    mut get_write_only_results: impl FnMut() -> WF,
    write_only_blobstores: impl Iterator<Item = BlobstoreId>,
    write_only: SrubWriteOnly,
) -> impl Iterator<Item = (bool, GetResult)>
where
    MF: Future<Output = Vec<GetResult>>,
    WF: Future<Output = Vec<GetResult>>,
{
    // Exit early if all mostly-write are ok, and don't check main blobstores
    if write_only == SrubWriteOnly::ScrubIfAbsent {
        let mut results = get_write_only_results().await.into_iter();
        if let Some((bs_id, Ok(Some(data)))) = results.next() {
            if results.all(|(_, r)| match r {
                Ok(Some(other_data)) => other_data == data,
                _ => false,
            }) {
                return Either::Left(std::iter::once((true, (bs_id, Ok(Some(data))))));
            }
        }
    }

    let write_only_results = async {
        match write_only {
            SrubWriteOnly::Scrub | SrubWriteOnly::SkipMissing => get_write_only_results().await,
            SrubWriteOnly::PopulateIfAbsent | SrubWriteOnly::ScrubIfAbsent => {
                write_only_blobstores.map(|id| (id, Ok(None))).collect()
            }
        }
    };
    let (normal_results, write_only_results) =
        future::join(get_main_results(), write_only_results).await;

    Either::Right(
        normal_results
            .into_iter()
            .map(|r| (false, r))
            .chain(write_only_results.into_iter().map(|r| (true, r))),
    )
}

pub fn scrub_parse_results(
    results: impl Iterator<Item = (bool, GetResult)>,
    all_main: impl Iterator<Item = BlobstoreId>,
) -> Result<Option<BlobstoreGetData>, ErrorKind> {
    let mut missing_main = HashSet::new();
    let mut missing_write_only = HashSet::new();
    let mut get_data = None;
    let mut main_errors = HashMap::new();
    let mut write_only_errors = HashMap::new();

    for (is_write_only, (blobstore_id, result)) in results {
        match result {
            Ok(None) => {
                if is_write_only {
                    missing_write_only.insert(blobstore_id);
                } else {
                    missing_main.insert(blobstore_id);
                }
            }
            Ok(Some(value)) => {
                let mut content_hash = XxHash::with_seed(0);
                content_hash.write(value.as_raw_bytes());
                let content_hash = content_hash.finish();
                let (all_values, _) = get_data.get_or_insert_with(|| (HashMap::new(), value));
                all_values
                    .entry(content_hash)
                    .or_insert_with(HashSet::new)
                    .insert(blobstore_id);
            }
            Err(err) => {
                if is_write_only {
                    write_only_errors.insert(blobstore_id, err);
                } else {
                    main_errors.insert(blobstore_id, err);
                }
            }
        }
    }
    match get_data {
        None => {
            if main_errors.is_empty() && write_only_errors.is_empty() {
                Ok(None)
            } else {
                Err(blobstores_failed_error(
                    all_main,
                    main_errors,
                    write_only_errors,
                ))
            }
        }
        Some((all_values, value)) if all_values.len() == 1 => {
            if missing_main.is_empty() && missing_write_only.is_empty() {
                Ok(Some(value))
            } else {
                // This silently ignores failed blobstores if at least one has a value
                Err(ErrorKind::SomeMissingItem {
                    missing_main: Arc::new(missing_main),
                    missing_write_only: Arc::new(missing_write_only),
                    value,
                })
            }
        }
        Some((all_values, _)) => {
            let answered = all_values.into_values().collect();
            let mut all_missing = HashSet::new();
            all_missing.extend(missing_main.into_iter());
            all_missing.extend(missing_write_only.into_iter());
            Err(ErrorKind::ValueMismatch(
                Arc::new(answered),
                Arc::new(all_missing),
            ))
        }
    }
}

fn remap_timeout_result<O>(
    timeout_or_result: Result<Result<O, Error>, tokio::time::error::Elapsed>,
) -> Result<O, Error> {
    timeout_or_result.unwrap_or_else(|_| Err(Error::msg("blobstore operation timeout")))
}

pub async fn inner_put(
    ctx: &CoreContext,
    mut scuba: MononokeScubaSampleBuilder,
    write_order: &AtomicUsize,
    blobstore_id: BlobstoreId,
    blobstore: &dyn BlobstorePutOps,
    key: String,
    value: BlobstoreBytes,
    put_behaviour: Option<PutBehaviour>,
) -> (BlobstoreId, Result<OverwriteStatus, Error>) {
    let size = value.len();
    let (pc, (stats, timeout_or_res)) = {
        let mut ctx = ctx.clone();
        let pc = ctx.fork_perf_counters();
        let ret = timeout(
            REQUEST_TIMEOUT,
            if let Some(put_behaviour) = put_behaviour {
                blobstore.put_explicit(&ctx, key.clone(), value, put_behaviour)
            } else {
                blobstore.put_with_status(&ctx, key.clone(), value)
            },
        )
        .timed()
        .await;
        (pc, ret)
    };
    let result = remap_timeout_result(timeout_or_res);
    record_put_stats(
        &mut scuba,
        &pc,
        stats,
        result.as_ref(),
        &key,
        ctx.metadata().session_id().as_str(),
        size,
        Some(blobstore_id),
        blobstore,
        Some(write_order.fetch_add(1, Ordering::Relaxed) + 1),
    );
    (blobstore_id, result)
}


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

pub mod base;
pub mod scrub;
pub mod scuba;

pub use crate::scrub::LoggingScrubHandler;
pub use crate::scrub::ScrubAction;
pub use crate::scrub::ScrubHandler;
pub use crate::scrub::ScrubOptions;
pub use crate::scrub::SrubWriteOnly;


/*
 * 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::cmp::max;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreMetadata;
use blobstore::BlobstorePutOps;
use blobstore::PutBehaviour;
use chrono::Duration as ChronoDuration;
use clap::ValueEnum;
use context::CoreContext;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use futures::Future;
use metaconfig_types::BlobstoreId;
use mononoke_types::Timestamp;
use once_cell::sync::Lazy;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::info;
use slog::warn;
use strum::EnumString;
use strum::EnumVariantNames;
use strum::IntoStaticStr;

use crate::base::inner_put;

static HEAL_MAX_BACKLOG: Lazy<Duration> =
    Lazy::new(|| Duration::from_secs(ChronoDuration::days(7).num_seconds() as u64));

/// What to do when the ScrubBlobstore finds a problem
#[derive(
    Debug,
    Clone,
    Copy,
    Eq,
    PartialEq,
    Hash,
    ValueEnum,
    EnumString,
    EnumVariantNames,
    IntoStaticStr
)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum ScrubAction {
    /// Log items needing repair
    ReportOnly,
    /// Do repairs
    Repair,
}

// How to treat write only stores during the scrub
#[derive(
    Debug,
    Clone,
    Copy,
    Eq,
    PartialEq,
    Hash,
    ValueEnum,
    EnumString,
    EnumVariantNames,
    IntoStaticStr
)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum SrubWriteOnly {
    /// don't take action on scrub missing keys from write only stores
    SkipMissing,
    /// take the normal scrub action for write only stores
    Scrub,
    /// Mode for populating empty stores.  Assumes its already missing. Don't attempt to read. Write with IfAbsent so won't overwrite if run incorrectluy.
    /// More efficient than the above if thes store is totally empty.
    PopulateIfAbsent,
    /// Mode for rescrubbing write-only stores before enabling them. Assumes that the data in them is correct,
    /// and won't read from the main stores unless the write-only stores have missing data or read failures
    /// This ensures that load on the main stores is kept to a minimum
    ScrubIfAbsent,
}

#[derive(Clone, Debug)]
pub struct ScrubOptions {
    pub scrub_action: ScrubAction,
    pub scrub_grace: Option<Duration>,
    pub scrub_action_on_missing_write_only: SrubWriteOnly,
    pub queue_peek_bound: Duration,
}

impl Default for ScrubOptions {
    fn default() -> Self {
        Self {
            scrub_action: ScrubAction::ReportOnly,
            scrub_grace: None,
            scrub_action_on_missing_write_only: SrubWriteOnly::Scrub,
            queue_peek_bound: *HEAL_MAX_BACKLOG,
        }
    }
}

pub fn default_scrub_handler() -> Arc<dyn ScrubHandler> {
    Arc::new(LoggingScrubHandler::new(false))
}

pub trait ScrubHandler: Send + Sync + fmt::Debug {
    /// Called when one of the inner stores required repair.
    fn on_repair(
        &self,
        ctx: &CoreContext,
        blobstore_id: BlobstoreId,
        key: &str,
        is_repaired: bool,
        meta: &BlobstoreMetadata,
    );
}

#[derive(Debug)]
pub struct LoggingScrubHandler {
    quiet: bool,
}

impl LoggingScrubHandler {
    pub fn new(quiet: bool) -> Self {
        Self { quiet }
    }
}

impl ScrubHandler for LoggingScrubHandler {
    fn on_repair(
        &self,
        ctx: &CoreContext,
        blobstore_id: BlobstoreId,
        key: &str,
        is_repaired: bool,
        _meta: &BlobstoreMetadata,
    ) {
        if !self.quiet {
            if is_repaired {
                info!(
                    ctx.logger(),
                    "scrub: blobstore_id {:?} repaired for {}", &blobstore_id, &key
                );
            } else {
                warn!(
                    ctx.logger(),
                    "scrub: blobstore_id {:?} not repaired for {}", &blobstore_id, &key
                );
            }
        }
    }
}

// Would be a closure, but async closures are unstable
async fn put_and_mark_repaired(
    ctx: &CoreContext,
    scuba: &MononokeScubaSampleBuilder,
    order: &AtomicUsize,
    id: BlobstoreId,
    store: &dyn BlobstorePutOps,
    key: &str,
    value: &BlobstoreGetData,
    scrub_handler: &dyn ScrubHandler,
    put_behaviour: PutBehaviour,
) -> Result<()> {
    let (_, res) = inner_put(
        ctx,
        scuba.clone(),
        order,
        id,
        store,
        key.to_owned(),
        value.as_bytes().clone(),
        Some(put_behaviour),
    )
    .await;
    scrub_handler.on_repair(ctx, id, key, res.is_ok(), value.as_meta());
    res.map(|_status| ())
}

pub async fn maybe_repair<F: Future<Output = Result<bool>>>(
    ctx: &CoreContext,
    key: &str,
    value: BlobstoreGetData,
    missing_main: Arc<HashSet<BlobstoreId>>,
    missing_write_only: Arc<HashSet<BlobstoreId>>,
    scrub_stores: &HashMap<BlobstoreId, Arc<dyn BlobstorePutOps>>,
    scrub_handler: &dyn ScrubHandler,
    scrub_options: &ScrubOptions,
    scuba: &MononokeScubaSampleBuilder,
    already_healed: impl FnOnce() -> F,
) -> Result<Option<BlobstoreGetData>> {
    let ctime_age = value.as_meta().ctime().map(|ctime| {
        let age_secs = max(0, Timestamp::from_timestamp_secs(ctime).since_seconds());
        Duration::from_secs(age_secs as u64)
    });

    match (ctime_age, scrub_options.scrub_grace) {
        // value written recently, within the grace period, so don't attempt repair
        (Some(ctime_age), Some(scrub_grace)) if ctime_age < scrub_grace => {
            return Ok(Some(value));
        }
        _ => {}
    }

    let mut needs_repair: HashMap<BlobstoreId, (PutBehaviour, &dyn BlobstorePutOps)> =
        HashMap::new();

    // For write only stores we can chose not to do the scrub action
    // e.g. if store is still being populated, a checking scrub wouldn't want to raise alarm on the store
    if scrub_options.scrub_action_on_missing_write_only != SrubWriteOnly::SkipMissing
        || !missing_main.is_empty()
    {
        // Only peek the queue if needed
        let already_healed = match ctime_age.as_ref() {
            // Avoid false alarms for recently written items still on the healer queue
            Some(ctime_age) if ctime_age < &scrub_options.queue_peek_bound => {
                already_healed().await?
            }
            _ => true,
        };

        // Only attempt the action if we don't know of pending writes from the queue
        if already_healed {
            for k in missing_main.iter() {
                if let Some(s) = scrub_stores.get(k) {
                    // We are repairing, overwrite is right thing to do as
                    // bad keys may be is_present, but not retrievable.
                    needs_repair.insert(*k, (PutBehaviour::Overwrite, s.as_ref()));
                }
            }
            for k in missing_write_only.iter() {
                if let Some(s) = scrub_stores.get(k) {
                    let put_behaviour = match scrub_options.scrub_action_on_missing_write_only {
                        SrubWriteOnly::SkipMissing => None,
                        SrubWriteOnly::Scrub => Some(PutBehaviour::Overwrite),
                        SrubWriteOnly::PopulateIfAbsent | SrubWriteOnly::ScrubIfAbsent => {
                            Some(PutBehaviour::IfAbsent)
                        }
                    };
                    if let Some(put_behaviour) = put_behaviour {
                        needs_repair.insert(*k, (put_behaviour, s.as_ref()));
                    }
                }
            }
        }
    }

    if scrub_options.scrub_action == ScrubAction::ReportOnly {
        for id in needs_repair.keys() {
            scrub_handler.on_repair(ctx, *id, key, false, value.as_meta());
        }
    } else {
        // inner_put to the stores that need it.
        let order = AtomicUsize::new(0);
        let repair_puts: FuturesUnordered<_> = needs_repair
            .into_iter()
            .map(|(id, (put_behaviour, store))| {
                put_and_mark_repaired(
                    ctx,
                    scuba,
                    &order,
                    id,
                    store,
                    key,
                    &value,
                    scrub_handler,
                    put_behaviour,
                )
            })
            .collect();

        repair_puts.try_for_each(|_| async { Ok(()) }).await?;
    }
    Ok(Some(value))
}


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

//! Logic for logging to mononoke_multiplex scuba table.
//! Maybe we should be using `LogBlob` here, but that might need
//! some modification there.

use anyhow::Result;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore_stats::add_completion_time;
use blobstore_stats::OperationType;
use blobstore_stats::BLOB_PRESENT;
use blobstore_stats::ERROR;
use blobstore_stats::KEY;
use blobstore_stats::OPERATION;
use context::CoreContext;
use futures_stats::FutureStats;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use scuba_ext::MononokeScubaSampleBuilder;

const MULTIPLEX_ID: &str = "multiplex_id";
const BLOB_SIZE: &str = "blob_size";
const SUCCESS: &str = "success";
const SYNC_QUEUE: &str = "mysql_sync_queue";

fn record_scuba_common(
    mut ctx: CoreContext,
    scuba: &mut MononokeScubaSampleBuilder,
    multiplex_id: &MultiplexId,
    key: &str,
    stats: FutureStats,
    operation: OperationType,
) {
    let pc = ctx.fork_perf_counters();
    pc.insert_nonzero_perf_counters(scuba);

    add_completion_time(scuba, ctx.metadata().session_id().as_str(), stats);

    scuba.add(KEY, key);
    scuba.add(OPERATION, operation);
    scuba.add(MULTIPLEX_ID, multiplex_id.clone());
}

pub fn record_put<T>(
    ctx: &CoreContext,
    scuba: &mut MononokeScubaSampleBuilder,
    multiplex_id: &MultiplexId,
    key: &str,
    blob_size: usize,
    stats: FutureStats,
    result: &Result<T>,
) {
    let op = OperationType::Put;
    record_scuba_common(ctx.clone(), scuba, multiplex_id, key, stats, op);

    scuba.add(BLOB_SIZE, blob_size);

    if let Err(error) = result.as_ref() {
        scuba.unsampled();
        scuba.add(ERROR, format!("{:#}", error)).add(SUCCESS, false);
    } else {
        scuba.add(SUCCESS, true);
    }
    scuba.log();
}

pub fn record_get(
    ctx: &CoreContext,
    scuba: &mut MononokeScubaSampleBuilder,
    multiplex_id: &MultiplexId,
    key: &str,
    stats: FutureStats,
    result: &Result<Option<BlobstoreGetData>>,
) {
    let op = OperationType::Get;
    record_scuba_common(ctx.clone(), scuba, multiplex_id, key, stats, op);

    match result.as_ref() {
        Err(error) => {
            scuba.unsampled();
            scuba.add(ERROR, format!("{:#}", error)).add(SUCCESS, false);
        }
        Ok(mb_blob) => {
            let blob_present = mb_blob.is_some();
            scuba.add(BLOB_PRESENT, blob_present).add(SUCCESS, true);

            if let Some(blob) = mb_blob.as_ref() {
                scuba.add(BLOB_SIZE, blob.len());
            }
        }
    }
    scuba.log();
}

pub fn record_is_present(
    ctx: &CoreContext,
    scuba: &mut MononokeScubaSampleBuilder,
    multiplex_id: &MultiplexId,
    key: &str,
    stats: FutureStats,
    result: &Result<BlobstoreIsPresent>,
) {
    let op = OperationType::IsPresent;
    record_scuba_common(ctx.clone(), scuba, multiplex_id, key, stats, op);

    let outcome = result.as_ref().map(|is_present| match is_present {
        BlobstoreIsPresent::Present => Some(true),
        BlobstoreIsPresent::Absent => Some(false),
        BlobstoreIsPresent::ProbablyNotPresent(_) => None,
    });

    match outcome {
        Err(error) => {
            scuba.unsampled();
            scuba.add(ERROR, format!("{:#}", error)).add(SUCCESS, false);
        }
        Ok(is_present) => {
            if let Some(is_present) = is_present {
                scuba.add(BLOB_PRESENT, is_present).add(SUCCESS, true);
            }
        }
    }
    scuba.log();
}

pub fn record_queue_stats(
    ctx: &CoreContext,
    scuba: &mut MononokeScubaSampleBuilder,
    key: &str,
    stats: FutureStats,
    blobstore_id: Option<BlobstoreId>,
    blobstore_type: String,
    result: Result<&(), &anyhow::Error>,
) {
    let pc = ctx.clone().fork_perf_counters();
    blobstore_stats::record_queue_stats(
        scuba,
        &pc,
        stats,
        result,
        key,
        ctx.metadata().session_id().as_str(),
        OperationType::Put,
        blobstore_id,
        blobstore_type,
        SYNC_QUEUE,
    );
}


/*
 * 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(crate) mod multiplex;
pub mod scrub;
#[cfg(test)]
mod test;
mod timed;

pub use multiplex::MultiplexQuorum;
pub use multiplex::Scuba;
pub use multiplex::WalMultiplexedBlobstore;
pub use timed::MultiplexTimeout;


/*
 * 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::fmt;
use std::num::NonZeroU64;
use std::num::NonZeroUsize;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Context as _;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_stats::OperationType;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use cloned::cloned;
use context::CoreContext;
use context::PerfCounterType;
use fbinit::FacebookInit;
use futures::future;
use futures::stream::FuturesUnordered;
use futures::Future;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use futures_stats::TimedFutureExt;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::Timestamp;
use multiplexedblob::scuba;
use scuba_ext::MononokeScubaSampleBuilder;
use thiserror::Error;
use time_ext::DurationExt;
use tokio::task::JoinHandle;

use crate::timed::with_timed_stores;
use crate::timed::MultiplexTimeout;
use crate::timed::TimedStore;
type BlobstoresReturnedError = HashMap<BlobstoreId, Error>;

#[derive(Error, Debug, Clone)]
pub enum ErrorKind {
    #[error("All blobstores failed: {0:?}")]
    AllFailed(Arc<BlobstoresReturnedError>),
    #[error("Failures on put in underlying single blobstores: {0:?}")]
    SomePutsFailed(Arc<BlobstoresReturnedError>),
    #[error("Failures on get in underlying single blobstores: {0:?}")]
    SomeGetsFailed(Arc<BlobstoresReturnedError>),
    #[error("Failures on is_present in underlying single blobstores: {0:?}")]
    SomeIsPresentsFailed(Arc<BlobstoresReturnedError>),
}

#[derive(Clone, Debug)]
pub struct MultiplexQuorum {
    pub(crate) read: NonZeroUsize,
    pub(crate) write: NonZeroUsize,
}

impl MultiplexQuorum {
    fn new(num_stores: usize, write: usize) -> Result<Self> {
        if write > num_stores {
            return Err(anyhow!(
                "Not enough blobstores for configured put or get needs. Have {}, need {} puts",
                num_stores,
                write,
            ));
        }

        Ok(Self {
            write: NonZeroUsize::new(write).ok_or_else(|| anyhow!("Write quorum cannot be 0"))?,
            read: NonZeroUsize::new(num_stores - write + 1).unwrap(),
        })
    }
}

#[derive(Clone)]
pub struct Scuba {
    pub(crate) inner_blobstores_scuba: MononokeScubaSampleBuilder,
    multiplex_scuba: MononokeScubaSampleBuilder,
    sample_rate: NonZeroU64,
}

impl Scuba {
    pub fn new_from_raw(
        fb: FacebookInit,
        inner_blobstores_scuba_table: Option<String>,
        multiplex_scuba_table: Option<String>,
        sample_rate: NonZeroU64,
    ) -> Result<Self> {
        let inner = inner_blobstores_scuba_table.map_or_else(
            || Ok(MononokeScubaSampleBuilder::with_discard()),
            |table| MononokeScubaSampleBuilder::new(fb, &table),
        )?;
        let multiplex = multiplex_scuba_table.map_or_else(
            || Ok(MononokeScubaSampleBuilder::with_discard()),
            |table| MononokeScubaSampleBuilder::new(fb, &table),
        )?;

        Self::new(inner, multiplex, sample_rate)
    }

    pub fn new(
        mut inner_blobstores_scuba: MononokeScubaSampleBuilder,
        mut multiplex_scuba: MononokeScubaSampleBuilder,
        sample_rate: NonZeroU64,
    ) -> Result<Self> {
        inner_blobstores_scuba.add_common_server_data();
        multiplex_scuba.add_common_server_data();
        Ok(Self {
            inner_blobstores_scuba,
            multiplex_scuba,
            sample_rate,
        })
    }

    pub fn sampled(&mut self) {
        self.inner_blobstores_scuba.sampled(self.sample_rate);
        self.multiplex_scuba.sampled(self.sample_rate);
    }
}

#[derive(Clone)]
pub struct WalMultiplexedBlobstore {
    /// Multiplexed blobstore configuration.
    pub(crate) multiplex_id: MultiplexId,
    /// Write-ahead log used to keep data consistent across blobstores.
    pub(crate) wal_queue: Arc<dyn BlobstoreWal>,

    pub(crate) quorum: MultiplexQuorum,
    /// These are the "normal" blobstores, which are read from on `get`, and written to on `put`
    /// as part of normal operation.
    pub(crate) blobstores: Arc<[TimedStore]>,
    /// Write-mostly blobstores are not normally read from on `get`, but take part in writes
    /// like a normal blobstore.
    pub(crate) write_only_blobstores: Arc<[TimedStore]>,

    /// Scuba table to log status of the underlying single blobstore queries.
    pub(crate) scuba: Scuba,

    /// Counter keeping track of the yet-to-complete blobstore operations in flight.
    pub(crate) inflight_ops_counter: Arc<AtomicU64>,
}

impl Drop for WalMultiplexedBlobstore {
    fn drop(&mut self) {
        // If there are any inflight get/put/is_present request to the blobstores,
        // wait a while before exiting. This is required to prevent the UAF issue
        // in Manifold. Reference: https://fburl.com/y5o44ed6
        if self.inflight_ops_counter.load(Ordering::Relaxed) > 0 {
            std::thread::sleep(std::time::Duration::from_secs(5))
        }
    }
}

impl std::fmt::Display for WalMultiplexedBlobstore {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "WAL MultiplexedBlobstore[normal {:?}, write only {:?}]",
            self.blobstores, self.write_only_blobstores
        )
    }
}

impl fmt::Debug for WalMultiplexedBlobstore {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "WalMultiplexedBlobstore: multiplex_id: {}",
            &self.multiplex_id
        )?;
        f.debug_map()
            .entries(self.blobstores.iter().map(|v| (v.id(), v)))
            .finish()
    }
}

impl WalMultiplexedBlobstore {
    pub fn new(
        multiplex_id: MultiplexId,
        wal_queue: Arc<dyn BlobstoreWal>,
        blobstores: Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
        write_only_blobstores: Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
        write_quorum: usize,
        timeout: Option<MultiplexTimeout>,
        scuba: Scuba,
    ) -> Result<Self> {
        let quorum = MultiplexQuorum::new(blobstores.len(), write_quorum)?;

        let to = timeout.unwrap_or_default();
        let blobstores = with_timed_stores(blobstores, to.clone()).into();
        let write_only_blobstores = with_timed_stores(write_only_blobstores, to).into();
        let inflight_ops_counter = Arc::new(AtomicU64::new(0));
        Ok(Self {
            multiplex_id,
            wal_queue,
            blobstores,
            write_only_blobstores,
            quorum,
            scuba,
            inflight_ops_counter,
        })
    }

    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
        scuba: &Scuba,
    ) -> Result<OverwriteStatus> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobPuts);

        let blob_size = value.len() as u64;

        // Log the blobstore key and wait till it succeeds
        let ts = Timestamp::now();
        let log_entry = BlobstoreWalEntry::new(key.clone(), self.multiplex_id, ts, blob_size);
        let (stats, result) = self.wal_queue.log(ctx, log_entry).timed().await;

        scuba::record_queue_stats(
            ctx,
            &mut scuba.multiplex_scuba.clone(),
            &key,
            stats,
            None,
            self.to_string(),
            result.as_ref().map(|_| &()),
        );

        let entry = result.with_context(|| {
            format!(
                "WAL Multiplexed Blobstore: Failed writing to the WAL: key {}",
                key
            )
        })?;

        // Prepare underlying main blobstores puts
        let mut put_futs = inner_multi_put(
            ctx,
            self.blobstores.clone(),
            &key,
            &value,
            put_behaviour,
            scuba,
            self.inflight_ops_counter.clone(),
        );

        // Wait for the quorum successful writes
        let mut quorum: usize = self.quorum.write.get();
        let mut put_errors = HashMap::new();
        let (stats, result) = async move {
            while let Some(result) = put_futs.next().await {
                match result {
                    Ok(_overwrite_status) => {
                        quorum = quorum.saturating_sub(1);
                        if quorum == 0 {
                            // Quorum blobstore writes succeeded, we can spawn the rest
                            // of the writes and not wait for them.
                            let main_puts =
                                spawn_stream_completion(put_futs.map_err(|(_id, err)| err));

                            // Spawn the write-only blobstore writes, we don't want to wait for them
                            let write_only_puts = inner_multi_put(
                                ctx,
                                self.write_only_blobstores.clone(),
                                &key,
                                &value,
                                put_behaviour,
                                scuba,
                                self.inflight_ops_counter.clone(),
                            );
                            let write_only_puts =
                                spawn_stream_completion(write_only_puts.map_err(|(_id, err)| err));

                            cloned!(ctx, self.wal_queue);
                            if put_errors.is_empty() {
                                // Optimisation: It put fully succeeded on all blobstores, we can remove
                                // it from queue and healer doesn't need to deal with it.
                                tokio::spawn(async move {
                                    let (r1, r2) = futures::join!(main_puts, write_only_puts);
                                    r1??;
                                    r2??;
                                    // TODO(yancouto): Batch deletes together.
                                    wal_queue.delete_by_key(&ctx, &[entry]).await?;
                                    anyhow::Ok(())
                                });
                            }

                            return Ok(OverwriteStatus::NotChecked);
                        }
                    }
                    Err((bs_id, err)) => {
                        put_errors.insert(bs_id, err);
                    }
                }
            }
            Err(put_errors)
        }
        .timed()
        .await;

        ctx.perf_counters().set_max_counter(
            PerfCounterType::BlobPutsMaxLatency,
            stats.completion_time.as_millis_unchecked() as i64,
        );
        ctx.perf_counters()
            .add_to_counter(PerfCounterType::BlobPutsTotalSize, blob_size as i64);

        result.map_err(|put_errors| {
            let errors = Arc::new(put_errors);
            let result_err = if errors.len() == self.blobstores.len() {
                // all main writes failed
                ErrorKind::AllFailed(errors)
            } else {
                // some main writes failed
                ErrorKind::SomePutsFailed(errors)
            };
            result_err.into()
        })
    }

    async fn get_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
        scuba: &Scuba,
    ) -> Result<Option<BlobstoreGetData>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobGets);

        let mut get_futs = inner_multi_get(
            ctx,
            self.blobstores.clone(),
            key,
            OperationType::Get,
            scuba,
            self.inflight_ops_counter.clone(),
        );

        // Wait for the quorum successful "Not Found" reads before
        // returning Ok(None).
        let mut quorum: usize = self.quorum.read.get();
        let mut get_errors = HashMap::with_capacity(get_futs.len());
        let (stats, result) = async move {
            while let Some((bs_id, result)) = get_futs.next().await {
                match result {
                    Ok(Some(get_data)) => {
                        return Ok(Some(get_data));
                    }
                    Ok(None) => {
                        quorum = quorum.saturating_sub(1);
                        if quorum == 0 {
                            // quorum blobstores couldn't find the given key in the blobstores
                            // let's trust them
                            return Ok(None);
                        }
                    }
                    Err(err) => {
                        get_errors.insert(bs_id, err);
                    }
                }
            }
            Err(get_errors)
        }
        .timed()
        .await;

        ctx.perf_counters().set_max_counter(
            PerfCounterType::BlobGetsMaxLatency,
            stats.completion_time.as_millis_unchecked() as i64,
        );

        let result = result.map_err(|get_errors| {
            let errors = Arc::new(get_errors);
            let result_err = if errors.len() == self.blobstores.len() {
                // all main reads failed
                ErrorKind::AllFailed(errors)
            } else {
                // some main reads failed
                ErrorKind::SomeGetsFailed(errors)
            };
            result_err.into()
        });

        match result {
            Ok(Some(ref data)) => {
                ctx.perf_counters()
                    .add_to_counter(PerfCounterType::BlobGetsTotalSize, data.len() as i64);
            }
            Ok(None) => {
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::BlobGetsNotFound);
            }
            _ => {}
        }
        result
    }

    async fn is_present_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
        scuba: &Scuba,
    ) -> Result<BlobstoreIsPresent> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::BlobPresenceChecks);

        let mut futs = inner_multi_is_present(
            ctx,
            self.blobstores.clone(),
            key,
            scuba,
            self.inflight_ops_counter.clone(),
        );

        // Wait for the quorum successful "Not Found" reads before
        // returning Ok(None).
        let mut quorum: usize = self.quorum.read.get();
        let mut errors = HashMap::with_capacity(futs.len());
        let (stats, result) = async move {
            while let Some(result) = futs.next().await {
                match result {
                    (_, Ok(BlobstoreIsPresent::Present)) => {
                        return Ok(BlobstoreIsPresent::Present);
                    }
                    (_, Ok(BlobstoreIsPresent::Absent)) => {
                        quorum = quorum.saturating_sub(1);
                        // we return if there is either quorum on missing
                        if quorum == 0 {
                            return Ok(BlobstoreIsPresent::Absent);
                        }
                    }
                    (bs_id, Ok(BlobstoreIsPresent::ProbablyNotPresent(err))) => {
                        // Treat this like an error from the underlying blobstore.
                        // In reality, this won't happen as multiplexed operates over sinle
                        // standard blobstores, which always can answer if the blob is present.
                        errors.insert(bs_id, err);
                    }
                    (bs_id, Err(err)) => {
                        errors.insert(bs_id, err);
                    }
                }
            }
            Err(errors)
        }
        .timed()
        .await;

        ctx.perf_counters().set_max_counter(
            PerfCounterType::BlobPresenceChecksMaxLatency,
            stats.completion_time.as_millis_unchecked() as i64,
        );

        let errors = match result {
            Ok(is_present) => {
                return Ok(is_present);
            }
            Err(errs) => errs,
        };

        // At this point the multiplexed is_present either failed or cannot say for sure
        // if the blob is present:
        // - no blob was found, but some of the blobstore `is_present` calls failed
        // - there was no read quorum on "not found" result
        let errors = Arc::new(errors);
        if errors.len() == self.blobstores.len() {
            // all main reads failed -> is_present failed
            return Err(ErrorKind::AllFailed(errors).into());
        }

        Ok(BlobstoreIsPresent::ProbablyNotPresent(
            ErrorKind::SomeIsPresentsFailed(errors).into(),
        ))
    }
}

#[async_trait]
impl Blobstore for WalMultiplexedBlobstore {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let mut scuba = self.scuba.clone();
        scuba.sampled();
        let (stats, result) = self.get_impl(ctx, key, &scuba).timed().await;
        scuba::record_get(
            ctx,
            &mut scuba.multiplex_scuba,
            &self.multiplex_id,
            key,
            stats,
            &result,
        );
        result
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let mut scuba = self.scuba.clone();
        scuba.sampled();
        let (stats, result) = self.is_present_impl(ctx, key, &scuba).timed().await;
        scuba::record_is_present(
            ctx,
            &mut scuba.multiplex_scuba,
            &self.multiplex_id,
            key,
            stats,
            &result,
        );
        result
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        BlobstorePutOps::put_with_status(self, ctx, key, value).await?;
        Ok(())
    }
}

#[async_trait]
impl BlobstorePutOps for WalMultiplexedBlobstore {
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        let size = value.len();
        let (stats, result) = self
            .put_impl(ctx, key.clone(), value, Some(put_behaviour), &self.scuba)
            .timed()
            .await;
        scuba::record_put(
            ctx,
            &mut self.scuba.multiplex_scuba.clone(),
            &self.multiplex_id,
            &key,
            size,
            stats,
            &result,
        );
        result
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        let size = value.len();
        let (stats, result) = self
            .put_impl(ctx, key.clone(), value, None, &self.scuba)
            .timed()
            .await;
        scuba::record_put(
            ctx,
            &mut self.scuba.multiplex_scuba.clone(),
            &self.multiplex_id,
            &key,
            size,
            stats,
            &result,
        );
        result
    }
}

fn spawn_stream_completion<T>(
    s: impl Stream<Item = Result<T>> + Send + 'static,
) -> JoinHandle<Result<()>> {
    tokio::spawn(s.try_for_each(|_| future::ok(())))
}

fn inner_multi_put(
    ctx: &CoreContext,
    blobstores: Arc<[TimedStore]>,
    key: &str,
    value: &BlobstoreBytes,
    put_behaviour: Option<PutBehaviour>,
    scuba: &Scuba,
    counter: Arc<AtomicU64>,
) -> FuturesUnordered<impl Future<Output = Result<OverwriteStatus, (BlobstoreId, Error)>>> {
    let put_futs: FuturesUnordered<_> = blobstores
        .iter()
        .map(|bs| {
            // Note: the key used to be passed in as a `&String` and `cloned!` and that triggered
            // a clippy crash for some reason. See D44027145 for context.
            let key = key.to_string();
            cloned!(
                bs,
                ctx,
                value,
                put_behaviour,
                scuba.inner_blobstores_scuba,
                counter
            );
            async move {
                counter.fetch_add(1, Ordering::Relaxed);
                let result = bs
                    .put(&ctx, key, value, put_behaviour, inner_blobstores_scuba)
                    .await;
                counter.fetch_sub(1, Ordering::Relaxed);
                result
            }
        })
        .collect();
    put_futs
}

pub(crate) type GetResult = (BlobstoreId, Result<Option<BlobstoreGetData>, Error>);

pub(crate) fn inner_multi_get<'a>(
    ctx: &'a CoreContext,
    blobstores: Arc<[TimedStore]>,
    key: &'a str,
    operation: OperationType,
    scuba: &Scuba,
    counter: Arc<AtomicU64>,
) -> FuturesUnordered<impl Future<Output = GetResult> + 'a> {
    let get_futs: FuturesUnordered<_> = blobstores
        .iter()
        .map(|bs| {
            cloned!(bs, scuba.inner_blobstores_scuba, counter);
            async move {
                (*bs.id(), {
                    counter.fetch_add(1, Ordering::Relaxed);
                    let result = bs.get(ctx, key, operation, inner_blobstores_scuba).await;
                    counter.fetch_sub(1, Ordering::Relaxed);
                    result
                })
            }
        })
        .collect();
    get_futs
}

fn inner_multi_is_present<'a>(
    ctx: &'a CoreContext,
    blobstores: Arc<[TimedStore]>,
    key: &'a str,
    scuba: &Scuba,
    counter: Arc<AtomicU64>,
) -> FuturesUnordered<impl Future<Output = (BlobstoreId, Result<BlobstoreIsPresent, Error>)> + 'a> {
    let futs: FuturesUnordered<_> = blobstores
        .iter()
        .map(|bs| {
            cloned!(bs, scuba.inner_blobstores_scuba, counter);
            async move {
                counter.fetch_add(1, Ordering::Relaxed);
                let result = bs.is_present(ctx, key, inner_blobstores_scuba).await;
                counter.fetch_sub(1, Ordering::Relaxed);
                result
            }
        })
        .collect();
    futs
}


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

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

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_stats::OperationType;
use blobstore_sync_queue::BlobstoreWal;
use context::CoreContext;
use futures::stream::StreamExt;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use multiplexedblob::base::ErrorKind;
use multiplexedblob::ScrubHandler;
use multiplexedblob::ScrubOptions;
use multiplexedblob::SrubWriteOnly;

use crate::multiplex;
use crate::MultiplexTimeout;
use crate::Scuba;
use crate::WalMultiplexedBlobstore;

impl WalMultiplexedBlobstore {
    async fn scrub_get(
        &self,
        ctx: &CoreContext,
        key: &str,
        write_only: SrubWriteOnly,
    ) -> Result<Option<BlobstoreGetData>, ErrorKind> {
        let mut scuba = self.scuba.clone();
        scuba.sampled();

        let results = multiplexedblob::base::scrub_get_results(
            || {
                multiplex::inner_multi_get(
                    ctx,
                    self.blobstores.clone(),
                    key,
                    OperationType::ScrubGet,
                    &scuba,
                    self.inflight_ops_counter.clone(),
                )
                .collect::<Vec<_>>()
            },
            || {
                multiplex::inner_multi_get(
                    ctx,
                    self.write_only_blobstores.clone(),
                    key,
                    OperationType::ScrubGet,
                    &scuba,
                    self.inflight_ops_counter.clone(),
                )
                .collect::<Vec<_>>()
            },
            self.write_only_blobstores.iter().map(|b| *b.id()),
            write_only,
        )
        .await;

        multiplexedblob::base::scrub_parse_results(results, self.blobstores.iter().map(|b| *b.id()))
    }
}

#[derive(Clone, Debug)]
pub struct WalScrubBlobstore {
    inner: WalMultiplexedBlobstore,
    all_blobstores: Arc<HashMap<BlobstoreId, Arc<dyn BlobstorePutOps>>>,
    scrub_options: ScrubOptions,
    scrub_handler: Arc<dyn ScrubHandler>,
}

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

impl WalScrubBlobstore {
    pub fn new(
        multiplex_id: MultiplexId,
        wal_queue: Arc<dyn BlobstoreWal>,
        blobstores: Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
        write_only_blobstores: Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
        write_quorum: usize,
        timeout: Option<MultiplexTimeout>,
        scuba: Scuba,
        scrub_options: ScrubOptions,
        scrub_handler: Arc<dyn ScrubHandler>,
    ) -> Result<Self> {
        let all_blobstores = Arc::new(
            blobstores
                .iter()
                .cloned()
                .chain(write_only_blobstores.iter().cloned())
                .collect(),
        );
        let inner = WalMultiplexedBlobstore::new(
            multiplex_id,
            wal_queue,
            blobstores,
            write_only_blobstores,
            write_quorum,
            timeout,
            scuba,
        )?;
        Ok(Self {
            inner,
            all_blobstores,
            scrub_options,
            scrub_handler,
        })
    }
}

#[async_trait]
impl Blobstore for WalScrubBlobstore {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let write_only = self.scrub_options.scrub_action_on_missing_write_only;
        match self.inner.scrub_get(ctx, key, write_only).await {
            Ok(value) => Ok(value),
            Err(ErrorKind::SomeFailedOthersNone {
                main_errors,
                write_only_errors,
            }) => {
                if self.inner.blobstores.len() - main_errors.len() < self.inner.quorum.read.get() {
                    Err(ErrorKind::SomeFailedOthersNone {
                        main_errors,
                        write_only_errors,
                    })
                    .context("Can't tell if blob exists or not due to failing blobstores")
                } else {
                    // We have a quorum of blobstores saying the key doesn't exist. That's enough.
                    Ok(None)
                }
            }
            Err(ErrorKind::SomeMissingItem {
                missing_main,
                missing_write_only,
                value,
            }) => {
                multiplexedblob::scrub::maybe_repair(
                    ctx,
                    key,
                    value,
                    missing_main,
                    missing_write_only,
                    self.all_blobstores.as_ref(),
                    self.scrub_handler.as_ref(),
                    &self.scrub_options,
                    &self.inner.scuba.inner_blobstores_scuba,
                    // On WAL we never look into queue except on healer
                    || futures::future::ok(true),
                )
                .await
                .with_context(|| anyhow!("While repairing blobstore key {}", key))
            }
            Err(err) => Err(err.into()),
        }
    }

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

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.inner.put(ctx, key, value).await
    }
}

#[async_trait]
impl BlobstorePutOps for WalScrubBlobstore {
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        self.inner
            .put_explicit(ctx, key, value, put_behaviour)
            .await
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.inner.put_with_status(ctx, key, value).await
    }
}