/*
 * 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 ratelimit;
mod shard;

use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::hash::Hasher;
use std::num::NonZeroU64;
use std::num::NonZeroUsize;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreMetadata;
use bytes::Buf;
use bytes::BufMut;
use bytes::Bytes;
use bytes::BytesMut;
use cacheblob::CachelibBlobstoreOptions;
use cachelib::VolatileLruCachePool;
use cloned::cloned;
use context::CoreContext;
use context::PerfCounterType;
use futures::future::WeakShared;
use futures::FutureExt;
use mononoke_types::BlobstoreBytes;
use shard::SemaphoreAcquisition;
use shard::Shards;
use shared_error::anyhow::IntoSharedError;
use shared_error::anyhow::SharedError;
use stats::prelude::*;
use tunables::tunables;
use twox_hash::XxHash;

use crate::ratelimit::AccessReason;
use crate::ratelimit::Ticket;

define_stats! {
    prefix = "mononoke.virtually_sharded_blobstore";
    gets: timeseries(Sum),
    gets_deduped: timeseries(Sum),
    gets_not_storable: timeseries(Sum),
    gets_not_storable_deduped: timeseries(Sum),
    puts: timeseries(Sum),
    puts_deduped: timeseries(Sum),
}

// 4MiB, minus a little space for the STORED prefix and the key. Note that we also arbitrarily add
// 128 bytes because 4 * 1024 * 1024 - 1024 is also our Filestore threshold so it's a good idea to
// not attempt to recompress that given it was chosen to fit in cachelib.
const MAX_CACHELIB_VALUE_SIZE: u64 = 4 * 1024 * 1024 - 1024 + 128;

const NOT_STORABLE: &[u8] = &[0];
const STORED: &[u8] = &[1];

#[derive(Debug)]
struct CacheKey(String);

impl CacheKey {
    fn from_key(key: &str) -> Self {
        Self(format!("vsb.{}", key))
    }
}

impl AsRef<[u8]> for CacheKey {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

#[derive(Debug)]
enum CacheData {
    /// Represents data that was found in cache.
    Stored(BlobstoreGetData),
    /// Represents data that is known to not be storable in cache (because it's too large,
    /// presumably). For this data, we skip semaphore access.
    NotStorable,
}

impl CacheData {
    fn deserialize(mut val: Bytes) -> Result<Self> {
        let prefix = val.split_to(1);

        if prefix.as_ref() == NOT_STORABLE {
            return Ok(Self::NotStorable);
        }

        if prefix.as_ref() == STORED {
            let val = BlobstoreBytes::decode(val)
                .ok_or_else(|| anyhow!("Invalid data in blob cache"))?
                .into();
            return Ok(Self::Stored(val));
        }

        Err(anyhow!("Invalid prefix: {:?}", prefix))
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug)]
enum PresenceData {
    /// We read this at some point. We will not dedupe based on this.
    Get,
    /// We wrote this ourselves, and this is the hash of the data. We will dedupe based on this.
    Put(u64),
}

impl PresenceData {
    const GET: &'static [u8] = &[0];
    const PUT: &'static [u8] = &[1];

    fn from_put(v: &BlobstoreBytes) -> Self {
        let mut hasher = XxHash::with_seed(0);
        hasher.write(v.as_bytes().as_ref());
        Self::Put(hasher.finish())
    }

    fn deserialize(mut val: Bytes) -> Result<Self> {
        let prefix = val.split_to(1);

        if prefix.as_ref() == Self::GET {
            return Ok(Self::Get);
        }

        if prefix.as_ref() == Self::PUT {
            let bytes: [u8; 8] = val
                .as_ref()
                .try_into()
                .context("Invalid put data in presence cache")?;

            return Ok(Self::Put(u64::from_ne_bytes(bytes)));
        }

        Err(anyhow!("Invalid prefix: {:?}", prefix))
    }

    fn serialize(&self) -> Bytes {
        match self {
            Self::Get => Self::GET.into(),
            Self::Put(v) => {
                let mut buff = BytesMut::with_capacity(1 + std::mem::size_of::<u64>());
                buff.put(Self::PUT);
                buff.put(v.to_ne_bytes().as_ref());
                buff.freeze()
            }
        }
    }

    fn is_put(&self) -> bool {
        match self {
            Self::Get => false,
            Self::Put(..) => true,
        }
    }
}

/// A type representing that a blob is known to exist. Used in matches below.
struct KnownToExist;

/// We allow filtering cache writes to make testing easier. This function is a default that does
/// not filter.
fn allow_all_filter(_: &Bytes) -> Result<()> {
    Ok(())
}

/// A layer over an existing blobstore that serializes access to virtual slices of the blobstore,
/// indexed by key. It also deduplicates writes for data that is already present.
#[derive(Clone)]
pub struct VirtuallyShardedBlobstore<T> {
    inner: Arc<Inner<T>>,
}

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

impl<T> VirtuallyShardedBlobstore<T> {
    pub fn new(
        blobstore: T,
        blob_pool: VolatileLruCachePool,
        presence_pool: VolatileLruCachePool,
        shards: NonZeroUsize,
        cachelib_options: CachelibBlobstoreOptions,
    ) -> Self {
        let cache = Cache {
            blob_pool,
            presence_pool,
            cache_filter: allow_all_filter,
            cachelib_options,
        };

        let inner = Inner::new(blobstore, shards, cache);

        Self {
            inner: Arc::new(inner),
        }
    }
}

impl<T: fmt::Debug> fmt::Debug for VirtuallyShardedBlobstore<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("VirtuallyShardedBlobstore")
            .field("blobstore", &self.inner.blobstore)
            .field("write_shards", &self.inner.write_shards.len())
            .field("read_shards", &self.inner.read_shards.len())
            .finish()
    }
}

struct Cache {
    presence_pool: VolatileLruCachePool,
    blob_pool: VolatileLruCachePool,
    cache_filter: fn(&Bytes) -> Result<()>,
    cachelib_options: CachelibBlobstoreOptions,
}

impl Cache {
    fn get_from_cache(&self, key: &CacheKey) -> Result<Option<CacheData>> {
        let val = match self.blob_pool.get(key)? {
            Some(val) => val,
            None => return Ok(None),
        };

        Ok(Some(CacheData::deserialize(val)?))
    }

    /// Set presence for this cache key.
    fn set_is_present(&self, key: &CacheKey, value: PresenceData) -> Result<()> {
        // If it's a put, then we overwrite existing data in cache, to record the new value.
        if value.is_put() {
            self.presence_pool.set_or_replace(key, value.serialize())?;
            return Ok(());
        }

        // If it's a get, we leave existing data alone. If existing data was a get then there is no
        // use updating it, and if it was a put, we don't *want* to update it (see
        // test_dedupe_writes_different_data).
        self.presence_pool.set(key, value.serialize())?;

        Ok(())
    }

    fn set_in_cache(
        &self,
        key: &CacheKey,
        presence: PresenceData,
        value: BlobstoreGetData,
    ) -> Result<()> {
        self.set_is_present(key, presence)?;

        let encode_limit = if self.cachelib_options.attempt_zstd {
            Some(MAX_CACHELIB_VALUE_SIZE)
        } else {
            None
        };
        let stored = value
            .into_bytes()
            .encode(encode_limit)
            .ok_or_else(|| anyhow!("Could not encode"))
            .and_then(|encoded| {
                (self.cache_filter)(&encoded)?;
                self.blob_pool.set(key, STORED.chain(encoded))
            })
            .unwrap_or(false);

        // NOTE: If a transient error occured while setting in cache, then we might store
        // NOT_STORABLE, even if the key is in fact storable. That's OK: it just means the next
        // gets will bypass the semaphore, but if the key does turn out to be cacheable, then it'll
        // get cached on the next read.
        if !stored {
            self.blob_pool.set(key, NOT_STORABLE)?;
        }

        Ok(())
    }

    /// Ask the cache if it knows whether the backing store has a given request PresenceData for
    /// this key. Returns Some(KnownToExist) if so. The request argument lets the caller control
    /// what they want to see: PresenceData::Get means we just want to know this key is readable.
    /// PresenceData::Put(v) means we want to know v specifically was last written by us.
    fn check_presence(
        &self,
        key: &CacheKey,
        request: PresenceData,
    ) -> Result<Option<KnownToExist>> {
        let stored = self
            .presence_pool
            .get(key)?
            .map(PresenceData::deserialize)
            .transpose()?;

        let r = match (stored, request) {
            (None, _) => {
                // Nothing in the presence cache at all (see test_read_after_write).
                None
            }
            (Some(_), PresenceData::Get) => {
                // Something in the cache, and we're not really looking for a specific value. This
                // works (see test_dedupe_reads).
                Some(KnownToExist)
            }
            (Some(PresenceData::Get), PresenceData::Put(..)) => {
                // We want a specific value, but we don't have one in cache. This does not work
                // (see test_dedupe_writes_different_data).
                None
            }
            (Some(PresenceData::Put(v1)), PresenceData::Put(v2)) => {
                // We have a put in the cache, and we want to check for a put. If they're the same,
                // this is a match (see test_dedupe_writes and test_dedupe_writes_different_data).
                if v1 == v2 { Some(KnownToExist) } else { None }
            }
        };

        Ok(r)
    }
}

type SharedReadFuture =
    WeakShared<Pin<Box<dyn Future<Output = Result<Option<BlobstoreGetData>, SharedError>> + Send>>>;

struct Inner<T> {
    blobstore: T,
    write_shards: Shards,
    read_shards: Shards,
    cache: Cache,
    large_inflight_reads: Mutex<HashMap<String, SharedReadFuture>>,
}

impl<T> Inner<T> {
    fn new(blobstore: T, shards: NonZeroUsize, cache: Cache) -> Self {
        Self {
            blobstore,
            write_shards: Shards::new(shards, PerfCounterType::BlobPutsShardAccessWait),
            read_shards: Shards::new(shards, PerfCounterType::BlobGetsShardAccessWait),
            cache,
            large_inflight_reads: Default::default(),
        }
    }
}

async fn shared_read<T: Blobstore + 'static>(
    inner: &Arc<Inner<T>>,
    ctx: &CoreContext,
    key: String,
    ticket: Ticket<'_>,
) -> Result<Option<BlobstoreGetData>> {
    let inflight_read = {
        let mut large_inflight_reads = inner.large_inflight_reads.lock().unwrap();
        if let Some(inflight_read) = large_inflight_reads.get(&key).and_then(WeakShared::upgrade) {
            // A read is already in flight.  Share it.
            STATS::gets_not_storable_deduped.add_value(1);
            ctx.perf_counters()
                .increment_counter(PerfCounterType::BlobGetsDeduplicatedLarge);
            ticket.cancel();
            inflight_read
        } else {
            // There is no read for this key in flight.  Create one.
            let inflight_read = {
                cloned!(ctx, inner, key);
                // Construct a scopeguard that will ensure we remove the
                // shared future from the inflight reads collection regardless
                // of whether it succeeds, fails, or is cancelled.
                let guard = scopeguard::guard((inner.clone(), key.clone()), |(inner, key)| {
                    let mut large_inflight_reads = inner.large_inflight_reads.lock().unwrap();
                    large_inflight_reads.remove(&key);
                });
                let ticket = ticket.into_owned();
                async move {
                    // Wait for our ticket.
                    ticket.finish().await.shared_error()?;
                    // If we were successful at getting the ticket, get data from the blobstore.
                    let result = inner.blobstore.get(&ctx, &key).await;
                    // Ensure the guard is moved into the future, and drop it
                    // now to remove ourself from the in flight reads.  This
                    // will happen regardless of whether the above steps were
                    // successful or not, or even if the future is cancelled.
                    drop(guard);
                    result.shared_error()
                }
                .boxed()
                .shared()
            };
            if let Some(inflight_read) = inflight_read.downgrade() {
                large_inflight_reads.insert(key, inflight_read);
            } else {
                // This shouldn't happen, as we just constructed the shared
                // future and haven't polled it yet.  But if it does, ensure
                // there is no entry for this key so the next read starts from
                // scratch.
                large_inflight_reads.remove(&key);
            }
            inflight_read
        }
    };
    Ok(inflight_read.await?)
}

fn report_deduplicated_put(ctx: &CoreContext, key: &str) {
    STATS::puts_deduped.add_value(1);

    let mut scuba = ctx.scuba().clone();
    if let Ok(Some(v)) = tunables()
        .deduplicated_put_sampling_rate()
        .unwrap_or_default()
        .try_into()
        .map(NonZeroU64::new)
    {
        scuba.sampled(v);
    }
    scuba.add("key", key).log_with_msg("Put deduplicated", None);

    ctx.perf_counters()
        .increment_counter(PerfCounterType::BlobPutsDeduplicated);
}

#[async_trait]
impl<T: Blobstore + 'static> Blobstore for VirtuallyShardedBlobstore<T> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        cloned!(self.inner, ctx);

        STATS::gets.add_value(1);
        let cache_key = CacheKey::from_key(key);

        // First, check the cache, and acquire a permit for this key if necessary.

        let take_lease = match inner.cache.get_from_cache(&cache_key)? {
            Some(CacheData::Stored(v)) => {
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::CachelibHits);
                return Ok(Some(v));
            }
            Some(CacheData::NotStorable) => {
                // We know for sure this data isn't cacheable. Don't try to acquire a permit
                // for it.  We will attempt to share reads with other readers of this data.
                false
            }
            None => true,
        };

        let key = key.to_owned();
        let fut = async move {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::CachelibMisses);

            let ticket = Ticket::new(&ctx, AccessReason::Read);

            let permit = if take_lease {
                let acq = inner
                    .read_shards
                    .acquire(&ctx, &key, ticket, || {
                        inner.cache.get_from_cache(&cache_key)
                    })
                    .await?;

                match acq {
                    SemaphoreAcquisition::Cancelled(CacheData::Stored(v), ticket) => {
                        // The data is cached, that's great. Return it. We're not going to hit
                        // the blobstore, so also return out ticket.
                        STATS::gets_deduped.add_value(1);
                        ctx.perf_counters()
                            .increment_counter(PerfCounterType::BlobGetsDeduplicated);
                        ticket.cancel();
                        return Ok(Some(v));
                    }
                    SemaphoreAcquisition::Cancelled(CacheData::NotStorable, ticket) => {
                        // The data cannot be cached. We'll have to go to the blobstore.
                        STATS::gets_not_storable.add_value(1);
                        if !tunables()
                            .disable_large_blob_read_deduplication()
                            .unwrap_or_default()
                        {
                            // Attempt to share with other readers of this blob.
                            return shared_read(&inner, &ctx, key, ticket).await;
                        } else {
                            ticket.finish().await?;
                            None
                        }
                    }
                    SemaphoreAcquisition::Acquired(permit) => Some(permit),
                }
            } else {
                // We already know the data can't be cached, so we'll have to go to the
                // blobstore.
                if !tunables()
                    .disable_large_blob_read_deduplication()
                    .unwrap_or_default()
                {
                    // Attempt to share with other reads of this blob.
                    return shared_read(&inner, &ctx, key, ticket).await;
                } else {
                    ticket.finish().await?;
                    None
                }
            };

            // NOTE: This is a no-op, but it's here to ensure permit is still in scope at this
            // point (which it should: if it doesn't, then that means we unconditionally released
            // the semaphore before doing the get, and that's wrong).
            scopeguard::defer! { drop(permit) };

            // Now, actually go the underlying blobstore.
            let res = inner.blobstore.get(&ctx, &key).await?;

            // And finally, attempt to cache what we got back.
            if let Some(ref data) = res {
                let _ = inner
                    .cache
                    .set_in_cache(&cache_key, PresenceData::Get, data.clone());
            }

            Ok(res)
        };

        tokio::spawn(fut).await?
    }

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

        STATS::puts.add_value(1);
        let cache_key = CacheKey::from_key(&key);
        let presence = PresenceData::from_put(&value);

        if let Ok(Some(KnownToExist)) = inner.cache.check_presence(&cache_key, presence) {
            report_deduplicated_put(&ctx, &key);
            return Ok(());
        }

        let fut = async move {
            let ticket = Ticket::new(&ctx, AccessReason::Write);

            let acq = inner
                .write_shards
                .acquire(&ctx, &key, ticket, || {
                    inner.cache.check_presence(&cache_key, presence)
                })
                .await?;

            let permit = match acq {
                SemaphoreAcquisition::Cancelled(KnownToExist, ticket) => {
                    report_deduplicated_put(&ctx, &key);
                    ticket.cancel();
                    return Ok(());
                }
                SemaphoreAcquisition::Acquired(permit) => permit,
            };

            scopeguard::defer! { drop(permit) };

            let res = inner.blobstore.put(&ctx, key, value.clone()).await?;

            let value = BlobstoreGetData::new(BlobstoreMetadata::default(), value);
            let _ = inner.cache.set_in_cache(&cache_key, presence, value);

            Ok(res)
        };

        tokio::spawn(fut).await?
    }

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

        let cache_key = CacheKey::from_key(key);
        let presence = PresenceData::Get;

        if let Ok(Some(KnownToExist)) = inner.cache.check_presence(&cache_key, presence) {
            return Ok(BlobstoreIsPresent::Present);
        }

        Ticket::new(ctx, AccessReason::Read).finish().await?;

        let result = inner.blobstore.is_present(ctx, key).await?;
        if let BlobstoreIsPresent::Present = &result {
            let _ = inner.cache.set_is_present(&cache_key, presence);
        }

        Ok(result)
    }
}

#[cfg(all(test, fbcode_build))]
mod test {
    use fbinit::FacebookInit;
    use futures_stats::TimedTryFutureExt;
    use nonzero_ext::nonzero;
    use once_cell::sync::OnceCell;
    use time_ext::DurationExt;

    use super::*;

    fn make_blobstore<B: Blobstore>(
        fb: FacebookInit,
        blobstore: B,
        blob_pool_name: &str,
        cache_shards: NonZeroUsize,
        cache_filter: fn(&Bytes) -> Result<()>,
    ) -> Result<VirtuallyShardedBlobstore<B>> {
        let cache = make_cache(fb, blob_pool_name, cache_filter)?;

        let inner = Inner::new(blobstore, cache_shards, cache);

        Ok(VirtuallyShardedBlobstore {
            inner: Arc::new(inner),
        })
    }

    fn make_cache(
        fb: FacebookInit,
        blob_pool_name: &str,
        cache_filter: fn(&Bytes) -> Result<()>,
    ) -> Result<Cache> {
        static INSTANCE: OnceCell<()> = OnceCell::new();
        INSTANCE.get_or_init(|| {
            let config = cachelib::LruCacheConfig::new(64 * 1024 * 1024);
            cachelib::init_cache(fb, config).unwrap();
        });

        let blob_pool = cachelib::get_or_create_volatile_pool(blob_pool_name, 8 * 1024 * 1024)?;
        let presence_pool = cachelib::get_or_create_volatile_pool("presence", 8 * 1024 * 1024)?;

        Ok(Cache {
            presence_pool,
            blob_pool,
            cache_filter,
            cachelib_options: CachelibBlobstoreOptions::default(),
        })
    }

    fn reject_all_filter(_: &Bytes) -> Result<()> {
        Err(anyhow!("Rejected!"))
    }

    mod caching {
        use mononoke_types::content_chunk::new_blob_and_pointer;
        use mononoke_types::BlobstoreKey;

        use super::*;

        #[fbinit::test]
        fn test_filestore_chunk_is_not_compressed(fb: FacebookInit) -> Result<()> {
            fn assert_below_threshold(b: &Bytes) -> Result<()> {
                assert!(
                    (b.len() as u64) < MAX_CACHELIB_VALUE_SIZE,
                    "Blob of size {} would have triggered compression!",
                    b.len()
                );
                Ok(())
            }

            let mut cache = make_cache(fb, "blobs", assert_below_threshold)?;
            cache.cachelib_options.attempt_zstd = false;

            let (blob, _) = new_blob_and_pointer(vec![0; 4193280]);

            // Use a key that looks reasonably close to what we'd actually use in prod.
            let key = format!("repo1234.{}", blob.id().blobstore_key());
            let key = CacheKey::from_key(&key);

            let bytes = BlobstoreBytes::from(blob);
            cache.set_in_cache(&key, PresenceData::Get, bytes.clone().into())?;

            let cached = cache.get_from_cache(&key)?.context("Blob not in cache")?;

            assert_matches::assert_matches!(cached, CacheData::Stored(cached) => {
                let expected = BlobstoreGetData::new(BlobstoreMetadata::default(), bytes);
                assert_eq!(expected, cached);
            });

            Ok(())
        }
    }

    mod sharding {
        use std::collections::HashMap;
        use std::sync::Mutex;
        use std::time::Duration;

        use borrowed::borrowed;
        use tokio::sync::broadcast;
        use tokio::sync::broadcast::Receiver;
        use tokio::sync::broadcast::Sender;

        use super::*;

        const TIMEOUT_MS: u64 = 100;

        /// Represents data stored in our TestBlobstore
        #[derive(Debug)]
        enum BlobData {
            Bytes(BlobstoreBytes),
            Channel(Sender<BlobstoreBytes>),
        }

        impl BlobData {
            /// Obtain a handle for a new get
            fn handle(&self) -> BlobDataHandle {
                match self {
                    BlobData::Bytes(ref b) => BlobDataHandle::Bytes(b.clone()),
                    BlobData::Channel(ref s) => BlobDataHandle::Channel(s.subscribe()),
                }
            }
        }

        /// Represents a handle for a single get from our TestBlobstore
        enum BlobDataHandle {
            Bytes(BlobstoreBytes),
            Channel(Receiver<BlobstoreBytes>),
        }

        impl BlobDataHandle {
            /// Obtain the bytes for this get.
            async fn bytes(self) -> Result<BlobstoreBytes> {
                let b = match self {
                    BlobDataHandle::Bytes(b) => b,
                    BlobDataHandle::Channel(mut r) => r.recv().await?,
                };

                Ok(b)
            }
        }

        #[derive(Default, Debug)]
        struct Blob {
            puts: u64,
            gets: u64,
            data: Option<BlobData>,
        }

        #[derive(Debug, Clone)]
        struct TestBlobstore {
            data: Arc<Mutex<HashMap<String, Blob>>>,
        }

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

        impl TestBlobstore {
            fn new() -> Self {
                Self {
                    data: Arc::new(Mutex::new(HashMap::new())),
                }
            }
        }

        #[async_trait]
        impl Blobstore for TestBlobstore {
            async fn put<'a>(
                &'a self,
                _ctx: &'a CoreContext,
                key: String,
                value: BlobstoreBytes,
            ) -> Result<()> {
                let mut data = self.data.lock().unwrap();
                let blob = data.entry(key).or_default();
                blob.puts += 1;
                blob.data = Some(BlobData::Bytes(value));
                Ok(())
            }

            async fn get<'a>(
                &'a self,
                _ctx: &'a CoreContext,
                key: &'a str,
            ) -> Result<Option<BlobstoreGetData>> {
                let handle = {
                    let mut data = self.data.lock().unwrap();
                    let blob = data.entry(key.to_owned()).or_default();
                    blob.gets += 1;
                    blob.data.as_ref().map(BlobData::handle)
                };

                let handle = match handle {
                    Some(handle) => handle,
                    None => {
                        return Ok(None);
                    }
                };

                let bytes = handle.bytes().await?;

                Ok(Some(BlobstoreGetData::new(
                    BlobstoreMetadata::default(),
                    bytes,
                )))
            }
        }

        #[fbinit::test]
        async fn test_dedupe_reads(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            let key = "foo";

            futures::future::try_join_all((0..10usize).map(|_| blobstore.get(ctx, key))).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 10);
                blob.data = Some(BlobData::Bytes(BlobstoreBytes::from_bytes("foo")));
            }

            futures::future::try_join_all((0..10usize).map(|_| blobstore.get(ctx, key))).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 11);
            }

            futures::future::try_join_all((0..10usize).map(|_| blobstore.is_present(ctx, key)))
                .await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 11);
            }

            Ok(())
        }

        #[fbinit::test]
        async fn test_cache_read(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            let key = "foo";
            let val = BlobstoreBytes::from_bytes("foo");

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                blob.data = Some(BlobData::Bytes(val.clone()));
            }

            let v1 = blobstore.get(ctx, key).await?;
            let v2 = blobstore.get(ctx, key).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 1);
            }

            assert_eq!(v1.unwrap().as_bytes(), &val);
            assert_eq!(v2.unwrap().as_bytes(), &val);

            Ok(())
        }

        #[fbinit::test]
        async fn test_read_after_write(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            let key = "foo";
            let val = BlobstoreBytes::from_bytes("foo");

            blobstore.put(ctx, key.to_owned(), val.clone()).await?;
            let v1 = blobstore.get(ctx, key).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 0);
            }

            assert_eq!(v1.unwrap().as_bytes(), &val);

            Ok(())
        }

        #[fbinit::test]
        async fn test_dedupe_reads_not_storable(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                reject_all_filter,
            )?;

            let key = "foo";
            let val = BlobstoreBytes::from_bytes("foo");

            let (sender, _) = broadcast::channel(1);
            assert_eq!(sender.receiver_count(), 0); // Nothing is waiting here yet

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                blob.data = Some(BlobData::Channel(sender.clone()));
            }

            // Spawn a bunch of reads
            let futs = tokio::spawn(futures::future::try_join_all((0..10usize).map(|_| {
                cloned!(blobstore, ctx);
                async move { blobstore.get(&ctx, key).await }
            })));

            tokio::time::timeout(Duration::from_millis(TIMEOUT_MS), async {
                // Wait for the first request to arrive. It'll be alone, since at this point we don't
                // know this is not cacheable.
                loop {
                    tokio::task::yield_now().await;
                    let count = sender.receiver_count();

                    if count > 1 {
                        return Err(anyhow!("Too many first receivers: {}", count));
                    }

                    if count > 0 {
                        sender
                            .send(val.clone())
                            .map_err(|_| anyhow!("First send failed"))?;

                        break;
                    }
                }

                // Wait for the next requests to arrive. At this point, we know this is not cacheable,
                // but they should join together and share the response.
                loop {
                    tokio::task::yield_now().await;
                    let count = sender.receiver_count();

                    if count > 1 {
                        return Err(anyhow!("Too many second receivers: {}", count));
                    }

                    if count > 0 {
                        sender
                            .send(val.clone())
                            .map_err(|_| anyhow!("Second send failed"))?;
                        break;
                    }
                }

                // Now, spawn a bunch more tasks.  Once again they'll share the request so there
                // should only be a single receiver.
                //
                // These tasks are a bit different from the ones we had already spawned, since they'll
                // check the cache *before* acquiring the semaphore, and won't ever try to acquire it
                // (whereas the other ones would have acquired it, and been released by the first task
                // afterwards).
                let futs = tokio::spawn(futures::future::try_join_all((0..10usize).map(|_| {
                    cloned!(blobstore, ctx);
                    async move { blobstore.get(&ctx, key).await }
                })));

                // Finally, wait for those requests to arrive.
                loop {
                    tokio::task::yield_now().await;
                    let count = sender.receiver_count();

                    if count > 1 {
                        return Err(anyhow!("Too many third receivers: {}", count));
                    }

                    if count > 0 {
                        sender
                            .send(val.clone())
                            .map_err(|_| anyhow!("Third send failed"))?;
                        break;
                    }
                }

                // Check our results
                let res = futs.await??;
                assert_eq!(res.len(), 10);
                for v in res {
                    assert_eq!(v.unwrap().as_bytes(), &val);
                }

                Result::<_>::Ok(())
            })
            .await??;

            // Check our results for the earlier calls.
            let res = futs.await??;
            assert_eq!(res.len(), 10);
            for v in res {
                assert_eq!(v.unwrap().as_bytes(), &val);
            }

            Ok(())
        }

        #[fbinit::test]
        async fn test_dedupe_writes(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            let key = "foo";
            let val = BlobstoreBytes::from_bytes("foo");

            futures::future::try_join_all(
                (0..10usize).map(|_| blobstore.put(ctx, key.to_owned(), val.clone())),
            )
            .await?;

            let handle = {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.puts, 1);
                blob.data.as_ref().unwrap().handle()
            };
            assert_eq!(handle.bytes().await?, val);

            futures::future::try_join_all((0..10usize).map(|_| blobstore.get(ctx, key))).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 0);
            }

            Ok(())
        }

        #[fbinit::test]
        async fn test_dedupe_writes_different_data(fb: FacebookInit) -> Result<()> {
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);
            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            let key = "foo";
            let val0 = BlobstoreBytes::from_bytes("foo");
            let val1 = BlobstoreBytes::from_bytes("foo");
            let val2 = BlobstoreBytes::from_bytes("bar");

            // First, populate the presence cache with some get data.

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                blob.data = Some(BlobData::Bytes(val0.clone()));
            }

            blobstore.get(ctx, key).await?;

            // Now, check that a put still goes through.

            blobstore.put(ctx, key.to_owned(), val1.clone()).await?;

            let handle1 = {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.puts, 1);
                blob.data.as_ref().unwrap().handle()
            };
            assert_eq!(handle1.bytes().await?, val1);

            // Put it again. This time, we don't expect a put to make it through to the blobstore.
            blobstore.put(ctx, key.to_owned(), val1.clone()).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.puts, 1);
            };

            // Now, check that a put for different data also goes through.

            blobstore.put(ctx, key.to_owned(), val2.clone()).await?;

            let handle2 = {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.puts, 2);
                blob.data.as_ref().unwrap().handle()
            };
            assert_eq!(handle2.bytes().await?, val2);

            // Finally, "evict" all the cached data by creating a new blobstore with a new blob
            // pool (but with the same presence pool), and check that a get doesn't overwrite the
            // put state after a get.

            let blobstore = make_blobstore(
                fb,
                TestBlobstore::new(),
                "blobs2",
                nonzero!(2usize),
                allow_all_filter,
            )?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                blob.data = Some(BlobData::Bytes(val0.clone()));
            };

            blobstore.get(ctx, key).await?;

            blobstore.put(ctx, key.to_owned(), val2.clone()).await?;

            {
                let mut data = blobstore.inner.blobstore.data.lock().unwrap();
                let blob = data.entry(key.to_owned()).or_default();
                assert_eq!(blob.gets, 1);
                assert_eq!(blob.puts, 0);
            };

            Ok(())
        }
    }

    mod ratelimiting {
        use std::time::Duration;

        use async_limiter::AsyncLimiter;
        use borrowed::borrowed;
        use context::SessionContainer;
        use fbinit::FacebookInit;
        use governor::Quota;
        use governor::RateLimiter;
        use nonzero_ext::nonzero;
        use scuba_ext::MononokeScubaSampleBuilder;
        use slog::o;
        use slog::Drain;
        use slog::Level;
        use slog::Logger;
        use slog_glog_fmt::default_drain;

        use super::*;

        #[derive(Clone, Debug)]
        struct DummyBlob;

        impl DummyBlob {
            fn new() -> Self {
                Self
            }
        }

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

        #[async_trait]
        impl Blobstore for DummyBlob {
            async fn get<'a>(
                &'a self,
                _ctx: &'a CoreContext,
                _key: &'a str,
            ) -> Result<Option<BlobstoreGetData>> {
                Ok(Some(BlobstoreGetData::new(
                    BlobstoreMetadata::default(),
                    BlobstoreBytes::from_bytes("foo"),
                )))
            }

            async fn put<'a>(
                &'a self,
                _ctx: &'a CoreContext,
                _key: String,
                _value: BlobstoreBytes,
            ) -> Result<()> {
                Ok(())
            }

            async fn is_present<'a>(
                &'a self,
                _ctx: &'a CoreContext,
                _key: &'a str,
            ) -> Result<BlobstoreIsPresent> {
                Ok(BlobstoreIsPresent::Present)
            }
        }

        fn logger() -> Logger {
            let drain = default_drain().filter_level(Level::Debug).ignore_res();
            Logger::root(drain, o![])
        }

        #[fbinit::test]
        async fn test_qps(fb: FacebookInit) -> Result<()> {
            let l1 = RateLimiter::direct(Quota::with_period(Duration::from_millis(10)).unwrap());
            let l1 = AsyncLimiter::new(l1).await;
            let l2 = RateLimiter::direct(Quota::with_period(Duration::from_millis(10)).unwrap());
            let l2 = AsyncLimiter::new(l2).await;
            let session = SessionContainer::builder(fb)
                .blobstore_read_limiter(l1)
                .blobstore_write_limiter(l2)
                .build();
            let ctx = session.new_context(logger(), MononokeScubaSampleBuilder::with_discard());

            let blobstore = make_blobstore(
                fb,
                DummyBlob::new(),
                "blobs",
                nonzero!(100usize),
                allow_all_filter,
            )?;

            borrowed!(ctx, blobstore);

            // get
            let (stats, _) = futures::future::try_join_all((0..10u64).map(|i| {
                let key = format!("get{}", i);
                async move { blobstore.get(ctx, &key).await }
            }))
            .try_timed()
            .await?;
            assert!(stats.completion_time.as_millis_unchecked() > 50);

            // is_present
            let (stats, _) = futures::future::try_join_all((0..10u64).map(|i| {
                let key = format!("present{}", i);
                async move { blobstore.is_present(ctx, &key).await }
            }))
            .try_timed()
            .await?;
            assert!(stats.completion_time.as_millis_unchecked() > 50);

            // put
            let bytes = BlobstoreBytes::from_bytes("test foobar");
            let (stats, _) = futures::future::try_join_all(
                (0..10u64).map(|i| blobstore.put(ctx, format!("put{}", i), bytes.clone())),
            )
            .try_timed()
            .await?;
            assert!(stats.completion_time.as_millis_unchecked() > 50);

            Ok(())
        }

        #[fbinit::test]
        async fn test_early_cache_hits_do_not_count(fb: FacebookInit) -> Result<()> {
            let l1 = RateLimiter::direct(
                Quota::per_second(nonzero!(100u32)).allow_burst(nonzero!(10u32)),
            );
            let l1 = AsyncLimiter::new(l1).await;
            let l2 = RateLimiter::direct(
                Quota::per_second(nonzero!(100u32)).allow_burst(nonzero!(10u32)),
            );
            let l2 = AsyncLimiter::new(l2).await;

            let session = SessionContainer::builder(fb)
                .blobstore_read_limiter(l1)
                .blobstore_write_limiter(l2)
                .build();
            let ctx = &session.new_context(logger(), MononokeScubaSampleBuilder::with_discard());
            borrowed!(ctx);

            let blobstore = &make_blobstore(
                fb,
                DummyBlob::new(),
                "blobs",
                nonzero!(100usize),
                allow_all_filter,
            )?;

            // get
            let (stats, _) = futures::future::try_join_all((0..10u64).flat_map(|i| {
                (0..10u64).map(move |_| {
                    let key = format!("get{}", i);
                    async move { blobstore.get(ctx, &key).await }
                })
            }))
            .try_timed()
            .await?;
            assert!(stats.completion_time.as_millis_unchecked() <= 100);

            // put
            let bytes = &BlobstoreBytes::from_bytes("test foobar");
            let (stats, _) = futures::future::try_join_all((0..10u64).flat_map(|i| {
                (0..10u64).map(move |_| blobstore.put(ctx, format!("put{}", i), bytes.clone()))
            }))
            .try_timed()
            .await?;
            assert!(stats.completion_time.as_millis_unchecked() <= 100);

            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::borrow::Cow;

use anyhow::Error;
use async_limiter::AsyncLimiter;
use context::CoreContext;
use context::PerfCounterType;
use futures::future::BoxFuture;
use futures::future::Fuse;
use futures::future::FutureExt;
use futures_stats::TimedTryFutureExt;
use time_ext::DurationExt;

#[derive(Copy, Clone)]
pub enum AccessReason {
    Read,
    Write,
}

/// Keep track of whether the ticket has been awaited.  This is used in tests
/// to validate that we always await the ticket.
#[derive(Clone, Eq, PartialEq, Default)]
pub struct CheckAwaited {
    awaited: bool,
}

impl CheckAwaited {
    /// Mark that the ticket that contains this `CheckAwaited` has been
    /// awaited and thus is now safe to drop.
    fn mark(&mut self) {
        self.awaited = true;
    }
}

/// Implementation of drop to assist with validation during tests.  If the
/// `Ticket` that owns this `CheckAwaited` has not been awaited before being
/// dropped, we fail the test.
#[cfg(test)]
impl Drop for CheckAwaited {
    fn drop(&mut self) {
        if !self.awaited {
            panic!("Dropped a Pending ticket. This should normally not happen");
        }
    }
}

/// A state machine representing access to the underlying blobstore. This lets us acquire access,
/// as well as attempt to cancel requests.
pub enum Ticket<'a> {
    /// No ticket was requested or awaited.
    NoTicket,
    /// A ticket was requested, but may not have been awaited yet.
    Pending {
        ctx: Cow<'a, CoreContext>,
        reason: AccessReason,
        access: Fuse<BoxFuture<'static, Result<(), Error>>>,
        limiter: Cow<'a, AsyncLimiter>,
        /// We keep track of whether this was awaited so that in tests we can warn if we're failing
        /// to await things we should be awaiting. This isn't used outside of tests (and shouldn't:
        /// if the runtime is shutting down, our futures can be dropped with those unused tickets),
        /// and is only used to assist in debugging / validation.
        awaited: CheckAwaited,
    },
    /// A ticket was requested and awaited.
    Acquired(Cow<'a, AsyncLimiter>),
}

impl<'a> Ticket<'a> {
    pub fn new(ctx: &'a CoreContext, reason: AccessReason) -> Self {
        let limiter = match reason {
            AccessReason::Read => ctx.session().blobstore_read_limiter(),
            AccessReason::Write => ctx.session().blobstore_write_limiter(),
        };

        let limiter = match limiter {
            Some(limiter) => limiter,
            None => {
                return Self::NoTicket;
            }
        };

        Self::Pending {
            ctx: Cow::Borrowed(ctx),
            reason,
            access: limiter.access().boxed().fuse(),
            limiter: Cow::Borrowed(limiter),
            awaited: CheckAwaited::default(),
        }
    }

    /// Check if this ticket is ready. This might kick off a request to acquire access.
    pub async fn is_ready(&mut self) -> Result<bool, Error> {
        match self {
            Self::NoTicket => Ok(true),
            Self::Pending {
                ref mut access,
                ref mut awaited,
                ..
            } => {
                // Race access future against one that is always ready. Since the select is biased, it will
                // favor the access future if it is indeed ready.
                let ready = futures::select_biased! {
                    r = &mut *access => r.map(|()| true),
                    default => Ok(false)
                }?;
                if ready {
                    awaited.mark();
                }
                Ok(ready)
            }
            Self::Acquired(_) => Ok(true),
        }
    }

    /// Wait for this ticket. Calling finish again on a ticket that has already finished will just
    /// return immediately.
    pub async fn finish(self) -> Result<Ticket<'a>, Error> {
        match self {
            Self::Pending {
                ctx,
                reason,
                access,
                limiter,
                mut awaited,
            } => {
                let (stats, ()) = access.try_timed().await?;
                awaited.mark();

                let counter = match reason {
                    AccessReason::Read => PerfCounterType::BlobGetsAccessWait,
                    AccessReason::Write => PerfCounterType::BlobPutsAccessWait,
                };
                ctx.perf_counters()
                    .add_to_counter(counter, stats.completion_time.as_millis_unchecked() as i64);

                Ok(Self::Acquired(limiter))
            }
            x => Ok(x),
        }
    }

    /// Attempt to relinquish this ticket, if anything had been acquired.
    pub fn cancel(mut self) {
        match self {
            Self::NoTicket => {
                // Nothing to do here: we never acquired anything.
            }
            Self::Pending {
                ref mut awaited, ..
            } => {
                // If we never polled this ticket, there is nothing to do. If we did, then we
                // cannot cancel synchronously (since we didn't wait)
                awaited.mark();
            }
            Self::Acquired(l) => {
                l.cancel();
            }
        };
    }

    /// Convert this ticket to an owned variant so that it can be moved into a
    /// shared future.
    pub fn into_owned(self) -> Ticket<'static> {
        match self {
            Self::NoTicket => Ticket::NoTicket,
            Self::Pending {
                ctx,
                reason,
                access,
                limiter,
                awaited,
            } => Ticket::Pending {
                ctx: Cow::Owned(ctx.into_owned()),
                reason,
                access,
                limiter: Cow::Owned(limiter.into_owned()),
                awaited,
            },
            Self::Acquired(l) => Ticket::Acquired(Cow::Owned(l.into_owned())),
        }
    }
}


/*
 * 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::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::num::NonZeroUsize;

use anyhow::Error;
use context::CoreContext;
use context::PerfCounterType;
use futures_stats::TimedTryFutureExt;
use time_ext::DurationExt;
use tokio::sync::Semaphore;
use tokio::sync::SemaphorePermit;

use crate::ratelimit::Ticket;

/// Returned by acquire. Indicates if a semaphore was acquired or not.
pub enum SemaphoreAcquisition<'a, T> {
    Acquired(SemaphorePermit<'a>),
    Cancelled(T, Ticket<'a>),
}

struct ShardHandle<'a> {
    ctx: &'a CoreContext,
    semaphore: &'a Semaphore,
    perf_counter_type: PerfCounterType,
}

impl<'a> ShardHandle<'a> {
    async fn acquire(&self) -> Result<SemaphorePermit<'a>, Error> {
        let (stats, permit) = self.semaphore.acquire().try_timed().await?;

        self.ctx.perf_counters().add_to_counter(
            self.perf_counter_type,
            stats.completion_time.as_millis_unchecked() as i64,
        );

        Ok(permit)
    }
}

pub struct Shards {
    semaphores: Vec<Semaphore>,
    perf_counter_type: PerfCounterType,
}

impl Shards {
    pub fn new(shard_count: NonZeroUsize, perf_counter_type: PerfCounterType) -> Self {
        let semaphores = (0..shard_count.get()).map(|_| Semaphore::new(1)).collect();

        Self {
            semaphores,
            perf_counter_type,
        }
    }

    pub fn len(&self) -> usize {
        self.semaphores.len()
    }

    fn handle<'a>(&'a self, ctx: &'a CoreContext, key: &str) -> ShardHandle<'a> {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        let semaphore = &self.semaphores[(hasher.finish() % self.semaphores.len() as u64) as usize];
        ShardHandle {
            ctx,
            semaphore,
            perf_counter_type: self.perf_counter_type,
        }
    }

    pub async fn acquire<'a, T, D>(
        &'a self,
        ctx: &'a CoreContext,
        key: &str,
        mut ticket: Ticket<'a>,
        determinator: D,
    ) -> Result<SemaphoreAcquisition<'a, T>, Error>
    where
        D: Fn() -> Result<Option<T>, Error>,
    {
        let handle = self.handle(ctx, key);

        //  Await the semaphore once, then check if our determinator says we should proceed.

        let permit = handle.acquire().await?;

        if let Some(r) = determinator()? {
            return Ok(SemaphoreAcquisition::Cancelled(r, ticket));
        }

        // We need to proceed. Check our rate limit.

        if ticket.is_ready().await? {
            return Ok(SemaphoreAcquisition::Acquired(permit));
        }

        // We are rate limited. Drop our permit, and wait for the rate limit.

        drop(permit);
        let ticket = ticket.finish().await?;

        // Check our determinator again. This isn't stricly speaking necessary, but it makes it
        // easier to implement callsites so that they don't have to bother with where exactly we
        // chose to cancel.

        let permit = handle.acquire().await?;

        if let Some(r) = determinator()? {
            return Ok(SemaphoreAcquisition::Cancelled(r, ticket));
        }

        Ok(SemaphoreAcquisition::Acquired(permit))
    }
}


/*
 * 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(slice_group_by)]

mod write_ahead_log;

use sql::mysql;
use sql::mysql_async::prelude::ConvIr;
use sql::mysql_async::prelude::FromValue;
use sql::mysql_async::FromValueError;
use sql::mysql_async::Value;
use uuid::Uuid;
pub use write_ahead_log::BlobstoreWal;
pub use write_ahead_log::BlobstoreWalEntry;
pub use write_ahead_log::SqlBlobstoreWal;

// Identifier for given blobstore operation to faciliate correlating same operation
// across multiple blobstores.
#[derive(Clone, Debug, Eq, PartialEq, Hash, mysql::OptTryFromRowField)]
pub struct OperationKey(pub Uuid);
impl OperationKey {
    pub fn gen() -> OperationKey {
        OperationKey(Uuid::new_v4())
    }

    pub fn is_null(&self) -> bool {
        self == &OperationKey(Uuid::nil())
    }
}

impl From<OperationKey> for Value {
    fn from(id: OperationKey) -> Self {
        let OperationKey(uuid) = id;
        Value::Bytes(uuid.as_bytes().to_vec())
    }
}

impl ConvIr<OperationKey> for OperationKey {
    fn new(v: Value) -> Result<Self, FromValueError> {
        match v {
            Value::Bytes(bytes) => Ok(OperationKey(
                Uuid::from_slice(&bytes[..])
                    .map_err(move |_| FromValueError(Value::Bytes(bytes)))?,
            )),
            v => Err(FromValueError(v)),
        }
    }

    fn commit(self) -> Self {
        self
    }

    fn rollback(self) -> Value {
        self.into()
    }
}

impl FromValue for OperationKey {
    type Intermediate = OperationKey;
}


/*
 * 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 metaconfig_types::MultiplexId;
use mononoke_types::Timestamp;
use sql_ext::mononoke_queries;

use crate::OperationKey;

mononoke_queries! {
    pub(crate) write WalInsertEntry(values: (
        blobstore_key: String,
        multiplex_id: MultiplexId,
        timestamp: Timestamp,
        operation_key: OperationKey,
        blob_size: Option<u64>,
    )) {
        none,
        "INSERT INTO blobstore_write_ahead_log (blobstore_key, multiplex_id, timestamp, operation_key, blob_size)
         VALUES {values}"
    }

    // In comparison to the sync-queue, we write blobstore keys to the WAL only once
    // during the `put` operation. This way when the healer reads entries from the WAL,
    // it doesn't need to filter out distinct operation keys and then blobstore keys
    // (because each blobstore key can have multiple appearances with the same and
    // with different operation keys).
    // The healer can just read all the entries older than the timestamp and they will
    // represent a set of different put opertions by design.
    pub(crate) read WalReadEntries(multiplex_id: MultiplexId, older_than: Timestamp, limit: usize) -> (
        String,
        MultiplexId,
        Timestamp,
        OperationKey,
        u64,
        Option<u64>,
    ) {
        "SELECT blobstore_key, multiplex_id, timestamp, operation_key, id, blob_size
         FROM blobstore_write_ahead_log
         WHERE multiplex_id = {multiplex_id} AND timestamp <= {older_than}
         LIMIT {limit}
         "
    }
}


/*
 * 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::ops::Index;
use std::slice::SliceIndex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use auto_impl::auto_impl;
use context::CoreContext;
use futures::channel::mpsc;
use futures::channel::oneshot;
use futures::future;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::future::Shared;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use metaconfig_types::MultiplexId;
use mononoke_types::Timestamp;
use rand::Rng;
use rendezvous::ConfigurableRendezVousController;
use rendezvous::RendezVous;
use rendezvous::RendezVousOptions;
use rendezvous::RendezVousStats;
use shared_error::anyhow::IntoSharedError;
use shared_error::anyhow::SharedError;
use sql::Connection;
use sql::WriteResult;
use sql_common::mysql::IsolationLevel;
use sql_construct::SqlShardedConstruct;
use sql_ext::mononoke_queries;
use sql_ext::SqlShardedConnections;
use tunables::tunables;
use vec1::Vec1;

const SQL_WAL_WRITE_BUFFER_SIZE: usize = 1000;

/// Row id of the entry, and which SQL shard it belongs to.
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct ReadInfo {
    /// Present if this entry was obtained from reading from SQL
    pub id: Option<u64>,
    /// Present on reads and also updated on writes
    pub shard_id: Option<usize>,
}

impl ReadInfo {
    fn into_present_tuple(self) -> Option<(u64, usize)> {
        self.id.zip(self.shard_id)
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct BlobstoreWalEntry {
    pub blobstore_key: String,
    pub multiplex_id: MultiplexId,
    pub timestamp: Timestamp,
    pub read_info: ReadInfo,
    pub blob_size: u64,
    pub retry_count: u32,
}

impl BlobstoreWalEntry {
    pub fn new(
        blobstore_key: String,
        multiplex_id: MultiplexId,
        timestamp: Timestamp,
        blob_size: u64,
    ) -> Self {
        Self {
            blobstore_key,
            multiplex_id,
            timestamp,
            blob_size,
            read_info: ReadInfo {
                id: None,
                shard_id: None,
            },
            retry_count: 0,
        }
    }

    pub fn increment_retry(&mut self) {
        self.retry_count += 1;
    }

    fn into_sql_tuple(self) -> (String, MultiplexId, Timestamp, u64, u32) {
        let Self {
            blobstore_key,
            multiplex_id,
            timestamp,
            blob_size,
            retry_count,
            ..
        } = self;
        (
            blobstore_key,
            multiplex_id,
            timestamp,
            blob_size,
            retry_count,
        )
    }

    fn from_row(shard_id: usize, row: (String, MultiplexId, Timestamp, u64, u64, u32)) -> Self {
        let (blobstore_key, multiplex_id, timestamp, id, blob_size, retry_count) = row;
        Self {
            blobstore_key,
            multiplex_id,
            timestamp,
            read_info: ReadInfo {
                id: Some(id),
                shard_id: Some(shard_id),
            },
            blob_size,
            retry_count,
        }
    }
}

type QueueResult = Result<BlobstoreWalEntry, SharedError>;
type EnqueueSender = mpsc::UnboundedSender<(oneshot::Sender<QueueResult>, BlobstoreWalEntry)>;

const DEL_CHUNK: usize = 10_000;

#[derive(Clone)]
pub struct SqlBlobstoreWal {
    read_master_connections: Vec1<Connection>,
    write_connections: Arc<Vec1<Connection>>,
    /// Sending entry over the channel allows it to be queued till
    /// the worker is free and able to write new entries to Mysql.
    enqueue_entry_sender: EnqueueSender,
    /// Worker allows to enqueue new entries while there is already
    /// a write query to Mysql in-fight.
    ensure_worker_scheduled: Shared<BoxFuture<'static, ()>>,
    /// Used to cycle through shards when reading
    conn_idx: Arc<AtomicUsize>,
    /// Used to batch deletions together and not overwhelm db with too many queries
    delete_rendezvous: RendezVous<BlobstoreWalEntry, (), ConfigurableRendezVousController>,
}

impl SqlBlobstoreWal {
    fn setup_worker(
        write_connections: Arc<Vec1<Connection>>,
    ) -> (EnqueueSender, Shared<BoxFuture<'static, ()>>) {
        // The mpsc channel needed as a way to enqueue new entries while there is an
        // in-flight write query to Mysql.
        // - queue_sender will be used to add new entries to the queue (channel),
        // - receiver - to read a new batch of entries and write them to Mysql.
        //
        // To notify the clients back that the entry was successfully written to Mysql,
        // a oneshot channel is used. When the enqueued entries are written, the clients
        // receive result of the operation:
        // error if something went wrong and nothing if it's ok.
        let (queue_sender, receiver) =
            mpsc::unbounded::<(oneshot::Sender<QueueResult>, BlobstoreWalEntry)>();

        let worker = async move {
            let mut conn_idx = rand::thread_rng().gen_range(0..write_connections.len());
            let enqueued_writes = receiver.ready_chunks(SQL_WAL_WRITE_BUFFER_SIZE).for_each(
                move |batch /* (Sender, BlobstoreWalEntry) */| {
                    conn_idx += 1;
                    let shard_id = conn_idx % write_connections.len();
                    let write_connection = write_connections[shard_id].clone();
                    async move {
                        let (senders, entries): (Vec<_>, Vec<_>) = batch.into_iter().unzip();

                        let result = insert_entries(&write_connection, &entries).await;
                        let result = result
                            .map_err(|err| err.context("Failed to insert to WAL").shared_error());
                        // We don't really need WriteResult data as we write in batches
                        let result = result.map(|_write_result| ());

                        // Update the clients
                        senders.into_iter().zip(entries).for_each(|(s, mut entry)| {
                            entry.read_info.shard_id = Some(shard_id);
                            // ignore the error, because receiver might have gone
                            let _ = s.send(result.clone().map(|()| entry));
                        });
                    }
                },
            );
            tokio::spawn(enqueued_writes);
        }
        .boxed()
        .shared();

        (queue_sender, worker)
    }

    /// Only read from a certain range of shards. Notice that writes still go to all shards.
    /// Fails if range is empty. Panics if range is out of bounds.
    pub fn with_read_range(
        self,
        range: impl SliceIndex<[Connection], Output = [Connection]>,
    ) -> Result<Self> {
        let read = self.read_master_connections.index(range);
        Ok(Self {
            read_master_connections: read.to_vec().try_into()?,
            ..self
        })
    }

    // This doesn't do any automatic rendezvous/queueing
    async fn inner_delete_by_key(
        write_connections: &[Connection],
        entries: HashSet<BlobstoreWalEntry>,
    ) -> Result<()> {
        // We're grouping by shard id AND multiplex id
        type GroupBy = (usize, MultiplexId);
        let mut del_info: Vec<(GroupBy, &String)> = entries
            .iter()
            .map(|entry| {
                let shard_id = entry
                    .read_info
                    .shard_id
                    .context("BlobstoreWalEntry must have `shard_id` to delete by key")?;
                Ok(((shard_id, entry.multiplex_id), &entry.blobstore_key))
            })
            .collect::<Result<_>>()?;
        del_info.sort_unstable_by_key(|(group, _)| *group);
        stream::iter(
            del_info
                .group_by(|(group1, _), (group2, _)| group1 == group2)
                .map(|batch| async move {
                    let (shard_id, multiplex_id) = batch[0].0;
                    let del_entries: Vec<String> =
                        batch.iter().map(|(_, key)| (*key).clone()).collect();
                    for chunk in del_entries.chunks(DEL_CHUNK) {
                        WalDeleteKeys::query(&write_connections[shard_id], &multiplex_id, chunk)
                            .await?;
                    }
                    anyhow::Ok(())
                })
                .collect::<Vec<_>>(), // prevents compiler bug
        )
        .buffered(10)
        .try_collect::<()>()
        .await
    }
}

#[async_trait]
#[auto_impl(Arc, Box)]
pub trait BlobstoreWal: Send + Sync {
    async fn log<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entry: BlobstoreWalEntry,
    ) -> Result<BlobstoreWalEntry> {
        self.log_many(ctx, vec![entry])
            .await?
            .into_iter()
            .next()
            .context("Missing entry")
    }

    async fn log_many<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entry: Vec<BlobstoreWalEntry>,
    ) -> Result<Vec<BlobstoreWalEntry>>;

    async fn read<'a>(
        &'a self,
        ctx: &'a CoreContext,
        multiplex_id: &MultiplexId,
        older_than: &Timestamp,
        limit: usize,
    ) -> Result<Vec<BlobstoreWalEntry>>;

    /// Entries must have `id` and `shard_id` set (automatic when they are obtained from `read`)
    async fn delete<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entries: &'a [BlobstoreWalEntry],
    ) -> Result<()>;

    /// Entries must have `shard_id` set (automatic when obtained from `log` or `log_many`)
    /// Will delete ALL entries with the same multiplex and key, independent of other fields.
    async fn delete_by_key(&self, ctx: &CoreContext, entries: &[BlobstoreWalEntry]) -> Result<()>;
}

#[async_trait]
impl BlobstoreWal for SqlBlobstoreWal {
    async fn log_many<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        entries: Vec<BlobstoreWalEntry>,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        self.ensure_worker_scheduled.clone().await;

        // If we want to optimize, we can avoid creating a oneshot for each entry by batching together.
        let write_futs = entries
            .iter()
            .cloned()
            .map(|entry| {
                let (sender, receiver) = oneshot::channel();
                self.enqueue_entry_sender.unbounded_send((sender, entry))?;
                Ok(receiver)
            })
            .collect::<Result<Vec<oneshot::Receiver<QueueResult>>>>()?;

        let write_results = future::try_join_all(write_futs)
            .map_err(|err| {
                // If one of the futures fails to receive result from the sql wal,
                // we cannot be sure that the entries were written Mysql WAL table.
                format_err!(
                    "Failed to receive results from the SqlBlobstoreWal: {:?}",
                    err
                )
            })
            .await?;

        write_results
            .into_iter()
            .collect::<Result<_, _>>()
            .context("Failed to write to the SqlBlobstoreWal")
    }

    async fn read<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        multiplex_id: &MultiplexId,
        older_than: &Timestamp,
        mut limit: usize,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        let mut entries = Vec::new();
        // Traverse shards in order fetching from them.
        // To optimise this you can start multiple jobs, one for each shard.
        let shards = self.read_master_connections.len();
        for _ in 0..shards {
            let cur_shard = self.conn_idx.fetch_add(1, Ordering::Relaxed) % shards;
            let rows = WalReadEntries::query(
                &self.read_master_connections[cur_shard],
                multiplex_id,
                older_than,
                &limit,
            )
            .await?;
            limit = limit.saturating_sub(rows.len());
            entries.extend(
                rows.into_iter()
                    .map(|r| BlobstoreWalEntry::from_row(cur_shard, r)),
            );
            if limit == 0 {
                break;
            }
        }
        Ok(entries)
    }

    async fn delete<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        entries: &'a [BlobstoreWalEntry],
    ) -> Result<()> {
        let mut entry_info: Vec<(u64, usize)> =
            entries
                .iter()
                .map(|entry| {
                    entry.read_info.clone().into_present_tuple().context(
                        "BlobstoreWalEntry must contain `read_info` to be able to delete it",
                    )
                })
                .collect::<Result<_, _>>()?;
        entry_info.sort_unstable_by_key(|(_, shard_id)| *shard_id);
        stream::iter(
            entry_info
                .group_by(|(_, shard_id1), (_, shard_id2)| shard_id1 == shard_id2)
                .map(|batch| async move {
                    let shard_id: usize = batch[0].1;
                    let ids: Vec<u64> = batch.iter().map(|(id, _)| *id).collect();
                    for chunk in ids.chunks(10_000) {
                        WalDeleteEntries::query(&self.write_connections[shard_id], chunk).await?;
                    }
                    anyhow::Ok(())
                })
                .collect::<Vec<_>>(), // prevents compiler bug
        )
        .buffered(10)
        .try_collect::<()>()
        .await
    }

    async fn delete_by_key(&self, ctx: &CoreContext, entries: &[BlobstoreWalEntry]) -> Result<()> {
        if !tunables()
            .wal_disable_rendezvous_on_deletes()
            .unwrap_or_default()
        {
            self.delete_rendezvous
                .dispatch(ctx.fb, entries.iter().cloned().collect(), || {
                    let connections = self.write_connections.clone();
                    |keys| async move {
                        Self::inner_delete_by_key(&connections, keys).await?;
                        // We don't care about results
                        Ok(HashMap::new())
                    }
                })
                .await?;
        } else {
            Self::inner_delete_by_key(&self.write_connections, entries.iter().cloned().collect())
                .await?;
        }
        Ok(())
    }
}

impl SqlShardedConstruct for SqlBlobstoreWal {
    const LABEL: &'static str = "blobstore_wal";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-blobstore-wal.sql");

    fn from_sql_shard_connections(connections: SqlShardedConnections) -> Self {
        let SqlShardedConnections {
            read_connections: _,
            mut read_master_connections,
            mut write_connections,
        } = connections;
        if !tunables().disable_wal_read_committed().unwrap_or_default() {
            for conn in read_master_connections
                .iter_mut()
                .chain(write_connections.iter_mut())
            {
                if let Connection::Mysql(conn) = conn {
                    // We don't care about strong locking of the queries, we just a simple "bag".
                    // Having stronger locking causes deadlocks. Let's use a weaker isolation level.
                    // See https://fb.workplace.com/groups/mysql.users/posts/24130604319894856
                    conn.set_isolation_level(Some(IsolationLevel::ReadCommitted));
                }
            }
        }
        let write_connections = Arc::new(write_connections);

        let (sender, ensure_worker_scheduled) =
            SqlBlobstoreWal::setup_worker(write_connections.clone());
        let conn_idx = rand::thread_rng().gen_range(0..read_master_connections.len());

        Self {
            write_connections,
            read_master_connections,
            enqueue_entry_sender: sender,
            ensure_worker_scheduled,
            conn_idx: Arc::new(AtomicUsize::new(conn_idx)),
            // For the delete rendezvous, we don't need to be super fast.
            // - 1 free connection just so we don't wait unnecessarily if traffic is very low
            // - It's fine to wait up to 5 secs to remove, though this likely won't happen.
            // - We're batching underlying requests at 10k
            delete_rendezvous: RendezVous::new(
                ConfigurableRendezVousController::new(
                    RendezVousOptions {
                        free_connections: 1,
                    },
                    || Duration::from_secs(5),
                    || DEL_CHUNK,
                ),
                Arc::new(RendezVousStats::new("wal_delete".to_owned())),
            ),
        }
    }
}

async fn insert_entries(
    write_connection: &Connection,
    entries: &[BlobstoreWalEntry],
) -> Result<WriteResult> {
    let entries: Vec<_> = entries
        .iter()
        .cloned()
        .map(|entry| entry.into_sql_tuple())
        .collect();
    let entries_ref: Vec<_> = entries
        .iter()
        .map(|(a, b, c, d, e)| (a, b, c, d, e)) // &(a, b, ...) into (&a, &b, ...)
        .collect();

    WalInsertEntry::query(write_connection, &entries_ref).await
}

mononoke_queries! {
    write WalDeleteEntries(>list ids: u64) {
        none,
        "DELETE FROM blobstore_write_ahead_log WHERE id in {ids}"
    }

    // This will delete ALL entries for given key. Used for the optimisation where we
    // remove keys from the queue if the write fully succeeded.
    // Ideally we wanted to use `(multiplex_id, blobstore_key) in {entries}` but that's not
    // possible using our SQL libraries, because everything must be a column value, and there are
    // no tuple/array column values (though you could write that query in raw SQL)
    write WalDeleteKeys(multiplex_id: MultiplexId, >list entries: String) {
        none,
        "DELETE FROM blobstore_write_ahead_log WHERE multiplex_id = {multiplex_id} AND blobstore_key in {entries}"
    }

    write WalInsertEntry(values: (
        blobstore_key: String,
        multiplex_id: MultiplexId,
        timestamp: Timestamp,
        blob_size: u64,
        retry_count: u32,
    )) {
        none,
        "INSERT INTO blobstore_write_ahead_log (blobstore_key, multiplex_id, timestamp, blob_size, retry_count)
         VALUES {values}"
    }

    // In comparison to the sync-queue, we write blobstore keys to the WAL only once
    // during the `put` operation. This way when the healer reads entries from the WAL,
    // it doesn't need to filter out distinct operation keys and then blobstore keys
    // (because each blobstore key can have multiple appearances with the same and
    // with different operation keys).
    // The healer can just read all the entries older than the timestamp and they will
    // represent a set of different put opertions by design.
    read WalReadEntries(multiplex_id: MultiplexId, older_than: Timestamp, limit: usize) -> (
        String,
        MultiplexId,
        Timestamp,
        u64,
        u64,
        u32,
    ) {
        "SELECT blobstore_key, multiplex_id, timestamp, id, blob_size, retry_count
         FROM blobstore_write_ahead_log
         WHERE multiplex_id = {multiplex_id} AND timestamp <= {older_than}
         LIMIT {limit}"
    }
}


/*
 * 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::format_err;
use anyhow::Error;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use blobstore_sync_queue::SqlBlobstoreWal;
use context::CoreContext;
use fbinit::FacebookInit;
use metaconfig_types::MultiplexId;
use mononoke_types::DateTime;
use sql_construct::SqlConstruct;

#[fbinit::test]
async fn test_write_ahead_log(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let wal = SqlBlobstoreWal::with_sqlite_in_memory()?;
    let mp = MultiplexId::new(1);

    let key0 = String::from("key0");
    let key1 = String::from("key1");
    let t0 = DateTime::from_rfc3339("2018-11-29T12:00:00.00Z")?.into();
    let t1 = DateTime::from_rfc3339("2018-11-29T12:01:00.00Z")?.into();
    let t2 = DateTime::from_rfc3339("2018-11-29T12:02:00.00Z")?.into();

    let entry0 = BlobstoreWalEntry::new(key0.clone(), mp, t0, 12);
    let entry1 = BlobstoreWalEntry::new(key0, mp, t1, 13);
    let entry2 = BlobstoreWalEntry::new(key1.clone(), mp, t1, 14);
    let entry3 = BlobstoreWalEntry::new(key1, mp, t2, 15);

    // add
    wal.log(&ctx, entry0.clone()).await.unwrap();
    wal.log_many(&ctx, vec![entry1, entry2.clone()])
        .await
        .unwrap();
    wal.log(&ctx, entry3.clone()).await.unwrap();

    // read different ranges of entries
    let validate = |entry: &BlobstoreWalEntry, expected: &BlobstoreWalEntry| {
        assert_eq!(entry.blobstore_key, expected.blobstore_key);
        assert_eq!(entry.multiplex_id, expected.multiplex_id);
        assert_eq!(entry.blob_size, expected.blob_size);
        assert_eq!(entry.retry_count, expected.retry_count);
        assert!(entry.read_info.id.is_some());
        assert!(entry.read_info.shard_id.is_some());
        // read_info is not compared
    };

    let some_entries = wal
        .read(&ctx, &mp, &t0, 1)
        .await
        .expect("DateTime range iteration failed");
    assert_eq!(some_entries.len(), 1);
    validate(
        some_entries
            .get(0)
            .ok_or_else(|| format_err!("must have entry"))?,
        &entry0,
    );

    let mut some_entries = wal
        .read(&ctx, &mp, &t1, 5)
        .await
        .expect("DateTime range iteration failed");
    assert_eq!(some_entries.len(), 3);
    some_entries.sort_by(|a, b| {
        if a.timestamp == b.timestamp {
            a.blobstore_key.cmp(&b.blobstore_key)
        } else {
            a.timestamp.cmp(&b.timestamp)
        }
    });
    validate(
        some_entries
            .get(2)
            .ok_or_else(|| format_err!("must have entry"))?,
        &entry2,
    );

    let mut some_entries = wal
        .read(&ctx, &mp, &t2, 7)
        .await
        .expect("DateTime range iteration failed");
    assert_eq!(some_entries.len(), 4);
    some_entries.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
    validate(
        some_entries
            .get(3)
            .ok_or_else(|| format_err!("must have entry"))?,
        &entry3,
    );

    let some_entries = wal
        .read(&ctx, &mp, &t2, 0)
        .await
        .expect("DateTime range iteration failed");
    assert!(some_entries.is_empty());

    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::HashMap;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Error;
use async_trait::async_trait;
use bonsai_git_mapping::extract_git_sha1_from_bonsai_extra;
use bonsai_git_mapping::BonsaiGitMapping;
use bonsai_git_mapping::BonsaiGitMappingEntry;
use bookmarks::BookmarkTransactionError;
use context::CoreContext;
use mononoke_types::hash::GitSha1;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use pushrebase_hook::PushrebaseCommitHook;
use pushrebase_hook::PushrebaseHook;
use pushrebase_hook::PushrebaseTransactionHook;
use pushrebase_hook::RebasedChangesets;
use sql::Transaction;

#[cfg(test)]
mod test;

#[derive(Clone)]
pub struct GitMappingPushrebaseHook {
    bonsai_git_mapping: Arc<dyn BonsaiGitMapping>,
}

impl GitMappingPushrebaseHook {
    pub fn new(bonsai_git_mapping: Arc<dyn BonsaiGitMapping>) -> Box<dyn PushrebaseHook> {
        Box::new(Self { bonsai_git_mapping })
    }
}

#[async_trait]
impl PushrebaseHook for GitMappingPushrebaseHook {
    async fn in_critical_section(&self) -> Result<Box<dyn PushrebaseCommitHook>, Error> {
        let hook = Box::new(GitMappingCommitHook {
            bonsai_git_mapping: self.bonsai_git_mapping.clone(),
            assignments: HashMap::new(),
        }) as Box<dyn PushrebaseCommitHook>;
        Ok(hook)
    }
}

struct GitMappingCommitHook {
    bonsai_git_mapping: Arc<dyn BonsaiGitMapping>,
    assignments: HashMap<ChangesetId, GitSha1>,
}

#[async_trait]
impl PushrebaseCommitHook for GitMappingCommitHook {
    fn post_rebase_changeset(
        &mut self,
        bcs_old: ChangesetId,
        bcs_new: &mut BonsaiChangesetMut,
    ) -> Result<(), Error> {
        let git_sha1 = extract_git_sha1_from_bonsai_extra(
            bcs_new
                .hg_extra
                .iter()
                .map(|(k, v)| (k.as_str(), v.as_slice())),
        )?;
        if let Some(git_sha1) = git_sha1 {
            self.assignments.insert(bcs_old, git_sha1);
        }
        Ok(())
    }

    async fn into_transaction_hook(
        self: Box<Self>,
        _ctx: &CoreContext,
        rebased: &RebasedChangesets,
    ) -> Result<Box<dyn PushrebaseTransactionHook>, Error> {
        // Let's tie assigned git hashes to rebased Bonsai changesets:
        let entries = self
            .assignments
            .iter()
            .map(|(cs_id, git_sha1)| {
                let replacement_bcs_id = rebased
                    .get(cs_id)
                    .ok_or_else(|| {
                        let e = format!(
                            "Commit was assigned a git hash, but is not found in rebased set: {}",
                            cs_id
                        );
                        Error::msg(e)
                    })?
                    .0;

                Ok(BonsaiGitMappingEntry::new(*git_sha1, replacement_bcs_id))
            })
            .collect::<Result<Vec<_>, Error>>()?;

        // NOTE: This check shouldn't be necessary as long as pushrebase hooks are bug-free, but
        // since they're a new addition, let's be conservative.
        if rebased.len() != self.assignments.len() {
            return Err(anyhow!(
                "Git mapping rebased set ({}) and assignments ({}) have different lengths!",
                rebased.len(),
                self.assignments.len(),
            ));
        }

        Ok(Box::new(GitMappingTransactionHook {
            bonsai_git_mapping: self.bonsai_git_mapping,
            entries,
        }) as Box<dyn PushrebaseTransactionHook>)
    }
}

struct GitMappingTransactionHook {
    bonsai_git_mapping: Arc<dyn BonsaiGitMapping>,
    entries: Vec<BonsaiGitMappingEntry>,
}

#[async_trait]
impl PushrebaseTransactionHook for GitMappingTransactionHook {
    async fn populate_transaction(
        &self,
        ctx: &CoreContext,
        txn: Transaction,
    ) -> Result<Transaction, BookmarkTransactionError> {
        let txn = self
            .bonsai_git_mapping
            .bulk_add_git_mapping_in_transaction(ctx, &self.entries[..], txn)
            .await
            .map_err(|e| BookmarkTransactionError::Other(e.into()))?;
        Ok(txn)
    }
}


/*
 * 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::Error;
use blobstore::Loadable;
use bonsai_git_mapping::BonsaiGitMapping;
use bonsai_git_mapping::BonsaiGitMappingArc;
use bonsai_git_mapping::BonsaiGitMappingRef;
use bonsai_hg_mapping::BonsaiHgMapping;
use bookmarks::Bookmarks;
use borrowed::borrowed;
use changeset_fetcher::ChangesetFetcher;
use changesets::Changesets;
use commit_graph::CommitGraph;
use context::CoreContext;
use fbinit::FacebookInit;
use filestore::FilestoreConfig;
use maplit::hashset;
use mononoke_types::RepositoryId;
use mononoke_types_mocks::hash::*;
use pushrebase::do_pushrebase_bonsai;
use repo_blobstore::RepoBlobstore;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedData;
use repo_identity::RepoIdentity;
use test_repo_factory::TestRepoFactory;
use tests_utils::bookmark;
use tests_utils::CreateCommitContext;

use crate::GitMappingPushrebaseHook;

#[facet::container]
#[derive(Clone)]
struct Repo {
    #[facet]
    bonsai_git_mapping: dyn BonsaiGitMapping,

    #[facet]
    bonsai_hg_mapping: dyn BonsaiHgMapping,

    #[facet]
    bookmarks: dyn Bookmarks,

    #[facet]
    changeset_fetcher: dyn ChangesetFetcher,

    #[facet]
    changesets: dyn Changesets,

    #[facet]
    commit_graph: CommitGraph,

    #[facet]
    filestore_config: FilestoreConfig,

    #[facet]
    repo_blobstore: RepoBlobstore,

    #[facet]
    repo_derived_data: RepoDerivedData,

    #[facet]
    repo_identity: RepoIdentity,
}

#[fbinit::test]
async fn pushrebase_populates_git_mapping(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo: Repo = TestRepoFactory::new(fb)?
        .with_id(RepositoryId::new(1))
        .build()
        .await?;
    borrowed!(ctx, repo);
    let mapping = repo.bonsai_git_mapping();

    let root = CreateCommitContext::new_root(ctx, repo).commit().await?;

    let cs1 = CreateCommitContext::new(ctx, repo, vec![root])
        .commit()
        .await?;

    let cs2 = CreateCommitContext::new(ctx, repo, vec![root])
        .add_extra("hg-git-rename-source".to_owned(), b"git".to_vec())
        .add_extra(
            "convert_revision".to_owned(),
            TWOS_GIT_SHA1.to_hex().as_bytes().to_owned(),
        )
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let book = bookmark(ctx, repo, "master").set_to(cs1).await?;

    let hooks = [GitMappingPushrebaseHook::new(repo.bonsai_git_mapping_arc())];

    let rebased = do_pushrebase_bonsai(
        ctx,
        repo,
        &Default::default(),
        &book,
        &hashset![cs2.clone()],
        &hooks,
    )
    .await?
    .rebased_changesets;

    let cs2_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs2.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs2"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs3 = CreateCommitContext::new(ctx, repo, vec![root])
        .add_extra("hg-git-rename-source".to_owned(), b"git".to_vec())
        .add_extra(
            "convert_revision".to_owned(),
            THREES_GIT_SHA1.to_hex().as_bytes().to_owned(),
        )
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs4 = CreateCommitContext::new(ctx, repo, vec![cs3.get_changeset_id()])
        .add_extra("hg-git-rename-source".to_owned(), b"git".to_vec())
        .add_extra(
            "convert_revision".to_owned(),
            FOURS_GIT_SHA1.to_hex().as_bytes().to_owned(),
        )
        .commit()
        .await?
        .load(ctx, repo.repo_blobstore())
        .await?;

    let rebased = do_pushrebase_bonsai(
        ctx,
        repo,
        &Default::default(),
        &book,
        &hashset![cs3.clone(), cs4.clone()],
        &hooks,
    )
    .await?
    .rebased_changesets;

    let cs3_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs3.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs3"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    let cs4_rebased = rebased
        .iter()
        .find(|e| e.id_old == cs4.get_changeset_id())
        .ok_or_else(|| Error::msg("missing cs4"))?
        .id_new
        .load(ctx, repo.repo_blobstore())
        .await?;

    assert_eq!(
        Some(TWOS_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(ctx, cs2_rebased.get_changeset_id())
            .await?,
    );
    assert_eq!(
        Some(THREES_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(ctx, cs3_rebased.get_changeset_id())
            .await?,
    );
    assert_eq!(
        Some(FOURS_GIT_SHA1),
        mapping
            .get_git_sha1_from_bonsai(ctx, cs4_rebased.get_changeset_id())
            .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 thiserror::Error;

use super::BonsaiGitMappingEntry;

#[derive(Debug, Error)]
pub enum AddGitMappingErrorKind {
    #[error(
        "Conflicting mapping {0:?} detected while inserting git mappings (tried inserting: {1:?})"
    )]
    Conflict(Option<BonsaiGitMappingEntry>, Vec<BonsaiGitMappingEntry>),
    #[error("Internal error occurred while inserting git mapping")]
    InternalError(#[source] anyhow::Error),
}

impl From<anyhow::Error> for AddGitMappingErrorKind {
    fn from(error: anyhow::Error) -> Self {
        AddGitMappingErrorKind::InternalError(error)
    }
}


/*
 * 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 ::sql::Transaction;
use anyhow::Result;
use ascii::AsciiStr;
use async_trait::async_trait;
use context::CoreContext;
use git_types::GitSha1Prefix;
use git_types::GitSha1sResolvedFromPrefix;
use mononoke_types::hash::GitSha1;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use slog::warn;

mod errors;
mod sql;

pub use crate::errors::AddGitMappingErrorKind;
pub use crate::sql::SqlBonsaiGitMapping;
pub use crate::sql::SqlBonsaiGitMappingBuilder;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiGitMappingEntry {
    pub git_sha1: GitSha1,
    pub bcs_id: ChangesetId,
}

impl BonsaiGitMappingEntry {
    pub fn new(git_sha1: GitSha1, bcs_id: ChangesetId) -> Self {
        BonsaiGitMappingEntry { git_sha1, bcs_id }
    }
}

pub enum BonsaisOrGitShas {
    Bonsai(Vec<ChangesetId>),
    GitSha1(Vec<GitSha1>),
}

impl BonsaisOrGitShas {
    pub fn is_empty(&self) -> bool {
        match self {
            BonsaisOrGitShas::Bonsai(v) => v.is_empty(),
            BonsaisOrGitShas::GitSha1(v) => v.is_empty(),
        }
    }
}

impl From<ChangesetId> for BonsaisOrGitShas {
    fn from(cs_id: ChangesetId) -> Self {
        BonsaisOrGitShas::Bonsai(vec![cs_id])
    }
}

impl From<Vec<ChangesetId>> for BonsaisOrGitShas {
    fn from(cs_ids: Vec<ChangesetId>) -> Self {
        BonsaisOrGitShas::Bonsai(cs_ids)
    }
}

impl From<GitSha1> for BonsaisOrGitShas {
    fn from(git_sha1: GitSha1) -> Self {
        BonsaisOrGitShas::GitSha1(vec![git_sha1])
    }
}

impl From<Vec<GitSha1>> for BonsaisOrGitShas {
    fn from(revs: Vec<GitSha1>) -> Self {
        BonsaisOrGitShas::GitSha1(revs)
    }
}

#[facet::facet]
#[async_trait]
pub trait BonsaiGitMapping: Send + Sync {
    async fn bulk_add(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGitMappingEntry],
    ) -> Result<(), AddGitMappingErrorKind>;

    async fn bulk_add_git_mapping_in_transaction(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiGitMappingEntry],
        transaction: Transaction,
    ) -> Result<Transaction, AddGitMappingErrorKind>;

    async fn get(
        &self,
        ctx: &CoreContext,
        field: BonsaisOrGitShas,
    ) -> Result<Vec<BonsaiGitMappingEntry>>;

    async fn get_git_sha1_from_bonsai(
        &self,
        ctx: &CoreContext,
        bcs_id: ChangesetId,
    ) -> Result<Option<GitSha1>> {
        let result = self
            .get(ctx, BonsaisOrGitShas::Bonsai(vec![bcs_id]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.git_sha1))
    }

    async fn get_bonsai_from_git_sha1(
        &self,
        ctx: &CoreContext,
        git_sha1: GitSha1,
    ) -> Result<Option<ChangesetId>> {
        let result = self
            .get(ctx, BonsaisOrGitShas::GitSha1(vec![git_sha1]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.bcs_id))
    }

    async fn get_many_git_sha1_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: GitSha1Prefix,
        limit: usize,
    ) -> Result<GitSha1sResolvedFromPrefix> {
        let mut fetched_cs = self
            .get_in_range(ctx, cs_prefix.min_cs(), cs_prefix.max_cs(), limit + 1)
            .await?;
        let res = match fetched_cs.len() {
            0 => GitSha1sResolvedFromPrefix::NoMatch,
            1 => GitSha1sResolvedFromPrefix::Single(fetched_cs[0].clone()),
            l if l <= limit => GitSha1sResolvedFromPrefix::Multiple(fetched_cs),
            _ => GitSha1sResolvedFromPrefix::TooMany({
                fetched_cs.pop();
                fetched_cs
            }),
        };
        Ok(res)
    }

    async fn bulk_import_from_bonsai(
        &self,
        ctx: &CoreContext,
        changesets: &[BonsaiChangeset],
    ) -> Result<()> {
        let mut entries = vec![];
        for bcs in changesets.iter() {
            match extract_git_sha1_from_bonsai_extra(bcs.hg_extra()) {
                Ok(Some(git_sha1)) => {
                    let entry = BonsaiGitMappingEntry::new(git_sha1, bcs.get_changeset_id());
                    entries.push(entry);
                }
                Ok(None) => {
                    warn!(
                        ctx.logger(),
                        "The git mapping is missing in bonsai commit extras: {:?}",
                        bcs.get_changeset_id()
                    );
                }
                Err(e) => {
                    warn!(ctx.logger(), "Couldn't fetch git mapping: {:?}", e);
                }
            }
        }
        self.bulk_add(ctx, &entries).await?;
        Ok(())
    }

    async fn get_in_range(
        &self,
        ctx: &CoreContext,
        low: GitSha1,
        high: GitSha1,
        limit: usize,
    ) -> Result<Vec<GitSha1>>;
}

pub const HGGIT_SOURCE_EXTRA: &str = "hg-git-rename-source";
pub const CONVERT_REVISION_EXTRA: &str = "convert_revision";

pub fn extract_git_sha1_from_bonsai_extra<'a, 'b, T>(extra: T) -> Result<Option<GitSha1>>
where
    T: Iterator<Item = (&'a str, &'b [u8])>,
{
    let (mut hggit_source_extra, mut convert_revision_extra) = (None, None);
    for (key, value) in extra {
        if key == HGGIT_SOURCE_EXTRA {
            hggit_source_extra = Some(value);
        }
        if key == CONVERT_REVISION_EXTRA {
            convert_revision_extra = Some(value);
        }
    }

    if hggit_source_extra == Some(b"git") {
        if let Some(convert_revision_extra) = convert_revision_extra {
            let git_sha1 = AsciiStr::from_ascii(convert_revision_extra)?;
            let git_sha1 = GitSha1::from_ascii_str(git_sha1)?;
            return Ok(Some(git_sha1));
        }
    }
    Ok(None)
}