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

#![feature(never_type)]

mod cachelib_utils;
mod factory;
mod memcache_utils;
mod mock_store;

use std::borrow::Borrow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;
use std::time::Duration;

use abomonation::Abomonation;
use anyhow::Context as _;
use anyhow::Error;
use async_trait::async_trait;
use bytes::Bytes;
use cloned::cloned;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use itertools::Itertools;
use memcache::KeyGen;
use memcache::MEMCACHE_VALUE_MAX_SIZE;
use stats::prelude::*;

pub use crate::cachelib_utils::CachelibHandler;
pub use crate::factory::CacheHandlerFactory;
pub use crate::memcache_utils::MemcacheHandler;
pub use crate::mock_store::MockStoreStats;

pub mod macro_reexport {
    pub use once_cell;
}

define_stats_struct! {
    CacheStats("mononoke.cache.{}", label: String),

    cachelib_hit: timeseries("cachelib.hit"; Rate, Sum),
    cachelib_miss: timeseries("cachelib.miss"; Rate, Sum),

    memcache_hit: timeseries("memcache.hit"; Rate, Sum),
    memcache_miss: timeseries("memcache.miss"; Rate, Sum),
    memcache_internal_err: timeseries("memcache.internal_err"; Rate, Sum),
    memcache_deserialize_err: timeseries("memcache.deserialize_err"; Rate, Sum),

    origin_hit: timeseries("origin.hit"; Rate, Sum),
    origin_miss: timeseries("origin.miss"; Rate, Sum),
}

#[macro_export]
macro_rules! impl_singleton_stats {
    ( $name:literal ) => {
        fn stats(&self) -> &$crate::CacheStats {
            use $crate::macro_reexport::once_cell::sync::Lazy;
            static STATS: Lazy<$crate::CacheStats> =
                Lazy::new(|| $crate::CacheStats::new(String::from($name)));
            &*STATS
        }
    };
}

/// Error type to help with proper reporting of memcache errors
pub enum McErrorKind {
    /// error came from calling memcache API
    MemcacheInternal,
    /// value returned from memcache was None
    Missing,
    /// deserialization of memcache data to Rust structures failed
    Deserialization,
}

const MEMCACHE_CONCURRENCY: usize = 100;

pub type McResult<T> = Result<T, McErrorKind>;

struct CachelibKey(String);
struct MemcacheKey(String);

/// TTL for caching an item
#[derive(Copy, Clone)]
pub enum CacheTtl {
    /// The item is valid forever, and can be cached indefinitely
    NoTtl,
    /// Fetch from backing store once the duration given expires
    Ttl(Duration),
}

/// Whether or not to cache an item
#[derive(Copy, Clone)]
pub enum CacheDisposition {
    /// Cache this item with the given TTL
    Cache(CacheTtl),
    /// Do not cache this item; re-fetch from backing store if it's requested again
    Ignore,
}

/// Implement this for a data item that can be cached. You will also need
/// #[derive(Abomonation)] on the data item.
pub trait MemcacheEntity: Sized {
    /// Convert the item to bytes that can live in Memcache and be deserialized
    /// in another process
    fn serialize(&self) -> Bytes;

    /// Deserialize the item from bytes into an object, or fail to do so
    fn deserialize(bytes: Bytes) -> McResult<Self>;
}

/// Implement this trait to indicate that you can cache values retrived through you
pub trait EntityStore<V> {
    /// Get the cachelib handler. This can be created with `.into()` on a `VolatileLruCachePool`
    fn cachelib(&self) -> &CachelibHandler<V>;

    /// Get the Memcache KeyGen, for creating Memcache keys. This has both code and site versions,
    /// as well as a prefix.
    fn keygen(&self) -> &KeyGen;

    /// Get the Memcache handler. This can be created with `into()` on a `MemcacheClient`.
    fn memcache(&self) -> &MemcacheHandler;

    /// Given a value `v`, decide whether or not to cache it.
    fn cache_determinator(&self, v: &V) -> CacheDisposition;

    /// Finds the cache stats for this handler
    ///
    /// Implement this method with `caching_ext::impl_singleton_stats!` macro, instead of by hand
    fn stats(&self) -> &CacheStats;
}

/// Implement this to make it possible to fetch keys via the cache
#[async_trait]
pub trait KeyedEntityStore<K, V>: EntityStore<V> {
    /// Given an item key, return the cachelib key to use.
    fn get_cache_key(&self, key: &K) -> String;

    /// Given a set of keys to fetch from backing store, return a map from keys to fetched values
    ///
    /// If a key has no value in the backing store, omit it from the result map. Only use an
    /// Error for a failure to fetch, not absence
    async fn get_from_db(&self, keys: HashSet<K>) -> Result<HashMap<K, V>, Error>;
}

/// Utility function to fetch all keys in a single chunk without parallelism
pub async fn get_or_fill<K, V>(
    store: &impl KeyedEntityStore<K, V>,
    keys: HashSet<K>,
) -> Result<HashMap<K, V>, Error>
where
    K: Hash + Eq + Clone,
    // TODO: We should relax the bounds on cachelib's set_cached. We don't need all of this:
    V: Abomonation + MemcacheEntity + Send + Clone + 'static,
{
    get_or_fill_chunked(store, keys, usize::MAX, 1).await
}

/// The core of caching with this module. Takes a store that implements
/// `KeyedEntityStore`, and a set of keys to fetch. Returns a map
/// of fetched values.
///
/// Your accessor functions for consumers should call this to get values
/// from cache or backing store, as this will do the job of keeping
/// cachelib filled from memcache, and memcache filled from your backing store
///
/// fetch_chunk and parallel_chunks are used to implement chunked
/// and parallel fetching. Keys to fetch from the backing store
/// will be split into `fetch_chunk` size groups, and at most `parallel_chunks`
/// groups will be in flight at once.
pub async fn get_or_fill_chunked<K, V>(
    store: &impl KeyedEntityStore<K, V>,
    keys: HashSet<K>,
    fetch_chunk: usize,
    parallel_chunks: usize,
) -> Result<HashMap<K, V>, Error>
where
    K: Hash + Eq + Clone,
    // TODO: We should relax the bounds on cachelib's set_cached. We don't need all of this:
    V: Abomonation + MemcacheEntity + Send + Clone + 'static,
{
    let mut ret = HashMap::<K, V>::with_capacity(keys.len());

    let stats = store.stats();

    let cachelib_keys: Vec<_> = keys
        .into_iter()
        .map(|key| {
            let cachelib_key = CachelibKey(store.get_cache_key(&key));
            (key, cachelib_key)
        })
        .collect();

    let (fetched_from_cachelib, to_fetch_from_memcache) = store
        .cachelib()
        .get_multiple_from_cachelib::<K>(cachelib_keys)
        .with_context(|| "Error reading from cachelib")?;

    stats
        .cachelib_hit
        .add_value(fetched_from_cachelib.len() as i64);
    stats
        .cachelib_miss
        .add_value(to_fetch_from_memcache.len() as i64);

    ret.extend(fetched_from_cachelib);

    let to_fetch_from_memcache: Vec<(K, CachelibKey, MemcacheKey)> = to_fetch_from_memcache
        .into_iter()
        .map(|(key, cachelib_key)| {
            let memcache_key = MemcacheKey(store.keygen().key(&cachelib_key.0));
            (key, cachelib_key, memcache_key)
        })
        .collect();

    let to_fetch_from_store = {
        let (fetched_from_memcache, to_fetch_from_store) =
            get_multiple_from_memcache(store.memcache(), to_fetch_from_memcache, stats).await;

        stats
            .memcache_hit
            .add_value(fetched_from_memcache.len() as i64);
        stats
            .memcache_miss
            .add_value(to_fetch_from_store.len() as i64);

        fill_multiple_cachelib(
            store.cachelib(),
            fetched_from_memcache
                .values()
                .filter_map(|(v, k)| match store.cache_determinator(v) {
                    CacheDisposition::Cache(ttl) => Some((k, ttl, v)),
                    _ => None,
                }),
        );

        ret.extend(fetched_from_memcache.into_iter().map(|(k, (v, _))| (k, v)));

        to_fetch_from_store
    };

    if !to_fetch_from_store.is_empty() {
        let to_fetch_from_store: Vec<_> = to_fetch_from_store
            .into_iter()
            .chunks(fetch_chunk)
            .into_iter()
            .map(|chunk| {
                let mut keys = HashSet::new();
                let mut key_mapping = HashMap::new();
                for (key, cachelib_key, memcache_key) in chunk {
                    keys.insert(key.clone());
                    key_mapping.insert(key.clone(), (cachelib_key, memcache_key));
                }
                fill_one_chunk(store, keys, key_mapping)
            })
            .collect();
        stream::iter(to_fetch_from_store)
            .buffer_unordered(parallel_chunks)
            .try_fold(&mut ret, |ret, chunk| async move {
                ret.extend(chunk);
                Ok::<_, Error>(ret)
            })
            .await?;
    }

    Ok(ret)
}

async fn fill_one_chunk<K, V>(
    store: &impl KeyedEntityStore<K, V>,
    keys: HashSet<K>,
    mut key_mapping: HashMap<K, (CachelibKey, MemcacheKey)>,
) -> Result<HashMap<K, V>, Error>
where
    K: Hash + Eq + Clone,
    // TODO: We should relax the bounds on cachelib's set_cached. We don't need all of this:
    V: Abomonation + MemcacheEntity + Send + Clone + 'static,
{
    let n_keys = keys.len();

    let stats = store.stats();
    let data = store
        .get_from_db(keys)
        .await
        .with_context(|| "Error reading from store")?;

    stats.origin_hit.add_value(data.len() as i64);
    stats.origin_miss.add_value((n_keys - data.len()) as i64);

    fill_caches_by_key(
        store,
        data.iter().map(|(key, v)| {
            let (cachelib_key, memcache_key) = key_mapping
                .remove(key)
                .expect("caching_ext: Missing entry in key_mapping, this should not happen");

            (cachelib_key, memcache_key, v)
        }),
    )
    .await;
    Ok(data)
}

/// Directly fill a cache from data you've prefetched outside the caching system
/// Allows things like microwave to avoid any backing store fetches
pub async fn fill_cache<'a, K, V>(
    store: &impl KeyedEntityStore<K, V>,
    data: impl IntoIterator<Item = (&'a K, &'a V)>,
) where
    K: Hash + Eq + Clone + 'a,
    V: Abomonation + MemcacheEntity + Send + Clone + 'static,
{
    fill_caches_by_key(
        store,
        data.into_iter().map(|(k, v)| {
            let cachelib_key = CachelibKey(store.get_cache_key(k));
            let memcache_key = MemcacheKey(store.keygen().key(&cachelib_key.0));
            (cachelib_key, memcache_key, v)
        }),
    )
    .await;
}

async fn fill_caches_by_key<'a, V>(
    store: &impl EntityStore<V>,
    data: impl IntoIterator<Item = (CachelibKey, MemcacheKey, &'a V)>,
) where
    V: Abomonation + MemcacheEntity + Send + Clone + 'static,
{
    let mut cachelib_keys = Vec::new();
    let mut memcache_keys = Vec::new();

    for (cachelib_key, memcache_key, v) in data.into_iter() {
        let ttl = match store.cache_determinator(v) {
            CacheDisposition::Cache(ttl) => ttl,
            CacheDisposition::Ignore => continue,
        };

        memcache_keys.push((memcache_key, ttl, v));
        cachelib_keys.push((cachelib_key, ttl, v));
    }

    fill_multiple_cachelib(store.cachelib(), cachelib_keys);

    fill_multiple_memcache(store.memcache(), memcache_keys).await;
}

async fn get_multiple_from_memcache<K, V>(
    memcache: &MemcacheHandler,
    keys: Vec<(K, CachelibKey, MemcacheKey)>,
    stats: &CacheStats,
) -> (
    HashMap<K, (V, CachelibKey)>,
    Vec<(K, CachelibKey, MemcacheKey)>,
)
where
    K: Eq + Hash,
    V: MemcacheEntity,
{
    let mc_fetch_futs = keys
        .into_iter()
        .map(move |(key, cachelib_key, memcache_key)| {
            cloned!(memcache);
            async move {
                let res = memcache
                    .get(memcache_key.0.clone())
                    .await
                    .map_err(|_| McErrorKind::MemcacheInternal)
                    .and_then(|maybe_bytes| maybe_bytes.ok_or(McErrorKind::Missing))
                    .and_then(V::deserialize);

                (key, cachelib_key, memcache_key, res)
            }
        });

    let mut entries = stream::iter(mc_fetch_futs).buffered(MEMCACHE_CONCURRENCY);

    let mut fetched = HashMap::new();
    let mut left_to_fetch = Vec::new();

    while let Some((key, cachelib_key, memcache_key, res)) = entries.next().await {
        match res {
            Ok(entity) => {
                fetched.insert(key, (entity, cachelib_key));
            }
            Err(e) => {
                match e {
                    McErrorKind::MemcacheInternal => stats.memcache_internal_err.add_value(1),
                    McErrorKind::Deserialization => stats.memcache_deserialize_err.add_value(1),
                    McErrorKind::Missing => {} // no op, we record missing at a higher level anyway.
                };

                left_to_fetch.push((key, cachelib_key, memcache_key));
            }
        }
    }

    (fetched, left_to_fetch)
}

fn fill_multiple_cachelib<'a, V>(
    cachelib: &'a CachelibHandler<V>,
    data: impl IntoIterator<Item = (impl Borrow<CachelibKey> + 'a, CacheTtl, &'a V)>,
) where
    V: Abomonation + Clone + Send + 'static,
{
    for (cachelib_key, ttl, v) in data {
        let cachelib_key = cachelib_key.borrow();

        let ttl = match ttl {
            CacheTtl::NoTtl => None,
            CacheTtl::Ttl(ttl) => Some(ttl),
        };

        // NOTE: We ignore failures to cache individual entries here.
        let _ = cachelib.set_cached(&cachelib_key.0, v, ttl);
    }
}

async fn fill_multiple_memcache<'a, V: 'a>(
    memcache: &'a MemcacheHandler,
    data: impl IntoIterator<Item = (MemcacheKey, CacheTtl, &'a V)>,
) where
    V: MemcacheEntity,
{
    if memcache.is_noop() {
        return;
    }

    let futs = data
        .into_iter()
        .filter_map(|(memcache_key, ttl, v)| {
            let bytes = v.serialize();

            if bytes.len() >= MEMCACHE_VALUE_MAX_SIZE {
                return None;
            }

            cloned!(memcache);

            Some(async move {
                match ttl {
                    CacheTtl::NoTtl => {
                        let _ = memcache.set(memcache_key.0, bytes).await;
                    }
                    CacheTtl::Ttl(ttl) => {
                        let _ = memcache.set_with_ttl(memcache_key.0, bytes, ttl).await;
                    }
                }
            })
        })
        .collect::<Vec<_>>();

    let fut = stream::iter(futs).for_each_concurrent(MEMCACHE_CONCURRENCY, |fut| fut);

    if memcache.is_async() {
        tokio::task::spawn(fut);
    } else {
        fut.await;
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering;

    use abomonation_derive::Abomonation;
    use maplit::hashmap;
    use maplit::hashset;

    use super::*;

    #[derive(Abomonation, Clone, Debug, PartialEq, Eq)]
    struct TestEntity(Vec<u8>);

    impl MemcacheEntity for TestEntity {
        fn serialize(&self) -> Bytes {
            Bytes::from(self.0.clone())
        }

        fn deserialize(bytes: Bytes) -> McResult<Self> {
            Ok(Self(bytes.to_vec()))
        }
    }

    struct TestStore {
        keygen: KeyGen,
        cachelib: CachelibHandler<TestEntity>,
        memcache: MemcacheHandler,
        calls: AtomicUsize,
        keys: AtomicUsize,
        data: HashMap<String, TestEntity>,
    }

    impl TestStore {
        pub fn new() -> Self {
            Self {
                keygen: KeyGen::new("", 0, 0),
                cachelib: CachelibHandler::create_mock(),
                memcache: MemcacheHandler::create_mock(),
                calls: AtomicUsize::new(0),
                keys: AtomicUsize::new(0),
                data: HashMap::new(),
            }
        }
    }

    impl EntityStore<TestEntity> for TestStore {
        fn cachelib(&self) -> &CachelibHandler<TestEntity> {
            &self.cachelib
        }

        fn keygen(&self) -> &KeyGen {
            &self.keygen
        }

        fn memcache(&self) -> &MemcacheHandler {
            &self.memcache
        }

        fn cache_determinator(&self, _: &TestEntity) -> CacheDisposition {
            CacheDisposition::Cache(CacheTtl::NoTtl)
        }

        impl_singleton_stats!("test");
    }

    #[async_trait]
    impl KeyedEntityStore<String, TestEntity> for TestStore {
        fn get_cache_key(&self, key: &String) -> String {
            format!("key:{}", key)
        }

        async fn get_from_db(
            &self,
            keys: HashSet<String>,
        ) -> Result<HashMap<String, TestEntity>, Error> {
            self.calls.fetch_add(1, Ordering::Relaxed);
            self.keys.fetch_add(keys.len(), Ordering::Relaxed);

            Ok(keys
                .into_iter()
                .filter_map(|k| {
                    let v = self.data.get(&k).cloned();
                    v.map(|v| (k, v))
                })
                .collect())
        }
    }

    #[tokio::test]
    async fn simple() -> Result<(), Error> {
        let store = TestStore::new();

        let res = get_or_fill(&store, hashset! {}).await?;
        assert_eq!(res.len(), 0);
        assert_eq!(store.cachelib.gets_count(), 0);
        assert_eq!(store.memcache.gets_count(), 0);

        let res = get_or_fill(&store, hashset! {"key".into()}).await?;
        assert_eq!(res.len(), 0);
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 1);
        assert_eq!(store.keys.load(Ordering::Relaxed), 1);

        Ok(())
    }

    #[tokio::test]
    async fn fetch_from_db_cachelib_memcache() -> Result<(), Error> {
        let mut store = TestStore::new();

        let e = TestEntity(vec![0]);
        store.data.insert("key".into(), e.clone());

        // Fetch from db
        let res = get_or_fill(&store, hashset! {"key".into()}).await?;
        assert_eq!(res, hashmap! { "key".into() => e.clone() });
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 1);
        assert_eq!(store.keys.load(Ordering::Relaxed), 1);

        // Now fetch from cachelib
        let res = get_or_fill(&store, hashset! {"key".into()}).await?;
        assert_eq!(res, hashmap! { "key".into() => e.clone() });
        assert_eq!(store.cachelib.gets_count(), 2);
        assert_eq!(store.memcache.gets_count(), 1);
        assert_eq!(store.keys.load(Ordering::Relaxed), 1);

        // Reset cachelib, fetch from memcache
        store.cachelib = CachelibHandler::create_mock();
        let res = get_or_fill(&store, hashset! {"key".into()}).await?;
        assert_eq!(res, hashmap! { "key".into() => e.clone() });
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 2);
        assert_eq!(store.keys.load(Ordering::Relaxed), 1);

        Ok(())
    }

    #[tokio::test]
    async fn fetch_from_db() -> Result<(), Error> {
        let mut store = TestStore::new();

        let e0 = TestEntity(vec![0]);
        let e1 = TestEntity(vec![1]);
        let e2 = TestEntity(vec![2]);

        store.data.insert("key0".into(), e0.clone());
        store.data.insert("key1".into(), e1.clone());
        store.data.insert("key2".into(), e2.clone());

        let res = get_or_fill(
            &store,
            hashset! { "key0".into(), "key1".into(), "key2".into() },
        )
        .await?;

        assert_eq!(
            res,
            hashmap! { "key0".into() => e0, "key1".into() => e1, "key2".into() => e2 }
        );
        assert_eq!(store.calls.load(Ordering::Relaxed), 1);
        assert_eq!(store.cachelib.gets_count(), 3);
        assert_eq!(store.memcache.gets_count(), 3);
        assert_eq!(store.keys.load(Ordering::Relaxed), 3);

        Ok(())
    }

    #[tokio::test]
    async fn fetch_from_db_chunked() -> Result<(), Error> {
        let mut store = TestStore::new();

        let e0 = TestEntity(vec![0]);
        let e1 = TestEntity(vec![1]);
        let e2 = TestEntity(vec![2]);

        store.data.insert("key0".into(), e0.clone());
        store.data.insert("key1".into(), e1.clone());
        store.data.insert("key2".into(), e2.clone());

        let res = get_or_fill_chunked(
            &store,
            hashset! { "key0".into(), "key1".into(), "key2".into() },
            1,
            3,
        )
        .await?;

        assert_eq!(
            res,
            hashmap! { "key0".into() => e0, "key1".into() => e1, "key2".into() => e2 }
        );
        assert_eq!(store.calls.load(Ordering::Relaxed), 3);
        assert_eq!(store.cachelib.gets_count(), 3);
        assert_eq!(store.memcache.gets_count(), 3);
        assert_eq!(store.keys.load(Ordering::Relaxed), 3);

        Ok(())
    }

    #[tokio::test]
    async fn fetch_from_all() -> Result<(), Error> {
        let mut store = TestStore::new();

        let e0 = TestEntity(vec![0]);
        let e1 = TestEntity(vec![1]);
        let e2 = TestEntity(vec![2]);

        store.data.insert("key0".into(), e0.clone());
        store.data.insert("key1".into(), e1.clone());
        store.data.insert("key2".into(), e2.clone());

        let res = get_or_fill(&store, hashset! { "key1".into() }).await?;
        assert_eq!(res, hashmap! { "key1".into() => e1.clone() });
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 1);
        assert_eq!(store.calls.load(Ordering::Relaxed), 1);

        // Reset cachelib
        store.cachelib = CachelibHandler::create_mock();
        let res = get_or_fill(&store, hashset! { "key0".into() }).await?;
        assert_eq!(res, hashmap! { "key0".into() => e0.clone() });
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 2);
        assert_eq!(store.calls.load(Ordering::Relaxed), 2);

        let res = get_or_fill(
            &store,
            hashset! { "key0".into(), "key1".into(), "key2".into() },
        )
        .await?;

        assert_eq!(
            res,
            hashmap! { "key0".into() => e0.clone(), "key1".into() => e1.clone(), "key2".into() => e2.clone() }
        );
        assert_eq!(store.cachelib.gets_count(), 1 + 3); // 3 new fetches from cachelib, 2 misses
        assert_eq!(store.memcache.gets_count(), 2 + 2); // 2 new fetches from memcache, 1 miss
        assert_eq!(store.calls.load(Ordering::Relaxed), 2 + 1); // 1 fetch from db

        // Only from cachelib
        let res = get_or_fill(
            &store,
            hashset! { "key0".into(), "key1".into(), "key2".into() },
        )
        .await?;

        assert_eq!(
            res,
            hashmap! { "key0".into() => e0.clone(), "key1".into() => e1.clone(), "key2".into() => e2.clone() }
        );
        assert_eq!(store.cachelib.gets_count(), 7);
        assert_eq!(store.memcache.gets_count(), 4);
        assert_eq!(store.calls.load(Ordering::Relaxed), 3);

        // // Reset cachelib, only from memcache
        store.cachelib = CachelibHandler::create_mock();
        let res = get_or_fill(
            &store,
            hashset! { "key0".into(), "key1".into(), "key2".into() },
        )
        .await?;

        assert_eq!(
            res,
            hashmap! { "key0".into() => e0.clone(), "key1".into() => e1.clone(), "key2".into() => e2.clone() }
        );
        assert_eq!(store.cachelib.gets_count(), 3); // 3 misses
        assert_eq!(store.memcache.gets_count(), 4 + 3); // 3 hits
        assert_eq!(store.calls.load(Ordering::Relaxed), 3);

        Ok(())
    }

    #[tokio::test]
    async fn get_from_db_elision() -> Result<(), Error> {
        let store = TestStore::new();

        get_or_fill(&store, hashset! {}).await?;
        assert_eq!(store.calls.load(Ordering::Relaxed), 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_fill_cache() -> Result<(), Error> {
        let store = TestStore::new();
        let e0 = TestEntity(vec![0]);
        fill_cache(&store, hashmap! { "key0".into() => e0.clone() }.iter()).await;

        let res = get_or_fill(&store, hashset! { "key0".into() }).await?;
        assert_eq!(res, hashmap! { "key0".into() => e0.clone() });
        assert_eq!(store.cachelib.gets_count(), 1);
        assert_eq!(store.memcache.gets_count(), 0);
        assert_eq!(store.calls.load(Ordering::Relaxed), 0);

        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::Result;
use bytes::Bytes;
use memcache::MemcacheClient;
use memcache::MemcacheSetType;

use crate::mock_store::MockStore;

#[derive(Clone)]
pub enum MemcacheHandler {
    Real(MemcacheClient),
    Mock(MockStore<Bytes>),
    Noop,
}

impl From<MemcacheClient> for MemcacheHandler {
    fn from(client: MemcacheClient) -> Self {
        MemcacheHandler::Real(client)
    }
}

impl MemcacheHandler {
    /// Returns true if this memcache handler is a no-op, and so operations
    /// can be entirely skipped.
    pub fn is_noop(&self) -> bool {
        match self {
            MemcacheHandler::Noop => true,
            MemcacheHandler::Real(_) | MemcacheHandler::Mock(_) => false,
        }
    }

    pub fn is_async(&self) -> bool {
        match self {
            MemcacheHandler::Real(_) => true,
            MemcacheHandler::Mock(_) | MemcacheHandler::Noop => false,
        }
    }

    pub async fn get(&self, key: String) -> Result<Option<Bytes>> {
        match self {
            MemcacheHandler::Real(ref client) => {
                client.get(key).await.map(|value| value.map(Bytes::from))
            }
            MemcacheHandler::Mock(store) => Ok(store.get(&key)),
            MemcacheHandler::Noop => Ok(None),
        }
    }

    pub async fn set<V>(&self, key: String, value: V) -> Result<()>
    where
        MemcacheSetType: From<V>,
        Bytes: From<V>,
        V: 'static,
    {
        match self {
            MemcacheHandler::Real(ref client) => client.set(key, value).await,
            MemcacheHandler::Mock(store) => {
                store.set(&key, value.into());
                Ok(())
            }
            MemcacheHandler::Noop => Ok(()),
        }
    }

    pub async fn set_with_ttl<V>(&self, key: String, value: V, duration: Duration) -> Result<()>
    where
        MemcacheSetType: From<V>,
        Bytes: From<V>,
        V: 'static,
    {
        match self {
            MemcacheHandler::Real(ref client) => client.set_with_ttl(key, value, duration).await,
            MemcacheHandler::Mock(_) => {
                // For now we ignore TTLs here
                self.set(key, value).await
            }
            MemcacheHandler::Noop => Ok(()),
        }
    }

    pub fn create_mock() -> Self {
        MemcacheHandler::Mock(MockStore::new())
    }

    pub fn create_noop() -> Self {
        MemcacheHandler::Noop
    }

    #[cfg(test)]
    pub(crate) fn gets_count(&self) -> usize {
        use std::sync::atomic::Ordering;
        match self {
            MemcacheHandler::Real(_) | MemcacheHandler::Noop => unimplemented!(),
            MemcacheHandler::Mock(MockStore { ref get_count, .. }) => {
                get_count.load(Ordering::SeqCst)
            }
        }
    }
}


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

#[derive(Debug, PartialEq)]
pub struct MockStoreStats {
    pub sets: usize,
    pub gets: usize,
    pub hits: usize,
    pub misses: usize,
}

#[derive(Clone, Debug)]
pub struct MockStore<T> {
    data: Arc<Mutex<HashMap<String, T>>>,
    pub(crate) set_count: Arc<AtomicUsize>,
    pub(crate) get_count: Arc<AtomicUsize>,
    pub(crate) hit_count: Arc<AtomicUsize>,
    pub(crate) miss_count: Arc<AtomicUsize>,
}

impl<T> MockStore<T> {
    pub(crate) fn new() -> Self {
        Self {
            data: Arc::new(Mutex::new(HashMap::new())),
            set_count: Arc::new(AtomicUsize::new(0)),
            get_count: Arc::new(AtomicUsize::new(0)),
            hit_count: Arc::new(AtomicUsize::new(0)),
            miss_count: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn stats(&self) -> MockStoreStats {
        MockStoreStats {
            sets: self.set_count.load(Ordering::SeqCst),
            gets: self.get_count.load(Ordering::SeqCst),
            misses: self.miss_count.load(Ordering::SeqCst),
            hits: self.hit_count.load(Ordering::SeqCst),
        }
    }
}

impl<T: Clone> MockStore<T> {
    pub fn get(&self, key: &String) -> Option<T> {
        self.get_count.fetch_add(1, Ordering::SeqCst);
        let value = self.data.lock().expect("poisoned lock").get(key).cloned();
        match &value {
            Some(..) => self.hit_count.fetch_add(1, Ordering::SeqCst),
            None => self.miss_count.fetch_add(1, Ordering::SeqCst),
        };
        value
    }

    pub fn set(&self, key: &str, value: T) {
        self.set_count.fetch_add(1, Ordering::SeqCst);
        self.data
            .lock()
            .expect("poisoned lock")
            .insert(key.to_owned(), value);
    }

    #[cfg(test)]
    pub(crate) fn data(&self) -> HashMap<String, T> {
        self.data.lock().expect("poisoned lock").clone()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_counts() {
        let store = MockStore::new();
        assert_eq!(
            store.stats(),
            MockStoreStats {
                sets: 0,
                gets: 0,
                misses: 0,
                hits: 0
            }
        );

        store.set("foo", &());
        assert_eq!(
            store.stats(),
            MockStoreStats {
                sets: 1,
                gets: 0,
                misses: 0,
                hits: 0
            }
        );

        store.get(&"foo".to_string());
        assert_eq!(
            store.stats(),
            MockStoreStats {
                sets: 1,
                gets: 1,
                misses: 0,
                hits: 1
            }
        );

        store.get(&"bar".to_string());
        assert_eq!(
            store.stats(),
            MockStoreStats {
                sets: 1,
                gets: 2,
                misses: 1,
                hits: 1
            }
        );
    }
}


/*
 * 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::fmt;
use std::io;
use std::str::FromStr;

use anyhow::bail;
use anyhow::Error;
use failure_ext::SlogKVErrorKey;
use slog::Drain;
use slog::Never;
use slog::OwnedKVList;
use slog::Record;
use slog::Serializer;
use slog::KV;
use slog_term::Decorator;

// Allow us to switch drain types without runtime check
enum EitherDrain<L, R> {
    Left(L),
    Right(R),
}

impl<O, E, L: Drain<Ok = O, Err = E>, R: Drain<Ok = O, Err = E>> Drain for EitherDrain<L, R> {
    type Ok = O;
    type Err = E;

    fn log(&self, record: &Record, values: &OwnedKVList) -> Result<Self::Ok, Self::Err> {
        match self {
            EitherDrain::Left(d) => d.log(record, values),
            EitherDrain::Right(d) => d.log(record, values),
        }
    }
}

// Filter in or out log messages based on the slog::Record::tag() (which in turn is mapped from Log::Record::target())
struct TagFilterDrain<D> {
    inner: D,
    include_tags: HashSet<String>,
    exclude_tags: HashSet<String>,
    pass_untagged: bool,
}

impl<D> TagFilterDrain<D> {
    fn should_log(&self, tag: &str) -> bool {
        if tag.is_empty() {
            return self.pass_untagged;
        }
        if self.exclude_tags.contains(tag) {
            return false;
        }
        if self.include_tags.is_empty() {
            return true;
        }
        self.include_tags.contains(tag)
    }
}

impl<D: Drain<Ok = (), Err = Never>> Drain for TagFilterDrain<D> {
    type Ok = ();
    type Err = Never;

    fn log(&self, record: &Record, values: &OwnedKVList) -> Result<Self::Ok, Self::Err> {
        if self.should_log(record.tag()) {
            self.inner.log(record, values)
        } else {
            Ok(())
        }
    }
}

pub fn make_tag_filter_drain<'a, D>(
    inner: D,
    include_tags: HashSet<String>,
    exclude_tags: HashSet<String>,
    pass_untagged: bool,
) -> Result<impl Drain<Ok = (), Err = Never> + 'a, Error>
where
    D: Drain<Ok = (), Err = Never> + 'a,
{
    if include_tags.is_empty() && exclude_tags.is_empty() {
        Ok(EitherDrain::Left(inner))
    } else {
        let intersection = include_tags
            .intersection(&exclude_tags)
            .collect::<HashSet<_>>();
        if !intersection.is_empty() {
            bail!(
                "Following tags are in both the include and exclude sets: {:?}",
                intersection
            );
        } else {
            Ok(EitherDrain::Right(TagFilterDrain {
                inner,
                include_tags,
                exclude_tags,
                pass_untagged,
            }))
        }
    }
}

/// Drain that only prints the message and newline plus error if present, nothing more
pub struct SimpleFormatWithError<D: Decorator> {
    decorator: D,
}

impl<D: Decorator> SimpleFormatWithError<D> {
    pub fn new(decorator: D) -> Self {
        Self { decorator }
    }
}

impl<D: Decorator> Drain for SimpleFormatWithError<D> {
    type Ok = ();
    type Err = io::Error;

    fn log(&self, record: &Record, values: &OwnedKVList) -> Result<Self::Ok, Self::Err> {
        self.decorator.with_record(record, values, |decorator| {
            write!(decorator, "{}\n", record.msg())?;

            let mut serializer = ErrorSerializer {
                error: None,
                root_cause: None,
                backtrace: None,
                causes: Vec::new(),
                error_debug: None,
            };
            record.kv().serialize(record, &mut serializer)?;
            values.serialize(record, &mut serializer)?;

            if let Some(error) = serializer.error {
                write!(decorator, "  Error:\n    {}\n", fix_indentation(error))?;

                if let Some(root_cause) = serializer.root_cause {
                    write!(decorator, "\n")?;
                    write!(
                        decorator,
                        "  Root cause:\n    {}\n",
                        fix_indentation(root_cause)
                    )?;
                }

                if let Some(backtrace) = serializer.backtrace {
                    write!(decorator, "\n")?;
                    write!(
                        decorator,
                        "  Backtrace:\n    {}\n",
                        fix_indentation(backtrace)
                    )?;
                }

                for (i, cause) in serializer.causes.into_iter().enumerate() {
                    if i == 0 {
                        write!(decorator, "\n")?;
                    }
                    write!(decorator, "  Caused by:\n    {}\n", fix_indentation(cause))?;
                }

                if let Some(error_debug) = serializer.error_debug {
                    write!(decorator, "\n")?;
                    write!(
                        decorator,
                        "  Debug context:\n    {}\n",
                        fix_indentation(error_debug)
                    )?;
                }
            }

            decorator.flush()?;
            Ok(())
        })
    }
}

struct ErrorSerializer {
    pub error: Option<String>,
    pub root_cause: Option<String>,
    pub backtrace: Option<String>,
    pub causes: Vec<String>,
    pub error_debug: Option<String>,
}

impl Serializer for ErrorSerializer {
    fn emit_arguments(&mut self, key: slog::Key, val: &fmt::Arguments) -> slog::Result {
        if let Ok(key) = SlogKVErrorKey::from_str(key) {
            use SlogKVErrorKey::*;
            match key {
                Error => self.error = non_empty_str_maybe(format!("{}", val)),
                RootCause => self.root_cause = non_empty_str_maybe(format!("{}", val)),
                Backtrace => self.backtrace = non_empty_str_maybe(format!("{}", val)),
                Cause => self.causes.push(format!("{}", val)),
                ErrorDebug => self.error_debug = non_empty_str_maybe(format!("{}", val)),
            }
        }

        Ok(())
    }
}

fn non_empty_str_maybe(s: String) -> Option<String> {
    if s.is_empty() { None } else { Some(s) }
}

fn fix_indentation(s: String) -> String {
    s.replace('\n', "\n    ")
}


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

mod mononoke_queries;
#[cfg(not(fbcode_build))]
mod oss;
pub mod replication;
mod sqlite;

pub use sql::SqlConnections;
pub use sql::SqlShardedConnections;
use sql::Transaction;
pub use sqlite::open_existing_sqlite_path;
pub use sqlite::open_sqlite_in_memory;
pub use sqlite::open_sqlite_path;

#[must_use]
pub enum TransactionResult {
    Succeeded(Transaction),
    Failed,
}

pub mod _macro_internal {
    pub use std::collections::hash_map::DefaultHasher;
    pub use std::hash::Hash;
    pub use std::hash::Hasher;

    pub use anyhow::Result;
    pub use paste;
    pub use sql::queries;
    pub use sql::Connection;
    pub use sql::WriteResult;
    pub use sql_query_config::SqlQueryConfig;
    pub use twox_hash::xxh3::Hash128;
    pub use twox_hash::xxh3::HasherExt;

    pub use crate::mononoke_queries::query_with_retry;
    pub use crate::mononoke_queries::query_with_retry_no_cache;
    pub use crate::mononoke_queries::CacheData;
    pub use crate::mononoke_queries::MemcacheWrapper;
}

pub mod facebook {
    #[cfg(fbcode_build)]
    mod r#impl;

    use std::fmt;
    use std::fmt::Debug;

    #[cfg(fbcode_build)]
    pub use r#impl::create_mysql_connections_sharded;
    #[cfg(fbcode_build)]
    pub use r#impl::create_mysql_connections_unsharded;
    #[cfg(fbcode_build)]
    pub use r#impl::myadmin::replication_status_chunked;
    #[cfg(fbcode_build)]
    pub use r#impl::myadmin::MyAdmin;
    #[cfg(fbcode_build)]
    pub use r#impl::myadmin::MyAdminLagMonitor;
    #[cfg(fbcode_build)]
    pub use r#impl::PoolConfig;
    #[cfg(fbcode_build)]
    pub use r#impl::SharedConnectionPool;

    #[cfg(not(fbcode_build))]
    pub use crate::oss::create_mysql_connections_sharded;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::create_mysql_connections_unsharded;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::MyAdmin;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::MyAdminLagMonitor;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::PoolConfig;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::SharedConnectionPool;

    /// MySQL global shared connection pool configuration.
    #[derive(Clone, Default)]
    pub struct MysqlOptions {
        pub pool: SharedConnectionPool,
        // pool config is used only once when the shared connection pool is being created
        pub pool_config: PoolConfig,
        pub read_connection_type: ReadConnectionType,
    }

    impl MysqlOptions {
        pub fn per_key_limit(&self) -> Option<usize> {
            #[cfg(not(fbcode_build))]
            {
                None
            }
            #[cfg(fbcode_build)]
            {
                Some(self.pool_config.per_key_limit as usize)
            }
        }
    }

    impl Debug for MysqlOptions {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(
                f,
                "MySQL pool with config {:?}, connection type: {:?}",
                self.pool_config, self.read_connection_type
            )
        }
    }

    /// Mirrors facebook::db::InstanceRequirement enum for DBLocator
    #[derive(Copy, Clone, Debug, Default)]
    pub enum ReadConnectionType {
        /// Choose master or replica, whatever is closest and available.
        /// Use this if both master and replica are in the same region, and reads
        /// should we served by both.
        Closest,
        /// Choose replicas only, avoiding the master, even if it means going to a
        /// remote region.
        #[default]
        ReplicaOnly,
        /// Choose master only (typically for writes). Will never connect to replica.
        Master,
        /// Choose closer first and inside the same region, replicas first.
        /// In case both master and replica in the same region - all reads
        /// will be routed to the replica.
        ReplicaFirst,
        /// Choose replicas that satisfy a lower bound HLC value in order to
        /// perform consistent read-your-writes operations
        ReadAfterWriteConsistency,
    }
}


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

use abomonation::Abomonation;
use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use caching_ext::*;
use itertools::Itertools;
use maplit::hashmap;
use maplit::hashset;
use memcache::KeyGen;
use retry::retry;
use retry::RetryLogic;
use sql_query_config::CachingConfig;
use tunables::tunables;

const RETRY_ATTEMPTS: usize = 2;

// This wraps around rust/shed/sql::queries, check that macro: https://fburl.com/code/semq9xm3
/// Define SQL queries that automatically retry on certain errors.
///
/// Caching can be enabled on a read query by:
/// - Adding "cacheable" keyword to your query.
/// - Make sure all parameters (input) to the query implement the Hash trait.
/// - Making sure the return values (output) implement Serialize, Deserialize, and Abomonation.
#[macro_export]
macro_rules! mononoke_queries {
    () => {};

    // Read query with a single expression. Redirect to read query with same expression for mysql and sqlite.
    (
        $vi:vis read $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) -> ($( $rtype:ty ),* $(,)*) { $q:expr }
        $( $rest:tt )*
    ) => {
        $crate::mononoke_queries! {
            $vi read $name (
                $( $pname: $ptype, )*
                $( >list $lname: $ltype )*
            ) -> ($( $rtype ),*) { mysql($q) sqlite($q) }
            $( $rest )*
        }
    };
    // Read query with a single expression and cache. Redirect to read query with same expression for mysql and sqlite.
    (
        $vi:vis cacheable read $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) -> ($( $rtype:ty ),* $(,)*) { $q:expr }
        $( $rest:tt )*
    ) => {
        $crate::mononoke_queries! {
            $vi cacheable read $name (
                $( $pname: $ptype, )*
                $( >list $lname: $ltype )*
            ) -> ($( $rtype ),*) { mysql($q) sqlite($q) }
            $( $rest )*
        }
    };

    // Full read query without cache. Call `sql::queries!` and re-export stuff, wrapped in retries, on a new module.
    (
        $vi:vis read $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) -> ($( $rtype:ty ),* $(,)*) { mysql($mysql_q:expr) sqlite($sqlite_q:expr) }
        $( $rest:tt )*
    ) => {
        $crate::_macro_internal::paste::item! {
            $crate::_macro_internal::queries! {
                pub read [<$name Impl>] (
                    $( $pname: $ptype, )*
                    $( >list $lname: $ltype )*
                ) -> ($( $rtype ),*) { mysql($mysql_q) sqlite($sqlite_q) }
            }

            #[allow(non_snake_case)]
            $vi mod $name {
                #[allow(unused_imports)]
                use super::*;

                #[allow(unused_imports)]
                use $crate::_macro_internal::*;

                // Not possible to retry query with transaction
                #[allow(unused_imports)]
                pub use [<$name Impl>]::query_with_transaction;

                #[allow(dead_code)]
                pub async fn query(
                    connection: &Connection,
                    $( $pname: & $ptype, )*
                    $( $lname: & [ $ltype ], )*
                ) -> Result<Vec<($( $rtype, )*)>> {
                    query_with_retry_no_cache(
                        || [<$name Impl>]::query(connection, $( $pname, )* $( $lname, )*),
                    ).await
                }
            }

            $crate::mononoke_queries! { $( $rest )* }
        }
    };

    // Full read query. Call `sql::queries!` and re-export stuff, wrapped in retries, on a new module.
    (
        $vi:vis cacheable read $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) -> ($( $rtype:ty ),* $(,)*) { mysql($mysql_q:expr) sqlite($sqlite_q:expr) }
        $( $rest:tt )*
    ) => {
        $crate::_macro_internal::paste::item! {
            $crate::_macro_internal::queries! {
                pub read [<$name Impl>] (
                    $( $pname: $ptype, )*
                    $( >list $lname: $ltype )*
                ) -> ($( $rtype ),*) { mysql($mysql_q) sqlite($sqlite_q) }
            }

            #[allow(non_snake_case)]
            $vi mod $name {
                #[allow(unused_imports)]
                use super::*;

                #[allow(unused_imports)]
                use $crate::_macro_internal::*;

                // Not possible to retry query with transaction
                #[allow(unused_imports)]
                pub use [<$name Impl>]::query_with_transaction;

                #[allow(dead_code)]
                pub async fn query(
                    config: &SqlQueryConfig,
                    connection: &Connection,
                    $( $pname: & $ptype, )*
                    $( $lname: & [ $ltype ], )*
                ) -> Result<Vec<($( $rtype, )*)>> {
                    let mut hasher = Hash128::with_seed(0);

                    $(
                        $pname.hash(&mut hasher);
                    )*
                    $(
                        $lname.hash(&mut hasher);
                    )*
                    stringify!($name).hash(&mut hasher);
                    stringify!($mysql_q).hash(&mut hasher);
                    stringify!($sqlite_q).hash(&mut hasher);
                    let key = hasher.finish_ext();
                    let data = CacheData {key, config: config.caching.as_ref()};

                    Ok(query_with_retry(
                        data,
                        || async move { Ok(MemcacheWrapper([<$name Impl>]::query(connection, $( $pname, )* $( $lname, )*).await?)) },
                    ).await?.0)
                }
            }

            $crate::mononoke_queries! { $( $rest )* }
        }
    };

    // Write query with a single expression. Redirect to write query with same expression for mysql and sqlite.
    (
        $vi:vis write $name:ident (
            values: ($( $vname:ident: $vtype:ty ),* $(,)*)
            $( , $pname:ident: $ptype:ty )* $(,)*
        ) { $qtype:ident, $q:expr }
        $( $rest:tt )*
    ) => {
        $crate::mononoke_queries! {
            $vi write $name (
                values: ( $( $vname: $vtype ),* )
                $( , $pname: $ptype )*
            ) { $qtype, mysql($q) sqlite($q) }
            $( $rest )*
        }
    };

    // Full write query with a list of values. Call `sql::queries!` and re-export stuff, wrapped in retries, on a new module.
    (
        $vi:vis write $name:ident (
            values: ($( $vname:ident: $vtype:ty ),* $(,)*)
            $( , $pname:ident: $ptype:ty )* $(,)*
        ) { $qtype:ident, mysql($mysql_q:expr) sqlite($sqlite_q:expr) }
        $( $rest:tt )*
    ) => {
        $crate::_macro_internal::paste::item! {
            $crate::_macro_internal::queries! {
                pub write [<$name Impl>] (
                    values: ( $( $vname: $vtype ),* )
                    $( , $pname: $ptype )*
                ) { $qtype, mysql($mysql_q) sqlite($sqlite_q) }
            }

            #[allow(non_snake_case)]
            $vi mod $name {
                #[allow(unused_imports)]
                use super::*;

                #[allow(unused_imports)]
                use $crate::_macro_internal::*;

                // Not possible to retry query with transaction
                #[allow(unused_imports)]
                pub use [<$name Impl>]::query_with_transaction;

                #[allow(dead_code)]
                pub async fn query(
                    connection: &Connection,
                    values: &[($( & $vtype, )*)],
                    $( $pname: & $ptype ),*
                ) -> Result<WriteResult> {
                    query_with_retry_no_cache(
                        || [<$name Impl>]::query(connection, values $( , $pname )* ),
                    ).await
                }
            }

            $crate::mononoke_queries! { $( $rest )* }
        }
    };

    // Write query with a single expression. Redirect to write query with same expression for mysql and sqlite.
    (
        $vi:vis write $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) { $qtype:ident, $q:expr }
        $( $rest:tt )*
    ) => {
        $crate::mononoke_queries! {
            $vi write $name (
                $( $pname: $ptype, )*
                $( >list $lname: $ltype )*
            ) { $qtype, mysql($q) sqlite($q) }
            $( $rest )*
        }
    };

    // Full write query without a list of values. Call `sql::queries!` and re-export stuff, wrapped in retries, on a new module.
    (
        $vi:vis write $name:ident (
            $( $pname:ident: $ptype:ty ),* $(,)*
            $( >list $lname:ident: $ltype:ty )*
        ) { $qtype:ident, mysql($mysql_q:expr) sqlite($sqlite_q:expr) }
        $( $rest:tt )*
    ) => {
        $crate::_macro_internal::paste::item! {
            $crate::_macro_internal::queries! {
                pub write [<$name Impl>] (
                    $( $pname: $ptype, )*
                    $( >list $lname: $ltype )*
                ) { $qtype, mysql($mysql_q) sqlite($sqlite_q) }
            }

            #[allow(non_snake_case)]
            $vi mod $name {
                #[allow(unused_imports)]
                use super::*;

                #[allow(unused_imports)]
                use $crate::_macro_internal::*;

                // Not possible to retry query with transaction
                #[allow(unused_imports)]
                pub use [<$name Impl>]::query_with_transaction;

                #[allow(dead_code)]
                pub async fn query(
                    connection: &Connection,
                    $( $pname: & $ptype, )*
                    $( $lname: & [ $ltype ], )*
                ) -> Result<WriteResult> {
                    query_with_retry_no_cache(
                        || [<$name Impl>]::query(connection, $( $pname, )* $( $lname, )*),
                    ).await
                }
            }

            $crate::mononoke_queries! { $( $rest )* }
        }
    };

}

#[cfg(fbcode_build)]
/// See https://fburl.com/sv/uk8w71td for error descriptions
fn retryable_mysql_errno(errno: u32) -> bool {
    match errno {
        // Deadlock error, advice is restarting transaction, so it's retryable
        1213 => true,
        // Admission control errors
        // Safe to retry on writes as well as the query didn't even start
        1914..=1916 => true,
        _ => false,
    }
}

#[cfg(fbcode_build)]
fn should_retry_mysql_query(err: &anyhow::Error) -> bool {
    use mysql_client::MysqlError;
    use MysqlError::*;
    match err.downcast_ref::<MysqlError>() {
        Some(ConnectionOperationError { mysql_errno, .. })
        | Some(QueryResultError { mysql_errno, .. }) => retryable_mysql_errno(*mysql_errno),
        _ => false,
    }
}

#[cfg(not(fbcode_build))]
fn should_retry_mysql_query(_err: &anyhow::Error) -> bool {
    false
}

type Key = u128;

pub struct CacheData<'a> {
    pub key: Key,
    pub config: Option<&'a CachingConfig>,
}

struct QueryCacheStore<'a, F, T> {
    key: Key,
    cache_config: &'a CachingConfig,
    cachelib: CachelibHandler<T>,
    memcache: MemcacheHandler,
    fetcher: F,
}

impl<F, V> EntityStore<V> for QueryCacheStore<'_, F, V> {
    fn cachelib(&self) -> &CachelibHandler<V> {
        &self.cachelib
    }

    fn keygen(&self) -> &KeyGen {
        &self.cache_config.keygen
    }

    fn memcache(&self) -> &MemcacheHandler {
        &self.memcache
    }

    fn cache_determinator(&self, _v: &V) -> CacheDisposition {
        CacheDisposition::Cache(CacheTtl::NoTtl)
    }

    caching_ext::impl_singleton_stats!("sql");
}

#[async_trait]
impl<V, F, Fut> KeyedEntityStore<Key, V> for QueryCacheStore<'_, F, V>
where
    V: Send + 'static,
    F: Fn() -> Fut + Send + Sync,
    Fut: Future<Output = Result<V>> + Send,
{
    fn get_cache_key(&self, key: &Key) -> String {
        // We just need a unique representation of the key as a String.
        // Let's use base64 as it's smaller than just .to_string()
        base64::encode(&key.to_ne_bytes())
    }

    async fn get_from_db(&self, keys: HashSet<Key>) -> Result<HashMap<Key, V>> {
        let key = keys.into_iter().exactly_one()?;
        anyhow::ensure!(key == self.key, "Fetched invalid key {}", key);
        let val = (self.fetcher)().await?;
        Ok(hashmap! { key => val })
    }
}

#[derive(Abomonation, Clone)]
pub struct MemcacheWrapper<T>(pub T);

impl<T> MemcacheEntity for MemcacheWrapper<T>
where
    T: serde::Serialize + for<'a> serde::Deserialize<'a>,
{
    fn serialize(&self) -> Bytes {
        serde_cbor::to_vec(&self.0)
            .expect("Should serialize cleanly")
            .into()
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        match serde_cbor::from_slice(bytes.as_ref()) {
            Ok(ok) => Ok(Self(ok)),
            Err(_) => Err(McErrorKind::Deserialization),
        }
    }
}

pub async fn query_with_retry_no_cache<T, Fut>(
    do_query: impl Fn() -> Fut + Send + Sync,
) -> Result<T>
where
    T: Send + 'static,
    Fut: Future<Output = Result<T>>,
{
    if tunables().disable_sql_auto_retries().unwrap_or_default() {
        return do_query().await;
    }
    Ok(retry(
        None,
        |_| do_query(),
        should_retry_mysql_query,
        // See https://fburl.com/7dmedu1u for backoff reasoning
        RetryLogic::ExponentialWithJitter {
            base: Duration::from_secs(10),
            factor: 1.2,
            jitter: Duration::from_secs(5),
        },
        RETRY_ATTEMPTS,
    )
    .await?
    .0)
}

pub async fn query_with_retry<T, Fut>(
    cache_data: CacheData<'_>,
    do_query: impl Fn() -> Fut + Send + Sync,
) -> Result<T>
where
    T: Send + Abomonation + MemcacheEntity + Clone + 'static,
    Fut: Future<Output = Result<T>> + Send,
{
    if tunables().disable_sql_auto_cache().unwrap_or_default() {
        return query_with_retry_no_cache(&do_query).await;
    }
    let fetch = || query_with_retry_no_cache(&do_query);
    let key = cache_data.key;
    if let Some(config) = cache_data.config.as_ref() {
        let store = QueryCacheStore {
            key: cache_data.key,
            cachelib: config.cache_handler_factory.cachelib(),
            memcache: config.cache_handler_factory.memcache(),
            cache_config: config,
            fetcher: fetch,
        };
        Ok(get_or_fill(&store, hashset! {key})
            .await?
            .into_iter()
            .exactly_one()
            .map_err(|_| anyhow!("Multiple values for a single key"))?
            .1)
    } else {
        fetch().await
    }
}

#[cfg(test)]
mod tests {
    mononoke_queries! {
        read TestQuery(param_str: String, param_uint: u64) -> (u64, Option<i32>, String, i64) {
            "SELECT 44, NULL, {param_str}, {param_uint}"
        }
        pub(crate) cacheable read TestQuery2() -> (u64, Option<String>) {
            "SELECT 44, NULL"
        }
        pub(super) write TestQuery3(values: (
            val1: i32,
        )) {
            none,
            "INSERT INTO my_table (num, str) VALUES {values}"
        }
        write TestQuery4(id: &str) {
            none,
            mysql("DELETE FROM my_table where id = {id}")
            sqlite("DELETE FROM mytable2 where id = {id}")
        }
    }

    #[allow(
        dead_code,
        unreachable_code,
        unused_variables,
        clippy::diverging_sub_expression
    )]
    async fn should_compile() -> anyhow::Result<()> {
        use sql_query_config::SqlQueryConfig;

        let config: &SqlQueryConfig = todo!();
        let connection: &sql::Connection = todo!();
        TestQuery::query(connection, todo!(), todo!()).await?;
        TestQuery::query_with_transaction(todo!(), todo!(), todo!()).await?;
        TestQuery2::query(config, connection).await?;
        TestQuery2::query_with_transaction(todo!()).await?;
        TestQuery3::query(connection, &[(&12,)]).await?;
        TestQuery3::query_with_transaction(todo!(), &[(&12,)]).await?;
        TestQuery4::query(connection, &"hello").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::Error;
use anyhow::Result;
use async_trait::async_trait;
use fbinit::FacebookInit;

use crate::facebook::*;
use crate::replication::ReplicaLag;
use crate::replication::ReplicaLagMonitor;
use crate::*;

macro_rules! fb_unimplemented {
    () => {
        unimplemented!("This is implemented only for fbcode_build!")
    };
}

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

impl PoolConfig {
    pub fn new(
        _size: usize,
        _threads_num: i32,
        _per_key_limit: u64,
        _conn_age_timeout_ms: u64,
        _conn_idle_timeout_ms: u64,
        _conn_open_timeout_ms: u64,
        _query_time_limit: Duration,
    ) -> Self {
        Self
    }

    pub fn default() -> Self {
        Self
    }
}

#[derive(Clone, Default)]
pub struct SharedConnectionPool;
impl SharedConnectionPool {
    pub fn new() -> Self {
        Self
    }
}

pub fn create_mysql_connections_unsharded(
    _fb: FacebookInit,
    _mysql_options: MysqlOptions,
    _label: String,
    _tier: String,
    _readonly: bool,
) -> Result<SqlConnections, Error> {
    fb_unimplemented!()
}

pub fn create_mysql_connections_sharded<S>(
    _fb: FacebookInit,
    _global_connection_pool: MysqlOptions,
    _label: String,
    _shardmap: String,
    _shards: S,
    _readonly: bool,
) -> Result<SqlShardedConnections, Error>
where
    S: IntoIterator<Item = usize> + Clone,
{
    fb_unimplemented!()
}

pub struct MyAdmin;
pub struct MyAdminLagMonitor;

impl MyAdmin {
    pub fn new(_: FacebookInit) -> Result<Self> {
        fb_unimplemented!()
    }

    pub fn single_shard_lag_monitor(&self, _: String) -> MyAdminLagMonitor {
        fb_unimplemented!()
    }

    pub fn shardmap_lag_monitor(&self, _: String) -> MyAdminLagMonitor {
        fb_unimplemented!()
    }
}

#[async_trait]
impl ReplicaLagMonitor for MyAdminLagMonitor {
    async fn get_replica_lag(&self) -> Result<Vec<ReplicaLag>> {
        fb_unimplemented!()
    }
}


/*
 * 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::fmt;
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use slog::info;
use slog::Logger;
use tokio::time;

const MAX_ALLOWED_REPLICATION_LAG_SECS: u64 = 5;
const REPLICATION_LAG_POLL_INTERVAL_SECS: u64 = 2;

// Laggable refers to an item that can lag.
// ReplicaLagMonitor can refer to a collection of Laggables or it can refer to services that
// will monitor replication lag.

// ---- ReplicaLagMonitor ----

#[async_trait]
pub trait ReplicaLagMonitor: Send + Sync {
    /// Returns the lag of all replicas configured.
    /// The Primary replica may or may not be returned, that is an implementation detail.
    async fn get_replica_lag(&self) -> Result<Vec<ReplicaLag>>;

    /// Returns the maximum lag of all replicas in cases that replicas are configured and available.
    /// In cases that no replica is configured, the result will be `ReplicaLag::no_delay()`.
    async fn get_max_replica_lag(&self) -> Result<ReplicaLag> {
        let max = self
            .get_replica_lag()
            .await?
            .into_iter()
            .max_by_key(|lag| lag.delay)
            .unwrap_or_else(ReplicaLag::no_delay);
        Ok(max)
    }

    /// Will poll periodically until the all replicas are below the given threshold of delay from
    /// the primary instance.
    /// Returns the last replica lag.
    async fn wait_for_replication(
        &self,
        config_getter: &(dyn Fn() -> WaitForReplicationConfig<'_> + Sync),
    ) -> Result<ReplicaLag> {
        loop {
            let max_lag = self.get_max_replica_lag().await?;
            let config = config_getter();
            if let Some(logger) = config.logger {
                info!(logger, "{}", max_lag);
            }
            if max_lag.delay < config.max_replication_lag_allowed {
                return Ok(max_lag);
            }
            // Wait before polling again.
            time::sleep(config.poll_interval).await;
        }
    }
}

pub struct NoReplicaLagMonitor();

#[async_trait]
impl ReplicaLagMonitor for NoReplicaLagMonitor {
    async fn get_replica_lag(&self) -> Result<Vec<ReplicaLag>> {
        Ok(vec![])
    }
}

pub struct ReplicaLag {
    pub delay: Duration,
    pub details: Option<String>,
}

impl ReplicaLag {
    pub fn new(delay: Duration, details: Option<String>) -> Self {
        ReplicaLag { delay, details }
    }

    pub fn no_delay() -> Self {
        ReplicaLag {
            delay: Duration::new(0, 0),
            details: None,
        }
    }
}

impl fmt::Display for ReplicaLag {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Replication lag is {}.{:03}s.",
            self.delay.as_secs(),
            self.delay.subsec_millis(),
        )?;
        if let Some(details) = &self.details {
            write!(f, " {}", details)?;
        }
        Ok(())
    }
}

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

#[derive(Clone)]
pub struct WaitForReplicationConfig<'a> {
    logger: Option<&'a Logger>,
    max_replication_lag_allowed: Duration,
    poll_interval: Duration,
}

impl<'a> Default for WaitForReplicationConfig<'a> {
    fn default() -> Self {
        WaitForReplicationConfig {
            logger: None,
            max_replication_lag_allowed: Duration::from_secs(MAX_ALLOWED_REPLICATION_LAG_SECS),
            poll_interval: Duration::from_secs(REPLICATION_LAG_POLL_INTERVAL_SECS),
        }
    }
}

impl<'a> WaitForReplicationConfig<'a> {
    pub fn new(
        max_replication_lag_allowed: Duration,
        poll_interval: Duration,
        logger: &'a Logger,
    ) -> Self {
        Self {
            max_replication_lag_allowed,
            poll_interval,
            logger: Some(logger),
        }
    }
    pub fn with_logger(mut self, logger: &'a Logger) -> Self {
        self.logger = Some(logger);
        self
    }
}

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

    use super::*;

    struct TestMonitor(u64);

    #[async_trait]
    impl ReplicaLagMonitor for TestMonitor {
        async fn get_replica_lag(&self) -> Result<Vec<ReplicaLag>> {
            Ok((1..self.0)
                .map(|lag| ReplicaLag::new(Duration::from_secs(lag), Some(format!("{}", lag))))
                .collect())
        }
    }

    #[tokio::test]
    async fn test_no_replica_lag_monitor() {
        let monitor = NoReplicaLagMonitor();
        let lag = monitor.get_max_replica_lag().await;
        // Linter gets confused here, says that expected is not used.
        let _expected = ReplicaLag::new(Duration::from_secs(0), None);
        assert_matches!(lag, Ok(_expected));
    }

    #[tokio::test]
    async fn test_max_lag() {
        let monitor = TestMonitor(5);
        let lag = monitor.get_max_replica_lag().await;
        // Linter gets confused here, says that expected is not used.
        let _expected = ReplicaLag::new(Duration::from_secs(5), Some("5".to_string()));
        assert_matches!(lag, Ok(_expected));
    }
}


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

use std::fs::create_dir_all;
use std::path::Path;
use std::time::Duration;

use anyhow::Result;
use sql::rusqlite::Connection as SqliteConnection;
use sql::rusqlite::OpenFlags as SqliteOpenFlags;

fn sqlite_setup_connection(con: &SqliteConnection) {
    // By default, when there's a read/write contention, SQLite will not wait,
    // but rather throw a `SQLITE_BUSY` error. See https://www.sqlite.org/lockingv3.html
    // This means that tests will fail in cases when production setup (e.g. one with MySQL)
    // would not. To change that, let's make sqlite wait for some time, before erroring out
    let _ = con.busy_timeout(Duration::from_secs(10));

    // By default, the `LIKE` operator is case-insensitive.  This doesn't
    // match MySQL, so change it to case-sensitive.
    let _ = con.pragma_update(None, "case_sensitive_like", true);
}

// Open a single sqlite connection to a new in memory database
pub fn open_sqlite_in_memory() -> Result<SqliteConnection> {
    let con = SqliteConnection::open_in_memory()?;
    sqlite_setup_connection(&con);
    Ok(con)
}

/// Open a single sqlite connection. The Sqlite DB will be created at path if necessary.
pub fn open_sqlite_path<P: AsRef<Path>>(path: P, readonly: bool) -> Result<SqliteConnection> {
    let path = path.as_ref();
    let con = {
        // Open a RW connection with create-on-open enabled, so that the Sqlite DB is initialized
        // even if we then re-open it with RO access. (You cannot open RO with create-on-open)
        // This is useful for tests that want to verify that nothing is written to the DB even when
        // it is empty.
        if let Some(parent) = path.parent() {
            create_dir_all(parent)?;
        }
        let flags = SqliteOpenFlags::SQLITE_OPEN_READ_WRITE | SqliteOpenFlags::SQLITE_OPEN_CREATE;

        SqliteConnection::open_with_flags(path, flags)?
    };

    let con = if readonly {
        let flags = SqliteOpenFlags::SQLITE_OPEN_READ_ONLY;
        SqliteConnection::open_with_flags(path, flags)?
    } else {
        con
    };

    sqlite_setup_connection(&con);
    Ok(con)
}

/// Open a single sqlite connection. The Sqlite DB must already exist at path.
pub fn open_existing_sqlite_path<P: AsRef<Path>>(
    path: P,
    readonly: bool,
) -> Result<SqliteConnection> {
    let path = path.as_ref();
    let flags = if readonly {
        SqliteOpenFlags::SQLITE_OPEN_READ_ONLY
    } else {
        // NB no creation flag, as db should already exist at this point
        SqliteOpenFlags::SQLITE_OPEN_READ_WRITE
    };
    let con = SqliteConnection::open_with_flags(path, flags)?;
    sqlite_setup_connection(&con);
    Ok(con)
}


/*
 * 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::fs::OpenOptions;
use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow as _;
use anyhow::anyhow;
use anyhow::Error;
use fbinit::FacebookInit;
#[cfg(fbcode_build)]
use scribe::ScribeClient; // oss uses anyhow

#[cfg(not(fbcode_build))]
mod oss;

#[cfg(not(fbcode_build))]
pub use oss::ScribeClientImplementation;
#[cfg(fbcode_build)]
pub use scuba::ScribeClientImplementation;

#[derive(Clone)]
pub enum Scribe {
    Client(Arc<ScribeClientImplementation>),
    LogToFile(Arc<Mutex<PathBuf>>),
}

impl ::std::fmt::Debug for Scribe {
    fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
        match self {
            Self::Client(_) => f.debug_struct("Scribe::Client").finish(),
            Self::LogToFile(_) => f.debug_struct("Scribe::LogToFile").finish(),
        }
    }
}

impl Scribe {
    pub fn new(fb: FacebookInit) -> Self {
        Self::Client(Arc::new(ScribeClientImplementation::new(fb)))
    }

    pub fn new_to_file(dir_path: PathBuf) -> Self {
        Self::LogToFile(Arc::new(Mutex::new(dir_path)))
    }

    pub fn offer(&self, category: &str, sample: &str) -> Result<(), Error> {
        use Scribe::*;

        match self {
            Client(client) => {
                #[cfg(fbcode_build)]
                let res = client.offer(category, sample.as_bytes());

                #[cfg(not(fbcode_build))]
                let res = {
                    let _ = client;
                    Ok(())
                };

                res
            }
            LogToFile(dir_path) => {
                let dir_path = dir_path.lock().unwrap();
                let is_valid_category = category
                    .chars()
                    .all(|c| char::is_alphanumeric(c) || c == '-' || c == '_');
                if !is_valid_category {
                    return Err(anyhow!("invalid category: {}", category));
                }
                let filename = dir_path.join(category);
                let mut file = OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(filename)?;
                ::std::writeln!(file, "{}", sample)?;
                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 anyhow::Result;
use fbinit::FacebookInit;

pub struct ScribeClientImplementation {}

impl ScribeClientImplementation {
    pub fn new(_fb: FacebookInit) -> Self {
        Self {}
    }

    pub fn offer(&self, _category: &str, _sample: &str) -> Result<()> {
        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::collections::hash_map::Entry;
use std::io::Error as IoError;
use std::num::NonZeroU64;
use std::path::Path;
use std::time::Duration;

use anyhow::Result;
use fbinit::FacebookInit;
use futures_stats::FutureStats;
use futures_stats::StreamStats;
use metadata::Metadata;
use nonzero_ext::nonzero;
use observability::ObservabilityContext;
use observability::ScubaLoggingDecisionFields;
pub use observability::ScubaVerbosityLevel;
use permission_checker::MononokeIdentitySetExt;
pub use scribe_ext::ScribeClientImplementation;
use scuba::builder::ServerData;
pub use scuba::Sampling;
use scuba::ScubaSample;
use scuba::ScubaSampleBuilder;
pub use scuba::ScubaValue;
use time_ext::DurationExt;
use tunables::tunables;

const FILE_PREFIX: &str = "file://";

/// An extensible wrapper struct around `ScubaSampleBuilder`
#[derive(Clone)]
pub struct MononokeScubaSampleBuilder {
    inner: ScubaSampleBuilder,
    maybe_observability_context: Option<ObservabilityContext>,
    // This field decides if sampled out requests should
    // still be logged when verbose logging is enabled
    fallback_sampled_out_to_verbose: bool,
}

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

enum ScubaLoggingType<'a> {
    ScubaTable(&'a str),
    LocalFile(&'a str),
}

fn get_scuba_logging_type(arg: &str) -> ScubaLoggingType {
    if let Some(path) = arg.strip_prefix(FILE_PREFIX) {
        ScubaLoggingType::LocalFile(path)
    } else {
        ScubaLoggingType::ScubaTable(arg)
    }
}

impl MononokeScubaSampleBuilder {
    pub fn new(fb: FacebookInit, scuba_table: &str) -> Result<Self> {
        Ok(Self {
            inner: Self::get_scuba_sample_builder(fb, get_scuba_logging_type(scuba_table))?,
            maybe_observability_context: None,
            fallback_sampled_out_to_verbose: false,
        })
    }

    pub fn with_discard() -> Self {
        Self {
            inner: ScubaSampleBuilder::with_discard(),
            maybe_observability_context: None,
            fallback_sampled_out_to_verbose: false,
        }
    }

    pub fn with_opt_table(fb: FacebookInit, scuba_table: Option<String>) -> Result<Self> {
        match scuba_table {
            None => Ok(Self::with_discard()),
            Some(scuba_table) => Self::new(fb, &scuba_table),
        }
    }

    pub fn with_observability_context(self, octx: ObservabilityContext) -> Self {
        Self {
            maybe_observability_context: Some(octx),
            ..self
        }
    }

    fn get_scuba_sample_builder(
        fb: FacebookInit,
        scuba_logging_type: ScubaLoggingType,
    ) -> Result<ScubaSampleBuilder, IoError> {
        Ok(match scuba_logging_type {
            ScubaLoggingType::ScubaTable(scuba_table) => ScubaSampleBuilder::new(fb, scuba_table),
            ScubaLoggingType::LocalFile(path) => {
                ScubaSampleBuilder::with_discard().with_log_file(path)?
            }
        })
    }

    fn get_logging_decision_fields(&self) -> ScubaLoggingDecisionFields {
        ScubaLoggingDecisionFields {
            maybe_session_id: self.get("session_uuid"),
            maybe_unix_username: self.get("unix_username"),
            maybe_source_hostname: self.get("source_hostname"),
        }
    }

    pub fn should_log_with_level(&self, level: ScubaVerbosityLevel) -> bool {
        match level {
            ScubaVerbosityLevel::Normal => true,
            ScubaVerbosityLevel::Verbose => self
                .maybe_observability_context
                .as_ref()
                .map_or(false, |octx| {
                    octx.should_log_scuba_sample(level, self.get_logging_decision_fields())
                }),
        }
    }

    pub fn add<K: Into<String>, V: Into<ScubaValue>>(&mut self, key: K, value: V) -> &mut Self {
        self.inner.add(key, value);
        self
    }

    pub fn add_metadata(&mut self, metadata: &Metadata) -> &mut Self {
        self.inner
            .add("session_uuid", metadata.session_id().to_string());

        self.inner.add(
            "client_identities",
            metadata
                .identities()
                .iter()
                .map(|i| i.to_string())
                .collect::<Vec<_>>(),
        );

        if let Some(client_hostname) = metadata.client_hostname() {
            // "source_hostname" to remain compatible with historical logging
            self.inner
                .add("source_hostname", client_hostname.to_owned());
        } else if let Some(client_ip) = metadata.client_ip() {
            self.inner.add("client_ip", client_ip.to_string());
        }
        if let Some(unix_name) = metadata.unix_name() {
            // "unix_username" to remain compatible with historical logging
            self.inner.add("unix_username", unix_name);
        }

        self.inner
            .add_opt("sandcastle_alias", metadata.sandcastle_alias());
        self.inner
            .add_opt("revproxy_region", metadata.revproxy_region().as_deref());
        self.inner
            .add_opt("sandcastle_nonce", metadata.sandcastle_nonce());
        self.inner
            .add_opt("clientinfo_tag", metadata.clientinfo_u64tag());
        self.inner
            .add_opt("client_tw_job", metadata.clientinfo_tw_job());
        self.inner
            .add_opt("client_tw_task", metadata.clientinfo_tw_task());

        self
    }

    pub fn sample_for_identities(&mut self, identities: &impl MononokeIdentitySetExt) {
        // Details of quicksand traffic aren't particularly interesting because all Quicksand tasks are
        // doing effectively the same thing at the same time. If we need real-time debugging, we can
        // always rely on updating the verbosity in real time.
        if identities.is_quicksand() {
            self.sampled_unless_verbose(nonzero!(100u64));
        }
    }

    pub fn log_with_msg<S: Into<Option<String>>>(&mut self, log_tag: &str, msg: S) {
        if self.fallback_sampled_out_to_verbose
            && self.should_log_with_level(ScubaVerbosityLevel::Verbose)
        {
            // We need to unsample before we log, so that
            // `sample_rate` field is not added, as we are about
            // to log everything.
            self.inner.unsampled();
        }

        self.inner.add("log_tag", log_tag);
        if let Some(mut msg) = msg.into() {
            match tunables()
                .max_scuba_msg_length()
                .unwrap_or_default()
                .try_into()
            {
                Ok(size) if size > 0 && msg.len() > size => {
                    msg.truncate(size);
                    msg.push_str(" (...)");
                }
                _ => {}
            };

            self.inner.add("msg", msg);
        }
        self.inner.log();
    }

    /// Same as `log_with_msg`, but sample is assumed to be verbose and is only logged
    /// if verbose logging conditions are met
    pub fn log_with_msg_verbose<S: Into<Option<String>>>(&mut self, log_tag: &str, msg: S) {
        if !self.should_log_with_level(ScubaVerbosityLevel::Verbose) {
            return;
        }

        self.log_with_msg(log_tag, msg)
    }

    pub fn add_stream_stats(&mut self, stats: &StreamStats) -> &mut Self {
        self.inner
            .add("poll_count", stats.poll_count)
            .add("poll_time_us", stats.poll_time.as_micros_unchecked())
            .add(
                "max_poll_time_us",
                stats.max_poll_time.as_micros_unchecked(),
            )
            .add_opt(
                "completion_time_us",
                stats
                    .completion_time
                    .as_ref()
                    .map(Duration::as_micros_unchecked),
            )
            .add("count", stats.count);

        self
    }

    pub fn add_prefixed_stream_stats(&mut self, stats: &StreamStats) -> &mut Self {
        self.inner
            .add("stream_poll_count", stats.poll_count)
            .add("stream_poll_time_us", stats.poll_time.as_micros_unchecked())
            .add(
                "stream_max_poll_time_us",
                stats.max_poll_time.as_micros_unchecked(),
            )
            .add("stream_count", stats.count)
            .add("stream_completed", stats.completed as u32)
            .add_opt(
                "stream_completion_time_us",
                stats
                    .completion_time
                    .as_ref()
                    .map(Duration::as_micros_unchecked),
            )
            .add_opt(
                "stream_first_item_time_us",
                stats
                    .first_item_time
                    .as_ref()
                    .map(Duration::as_micros_unchecked),
            );

        self
    }

    pub fn add_future_stats(&mut self, stats: &FutureStats) -> &mut Self {
        self.inner
            .add("poll_count", stats.poll_count)
            .add("poll_time_us", stats.poll_time.as_micros_unchecked())
            .add(
                "max_poll_time_us",
                stats.max_poll_time.as_micros_unchecked(),
            )
            .add(
                "completion_time_us",
                stats.completion_time.as_micros_unchecked(),
            );

        self
    }

    pub fn is_discard(&self) -> bool {
        self.inner.is_discard()
    }

    pub fn sampled(&mut self, sample_rate: NonZeroU64) -> &mut Self {
        self.fallback_sampled_out_to_verbose = false;
        self.inner.sampled(sample_rate);
        self
    }

    pub fn sampled_unless_verbose(&mut self, sample_rate: NonZeroU64) -> &mut Self {
        self.fallback_sampled_out_to_verbose = true;
        self.inner.sampled(sample_rate);
        self
    }

    pub fn unsampled(&mut self) -> &mut Self {
        self.inner.unsampled();
        self
    }

    pub fn log(&mut self) -> bool {
        self.inner.log()
    }

    /// Same as `log`, but sample is assumed to be verbose and is only logged
    /// if verbose logging conditions are met
    pub fn log_verbose(&mut self) -> bool {
        if !self.should_log_with_level(ScubaVerbosityLevel::Verbose) {
            // Return value of the `log` function indicates whether
            // the sample passed sampling. If it's too verbose, let's
            // return false
            return false;
        }

        self.log()
    }

    pub fn add_common_server_data(&mut self) -> &mut Self {
        self.inner.add_common_server_data();
        self
    }

    pub fn sampling(&self) -> &Sampling {
        self.inner.sampling()
    }

    pub fn add_mapped_common_server_data<F>(&mut self, mapper: F) -> &mut Self
    where
        F: Fn(ServerData) -> &'static str,
    {
        self.inner.add_mapped_common_server_data(mapper);
        self
    }

    pub fn with_log_file<L: AsRef<Path>>(mut self, log_file: L) -> Result<Self, IoError> {
        self.inner = self.inner.with_log_file(log_file)?;
        Ok(self)
    }

    pub fn with_seq(mut self, key: impl Into<String>) -> Self {
        self.inner = self.inner.with_seq(key);
        self
    }

    pub fn log_with_time(&mut self, time: u64) -> bool {
        self.inner.log_with_time(time)
    }

    pub fn entry<K: Into<String>>(&mut self, key: K) -> Entry<String, ScubaValue> {
        self.inner.entry(key)
    }

    pub fn flush(&self, timeout: Duration) {
        self.inner.flush(timeout)
    }

    pub fn get_sample(&self) -> &ScubaSample {
        self.inner.get_sample()
    }

    pub fn add_opt<K: Into<String>, V: Into<ScubaValue>>(
        &mut self,
        key: K,
        value: Option<V>,
    ) -> &mut Self {
        self.inner.add_opt(key, value);
        self
    }

    pub fn get<K: Into<String>>(&self, key: K) -> Option<&ScubaValue> {
        self.inner.get(key)
    }
}


/*
 * 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::anyhow;
use anyhow::Context;
use anyhow::Result;
use fbinit::FacebookInit;
use metaconfig_types::DatabaseConfig;
use metaconfig_types::LocalDatabaseConfig;
use metaconfig_types::MetadataDatabaseConfig;
use metaconfig_types::RemoteDatabaseConfig;
use metaconfig_types::RemoteMetadataDatabaseConfig;
use metaconfig_types::ShardableRemoteDatabaseConfig;
use metaconfig_types::ShardedDatabaseConfig;
use sql_ext::facebook::MysqlOptions;

use crate::construct::SqlConstruct;
use crate::facebook::FbSqlConstruct;
use crate::facebook::FbSqlShardedConstruct;

/// Trait that allows construction from database config.
pub trait SqlConstructFromDatabaseConfig: FbSqlConstruct + SqlConstruct {
    fn with_database_config(
        fb: FacebookInit,
        database_config: &DatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        match database_config {
            DatabaseConfig::Local(LocalDatabaseConfig { path }) => {
                Self::with_sqlite_path(path.join("sqlite_dbs"), readonly)
            }
            DatabaseConfig::Remote(config) => {
                Self::with_mysql(fb, config.db_address.clone(), mysql_options, readonly)
            }
        }
        .with_context(|| {
            format!(
                "While connecting to {:?} (with options {:?})",
                database_config, mysql_options
            )
        })
    }
}

impl<T: SqlConstruct + FbSqlConstruct> SqlConstructFromDatabaseConfig for T {}

/// Trait that allows construction from sharded database config.
pub trait SqlConstructFromShardedDatabaseConfig: FbSqlShardedConstruct {
    fn with_sharded_database_config(
        fb: FacebookInit,
        database_config: &ShardedDatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        match database_config {
            ShardedDatabaseConfig::Local(LocalDatabaseConfig { path }) => {
                Self::with_sqlite_path(path.join("sqlite_dbs"), readonly)
            }
            ShardedDatabaseConfig::Sharded(config) => Self::with_sharded_mysql(
                fb,
                config.shard_map.clone(),
                config.shard_num,
                mysql_options,
                readonly,
            ),
            ShardedDatabaseConfig::Unsharded(config) => {
                Self::with_mysql(fb, config.db_address.clone(), mysql_options, readonly)
            }
        }
        .with_context(|| {
            format!(
                "While connecting to {:?} (with options {:?})",
                database_config, mysql_options
            )
        })
    }
}

impl<T: FbSqlShardedConstruct> SqlConstructFromShardedDatabaseConfig for T {}

/// Trait that allows construction from the metadata database config.
pub trait SqlConstructFromMetadataDatabaseConfig: FbSqlConstruct + SqlConstruct {
    fn with_metadata_database_config(
        fb: FacebookInit,
        metadata_database_config: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        match metadata_database_config {
            MetadataDatabaseConfig::Local(LocalDatabaseConfig { path }) => {
                Self::with_sqlite_path(path.join("sqlite_dbs"), readonly)
            }
            MetadataDatabaseConfig::Remote(remote) => {
                Self::with_remote_metadata_database_config(fb, remote, mysql_options, readonly)
            }
        }
    }

    fn with_remote_metadata_database_config(
        fb: FacebookInit,
        remote: &RemoteMetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        let config = Self::remote_database_config(remote)
            .ok_or_else(|| anyhow!("no configuration available"))?;
        Self::with_mysql(fb, config.db_address.clone(), mysql_options, readonly)
    }

    /// Get the remote database config for this type.  Override this to use a database other than
    /// the primary database.
    fn remote_database_config(
        remote: &RemoteMetadataDatabaseConfig,
    ) -> Option<&RemoteDatabaseConfig> {
        Some(&remote.primary)
    }
}

/// Trait that allows construction of shardable databases from the metadata database config.
pub trait SqlShardableConstructFromMetadataDatabaseConfig:
    FbSqlConstruct + FbSqlShardedConstruct + SqlConstruct
{
    fn with_metadata_database_config(
        fb: FacebookInit,
        metadata_database_config: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        match metadata_database_config {
            MetadataDatabaseConfig::Local(LocalDatabaseConfig { path }) => {
                Self::with_sqlite_path(path.join("sqlite_dbs"), readonly)
            }
            MetadataDatabaseConfig::Remote(remote) => {
                Self::with_remote_metadata_database_config(fb, remote, mysql_options, readonly)
            }
        }
    }

    fn with_remote_metadata_database_config(
        fb: FacebookInit,
        remote: &RemoteMetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
        readonly: bool,
    ) -> Result<Self> {
        let config = Self::remote_database_config(remote)
            .ok_or_else(|| anyhow!("no configuration available"))?;
        match config {
            ShardableRemoteDatabaseConfig::Unsharded(config) => {
                Self::with_mysql(fb, config.db_address.clone(), mysql_options, readonly)
            }
            ShardableRemoteDatabaseConfig::Sharded(config) => Self::with_sharded_mysql(
                fb,
                config.shard_map.clone(),
                config.shard_num,
                mysql_options,
                readonly,
            ),
        }
    }

    /// Get the remote database config for this type.
    fn remote_database_config(
        remote: &RemoteMetadataDatabaseConfig,
    ) -> Option<&ShardableRemoteDatabaseConfig>;
}


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

use anyhow::Result;
use sql::Connection;
use sql::SqlConnections;
use sql::SqlShardedConnections;
use sql_ext::open_existing_sqlite_path;
use sql_ext::open_sqlite_in_memory;
use sql_ext::open_sqlite_path;
use vec1::vec1;

/// Construct a SQL data manager backed by a database
///
/// This trait should be implemented by any data manager that can be
/// constructed from a single backing database.
pub trait SqlConstruct: Sized + Send + Sync + 'static {
    /// Label used for statistics
    const LABEL: &'static str;

    /// Query used to create an empty instance of the database
    const CREATION_QUERY: &'static str;

    /// Construct an instance from SqlConnections
    ///
    /// This function may be called in an async context and must not block.
    fn from_sql_connections(connections: SqlConnections) -> Self;

    /// Construct an instance from an in-memory SQLite instance
    fn with_sqlite_in_memory() -> Result<Self> {
        let conn = open_sqlite_in_memory()?;
        conn.execute_batch(Self::CREATION_QUERY)?;
        let connections = SqlConnections::new_single(Connection::with_sqlite(conn));
        Ok(Self::from_sql_connections(connections))
    }

    /// Construct an instance from a SQLite database
    fn with_sqlite_path<P: AsRef<Path>>(path: P, readonly: bool) -> Result<Self> {
        let path = path.as_ref();
        let conn = open_sqlite_path(path, false)?;
        conn.execute_batch(Self::CREATION_QUERY)?;
        let write_connection = Connection::with_sqlite(conn);
        let read_connection = Connection::with_sqlite(open_existing_sqlite_path(path, true)?);
        let connections = SqlConnections {
            write_connection: if readonly {
                read_connection.clone()
            } else {
                write_connection
            },
            read_master_connection: read_connection.clone(),
            read_connection,
        };
        Ok(Self::from_sql_connections(connections))
    }
}

/// Construct a SQL data manager backed by a sharded database
///
/// This trait should be implemented by any data manager that can be
/// constructed from a sharded backing database.
pub trait SqlShardedConstruct: Sized + Send + Sync + 'static {
    /// Label used for statistics
    const LABEL: &'static str;

    /// Query used to create an empty instance of a shard
    const CREATION_QUERY: &'static str;

    /// Construct an instance from a vector of SqlConnections, one for each shard
    ///
    /// This function may be called in an async context and must not block.
    fn from_sql_shard_connections(shard_connections: SqlShardedConnections) -> Self;
}

impl<T> SqlConstruct for T
where
    T: SqlShardedConstruct,
{
    const LABEL: &'static str = <T as SqlShardedConstruct>::LABEL;
    const CREATION_QUERY: &'static str = <T as SqlShardedConstruct>::CREATION_QUERY;

    fn from_sql_connections(conns: SqlConnections) -> Self {
        Self::from_sql_shard_connections(SqlShardedConnections {
            write_connections: vec1![conns.write_connection],
            read_connections: vec1![conns.read_connection],
            read_master_connections: vec1![conns.read_master_connection],
        })
    }
}


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

//! Common construction utilities for SQL database managers.
//!
//! Mononoke data stores that are backed by SQL databases are managed by a database manager, like
//! `SqlChangesets`, `SqlBookmarks`, etc.  This crate provides common utilities for constructing
//! these database managers, backed by various database types.
//!
//! Database managers should implement `SqlConstruct` to define how to be constructed from
//! a set of `SqlConnections`.  This is sufficient to allow construction based on `DatabaseConfig`,
//! which is provided through the `SqlConstructFromDatabaseConfig` trait.
//!
//! Database managers that support sharding should additionally implement `SqlShardedConstruct` for
//! the sharded case.
//!
//! Database managers that would like to be constructed from repository metadata configuration
//! should implement the `SqlConstructFromMetadataDatabaseConfig` trait.  If their data is not
//! stored in the primary metadata database, they should implement the `remote_database_config`
//! method to define which configuration is used for remote database configuration.
//!
//! Database managers that support sharding should instead implement the
//! `SqlShardableConstructFromMetadataDatabaseConfig` trait, which allows them to return
//! either sharded or unsharded configuration from `remote_database_config`.

mod config;
mod construct;
#[cfg(not(fbcode_build))]
mod oss;

pub use config::SqlConstructFromDatabaseConfig;
pub use config::SqlConstructFromMetadataDatabaseConfig;
pub use config::SqlConstructFromShardedDatabaseConfig;
pub use config::SqlShardableConstructFromMetadataDatabaseConfig;
pub use construct::SqlConstruct;
pub use construct::SqlShardedConstruct;

pub mod facebook {
    #[cfg(fbcode_build)]
    mod r#impl;

    #[cfg(fbcode_build)]
    pub use r#impl::FbSqlConstruct;
    #[cfg(fbcode_build)]
    pub use r#impl::FbSqlShardedConstruct;

    #[cfg(not(fbcode_build))]
    pub use crate::oss::FbSqlConstruct;
    #[cfg(not(fbcode_build))]
    pub use crate::oss::FbSqlShardedConstruct;
}


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

use anyhow::Result;
use fbinit::FacebookInit;
use sql_ext::facebook::MysqlOptions;

use crate::SqlConstruct;
use crate::SqlShardedConstruct;

macro_rules! fb_unimplemented {
    () => {
        unimplemented!("This is implemented only for fbcode_build!")
    };
}

/// Construct a SQL data manager backed by Facebook infrastructure
pub trait FbSqlConstruct: SqlConstruct + Sized + Send + Sync + 'static {
    fn with_mysql<'a>(_: FacebookInit, _: String, _: &'a MysqlOptions, _: bool) -> Result<Self> {
        fb_unimplemented!()
    }
}

impl<T: SqlConstruct> FbSqlConstruct for T {}

/// Construct a sharded SQL data manager backed by Facebook infrastructure
pub trait FbSqlShardedConstruct: SqlShardedConstruct + Sized + Send + Sync + 'static {
    fn with_sharded_mysql(
        _: FacebookInit,
        _: String,
        _: NonZeroUsize,
        _: &MysqlOptions,
        _: bool,
    ) -> Result<Self> {
        fb_unimplemented!()
    }
}

impl<T: SqlShardedConstruct> FbSqlShardedConstruct for T {}


/*
 * 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 speed_sliding_window;
pub mod timeseries;