/*
 * 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::Context;
use anyhow::Error;
use blobstore::DEFAULT_PUT_BEHAVIOUR;
use borrowed::borrowed;
use bytes::Bytes;
use fbinit::FacebookInit;
use rand::distributions::Alphanumeric;
use rand::thread_rng;
use rand::Rng;
use rand::RngCore;
use strum::IntoEnumIterator;

use super::*;

async fn test_chunking_methods<Test, Fut>(
    fb: FacebookInit,
    put_behaviour: PutBehaviour,
    do_test: Test,
) -> Result<(), Error>
where
    Test: Fn(CoreContext, CountedSqlblob, Arc<TestSource>) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    for allow_inline in [true, false] {
        let (test_source, config_store) = get_test_config_store();
        let blobstore =
            Sqlblob::with_sqlite_in_memory(put_behaviour, &config_store, allow_inline, 0)?;
        let ctx = CoreContext::test_mock(fb);
        do_test(ctx, blobstore, test_source)
            .await
            .with_context(|| format_err!("while testing allow_inline {}", allow_inline))?;
    }
    Ok(())
}

async fn read_write_size(
    fb: FacebookInit,
    put_behaviour: PutBehaviour,
    blob_size: usize,
) -> Result<(), Error> {
    test_chunking_methods(fb, put_behaviour, |ctx, bs, _| async move {
        borrowed!(ctx);
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key = format!("manifoldblob_test_{}", suffix);

        let mut bytes_in = vec![0u8; blob_size];
        thread_rng().fill_bytes(&mut bytes_in);

        let blobstore_bytes = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in));

        assert!(
            !bs.is_present(ctx, &key).await?.assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        // Write a fresh blob
        bs.put(ctx, key.clone(), blobstore_bytes).await?;
        // Read back and verify
        let bytes_out = bs.get(ctx, &key).await?;
        assert_eq!(&bytes_in.to_vec(), bytes_out.unwrap().as_raw_bytes());

        assert!(
            bs.is_present(ctx, &key).await?.assume_not_found_if_unsure(),
            "Blob should exist now"
        );
        Ok(())
    })
    .await
}

#[fbinit::test]
async fn read_write(fb: FacebookInit) -> Result<(), Error> {
    for put_behaviour in PutBehaviour::iter() {
        // test a range of sizes that are inlineable and not inlineable
        for size in [0, 1, 2, 3, 64, MAX_INLINE_LEN, 254, 255, 256, 512] {
            let blob_size: usize = size.try_into()?;
            read_write_size(fb, put_behaviour, blob_size)
                .await
                .with_context(|| format_err!("while testing size {}", size))?;
        }
    }
    Ok(())
}

#[fbinit::test]
async fn double_put(fb: FacebookInit) -> Result<(), Error> {
    test_chunking_methods(fb, DEFAULT_PUT_BEHAVIOUR, |ctx, bs, _| async move {
        borrowed!(ctx);
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key = format!("manifoldblob_test_{}", suffix);

        let mut bytes_in = [0u8; 64];
        thread_rng().fill_bytes(&mut bytes_in);

        let blobstore_bytes = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in));

        assert!(
            !bs.is_present(ctx, &key).await?.assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        // Write a fresh blob
        bs.put(ctx, key.clone(), blobstore_bytes.clone()).await?;
        // Write it again
        bs.put(ctx, key.clone(), blobstore_bytes.clone()).await?;

        assert!(
            bs.is_present(ctx, &key).await?.assume_not_found_if_unsure(),
            "Blob should exist now"
        );
        Ok(())
    })
    .await
}

#[fbinit::test]
async fn overwrite(fb: FacebookInit) -> Result<(), Error> {
    test_chunking_methods(fb, PutBehaviour::Overwrite, |ctx, bs, _| async move {
        borrowed!(ctx);
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key = format!("manifoldblob_test_{}", suffix);

        let mut bytes_1 = [0u8; 64];
        thread_rng().fill_bytes(&mut bytes_1);
        let mut bytes_2 = [0u8; 32];
        thread_rng().fill_bytes(&mut bytes_2);

        let blobstore_bytes_1 = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_1));
        let blobstore_bytes_2 = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_2));

        // Write a fresh blob
        bs.put(ctx, key.clone(), blobstore_bytes_1.clone()).await?;
        // Overwrite it
        bs.put(ctx, key.clone(), blobstore_bytes_2.clone()).await?;

        assert_eq!(
            bs.get(ctx, &key).await?.map(|get| get.into_bytes()),
            Some(blobstore_bytes_2),
        );
        Ok(())
    })
    .await
}

#[fbinit::test]
async fn dedup(fb: FacebookInit) -> Result<(), Error> {
    test_chunking_methods(fb, DEFAULT_PUT_BEHAVIOUR, |ctx, bs, _| async move {
        borrowed!(ctx);
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key1 = format!("manifoldblob_test_{}", suffix);
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key2 = format!("manifoldblob_test_{}", suffix);

        let mut bytes_in = [0u8; 64];
        thread_rng().fill_bytes(&mut bytes_in);

        let blobstore_bytes = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in));

        assert!(
            !bs.is_present(ctx, &key1)
                .await?
                .assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        assert!(
            !bs.is_present(ctx, &key2)
                .await?
                .assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        // Write a fresh blob
        bs.put(ctx, key1.clone(), blobstore_bytes.clone()).await?;
        // Write it again under a different key
        bs.put(ctx, key2.clone(), blobstore_bytes.clone()).await?;

        // Reach inside the store and confirm it only stored the data once
        let data_store = bs.get_data_store();
        let row1 = data_store.get(&key1).await?.expect("Blob 1 not found");
        let row2 = data_store.get(&key2).await?.expect("Blob 2 not found");
        assert_eq!(row1.id, row2.id, "Chunk stored under different ids");
        assert_eq!(row1.count, row2.count, "Chunk count differs");
        assert_eq!(
            row1.chunking_method, row2.chunking_method,
            "Chunking method differs"
        );
        Ok(())
    })
    .await
}

#[fbinit::test]
async fn link(fb: FacebookInit) -> Result<(), Error> {
    test_chunking_methods(fb, DEFAULT_PUT_BEHAVIOUR, |ctx, bs, _| async move {
        borrowed!(ctx);
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key1 = format!("manifoldblob_test_{}", suffix);
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key2 = format!("manifoldblob_test_{}", suffix);

        let mut bytes_in = [0u8; 64];
        thread_rng().fill_bytes(&mut bytes_in);

        let blobstore_bytes = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in));

        assert!(
            !bs.is_present(ctx, &key1)
                .await?
                .assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        assert!(
            !bs.is_present(ctx, &key2)
                .await?
                .assume_not_found_if_unsure(),
            "Blob should not exist yet"
        );

        // Write a fresh blob
        bs.put(ctx, key1.clone(), blobstore_bytes.clone()).await?;
        // Link to a different key
        bs.copy(ctx, &key1, key2.clone()).await?;

        // Check that reads from the two keys match
        let bytes1 = bs.get(ctx, &key1).await?;
        let bytes2 = bs.get(ctx, &key2).await?;
        assert_eq!(
            bytes1.unwrap().as_raw_bytes(),
            bytes2.unwrap().as_raw_bytes()
        );

        // Reach inside the store and confirm it only stored the data once
        let data_store = bs.get_data_store();
        let row1 = data_store.get(&key1).await?.expect("Blob 1 not found");
        let row2 = data_store.get(&key2).await?.expect("Blob 2 not found");
        assert_eq!(row1.id, row2.id, "Chunk stored under different ids");
        assert_eq!(row1.count, row2.count, "Chunk count differs");
        assert_eq!(
            row1.chunking_method, row2.chunking_method,
            "Chunking method differs"
        );
        Ok(())
    })
    .await
}

#[fbinit::test]
async fn generations(fb: FacebookInit) -> Result<(), Error> {
    for auto_inline_puts in [true, false] {
        // Generate unique keys.
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key1 = format!("manifoldblob_test_{}", suffix);
        let suffix: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .map(char::from)
            .collect();
        let key2 = format!("manifoldblob_test_{}", suffix);

        let blobstore_bytes_inline = {
            let mut bytes_in = [0u8; MAX_INLINE_LEN as usize];
            thread_rng().fill_bytes(&mut bytes_in);
            BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in))
        };

        let blobstore_bytes = {
            let mut bytes_in = [0u8; 1024];
            thread_rng().fill_bytes(&mut bytes_in);
            BlobstoreBytes::from_bytes(Bytes::copy_from_slice(&bytes_in))
        };

        for blobstore_bytes in [blobstore_bytes_inline, blobstore_bytes] {
            let (test_source, config_store) = get_test_config_store();
            let bs = Sqlblob::with_sqlite_in_memory(
                DEFAULT_PUT_BEHAVIOUR,
                &config_store,
                auto_inline_puts,
                0, // no grace period for ctime updates,
            )?;
            let ctx = CoreContext::test_mock(fb);
            borrowed!(ctx);

            // Write a fresh blob
            bs.put_with_status(ctx, key1.clone(), blobstore_bytes.clone())
                .await?;

            let generations = bs.get_chunk_generations(&key1).await?;

            let value_len: u64 = blobstore_bytes.len().try_into()?;
            if !auto_inline_puts || value_len > MAX_INLINE_LEN {
                // Inspect, and determine that the generation number is present
                assert_eq!(
                    generations,
                    vec![Some(2)],
                    "Generation set to 2 {} {}",
                    auto_inline_puts,
                    value_len
                );
            } else {
                assert_eq!(generations, vec![], "No generations expected");
            }

            let mark_gen = set_test_generations(test_source.as_ref(), 4, 3, 0, INITIAL_VERSION + 1);

            config_store.force_update_configs();

            // Set the generation and confirm
            bs.set_generation(&key1, true, mark_gen).await?;
            let generations = bs.get_chunk_generations(&key1).await?;
            if value_len > MAX_INLINE_LEN {
                assert_eq!(generations, vec![Some(3)], "Generation set to 3");
            } else {
                // We should now have no generations as the set_generation(key, true) should have inlined
                assert_eq!(generations, vec![], "No generations expected");
            }

            // don't need to run these with variety of key lengths
            if value_len > MAX_INLINE_LEN {
                // Update via another key, confirm both have put generation
                set_test_generations(test_source.as_ref(), 5, 4, 2, INITIAL_VERSION + 2);
                config_store.force_update_configs();
                bs.put(ctx, key2.clone(), blobstore_bytes.clone()).await?;
                let generations = bs.get_chunk_generations(&key1).await?;
                assert_eq!(generations, vec![Some(5)], "key1 generation not updated");
                let generations = bs.get_chunk_generations(&key2).await?;
                assert_eq!(generations, vec![Some(5)], "key2 generation not updated");

                // Now update via the route GC uses, confirm it updates nicely and doesn't leap to
                // the wrong version.
                let mark_gen =
                    set_test_generations(test_source.as_ref(), 999, 10, 3, INITIAL_VERSION + 3);
                config_store.force_update_configs();
                bs.set_generation(&key1, true, mark_gen).await?;
                let generations = bs.get_chunk_generations(&key1).await?;
                assert_eq!(generations, vec![Some(10)], "key1 generation not updated");
                let generations = bs.get_chunk_generations(&key2).await?;
                assert_eq!(generations, vec![Some(10)], "key2 generation not updated");
            }
        }
    }
    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::fmt::Display;
use std::ops::Deref;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use stats::prelude::*;

use crate::Blobstore;
use crate::BlobstoreBytes;
use crate::BlobstoreEnumerationData;
use crate::BlobstoreGetData;
use crate::BlobstoreIsPresent;
use crate::BlobstoreKeyParam;
use crate::BlobstoreKeySource;
use crate::BlobstorePutOps;
use crate::BlobstoreUnlinkOps;
use crate::OverwriteStatus;
use crate::PutBehaviour;

define_stats_struct! {
    CountedBlobstoreStats("mononoke.blobstore.{}", prefix: String),
    get: timeseries(Rate, Sum),
    get_ok: timeseries(Rate, Sum),
    get_err: timeseries(Rate, Sum),
    put: timeseries(Rate, Sum),
    put_ok: timeseries(Rate, Sum),
    put_err: timeseries(Rate, Sum),
    put_not_checked: timeseries(Rate, Sum),
    put_new: timeseries(Rate, Sum),
    put_overwrote: timeseries(Rate, Sum),
    put_prevented: timeseries(Rate, Sum),
    is_present: timeseries(Rate, Sum),
    is_present_ok: timeseries(Rate, Sum),
    is_present_err: timeseries(Rate, Sum),
    copy: timeseries(Rate, Sum),
    copy_ok: timeseries(Rate, Sum),
    copy_err: timeseries(Rate, Sum),
    unlink: timeseries(Rate, Sum),
    unlink_ok: timeseries(Rate, Sum),
    unlink_err: timeseries(Rate, Sum),
    enumerate: timeseries(Rate, Sum),
    enumerate_ok: timeseries(Rate, Sum),
    enumerate_err: timeseries(Rate, Sum),
}

#[derive(Debug)]
pub struct CountedBlobstore<T> {
    blobstore: T,
    stats: CountedBlobstoreStats,
}

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

impl<T> CountedBlobstore<T> {
    pub fn new(name: String, blobstore: T) -> Self {
        Self {
            blobstore,
            stats: CountedBlobstoreStats::new(name),
        }
    }

    pub fn into_inner(self) -> T {
        self.blobstore
    }

    pub fn as_inner(&self) -> &T {
        &self.blobstore
    }
}

#[async_trait]
impl<T: Blobstore> Blobstore for CountedBlobstore<T> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.stats.get.add_value(1);
        let res = self.blobstore.get(ctx, key).await;
        match res {
            Ok(_) => self.stats.get_ok.add_value(1),
            Err(_) => self.stats.get_err.add_value(1),
        }
        res
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.stats.put.add_value(1);
        let res = self.blobstore.put(ctx, key, value).await;
        match res {
            Ok(()) => self.stats.put_ok.add_value(1),
            Err(_) => self.stats.put_err.add_value(1),
        }
        res
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.stats.is_present.add_value(1);
        let res = self.blobstore.is_present(ctx, key).await;
        match res {
            Ok(_) => self.stats.is_present_ok.add_value(1),
            Err(_) => self.stats.is_present_err.add_value(1),
        }
        res
    }

    async fn copy<'a>(
        &'a self,
        ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        self.stats.copy.add_value(1);
        let res = self.blobstore.copy(ctx, old_key, new_key).await;
        match res {
            Ok(()) => self.stats.copy_ok.add_value(1),
            Err(_) => self.stats.copy_err.add_value(1),
        }
        res
    }
}

impl<T: BlobstorePutOps> CountedBlobstore<T> {
    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
    ) -> Result<OverwriteStatus> {
        self.stats.put.add_value(1);
        let res = if let Some(put_behaviour) = put_behaviour {
            self.blobstore
                .put_explicit(ctx, key, value, put_behaviour)
                .await
        } else {
            self.blobstore.put_with_status(ctx, key, value).await
        };
        match res {
            Ok(status) => {
                self.stats.put_ok.add_value(1);
                match status {
                    OverwriteStatus::NotChecked => self.stats.put_not_checked.add_value(1),
                    OverwriteStatus::New => self.stats.put_new.add_value(1),
                    OverwriteStatus::Overwrote => self.stats.put_overwrote.add_value(1),
                    OverwriteStatus::Prevented => self.stats.put_prevented.add_value(1),
                };
            }
            Err(_) => self.stats.put_err.add_value(1),
        }
        res
    }
}

#[async_trait]
impl<T: BlobstorePutOps> BlobstorePutOps for CountedBlobstore<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_trait]
impl<T: BlobstoreUnlinkOps> BlobstoreUnlinkOps for CountedBlobstore<T> {
    async fn unlink<'a>(&'a self, ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        self.stats.unlink.add_value(1);
        let res = self.blobstore.unlink(ctx, key).await;
        match res {
            Ok(()) => self.stats.unlink_ok.add_value(1),
            Err(_) => self.stats.unlink_err.add_value(1),
        }
        res
    }
}

#[async_trait]
impl<T: BlobstoreKeySource> BlobstoreKeySource for CountedBlobstore<T> {
    async fn enumerate<'a>(
        &'a self,
        ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData> {
        self.stats.enumerate.add_value(1);
        let res = self.blobstore.enumerate(ctx, range).await;
        match res {
            Ok(_) => self.stats.enumerate_ok.add_value(1),
            Err(_) => self.stats.enumerate_err.add_value(1),
        }
        res
    }
}

impl<T: Blobstore> Deref for CountedBlobstore<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.as_inner()
    }
}


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

use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;

use super::Blobstore;
use super::BlobstoreBytes;
use super::BlobstoreGetData;
use super::BlobstorePutOps;
use super::BlobstoreUnlinkOps;
use super::OverwriteStatus;
use super::PutBehaviour;

/// Disabled blobstore which fails all operations with a reason. Primarily used as a
/// placeholder for administratively disabled blobstores.
#[derive(Debug)]
pub struct DisabledBlob {
    reason: String,
}

impl DisabledBlob {
    pub fn new(reason: impl Into<String>) -> Self {
        DisabledBlob {
            reason: reason.into(),
        }
    }
}

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

#[async_trait]
impl Blobstore for DisabledBlob {
    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        Err(anyhow!("Blobstore disabled: {}", self.reason))
    }

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

    async fn copy<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _old_key: &'a str,
        _new_key: String,
    ) -> Result<()> {
        Err(anyhow!("Blobstore disabled: {}", self.reason))
    }
}

#[async_trait]
impl BlobstorePutOps for DisabledBlob {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _key: String,
        _value: BlobstoreBytes,
        _put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        Err(anyhow!("Blobstore disabled: {}", self.reason))
    }

    async fn put_with_status<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        _key: String,
        _value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        Err(anyhow!("Blobstore disabled: {}", self.reason))
    }
}

#[async_trait]
impl BlobstoreUnlinkOps for DisabledBlob {
    async fn unlink<'a>(&'a self, _ctx: &'a CoreContext, _key: &'a str) -> Result<()> {
        Err(anyhow!("Blobstore disabled: {}", self.reason))
    }
}

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

    use super::*;

    #[fbinit::test]
    async fn test_disabled(fb: FacebookInit) {
        let disabled = DisabledBlob::new("test");
        let ctx = CoreContext::test_mock(fb);

        match disabled.get(&ctx, "foobar").await {
            Ok(_) => panic!("Unexpected success"),
            Err(err) => println!("Got error: {:?}", err),
        }

        match disabled
            .put(
                &ctx,
                "foobar".to_string(),
                BlobstoreBytes::from_bytes(vec![]),
            )
            .await
        {
            Ok(_) => panic!("Unexpected success"),
            Err(err) => println!("Got error: {:?}", 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 thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Blob {0} not found in blobstore")]
    NotFound(String),
    #[error("Error while opening state for blob store")]
    StateOpen,
}


/*
 * 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 counted_blobstore;
mod disabled;
mod errors;
pub mod macros;

use std::collections::HashSet;
use std::fmt;
use std::io::Cursor;
use std::ops::Bound;
use std::ops::RangeBounds;
use std::ops::RangeFrom;
use std::ops::RangeFull;
use std::ops::RangeInclusive;
use std::ops::RangeToInclusive;

use abomonation_derive::Abomonation;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use auto_impl::auto_impl;
use bytes::Bytes;
use clap::ValueEnum;
use context::CoreContext;
use serde_derive::Deserialize;
use serde_derive::Serialize;
use strum::AsRefStr;
use strum::Display;
use strum::EnumIter;
use strum::EnumString;
use strum::IntoStaticStr;
use thiserror::Error;
use trait_set::trait_set;

pub use crate::counted_blobstore::CountedBlobstore;
pub use crate::disabled::DisabledBlob;
pub use crate::errors::ErrorKind;

// This module exists to namespace re-exported
// imports, needed for macro exports.
pub mod private {
    pub use std::convert::TryFrom;
    pub use std::convert::TryInto;

    pub use anyhow::Error;
    pub use async_trait::async_trait;
    pub use context::CoreContext;
    pub use fbthrift::compact_protocol;

    pub use crate::Blobstore;
    pub use crate::BlobstoreBytes;
    pub use crate::BlobstoreGetData;
    pub use crate::Loadable;
    pub use crate::LoadableError;
    pub use crate::Storable;
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BlobstoreGetData {
    meta: BlobstoreMetadata,
    bytes: BlobstoreBytes,
}

const UNCOMPRESSED: u8 = b'0';
const COMPRESSED: u8 = b'1';

impl BlobstoreGetData {
    #[inline]
    pub fn new(meta: BlobstoreMetadata, bytes: BlobstoreBytes) -> Self {
        BlobstoreGetData { meta, bytes }
    }

    #[inline]
    pub fn as_meta(&self) -> &BlobstoreMetadata {
        &self.meta
    }

    #[inline]
    pub fn into_bytes(self) -> BlobstoreBytes {
        self.bytes
    }

    #[inline]
    pub fn as_bytes(&self) -> &BlobstoreBytes {
        &self.bytes
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.bytes.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.bytes.is_empty()
    }

    #[inline]
    pub fn into_raw_bytes(self) -> Bytes {
        self.into_bytes().into_bytes()
    }

    #[inline]
    pub fn as_raw_bytes(&self) -> &Bytes {
        self.as_bytes().as_bytes()
    }

    #[inline]
    pub fn from_bytes<B: Into<Bytes>>(bytes: B) -> Self {
        BlobstoreGetData {
            meta: BlobstoreMetadata::default(),
            bytes: BlobstoreBytes::from_bytes(bytes.into()),
        }
    }

    #[inline]
    pub fn remove_ctime(&mut self) {
        self.meta.ctime = None;
    }
}

impl From<BlobstoreBytes> for BlobstoreGetData {
    fn from(blob_bytes: BlobstoreBytes) -> Self {
        BlobstoreGetData {
            meta: BlobstoreMetadata::default(),
            bytes: blob_bytes,
        }
    }
}

impl From<BlobstoreGetData> for BlobstoreMetadata {
    fn from(blob: BlobstoreGetData) -> Self {
        blob.meta
    }
}

impl From<BlobstoreGetData> for BlobstoreBytes {
    fn from(blob: BlobstoreGetData) -> Self {
        blob.into_bytes()
    }
}

#[derive(Abomonation, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PackMetadata {
    /// Gives an idea what its packed with, if anything
    pub pack_key: String,
    /// How big the overall size of compressed forms was to reach this
    pub relevant_compressed_size: u64,
    /// How big the overall size of uncompressed forms was to reach this
    pub relevant_uncompressed_size: u64,
}

/// Optional information about the size of a value
#[derive(Abomonation, Clone, Debug, PartialEq, Eq, Hash)]
pub struct SizeMetadata {
    /// How much size this value has added to storage on its own.
    /// unique uncompressed size is already available from BlobstoreBytes::len()
    pub unique_compressed_size: u64,
    /// Info about packing, if its packed
    pub pack_meta: Option<PackMetadata>,
}

#[derive(Abomonation, Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct BlobstoreMetadata {
    ctime: Option<i64>,
    sizes: Option<SizeMetadata>,
}

impl BlobstoreMetadata {
    #[inline]
    pub fn new(ctime: Option<i64>, sizes: Option<SizeMetadata>) -> BlobstoreMetadata {
        BlobstoreMetadata { ctime, sizes }
    }

    #[inline]
    pub fn ctime(&self) -> Option<i64> {
        self.ctime
    }

    pub fn sizes(&self) -> Option<&SizeMetadata> {
        self.sizes.as_ref()
    }
}

/// A type representing bytes written to or read from a blobstore. The goal here is to ensure
/// that only types that implement `From<BlobstoreBytes>` and `Into<BlobstoreBytes>` can be
/// stored in the blob store.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct BlobstoreBytes(Bytes);

impl BlobstoreBytes {
    /// Construct an empty BlobstoreBytes
    pub fn empty() -> Self {
        BlobstoreBytes(Bytes::new())
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// This should only be used by blobstore and From/Into<BlobstoreBytes> implementations.
    #[inline]
    pub fn from_bytes<B: Into<Bytes>>(bytes: B) -> Self {
        BlobstoreBytes(bytes.into())
    }

    /// This should only be used by blobstore and From/Into<BlobstoreBytes> implementations.
    #[inline]
    pub fn into_bytes(self) -> Bytes {
        self.0
    }

    /// This should only be used by blobstore and From/Into<BlobstoreBytes> implementations.
    #[inline]
    pub fn as_bytes(&self) -> &Bytes {
        &self.0
    }

    pub fn encode(self, encode_limit: Option<u64>) -> Option<Bytes> {
        let mut bytes = vec![UNCOMPRESSED];
        let prepared = BlobstoreBytesSerialisable::from(self);
        unsafe { abomonation::encode(&prepared, &mut bytes).ok()? };

        match encode_limit {
            Some(encode_limit) if bytes.len() as u64 >= encode_limit => {
                let mut compressed = Vec::with_capacity(bytes.len());
                compressed.push(COMPRESSED);

                let mut cursor = Cursor::new(bytes);
                cursor.set_position(1);
                zstd::stream::copy_encode(cursor, &mut compressed, 0 /* use default */).ok()?;
                Some(Bytes::from(compressed))
            }
            _ => Some(Bytes::from(bytes)),
        }
    }

    pub fn decode(mut bytes: Bytes) -> Option<Self> {
        let prefix_size = 1;
        if bytes.len() < prefix_size {
            return None;
        }

        let is_compressed = bytes.split_to(prefix_size);
        let mut bytes: Vec<u8> = if is_compressed[0] == COMPRESSED {
            let cursor = Cursor::new(bytes);
            zstd::decode_all(cursor).ok()?
        } else {
            bytes.as_ref().into()
        };

        let get_data_serialisable =
            unsafe { abomonation::decode::<BlobstoreBytesSerialisable>(&mut bytes) };

        get_data_serialisable.and_then(|(get_data_serialisable, tail)| {
            if tail.is_empty() {
                Some(get_data_serialisable.clone().into())
            } else {
                None
            }
        })
    }
}

/// Serialisable counterpart of BlobstoreBytes fields mimic exactly its original
/// struct except in types that cannot be serialised
#[derive(Abomonation, Clone)]
struct BlobstoreBytesSerialisable(Vec<u8>);

impl From<BlobstoreBytes> for BlobstoreBytesSerialisable {
    fn from(blob_bytes: BlobstoreBytes) -> Self {
        BlobstoreBytesSerialisable(blob_bytes.into_bytes().as_ref().into())
    }
}

impl From<BlobstoreBytesSerialisable> for BlobstoreBytes {
    fn from(blob_bytes: BlobstoreBytesSerialisable) -> Self {
        BlobstoreBytes(blob_bytes.0.into())
    }
}

#[derive(Debug)]
pub enum BlobstoreIsPresent {
    // The blob is definitely present in the blobstore
    Present,
    /// The blob is definitely not present
    Absent,
    /// The blobstore has no evidence that the blob is present,
    /// however some of the operations resulted in errors.
    ProbablyNotPresent(Error),
}

impl BlobstoreIsPresent {
    pub fn assume_not_found_if_unsure(self) -> bool {
        match self {
            BlobstoreIsPresent::Present => true,
            BlobstoreIsPresent::Absent => false,
            BlobstoreIsPresent::ProbablyNotPresent(_) => false,
        }
    }

    pub fn fail_if_unsure(self) -> Result<bool, Error> {
        match self {
            BlobstoreIsPresent::Present => Ok(true),
            BlobstoreIsPresent::Absent => Ok(false),
            BlobstoreIsPresent::ProbablyNotPresent(err) => Err(err),
        }
    }
}

/// The blobstore interface, shared across all blobstores.
/// A blobstore must provide the following guarantees:
/// 1. `get` and `put` are atomic with respect to each other; a put will either put the entire
///    value, or not put anything, and a get will return either None, or the entire value that an
///    earlier put inserted.
/// 2. Once the future returned by `put` completes, the data is durably stored. This implies that
///    a permanent failure of the backend will not lose the data unless multiple replicas in the
///    backend are lost. For example, if you have replicas in multiple datacentres, you will
///    not lose data until you lose two or more datacentres. However, losing replicas can make the
///    data inaccessible for a time.
/// 3. Once the future returned by `put` completes, calling `get` from any process will get you a
///    future that will return the data that was saved in the blobstore; this is so that after the
///    `put` completes, Mononoke can update a database table and be confident that all Mononoke
///    instances can `get` the blobs that the database refers to.
///
/// Implementations of this trait can assume that the same value is supplied if two keys are
/// equal - in other words, each key is associated with at most one globally unique value.
/// In other words, `put(key, value).and_then(put(key, value2))` implies `value == value2` for the
/// `BlobstoreBytes` definition of equality. If `value != value2`, then the implementation's
/// behaviour is implementation defined (it can overwrite or not write at all, as long as it does
/// not break the atomicity guarantee, and does not have to be consistent in its behaviour).
///
/// Implementations of Blobstore must be `Clone` if they are to interoperate with other Mononoke
/// uses of Blobstores
#[async_trait]
#[auto_impl(&, Arc, Box)]
pub trait Blobstore: fmt::Display + fmt::Debug + Send + Sync {
    /// Fetch the value associated with `key`, or None if no value is present
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>>;
    /// Associate `value` with `key` for future gets; if `put` is called with different `value`s
    /// for the same key, the implementation may return any `value` it's been given in response
    /// to a `get` for that `key`.
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()>;
    /// Check that `get` will return a value for a given `key`, and not None. The provided
    /// implentation just calls `get`, and discards the return value; this can be overridden to
    /// avoid transferring data. In the absence of concurrent `put` calls, this must return
    /// `BlobstoreIsPresent::Absent` if `get` would return `None`, and `BlobstoreIsPresent::Present`
    /// if `get` would return `Some(_)`.
    /// In some cases, when it couldn't determine whether the key exists or not, it would
    /// return `BlobstoreIsPresent::ProbablyNotPresent`.
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        Ok(if self.get(ctx, key).await?.is_some() {
            BlobstoreIsPresent::Present
        } else {
            BlobstoreIsPresent::Absent
        })
    }
    /// Copy the value from one key to another. The default behaviour is to `get` and `put` the
    /// value, though some blobstores might have more efficient implementations that avoid
    /// transferring data.
    async fn copy<'a>(
        &'a self,
        ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        let value = self
            .get(ctx, old_key)
            .await?
            .with_context(|| format!("key {} not present", old_key))?;
        Ok(self.put(ctx, new_key, value.bytes).await?)
    }
}

/// Mononoke binaries will not overwrite existing blobstore keys by default
pub const DEFAULT_PUT_BEHAVIOUR: PutBehaviour = PutBehaviour::IfAbsent;

/// For blobstore implementors and advanced admin type users to control requested put behaviour
#[derive(
    IntoStaticStr,
    Clone,
    Copy,
    Debug,
    Display,
    EnumIter,
    EnumString,
    ValueEnum,
    Eq,
    PartialEq
)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum PutBehaviour {
    /// Blobstore::put will overwrite even if key is already present
    Overwrite,
    /// Blobstore::put will overwrite even if key is already present, plus log that and overwrite occured
    OverwriteAndLog,
    /// Blobstore::put will not overwrite if the key is already present.
    /// NB due to underlying stores TOCTOU limitations some puts might overwrite when when racing another put.
    /// This is expected, thus Blobstore::put() cannot reveal if the put wrote or not as behaviour other than
    /// logging/metrics should not depend on it.
    IfAbsent,
}

impl PutBehaviour {
    // For use inside BlobstoreWithPutBehaviour::put_with_behaviour implementations
    pub fn should_overwrite(&self) -> bool {
        match self {
            PutBehaviour::Overwrite | PutBehaviour::OverwriteAndLog => true,
            PutBehaviour::IfAbsent => false,
        }
    }

    // For use inside OverwriteHandler::on_put_overwrite implementations
    pub fn should_log(&self) -> bool {
        match self {
            PutBehaviour::Overwrite => false,
            PutBehaviour::IfAbsent | PutBehaviour::OverwriteAndLog => true,
        }
    }
}

/// For use from logging blobstores so they can record the overwrite status
/// `BlobstorePutOps::put_with_status`, and eventually `Blobstore::copy()` will return this.
#[derive(AsRefStr, Clone, Copy, Debug, Eq, PartialEq)]
pub enum OverwriteStatus {
    // We did not check if the key existed before writing it
    NotChecked,
    // This key did not exist before, therefore we did not overwrite
    New,
    // This did exist before, but we wrote over it
    Overwrote,
    // This did exist before, and the overwrite was prevented
    Prevented,
}

/// Lower level blobstore put api used by blobstore implementors and admin tooling
#[async_trait]
#[auto_impl(Arc, Box)]
pub trait BlobstorePutOps: Blobstore {
    /// Adds ability to specify the put behaviour explicitly so that even if per process default was
    /// IfAbsent  once could chose to OverwriteAndLog.  Expected to be used only in admin tools
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus>;

    /// Similar to `Blobstore::put`, but returns the OverwriteStatus as feedback rather than unit.
    /// Its here rather so we don't reveal the OverwriteStatus to regular put users.
    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus>;
}

/// Mixin trait for blobstores that support the `unlink()` operation
#[async_trait]
#[auto_impl(Arc, Box)]
pub trait BlobstoreUnlinkOps: Blobstore + BlobstorePutOps {
    /// Similar to unlink(2), this removes a key, resulting in content being removed if its the last key pointing to it.
    /// An error is returned if the key does not exist
    async fn unlink<'a>(&'a self, ctx: &'a CoreContext, key: &'a str) -> Result<()>;
}

/// BlobstoreKeySource Interface
/// Abstract for use with populate_healer
#[async_trait]
#[auto_impl(Arc, Box)]
pub trait BlobstoreKeySource: Blobstore {
    async fn enumerate<'a>(
        &'a self,
        ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData>;
}

trait_set! {
    /// A trait alias that represents blobstores that can be enumerated,
    /// updated and have their keys unlinked.
    #[auto_impl(Arc, Box)]
    pub trait BlobstoreEnumerableWithUnlink = BlobstoreKeySource + BlobstoreUnlinkOps;
}

/// Range of keys.  The range is inclusive (both start and end key are
/// included in the range), which matches Manifold behaviour.  If the key is
/// empty then the range is unbounded on that end.
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct BlobstoreKeyRange {
    pub begin_key: String,
    pub end_key: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub enum BlobstoreKeyToken {
    // For fileblob and manifold
    StringToken(String),
    // its an enum as other stores might have non-string tokens
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub enum BlobstoreKeyParam {
    Start(BlobstoreKeyRange),
    Continuation(BlobstoreKeyToken),
}

impl From<RangeInclusive<String>> for BlobstoreKeyParam {
    fn from(range: RangeInclusive<String>) -> Self {
        let (start, end) = range.into_inner();
        BlobstoreKeyParam::Start(BlobstoreKeyRange {
            begin_key: start,
            end_key: end,
        })
    }
}

impl From<RangeToInclusive<String>> for BlobstoreKeyParam {
    fn from(range: RangeToInclusive<String>) -> Self {
        BlobstoreKeyParam::Start(BlobstoreKeyRange {
            begin_key: String::from(""),
            end_key: range.end,
        })
    }
}

impl From<RangeFrom<String>> for BlobstoreKeyParam {
    fn from(range: RangeFrom<String>) -> Self {
        BlobstoreKeyParam::Start(BlobstoreKeyRange {
            begin_key: range.start,
            end_key: String::from(""),
        })
    }
}

impl From<RangeFull> for BlobstoreKeyParam {
    fn from(_: RangeFull) -> Self {
        BlobstoreKeyParam::Start(BlobstoreKeyRange {
            begin_key: String::from(""),
            end_key: String::from(""),
        })
    }
}

impl RangeBounds<String> for &BlobstoreKeyRange {
    fn start_bound(&self) -> Bound<&String> {
        if self.begin_key.is_empty() {
            Bound::Unbounded
        } else {
            Bound::Included(&self.begin_key)
        }
    }

    fn end_bound(&self) -> Bound<&String> {
        if self.end_key.is_empty() {
            Bound::Unbounded
        } else {
            Bound::Included(&self.end_key)
        }
    }
}

#[derive(Debug, Clone)]
pub struct BlobstoreEnumerationData {
    pub keys: HashSet<String>,
    /// current range being iterated, this range can be used to resume enumeration
    pub next_token: Option<BlobstoreKeyParam>,
}

#[derive(Debug, Error)]
pub enum LoadableError {
    #[error("Blobstore error")]
    Error(#[from] Error),
    #[error("Blob is missing: {0}")]
    Missing(String),
}

#[async_trait]
pub trait Loadable {
    type Value: Sized + 'static;

    async fn load<'a, B: Blobstore>(
        &'a self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Value, LoadableError>;
}

#[async_trait]
pub trait Storable: Sized {
    type Key: 'static;

    async fn store<'a, B: Blobstore>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Key>;
}

/// StoreLoadable represents an object that be loaded asynchronously through a given store of type
/// S. This offers a bit more flexibility over Blobstore's Loadable, which requires that the object
/// be asynchronously load loadable from a Blobstore. This level of indirection allows for using
/// Manifest's implementations with Manifests that are not backed by a Blobstore.
#[async_trait]
pub trait StoreLoadable<S> {
    type Value;

    async fn load<'a>(
        &'a self,
        ctx: &'a CoreContext,
        store: &'a S,
    ) -> Result<Self::Value, LoadableError>;
}

/// For convenience, all Blobstore Loadables are StoreLoadable through any Blobstore.
#[async_trait]
impl<L: Loadable + Sync, S: Blobstore> StoreLoadable<S> for L {
    type Value = <L as Loadable>::Value;

    async fn load<'a>(
        &'a self,
        ctx: &'a CoreContext,
        store: &'a S,
    ) -> Result<Self::Value, LoadableError> {
        self.load(ctx, store).await
    }
}

/// Trait to copy the value of the key between different blobstores. In the generic
/// case, we need to load from one blobstore to memory, then store in the other.
/// But it may be specialised to be more efficient in specific blobstore pairs.
#[async_trait]
pub trait BlobCopier {
    async fn copy(&self, ctx: &CoreContext, key: String) -> Result<()>;
}

/// Works for any pair of blobstores, does no optimisation at all.
pub struct GenericBlobstoreCopier<'a, A, B> {
    pub source: &'a A,
    pub target: &'a B,
}

#[async_trait]
impl<'a, A: Blobstore, B: Blobstore> BlobCopier for GenericBlobstoreCopier<'a, A, B> {
    async fn copy(&self, ctx: &CoreContext, key: String) -> Result<()> {
        let value = self
            .source
            .get(ctx, &key)
            .await?
            .with_context(|| format!("key {} not present", key))?;
        self.target.put(ctx, key, value.into_bytes()).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.
 */

#[macro_export]
macro_rules! impl_blobstore_conversions {
    ($ty:ident, $thrift_ty:ty) => {
        impl $crate::private::TryFrom<$crate::private::BlobstoreBytes> for $ty {
            type Error = $crate::private::Error;

            fn try_from(
                bytes: $crate::private::BlobstoreBytes,
            ) -> Result<Self, $crate::private::Error> {
                let t: $thrift_ty =
                    $crate::private::compact_protocol::deserialize(bytes.as_bytes().as_ref())?;
                t.try_into()
            }
        }

        impl From<$ty> for $crate::private::BlobstoreBytes {
            fn from(other: $ty) -> Self {
                let thrift: $thrift_ty = other.into();
                let data = $crate::private::compact_protocol::serialize(&thrift);
                Self::from_bytes(data)
            }
        }

        impl $crate::private::TryFrom<$crate::private::BlobstoreGetData> for $ty {
            type Error = $crate::private::Error;

            fn try_from(
                blob: $crate::private::BlobstoreGetData,
            ) -> Result<Self, $crate::private::Error> {
                blob.into_bytes().try_into()
            }
        }

        impl From<$ty> for $crate::private::BlobstoreGetData {
            fn from(other: $ty) -> Self {
                Into::<$crate::private::BlobstoreBytes>::into(other).into()
            }
        }
    };
}

/// You can use this macro under the following conditions:
/// 1. handle_type needs to implement TryFrom<handle_thrift_type> and Into<handle_thrift_type>
/// 2. same for value_type and value_thrift_type
/// 3. value_type has method `fn handle(&self) -> handle_type`
/// 4. handle_type has method `fb blobstore_key(&self) -> String`
#[macro_export]
macro_rules! impl_loadable_storable {
    (
        handle_type => $handle: ident,
        handle_thrift_type => $thrift_handle: ident,
        value_type => $ty:ident,
        value_thrift_type => $thrift_ty: ident,
    ) => {
        #[$crate::private::async_trait]
        impl $crate::private::Storable for $ty {
            type Key = $handle;

            async fn store<'a, B: $crate::private::Blobstore>(
                self,
                ctx: &'a $crate::private::CoreContext,
                blobstore: &'a B,
            ) -> Result<Self::Key, $crate::private::Error> {
                let handle = *self.handle();
                let key = handle.blobstore_key();
                blobstore.put(ctx, key, self.into()).await?;
                Ok(handle)
            }
        }

        #[$crate::private::async_trait]
        impl $crate::private::Loadable for $handle {
            type Value = $ty;

            async fn load<'a, B: $crate::private::Blobstore>(
                &'a self,
                ctx: &'a $crate::private::CoreContext,
                blobstore: &'a B,
            ) -> Result<Self::Value, $crate::private::LoadableError> {
                let id = *self;
                let bytes = blobstore.get(ctx, &id.blobstore_key()).await?;
                match bytes {
                    Some(bytes) => bytes
                        .try_into()
                        .map_err($crate::private::LoadableError::Error),
                    None => Err($crate::private::LoadableError::Missing(id.blobstore_key())),
                }
            }
        }

        $crate::impl_blobstore_conversions!($handle, $thrift_handle);
        $crate::impl_blobstore_conversions!($ty, $thrift_ty);
    };
}


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

//! Tests run against all blobstore implementations.

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

use std::sync::Arc;

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use borrowed::borrowed;
use bytes::Bytes;
use context::CoreContext;
use fbinit::FacebookInit;
use fileblob::Fileblob;
use memblob::Memblob;
use mononoke_types::BlobstoreBytes;
use sqlblob::get_test_config_store;
use sqlblob::Sqlblob;
use strum::IntoEnumIterator;
use tempdir::TempDir;

async fn overwrite<B: Blobstore + BlobstorePutOps>(
    fb: FacebookInit,
    blobstore: B,
    has_ctime: bool,
    put_behaviour: PutBehaviour,
) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);

    let key = "some_key".to_string() + &put_behaviour.to_string();
    let value = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(b"appleveldatav1"));

    let put_status1 = blobstore
        .put_with_status(ctx, key.clone(), value.clone())
        .await?;
    let expected_status1 = match put_behaviour {
        PutBehaviour::Overwrite => OverwriteStatus::NotChecked,
        PutBehaviour::OverwriteAndLog => OverwriteStatus::New,
        PutBehaviour::IfAbsent => OverwriteStatus::New,
    };
    assert_eq!(
        expected_status1, put_status1,
        "checking new {:?}",
        put_behaviour
    );

    let roundtrip1 = blobstore.get(ctx, &key).await?.unwrap();

    let ctime1 = roundtrip1.as_meta().ctime();

    let value2 = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(b"appleveldatav2"));

    let put_status2 = blobstore
        .put_with_status(ctx, key.clone(), value2.clone())
        .await?;

    let roundtrip2 = blobstore.get(ctx, &key).await?.unwrap();
    let ctime2 = roundtrip2.as_meta().ctime();
    if put_behaviour.should_overwrite() {
        assert_eq!(ctime2.is_some(), has_ctime);
        assert_eq!(
            value2,
            roundtrip2.into_bytes(),
            "checking overwrite value {:?}",
            put_behaviour
        );
    } else {
        assert_eq!(ctime1, ctime2);
        assert_eq!(
            value,
            roundtrip2.into_bytes(),
            "checking overwrite value {:?}",
            put_behaviour
        );
    }

    let expected_status2 = match put_behaviour {
        PutBehaviour::Overwrite => OverwriteStatus::NotChecked,
        PutBehaviour::OverwriteAndLog => OverwriteStatus::Overwrote,
        PutBehaviour::IfAbsent => OverwriteStatus::Prevented,
    };
    assert_eq!(
        expected_status2, put_status2,
        "checking overwrite {:?}",
        put_behaviour
    );

    Ok(())
}

async fn roundtrip_and_link<B: BlobstoreUnlinkOps>(
    fb: FacebookInit,
    blobstore: B,
    has_ctime: bool,
) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);

    let key = "randomkey";
    let value = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(b"appleveldata"));

    // Roundtrip
    blobstore.put(ctx, key.to_owned(), value.clone()).await?;

    let roundtrip = blobstore.get(ctx, key).await?.unwrap();

    let orig_ctime = roundtrip.as_meta().ctime();

    assert_eq!(orig_ctime.is_some(), has_ctime);
    assert_eq!(value, roundtrip.into_bytes());

    let newkey = "newkey";

    // And now the link
    blobstore.copy(ctx, key, newkey.to_owned()).await?;

    let newvalue = blobstore.get(ctx, newkey).await?.unwrap();

    let new_ctime = newvalue.as_meta().ctime();
    assert_eq!(new_ctime.is_some(), has_ctime);
    assert_eq!(orig_ctime, new_ctime);
    assert_eq!(value, newvalue.into_bytes());

    let newkey_is_present = blobstore
        .is_present(ctx, newkey)
        .await?
        .assume_not_found_if_unsure();

    assert!(newkey_is_present);

    // Try unlink
    blobstore.unlink(ctx, newkey).await?;
    let newkey_is_present2 = blobstore
        .is_present(ctx, newkey)
        .await?
        .assume_not_found_if_unsure();
    assert!(!newkey_is_present2);

    // Check we get error when unlinking an unknown key
    let unknown_key = "expected_missing_key";
    assert!(blobstore.unlink(ctx, unknown_key).await.is_err());

    Ok(())
}

async fn missing<B: Blobstore>(fb: FacebookInit, blobstore: B) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);

    let key = "missing";
    let out = blobstore.get(ctx, key).await?;

    assert!(out.is_none());
    Ok(())
}

macro_rules! blobstore_test_impl {
    ($mod_name: ident => {
        state: $state: expr,
        new: $new_cb: expr,
        persistent: $persistent: expr,
        has_ctime: $has_ctime: expr,
    }) => {
        mod $mod_name {
            use super::*;

            #[fbinit::test]
            async fn test_overwrite(fb: FacebookInit) -> Result<(), Error> {
                let state = $state;
                let has_ctime = $has_ctime;
                let factory = $new_cb;
                for b in PutBehaviour::iter() {
                    overwrite(fb, factory(state.clone(), b)?, has_ctime, b).await?
                }
                Ok(())
            }

            #[fbinit::test]
            async fn test_roundtrip_and_link(fb: FacebookInit) -> Result<(), Error> {
                let state = $state;
                let has_ctime = $has_ctime;
                let factory = $new_cb;
                roundtrip_and_link(
                    fb,
                    factory(state.clone(), PutBehaviour::Overwrite)?,
                    has_ctime,
                )
                .await
            }

            #[fbinit::test]
            async fn test_missing(fb: FacebookInit) -> Result<(), Error> {
                let state = $state;
                let factory = $new_cb;
                missing(fb, factory(state, PutBehaviour::Overwrite)?).await
            }

            #[fbinit::test]
            async fn test_boxable(_fb: FacebookInit) -> Result<(), Error> {
                let state = $state;
                let factory = $new_cb;
                // This is really just checking that the constructed type is Sized
                let _ = Box::new(factory(state, PutBehaviour::Overwrite)?);
                Ok(())
            }
        }
    };
}

blobstore_test_impl! {
    memblob_test => {
        state: (),
        new: move |_, put_behaviour| Ok::<_,Error>(Memblob::new(put_behaviour)),
        persistent: false,
        has_ctime: false,
    }
}

blobstore_test_impl! {
    box_blobstore_test => {
        state: (),
        new: move |_, put_behaviour| Ok::<_,Error>(Box::new(Memblob::new(put_behaviour))),
        persistent: false,
        has_ctime: false,
    }
}

blobstore_test_impl! {
    fileblob_test => {
        state: Arc::new(TempDir::new("fileblob_test").unwrap()),
        new: move |dir: Arc<TempDir>, put_behaviour,| Fileblob::open(&*dir, put_behaviour),
        persistent: true,
        has_ctime: true,
    }
}

blobstore_test_impl! {
    sqlblob_test_no_inline => {
        state: (),
        new: move |_, put_behaviour,| Sqlblob::with_sqlite_in_memory(put_behaviour, &(get_test_config_store().1), false, 0),
        persistent: true,
        has_ctime: true,
    }
}

blobstore_test_impl! {
    sqlblob_test_allow_inline => {
        state: (),
        new: move |_, put_behaviour,| Sqlblob::with_sqlite_in_memory(put_behaviour, &(get_test_config_store().1), true, 0),
        persistent: true,
        has_ctime: true,
    }
}

#[cfg(fbcode_build)]
fn create_cache(fb: FacebookInit) -> Result<(), Error> {
    let config = cachelib::LruCacheConfig::new(128 * 1024 * 1024);
    cachelib::init_cache(fb, config)?;
    Ok(())
}

#[cfg(fbcode_build)]
#[fbinit::test]
async fn test_cache_blob_maybe_zstd(fb: FacebookInit) -> Result<(), Error> {
    cache_blob_tests(fb, true).await
}

#[cfg(fbcode_build)]
#[fbinit::test]
async fn test_cache_blob_no_zstd(fb: FacebookInit) -> Result<(), Error> {
    cache_blob_tests(fb, false).await
}

#[cfg(fbcode_build)]
async fn cache_blob_tests(fb: FacebookInit, expect_zstd: bool) -> Result<(), Error> {
    let options = cacheblob::CachelibBlobstoreOptions::new_eager(Some(expect_zstd));
    let suffix = if expect_zstd { "_maybe_zstd" } else { "_raw" };

    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);
    create_cache(fb)?;
    let blob_pool = Arc::new(cachelib::get_or_create_pool(
        &["blob_pool", suffix].concat(),
        20 * 1024 * 1024,
    )?);
    let presence_pool = Arc::new(cachelib::get_or_create_pool(
        &["presence_pool", suffix].concat(),
        20 * 1024 * 1024,
    )?);

    let inner = Arc::new(Memblob::new(PutBehaviour::Overwrite));
    let cache_blob =
        cacheblob::new_cachelib_blobstore(inner.clone(), blob_pool.clone(), presence_pool, options);

    let small_key = "small_key";
    let value = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(b"smalldata"));
    cache_blob
        .put(ctx, small_key.to_owned(), value.clone())
        .await?;

    // Peek into cachelib to check its as expected
    let cachelib_len = blob_pool
        .get(small_key)
        .map(|bytes| bytes.map(|b| b.len()))?;
    assert!(cachelib_len.is_some());
    assert!(
        cachelib_len.unwrap() > value.as_bytes().len(),
        "Expected cachelib value to be larger due to framing"
    );

    assert_eq!(
        cache_blob
            .get(ctx, small_key)
            .await?
            .map(|bytes| bytes.into_bytes()),
        Some(value)
    );

    let large_key = "large_key";
    let size = 5 * 1024 * 1024;
    let mut large_value = Vec::with_capacity(size);
    large_value.resize(size, b'a');
    let large_value = BlobstoreBytes::from_bytes(large_value);

    cache_blob
        .put(ctx, large_key.to_owned(), large_value.clone())
        .await?;

    // Peek into cachelib to check its as expected
    let cachelib_len = blob_pool
        .get(large_key)
        .map(|bytes| bytes.map(|b| b.len()))?;

    if expect_zstd {
        assert!(cachelib_len.is_some());
        assert!(
            cachelib_len < Some(large_value.len()),
            "Expected cachelib value to be smaller due to compression"
        );
    } else {
        assert!(
            cachelib_len.is_none(),
            "Cachelib value is too large, so should not be in cachelib"
        );
    }

    // Check that inner blob is same size after put
    let inner_blob = inner.get(ctx, large_key).await?;
    assert_eq!(
        inner_blob.map(|b| b.as_bytes().len()),
        Some(large_value.len())
    );

    // Check that blob is the same when read through cache_blob's cachelib
    // layer
    assert_eq!(
        cache_blob
            .get(ctx, large_key)
            .await?
            .map(|bytes| bytes.into_bytes()),
        Some(large_value)
    );

    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::collections::VecDeque;
use std::fmt;
use std::future::Future;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::SystemTime;

use anyhow::bail;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreMetadata;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use context::CoreContext;
use futures::channel::oneshot;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use lock_ext::LockExt;
use metaconfig_types::MultiplexId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::Timestamp;
pub struct Tickable<T> {
    pub storage: Arc<Mutex<HashMap<String, T>>>,
    // queue of pending operations
    queue: Arc<Mutex<VecDeque<oneshot::Sender<Option<String>>>>>,
}

impl<T: fmt::Debug> fmt::Debug for Tickable<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Tickable")
            .field("storage", &self.storage)
            .field("pending", &self.queue.with(|q| q.len()))
            .finish()
    }
}

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

impl<T> Tickable<T> {
    pub fn new() -> Self {
        Self {
            storage: Default::default(),
            queue: Default::default(),
        }
    }

    // Broadcast either success or error to a set of outstanding futures, advancing the
    // overall state by one tick.
    pub fn tick(&self, error: Option<&str>) {
        let mut queue = self.queue.lock().unwrap();
        for send in queue.drain(..) {
            send.send(error.map(String::from)).unwrap();
        }
    }

    // Drain the queue of the first N pending requests. Helpful when the consumers
    // on the other end of the channels already dropped.
    pub fn drain(&self, num: usize) {
        let mut queue = self.queue.lock().unwrap();
        for _entry in queue.drain(0..num) {}
    }

    // Register this task on the tick queue and wait for it to progress.
    pub fn on_tick(&self) -> impl Future<Output = Result<()>> {
        let (send, recv) = oneshot::channel();
        let mut queue = self.queue.lock().unwrap();
        queue.push_back(send);
        async move {
            let error = tokio::time::timeout(std::time::Duration::from_secs(10), recv).await??;
            match error {
                None => Ok(()),
                Some(error) => bail!(error),
            }
        }
    }
}

impl Tickable<(BlobstoreBytes, u64)> {
    pub fn get_bytes(&self, key: &str) -> Option<BlobstoreBytes> {
        self.storage
            .with(|s| s.get(key).map(|(v, _ctime)| v).cloned())
    }

    pub fn add_bytes(&self, key: String, value: BlobstoreBytes) {
        let ctime = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        self.storage.with(|s| {
            s.insert(key, (value, ctime));
        })
    }
}

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

        on_tick.await?;
        Ok(storage.with(|s| {
            s.get(key).cloned().map(|(v, ctime)| {
                BlobstoreGetData::new(BlobstoreMetadata::new(Some(ctime as i64), None), v)
            })
        }))
    }

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

#[async_trait]
impl BlobstorePutOps for Tickable<(BlobstoreBytes, u64)> {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        self.on_tick().await?;
        if put_behaviour == PutBehaviour::IfAbsent {
            if self.storage.with(|s| s.contains_key(&key)) {
                return Ok(OverwriteStatus::Prevented);
            }
        }
        self.add_bytes(key, value);
        Ok(OverwriteStatus::NotChecked)
    }

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

#[async_trait]
impl BlobstoreWal for Tickable<BlobstoreWalEntry> {
    async fn log<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        entry: BlobstoreWalEntry,
    ) -> Result<BlobstoreWalEntry> {
        self.on_tick().await?;
        self.storage.with(|s| {
            s.insert(entry.blobstore_key.clone(), entry.clone());
        });
        Ok(entry)
    }

    async fn log_many<'a>(
        &'a self,
        ctx: &'a CoreContext,
        entries: Vec<BlobstoreWalEntry>,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        stream::iter(entries)
            .then(|entry| self.log(ctx, entry))
            .try_collect()
            .await
    }

    async fn read<'a>(
        &'a self,
        _c: &'a CoreContext,
        _u: &MultiplexId,
        _o: &Timestamp,
        _l: usize,
    ) -> Result<Vec<BlobstoreWalEntry>> {
        Ok(self.storage.with(|s| s.values().cloned().collect()))
    }

    async fn delete<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        entries: &'a [BlobstoreWalEntry],
    ) -> Result<()> {
        self.on_tick().await?;
        self.storage.with(|s| {
            for entry in entries {
                s.remove(&entry.blobstore_key);
            }
        });
        Ok(())
    }

    async fn delete_by_key(&self, ctx: &CoreContext, entries: &[BlobstoreWalEntry]) -> Result<()> {
        self.delete(ctx, entries).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::fmt;
use std::num::NonZeroU32;
use std::num::NonZeroUsize;
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 governor::clock::DefaultClock;
use governor::state::direct::NotKeyed;
use governor::state::InMemoryState;
use governor::Jitter;
use governor::Quota;
use governor::RateLimiter;
use mononoke_types::BlobstoreBytes;
use nonzero_ext::nonzero;

#[derive(Clone, Copy, Debug, Default)]
pub struct ThrottleOptions {
    pub read_qps: Option<NonZeroU32>,
    pub write_qps: Option<NonZeroU32>,
    pub read_bytes: Option<NonZeroUsize>,
    pub write_bytes: Option<NonZeroUsize>,
    pub read_burst_bytes: Option<NonZeroUsize>,
    pub write_burst_bytes: Option<NonZeroUsize>,
    pub bytes_min_count: Option<NonZeroUsize>,
}

impl ThrottleOptions {
    pub fn has_throttle(&self) -> bool {
        self.read_qps.is_some()
            || self.write_qps.is_some()
            || self.read_bytes.is_some()
            || self.write_bytes.is_some()
    }
}

fn bytes_to_count(bytes_min_count: usize, num_bytes: usize) -> NonZeroU32 {
    let count: u32 = (num_bytes / bytes_min_count).try_into().unwrap_or(u32::MAX);
    NonZeroU32::new(count).unwrap_or(nonzero!(1u32))
}

// The rate limiters use u32, so if we want to go > 4GiB/s need to scale the bytes to the rate limiter count.
// Default of 1_000 allows us to throttle in the range 1KB to 42.94GB/s (aka 40GiB/s). 40GiB/s should be enough for anybody :)
pub const DEFAULT_BYTES_MIN_COUNT: usize = 1_000;

// Any blobs over this the max don't attempt to throttle, instead they error.
// Default is set high as we'd rather throttle than error unless specified
pub const DEFAULT_BURST_BYTES_S: usize = 100_000_000;

/// A Blobstore that rate limits the number of read and write operations.
pub struct ThrottledBlob<T: fmt::Debug> {
    blobstore: T,
    read_qps_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
    write_qps_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
    read_bytes_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
    write_bytes_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
    bytes_min_count: usize,
    /// The options fields are used for Debug. They are not consulted at runtime.
    options: ThrottleOptions,
}

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

static JITTER_MAX: Duration = Duration::from_millis(5);

fn jitter() -> Jitter {
    Jitter::up_to(JITTER_MAX)
}

impl<T: fmt::Debug + Send + Sync> ThrottledBlob<T> {
    pub async fn new(blobstore: T, options: ThrottleOptions) -> Self {
        let qps_limiter =
            |qps: Option<NonZeroU32>| qps.map(|qps| RateLimiter::direct(Quota::per_second(qps)));
        let read_qps_limiter = qps_limiter(options.read_qps);
        let write_qps_limiter = qps_limiter(options.write_qps);

        let bytes_min_count = options
            .bytes_min_count
            .map_or(DEFAULT_BYTES_MIN_COUNT, |v| v.get());
        let bytes_limiter = |bytes_s: Option<NonZeroUsize>, burst_bytes_s: Option<NonZeroUsize>| {
            bytes_s.map(|bytes_s| {
                let count_s = bytes_to_count(bytes_min_count, bytes_s.get());
                RateLimiter::direct(Quota::per_second(count_s).allow_burst(
                    burst_bytes_s.map_or_else(
                        || bytes_to_count(bytes_min_count, DEFAULT_BURST_BYTES_S),
                        |burst_bytes_s| bytes_to_count(bytes_min_count, burst_bytes_s.get()),
                    ),
                ))
            })
        };
        let read_bytes_limiter = bytes_limiter(options.read_bytes, options.read_burst_bytes);
        let write_bytes_limiter = bytes_limiter(options.write_bytes, options.write_burst_bytes);

        Self {
            blobstore,
            read_qps_limiter,
            write_qps_limiter,
            read_bytes_limiter,
            write_bytes_limiter,
            bytes_min_count,
            options,
        }
    }

    // Convert from number of bytes to the count to request from until_n_ready
    fn count_n(&self, num_bytes: usize) -> NonZeroU32 {
        bytes_to_count(self.bytes_min_count, num_bytes)
    }
}

#[async_trait]
impl<T: Blobstore> Blobstore for ThrottledBlob<T> {
    // Thottling for get() bytes/s is approximate as we only know the size after the blob has been read
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        if let Some(limiter) = self.read_qps_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        if let Some(limiter) = self.read_bytes_limiter.as_ref() {
            // Only know we'll use some bytes. Access one count so we throttle if already over the limit
            limiter.until_ready_with_jitter(jitter()).await;
        }

        let get_data = self.blobstore.get(ctx, key).await?;

        if let Some(limiter) = self.read_bytes_limiter.as_ref() {
            // Now we know the size, request rest of the quota
            if let Some(data) = get_data.as_ref() {
                let count_n = self.count_n(data.as_bytes().len());
                let adjusted_n = NonZeroU32::new(count_n.get().saturating_sub(1));
                if let Some(adjusted_n) = adjusted_n {
                    limiter
                        .until_n_ready_with_jitter(adjusted_n, jitter())
                        .await?;
                }
            }
        }
        Ok(get_data)
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        if let Some(limiter) = self.write_qps_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        if let Some(limiter) = self.write_bytes_limiter.as_ref() {
            limiter
                .until_n_ready_with_jitter(self.count_n(value.len()), jitter())
                .await?;
        }
        self.blobstore.put(ctx, key, value).await
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        if let Some(limiter) = self.read_qps_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        // TODO(ahornby) would need to enhance Blobstore::is_present() to know how many bytes it transferred.
        // Some stores fetch just a flag, some fetch all the data then throw it away.
        if let Some(limiter) = self.read_bytes_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        self.blobstore.is_present(ctx, key).await
    }
}

#[async_trait]
impl<T: BlobstorePutOps> BlobstorePutOps for ThrottledBlob<T> {
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        if let Some(limiter) = self.write_qps_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        if let Some(limiter) = self.write_bytes_limiter.as_ref() {
            limiter
                .until_n_ready_with_jitter(self.count_n(value.len()), jitter())
                .await?;
        }
        self.blobstore
            .put_explicit(ctx, key, value, put_behaviour)
            .await
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        if let Some(limiter) = self.write_qps_limiter.as_ref() {
            limiter.until_ready_with_jitter(jitter()).await;
        }
        if let Some(limiter) = self.write_bytes_limiter.as_ref() {
            limiter
                .until_n_ready_with_jitter(self.count_n(value.len()), jitter())
                .await?;
        }
        self.blobstore.put_with_status(ctx, key, value).await
    }
}

impl<T: fmt::Debug> fmt::Debug for ThrottledBlob<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ThrottledBlob")
            .field("blobstore", &self.blobstore)
            .field("options", &self.options)
            .finish()
    }
}