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

use std::fmt;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::CountedBlobstore;
use cloned::cloned;
use context::CoreContext;
use context::PerfCounterType;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use mononoke_types::BlobstoreBytes;
use prefixblob::PrefixBlobstore;
use redactedblobstore::config::GET_OPERATION;
use redactedblobstore::RedactedBlobstore;
use stats::prelude::*;

define_stats! {
    prefix = "mononoke.blobstore.cacheblob";
    get_miss: dynamic_timeseries("{}.get_miss", (cache_name: &'static str); Rate, Sum),
    get_hit: dynamic_timeseries("{}.get_hit", (cache_name: &'static str); Rate, Sum),
    presence_hit: dynamic_timeseries("{}.presence_hit", (cache_name: &'static str); Rate, Sum),
    presence_miss: dynamic_timeseries("{}.presence_miss", (cache_name: &'static str); Rate, Sum),
}

/// Extra operations that can be performed on a cache. Other wrappers can implement this trait for
/// e.g. all `WrapperBlobstore<CacheBlobstore<T>>`.
///
/// This is primarily used by the admin command to manually check memcache.
#[async_trait]
pub trait CacheBlobstoreExt: Blobstore {
    async fn get_no_cache_fill<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>>;
    async fn get_cache_only<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>>;
}

/// The operations a cache must provide in order to be usable as the caching layer for a
/// caching blobstore that caches blob contents and blob presence.
/// For caches that do no I/O (e.g. in-memory caches), use Result::into_future() to create the
/// return types - it is up to CacheBlobstore to use future::lazy where this would be unacceptable
/// Errors returned by the cache are always ignored.
///
/// The cache is expected to act as-if each entry is in one of four states:
/// 1. Empty, implying that the cache has no knowledge of the backing store state for this key.
/// 2. Leased, implying that the cache is aware that there is an attempt being made to update the
///    backing store for this key.
/// 3. Present, implying that the cache is aware that a backing store entry exists for this key
///    but does not have a copy of the blob.
/// 4. Known, implying that the cache has a copy of the blob for this key.
///
/// When the cache engages in eviction, it demotes entries according to the following plan:
/// Present and Leased can only demote to Empty.
/// Known can demote to Present or Empty.
/// No state is permitted to demote to Leased.
/// Caches that do not support LeaseOps do not have the Leased state.
#[async_trait]
pub trait CacheOps: fmt::Display + fmt::Debug + Send + Sync {
    const HIT_COUNTER: Option<PerfCounterType> = None;
    const MISS_COUNTER: Option<PerfCounterType> = None;
    const CACHE_NAME: &'static str = "unknown";

    /// Fetch the blob from the cache, if possible. Return `None` if the cache does not have a
    /// copy of the blob (i.e. the cache entry is not in Known state).
    async fn get(&self, key: &str) -> Option<BlobstoreGetData>;

    /// Tell the cache that the backing store value for this `key` is `value`. This should put the
    /// cache entry for this `key` into Known state or a demotion of Known state (Present, Empty).
    async fn put(&self, key: &str, value: BlobstoreGetData);

    /// Ask the cache if it knows whether the backing store has a value for this key. Returns
    /// `true` if there is definitely a value (i.e. cache entry in Present or Known state), `false`
    /// otherwise (Empty or Leased states).
    async fn check_present(&self, key: &str) -> bool;
}

/// The operations a cache must provide to take part in the update lease protocol. This reduces the
/// thundering herd on writes by using the Leased state to ensure that only one user of this cache
/// can write to the backing store at any time. Note that this is not a guarantee that there will
/// be only one writer to the backing store for any given key - notably, the cache can demote
/// Leased to Empty, thus letting another writer that shares the same cache through to the backing
/// store.
#[async_trait]
pub trait LeaseOps: fmt::Display + fmt::Debug + Send + Sync {
    /// Ask the cache to attempt to lock out other users of this cache for a particular key.
    /// This is an atomic test-and-set of the cache entry; it tests that the entry is Empty, and if
    /// the entry is Empty, it changes it to the Leased state.
    /// The result is `true` if the test-and-set changed the entry to Leased state, `false`
    /// otherwise
    async fn try_add_put_lease(&self, key: &str) -> Result<bool>;

    /// Will keep the lease alive until `done` future resolves.
    /// Note that it should only be called after successful try_add_put_lease()
    fn renew_lease_until(&self, ctx: CoreContext, key: &str, done: BoxFuture<'static, ()>);

    /// Wait for a suitable (cache-defined) period between `try_add_put_lease` attempts.
    /// For caches without a notification method, this should just be a suitable delay.
    /// For caches that can notify on key change, this should wait for that notification.
    /// It is acceptable to return from this future without checking the state of the cache entry.
    async fn wait_for_other_leases(&self, key: &str);

    /// Releases any leases held on `key`. The entry must transition from Leased to Empty.
    async fn release_lease(&self, key: &str);
}

/// A caching layer over a blobstore, using a cache defined by its CacheOps. The idea is that
/// generic code that any caching layer needs is defined here, while code that's cache-specific
/// goes into CacheOps
#[derive(Clone)]
pub struct CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone,
    L: LeaseOps + Clone,
    T: Blobstore,
{
    blobstore: T,
    cache: C,
    lease: L,
    lazy_cache_put: bool,
}

impl<C, L, T> fmt::Display for CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone,
    L: LeaseOps + Clone,
    T: Blobstore,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "CacheBlobstore<{}, {}, {}>",
            &self.cache, &self.lease, &self.blobstore
        )
    }
}

impl<C, L, T> CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone,
    L: LeaseOps + Clone,
    T: Blobstore,
{
    pub fn new(cache: C, lease: L, blobstore: T, lazy_cache_put: bool) -> Self {
        Self {
            blobstore,
            cache,
            lease,
            lazy_cache_put,
        }
    }

    fn take_put_lease<'a>(&'a self, key: &'a str) -> BoxFuture<'a, bool> {
        async move {
            if self.lease.try_add_put_lease(key).await.map_err(|_| ()) == Ok(true) {
                return true;
            }
            if self.cache.check_present(key).await {
                return false;
            }
            self.lease.wait_for_other_leases(key).await;
            self.take_put_lease(key).await
        }
        .boxed()
    }
}

#[async_trait]
impl<C, L, T> Blobstore for CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone + 'static,
    L: LeaseOps + Clone + 'static,
    T: Blobstore,
{
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let blob = self.cache.get(key).await;
        if blob.is_some() {
            if let Some(counter) = C::HIT_COUNTER {
                ctx.perf_counters().increment_counter(counter);
            }
            STATS::get_hit.add_value(1, (C::CACHE_NAME,));
            Ok(blob)
        } else {
            if let Some(counter) = C::MISS_COUNTER {
                ctx.perf_counters().increment_counter(counter);
            }
            STATS::get_miss.add_value(1, (C::CACHE_NAME,));
            let blob = self.blobstore.get(ctx, key).await?;
            if let Some(ref blob) = blob {
                let key = key.to_owned();
                cloned!(self.cache, blob);
                tokio::spawn(async move { cache.put(&key, blob).await });
            }
            Ok(blob)
        }
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        let can_put = self.take_put_lease(&key).await;
        if can_put {
            self.blobstore.put(ctx, key.clone(), value.clone()).await?;

            cloned!(self.cache, self.lease);
            let cache_put = async move {
                cache.put(&key, value.into()).await;
                lease.release_lease(&key).await
            };
            if self.lazy_cache_put {
                tokio::spawn(cache_put);
            } else {
                let _ = cache_put.await;
            }
        }
        Ok(())
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let present = self.cache.check_present(key).await;
        if present {
            STATS::presence_hit.add_value(1, (C::CACHE_NAME,));
            Ok(BlobstoreIsPresent::Present)
        } else {
            STATS::presence_miss.add_value(1, (C::CACHE_NAME,));
            self.blobstore.is_present(ctx, key).await
        }
    }
}

#[async_trait]
impl<C, L, T> CacheBlobstoreExt for CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone + 'static,
    L: LeaseOps + Clone + 'static,
    T: Blobstore,
{
    async fn get_no_cache_fill<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let blob = self.cache.get(key).await;
        if blob.is_some() {
            Ok(blob)
        } else {
            self.blobstore.get(ctx, key).await
        }
    }

    async fn get_cache_only<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.cache.get(key).map(Ok).await
    }
}

impl<C, L, T> fmt::Debug for CacheBlobstore<C, L, T>
where
    C: CacheOps + Clone,
    L: LeaseOps + Clone,
    T: Blobstore,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CacheBlobstore")
            .field("blobstore", &self.blobstore)
            .field("cache", &self.cache)
            .field("lease", &self.lease)
            .finish()
    }
}

#[async_trait]
impl<T: CacheBlobstoreExt> CacheBlobstoreExt for CountedBlobstore<T> {
    #[inline]
    async fn get_no_cache_fill<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.as_inner().get_no_cache_fill(ctx, key).await
    }

    #[inline]
    async fn get_cache_only<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.as_inner().get_cache_only(ctx, key).await
    }
}

#[async_trait]
impl<T: CacheBlobstoreExt> CacheBlobstoreExt for PrefixBlobstore<T> {
    #[inline]
    async fn get_no_cache_fill<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.as_inner()
            .get_no_cache_fill(ctx, &self.prepend(key))
            .await
    }

    #[inline]
    async fn get_cache_only<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.as_inner()
            .get_cache_only(ctx, &self.prepend(key))
            .await
    }
}

#[async_trait]
impl<T: CacheBlobstoreExt> CacheBlobstoreExt for RedactedBlobstore<T> {
    #[inline]
    async fn get_no_cache_fill<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let blobstore = self.access_blobstore(ctx, key, GET_OPERATION)?;
        blobstore.get_no_cache_fill(ctx, key).await
    }

    #[inline]
    async fn get_cache_only<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let blobstore = self.access_blobstore(ctx, key, GET_OPERATION)?;
        blobstore.get_cache_only(ctx, key).await
    }
}


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

use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::CountedBlobstore;
use cloned::cloned;
use context::CoreContext;
use context::PerfCounterType;
use fbinit::FacebookInit;
use fbthrift::compact_protocol;
use futures::future::select;
use futures::future::BoxFuture;
use futures::future::Either;
use hostname::get_hostname;
use memcache::KeyGen;
use memcache::MemcacheClient;
use memcache_lock_thrift::LockState;
use slog::warn;
use stats::prelude::*;
use tunables::tunables;

use crate::dummy::DummyLease;
use crate::CacheBlobstore;
use crate::CacheOps;
use crate::LeaseOps;

define_stats! {
    prefix = "mononoke.blobstore.memcache";
    blob_put: timeseries("blob_put"; Rate, Sum),
    blob_put_err: timeseries("blob_put_err"; Rate, Sum),
    presence_put: timeseries("presence_put"; Rate, Sum),
    presence_put_err: timeseries("presence_put_err"; Rate, Sum),
    blob_presence: timeseries("blob_presence"; Rate, Sum),
    blob_presence_hit: timeseries("blob_presence_hit"; Rate, Sum),
    blob_presence_miss: timeseries("blob_presence_miss"; Rate, Sum),
    blob_presence_err: timeseries("blob_presence_err"; Rate, Sum),
    presence_get: timeseries("presence_get"; Rate, Sum),
    presence_check_hit: timeseries("presence_check_hit"; Rate, Sum),
    presence_check_miss: timeseries("presence_check_miss"; Rate, Sum),
    // This can come from leases as well as presence checking.
    presence_err: timeseries("presence_err"; Rate, Sum),
}

#[allow(non_snake_case)]
mod LEASE_STATS {
    use stats::define_stats;
    define_stats! {
        prefix = "mononoke.blobstore.memcache.lease";
        claim: dynamic_timeseries("{}.claim", (lease_type: &'static str); Rate, Sum),
        claim_err: dynamic_timeseries("{}.claim_err", (lease_type: &'static str); Rate, Sum),
        conflict: dynamic_timeseries("{}.conflict", (lease_type: &'static str); Rate, Sum),
        wait_ms: dynamic_timeseries("{}.wait_ms", (lease_type: &'static str); Rate, Sum),
        release: dynamic_timeseries("{}.release", (lease_type: &'static str); Rate, Sum),
        release_good: dynamic_timeseries("{}.release_good", (lease_type: &'static str); Rate, Sum),
        release_held_by_other: dynamic_timeseries("{}.release_held_by_other", (lease_type: &'static str); Rate, Sum),
        release_bad_key: dynamic_timeseries("{}.release_bad_key", (lease_type: &'static str); Rate, Sum),
        release_key_set: dynamic_timeseries("{}.release_key_set", (lease_type: &'static str); Rate, Sum),
        release_no_lease: dynamic_timeseries("{}.release_no_lease", (lease_type: &'static str); Rate, Sum),
    }
    pub use self::STATS::*;
}

/// A caching layer over an existing blobstore, backed by memcache
#[derive(Clone, Debug)]
pub struct MemcacheOps {
    lease_type: &'static str,
    memcache: MemcacheClient,
    keygen: KeyGen,
    presence_keygen: KeyGen,
    hostname: String,
}

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

const MEMCACHE_MAX_SIZE: usize = 1024000;
const MC_CODEVER: u32 = 0;
const MC_SITEVER: u32 = 1;

async fn mc_raw_put(
    memcache: MemcacheClient,
    orig_key: String,
    key: String,
    value: BlobstoreGetData,
    presence_key: String,
) {
    let uploaded = compact_protocol::serialize(LockState::uploaded_key(orig_key));

    STATS::presence_put.add_value(1);
    // This cache key is read by leases, and if it's set then lease can't be reacquired.
    // To be on the safe side let's add a ttl on this memcache key.
    let lock_ttl = Duration::from_secs(50);
    let res = memcache
        .set_with_ttl(presence_key, uploaded, lock_ttl)
        .await;
    if res.is_err() {
        STATS::presence_put_err.add_value(1);
    }

    if value.as_bytes().len() < MEMCACHE_MAX_SIZE {
        STATS::blob_put.add_value(1);
        let res = memcache.set(key, value.into_raw_bytes()).await;
        if res.is_err() {
            STATS::blob_put_err.add_value(1);
        }
    }
}

impl MemcacheOps {
    pub fn new(
        fb: FacebookInit,
        lease_type: &'static str,
        backing_store_params: impl ToString,
    ) -> Result<Self> {
        let hostname = get_hostname()?;

        let blob_key = format!(
            "scm.mononoke.blobstore.{}.{}",
            lease_type,
            backing_store_params.to_string()
        );

        let presence_key = format!(
            "scm.mononoke.blobstore.presence.{}.{}",
            lease_type,
            backing_store_params.to_string()
        );

        let sitever = if tunables().blobstore_memcache_sitever().unwrap_or_default() > 0 {
            tunables().blobstore_memcache_sitever().unwrap_or_default() as u32
        } else {
            MC_SITEVER
        };

        Ok(Self {
            lease_type,
            memcache: MemcacheClient::new(fb)?,
            keygen: KeyGen::new(blob_key, MC_CODEVER, sitever),
            presence_keygen: KeyGen::new(presence_key, MC_CODEVER, sitever),
            hostname,
        })
    }

    async fn get_lock_state(&self, key: String) -> Option<LockState> {
        let mc_key = self.presence_keygen.key(key.clone());
        STATS::presence_get.add_value(1);
        match self.memcache.get(mc_key.clone()).await {
            Ok(opt_blob) => {
                let blob = opt_blob?;
                #[cfg(fbcode_build)]
                let blob = Vec::from(blob);
                let state = compact_protocol::deserialize(blob).ok()?;
                if let LockState::uploaded_key(ref up_key) = state {
                    if key != *up_key {
                        // The lock state is invalid - fix it up by dropping the lock
                        let _ = self.memcache.del(mc_key).await;
                        return None;
                    }
                }
                Some(state)
            }
            Err(_) => {
                STATS::presence_err.add_value(1);
                None
            }
        }
    }
}

pub fn new_memcache_blobstore<T>(
    fb: FacebookInit,
    blobstore: T,
    backing_store_name: &'static str,
    backing_store_params: impl ToString,
) -> Result<CountedBlobstore<CacheBlobstore<MemcacheOps, DummyLease, T>>>
where
    T: Blobstore + Clone,
{
    let cache_ops = MemcacheOps::new(fb, backing_store_name, backing_store_params)?;
    Ok(CountedBlobstore::new(
        "memcache".to_string(),
        CacheBlobstore::new(cache_ops, DummyLease {}, blobstore, true),
    ))
}

#[async_trait]
impl CacheOps for MemcacheOps {
    const HIT_COUNTER: Option<PerfCounterType> = Some(PerfCounterType::MemcacheHits);
    const MISS_COUNTER: Option<PerfCounterType> = Some(PerfCounterType::MemcacheMisses);
    const CACHE_NAME: &'static str = "memcache";

    // Turns errors to Ok(None)
    async fn get(&self, key: &str) -> Option<BlobstoreGetData> {
        let mc_key = self.keygen.key(key);
        let buf = self.memcache.get(mc_key).await.ok()??;
        Some(BlobstoreGetData::from_bytes(buf))
    }

    async fn put(&self, key: &str, value: BlobstoreGetData) {
        let mc_key = self.keygen.key(key);
        let presence_key = self.presence_keygen.key(key);
        let orig_key = key.to_string();

        mc_raw_put(self.memcache.clone(), orig_key, mc_key, value, presence_key).await
    }

    async fn check_present(&self, key: &str) -> bool {
        let key = key.to_string();
        match self.get_lock_state(key.clone()).await {
            // get_lock_state will delete the lock and return None if there's a bad
            // uploaded_key
            Some(LockState::uploaded_key(_)) => {
                STATS::presence_check_hit.add_value(1);
                true
            }
            _ => {
                STATS::presence_check_miss.add_value(1);
                let mc_key = self.keygen.key(key);
                STATS::blob_presence.add_value(1);
                let blob_presence = self.memcache.get(mc_key).await;
                match blob_presence {
                    Ok(Some(_)) => STATS::blob_presence_hit.add_value(1),
                    Ok(None) => STATS::blob_presence_miss.add_value(1),
                    Err(_) => STATS::blob_presence_err.add_value(1),
                }
                blob_presence.unwrap_or(None).is_some()
            }
        }
    }
}

#[async_trait]
impl LeaseOps for MemcacheOps {
    async fn try_add_put_lease(&self, key: &str) -> Result<bool> {
        let mc_key = self.presence_keygen.key(key);
        let lockstate = compact_protocol::serialize(LockState::locked_by(self.hostname.clone()));
        let lock_ttl = Duration::from_secs(10);
        let lease_type = self.lease_type;
        let res = self
            .memcache
            .add_with_ttl(mc_key, lockstate, lock_ttl)
            .await;
        match res {
            Ok(true) => LEASE_STATS::claim.add_value(1, (lease_type,)),
            Ok(false) => LEASE_STATS::conflict.add_value(1, (lease_type,)),
            Err(_) => LEASE_STATS::claim_err.add_value(1, (lease_type,)),
        }
        res
    }

    fn renew_lease_until(&self, ctx: CoreContext, key: &str, mut done: BoxFuture<'static, ()>) {
        let lockstate = compact_protocol::serialize(LockState::locked_by(self.hostname.clone()));
        let lock_ttl = Duration::from_secs(10);
        let mc_key = self.presence_keygen.key(key);
        let key = key.to_string();
        cloned!(self.memcache);

        let this = self.clone();
        tokio::spawn(async move {
            loop {
                let res = memcache
                    .set_with_ttl(mc_key.clone(), lockstate.clone(), lock_ttl)
                    .await;
                if res.is_err() {
                    warn!(ctx.logger(), "failed to renew lease for {}", mc_key);
                }

                let sleep = tokio::time::sleep(Duration::from_secs(1));
                futures::pin_mut!(sleep);
                let res = select(sleep, done).await;
                match res {
                    Either::Left((_, new_done)) => {
                        done = new_done;
                    }
                    Either::Right(..) => {
                        break;
                    }
                }
            }

            this.release_lease(&key).await;
        });
    }

    async fn wait_for_other_leases(&self, _key: &str) {
        let retry_millis = 200;
        let retry_delay = Duration::from_millis(retry_millis);
        LEASE_STATS::wait_ms.add_value(retry_millis as i64, (self.lease_type,));
        tokio::time::sleep(retry_delay).await;
    }

    async fn release_lease(&self, key: &str) {
        let mc_key = self.presence_keygen.key(key);
        LEASE_STATS::release.add_value(1, (self.lease_type,));
        cloned!(self.memcache, self.hostname, self.lease_type);

        // This future checks the state of the lease, and releases it only
        // if it's locked by us right now.
        let f = async move {
            let bytes = match memcache.get(mc_key.clone()).await {
                Ok(Some(bytes)) => bytes,
                Ok(None) => {
                    LEASE_STATS::release_no_lease.add_value(1, (lease_type,));
                    return;
                }
                Err(_) => return,
            };

            #[cfg(fbcode_build)]
            let bytes = Vec::from(bytes);
            let state: LockState = match compact_protocol::deserialize(bytes) {
                Ok(state) => state,
                Err(_) => {
                    LEASE_STATS::release_bad_key.add_value(1, (lease_type,));
                    // Fix up invalid value
                    let _ = memcache.del(mc_key).await;
                    return;
                }
            };

            match state {
                LockState::locked_by(locked_by) => {
                    if locked_by == hostname {
                        LEASE_STATS::release_good.add_value(1, (lease_type,));
                        // The lease is held by us - just remove it
                        let _ = memcache.del(mc_key).await;
                    } else {
                        LEASE_STATS::release_held_by_other.add_value(1, (lease_type,));
                        // Someone else grabbed a lease, leave it alone
                    }
                }
                LockState::uploaded_key(up_key) => {
                    if up_key != mc_key {
                        LEASE_STATS::release_bad_key.add_value(1, (lease_type,));
                        // Invalid key - fix it up. Normally that shouldn't
                        // ever occur
                        let _ = memcache.del(mc_key).await;
                    } else {
                        LEASE_STATS::release_key_set.add_value(1, (lease_type,));
                        // Key is valid, and is most likely set by
                        // cache.put(...). Lease is already release,
                        // no need to do anything here
                    }
                }
                LockState::UnknownField(_) => {
                    LEASE_STATS::release_bad_key.add_value(1, (lease_type,));
                    // Possibly a newer version of the server enabled it?
                    // Don't touch it just in case
                }
            }
        };
        // We don't have to wait for the releasing to finish, it can be done in background
        // because leases have a timeout. So even if they haven't been released explicitly they
        // will be released after a timeout.
        tokio::spawn(f);
    }
}


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

use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use context::CoreContext;
use futures::future;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use lock_ext::LockExt;
use mononoke_types::BlobstoreBytes;
use tokio::sync::Mutex as AsyncMutex;

#[derive(Default, Debug)]
pub struct Cache {
    live: HashMap<String, BlobstoreBytes>,
    flushing: Option<Arc<HashMap<String, BlobstoreBytes>>>,
}

impl Cache {
    pub fn len(&self) -> usize {
        self.live.len() + self.flushing.as_ref().map_or(0, |flushing| flushing.len())
    }
}

/// A blobstore wrapper that reads from the underlying blobstore but writes to memory.
#[derive(Clone, Debug)]
pub struct MemWritesBlobstore<T> {
    inner: T,
    cache: Arc<Mutex<Cache>>,
    // Mutex to ensure only one task is flushing the cache at a time.
    // Note: this doesn't wrap the cache as read access is permitted while
    // the mutex is held.
    flush_mutex: Arc<AsyncMutex<()>>,
    no_access_to_inner: Arc<AtomicBool>,
}

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

impl<T: Blobstore + Clone> MemWritesBlobstore<T> {
    pub fn new(blobstore: T) -> Self {
        Self {
            inner: blobstore,
            cache: Default::default(),
            flush_mutex: Default::default(),
            no_access_to_inner: Default::default(),
        }
    }

    /// Write all in-memory entries to underlying blobstore.
    ///
    /// NOTE: In case of error all pending changes will be lost.
    pub async fn persist<'a>(&'a self, ctx: &'a CoreContext) -> Result<()> {
        if self.no_access_to_inner.load(Ordering::Relaxed) {
            return Err(anyhow!(
                "unexpected write to memory blobstore when access to inner blobstore was disabled"
            ));
        }

        // Obtain the flush mutex.  This should ensure that only one persist
        // happens at a time.
        let _flush_guard = self.flush_mutex.lock().await;

        let items = self.cache.with(|cache| {
            if cache.flushing.is_some() {
                // This should be prevented by the flush guard.
                return Err(anyhow!(
                    "unexpected persist while another persist is ongoing"
                ));
            }
            let flushing = Arc::new(mem::take(&mut cache.live));
            cache.flushing = Some(flushing.clone());
            Ok(flushing)
        })?;

        let flush = async_stream::stream! {
            for (key, value) in items.iter() {
                 yield self.inner.put(ctx, key.clone(), value.clone());
            }
        };

        let result = flush
            .buffered(4096)
            .try_for_each(|_| future::ready(Ok(())))
            .await;

        // Discard flushing items, whether or not we were successful at
        // flushing the cache.
        self.cache.with(|cache| {
            cache.flushing = None;
        });

        result
    }

    pub fn get_inner(&self) -> T {
        self.inner.clone()
    }

    pub fn get_cache(&self) -> &Arc<Mutex<Cache>> {
        &self.cache
    }

    pub fn set_no_access_to_inner(&self, no_access_to_inner: bool) {
        self.no_access_to_inner
            .store(no_access_to_inner, Ordering::Relaxed);
    }
}

#[async_trait]
impl<T: Blobstore + Clone> Blobstore for MemWritesBlobstore<T> {
    async fn put<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.cache.with(|cache| cache.live.insert(key, value));
        Ok(())
    }

    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let value = self.cache.with(|cache| {
            if let Some(value) = cache.live.get(key) {
                Some(value.clone())
            } else if let Some(flushing) = cache.flushing.as_ref() {
                flushing.get(key).cloned()
            } else {
                None
            }
        });

        match value {
            Some(value) => Ok(Some(value.into())),
            None => {
                if self.no_access_to_inner.load(Ordering::Relaxed) {
                    return Ok(None);
                }

                Ok(self.inner.get(ctx, key).await?.map(Into::into))
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::time::Duration;

    use blobstore::PutBehaviour;
    use borrowed::borrowed;
    use bytes::Bytes;
    use cloned::cloned;
    use fbinit::FacebookInit;
    use memblob::Memblob;
    use tokio::sync::watch;

    use super::*;

    #[fbinit::test]
    async fn basic_read(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let inner = Memblob::default();
        let foo_key = "foo";
        inner
            .put(
                ctx,
                foo_key.to_owned(),
                BlobstoreBytes::from_bytes("foobar"),
            )
            .await
            .expect("initial put should work");
        let outer = MemWritesBlobstore::new(inner.clone());

        assert!(
            outer
                .is_present(ctx, foo_key)
                .await
                .expect("is_present to inner should work")
                .assume_not_found_if_unsure()
        );

        assert_eq!(
            outer
                .get(ctx, foo_key)
                .await
                .expect("get to inner should work")
                .expect("value should be present")
                .into_raw_bytes(),
            Bytes::from("foobar"),
        );
    }

    #[fbinit::test]
    async fn redirect_writes(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let inner = Memblob::default();
        let foo_key = "foo";

        let outer = MemWritesBlobstore::new(inner.clone());
        outer
            .put(
                ctx,
                foo_key.to_owned(),
                BlobstoreBytes::from_bytes("foobar"),
            )
            .await
            .expect("put should work");

        assert!(
            !inner
                .is_present(ctx, foo_key)
                .await
                .expect("is_present on inner should work")
                .assume_not_found_if_unsure(),
            "foo should not be present in inner",
        );

        assert!(
            outer
                .is_present(ctx, foo_key)
                .await
                .expect("is_present on outer should work")
                .assume_not_found_if_unsure(),
            "foo should be present in outer",
        );

        assert_eq!(
            outer
                .get(ctx, foo_key)
                .await
                .expect("get to outer should work")
                .expect("value should be present")
                .into_raw_bytes(),
            Bytes::from("foobar"),
        );
    }

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

        let inner = Memblob::default();
        let outer = MemWritesBlobstore::new(inner.clone());

        let key = "key";
        let value = BlobstoreBytes::from_bytes("value");

        outer.put(ctx, key.to_owned(), value.clone()).await?;

        assert!(inner.get(ctx, key).await?.is_none());

        outer.persist(ctx).await?;

        assert_eq!(inner.get(ctx, key).await?, Some(value.into()));

        Ok(())
    }

    /// A blobstore wrapper that prevents writes until a flag is set.
    #[derive(Clone, Debug)]
    pub struct GatedBlobstore<T> {
        inner: T,
        allow: watch::Receiver<bool>,
    }

    impl<T: Blobstore + Clone> GatedBlobstore<T> {
        pub fn new(blobstore: T, allow: watch::Receiver<bool>) -> Self {
            Self {
                inner: blobstore,
                allow,
            }
        }
    }
    impl<T: std::fmt::Display> std::fmt::Display for GatedBlobstore<T> {
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
            write!(f, "GatedBlobstore<{}>", self.inner)
        }
    }

    #[async_trait]
    impl<T: Blobstore + Clone> Blobstore for GatedBlobstore<T> {
        async fn put<'a>(
            &'a self,
            ctx: &'a CoreContext,
            key: String,
            value: BlobstoreBytes,
        ) -> Result<()> {
            let mut allow = self.allow.clone();
            while !*allow.borrow() {
                // Wait until we are allowed to continue.
                allow.changed().await?;
            }
            self.inner.put(ctx, key, value).await
        }

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

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

        let inner = Memblob::new(PutBehaviour::Overwrite);
        let (allow_tx, allow_rx) = watch::channel(false);
        let delay = Arc::new(GatedBlobstore::new(inner.clone(), allow_rx));
        let outer = MemWritesBlobstore::new(delay);

        let key = "key";
        let value = BlobstoreBytes::from_bytes("value");

        outer.put(ctx, key.to_owned(), value.clone()).await?;
        assert!(inner.get(ctx, key).await?.is_none());

        let persist = tokio::spawn({
            cloned!(ctx, outer);
            async move { outer.persist(&ctx).await }
        });

        // Wait until a write is in flight (there will be more than 1
        // receiver, as each write clones the receiver)
        while allow_tx.receiver_count() <= 1 {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        assert_eq!(inner.get(ctx, key).await?, None);
        assert_eq!(outer.get(ctx, key).await?, Some(value.clone().into()));

        let key2 = "key2";
        let value2 = BlobstoreBytes::from_bytes("value2");

        outer.put(ctx, key.to_owned(), value2.clone()).await?;
        outer.put(ctx, key2.to_owned(), value2.clone()).await?;

        assert_eq!(inner.get(ctx, key).await?, None);
        assert_eq!(inner.get(ctx, key2).await?, None);
        assert_eq!(outer.get(ctx, key).await?, Some(value2.clone().into()));
        assert_eq!(outer.get(ctx, key2).await?, Some(value2.clone().into()));

        allow_tx.send(true)?;
        persist.await??;

        assert_eq!(inner.get(ctx, key).await?, Some(value.into()));
        assert_eq!(inner.get(ctx, key2).await?, None);
        assert_eq!(outer.get(ctx, key).await?, Some(value2.clone().into()));
        assert_eq!(outer.get(ctx, key2).await?, Some(value2.clone().into()));

        outer.persist(ctx).await?;
        assert_eq!(inner.get(ctx, key).await?, Some(value2.clone().into()));
        assert_eq!(inner.get(ctx, key2).await?, Some(value2.into()));

        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;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Injected failure in get to ChaosBlobstore for key {0}")]
    InjectedChaosGet(String),
    #[error("Injected failure in put to ChaosBlobstore for key {0}")]
    InjectedChaosPut(String),
    #[error("Injected failure in is_present to ChaosBlobstore for key {0}")]
    InjectedChaosIsPresent(String),
}


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

use std::num::NonZeroU32;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use rand::thread_rng;
use rand::Rng;

mod errors;
pub use crate::errors::ErrorKind;

const NEVER_CHAOS_THRESHOLD: f32 = 1.0;
const ALWAYS_CHAOS_THRESHOLD: f32 = -1.0;

#[derive(Clone, Copy, Debug, Default)]
pub struct ChaosOptions {
    error_sample_read: Option<NonZeroU32>,
    error_sample_write: Option<NonZeroU32>,
}

impl ChaosOptions {
    /// Pass `error_sample_read` or `error_sample_write` value
    /// from Some(1) for always chaos to Some(N) to get 1/N chance of failure.
    pub fn new(
        error_sample_read: Option<NonZeroU32>,
        error_sample_write: Option<NonZeroU32>,
    ) -> Self {
        Self {
            error_sample_read,
            error_sample_write,
        }
    }

    pub fn has_chaos(&self) -> bool {
        self.error_sample_read.is_some() || self.error_sample_write.is_some()
    }
}

/// A layer over an existing blobstore that errors randomly
#[derive(Clone, Debug)]
pub struct ChaosBlobstore<T> {
    blobstore: T,
    sample_threshold_read: f32,
    sample_threshold_write: f32,
    #[allow(dead_code)]
    options: ChaosOptions,
}

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

fn derive_threshold(sample_rate: Option<NonZeroU32>) -> f32 {
    sample_rate.map_or(NEVER_CHAOS_THRESHOLD, |rate| {
        match rate.get() {
            // Avoid chance of rng returning 0.0 and threshold being 0.0
            1 => ALWAYS_CHAOS_THRESHOLD,
            // If rate 100, then rng must generate over 0.99 to trigger error
            n => 1.0 - (1.0 / (n as f32)),
        }
    })
}

impl<T> ChaosBlobstore<T> {
    pub fn new(blobstore: T, options: ChaosOptions) -> Self {
        let sample_threshold_read = derive_threshold(options.error_sample_read);
        let sample_threshold_write = derive_threshold(options.error_sample_write);
        Self {
            blobstore,
            sample_threshold_read,
            sample_threshold_write,
            options,
        }
    }
}

#[async_trait]
impl<T: Blobstore + BlobstorePutOps> Blobstore for ChaosBlobstore<T> {
    #[inline]
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let should_error = thread_rng().gen::<f32>() > self.sample_threshold_read;
        let get = self.blobstore.get(ctx, key);
        if should_error {
            Err(ErrorKind::InjectedChaosGet(key.to_owned()).into())
        } else {
            get.await
        }
    }

    #[inline]
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.put_impl(ctx, key, value, None).await?;
        Ok(())
    }

    #[inline]
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let should_error = thread_rng().gen::<f32>() > self.sample_threshold_read;
        let is_present = self.blobstore.is_present(ctx, key);
        if should_error {
            Err(ErrorKind::InjectedChaosIsPresent(key.to_owned()).into())
        } else {
            is_present.await
        }
    }
}

impl<T: BlobstorePutOps> ChaosBlobstore<T> {
    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
    ) -> Result<OverwriteStatus> {
        let should_error = thread_rng().gen::<f32>() > self.sample_threshold_write;
        let put = if should_error {
            None
        } else {
            let put = if let Some(put_behaviour) = put_behaviour {
                self.blobstore
                    .put_explicit(ctx, key.clone(), value, put_behaviour)
            } else {
                self.blobstore.put_with_status(ctx, key.clone(), value)
            };
            Some(put)
        };
        match put {
            None => Err(ErrorKind::InjectedChaosPut(key).into()),
            Some(put) => put.await,
        }
    }
}

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

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

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

    use super::*;

    #[fbinit::test]
    async fn test_error_on_write(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let wrapper =
            ChaosBlobstore::new(base.clone(), ChaosOptions::new(None, NonZeroU32::new(1)));
        let key = "foobar";

        let r = wrapper
            .put(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_err());
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(!base_present);
    }

    #[fbinit::test]
    async fn test_error_on_write_with_status(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let wrapper =
            ChaosBlobstore::new(base.clone(), ChaosOptions::new(None, NonZeroU32::new(1)));
        let key = "foobar";

        let r = wrapper
            .put_with_status(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_err());
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(!base_present);
    }

    #[fbinit::test]
    async fn test_error_on_read(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let wrapper =
            ChaosBlobstore::new(base.clone(), ChaosOptions::new(NonZeroU32::new(1), None));
        let key = "foobar";

        let r = wrapper
            .put(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_ok());
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(base_present);
        let r = wrapper.get(ctx, key).await;
        assert!(r.is_err());
    }
}


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

use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use rand::Rng;
use rand_distr::Distribution;

pub type Normal = rand_distr::Normal<f64>;

#[derive(Clone, Copy, Default, Debug)]
pub struct DelayOptions {
    pub get_dist: Option<Normal>,
    pub put_dist: Option<Normal>,
}

impl DelayOptions {
    pub fn has_delay(&self) -> bool {
        self.get_dist.is_some() || self.put_dist.is_some()
    }
}

#[derive(Debug)]
pub struct DelayedBlobstore<B> {
    inner: B,
    get_dist: Option<Normal>,
    put_dist: Option<Normal>,
}

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

impl<B> DelayedBlobstore<B> {
    pub fn new(inner: B, get_dist: Normal, put_dist: Normal) -> Self {
        Self {
            inner,
            get_dist: Some(get_dist),
            put_dist: Some(put_dist),
        }
    }

    pub fn from_options(inner: B, options: DelayOptions) -> Self {
        Self {
            inner,
            get_dist: options.get_dist,
            put_dist: options.put_dist,
        }
    }
}

#[async_trait]
impl<B: BlobstorePutOps> Blobstore for DelayedBlobstore<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        delay(self.get_dist).await;
        self.inner.get(ctx, key).await
    }

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

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

impl<T: BlobstorePutOps> DelayedBlobstore<T> {
    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
    ) -> Result<OverwriteStatus> {
        delay(self.put_dist).await;

        if let Some(put_behaviour) = put_behaviour {
            self.inner
                .put_explicit(ctx, key.clone(), value, put_behaviour)
                .await
        } else {
            self.inner.put_with_status(ctx, key.clone(), value).await
        }
    }
}

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

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

async fn delay<D>(distribution: Option<D>)
where
    D: Distribution<f64>,
{
    if let Some(distribution) = distribution {
        let seconds = rand::thread_rng().sample(distribution).abs();
        tokio::time::sleep(Duration::new(
            seconds.trunc() as u64,
            (seconds.fract() * 1e+9) as u32,
        ))
        .await;
    }
}


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

//! Ephemeral Blobstore Bubbles

use std::fmt;
use std::num::NonZeroU64;
use std::str::FromStr;
use std::sync::Arc;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Result;
use async_stream::try_stream;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use blobstore::BlobstoreEnumerableWithUnlink;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreKeyParam;
use blobstore::BlobstoreKeySource;
use blobstore::BlobstoreUnlinkOps;
use changesets::ChangesetsArc;
use context::CoreContext;
use derivative::Derivative;
use futures::future::try_join_all;
use futures::pin_mut;
use futures::stream::TryStreamExt;
use futures::Stream;
use futures_lazy_shared::LazyShared;
use metaconfig_types::RepoConfigArc;
use mononoke_types::repo::EPH_ID_PREFIX;
use mononoke_types::repo::EPH_ID_SUFFIX;
use mononoke_types::DateTime;
use prefixblob::PrefixBlobstore;
use repo_blobstore::RepoBlobstore;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityArc;
use repo_identity::RepoIdentityRef;
use serde::Deserialize;
use serde::Serialize;
use sql::mysql_async::prelude::ConvIr;
use sql::mysql_async::prelude::FromValue;
use sql::mysql_async::FromValueError;
use sql::mysql_async::Value;
use sql::sql_common::mysql::opt_try_from_rowfield;
use sql::sql_common::mysql::OptionalTryFromRowField;
use sql::sql_common::mysql::RowField;
use sql::sql_common::mysql::ValueError;
use sql_ext::SqlConnections;

use crate::changesets::EphemeralChangesets;
use crate::error::EphemeralBlobstoreError;
use crate::handle::EphemeralHandle;
use crate::view::EphemeralRepoView;

pub enum StorageLocation {
    // This is not ephemeral
    Persistent,
    // This is ephemeral, but its bubble is not yet known
    UnknownBubble,
    // This is ephemeral and located at given bubble
    Bubble(BubbleId),
}

impl From<BubbleId> for StorageLocation {
    fn from(bubble_id: BubbleId) -> Self {
        Self::Bubble(bubble_id)
    }
}

impl StorageLocation {
    pub fn ephemeral(maybe_bubble: Option<BubbleId>) -> Self {
        match maybe_bubble {
            None => Self::UnknownBubble,
            Some(id) => Self::Bubble(id),
        }
    }
}

/// Ephemeral Blobstore Bubble ID.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct BubbleId(NonZeroU64);

impl fmt::Display for BubbleId {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(fmt)
    }
}

impl From<BubbleId> for Value {
    fn from(bubble_id: BubbleId) -> Self {
        Value::UInt(bubble_id.0.into())
    }
}

impl From<BubbleId> for NonZeroU64 {
    fn from(bubble_id: BubbleId) -> Self {
        bubble_id.0
    }
}

impl ConvIr<BubbleId> for BubbleId {
    fn new(v: Value) -> Result<Self, FromValueError> {
        let from_u64 = |id, v| match NonZeroU64::new(id) {
            Some(id) => Ok(BubbleId(id)),
            None => Err(FromValueError(v)),
        };
        match v {
            Value::UInt(id) => from_u64(id, v),
            Value::Int(id) => from_u64(id as u64, v),
            v => Err(FromValueError(v)),
        }
    }

    fn commit(self) -> BubbleId {
        self
    }

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

impl FromValue for BubbleId {
    type Intermediate = BubbleId;
}

impl OptionalTryFromRowField for BubbleId {
    fn try_from_opt(field: RowField) -> Result<Option<Self>, ValueError> {
        opt_try_from_rowfield(field)
    }
}

impl TryFrom<i64> for BubbleId {
    type Error = ();

    fn try_from(value: i64) -> Result<Self, Self::Error> {
        u64::try_from(value)
            .ok()
            .and_then(NonZeroU64::new)
            .map(BubbleId::new)
            .ok_or(())
    }
}

impl FromStr for BubbleId {
    type Err = <NonZeroU64 as FromStr>::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        NonZeroU64::from_str(s).map(Self)
    }
}

impl BubbleId {
    pub fn new(id: NonZeroU64) -> Self {
        BubbleId(id)
    }

    /// Generate the blobstore prefix for this bubble.
    pub fn prefix(&self) -> String {
        format!("{}{}{}", EPH_ID_PREFIX, self.0, EPH_ID_SUFFIX,)
    }
}

type RawBubbleBlobstore = PrefixBlobstore<Arc<dyn BlobstoreEnumerableWithUnlink>>;

/// Enum representing the expiry status of a bubble in the backing store.
#[derive(Copy, Debug, Clone, PartialEq, Serialize, Deserialize, Abomonation)]
pub enum ExpiryStatus {
    Active = 0,
    Expired = 1,
}

impl fmt::Display for ExpiryStatus {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ExpiryStatus::Expired => write!(fmt, "Expired"),
            ExpiryStatus::Active => write!(fmt, "Active"),
        }
    }
}

impl From<ExpiryStatus> for Value {
    fn from(status: ExpiryStatus) -> Self {
        let val = match status {
            ExpiryStatus::Expired => 1,
            ExpiryStatus::Active => 0,
        };
        Value::Int(val)
    }
}

impl ConvIr<ExpiryStatus> for ExpiryStatus {
    fn new(v: Value) -> Result<Self, FromValueError> {
        match v {
            Value::Int(1) => Ok(ExpiryStatus::Expired),
            Value::Int(0) => Ok(ExpiryStatus::Active),
            v => Err(FromValueError(v)),
        }
    }

    fn commit(self) -> ExpiryStatus {
        self
    }

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

impl FromValue for ExpiryStatus {
    type Intermediate = ExpiryStatus;
}

impl OptionalTryFromRowField for ExpiryStatus {
    fn try_from_opt(field: RowField) -> Result<Option<Self>, ValueError> {
        opt_try_from_rowfield(field)
    }
}

impl TryFrom<i64> for ExpiryStatus {
    type Error = ();

    fn try_from(value: i64) -> Result<Self, Self::Error> {
        match value {
            1 => Ok(ExpiryStatus::Expired),
            0 => Ok(ExpiryStatus::Active),
            _ => Err(()),
        }
    }
}

/// An opened ephemeral blobstore bubble.  This is a miniature blobstore
/// that stores blobs just for this ephemeral bubble in a particular repo.
#[derive(Derivative, Clone)]
#[derivative(Debug)]
pub struct Bubble {
    /// ID of the current bubble.
    bubble_id: BubbleId,

    /// Expiration time.  After this time, the bubble no longer exists.
    /// This includes the grace period from the ephemeral blobstore.
    expires_at: DateTime,

    /// Blobstore to use for accessing blobs in this bubble, without redaction
    /// or repo prefix wrappers.
    pub(crate) blobstore: RawBubbleBlobstore,

    expired: ExpiryStatus,

    /// SQL connection
    #[derivative(Debug = "ignore")]
    connections: SqlConnections,

    /// The labels associated with the bubble. Currently used for lifetime extension.
    #[derivative(Debug = "ignore")]
    labels: LazyShared<Result<Vec<String>, EphemeralBlobstoreError>>,
}

impl fmt::Display for Bubble {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Bubble({})<{}>", self.bubble_id, self.blobstore)
    }
}

impl Bubble {
    pub(crate) fn new(
        bubble_id: BubbleId,
        expires_at: DateTime,
        blobstore: Arc<dyn BlobstoreEnumerableWithUnlink>,
        connections: SqlConnections,
        expired: ExpiryStatus,
        labels: LazyShared<Result<Vec<String>, EphemeralBlobstoreError>>,
    ) -> Self {
        let blobstore = PrefixBlobstore::new(blobstore, bubble_id.prefix());
        Self {
            bubble_id,
            expires_at,
            blobstore,
            connections,
            expired,
            labels,
        }
    }

    pub(crate) fn check_unexpired(&self) -> Result<()> {
        if self.expires_at >= DateTime::now() && self.expired != ExpiryStatus::Expired {
            Ok(())
        } else {
            Err(EphemeralBlobstoreError::BubbleExpired(self.bubble_id).into())
        }
    }

    pub(crate) async fn delete_blobs_in_bubble(&self, ctx: &CoreContext) -> Result<usize> {
        let key_stream = self.get_keys_in_bubble(ctx, None).await;
        let mut keys_deleted = 0;
        pin_mut!(key_stream);
        while let Some(keys) = key_stream.try_next().await? {
            // As long as the unlink operation on underlying blobstores is "truly" async
            // (i.e. they yield on cross service I/O), the below will execute the unlinking
            // concurrently and terminate early on the first error encountered.
            let unlinked_keys =
                try_join_all(keys.iter().map(|key| self.blobstore.unlink(ctx, key))).await?;
            keys_deleted += unlinked_keys.len();
        }
        // If the unlinking of all blobs within the bubble was successful, return the
        // number of blobs unlinked.
        Ok(keys_deleted)
    }

    pub(crate) async fn keys_in_bubble(
        &self,
        ctx: &CoreContext,
        start_from: Option<String>,
        limit: u32,
    ) -> Result<Vec<String>> {
        let key_stream = self.get_keys_in_bubble(ctx, start_from).await;
        let mut collected_keys = vec![];
        let limit = limit.try_into().unwrap();
        pin_mut!(key_stream);
        // Executing the below sequentially since we want to maintain
        // the ordering of the elements returned. Plus, we want to exit
        // as soon as we have the required number of keys.
        while let Some(keys) = key_stream.try_next().await? {
            collected_keys.extend(keys);
            if collected_keys.len() >= limit {
                break;
            }
        }
        // In cases where limit % batch_size != 0, we would have fetched
        // more than required, trim the extra keys.
        collected_keys.truncate(limit);
        Ok(collected_keys)
    }

    async fn get_keys_in_bubble<'a>(
        &'a self,
        ctx: &'a CoreContext,
        start_from: Option<String>,
    ) -> impl Stream<Item = Result<Vec<String>>> + 'a {
        let search_range = match start_from {
            Some(start) => BlobstoreKeyParam::from(start..),
            None => BlobstoreKeyParam::from(..),
        };
        let mut token = Arc::new(search_range);
        try_stream! {
            loop {
                let result = self.blobstore.enumerate(ctx, &token).await?;
                yield Vec::from_iter(result.keys);
                token = match result.next_token {
                    Some(next_token) => Arc::new(next_token),
                    None => break,
                };
            }
        }
    }

    pub fn bubble_id(&self) -> BubbleId {
        self.bubble_id
    }

    pub fn expired(&self) -> ExpiryStatus {
        self.expired
    }

    pub fn expires_at(&self) -> DateTime {
        self.expires_at
    }

    pub async fn labels(&self) -> Result<Vec<String>> {
        self.labels
            .get_or_init(|| async { Ok(Vec::new()) })
            .await
            .map_err(|e| anyhow!(e))
    }

    /// Return a blobstore that gives priority to accessing the bubble, but falls back
    /// to the main blobstore.
    pub fn wrap_repo_blobstore(&self, main_blobstore: RepoBlobstore) -> RepoBlobstore {
        // Repo prefix/redaction is added only once by RepoBlobstore
        RepoBlobstore::new_with_wrapped_inner_blobstore(main_blobstore, |bs| {
            Arc::new(EphemeralHandle::new(self.clone(), bs))
        })
    }

    fn changesets_with_blobstore(
        &self,
        repo_blobstore: RepoBlobstore,
        container: &(impl ChangesetsArc + RepoIdentityRef),
    ) -> EphemeralChangesets {
        EphemeralChangesets::new(
            container.repo_identity().id(),
            self.bubble_id(),
            repo_blobstore,
            self.connections.clone(),
            container.changesets_arc(),
        )
    }

    pub fn changesets(
        &self,
        container: &(impl ChangesetsArc + RepoIdentityRef + RepoBlobstoreRef),
    ) -> EphemeralChangesets {
        let repo_blobstore = self.wrap_repo_blobstore(container.repo_blobstore().clone());
        self.changesets_with_blobstore(repo_blobstore, container)
    }

    pub fn repo_view(
        &self,
        container: &(
             impl RepoBlobstoreRef + RepoIdentityRef + RepoIdentityArc + ChangesetsArc + RepoConfigArc
         ),
    ) -> EphemeralRepoView {
        let repo_blobstore = self.wrap_repo_blobstore(container.repo_blobstore().clone());
        let repo_identity = container.repo_identity_arc();
        let repo_config = container.repo_config_arc();
        EphemeralRepoView {
            repo_blobstore: Arc::new(repo_blobstore.clone()),
            changesets: Arc::new(self.changesets_with_blobstore(repo_blobstore, container)),
            repo_identity,
            repo_config,
        }
    }
}

// These blobstore methods are not to be used directly as they bypass redaction.
// Instead use .wrap_repo_blobstore
impl Bubble {
    pub(crate) async fn get(
        &self,
        ctx: &CoreContext,
        key: &str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.check_unexpired()?;
        self.blobstore.get(ctx, key).await
    }

    pub(crate) async fn put(
        &self,
        ctx: &CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.check_unexpired()?;
        self.blobstore.put(ctx, key, value).await
    }

    pub(crate) async fn is_present(
        &self,
        ctx: &CoreContext,
        key: &str,
    ) -> Result<BlobstoreIsPresent> {
        self.check_unexpired()?;
        self.blobstore.is_present(ctx, key).await
    }
}


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

//! Ephemeral Store Builder

use std::sync::Arc;
use std::time::Duration;

use blobstore::BlobstoreEnumerableWithUnlink;
use metaconfig_types::BubbleDeletionMode;
use mononoke_types::RepositoryId;
use sql_construct::SqlConstruct;
use sql_ext::SqlConnections;
use sql_query_config::SqlQueryConfig;

use crate::store::RepoEphemeralStore;

/// Ephemeral Store Builder.
pub struct RepoEphemeralStoreBuilder {
    /// Database used to manage the ephemeral blobstore metadata.
    connections: SqlConnections,
}

impl SqlConstruct for RepoEphemeralStoreBuilder {
    const LABEL: &'static str = "ephemeral_blobstore";

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

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl RepoEphemeralStoreBuilder {
    pub fn build(
        self,
        repo_id: RepositoryId,
        blobstore: Arc<dyn BlobstoreEnumerableWithUnlink>,
        sql_query_config: Arc<SqlQueryConfig>,
        initial_bubble_lifespan: Duration,
        bubble_expiration_grace: Duration,
        bubble_deletion_mode: BubbleDeletionMode,
    ) -> RepoEphemeralStore {
        RepoEphemeralStore::new(
            repo_id,
            self.connections,
            blobstore,
            sql_query_config,
            initial_bubble_lifespan,
            bubble_expiration_grace,
            bubble_deletion_mode,
        )
    }
}


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

use std::sync::Arc;

use anyhow::bail;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Loadable;
use changesets::ChangesetEntry;
use changesets::ChangesetInsert;
use changesets::Changesets;
use changesets::SortOrder;
use context::CoreContext;
use derivative::Derivative;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::try_join;
use itertools::Itertools;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::Generation;
use mononoke_types::RepositoryId;
use repo_blobstore::RepoBlobstore;
use sorted_vector_map::SortedVectorMap;
use sql::Connection;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use vec1::Vec1;

use crate::bubble::BubbleId;

// Knows how to fetch snapshot changesets. Since initially there will
// be a single snapshot per bubble, there's no need to optimise anything
// on this class. As the need arises, we can tweak this, for example by
// having an extra table that stores parent information to avoid looking
// at the blobstore.
#[derive(Derivative, Clone)]
#[derivative(Debug)]
pub struct EphemeralChangesets {
    repo_id: RepositoryId,
    bubble_id: BubbleId,
    repo_blobstore: RepoBlobstore,
    #[derivative(Debug = "ignore")]
    connections: SqlConnections,
    #[derivative(Debug = "ignore")]
    persistent_changesets: Arc<dyn Changesets>,
}

mononoke_queries! {
    read SelectChangesets(
        repo_id: RepositoryId,
        bubble_id: BubbleId,
        >list cs_id: ChangesetId
    ) -> (ChangesetId, u64) {
        "SELECT cs_id, gen
         FROM ephemeral_bubble_changeset_mapping
         WHERE repo_id = {repo_id} AND bubble_id = {bubble_id} AND cs_id IN {cs_id}"
    }

    write InsertChangeset(
        values: (repo_id: RepositoryId, cs_id: ChangesetId, bubble_id: BubbleId, gen: u64)
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO ephemeral_bubble_changeset_mapping
        (repo_id, cs_id, bubble_id, gen)
        VALUES {values}"
    }
}

impl EphemeralChangesets {
    pub(crate) fn new(
        repo_id: RepositoryId,
        bubble_id: BubbleId,
        repo_blobstore: RepoBlobstore,
        connections: SqlConnections,
        persistent_changesets: Arc<dyn Changesets>,
    ) -> Self {
        Self {
            repo_id,
            bubble_id,
            repo_blobstore,
            connections,
            persistent_changesets,
        }
    }

    async fn fetch_gens_with_connection(
        &self,
        cs_ids: &[ChangesetId],
        connection: &Connection,
    ) -> Result<Vec<(ChangesetId, u64)>> {
        SelectChangesets::query(connection, &self.repo_id, &self.bubble_id, cs_ids).await
    }

    pub async fn fetch_gens(
        &self,
        cs_ids: &[ChangesetId],
    ) -> Result<SortedVectorMap<ChangesetId, u64>> {
        let mut gens: SortedVectorMap<_, _> = self
            .fetch_gens_with_connection(cs_ids, &self.connections.read_connection)
            .await?
            .into_iter()
            .collect();
        if gens.len() != cs_ids.len() {
            let missing: Vec<_> = cs_ids
                .iter()
                .cloned()
                .filter(|id| !gens.contains_key(id))
                .collect();
            let mut extra = self
                .fetch_gens_with_connection(
                    missing.as_slice(),
                    &self.connections.read_master_connection,
                )
                .await?
                .into_iter()
                .collect();
            gens.append(&mut extra);
        }
        Ok(gens)
    }

    async fn get_ephemeral(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
    ) -> Result<Vec<ChangesetEntry>> {
        let gens = self.fetch_gens(cs_ids).await?;
        let changesets: Vec<_> = stream::iter(
            cs_ids
                .iter()
                .filter(|id| gens.get(id).is_some())
                .map(|id| id.load(ctx, &self.repo_blobstore))
                .collect::<Vec<_>>(), // without this we get compile errors
        )
        .buffered(100)
        .try_collect()
        .await?;
        Ok(changesets
            .into_iter()
            .filter_map(|cs| {
                let cs_id = cs.get_changeset_id();
                gens.get(&cs_id).map(|gen| ChangesetEntry {
                    repo_id: self.repo_id,
                    cs_id,
                    parents: cs.parents().collect(),
                    gen: *gen,
                })
            })
            .collect())
    }
}

#[async_trait]
impl Changesets for EphemeralChangesets {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn add(&self, ctx: &CoreContext, cs: ChangesetInsert) -> Result<bool> {
        let parents_len = cs.parents.len();
        let parents = self.get_many(ctx, cs.parents.clone()).await?;
        if parents.len() != parents_len {
            bail!(
                "Not all parents found, expected [{}], found [{}]",
                cs.parents.into_iter().map(|id| id.to_string()).join(", "),
                parents
                    .into_iter()
                    .map(|entry| entry.cs_id.to_string())
                    .join(", ")
            );
        }
        let gen = parents
            .into_iter()
            .map(|entry| entry.gen)
            .max()
            .unwrap_or(0)
            + 1;
        let result = InsertChangeset::query(
            &self.connections.write_connection,
            &[(&self.repo_id, &cs.cs_id, &self.bubble_id, &gen)],
        )
        .await?;
        Ok(result.last_insert_id().is_some())
    }

    async fn add_many(
        &self,
        ctx: &CoreContext,
        css: Vec1<(ChangesetInsert, Generation)>,
    ) -> Result<()> {
        // If necessary, this can be optimised.
        for (cs, _) in css {
            self.add(ctx, cs).await?;
        }
        Ok(())
    }

    async fn get(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<Option<ChangesetEntry>> {
        Ok(self.get_many(ctx, vec![cs_id]).await?.into_iter().next())
    }

    async fn get_many(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetEntry>> {
        let ephemeral = self.get_ephemeral(ctx, &cs_ids);
        let persistent = self.persistent_changesets.get_many(ctx, cs_ids.clone());
        let (mut ephemeral, persistent) = try_join!(ephemeral, persistent)?;
        ephemeral.extend(persistent);
        Ok(ephemeral)
    }

    /// Use caching for the full changeset ids and slower path otherwise.
    async fn get_many_by_prefix(
        &self,
        _ctx: &CoreContext,
        _cs_prefix: ChangesetIdPrefix,
        _limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        unimplemented!()
    }

    fn prime_cache(&self, _ctx: &CoreContext, _changesets: &[ChangesetEntry]) {
        // no caching involved
    }

    async fn enumeration_bounds(
        &self,
        _ctx: &CoreContext,
        _read_from_master: bool,
        _known_heads: Vec<ChangesetId>,
    ) -> Result<Option<(u64, u64)>> {
        unimplemented!()
    }

    fn list_enumeration_range(
        &self,
        _ctx: &CoreContext,
        _min_id: u64,
        _max_id: u64,
        _sort_and_limit: Option<(SortOrder, u64)>,
        _read_from_master: bool,
    ) -> BoxStream<'_, Result<(ChangesetId, u64)>> {
        unimplemented!()
    }
}


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

use mononoke_types::RepositoryId;
use thiserror::Error;

use crate::bubble::BubbleId;

#[derive(Debug, Error, Clone)]
pub enum EphemeralBlobstoreError {
    /// The repository does not have an ephemeral blobstore.
    #[error("repo {0} does not have an ephemeral blobstore")]
    NoEphemeralBlobstore(RepositoryId),

    /// A new bubble could not be created.
    #[error("failed to create a new bubble")]
    CreateBubbleFailed,

    /// The requested bubble does not exist.  Either it was never created or has expired.
    #[error("bubble {0} does not exist, or has expired")]
    NoSuchBubble(BubbleId),

    /// An in-use bubble has expired.
    #[error("bubble {0} has expired")]
    BubbleExpired(BubbleId),

    /// The requested bubble could not be deleted.
    #[error("failed to delete bubble {0}")]
    DeleteBubbleFailed(BubbleId),

    /// The bubble deletion action is disabled
    #[error("bubble deletion is disabled")]
    DeleteBubbleDisabled,

    /// Failed to fetch labels associated with the bubble
    #[error("failed to fetch labels for bubble {0}")]
    FetchBubbleLabelsFailed(BubbleId),
}


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

use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::BlobCopier;
use blobstore::Blobstore;
use context::CoreContext;
use filestore::FetchKey;
use filestore::FilestoreConfig;
use mononoke_types::ContentMetadataV2;
use mononoke_types::RepositoryId;
use repo_blobstore::RepoBlobstore;

use crate::Bubble;
use crate::BubbleId;

struct Copier {
    /// Blobstore that writes to ephemeral blobstore but has no bubble prefix
    /// and doesn't fallback to persistent blobstore
    raw_eph_blobstore: Arc<dyn Blobstore>,
    repo_prefix: String,
    prefix_bubble1: String,
    prefix_bubble2: String,
}

#[async_trait]
impl BlobCopier for Copier {
    async fn copy(&self, ctx: &CoreContext, key: String) -> Result<()> {
        let old_key = [
            self.prefix_bubble1.as_str(),
            self.repo_prefix.as_str(),
            &key,
        ]
        .concat();
        let new_key = [
            self.prefix_bubble2.as_str(),
            self.repo_prefix.as_str(),
            &key,
        ]
        .concat();
        self.raw_eph_blobstore.copy(ctx, &old_key, new_key).await?;
        Ok(())
    }
}

impl Bubble {
    pub async fn copy_file_to_bubble(
        &self,
        ctx: &CoreContext,
        repo_id: RepositoryId,
        repo_blobstore: RepoBlobstore,
        other: BubbleId,
        config: FilestoreConfig,
        id: FetchKey,
    ) -> Result<Option<ContentMetadataV2>> {
        self.check_unexpired()?;
        let blobstore = self.wrap_repo_blobstore(repo_blobstore);
        let data = match filestore::get_metadata(&blobstore, ctx, &id).await? {
            Some(data) => data,
            None => return Ok(None),
        };

        let raw_eph_blobstore = Arc::new(self.blobstore.clone().into_inner()) as Arc<dyn Blobstore>;

        filestore::copy(
            &blobstore,
            &Copier {
                raw_eph_blobstore,
                repo_prefix: repo_id.prefix(),
                prefix_bubble1: self.bubble_id().prefix(),
                prefix_bubble2: other.prefix(),
            },
            config,
            ctx,
            &data,
        )
        .await?;
        Ok(Some(data))
    }
}


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

use std::fmt;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreBytes;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use context::CoreContext;

use crate::bubble::Bubble;

/// EphemeralHandle is a blobstore that wraps both a bubble blobstore and a
/// backing "persistent" blobstore. First, it queries the bubble blobstore
/// and if a blob is not present, it queries the persistent one.
#[derive(Debug)]
pub struct EphemeralHandle<B: Blobstore> {
    bubble: Bubble,
    main_blobstore: B,
}

impl<B: Blobstore> EphemeralHandle<B> {
    pub(crate) fn new(bubble: Bubble, main_blobstore: B) -> Self {
        Self {
            bubble,
            main_blobstore,
        }
    }
}

impl<B: Blobstore> fmt::Display for EphemeralHandle<B> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "EphemeralHandle<{}, {}>",
            self.bubble, self.main_blobstore
        )
    }
}

#[async_trait]
impl<B: Blobstore> Blobstore for EphemeralHandle<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        Ok(match self.bubble.get(ctx, key).await? {
            Some(content) => Some(content),
            None => self.main_blobstore.get(ctx, key).await?,
        })
    }

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

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        Ok(match self.bubble.is_present(ctx, key).await? {
            BlobstoreIsPresent::Absent | BlobstoreIsPresent::ProbablyNotPresent(_) => {
                self.main_blobstore.is_present(ctx, key).await?
            }
            BlobstoreIsPresent::Present => BlobstoreIsPresent::Present,
        })
    }
}


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

//! Ephemeral Blobstore
//!
//! Unlike other blobstores, this blobstore organises blobs into ephemeral
//! "bubbles".  Blobs placed within a bubble are deleted when the bubble
//! expires.  A bubble's lifespan can be extended at any time before it
//! expires.

mod bubble;
mod builder;
mod changesets;
mod error;
mod file;
mod handle;
mod store;
mod view;

pub use crate::bubble::Bubble;
pub use crate::bubble::BubbleId;
pub use crate::bubble::StorageLocation;
pub use crate::builder::RepoEphemeralStoreBuilder;
pub use crate::changesets::EphemeralChangesets;
pub use crate::error::EphemeralBlobstoreError;
pub use crate::handle::EphemeralHandle;
pub use crate::store::ArcRepoEphemeralStore;
pub use crate::store::RepoEphemeralStore;
pub use crate::store::RepoEphemeralStoreArc;
pub use crate::store::RepoEphemeralStoreRef;
pub use crate::view::EphemeralRepoView;