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

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

use anyhow::Result;
use blobstore::BlobstoreEnumerableWithUnlink;
use chrono::Duration as ChronoDuration;
use context::CoreContext;
use derivative::Derivative;
use futures_lazy_shared::LazyShared;
use metaconfig_types::BubbleDeletionMode;
use mononoke_types::ChangesetId;
use mononoke_types::DateTime;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use sql_query_config::SqlQueryConfig;

use crate::bubble::Bubble;
use crate::bubble::BubbleId;
use crate::bubble::ExpiryStatus;
use crate::error::EphemeralBlobstoreError;

/// Ephemeral Store.
#[derive(Derivative)]
#[derivative(Debug)]
struct RepoEphemeralStoreInner {
    /// The backing blobstore where blobs are stored, without any redaction
    /// or repo prefix wrappers.
    pub(crate) blobstore: Arc<dyn BlobstoreEnumerableWithUnlink>,

    #[derivative(Debug = "ignore")]
    /// Config used to do SQL queries to underlying DB
    pub(crate) sql_config: Arc<SqlQueryConfig>,

    #[derivative(Debug = "ignore")]
    /// Database used to manage the ephemeral store.
    pub(crate) connections: SqlConnections,

    /// Initial value of the lifespan for bubbles in this store, i.e. the
    /// amount of time they last from the call to create.
    pub(crate) initial_bubble_lifespan: ChronoDuration,

    /// Grace period after bubbles expire during which requests which have
    /// already opened a bubble can continue to access them.  The bubble
    /// contents will not be deleted until after the grace period.
    pub(crate) bubble_expiration_grace: ChronoDuration,

    /// The mode of deletion to be used when cleaning up expired bubbles.
    /// The value determines if the bubbles need to be simply marked as
    /// expired or actually deleted from the physical store.
    pub(crate) bubble_deletion_mode: BubbleDeletionMode,
}

/// Ephemeral Store
#[facet::facet]
#[derive(Debug, Clone)]
pub struct RepoEphemeralStore {
    /// Repo this belongs to
    repo_id: RepositoryId,
    inner: Option<Arc<RepoEphemeralStoreInner>>,
}

mononoke_queries! {
    write CreateBubble(
        created_at: Timestamp,
        expires_at: Timestamp,
        owner_identity: Option<&str>,
    ) {
        none,
        "INSERT INTO ephemeral_bubbles (created_at, expires_at, owner_identity)
         VALUES ({created_at}, {expires_at}, {owner_identity})"
    }

    write AddBubbleLabels(
        values: (
            bubble_id: BubbleId,
            label: str,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO ephemeral_bubble_labels (bubble_id, label)
        VALUES {values}"
    }

    write DeleteBubbleLabels(
        id: BubbleId,
        >list labels: &str
    ) {
        none,
        "DELETE FROM ephemeral_bubble_labels WHERE
        bubble_id = {id} AND label IN {labels}"
    }

    write DeleteBubbleLabelsById(
        id: BubbleId,
    ) {
        none,
        "DELETE FROM ephemeral_bubble_labels WHERE
        bubble_id = {id}"
    }

    cacheable read SelectBubbleById(
        id: BubbleId,
    ) -> (Timestamp, ExpiryStatus, Option<String>) {
        "SELECT expires_at, expired, owner_identity FROM ephemeral_bubbles
         WHERE id = {id}"
    }

    read SelectBubbleLabelsById(
        id: BubbleId,
    ) -> (String, ) {
        "SELECT label FROM ephemeral_bubble_labels
        WHERE bubble_id = {id}"
    }

    read SelectBubbleFromChangeset(
        repo_id: RepositoryId,
        cs_id: ChangesetId,
    ) -> (BubbleId,) {
        "SELECT bubble_id
        FROM ephemeral_bubble_changeset_mapping
        WHERE repo_id = {repo_id} AND cs_id = {cs_id}"
    }

    read SelectChangesetFromBubble(
        id: BubbleId,
    ) -> (ChangesetId,) {
        "SELECT cs_id
        FROM ephemeral_bubble_changeset_mapping
        WHERE bubble_id = {id}"
    }

    read SelectBubblesWithExpiry(
        expires_at: Timestamp,
        limit: u32,
    ) -> (BubbleId,) {
        "SELECT B.id
        FROM ephemeral_bubbles B
        WHERE B.expires_at < {expires_at}
        AND NOT EXISTS (SELECT id from ephemeral_bubble_labels L WHERE B.id = L.bubble_id LIMIT 1)
        LIMIT {limit}"
    }

    read SelectBubblesWithExpiryAndStatus(
        expires_at: Timestamp,
        limit: u32,
        expiry_status: ExpiryStatus,
    ) -> (BubbleId,) {
        "SELECT B.id
        FROM ephemeral_bubbles B
        WHERE B.expires_at < {expires_at} AND B.expired = {expiry_status}
        AND NOT EXISTS (SELECT id FROM ephemeral_bubble_labels L WHERE B.id = L.bubble_id LIMIT 1)
        LIMIT {limit}"
    }

    write UpdateExpired(
        expired: ExpiryStatus,
        id: BubbleId
    ) {
        none,
        "UPDATE ephemeral_bubbles
        SET expired={expired}
        WHERE id={id}"
    }

    write DeleteBubble(
        id: BubbleId,
    ) {
        none,
        "DELETE
        FROM ephemeral_bubbles
        WHERE id={id} AND expired"
    }

    write DeleteBubbleChangesetMapping(
        id: BubbleId,
    ) {
        none,
        "DELETE
        FROM ephemeral_bubble_changeset_mapping
        WHERE bubble_id IN (SELECT id FROM ephemeral_bubbles WHERE id = {id} AND expired)"
    }

    write DeleteExpiredBubbleLabels(
        id: BubbleId,
    ) {
        none,
        "DELETE
        FROM ephemeral_bubble_labels
        WHERE bubble_id IN (SELECT id FROM ephemeral_bubbles WHERE id = {id} AND expired)"
    }
}

// Approximating a big duration to max chrono duration (~10^11 years) is good enough
fn to_chrono(duration: Duration) -> ChronoDuration {
    ChronoDuration::from_std(duration).unwrap_or_else(|_| ChronoDuration::max_value())
}

impl RepoEphemeralStoreInner {
    async fn create_bubble(
        &self,
        custom_duration: Option<Duration>,
        labels: Vec<String>,
    ) -> Result<Bubble> {
        let created_at = DateTime::now();
        let duration = match custom_duration {
            None => self.initial_bubble_lifespan,
            Some(duration) => to_chrono(duration),
        };
        let expires_at = created_at + duration;
        let txn = self
            .connections
            .write_connection
            .start_transaction()
            .await?;
        let (txn, res) = CreateBubble::query_with_transaction(
            txn,
            &Timestamp::from(created_at),
            &Timestamp::from(expires_at),
            &None,
        )
        .await?;

        match res.last_insert_id() {
            Some(id) if res.affected_rows() == 1 => {
                let bubble_id = BubbleId::new(
                    std::num::NonZeroU64::new(id)
                        .ok_or(EphemeralBlobstoreError::CreateBubbleFailed)?,
                );
                if !labels.is_empty() {
                    let bubble_labels = labels
                        .iter()
                        .map(|label| (&bubble_id, label as &str))
                        .collect::<Vec<_>>();
                    let (txn, _res) =
                        AddBubbleLabels::query_with_transaction(txn, bubble_labels.as_slice())
                            .await?;
                    txn.commit().await?;
                } else {
                    txn.commit().await?;
                };
                Ok(Bubble::new(
                    bubble_id,
                    expires_at + self.bubble_expiration_grace,
                    self.blobstore.clone(),
                    self.connections.clone(),
                    ExpiryStatus::Active,
                    LazyShared::new_ready(Ok(labels)),
                ))
            }
            _ => Err(EphemeralBlobstoreError::CreateBubbleFailed.into()),
        }
    }

    /// Add labels to an existing bubble
    async fn add_bubble_labels(&self, bubble_id: BubbleId, labels: Vec<String>) -> Result<()> {
        // Open the bubble to validate if the bubble exists and has not expired.
        self.open_bubble(bubble_id).await?;
        let bubble_labels = labels
            .iter()
            .map(|label| (&bubble_id, label.as_str()))
            .collect::<Vec<_>>();
        // The bubble exists, add labels to it.
        AddBubbleLabels::query(&self.connections.write_connection, bubble_labels.as_slice())
            .await?;
        Ok(())
    }

    /// Remove labels associated with an existing bubble
    async fn remove_bubble_labels(&self, bubble_id: BubbleId, labels: Vec<String>) -> Result<()> {
        // Open the bubble to validate if the bubble exists and has not expired.
        let bubble = self.open_bubble(bubble_id).await?;
        let labels = labels
            .iter()
            .map(|label| label.as_str())
            .collect::<Vec<_>>();
        // The bubble exists, remove labels from it.
        if labels.is_empty() {
            // No input labels were provided. In this case remove_labels should remove all
            // labels associated with the bubble.
            DeleteBubbleLabelsById::query(&self.connections.write_connection, &bubble.bubble_id())
                .await?;
        } else {
            // Specific labels were provided as input. Only remove those labels.
            DeleteBubbleLabels::query(
                &self.connections.write_connection,
                &bubble.bubble_id(),
                labels.as_slice(),
            )
            .await?;
        }
        Ok(())
    }

    /// Given a changeset ID, fetch the corresponding Bubble ID.
    async fn bubble_from_changeset(
        &self,
        repo_id: &RepositoryId,
        cs_id: &ChangesetId,
    ) -> Result<Option<BubbleId>> {
        let rows =
            SelectBubbleFromChangeset::query(&self.connections.read_connection, repo_id, cs_id)
                .await?;
        Ok(rows.into_iter().next().map(|b| b.0))
    }

    /// Given a Bubble ID, fetch labels corresponding to that bubble.
    async fn labels_from_bubble(&self, bubble_id: &BubbleId) -> Result<Vec<String>> {
        let rows =
            SelectBubbleLabelsById::query(&self.connections.read_connection, bubble_id).await?;
        Ok(rows.into_iter().map(|l| l.0).collect::<Vec<_>>())
    }

    async fn changesets_from_bubble(&self, bubble_id: &BubbleId) -> Result<Vec<ChangesetId>> {
        let rows =
            SelectChangesetFromBubble::query(&self.connections.read_connection, bubble_id).await?;
        Ok(rows.into_iter().map(|b| b.0).collect::<Vec<_>>())
    }

    /// Gets the vector of bubbles that are past their expiry period
    /// by atleast a duration of expiry_offset + bubble_expiration_grace
    async fn get_expired_bubbles(
        &self,
        expiry_offset: Duration,
        max_bubbles: u32,
    ) -> Result<Vec<BubbleId>> {
        let expiry_cutoff = DateTime::now() - to_chrono(expiry_offset);
        let rows = match self.bubble_deletion_mode {
            // If deletion mode is MarkOnly, we want to fetch only those
            // bubbles that are past expiry period but NOT marked as
            // expired yet (i.e. are active)
            BubbleDeletionMode::MarkOnly => {
                SelectBubblesWithExpiryAndStatus::query(
                    &self.connections.write_connection,
                    &Timestamp::from(expiry_cutoff - self.bubble_expiration_grace),
                    &max_bubbles,
                    &ExpiryStatus::Active,
                )
                .await?
            }
            // If hard delete is required, we want to fetch bubbles regardless
            // of their expiry status as long as they are past expiry date.
            _ => {
                SelectBubblesWithExpiry::query(
                    &self.connections.write_connection,
                    &Timestamp::from(expiry_cutoff - self.bubble_expiration_grace),
                    &max_bubbles,
                )
                .await?
            }
        };
        Ok(rows.into_iter().map(|b| b.0).collect::<Vec<_>>())
    }

    async fn keys_in_bubble(
        &self,
        bubble_id: BubbleId,
        ctx: &CoreContext,
        start_from: Option<String>,
        max: u32,
    ) -> Result<Vec<String>> {
        let bubble = self.open_bubble_raw(bubble_id, false).await?;
        bubble.keys_in_bubble(ctx, start_from, max).await
    }

    /// Method responsible for deleting the bubble and all the data contained within.
    /// Returns the number of blobs deleted from the bubble.
    async fn delete_bubble(&self, bubble_id: BubbleId, ctx: &CoreContext) -> Result<usize> {
        // Step 0: Validate if bubble deletion is enabled.
        if let BubbleDeletionMode::Disabled = self.bubble_deletion_mode {
            return Err(EphemeralBlobstoreError::DeleteBubbleDisabled.into());
        }
        // Step 1: Mark the bubble as expired in the backing SQL Store.
        let res = UpdateExpired::query(
            &self.connections.write_connection,
            &ExpiryStatus::Expired,
            &bubble_id,
        )
        .await?;
        if res.affected_rows() != 1 {
            return Err(EphemeralBlobstoreError::DeleteBubbleFailed(bubble_id).into());
        }
        // If only marking is required, exit now.
        if let BubbleDeletionMode::MarkOnly = self.bubble_deletion_mode {
            return Ok(0); // Since 0 blob items were unlinked/removed.
        }
        // Step 2: Delete the blob content within the expired bubble.
        let bubble = self.open_bubble_raw(bubble_id, false).await?;
        let count = bubble.delete_blobs_in_bubble(ctx).await?;

        // Step 3: Delete the metadata associated with the bubble from
        // the backing SQL store.
        let res =
            DeleteBubbleChangesetMapping::query(&self.connections.write_connection, &bubble_id)
                .await?;
        if res.affected_rows() > 1 {
            return Err(EphemeralBlobstoreError::DeleteBubbleFailed(bubble_id).into());
        }
        // When invoked through processes, this will always be a no-op. This is useful for
        // manually deleting a bubble regardless of its expiry status.
        DeleteExpiredBubbleLabels::query(&self.connections.write_connection, &bubble_id).await?;

        // Step 4: Delete the bubble itself from the backing SQL store.
        let res = DeleteBubble::query(&self.connections.write_connection, &bubble_id).await?;
        if res.affected_rows() > 1 {
            return Err(EphemeralBlobstoreError::DeleteBubbleFailed(bubble_id).into());
        }
        Ok(count)
    }

    async fn open_bubble_raw(&self, bubble_id: BubbleId, fail_on_expired: bool) -> Result<Bubble> {
        let mut bubble_rows = SelectBubbleById::query(
            self.sql_config.as_ref(),
            &self.connections.read_connection,
            &bubble_id,
        )
        .await?;
        let labels_future = {
            let connection = self.connections.read_connection.clone();
            let bubble_id = bubble_id.clone();
            async move {
                let label_rows = SelectBubbleLabelsById::query(&connection, &bubble_id)
                    .await
                    .map_err(|_| EphemeralBlobstoreError::FetchBubbleLabelsFailed(bubble_id))?;
                let labels = label_rows.into_iter().map(|l| l.0).collect::<Vec<_>>();
                Ok(labels)
            }
        };
        let labels_lazy = LazyShared::new_future(labels_future);

        if bubble_rows.is_empty() {
            // Perhaps the bubble hasn't showed up yet due to replication lag.
            // Let's retry on master just in case.
            bubble_rows = SelectBubbleById::query(
                self.sql_config.as_ref(),
                &self.connections.read_master_connection,
                &bubble_id,
            )
            .await?;
            if bubble_rows.is_empty() {
                return Err(EphemeralBlobstoreError::NoSuchBubble(bubble_id).into());
            }
        }
        // TODO(mbthomas): check owner_identity
        let (expires_at, expiry_status, ref _owner_identity) = bubble_rows[0];
        let expires_at: DateTime = expires_at.into();
        // Checking only the expiry_status and not the expired_date since a bubble could
        // be active even if its past its expiry date if it has labels associated with it.
        if fail_on_expired && expiry_status == ExpiryStatus::Expired {
            return Err(EphemeralBlobstoreError::NoSuchBubble(bubble_id).into());
        }

        Ok(Bubble::new(
            bubble_id,
            expires_at + self.bubble_expiration_grace,
            self.blobstore.clone(),
            self.connections.clone(),
            expiry_status,
            labels_lazy,
        ))
    }

    async fn open_bubble(&self, bubble_id: BubbleId) -> Result<Bubble> {
        self.open_bubble_raw(bubble_id, true).await
    }
}

impl RepoEphemeralStore {
    pub(crate) fn new(
        repo_id: RepositoryId,
        connections: SqlConnections,
        blobstore: Arc<dyn BlobstoreEnumerableWithUnlink>,
        sql_config: Arc<SqlQueryConfig>,
        initial_bubble_lifespan: Duration,
        bubble_expiration_grace: Duration,
        bubble_deletion_mode: BubbleDeletionMode,
    ) -> Self {
        Self {
            repo_id,
            inner: Some(Arc::new(RepoEphemeralStoreInner {
                blobstore,
                sql_config,
                connections,
                initial_bubble_lifespan: to_chrono(initial_bubble_lifespan),
                bubble_expiration_grace: to_chrono(bubble_expiration_grace),
                bubble_deletion_mode,
            })),
        }
    }

    pub fn disabled(repo_id: RepositoryId) -> Self {
        Self {
            inner: None,
            repo_id,
        }
    }

    fn inner(&self) -> Result<&RepoEphemeralStoreInner, EphemeralBlobstoreError> {
        self.inner
            .as_deref()
            .ok_or(EphemeralBlobstoreError::NoEphemeralBlobstore(self.repo_id))
    }

    pub async fn create_bubble(
        &self,
        custom_duration: Option<Duration>,
        labels: Vec<String>,
    ) -> Result<Bubble> {
        self.inner()?.create_bubble(custom_duration, labels).await
    }

    /// Method responsible for deleting the bubble and all the data contained within.
    /// Returns the number of blobs deleted from the bubble.
    /// NOTE: Deletes the bubble regardless of its expiry status. Make sure the bubble
    /// is suitable for deletion.
    pub async fn delete_bubble(&self, bubble_id: BubbleId, ctx: &CoreContext) -> Result<usize> {
        self.inner()?.delete_bubble(bubble_id, ctx).await
    }

    /// Gets the vector of bubbles that are past their expiry period
    /// by atleast a duration of expiry_offset + bubble_expiration_grace
    pub async fn get_expired_bubbles(
        &self,
        expiry_offset: Duration,
        max_bubbles: u32,
    ) -> Result<Vec<BubbleId>> {
        self.inner()?
            .get_expired_bubbles(expiry_offset, max_bubbles)
            .await
    }

    /// Gets the blob keys stored within the bubble, optionally starting
    /// from 'start_from' and upto 'max' in count.
    pub async fn keys_in_bubble(
        &self,
        bubble_id: BubbleId,
        ctx: &CoreContext,
        start_from: Option<String>,
        max: u32,
    ) -> Result<Vec<String>> {
        self.inner()?
            .keys_in_bubble(bubble_id, ctx, start_from, max)
            .await
    }

    /// Open the bubble corresponding to the given bubble ID if the bubble
    /// exists and has not yet expired.
    pub async fn open_bubble(&self, bubble_id: BubbleId) -> Result<Bubble> {
        self.inner()?.open_bubble(bubble_id).await
    }

    /// Open the bubble corresponding to the given bubble ID regardless
    /// of the expiry status or date.
    /// NOTE: To be used only for debugging, use open_bubble for other
    /// production use cases.
    pub async fn open_bubble_raw(&self, bubble_id: BubbleId) -> Result<Bubble> {
        self.inner()?.open_bubble_raw(bubble_id, false).await
    }

    /// Given a changeset ID, fetch the corresponding bubble ID.
    pub async fn bubble_from_changeset(&self, cs_id: &ChangesetId) -> Result<Option<BubbleId>> {
        self.inner()?
            .bubble_from_changeset(&self.repo_id, cs_id)
            .await
    }

    /// Given a bubble ID, fetch the corresponding changeset ID within the
    /// repository associated with the bubble.
    pub async fn changesets_from_bubble(&self, bubble_id: &BubbleId) -> Result<Vec<ChangesetId>> {
        self.inner()?.changesets_from_bubble(bubble_id).await
    }

    /// Associate the given labels with the bubble corresponding to the input
    /// bubble ID.
    pub async fn add_bubble_labels(&self, bubble_id: BubbleId, labels: Vec<String>) -> Result<()> {
        self.inner()?.add_bubble_labels(bubble_id, labels).await
    }

    /// Disassociate the given labels from the bubble corresponding to the input
    /// bubble ID.
    pub async fn remove_bubble_labels(
        &self,
        bubble_id: BubbleId,
        labels: Vec<String>,
    ) -> Result<()> {
        self.inner()?.remove_bubble_labels(bubble_id, labels).await
    }

    /// Given a bubble ID, fetches the labels corresponding to that bubble.
    pub async fn labels_from_bubble(&self, bubble_id: &BubbleId) -> Result<Vec<String>> {
        self.inner()?.labels_from_bubble(bubble_id).await
    }
}

#[cfg(test)]
mod test {
    use anyhow::anyhow;
    use blobstore::Blobstore;
    use blobstore::BlobstoreBytes;
    use blobstore::BlobstoreEnumerableWithUnlink;
    use blobstore::BlobstoreKeyParam;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use maplit::hashset;
    use memblob::Memblob;
    use metaconfig_types::BubbleDeletionMode;
    use metaconfig_types::PackFormat;
    use mononoke_types_mocks::repo::REPO_ZERO;
    use packblob::PackBlob;
    use repo_blobstore::RepoBlobstore;
    use scuba_ext::MononokeScubaSampleBuilder;
    use sql_construct::SqlConstruct;

    use super::*;
    use crate::builder::RepoEphemeralStoreBuilder;

    fn bootstrap(
        fb: FacebookInit,
        initial_lifespan: Duration,
        grace_period: Duration,
        deletion_mode: BubbleDeletionMode,
    ) -> Result<(
        CoreContext,
        Arc<dyn BlobstoreEnumerableWithUnlink>,
        RepoBlobstore,
        RepoEphemeralStore,
    )> {
        let ctx = CoreContext::test_mock(fb);
        // The ephemeral blobstore will normally be used stacked on top of
        // packblob, so use this in the test, too.
        let blobstore = Arc::new(PackBlob::new(
            Memblob::default(),
            PackFormat::ZstdIndividual(0),
        )) as Arc<dyn BlobstoreEnumerableWithUnlink>;
        let repo_blobstore = RepoBlobstore::new(
            Arc::new(Memblob::default()),
            None,
            REPO_ZERO,
            MononokeScubaSampleBuilder::with_discard(),
        );
        let sql_config = Arc::new(SqlQueryConfig { caching: None });
        let eph = RepoEphemeralStoreBuilder::with_sqlite_in_memory()?.build(
            REPO_ZERO,
            blobstore.clone(),
            sql_config,
            initial_lifespan,
            grace_period,
            deletion_mode,
        );
        Ok((ctx, blobstore, repo_blobstore, eph))
    }

    #[fbinit::test]
    async fn basic_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (ctx, blobstore, repo_blobstore, eph) =
            bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let key = "test_key".to_string();

        // Create a bubble and put data in it.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        let bubble1_id = bubble1.bubble_id();
        let bubble1 = bubble1.wrap_repo_blobstore(repo_blobstore.clone());
        bubble1
            .put(&ctx, key.clone(), BlobstoreBytes::from_bytes("test data"))
            .await?;
        let data = bubble1.get(&ctx, &key).await?.unwrap().into_bytes();
        assert_eq!(data.as_bytes().as_ref(), b"test data");

        // Re-open the bubble and confirm we can read the data.
        let bubble1_read = eph
            .open_bubble(bubble1_id)
            .await?
            .wrap_repo_blobstore(repo_blobstore.clone());
        let data = bubble1_read.get(&ctx, &key).await?.unwrap().into_bytes();
        assert_eq!(data.as_bytes().as_ref(), b"test data");

        // Enumerate the blobstore and check the key got its prefix.
        let enumerated = blobstore
            .enumerate(&ctx, &BlobstoreKeyParam::from(..))
            .await?;
        assert_eq!(
            enumerated.keys,
            hashset! { format!("eph{}.repo0000.{}", bubble1_id, key) }
        );

        // Create a new bubble and put data in it.
        let bubble2 = eph.create_bubble(None, vec![]).await?;
        let bubble2_id = bubble2.bubble_id();
        let bubble2 = bubble2.wrap_repo_blobstore(repo_blobstore.clone());
        bubble2
            .put(
                &ctx,
                key.clone(),
                BlobstoreBytes::from_bytes("other test data"),
            )
            .await?;
        let data = bubble2.get(&ctx, &key).await?.unwrap().into_bytes();
        assert_eq!(data.as_bytes().as_ref(), b"other test data");

        let data = bubble1.get(&ctx, &key).await?.unwrap().into_bytes();
        assert_eq!(data.as_bytes().as_ref(), b"test data");

        // There should now be two separate keys.
        let enumerated = blobstore
            .enumerate(&ctx, &BlobstoreKeyParam::from(..))
            .await?;
        assert_eq!(
            enumerated.keys,
            hashset! {
                format!("eph{}.repo0000.{}", bubble1_id, key),
                format!("eph{}.repo0000.{}", bubble2_id, key),
            }
        );
        Ok(())
    }

    #[fbinit::test]
    async fn basic_test_with_labels(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        // Create a bubble with labels associated to it.
        let bubble1 = eph.create_bubble(None, labels.clone()).await?;
        // Verify all the labels are associated with the bubble.
        assert!(
            bubble1
                .labels()
                .await?
                .iter()
                .zip(&labels)
                .filter(|&(l, r)| l == r)
                .count()
                == labels.len()
        );
        Ok(())
    }

    #[fbinit::test]
    async fn add_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Ensure that the bubble is created with no labels.
        assert!(bubble.labels().await?.is_empty());
        // Add labels to the newly created bubble.
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        eph.add_bubble_labels(bubble.bubble_id(), labels.clone())
            .await?;
        // Reopen bubble from the store and verify it has the added labels.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?.len(), labels.len());
        assert!(
            bubble_read
                .labels()
                .await?
                .iter()
                .all(|label| labels.contains(label))
        );
        Ok(())
    }

    #[fbinit::test]
    async fn add_duplicate_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Ensure that the bubble is created with no labels.
        assert!(bubble.labels().await?.is_empty());
        // Add unique + duplicate labels to the newly created bubble.
        let labels = vec![
            "workspace".to_string(),
            "debug_version".to_string(),
            "workspace".to_string(),
            "workspace".to_string(),
        ];
        eph.add_bubble_labels(bubble.bubble_id(), labels).await?;
        // Reopen bubble from the store and verify it has the added labels
        // and the labels are not duplicated.
        let unique_labels = vec!["workspace".to_string(), "debug_version".to_string()];
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?.len(), unique_labels.len());
        assert!(
            bubble_read
                .labels()
                .await?
                .iter()
                .all(|label| unique_labels.contains(label))
        );
        Ok(())
    }

    #[fbinit::test]
    async fn add_empty_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Ensure that the bubble is created with no labels.
        assert!(bubble.labels().await?.is_empty());
        // Add an empty vec of labels.
        let labels = vec![];
        eph.add_bubble_labels(bubble.bubble_id(), labels).await?;
        // Reopen bubble from the store and verify it still doesn't
        // have any labels since we used an empty vec of labels as input.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert!(bubble_read.labels().await?.is_empty());
        Ok(())
    }

    #[fbinit::test]
    async fn empty_labels_from_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with no labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Fetch the labels associated with the newly created bubble.
        let id = bubble.bubble_id();
        let labels = eph.labels_from_bubble(&id).await?;
        // Validate that no labels are associated with the bubble.
        assert!(labels.is_empty());
        Ok(())
    }

    #[fbinit::test]
    async fn non_empty_labels_from_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let labels = vec!["workspace".to_string(), "debug".to_string()];
        // Create a bubble with some labels associated to it.
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Fetch the labels associated with the newly created bubble.
        let id = bubble.bubble_id();
        let returned_labels = eph.labels_from_bubble(&id).await?;
        // Validate that the labels returned are the same as the stored labels.
        assert_eq!(returned_labels.len(), labels.len());
        assert!(returned_labels.iter().all(|label| labels.contains(label)));
        Ok(())
    }

    #[fbinit::test]
    async fn added_removed_labels_from_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with no labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Add labels to the newly created bubble.
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        eph.add_bubble_labels(bubble.bubble_id(), labels.clone())
            .await?;
        // Fetch the labels associated with the bubble.
        let id = bubble.bubble_id();
        let returned_labels = eph.labels_from_bubble(&id).await?;
        // Validate that the labels returned are the same as the added labels.
        assert_eq!(returned_labels.len(), labels.len());
        assert!(returned_labels.iter().all(|label| labels.contains(label)));
        // Remove all labels associated with the bubble.
        eph.remove_bubble_labels(bubble.bubble_id(), labels).await?;
        let returned_labels = eph.labels_from_bubble(&id).await?;
        // Validate that no labels are returned since all labels were deleted.
        assert!(returned_labels.is_empty());
        Ok(())
    }

    #[fbinit::test]
    async fn remove_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Reopen bubble and verify all the labels are associated with the bubble.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?.len(), labels.len());
        assert!(
            bubble_read
                .labels()
                .await?
                .iter()
                .all(|label| labels.contains(label))
        );
        // Remove all labels associated with the bubble.
        eph.remove_bubble_labels(bubble.bubble_id(), labels).await?;
        // Reopen the bubble and validate that it has no labels associated with it.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert!(bubble_read.labels().await?.is_empty());
        Ok(())
    }

    #[fbinit::test]
    async fn partial_remove_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let bubble = eph
            .create_bubble(
                None,
                vec!["workspace".to_string(), "debug_version".to_string()],
            )
            .await?;
        // Remove a subset of the labels associated with the bubble.
        eph.remove_bubble_labels(bubble.bubble_id(), vec!["workspace".to_string()])
            .await?;
        // Reopen the bubble and validate it has the remaining labels associated with it.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(
            bubble_read.labels().await?,
            vec!["debug_version".to_string()]
        );
        Ok(())
    }

    #[fbinit::test]
    async fn remove_absent_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Remove labels that are not part of the bubble. This should be a no-op
        // from the bubble's perspective.
        eph.remove_bubble_labels(bubble.bubble_id(), vec!["some_random_label".to_string()])
            .await?;
        // Reopen the bubble and validate it has the same labels as before since the
        // no existing labels were deleted.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?.len(), labels.len());
        assert!(
            bubble_read
                .labels()
                .await?
                .iter()
                .all(|label| labels.contains(label))
        );
        Ok(())
    }

    #[fbinit::test]
    async fn remove_duplicate_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Remove labels that are not part of the bubble. This should be a no-op
        // from the bubble's perspective.
        eph.remove_bubble_labels(
            bubble.bubble_id(),
            vec![
                "workspace".to_string(),
                "workspace".to_string(),
                "workspace".to_string(),
            ],
        )
        .await?;
        // Reopen the bubble and validate it has only "debug_version" label since
        // the other label was deleted. Passing input vec with duplicate labels
        // should not have any other effect.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(
            bubble_read.labels().await?,
            vec!["debug_version".to_string()]
        );
        Ok(())
    }

    #[fbinit::test]
    async fn remove_all_bubble_labels_without_input_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Remove bubble labels without specifying the exact labels to remove.
        eph.remove_bubble_labels(bubble.bubble_id(), Vec::new())
            .await?;
        // Reopen the bubble and validate that all the labels have been removed.
        // Calling remove_bubble_labels without any input labels should remove
        // all the associated labels.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?, Vec::<String>::new());
        Ok(())
    }

    #[fbinit::test]
    async fn remove_all_bubble_labels_with_empty_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with no labels associated to it.
        let bubble = eph.create_bubble(None, Vec::new()).await?;
        // Remove bubble labels without specifying the exact labels to remove.
        // Even though no labels are associated with the bubble, operation of
        // removing all labels still succeeds as a no-op.
        eph.remove_bubble_labels(bubble.bubble_id(), Vec::new())
            .await?;
        // Reopen the bubble and validate that the state of labels is the
        // same as before
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?, Vec::<String>::new());
        Ok(())
    }

    #[fbinit::test]
    async fn add_and_remove_bubble_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Ensure that the bubble is created with no labels.
        assert!(bubble.labels().await?.is_empty());
        // Add labels to the newly created bubble.
        let labels = vec![
            "workspace".to_string(),
            "debug_version".to_string(),
            "some_label".to_string(),
            "important_snapshot".to_string(),
        ];
        eph.add_bubble_labels(bubble.bubble_id(), labels).await?;
        // Remove a partial subset of the labels associated to the bubble.
        eph.remove_bubble_labels(
            bubble.bubble_id(),
            vec!["workspace".to_string(), "debug_version".to_string()],
        )
        .await?;
        // Reopen the bubble and verify that it has the right set of labels after the
        // add and remove operation.
        let remaining_labels = vec!["some_label".to_string(), "important_snapshot".to_string()];
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        assert_eq!(bubble_read.labels().await?.len(), remaining_labels.len());
        assert!(
            bubble_read
                .labels()
                .await?
                .iter()
                .all(|label| remaining_labels.contains(label))
        );
        Ok(())
    }

    #[fbinit::test]
    async fn create_and_fetch_labels_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let labels = vec!["workspace".to_string(), "debug_version".to_string()];
        // Create a bubble with labels associated to it.
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Re-opening the bubble from storage.
        let bubble_read = eph.open_bubble(bubble.bubble_id()).await?;
        // Validate that the labels in the stored bubble and the retrieved bubble are the same.
        let bubble_labels = bubble.labels().await?;
        let bubble_read_labels = bubble_read.labels().await?;
        assert_eq!(bubble_labels.len(), bubble_read_labels.len());
        assert!(
            bubble_labels
                .iter()
                .all(|label| bubble_read_labels.contains(label))
        );
        Ok(())
    }

    #[fbinit::test]
    async fn create_and_fetch_active_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Ensure a newly created bubble exists in Active status
        assert_eq!(bubble1.expired(), ExpiryStatus::Active);
        // Re-opening the bubble from storage returns the same status
        let bubble1_read = eph.open_bubble(bubble1.bubble_id()).await?;
        assert_eq!(bubble1_read.expired(), bubble1.expired());
        Ok(())
    }

    #[fbinit::test]
    async fn deletion_mode_disabled_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        // We want an ephemeral store where deletion is disabled.
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::Disabled)?;
        // Create an empty bubble.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Attempt to delete the bubble
        let res = eph.delete_bubble(bubble1.bubble_id(), &ctx).await;
        // Since the bubble is deleted, reopening the bubble should
        // throw the "no such bubble" error
        match res {
            Err(e) => match e.downcast_ref::<EphemeralBlobstoreError>() {
                Some(EphemeralBlobstoreError::DeleteBubbleDisabled) => Ok(()),
                _ => Err(anyhow!("Invalid error during bubble deletion")),
            },
            _ => Err(anyhow!(
                "Bubble deletion should be disabled but it is enabled"
            )),
        }
    }

    #[fbinit::test]
    async fn deletion_mode_markonly_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        // We want an ephemeral store where deletion mode is mark only.
        let (ctx, blobstore, repo_blobstore, eph) =
            bootstrap(fb, initial, grace, BubbleDeletionMode::MarkOnly)?;
        // Create a bubble and add data to it.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        let bubble1_repo = bubble1.wrap_repo_blobstore(repo_blobstore.clone());
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_1"),
                BlobstoreBytes::from_bytes("test data 1"),
            )
            .await?;
        // Add more data to it
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_2"),
                BlobstoreBytes::from_bytes("test data 2"),
            )
            .await?;
        // Enumerate the blobstore and check the required data is present.
        let enumerated = blobstore
            .enumerate(&ctx, &BlobstoreKeyParam::from(..))
            .await?;
        // Should contain two keys for now
        assert_eq!(enumerated.keys.len(), 2);
        // Delete the bubble
        let deleted = eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        // Should be 0 since deletion mode is MarkOnly
        assert_eq!(deleted, 0);
        // Even though the bubble hasn't been deleted from the backing physical store
        // it should appear as deleted to the user and hence not be accessible.
        let res = eph.open_bubble(bubble1.bubble_id()).await;
        // Since the bubble is deleted, reopening the bubble should
        // throw the "no such bubble" error
        match res {
            Err(e) => match e.downcast_ref::<EphemeralBlobstoreError>() {
                Some(EphemeralBlobstoreError::NoSuchBubble(_)) => Ok(()),
                _ => Err(anyhow!("Invalid error post bubble deletion")),
            },
            _ => Err(anyhow!(
                "Bubble expected to be (soft) deleted but it still exists for the user"
            )),
        }
    }

    #[fbinit::test]
    async fn delete_empty_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Delete the bubble
        let deleted = eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        // Should be 0 since the bubble was empty
        assert_eq!(deleted, 0);
        Ok(())
    }

    #[fbinit::test]
    async fn delete_empty_bubble_with_label_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble with labels.
        let labels = vec!["workspace".to_string(), "test".to_string()];
        let bubble1 = eph.create_bubble(None, labels).await?;
        // Delete the bubble
        let deleted = eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        // Should be 0 since the bubble was empty
        assert_eq!(deleted, 0);
        Ok(())
    }

    #[fbinit::test]
    async fn delete_nonempty_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (ctx, blobstore, repo_blobstore, eph) =
            bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble and add data to it.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        let bubble1_repo = bubble1.wrap_repo_blobstore(repo_blobstore.clone());
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_1"),
                BlobstoreBytes::from_bytes("test data 1"),
            )
            .await?;
        // Add more data to it
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_2"),
                BlobstoreBytes::from_bytes("test data 2"),
            )
            .await?;
        // Add some more data to it
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_3"),
                BlobstoreBytes::from_bytes("test data 3"),
            )
            .await?;
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_4"),
                BlobstoreBytes::from_bytes("test data 4"),
            )
            .await?;
        // Enumerate the blobstore and check the required data is present.
        let enumerated = blobstore
            .enumerate(&ctx, &BlobstoreKeyParam::from(..))
            .await?;
        // Should contain four keys for now
        assert_eq!(enumerated.keys.len(), 4);
        // Delete the bubble
        let deleted = eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        // Should be 4 based on the input data added
        assert_eq!(deleted, 4);
        Ok(())
    }

    #[fbinit::test]
    async fn reopen_deleted_bubble_test(fb: FacebookInit) -> Result<()> {
        let initial = Duration::from_secs(30 * 24 * 60 * 60);
        let grace = Duration::from_secs(6 * 60 * 60);
        let (ctx, _, repo_blobstore, eph) =
            bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create a bubble and add data to it.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        let bubble1_repo = bubble1.wrap_repo_blobstore(repo_blobstore.clone());
        bubble1_repo
            .put(
                &ctx,
                String::from("test_key_1"),
                BlobstoreBytes::from_bytes("test data 1"),
            )
            .await?;
        // Delete the bubble
        eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        let res = eph.open_bubble(bubble1.bubble_id()).await;
        // Since the bubble is deleted, reopening the bubble should
        // throw the "no such bubble" error
        match res {
            Err(e) => match e.downcast_ref::<EphemeralBlobstoreError>() {
                Some(EphemeralBlobstoreError::NoSuchBubble(_)) => Ok(()),
                _ => Err(anyhow!("Invalid error post bubble deletion")),
            },
            _ => Err(anyhow!("Bubble expected to be deleted but it still exists")),
        }
    }

    #[fbinit::test]
    async fn get_expired_bubbles_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble that would expire immediately.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Validate bubble is created in active state
        assert_eq!(bubble1.expired(), ExpiryStatus::Active);
        // Create an empty bubble that won't expire anytime soon.
        let bubble2 = eph
            .create_bubble(Some(Duration::from_secs(10000)), vec![])
            .await?;
        // Validate bubble is created in active state
        assert_eq!(bubble2.expired(), ExpiryStatus::Active);
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        // Only one bubble should be returned since only one has expired so far
        assert_eq!(res.len(), 1);
        let res_bubble_id = res
            .first()
            .expect("Invalid number of expired bubbles")
            .clone();
        // The first bubble should be the only one that's expired
        assert_eq!(res_bubble_id, bubble1.bubble_id());
        Ok(())
    }

    #[fbinit::test]
    async fn get_expired_bubbles_with_labels_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble that would expire immediately but has labels
        // associated with it.
        let labels = vec!["workspace".to_string(), "test".to_string()];
        eph.create_bubble(None, labels).await?;
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        // Even though there exists one bubble beyond its expiry time, it should
        // not be returned since it has labels associated with it.
        assert_eq!(res.len(), 0);
        Ok(())
    }

    #[fbinit::test]
    async fn get_expired_bubbles_with_removed_labels_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble that would expire immediately but has labels
        // associated with it.
        let labels = vec!["workspace".to_string(), "test".to_string()];
        let bubble = eph.create_bubble(None, labels.clone()).await?;
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        // Even though there exists one bubble beyond its expiry time, it should
        // not be returned since it has labels associated with it.
        assert_eq!(res.len(), 0);
        // Remove the labels associated with the bubble.
        eph.remove_bubble_labels(bubble.bubble_id(), labels).await?;
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        // Now that the labels are removed, we should get the above bubble as
        // an expired bubble.
        assert_eq!(res.len(), 1);
        let res_bubble_id = res.first().expect("Invalid number of expired bubbles");
        assert_eq!(res_bubble_id, &bubble.bubble_id());
        Ok(())
    }

    #[fbinit::test]
    async fn add_labels_to_expired_bubble_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        // Create in MarkOnly mode since we want to soft delete to test expiration
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkOnly)?;
        // Create an empty bubble that would expire immediately.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Mark bubble as expired (i.e. soft delete)
        eph.delete_bubble(bubble.bubble_id(), &ctx).await?;
        // Add new labels to the bubble. This operation should fail since adding or
        // removing labels from an expired bubble is not permitted.
        let res = eph
            .add_bubble_labels(
                bubble.bubble_id(),
                vec!["workspace".to_string(), "test".to_string()],
            )
            .await;
        assert!(res.is_err());
        Ok(())
    }

    #[fbinit::test]
    async fn remove_labels_from_expired_bubble_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        // Create in MarkOnly mode since we want to soft delete to test expiration
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkOnly)?;
        // Create an empty bubble that would expire immediately.
        let bubble = eph.create_bubble(None, vec![]).await?;
        // Mark bubble as expired (i.e. soft delete)
        eph.delete_bubble(bubble.bubble_id(), &ctx).await?;
        let res = eph
            .remove_bubble_labels(
                bubble.bubble_id(),
                vec!["workspace".to_string(), "test".to_string()],
            )
            .await;
        assert!(res.is_err());
        Ok(())
    }

    #[fbinit::test]
    async fn get_expired_bubbles_offset_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble that would expire immediately.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Validate bubble is created in active state
        assert_eq!(bubble1.expired(), ExpiryStatus::Active);
        // Create an empty bubble that won't expire anytime soon.
        let bubble2 = eph
            .create_bubble(Some(Duration::from_secs(10000)), vec![])
            .await?;
        // Validate bubble is created in active state
        assert_eq!(bubble2.expired(), ExpiryStatus::Active);
        // Get expired bubbles
        let res = eph
            .get_expired_bubbles(Duration::from_secs(1000), 10)
            .await?;
        // No items should be returned since there aren't any bubbles
        // that have been expired for atleast the past 1000 seconds
        assert_eq!(res.len(), 0);
        Ok(())
    }

    #[fbinit::test]
    async fn get_expired_bubbles_markonly_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        // We want an ephemeral store that only marks the bubbles as expired.
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkOnly)?;
        // Create empty bubbles that would expire immediately.
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        // All 3 bubbles qualify for expiry and currently have their
        // status as ACTIVE, hence should be present in res
        assert_eq!(res.len(), 3);
        // Delete the bubbles in mark-only mode
        for id in res.into_iter() {
            eph.delete_bubble(id, &ctx).await?;
        }
        // The bubbles should now be soft-deleted but still be present in the
        // backing table. However, get_expired_bubbles should not return those
        // bubbles since their status will be EXPIRED.
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 10).await?;
        assert_eq!(res.len(), 0);
        Ok(())
    }

    #[fbinit::test]
    async fn get_n_expired_bubbles_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create empty bubbles that would expire immediately.
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        eph.create_bubble(None, vec![]).await?;
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 2).await?;
        // Only 2 bubbles should be returned given the input limit of 2
        assert_eq!(res.len(), 2);
        // Get expired bubbles
        let res = eph.get_expired_bubbles(Duration::from_secs(0), 0).await?;
        // No bubbles should be returned since limit is 0
        assert_eq!(res.len(), 0);
        Ok(())
    }

    #[fbinit::test]
    async fn reopen_expired_bubble_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        // Using MarkOnly mode since we want to mark the bubble as expired (Soft Delete)
        let (ctx, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkOnly)?;
        // Create an empty bubble that would expire immediately.
        let bubble1 = eph.create_bubble(None, vec![]).await?;
        // Mark the bubble as expired
        eph.delete_bubble(bubble1.bubble_id(), &ctx).await?;
        // Opening the expired bubble should give a
        // "No such bubble" error
        let res = eph.open_bubble(bubble1.bubble_id()).await;
        match res {
            Err(e) => match e.downcast_ref::<EphemeralBlobstoreError>() {
                Some(EphemeralBlobstoreError::NoSuchBubble(_)) => Ok(()),
                _ => Err(anyhow!("Invalid error post bubble deletion")),
            },
            _ => Err(anyhow!("Bubble expected to be deleted but it still exists")),
        }
    }

    #[fbinit::test]
    async fn reopen_expired_bubble_with_labels_test(fb: FacebookInit) -> Result<()> {
        // We want immediately expiring bubbles
        let initial = Duration::from_secs(0);
        let grace = Duration::from_secs(0);
        let (_, _, _, eph) = bootstrap(fb, initial, grace, BubbleDeletionMode::MarkAndDelete)?;
        // Create an empty bubble with labels that would expire immediately.
        let labels = vec!["workspace".to_string()];
        let bubble1 = eph.create_bubble(None, labels).await?;
        let opened_bubble = eph.open_bubble(bubble1.bubble_id()).await?;
        // Bubble should be reopened successfully since even though its expired by time
        // ,having labels associated with it should mark it as active.
        assert_eq!(opened_bubble.bubble_id(), bubble1.bubble_id());
        assert_eq!(opened_bubble.labels().await?, bubble1.labels().await?);
        Ok(())
    }
}


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

use changesets::Changesets;
use metaconfig_types::RepoConfig;
use repo_blobstore::RepoBlobstore;
use repo_identity::RepoIdentity;

#[facet::container]
pub struct EphemeralRepoView {
    #[facet]
    pub(crate) repo_blobstore: RepoBlobstore,

    #[facet]
    pub(crate) changesets: dyn Changesets,

    #[facet]
    pub(crate) repo_identity: RepoIdentity,

    #[facet]
    pub(crate) repo_config: RepoConfig,
}


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

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use arg_extensions::ArgDefaults;
use clap::ArgAction;
use clap::Args;
use metaconfig_types::PackFormat;
use rand_distr::Normal;

use crate::PutBehaviour;

/// Options for controlling the blobstore
#[derive(Args, Debug)]
pub struct BlobstoreArgs {
    /// Read QPS limit
    #[clap(long)]
    pub blobstore_read_qps: Option<NonZeroU32>,

    /// Write QPS limit
    #[clap(long)]
    pub blobstore_write_qps: Option<NonZeroU32>,

    /// Read bytes/s limit
    #[clap(long)]
    pub blobstore_read_bytes_s: Option<NonZeroUsize>,

    /// Write bytes/s limit
    #[clap(long)]
    pub blobstore_write_bytes_s: Option<NonZeroUsize>,

    /// Read burst bytes/s limit
    #[clap(long)]
    pub blobstore_read_burst_bytes_s: Option<NonZeroUsize>,

    /// Write burst bytes/s limit
    #[clap(long)]
    pub blobstore_write_burst_bytes_s: Option<NonZeroUsize>,

    /// Minimum number of bytes ThrottledBlob can count
    #[clap(long)]
    pub blobstore_bytes_min_throttle: Option<NonZeroUsize>,

    /// Rate of errors on reads.  For value N, it will error randomly
    /// 1/N times.  For multiplexed stores, this will only apply to the
    /// first store in the multiplex.
    #[clap(long)]
    pub blobstore_read_chaos_rate: Option<NonZeroU32>,

    /// Rate of errors on writes.  For value N, it will error randomly
    /// 1/N times.  For multiplexed stores, this will only apply to the
    /// first store in the multiplex.
    #[clap(long)]
    pub blobstore_write_chaos_rate: Option<NonZeroU32>,

    /// Mean value of additional delay for blobstore put calls.
    #[clap(long, requires = "blobstore_put_stddev_delay_secs")]
    pub blobstore_put_mean_delay_secs: Option<f64>,

    /// Standard devation of additional delay for blobstore put calls.
    #[clap(long, requires = "blobstore_put_mean_delay_secs")]
    pub blobstore_put_stddev_delay_secs: Option<f64>,

    /// Mean value of additional delay for blobstore get calls.
    #[clap(long, requires = "blobstore_get_stddev_delay_secs")]
    pub blobstore_get_mean_delay_secs: Option<f64>,

    /// Standard deviation of additional delay for blobstore get calls.
    #[clap(long, requires = "blobstore_get_mean_delay_secs")]
    pub blobstore_get_stddev_delay_secs: Option<f64>,

    /// Override config to enable or disable zstd compression on write
    /// via PackBlob
    // This is Option<bool> as we distinguish between option being
    // not present vs being set to false.
    #[clap(long, value_name = "BOOL")]
    pub blobstore_write_zstd: Option<bool>,

    /// Override config to set the zstd compression level for writes
    /// via PackBlob.
    #[clap(long)]
    pub blobstore_write_zstd_level: Option<i32>,

    /// Whether to attempt zstd compression when the blobstore is putting
    /// things into cachelib over the threshold size.
    // For compatibility with existing usage, this arg takes value,
    // for example `--blobstore-cachelib-attempt-zstd=true`.
    #[clap(long, default_value_t = false, value_name = "BOOL", action = ArgAction::Set)]
    pub blobstore_cachelib_attempt_zstd: bool,

    /// Desired blobstore behaviour when a put is made to an existing key.
    #[clap(long)]
    pub blobstore_put_behaviour: Option<PutBehaviour>,
}

impl BlobstoreArgs {
    pub fn put_format_override(&self) -> Result<Option<PackFormat>> {
        match (self.blobstore_write_zstd, self.blobstore_write_zstd_level) {
            (None, None) => Ok(None),
            (Some(false), None) => Ok(Some(PackFormat::Raw)),
            (Some(true), Some(level)) => Ok(Some(PackFormat::ZstdIndividual(level))),
            (Some(true), None) => Err(anyhow!(
                "Invalid arguments: --blobstore-write-zstd=true requires --blobstore-write-zstd-level"
            )),
            (_, Some(level)) => Err(anyhow!(
                "Invalid arguments: --blobstore-write-zstd-level={} requires --blobstore-write-zstd=true",
                level
            )),
        }
    }

    pub fn get_delay_distribution(&self) -> Result<Option<Normal<f64>>> {
        delay_distribution(
            self.blobstore_get_mean_delay_secs,
            self.blobstore_get_stddev_delay_secs,
        )
        .context("Failed to create blobstore get delay distribution")
    }

    pub fn put_delay_distribution(&self) -> Result<Option<Normal<f64>>> {
        delay_distribution(
            self.blobstore_put_mean_delay_secs,
            self.blobstore_put_stddev_delay_secs,
        )
        .context("Failed to create blobstore put delay distribution")
    }
}

#[derive(Default, Debug)]
pub struct BlobstoreArgDefaults {
    pub read_qps: Option<NonZeroU32>,
    pub write_qps: Option<NonZeroU32>,
    pub cachelib_attempt_zstd: Option<bool>,
    pub put_behaviour: Option<PutBehaviour>,
}

impl ArgDefaults for BlobstoreArgDefaults {
    fn arg_defaults(&self) -> Vec<(&'static str, String)> {
        let mut defaults = Vec::new();
        if let Some(read_qps) = self.read_qps {
            defaults.push(("blobstore_read_qps", read_qps.to_string()));
        }
        if let Some(write_qps) = self.write_qps {
            defaults.push(("blobstore_write_qps", write_qps.to_string()));
        }
        if let Some(cachelib_attempt_zstd) = self.cachelib_attempt_zstd {
            defaults.push((
                "blobstore_cachelib_attempt_zstd",
                cachelib_attempt_zstd.to_string(),
            ));
        }
        if let Some(put_behaviour) = self.put_behaviour {
            defaults.push((
                "blobstore_put_behaviour",
                <&'static str>::from(put_behaviour).to_string(),
            ));
        }

        defaults
    }
}

fn delay_distribution(mean: Option<f64>, stddev: Option<f64>) -> Result<Option<Normal<f64>>> {
    match (mean, stddev) {
        (Some(mean), Some(stddev)) => {
            let dist = Normal::new(mean, stddev)?;
            Ok(Some(dist))
        }
        _ => Ok(None),
    }
}


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

use std::num::NonZeroU64;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use blobstore::Blobstore;
use blobstore::BlobstoreEnumerableWithUnlink;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::DisabledBlob;
use blobstore::ErrorKind;
use blobstore::PutBehaviour;
use blobstore::DEFAULT_PUT_BEHAVIOUR;
use blobstore_sync_queue::SqlBlobstoreWal;
use cacheblob::CachelibBlobstoreOptions;
use cached_config::ConfigStore;
use chaosblob::ChaosBlobstore;
use chaosblob::ChaosOptions;
use delayblob::DelayOptions;
use delayblob::DelayedBlobstore;
use fbinit::FacebookInit;
use fileblob::Fileblob;
use futures::future;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures_watchdog::WatchdogExt;
use logblob::LogBlob;
#[cfg(fbcode_build)]
use manifoldblob::ManifoldOptions;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use metaconfig_types::MultiplexedStoreType;
use metaconfig_types::PackConfig;
use metaconfig_types::ShardableRemoteDatabaseConfig;
use metaconfig_types::ShardedDatabaseConfig;
use multiplexedblob::ScrubAction;
use multiplexedblob::ScrubHandler;
use multiplexedblob::ScrubOptions;
use multiplexedblob::SrubWriteOnly;
use multiplexedblob_wal::scrub::WalScrubBlobstore;
use multiplexedblob_wal::Scuba as WalScuba;
use multiplexedblob_wal::WalMultiplexedBlobstore;
use packblob::PackBlob;
use packblob::PackOptions;
use readonlyblob::ReadOnlyBlobstore;
use samplingblob::ComponentSamplingHandler;
use samplingblob::SamplingBlobstorePutOps;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::Logger;
use sql_construct::SqlConstructFromShardedDatabaseConfig;
use sql_ext::facebook::MysqlOptions;
use sqlblob::CountedSqlblob;
use sqlblob::Sqlblob;
use throttledblob::ThrottleOptions;
use throttledblob::ThrottledBlob;

use crate::ReadOnlyStorage;

#[derive(Clone, Debug)]
pub struct BlobstoreOptions {
    pub chaos_options: ChaosOptions,
    pub delay_options: DelayOptions,
    pub throttle_options: ThrottleOptions,
    #[cfg(fbcode_build)]
    pub manifold_options: ManifoldOptions,
    pub pack_options: PackOptions,
    pub cachelib_options: CachelibBlobstoreOptions,
    pub put_behaviour: PutBehaviour,
    pub scrub_options: Option<ScrubOptions>,
    pub sqlblob_mysql_options: MysqlOptions,
}

impl BlobstoreOptions {
    pub fn new(
        chaos_options: ChaosOptions,
        delay_options: DelayOptions,
        throttle_options: ThrottleOptions,
        #[cfg(fbcode_build)] manifold_options: ManifoldOptions,
        pack_options: PackOptions,
        cachelib_options: CachelibBlobstoreOptions,
        put_behaviour: Option<PutBehaviour>,
        sqlblob_mysql_options: MysqlOptions,
    ) -> Self {
        Self {
            chaos_options,
            delay_options,
            throttle_options,
            #[cfg(fbcode_build)]
            manifold_options,
            pack_options,
            cachelib_options,
            // If not specified, maintain status quo, which is overwrite
            put_behaviour: put_behaviour.unwrap_or(DEFAULT_PUT_BEHAVIOUR),
            // These are added via the builder methods
            scrub_options: None,
            sqlblob_mysql_options,
        }
    }

    #[cfg(fbcode_build)]
    pub fn manifold_blobstore(
        manifold_options: ManifoldOptions,
        put_behaviour: PutBehaviour,
    ) -> Self {
        Self {
            manifold_options,
            put_behaviour,
            chaos_options: Default::default(),
            delay_options: Default::default(),
            throttle_options: Default::default(),
            pack_options: Default::default(),
            cachelib_options: Default::default(),
            scrub_options: None,
            sqlblob_mysql_options: Default::default(),
        }
    }

    pub fn set_scrub_options(&mut self, scrub_options: ScrubOptions) {
        self.scrub_options = Some(scrub_options);
    }

    pub fn with_scrub_action(self, scrub_action: Option<ScrubAction>) -> Self {
        if let Some(scrub_action) = scrub_action {
            let mut scrub_options = self.scrub_options.unwrap_or_default();
            scrub_options.scrub_action = scrub_action;
            Self {
                scrub_options: Some(scrub_options),
                ..self
            }
        } else {
            self
        }
    }

    pub fn with_scrub_grace(self, scrub_grace: Option<u64>) -> Self {
        if let Some(mut scrub_options) = self.scrub_options {
            scrub_options.scrub_grace = scrub_grace.map(Duration::from_secs);
            Self {
                scrub_options: Some(scrub_options),
                ..self
            }
        } else {
            self
        }
    }

    pub fn with_scrub_action_on_missing_write_only(self, scrub_missing: SrubWriteOnly) -> Self {
        if let Some(mut scrub_options) = self.scrub_options {
            scrub_options.scrub_action_on_missing_write_only = scrub_missing;
            Self {
                scrub_options: Some(scrub_options),
                ..self
            }
        } else {
            self
        }
    }

    pub fn with_scrub_queue_peek_bound(self, queue_peek_bound_secs: u64) -> Self {
        if let Some(mut scrub_options) = self.scrub_options {
            scrub_options.queue_peek_bound = Duration::from_secs(queue_peek_bound_secs);
            Self {
                scrub_options: Some(scrub_options),
                ..self
            }
        } else {
            self
        }
    }
}

/// Construct a blobstore according to the specification. The multiplexed blobstore
/// needs an SQL DB for its queue, as does the MySQL blobstore.
/// If `throttling.read_qps` or `throttling.write_qps` are Some then ThrottledBlob will be used to limit
/// QPS to the underlying blobstore
pub fn make_blobstore<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    mysql_options: &'a MysqlOptions,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
    scrub_handler: &'a Arc<dyn ScrubHandler>,
    component_sampler: Option<&'a Arc<dyn ComponentSamplingHandler>>,
) -> BoxFuture<'a, Result<Arc<dyn Blobstore>, Error>> {
    async move {
        let store = make_blobstore_put_ops(
            fb,
            blobconfig,
            mysql_options,
            readonly_storage,
            blobstore_options,
            logger,
            config_store,
            scrub_handler,
            component_sampler,
            None,
        )
        .await?;
        // Workaround for trait A {} trait B:A {} but Arc<dyn B> is not a Arc<dyn A>
        // See https://github.com/rust-lang/rfcs/issues/2765 if interested
        Ok(Arc::new(store) as Arc<dyn Blobstore>)
    }
    .boxed()
}

pub async fn make_sql_blobstore<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    config_store: &'a ConfigStore,
) -> Result<CountedSqlblob, Error> {
    use BlobConfig::*;
    match blobconfig {
        Sqlite { path } => Sqlblob::with_sqlite_path(
            path.join("blobs"),
            readonly_storage.0,
            blobstore_options.put_behaviour,
            config_store,
        )
        .context(ErrorKind::StateOpen),
        Mysql { remote } => {
            let (tier_name, shard_count) = match remote {
                ShardableRemoteDatabaseConfig::Unsharded(config) => (config.db_address, None),
                ShardableRemoteDatabaseConfig::Sharded(config) => {
                    (config.shard_map.clone(), Some(config.shard_num))
                }
            };
            make_sql_blobstore_xdb(
                fb,
                tier_name,
                shard_count,
                blobstore_options,
                readonly_storage,
                blobstore_options.put_behaviour,
                config_store,
            )
            .await
        }
        _ => bail!("Not an SQL blobstore"),
    }
}

// Most users should call `make_sql_blobstore` instead, however its useful to expose this to reduce duplication with benchmark tools.
pub async fn make_sql_blobstore_xdb<'a>(
    fb: FacebookInit,
    tier_name: String,
    shard_count: Option<NonZeroUsize>,
    blobstore_options: &'a BlobstoreOptions,
    readonly_storage: ReadOnlyStorage,
    put_behaviour: PutBehaviour,
    config_store: &'a ConfigStore,
) -> Result<CountedSqlblob, Error> {
    let mysql_options = blobstore_options.sqlblob_mysql_options.clone();
    match shard_count {
        None => {
            Sqlblob::with_mysql_unsharded(
                fb,
                tier_name,
                mysql_options,
                readonly_storage.0,
                put_behaviour,
                config_store,
            )
            .await
        }
        Some(shard_num) => {
            Sqlblob::with_mysql(
                fb,
                tier_name,
                shard_num,
                mysql_options,
                readonly_storage.0,
                put_behaviour,
                config_store,
            )
            .await
        }
    }
}

pub fn make_packblob_wrapper<'a, T>(
    pack_config: Option<PackConfig>,
    blobstore_options: &'a BlobstoreOptions,
    store: T,
) -> Result<PackBlob<T>, Error> {
    // Take the user specified option if provided, otherwise use the config
    let put_format = if let Some(put_format) = blobstore_options.pack_options.override_put_format {
        put_format
    } else {
        pack_config.map(|c| c.put_format).unwrap_or_default()
    };

    Ok(PackBlob::new(store, put_format))
}

/// Construct a PackBlob according to the spec; you are responsible for
/// finding a PackBlob config
pub async fn make_packblob<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
) -> Result<PackBlob<Arc<dyn BlobstoreUnlinkOps>>, Error> {
    if let BlobConfig::Pack {
        pack_config,
        blobconfig,
    } = blobconfig
    {
        let store = make_blobstore_with_link(
            fb,
            *blobconfig,
            readonly_storage,
            blobstore_options,
            logger,
            config_store,
        )
        .watched(logger)
        .await?;

        Ok(make_packblob_wrapper(
            pack_config,
            blobstore_options,
            store,
        )?)
    } else {
        bail!("Not a PackBlob")
    }
}

#[cfg(fbcode_build)]
pub async fn make_manifold_blobstore(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    blobstore_options: &BlobstoreOptions,
) -> Result<Arc<dyn BlobstoreEnumerableWithUnlink>, Error> {
    use BlobConfig::*;
    let (bucket, prefix, ttl) = match blobconfig {
        Manifold { bucket, prefix } => (bucket, prefix, None),
        ManifoldWithTtl {
            bucket,
            prefix,
            ttl,
        } => (bucket, prefix, Some(ttl)),
        _ => bail!("Not a Manifold blobstore"),
    };
    crate::facebook::make_manifold_blobstore(
        fb,
        &prefix,
        &bucket,
        ttl,
        &blobstore_options.manifold_options,
        blobstore_options.put_behaviour,
    )
}

#[cfg(not(fbcode_build))]
pub async fn make_manifold_blobstore(
    _fb: FacebookInit,
    _blobconfig: BlobConfig,
    _blobstore_options: &BlobstoreOptions,
) -> Result<Arc<dyn BlobstoreEnumerableWithUnlink>, Error> {
    unimplemented!("This is implemented only for fbcode_build")
}

pub async fn make_files_blobstore(
    blobconfig: BlobConfig,
    blobstore_options: &BlobstoreOptions,
) -> Result<Fileblob, Error> {
    if let BlobConfig::Files { path } = blobconfig {
        Fileblob::create(path.join("blobs"), blobstore_options.put_behaviour)
            .context(ErrorKind::StateOpen)
    } else {
        bail!("Not a file blobstore")
    }
}

async fn make_blobstore_with_link<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
) -> Result<Arc<dyn BlobstoreUnlinkOps>, Error> {
    use BlobConfig::*;
    match blobconfig {
        Sqlite { .. } | Mysql { .. } => make_sql_blobstore(
            fb,
            blobconfig,
            readonly_storage,
            blobstore_options,
            config_store,
        )
        .watched(logger)
        .await
        .map(|store| Arc::new(store) as Arc<dyn BlobstoreUnlinkOps>),
        Manifold { .. } | ManifoldWithTtl { .. } => {
            make_manifold_blobstore(fb, blobconfig, blobstore_options)
                .watched(logger)
                .await
                .map(|store| Arc::new(store) as Arc<dyn BlobstoreUnlinkOps>)
        }
        Files { .. } => make_files_blobstore(blobconfig, blobstore_options)
            .await
            .map(|store| Arc::new(store) as Arc<dyn BlobstoreUnlinkOps>),
        _ => bail!("Not a physical blobstore"),
    }
}

// Constructs the BlobstoreEnumerableWithUnlink store implementations
// for blobstores. If the blobstore is a wrapper blobstore, the inner
// physical blobstore construction is delegated to another function
// and the result is wrapped up in this function.
pub async fn make_blobstore_enumerable_with_unlink<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
) -> Result<Arc<dyn BlobstoreEnumerableWithUnlink>, Error> {
    use BlobConfig::*;
    match blobconfig {
        Pack {
            pack_config,
            blobconfig,
        } => {
            let store =
                raw_blobstore_enumerable_with_unlink(fb, *blobconfig, blobstore_options, logger)
                    .watched(logger)
                    .await?;
            let pack_store = make_packblob_wrapper(pack_config, blobstore_options, store)?;
            Ok(Arc::new(pack_store) as Arc<dyn BlobstoreEnumerableWithUnlink>)
        }
        _ => raw_blobstore_enumerable_with_unlink(fb, blobconfig, blobstore_options, logger).await,
    }
}

// Constructs the raw BlobstoreEnumerableWithUnlink store implementations for low level
// blobstore access. The blobstore created is NOT a wrapper (e.g. PackBlob)
pub async fn raw_blobstore_enumerable_with_unlink<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
) -> Result<Arc<dyn BlobstoreEnumerableWithUnlink>, Error> {
    use BlobConfig::*;
    match blobconfig {
        Manifold { .. } | ManifoldWithTtl { .. } => {
            make_manifold_blobstore(fb, blobconfig, blobstore_options)
                .watched(logger)
                .await
        }
        Files { .. } => make_files_blobstore(blobconfig, blobstore_options)
            .await
            .map(|store| Arc::new(store) as Arc<dyn BlobstoreEnumerableWithUnlink>),
        _ => bail!("Not a physical blobstore that supports unlink + keysource + putops"),
    }
}

// Constructs the BlobstorePutOps store implementations for low level blobstore access
fn make_blobstore_put_ops<'a>(
    fb: FacebookInit,
    blobconfig: BlobConfig,
    mysql_options: &'a MysqlOptions,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
    scrub_handler: &'a Arc<dyn ScrubHandler>,
    component_sampler: Option<&'a Arc<dyn ComponentSamplingHandler>>,
    blobstore_id: Option<BlobstoreId>,
) -> BoxFuture<'a, Result<Arc<dyn BlobstorePutOps>, Error>> {
    // NOTE: This needs to return a BoxFuture because it recurses.
    async move {
        use BlobConfig::*;

        let mut needs_wrappers = true;
        let store = match blobconfig {
            // Physical blobstores
            Sqlite { .. } | Mysql { .. } => make_sql_blobstore(
                fb,
                blobconfig,
                readonly_storage,
                blobstore_options,
                config_store,
            )
            .watched(logger)
            .await
            .map(|store| Arc::new(store) as Arc<dyn BlobstorePutOps>)?,
            Manifold { .. } | ManifoldWithTtl { .. } => {
                make_manifold_blobstore(fb, blobconfig, blobstore_options)
                    .watched(logger)
                    .await
                    .map(|store| Arc::new(store) as Arc<dyn BlobstorePutOps>)?
            }
            Files { .. } => make_files_blobstore(blobconfig, blobstore_options)
                .await
                .map(|store| Arc::new(store) as Arc<dyn BlobstorePutOps>)?,
            S3 {
                bucket,
                keychain_group,
                region_name,
                endpoint,
                num_concurrent_operations,
                secret_name,
            } => {
                #[cfg(fbcode_build)]
                {
                    ::s3blob::S3Blob::new(
                        fb,
                        bucket,
                        keychain_group,
                        secret_name,
                        region_name,
                        endpoint,
                        blobstore_options.put_behaviour,
                        logger,
                        num_concurrent_operations,
                    )
                    .watched(logger)
                    .await
                    .context(ErrorKind::StateOpen)
                    .map(|store| Arc::new(store) as Arc<dyn BlobstorePutOps>)?
                }
                #[cfg(not(fbcode_build))]
                {
                    let _ = (
                        bucket,
                        keychain_group,
                        secret_name,
                        region_name,
                        endpoint,
                        num_concurrent_operations,
                    );
                    unimplemented!("This is implemented only for fbcode_build")
                }
            }

            // Special case
            Disabled => {
                Arc::new(DisabledBlob::new("Disabled by configuration")) as Arc<dyn BlobstorePutOps>
            }

            // Wrapper blobstores
            MultiplexedWal {
                multiplex_id,
                blobstores,
                write_quorum,
                queue_db,
                inner_blobstores_scuba_table,
                multiplex_scuba_table,
                scuba_sample_rate,
            } => {
                needs_wrappers = false;
                make_multiplexed_wal(
                    fb,
                    multiplex_id,
                    queue_db,
                    inner_blobstores_scuba_table,
                    multiplex_scuba_table,
                    scuba_sample_rate,
                    blobstores,
                    write_quorum,
                    mysql_options,
                    readonly_storage,
                    blobstore_options,
                    logger,
                    config_store,
                    scrub_handler,
                    component_sampler,
                )
                .watched(logger)
                .await?
            }
            Logging {
                blobconfig,
                scuba_table,
                scuba_sample_rate,
            } => {
                needs_wrappers = false;
                let store = make_blobstore_put_ops(
                    fb,
                    *blobconfig,
                    mysql_options,
                    readonly_storage,
                    blobstore_options,
                    logger,
                    config_store,
                    scrub_handler,
                    component_sampler,
                    None,
                )
                .watched(logger)
                .await?;

                let scuba = scuba_table
                    .map_or(Ok(MononokeScubaSampleBuilder::with_discard()), |table| {
                        MononokeScubaSampleBuilder::new(fb, &table)
                    })?;
                Arc::new(LogBlob::new(store, scuba, scuba_sample_rate)) as Arc<dyn BlobstorePutOps>
            }
            Pack { .. } => {
                // NB packblob does not apply the wrappers internally
                make_packblob(
                    fb,
                    blobconfig,
                    readonly_storage,
                    blobstore_options,
                    logger,
                    config_store,
                )
                .watched(logger)
                .await
                .map(|store| Arc::new(store) as Arc<dyn BlobstorePutOps>)?
            }
        };

        let store = if needs_wrappers {
            let store = if let Some(component_sampler) = component_sampler {
                Arc::new(SamplingBlobstorePutOps::new(
                    store,
                    blobstore_id,
                    component_sampler.clone(),
                )) as Arc<dyn BlobstorePutOps>
            } else {
                store
            };

            let store = if readonly_storage.0 {
                Arc::new(ReadOnlyBlobstore::new(store)) as Arc<dyn BlobstorePutOps>
            } else {
                store
            };

            let store = if blobstore_options.throttle_options.has_throttle() {
                Arc::new(
                    ThrottledBlob::new(store, blobstore_options.throttle_options)
                        .watched(logger)
                        .await,
                ) as Arc<dyn BlobstorePutOps>
            } else {
                store
            };

            let store = if blobstore_options.chaos_options.has_chaos() {
                Arc::new(ChaosBlobstore::new(store, blobstore_options.chaos_options))
                    as Arc<dyn BlobstorePutOps>
            } else {
                store
            };

            if blobstore_options.delay_options.has_delay() {
                Arc::new(DelayedBlobstore::from_options(
                    store,
                    blobstore_options.delay_options,
                )) as Arc<dyn BlobstorePutOps>
            } else {
                store
            }
        } else {
            // Already applied the wrappers inside the store
            store
        };

        // NOTE: Do not add wrappers here that should only be added once per repository, since this
        // function will get called recursively for each member of a Multiplex! For those, use
        // RepoBlobstore::new instead.

        Ok(store)
    }
    .boxed()
}

async fn make_multiplexed_wal<'a>(
    fb: FacebookInit,
    multiplex_id: MultiplexId,
    queue_db: ShardedDatabaseConfig,
    inner_blobstores_scuba_table: Option<String>,
    multiplex_scuba_table: Option<String>,
    scuba_sample_rate: NonZeroU64,
    inner_config: Vec<(BlobstoreId, MultiplexedStoreType, BlobConfig)>,
    write_quorum: usize,
    mysql_options: &'a MysqlOptions,
    readonly_storage: ReadOnlyStorage,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
    scrub_handler: &'a Arc<dyn ScrubHandler>,
    component_sampler: Option<&'a Arc<dyn ComponentSamplingHandler>>,
) -> Result<Arc<dyn BlobstorePutOps>, Error> {
    let (normal_components, write_only_components) = setup_inner_blobstores(
        fb,
        inner_config,
        mysql_options,
        blobstore_options,
        logger,
        config_store,
        scrub_handler,
        component_sampler,
    )
    .await?;

    let wal_queue = Arc::new(SqlBlobstoreWal::with_sharded_database_config(
        fb,
        &queue_db,
        mysql_options,
        readonly_storage.0,
    )?);
    let scuba = WalScuba::new_from_raw(
        fb,
        inner_blobstores_scuba_table,
        multiplex_scuba_table,
        scuba_sample_rate,
    )?;

    let blobstore = match &blobstore_options.scrub_options {
        Some(scrub_options) => {
            Arc::new(WalScrubBlobstore::new(
                multiplex_id,
                wal_queue,
                normal_components,
                write_only_components,
                write_quorum,
                None, // use default timeouts
                scuba,
                scrub_options.clone(),
                scrub_handler.clone(),
            )?) as Arc<dyn BlobstorePutOps>
        }
        None => Arc::new(WalMultiplexedBlobstore::new(
            multiplex_id,
            wal_queue,
            normal_components,
            write_only_components,
            write_quorum,
            None, // use default timeouts
            scuba,
        )?) as Arc<dyn BlobstorePutOps>,
    };

    Ok(blobstore)
}

type InnerBlobstore = (BlobstoreId, Arc<dyn BlobstorePutOps>);

async fn setup_inner_blobstores<'a>(
    fb: FacebookInit,
    inner_config: Vec<(BlobstoreId, MultiplexedStoreType, BlobConfig)>,
    mysql_options: &'a MysqlOptions,
    blobstore_options: &'a BlobstoreOptions,
    logger: &'a Logger,
    config_store: &'a ConfigStore,
    scrub_handler: &'a Arc<dyn ScrubHandler>,
    component_sampler: Option<&'a Arc<dyn ComponentSamplingHandler>>,
) -> Result<(Vec<InnerBlobstore>, Vec<InnerBlobstore>), Error> {
    let component_readonly = blobstore_options
        .scrub_options
        .as_ref()
        .map_or(ReadOnlyStorage(false), |v| {
            ReadOnlyStorage(v.scrub_action != ScrubAction::Repair)
        });

    let mut applied_chaos = false;
    let components = future::try_join_all(inner_config.into_iter().map({
        move |(blobstoreid, store_type, config)| {
            let mut blobstore_options = blobstore_options.clone();

            if blobstore_options.chaos_options.has_chaos() {
                if applied_chaos {
                    blobstore_options = BlobstoreOptions {
                        chaos_options: ChaosOptions::new(None, None),
                        ..blobstore_options
                    };
                } else {
                    applied_chaos = true;
                }
            }

            async move {
                let store = make_blobstore_put_ops(
                    fb,
                    config,
                    mysql_options,
                    component_readonly,
                    &blobstore_options,
                    logger,
                    config_store,
                    scrub_handler,
                    component_sampler,
                    Some(blobstoreid),
                )
                .watched(logger)
                .await?;

                Result::<_, Error>::Ok((blobstoreid, store_type, store))
            }
        }
    }))
    .await?;

    // For now, `partition` could do this, but this will be easier to extend when we introduce more store types
    let mut normal_components = vec![];
    let mut write_only_components = vec![];
    for (blobstore_id, store_type, store) in components.into_iter() {
        match store_type {
            MultiplexedStoreType::Normal => normal_components.push((blobstore_id, store)),
            MultiplexedStoreType::WriteOnly => write_only_components.push((blobstore_id, store)),
        }
    }
    Ok((normal_components, write_only_components))
}


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

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

use arg_extensions::ArgDefaults;
use clap::ArgAction;
use clap::Args;

mod args;
mod blobstore;
#[cfg(fbcode_build)]
mod facebook;
mod sql;

pub use ::blobstore::PutBehaviour;
pub use ::blobstore::DEFAULT_PUT_BEHAVIOUR;
pub use blobstore_stats::OperationType;
pub use cacheblob::CachelibBlobstoreOptions;
pub use chaosblob::ChaosOptions;
pub use delayblob::DelayOptions;
#[cfg(fbcode_build)]
pub use facebook::ManifoldArgs;
#[cfg(fbcode_build)]
pub use manifoldblob::ManifoldOptions;
pub use multiplexedblob::scrub::default_scrub_handler;
pub use multiplexedblob::scrub::ScrubOptions;
pub use multiplexedblob::scrub::SrubWriteOnly;
pub use multiplexedblob::ScrubAction;
pub use multiplexedblob::ScrubHandler;
pub use packblob::PackOptions;
pub use samplingblob::ComponentSamplingHandler;
pub use throttledblob::ThrottleOptions;

pub use crate::args::BlobstoreArgDefaults;
pub use crate::args::BlobstoreArgs;
pub use crate::blobstore::make_blobstore;
pub use crate::blobstore::make_blobstore_enumerable_with_unlink;
pub use crate::blobstore::make_files_blobstore;
pub use crate::blobstore::make_manifold_blobstore;
pub use crate::blobstore::make_packblob;
pub use crate::blobstore::make_sql_blobstore;
pub use crate::blobstore::make_sql_blobstore_xdb;
pub use crate::blobstore::BlobstoreOptions;
pub use crate::sql::MetadataSqlFactory;
pub use crate::sql::SqlTierInfo;

#[derive(Copy, Clone, PartialEq)]
pub struct ReadOnlyStorage(pub bool);

impl ArgDefaults for ReadOnlyStorage {
    fn arg_defaults(&self) -> Vec<(&'static str, String)> {
        vec![("with_readonly_storage", self.0.to_string())]
    }
}

impl ReadOnlyStorage {
    pub fn from_args(args: &ReadOnlyStorageArgs) -> Self {
        ReadOnlyStorage(args.with_readonly_storage)
    }
}

/// Command line arguments for controlling read-only storage
#[derive(Args, Debug)]
pub struct ReadOnlyStorageArgs {
    /// Error on any attempts to write to storage if set to true
    // For compatibility with existing usage, allows usage as
    // `--with-readonly-storage=true`.
    #[clap(
        long,
        value_name = "BOOL",
        default_value = "false",
        default_missing_value = "true",
        action = ArgAction::Set,
        num_args = 0..=1,
    )]
    pub with_readonly_storage: bool,
}