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

use std::fmt::Debug;
use std::future::Future;
use std::panic;
use std::sync::Arc;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore_sync_queue::BlobstoreWal;
use blobstore_sync_queue::BlobstoreWalEntry;
use blobstore_sync_queue::SqlBlobstoreWal;
use blobstore_test_utils::Tickable;
use borrowed::borrowed;
use bytes::Bytes;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::future::FutureExt;
use futures::task::Poll;
use lock_ext::LockExt;
use metaconfig_types::BlobstoreId;
use metaconfig_types::MultiplexId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::Timestamp;
use multiplexedblob::LoggingScrubHandler;
use multiplexedblob::ScrubAction;
use multiplexedblob::ScrubHandler;
use multiplexedblob::ScrubOptions;
use multiplexedblob::SrubWriteOnly;
use nonzero_ext::nonzero;
use scuba_ext::MononokeScubaSampleBuilder;
use sql_construct::SqlConstruct;

use crate::scrub::WalScrubBlobstore;
use crate::MultiplexTimeout;
use crate::Scuba;
use crate::WalMultiplexedBlobstore;

#[fbinit::test]
async fn test_quorum_is_valid(_fb: FacebookInit) -> Result<()> {
    let scuba = Scuba::new(
        MononokeScubaSampleBuilder::with_discard(),
        MononokeScubaSampleBuilder::with_discard(),
        nonzero!(1u64),
    )?;
    let wal = Arc::new(SqlBlobstoreWal::with_sqlite_in_memory()?);

    // Check the quorum cannot be zero
    {
        // no main-stores, no write-only
        assert!(setup_multiplex(0, 0, None).is_err());
    }

    // Check creating multiplex fails if there are no enough main blobstores
    {
        let stores = (0..2)
            .map(|id| {
                (
                    BlobstoreId::new(id),
                    Arc::new(Tickable::new()) as Arc<dyn BlobstorePutOps>,
                )
            })
            .collect();
        // write-only don't count into the quorum
        let write_only = (2..4)
            .map(|id| {
                (
                    BlobstoreId::new(id),
                    Arc::new(Tickable::new()) as Arc<dyn BlobstorePutOps>,
                )
            })
            .collect();
        let quorum = 3;
        let result = WalMultiplexedBlobstore::new(
            MultiplexId::new(0),
            wal,
            stores,
            write_only,
            quorum,
            None,
            scuba,
        );

        assert!(result.is_err());
    }

    // Check creating multiplex succeeds with the same amount of stores as the quorum
    {
        assert!(setup_multiplex(3, 3, None).is_ok());
    }

    Ok(())
}

#[fbinit::test]
async fn test_put_wal_fails(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    let v = make_value("v");
    let k = "k";

    let mut put_fut = multiplex.put(&ctx, k.to_owned(), v).boxed();
    assert_pending(&mut put_fut).await;

    // wal queue write fails
    tickable_queue.tick(Some("wal queue failed"));

    // multiplex put should fail
    assert!(put_fut.await.is_err());

    // check there is no blob in the storage
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        // blobstore gets succeed
        for (_id, store) in tickable_blobstores.iter() {
            store.tick(None);
        }
        validate_blob(get_fut.await, Ok(None));
    }

    Ok(())
}

#[fbinit::test]
async fn test_put_fails(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    // All puts fail, the multiplex put should fail: [x] [x] [x]
    {
        let v = make_value("v0");
        let k = "k0";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        // all blobstores should fail
        for (id, store) in tickable_blobstores.iter() {
            store.tick(Some(format!("all fail: bs{} failed", id).as_str()));
        }

        assert!(put_fut.await.is_err());

        // No `put` succeeded, there is no blob in the storage
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // blobstore gets succeed
            for (_id, store) in tickable_blobstores.iter() {
                store.tick(None);
            }
            validate_blob(get_fut.await, Ok(None));
        }
    }

    // Second blobstore put succeeded but no quorum was achieved, multiplex put fails: [x] [ ] [x]
    {
        let v = make_value("v1");
        let k = "k1";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        // first blobstore fails
        tickable_blobstores[0].1.tick(Some("all fail: bs0 failed"));
        assert_pending(&mut put_fut).await;
        // second blobstore succeeds
        tickable_blobstores[1].1.tick(None);
        assert_pending(&mut put_fut).await;
        // third blobstore fails
        tickable_blobstores[2].1.tick(Some("all fail: bs2 failed"));

        // the multiplex put fails
        assert!(put_fut.await.is_err());

        // The blob was written to the second blobstore even if the multiplex put failed.
        // That means `get` result is undefined: it can either return `None` or `Some`
        // depending on which blobstores' put returned first.
        //
        // This is fine because if multiplex `put` failed, multiplex blobstore doesn't provide
        // any guarantees whether the blob is present in the storage or not.
        // There is only guarantee: if `put` succeeded, the blob will always be present,
        // i.e. `get` will always return `Some` (if it didn't fail for some other reason).

        // 1st and 3rd blobstore returned `None`, before 2nd blobstore returned anything
        // there is a read quorum on `None`, so multiplex `get` should also return `None`
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // first and third blobstores don't have the blob
            tickable_blobstores[0].1.tick(None);
            tickable_blobstores[2].1.tick(None);

            // the result is ready
            validate_blob(get_fut.await, Ok(None));
            tickable_blobstores[1].1.drain(1);
        }

        // 2nd blobstore returned before the read quorum on `None` was achieved, so
        // multiplex `get` should also return `Some`
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // first blobstore doesn't have the blob
            tickable_blobstores[0].1.tick(None);
            // second blobstore has
            tickable_blobstores[1].1.tick(None);

            // the result is ready
            validate_blob(get_fut.await, Ok(Some(&v)));
            tickable_blobstores[2].1.drain(1);
        }
    }

    Ok(())
}

async fn queue_keys(ctx: &CoreContext, multiplex: &WalMultiplexedBlobstore) -> Result<Vec<String>> {
    let mut entries: Vec<_> = multiplex
        .wal_queue
        .read(ctx, &multiplex.multiplex_id, &Timestamp::now(), 100)
        .await?
        .into_iter()
        .map(|e| e.blobstore_key)
        .collect();
    entries.sort_unstable();
    Ok(entries)
}

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

    let quorum = 2;
    let (tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, quorum, None)?;

    // Quorum puts succeed, the multiplex put succeeds: [ ] [x] [ ]
    // Should wait for the third put to complete.
    {
        let v = make_value("v2");
        let k = "k2";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;
        assert_eq!(&queue_keys(&ctx, &multiplex).await?, &["k2"]);

        // first blobstore succeeds
        tickable_blobstores[0].1.tick(None);
        assert_pending(&mut put_fut).await;
        // second blobstore fails
        tickable_blobstores[1].1.tick(Some("all fail: bs1 failed"));
        assert_pending(&mut put_fut).await;
        // third blobstore succeeds
        tickable_blobstores[2].1.tick(None);

        assert!(put_fut.await.is_ok());

        // check we can read the blob from the 1st store
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // first blobstore returns the blob
            tickable_blobstores[0].1.tick(None);
            validate_blob(get_fut.await, Ok(Some(&v)));

            // drain the tickables of the pending requests, as they won't be claimed
            for (_id, store) in &tickable_blobstores[1..3] {
                store.drain(1);
            }
        }

        // check we can read the blob from the 3rd store
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // first blobstore fails
            tickable_blobstores[0].1.tick(Some("bs0 failed"));
            assert_pending(&mut get_fut).await;

            // second blobstore doesn't have the blob
            tickable_blobstores[1].1.tick(None);
            assert_pending(&mut get_fut).await;

            // third blobstore succeeds
            tickable_blobstores[2].1.tick(None);
            validate_blob(get_fut.await, Ok(Some(&v)));
        }
    }

    // All puts succeed, the multiplex put succeeds: [ ] [ ] [ ]
    // Should not wait for the third put to complete.
    {
        let v = make_value("v3");
        let k = "k3";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        // first quorum blobstore puts succeed, multiplex doesn't wait for the rest
        // of the puts
        for (_id, store) in &tickable_blobstores[0..quorum] {
            store.tick(None);
        }

        assert!(put_fut.await.is_ok());

        // check we can read the blob
        {
            let mut get_fut = multiplex.get(&ctx, k).boxed();
            assert_pending(&mut get_fut).await;

            // blobstore gets succeed
            for (_id, store) in tickable_blobstores.iter() {
                store.tick(None);
            }
            validate_blob(get_fut.await, Ok(Some(&v)));
        }
        // check the optimisation to delete the item from the queue
        {
            tokio::task::yield_now().await;
            // 2 item in the queue (k2 is from the previous test)
            assert_eq!(&queue_keys(&ctx, &multiplex).await?, &["k2", "k3"]);
            // Tick the rest of the blobstores
            for (_id, store) in &tickable_blobstores[quorum..] {
                store.tick(None);
            }
            tokio::task::yield_now().await;
            // Tick deletion
            tickable_queue.tick(None);
            tokio::task::yield_now().await;
            // k2 correctly doesn't leave the key because we didn't succeed all writes
            assert_eq!(&queue_keys(&ctx, &multiplex).await?, &["k2"]);
        }
    }

    Ok(())
}

#[fbinit::test]
async fn test_get_on_missing(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (_tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    // No blobstores have the key
    let k = "k1";

    // all gets succeed, but multiplexed returns `None`
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        tickable_blobstores[0].1.tick(None);
        assert_pending(&mut get_fut).await;
        tickable_blobstores[1].1.tick(None);

        // the read-quorum on `None` achieved, multiplexed get returns `None`
        validate_blob(get_fut.await, Ok(None));
        tickable_blobstores[2].1.drain(1);
    }

    // two gets succeed, but multiplexed returns `None`
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        tickable_blobstores[0].1.tick(None);
        tickable_blobstores[1].1.tick(Some("bs1 failed"));
        // muliplexed get waits on the third
        assert_pending(&mut get_fut).await;
        tickable_blobstores[2].1.tick(None);

        // the read-quorum on `None` achieved, multiplexed get returns `None`
        validate_blob(get_fut.await, Ok(None));
    }

    // two gets fail, multiplexed get fails, because no read quorum
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        tickable_blobstores[0].1.tick(Some("bs0 failed"));
        tickable_blobstores[1].1.tick(Some("bs1 failed"));
        // muliplexed get waits on the third, which returns `None`
        assert_pending(&mut get_fut).await;
        tickable_blobstores[2].1.tick(None);

        // no read-quorum, multiplexed get fails
        validate_blob(get_fut.await, Err(anyhow::anyhow!("error")));
    }

    Ok(())
}

#[fbinit::test]
async fn test_get_on_existing(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    // Two blobstores have the key, one failed to write: [ ] [x] [ ]

    let v = make_value("v1");
    let k = "k1";

    let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
    assert_pending(&mut put_fut).await;

    // wal queue write succeeds
    tickable_queue.tick(None);
    assert_pending(&mut put_fut).await;

    tickable_blobstores[0].1.tick(None);
    tickable_blobstores[1].1.tick(Some("bs1 failed"));
    tickable_blobstores[2].1.tick(None);

    // multiplexed put succeeds: write quorum achieved
    assert!(put_fut.await.is_ok());

    // all gets succeed, but multiplexed returns on the first successful `Some`
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        // first blobstore returns the blob
        tickable_blobstores[0].1.tick(None);
        validate_blob(get_fut.await, Ok(Some(&v)));

        // drain the tickables of the pending requests, as they won't be claimed
        for (_id, store) in &tickable_blobstores[1..3] {
            store.drain(1);
        }
    }

    // first get fails, but multiplexed returns on the third get
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        // first blobstore get fails
        tickable_blobstores[0].1.tick(Some("bs1 failed!"));
        assert_pending(&mut get_fut).await;

        // second blobstore get returns None
        tickable_blobstores[1].1.tick(None);
        assert_pending(&mut get_fut).await;

        // third blobstore get returns Some
        tickable_blobstores[2].1.tick(None);
        validate_blob(get_fut.await, Ok(Some(&v)));
    }

    // 2 first gets fail, but multiplexed returns `Some` on the third get
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        // first blobstore get fails
        tickable_blobstores[0].1.tick(Some("bs1 failed!"));
        assert_pending(&mut get_fut).await;

        // second blobstore get fails
        tickable_blobstores[1].1.tick(Some("bs2 failed!"));
        assert_pending(&mut get_fut).await;

        // third blobstore get returns Some
        tickable_blobstores[2].1.tick(None);
        validate_blob(get_fut.await, Ok(Some(&v)));
    }

    // all blobstores that have the blob fail, multiplexed get fail:
    // no read quorum on `None` was achieved
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        // first and third blobstore gets fail
        tickable_blobstores[0].1.tick(Some("bs1 failed!"));
        tickable_blobstores[2].1.tick(Some("bs3 failed!"));
        assert_pending(&mut get_fut).await;

        // second blobstore get returns None
        tickable_blobstores[1].1.tick(None);
        validate_blob(get_fut.await, Err(anyhow::anyhow!("error")));
    }

    // all blobstores gets fail, multiplexed get fail
    {
        let mut get_fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut get_fut).await;

        for (id, store) in tickable_blobstores {
            store.tick(Some(format!("bs{} failed!", id).as_str()));
        }

        validate_blob(get_fut.await, Err(anyhow::anyhow!("error")));
    }

    Ok(())
}

#[fbinit::test]
async fn test_is_present_missing(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (_tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    // No blobstores have the key
    let k = "k1";

    // all `is_present` succeed, multiplexed returns `Absent`
    {
        let mut fut = multiplex.is_present(&ctx, k).boxed();
        assert_pending(&mut fut).await;

        tickable_blobstores[0].1.tick(None);
        assert_pending(&mut fut).await;
        tickable_blobstores[1].1.tick(None);

        // the read-quorum on `None` achieved, multiplexed returns `Absent`
        assert_is_present_ok(fut.await, BlobstoreIsPresent::Absent);
        tickable_blobstores[2].1.drain(1);
    }

    // two `is_present`s succeed, multiplexed returns `Absent`
    {
        let mut fut = multiplex.is_present(&ctx, k).boxed();
        assert_pending(&mut fut).await;

        tickable_blobstores[0].1.tick(None);
        tickable_blobstores[1].1.tick(Some("bs1 failed"));
        // muliplexed is_present waits on the third
        assert_pending(&mut fut).await;
        tickable_blobstores[2].1.tick(None);

        // the read-quorum achieved, multiplexed returns `Absent`
        assert_is_present_ok(fut.await, BlobstoreIsPresent::Absent);
    }

    // two `is_present`s fail, multiplexed returns `ProbablyNotPresent`
    {
        let mut fut = multiplex.is_present(&ctx, k).boxed();
        assert_pending(&mut fut).await;

        tickable_blobstores[0].1.tick(Some("bs0 failed"));
        tickable_blobstores[1].1.tick(None);
        // muliplexed is_present waits on the third
        assert_pending(&mut fut).await;
        tickable_blobstores[2].1.tick(Some("bs2 failed"));

        assert_is_present_ok(
            fut.await,
            BlobstoreIsPresent::ProbablyNotPresent(anyhow!("some failed!")),
        );
    }

    // all `is_present`s fail, multiplexed fails
    {
        let mut fut = multiplex.is_present(&ctx, k).boxed();
        assert_pending(&mut fut).await;

        for (id, store) in tickable_blobstores {
            store.tick(Some(format!("bs{} failed!", id).as_str()));
        }
        assert!(fut.await.is_err());
    }

    Ok(())
}

#[fbinit::test]
async fn test_is_present_existing(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let (tickable_queue, tickable_blobstores, multiplex) = setup_multiplex(3, 2, None)?;

    // Two blobstores have the key, one failed to write: [ ] [x] [ ]
    {
        let v = make_value("v1");
        let k = "k1";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        tickable_blobstores[0].1.tick(None);
        tickable_blobstores[1].1.tick(Some("bs1 failed"));
        tickable_blobstores[2].1.tick(None);

        // multiplexed put succeeds: write quorum achieved
        assert!(put_fut.await.is_ok());

        // first `is_present` succeed with `Present`, multiplexed returns `Present`
        {
            let mut fut = multiplex.is_present(&ctx, k).boxed();
            assert_pending(&mut fut).await;

            tickable_blobstores[0].1.tick(None);
            assert_is_present_ok(fut.await, BlobstoreIsPresent::Present);

            for (_id, store) in &tickable_blobstores[1..] {
                store.drain(1);
            }
        }

        // first `is_present` fails, second succeed with `Absent`, third returns `Present`
        // multiplexed returns `Present`
        {
            let mut fut = multiplex.is_present(&ctx, k).boxed();
            assert_pending(&mut fut).await;

            tickable_blobstores[0].1.tick(Some("bs0 failed"));
            tickable_blobstores[1].1.tick(None);
            assert_pending(&mut fut).await;

            tickable_blobstores[2].1.tick(None);
            assert_is_present_ok(fut.await, BlobstoreIsPresent::Present);
        }
    }

    // Two blobstores failed to write, one succeeded: [x] [ ] [x]
    {
        let v = make_value("v2");
        let k = "k2";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        tickable_blobstores[0].1.tick(Some("bs0 failed"));
        tickable_blobstores[1].1.tick(None);
        tickable_blobstores[2].1.tick(Some("bs2 failed"));

        // multiplexed put failed: no write quorum
        assert!(put_fut.await.is_err());

        // the first `is_present` to succeed returns `Present`, multiplexed returns `Present`
        {
            let mut fut = multiplex.is_present(&ctx, k).boxed();
            assert_pending(&mut fut).await;

            tickable_blobstores[1].1.tick(None);
            assert_is_present_ok(fut.await, BlobstoreIsPresent::Present);

            tickable_blobstores[0].1.drain(1);
            tickable_blobstores[2].1.drain(1);
        }

        // if the first two `is_present` calls return `Absent`, multiplexed returns `Absent`
        {
            let mut fut = multiplex.is_present(&ctx, k).boxed();
            assert_pending(&mut fut).await;

            tickable_blobstores[0].1.tick(None);
            tickable_blobstores[2].1.tick(None);

            assert_is_present_ok(fut.await, BlobstoreIsPresent::Absent);
            tickable_blobstores[1].1.drain(1);
        }

        // if one `is_present` returns `Absent`, another 2 fail, multiplexed is unsure
        {
            let mut fut = multiplex.is_present(&ctx, k).boxed();
            assert_pending(&mut fut).await;

            tickable_blobstores[0].1.tick(None);
            for (id, store) in &tickable_blobstores[1..] {
                store.tick(Some(format!("bs{} failed", id).as_str()));
            }

            assert_is_present_ok(
                fut.await,
                BlobstoreIsPresent::ProbablyNotPresent(anyhow!("some failed!")),
            );
        }
    }

    Ok(())
}

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

    // Ensure that even if the quorum writes succeeded, the rest of the writes are not
    // being dropped.
    {
        let timeout = MultiplexTimeout {
            // we want writes to succeed
            write: Duration::from_secs(10),
            // and reads to fail because of timeout
            read: Duration::from_millis(5),
        };
        let (tickable_queue, tickable_blobstores, multiplex) =
            setup_multiplex(3, 2, Some(timeout))?;

        let v = make_value("v1");
        let k = "k1";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        // quorum puts succeed -> the multiplexed put succeeds
        tickable_blobstores[0].1.tick(None);
        tickable_blobstores[1].1.tick(None);
        assert!(put_fut.await.is_ok());

        // Now we'll try to respond to the last put, that haven't completed.
        // It should not panic, as the future is still in flight and still
        // waits for the response.
        tokio::time::sleep(Duration::from_millis(5)).await;
        let result = panic::catch_unwind(|| tickable_blobstores[2].1.tick(None));
        assert!(result.is_ok());

        // We'll try to read with a delayed response from the blobstores.
        // Get should fail.
        let mut fut = multiplex.get(&ctx, k).boxed();
        assert_pending(&mut fut).await;
        tokio::time::sleep(Duration::from_millis(25)).await;
        assert!(fut.await.is_err());
    }

    // Ensure that after the quorum writes succeeded, the rest of the writes are
    // dropped because of the timeout.
    {
        // set write timeout to a very low value
        let timeout = MultiplexTimeout::new(
            None,                           /* read */
            Some(Duration::from_millis(5)), /* write */
        );
        let (tickable_queue, tickable_blobstores, multiplex) =
            setup_multiplex(3, 2, Some(timeout))?;

        let v = make_value("v2");
        let k = "k2";

        let mut put_fut = multiplex.put(&ctx, k.to_owned(), v.clone()).boxed();
        assert_pending(&mut put_fut).await;

        // wal queue write succeeds
        tickable_queue.tick(None);
        assert_pending(&mut put_fut).await;

        // quorum puts succeed -> the multiplexed put succeeds
        tickable_blobstores[0].1.tick(None);
        tickable_blobstores[1].1.tick(None);
        assert!(put_fut.await.is_ok());

        // Now we'll delay the response again and respond to the last put,
        // that haven't completed.
        // It *should* panic, as the future was dropped due to the timeout.
        tokio::time::sleep(Duration::from_millis(25)).await;
        let result = panic::catch_unwind(|| tickable_blobstores[2].1.tick(None));
        assert!(result.is_err());
    }

    Ok(())
}

async fn assert_pending<T: Debug>(fut: &mut (impl Future<Output = T> + Unpin)) {
    match futures::poll!(fut) {
        Poll::Pending => {}
        state => {
            panic!("future must be pending, received: {:?}", state);
        }
    }
}

fn setup_multiplex(
    num: u64,
    quorum: usize,
    timeout: Option<MultiplexTimeout>,
) -> Result<(
    Arc<Tickable<BlobstoreWalEntry>>,
    Vec<(BlobstoreId, Arc<Tickable<(BlobstoreBytes, u64)>>)>,
    WalMultiplexedBlobstore,
)> {
    let (tickable_queue, wal_queue) = setup_queue();
    let (tickable_blobstores, blobstores) = setup_blobstores(num);
    let scuba = Scuba::new(
        MononokeScubaSampleBuilder::with_discard(),
        MononokeScubaSampleBuilder::with_discard(),
        nonzero!(1u64),
    )?;
    let multiplex = WalMultiplexedBlobstore::new(
        MultiplexId::new(1),
        wal_queue,
        blobstores,
        vec![],
        quorum,
        timeout,
        scuba,
    )?;

    Ok((tickable_queue, tickable_blobstores, multiplex))
}

type TickableBytes = Tickable<(BlobstoreBytes, u64)>;

fn setup_blobstores(
    num: u64,
) -> (
    Vec<(BlobstoreId, Arc<TickableBytes>)>,
    Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
) {
    let tickable_blobstores: Vec<_> = (0..num)
        .map(|id| (BlobstoreId::new(id), Arc::new(TickableBytes::new())))
        .collect();
    let blobstores = tickable_blobstores
        .clone()
        .into_iter()
        .map(|(id, store)| (id, store as Arc<dyn BlobstorePutOps>))
        .collect();
    (tickable_blobstores, blobstores)
}

fn setup_queue() -> (Arc<Tickable<BlobstoreWalEntry>>, Arc<dyn BlobstoreWal>) {
    let tickable_queue = Arc::new(Tickable::new());
    let wal_queue = tickable_queue.clone() as Arc<dyn BlobstoreWal>;
    (tickable_queue, wal_queue)
}

fn make_value(value: &str) -> BlobstoreBytes {
    BlobstoreBytes::from_bytes(Bytes::copy_from_slice(value.as_bytes()))
}

fn validate_blob(
    get_data: Result<Option<BlobstoreGetData>>,
    expected: Result<Option<&BlobstoreBytes>>,
) {
    assert_eq!(get_data.is_ok(), expected.is_ok());
    if let Ok(expected) = expected {
        let get_data = get_data.unwrap().map(|data| data.into_bytes());
        assert_eq!(get_data.as_ref(), expected);
    }
}

fn assert_is_present_ok(result: Result<BlobstoreIsPresent>, expected: BlobstoreIsPresent) {
    assert!(result.is_ok());
    match (result.unwrap(), expected) {
        (BlobstoreIsPresent::Absent, BlobstoreIsPresent::Absent)
        | (BlobstoreIsPresent::Present, BlobstoreIsPresent::Present)
        | (BlobstoreIsPresent::ProbablyNotPresent(_), BlobstoreIsPresent::ProbablyNotPresent(_)) => {
        }
        (res, exp) => {
            panic!(
                "`is_present` call must return {:?}, received: {:?}",
                exp, res
            );
        }
    }
}

async fn scrub_none(
    fb: FacebookInit,
    scrub_action_on_missing_write_only: SrubWriteOnly,
) -> Result<()> {
    let bid0 = BlobstoreId::new(0);
    let bs0 = Arc::new(Tickable::new());
    let bid1 = BlobstoreId::new(1);
    let bs1 = Arc::new(Tickable::new());
    let bid2 = BlobstoreId::new(2);
    let bs2 = Arc::new(Tickable::new());

    let (_, queue) = setup_queue();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);
    let bs = WalScrubBlobstore::new(
        MultiplexId::new(1),
        queue,
        vec![(bid0, bs0.clone()), (bid1, bs1.clone())],
        vec![(bid2, bs2.clone())],
        1,
        None,
        Scuba::new_from_raw(fb, None, None, nonzero!(1u64))?,
        ScrubOptions {
            scrub_action_on_missing_write_only,
            ..ScrubOptions::default()
        },
        Arc::new(LoggingScrubHandler::new(false)) as Arc<dyn ScrubHandler>,
    )?;

    let mut fut = bs.get(ctx, "key");
    assert!(futures::poll!(&mut fut).is_pending());

    // No entry for "key" - blobstores return None...
    bs0.tick(None);
    bs1.tick(None);
    // Expect a read from writemostly stores regardless
    bs2.tick(None);

    fut.await?;

    Ok(())
}

#[fbinit::test]
async fn scrub_blobstore_fetch_none(fb: FacebookInit) -> Result<()> {
    scrub_none(fb, SrubWriteOnly::Scrub).await?;
    scrub_none(fb, SrubWriteOnly::SkipMissing).await?;
    scrub_none(fb, SrubWriteOnly::PopulateIfAbsent).await
}
async fn scrub_scenarios(
    fb: FacebookInit,
    scrub_action_on_missing_write_only: SrubWriteOnly,
) -> Result<()> {
    use SrubWriteOnly::*;
    println!("{:?}", scrub_action_on_missing_write_only);
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx);
    let (tick_queue, queue) = setup_queue();
    let scrub_handler = Arc::new(LoggingScrubHandler::new(false)) as Arc<dyn ScrubHandler>;
    let bid0 = BlobstoreId::new(0);
    let bs0 = Arc::new(Tickable::new());
    let bid1 = BlobstoreId::new(1);
    let bs1 = Arc::new(Tickable::new());
    let bid2 = BlobstoreId::new(2);
    let bs2 = Arc::new(Tickable::new());
    let scuba = Scuba::new_from_raw(fb, None, None, nonzero!(1u64))?;
    let bs = WalScrubBlobstore::new(
        MultiplexId::new(1),
        queue.clone(),
        vec![(bid0, bs0.clone()), (bid1, bs1.clone())],
        vec![(bid2, bs2.clone())],
        1,
        None,
        scuba.clone(),
        ScrubOptions {
            scrub_action: ScrubAction::ReportOnly,
            scrub_action_on_missing_write_only,
            ..ScrubOptions::default()
        },
        scrub_handler.clone(),
    )?;

    // non-existing key when one main blobstore failing
    {
        let k0 = "k0";

        let mut get_fut = bs.get(ctx, k0).boxed();
        assert_pending(&mut get_fut).await;

        bs0.tick(None);
        assert_pending(&mut get_fut).await;

        bs1.tick(Some("bs1 failed"));

        bs2.tick(None);
        assert!(get_fut.await.is_err(), "SomeNone + Err expected Err");
    }

    // non-existing key when one write only blobstore failing
    {
        let k0 = "k0";

        let mut get_fut = bs.get(ctx, k0).boxed();
        assert_pending(&mut get_fut).await;

        bs0.tick(None);
        assert_pending(&mut get_fut).await;

        bs1.tick(None);

        match scrub_action_on_missing_write_only {
            PopulateIfAbsent | ScrubIfAbsent => {
                // bs2 is ignored because it's write_only and the result of the normal
                // blobstores is failing
                assert_eq!(get_fut.await.unwrap(), None, "SomeNone + Err expected None");
            }
            _ => {
                bs2.tick(Some("bs2 failed"));
                // Write mostly blobstore failed but we still have quorum to say
                // the get succeeded
                assert_eq!(get_fut.await.unwrap(), None, "SomeNone + Err expected None");
            }
        }
    }

    // fail all but one store on put to make sure only one has the data
    // only replica containing key fails on read.
    {
        let v1 = make_value("v1");
        let k1 = "k1";

        let mut put_fut = bs.put(ctx, k1.to_owned(), v1.clone()).boxed();
        assert_pending(&mut put_fut).await;
        tick_queue.tick(None);
        assert_pending(&mut put_fut).await;
        bs0.tick(None);
        bs1.tick(Some("bs1 failed"));
        bs2.tick(Some("bs2 failed"));
        put_fut.await.unwrap();

        assert_eq!(
            queue
                .read(ctx, &MultiplexId::new(1), &Timestamp::now(), 2)
                .await
                .unwrap()
                .len(),
            1
        );

        let mut get_fut = bs.get(ctx, k1).boxed();
        assert_pending(&mut get_fut).await;

        bs0.tick(Some("bs0 failed"));
        assert_pending(&mut get_fut).await;

        bs1.tick(None);
        bs2.tick(None);
        assert!(get_fut.await.is_err(), "None/Err while replicating");
    }

    // all replicas fail
    {
        let k2 = "k2";

        let mut get_fut = bs.get(ctx, k2).boxed();
        assert_pending(&mut get_fut).await;
        bs0.tick(Some("bs0 failed"));
        bs1.tick(Some("bs1 failed"));
        bs2.tick(Some("bs1 failed"));
        assert!(get_fut.await.is_err(), "Err/Err");
    }

    // Now replace bs1 & bs2 with an empty blobstore, and see the scrub work
    let bid1 = BlobstoreId::new(1);
    let bs1 = Arc::new(Tickable::new());
    let bid2 = BlobstoreId::new(2);
    let bs2 = Arc::new(Tickable::new());
    let bs = WalScrubBlobstore::new(
        MultiplexId::new(1),
        queue.clone(),
        vec![(bid0, bs0.clone()), (bid1, bs1.clone())],
        vec![(bid2, bs2.clone())],
        1,
        None,
        scuba.clone(),
        ScrubOptions {
            scrub_action: ScrubAction::Repair,
            scrub_action_on_missing_write_only,
            ..ScrubOptions::default()
        },
        scrub_handler.clone(),
    )?;

    // Non-existing key in all blobstores, new blobstore failing
    {
        let k0 = "k0";

        let mut get_fut = bs.get(ctx, k0).boxed();
        assert_pending(&mut get_fut).await;

        bs0.tick(None);
        assert_pending(&mut get_fut).await;

        bs1.tick(Some("bs1 failed"));

        bs2.tick(None);
        assert!(get_fut.await.is_err(), "None/Err after replacement");
    }

    // only replica containing key replaced after failure - DATA LOST
    {
        let k1 = "k1";

        let mut get_fut = bs.get(ctx, k1).boxed();
        assert_pending(&mut get_fut).await;
        bs0.tick(Some("bs0 failed"));
        bs1.tick(None);
        bs2.tick(None);
        assert!(get_fut.await.is_err(), "Empty replacement against error");
    }

    // One working replica after failure.
    {
        let v1 = make_value("v1");
        let k1 = "k1";

        match queue
            .read(ctx, &MultiplexId::new(1), &Timestamp::now(), 2)
            .await
            .unwrap()
            .as_slice()
        {
            [entry] => {
                let to_del = [entry.clone()];
                let mut fut = queue.delete(ctx, &to_del).boxed();
                assert!(futures::poll!(&mut fut).is_pending());
                tick_queue.tick(None);
                fut.await.unwrap()
            }
            _ => panic!("only one entry expected"),
        }

        // bs1 and bs2 empty at this point
        assert_eq!(bs0.get_bytes(k1), Some(v1.clone()));
        assert!(bs1.storage.with(|s| s.is_empty()));
        assert!(bs2.storage.with(|s| s.is_empty()));

        let mut get_fut = bs.get(ctx, k1).boxed();
        assert_pending(&mut get_fut).await;
        // tick the gets
        bs0.tick(None);
        assert_pending(&mut get_fut).await;
        bs1.tick(None);
        if scrub_action_on_missing_write_only != SrubWriteOnly::PopulateIfAbsent {
            // this read doesn't happen in this mode
            bs2.tick(None);
        }
        assert_pending(&mut get_fut).await;
        // Tick the repairs
        bs1.tick(None);
        bs2.tick(None);

        // Succeeds
        assert_eq!(get_fut.await.unwrap().map(|v| v.into()), Some(v1.clone()));
        // Now all populated.
        assert_eq!(bs0.get_bytes(k1), Some(v1.clone()));
        assert_eq!(bs1.get_bytes(k1), Some(v1.clone()));
        match scrub_action_on_missing_write_only {
            SrubWriteOnly::Scrub
            | SrubWriteOnly::PopulateIfAbsent
            | SrubWriteOnly::ScrubIfAbsent => {
                assert_eq!(bs2.get_bytes(k1), Some(v1.clone()))
            }
            SrubWriteOnly::SkipMissing => {
                assert_eq!(bs2.get_bytes(k1), None)
            }
        }
    }

    // Main blobstore gets ignored on failure
    {
        let v2 = make_value("v2");
        let k2 = "k2";

        bs0.add_bytes(k2.to_string(), v2.clone());
        // bs1 and bs2 don't have k2
        assert_eq!(bs0.get_bytes(k2), Some(v2.clone()));
        assert_eq!(bs1.get_bytes(k2), None);
        assert_eq!(bs2.get_bytes(k2), None);

        let mut get_fut = bs.get(ctx, k2).boxed();
        assert_pending(&mut get_fut).await;
        // gets
        bs0.tick(None);
        bs1.tick(Some("bs1 get failed"));
        if scrub_action_on_missing_write_only != SrubWriteOnly::PopulateIfAbsent {
            // this read doesn't happen in this mode
            bs2.tick(None);
        }
        if scrub_action_on_missing_write_only != SrubWriteOnly::SkipMissing {
            assert_pending(&mut get_fut).await;
            // repair bs2
            bs2.tick(None);
        }
        assert_eq!(get_fut.await.unwrap().map(|v| v.into()), Some(v2.clone()));
        // bs1 still doesn't have the value. This is expected because we don't
        // want a single failing blobstore blocking the scrub.
        assert_eq!(bs1.get_bytes(k2), None);
        if scrub_action_on_missing_write_only != SrubWriteOnly::SkipMissing {
            // bs2 got repaired successfully
            assert_eq!(bs2.get_bytes(k2), Some(v2.clone()));
        } else {
            assert_eq!(bs2.get_bytes(k2), None);
        }
    }
    Ok(())
}

#[fbinit::test]
async fn scrubbed(fb: FacebookInit) {
    scrub_scenarios(fb, SrubWriteOnly::Scrub).await.unwrap();
    scrub_scenarios(fb, SrubWriteOnly::SkipMissing)
        .await
        .unwrap();
    scrub_scenarios(fb, SrubWriteOnly::PopulateIfAbsent)
        .await
        .unwrap();
}


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

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

use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use blobstore_stats::record_get_stats;
use blobstore_stats::record_is_present_stats;
use blobstore_stats::record_put_stats;
use blobstore_stats::OperationType;
use context::CoreContext;
use futures::Future;
use futures_stats::TimedFutureExt;
use metaconfig_types::BlobstoreId;
use mononoke_types::BlobstoreBytes;
use scuba_ext::MononokeScubaSampleBuilder;
use tokio::time::timeout;

// inferred from the current timeout, see https://fburl.com/code/rgj8497o
const GET_REQUEST_TIMEOUT: Duration = Duration::from_secs(100);
const PUT_REQUEST_TIMEOUT: Duration = Duration::from_secs(600);

#[derive(Clone, Debug)]
pub struct MultiplexTimeout {
    pub read: Duration,
    pub write: Duration,
}

impl Default for MultiplexTimeout {
    fn default() -> Self {
        Self::new(None, None)
    }
}

impl MultiplexTimeout {
    /// This allows to set either both timeouts or only one of them
    pub fn new(read: Option<Duration>, write: Option<Duration>) -> Self {
        Self {
            read: read.unwrap_or(GET_REQUEST_TIMEOUT),
            write: write.unwrap_or(PUT_REQUEST_TIMEOUT),
        }
    }
}

#[derive(Clone)]
pub(crate) struct TimedStore {
    id: BlobstoreId,
    inner: Arc<dyn BlobstorePutOps>,
    /// Timeout enforced on the read/write futures, including those running in the background
    timeout: MultiplexTimeout,
}

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

impl std::fmt::Display for TimedStore {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.inner.fmt(f)
    }
}

impl TimedStore {
    pub(crate) fn new(
        id: BlobstoreId,
        inner: Arc<dyn BlobstorePutOps>,
        timeout: MultiplexTimeout,
    ) -> Self {
        Self { id, inner, timeout }
    }

    pub(crate) fn id(&self) -> &BlobstoreId {
        &self.id
    }

    pub(crate) async fn put(
        &self,
        ctx: &CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
        mut scuba: MononokeScubaSampleBuilder,
    ) -> Result<OverwriteStatus, (BlobstoreId, Error)> {
        let size = value.len();
        let put_fut = if let Some(put_behaviour) = put_behaviour {
            self.inner
                .put_explicit(ctx, key.clone(), value, put_behaviour)
        } else {
            self.inner.put_with_status(ctx, key.clone(), value)
        };

        let pc = ctx.clone().fork_perf_counters();
        let (stats, result) = with_timeout(put_fut, self.timeout.write).timed().await;

        record_put_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            &key,
            ctx.metadata().session_id().as_str(),
            size,
            Some(self.id.clone()),
            self.inner.clone(),
            None,
        );

        result.map_err(|er| (self.id.clone(), er))
    }

    pub(crate) async fn get(
        &self,
        ctx: &CoreContext,
        key: &str,
        operation: OperationType,
        mut scuba: MononokeScubaSampleBuilder,
    ) -> Result<Option<BlobstoreGetData>, Error> {
        let pc = ctx.clone().fork_perf_counters();
        let (stats, result) = with_timeout(self.inner.get(ctx, key), self.timeout.read)
            .timed()
            .await;

        record_get_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            key,
            ctx.metadata().session_id().as_str(),
            operation,
            Some(self.id.clone()),
            self.inner.clone(),
        );

        result
    }

    pub(crate) async fn is_present(
        &self,
        ctx: &CoreContext,
        key: &str,
        mut scuba: MononokeScubaSampleBuilder,
    ) -> (BlobstoreId, Result<BlobstoreIsPresent>) {
        let pc = ctx.clone().fork_perf_counters();
        let (stats, result) = with_timeout(self.inner.is_present(ctx, key), self.timeout.read)
            .timed()
            .await;

        record_is_present_stats(
            &mut scuba,
            &pc,
            stats,
            result.as_ref(),
            key,
            ctx.metadata().session_id().as_str(),
            Some(self.id.clone()),
            self.inner.clone(),
        );

        (self.id.clone(), result)
    }
}

pub(crate) fn with_timed_stores(
    blobstores: Vec<(BlobstoreId, Arc<dyn BlobstorePutOps>)>,
    to: MultiplexTimeout,
) -> Vec<TimedStore> {
    blobstores
        .into_iter()
        .map(|(id, bs)| TimedStore::new(id, bs, to.clone()))
        .collect()
}

async fn with_timeout<T>(fut: impl Future<Output = Result<T>>, to: Duration) -> Result<T> {
    let timeout_or_result = timeout(to, fut).await;
    timeout_or_result.unwrap_or_else(|_| Err(Error::msg("blobstore operation timeout")))
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "packblob crate",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=packblob-thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "packblob.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


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

use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use blobstore::SizeMetadata;
use bufsize::SizeCounter;
use bytes::Buf;
use bytes::BufMut;
use bytes::Bytes;
use bytes::BytesMut;
use fbthrift::compact_protocol;
use fbthrift::compact_protocol::CompactProtocolSerializer;
use fbthrift::serialize::Serialize as ThriftSerialize;
use mononoke_types::BlobstoreBytes;
use packblob_thrift::StorageEnvelope;
use packblob_thrift::StorageFormat;

use crate::pack;

enum HeaderType {
    PackBlobCompactFormat,
}

impl TryFrom<u32> for HeaderType {
    type Error = Error;
    fn try_from(value: u32) -> Result<Self, Self::Error> {
        match value {
            // 0 is thrift compact_protocol.  We can use other values for other encodings in future
            0 => Ok(HeaderType::PackBlobCompactFormat),
            _ => Err(format_err!("Unknown header value for packblob {}", value))?,
        }
    }
}

impl From<HeaderType> for u32 {
    fn from(value: HeaderType) -> u32 {
        match value {
            HeaderType::PackBlobCompactFormat => 0,
        }
    }
}

// Serialize a Thrift value using the compact protocol with a header prefixed
fn compact_serialize_with_header<T>(header: u32, v: T) -> Bytes
where
    T: ThriftSerialize<CompactProtocolSerializer<SizeCounter>>
        + ThriftSerialize<CompactProtocolSerializer<BytesMut>>,
{
    // Get the size for the compact serialization
    let sz = compact_protocol::serialize_size(&v) + size_of::<u32>();

    // Use BytesMut split to get two views on the buffer
    let mut overall_buf = BytesMut::with_capacity(sz);
    overall_buf.put_u32(header);
    let body_buf = overall_buf.split();

    // Thrift serialize
    let body = compact_protocol::serialize_to_buffer(&v, body_buf).into_inner();

    // Recombine header and body. This is O(1) as buffers are contiguous
    overall_buf.unsplit(body);

    overall_buf.freeze()
}

// new type so can implement conversions
pub(crate) struct PackEnvelope(pub packblob_thrift::StorageEnvelope);

impl PackEnvelope {
    pub fn decode(self, key: &str) -> Result<(BlobstoreBytes, SizeMetadata), Error> {
        Ok(match self.0.storage {
            StorageFormat::Single(single) => {
                let (decoded, unique_compressed_size) = pack::decode_independent(single)
                    .with_context(|| format!("While decoding independent {:?}", key))?;
                let sizing = SizeMetadata {
                    unique_compressed_size,
                    pack_meta: None,
                };
                (decoded, sizing)
            }
            StorageFormat::Packed(packed) => pack::decode_pack(packed, key)
                .with_context(|| format!("While decoding pack for {:?}", key))?,
            StorageFormat::UnknownField(e) => {
                return Err(format_err!("StorageFormat::UnknownField {:?}", e));
            }
        })
    }
}

impl TryFrom<BlobstoreBytes> for PackEnvelope {
    type Error = Error;

    fn try_from(bytes: BlobstoreBytes) -> Result<Self, Error> {
        let mut bytes = bytes.into_bytes();
        let header: HeaderType = HeaderType::try_from(bytes.get_u32())?;
        let t: StorageEnvelope = match header {
            HeaderType::PackBlobCompactFormat => compact_protocol::deserialize(&bytes)?,
        };
        Ok(PackEnvelope(t))
    }
}

impl From<PackEnvelope> for BlobstoreBytes {
    fn from(p: PackEnvelope) -> Self {
        let data = compact_serialize_with_header(HeaderType::PackBlobCompactFormat.into(), &p.0);
        BlobstoreBytes::from_bytes(data)
    }
}

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

    #[test]
    fn serialize_roundtrip() -> Result<(), Error> {
        let value = "hello world!".to_string();
        let magic = 42;
        let mut bytes = compact_serialize_with_header(magic, value.clone());
        let header = bytes.get_u32();
        assert_eq!(magic, header);
        let roundtripped: String = compact_protocol::deserialize(bytes)?;
        assert_eq!(value, roundtripped);
        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.
 */

mod envelope;
mod pack;
mod store;

pub use pack::get_entry_compressed_size;
pub use pack::EmptyPack;
pub use pack::Pack;
pub use pack::SingleCompressed;
pub use store::PackBlob;
pub use store::PackOptions;


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

use std::collections::HashMap;
use std::io;
use std::io::Cursor;
use std::io::Write;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use ascii::AsciiString;
use blobstore::PackMetadata;
use blobstore::SizeMetadata;
use bytes::Buf;
use bytes::BufMut;
use bytes::Bytes;
use bytes::BytesMut;
use mononoke_types::hash::Context as HashContext;
use mononoke_types::repo::EPH_REPO_PREFIX_REGEX;
use mononoke_types::repo::REPO_PREFIX_REGEX;
use mononoke_types::BlobstoreBytes;
use packblob_thrift::PackedEntry;
use packblob_thrift::PackedFormat;
use packblob_thrift::PackedValue;
use packblob_thrift::SingleValue;
use packblob_thrift::StorageEnvelope;
use packblob_thrift::StorageFormat;
use packblob_thrift::ZstdFromDictValue;
use zstd::bulk::Compressor;
use zstd::dict::EncoderDictionary;
use zstd::stream::read::Decoder as ZstdDecoder;
use zstd::stream::write::Encoder as ZstdEncoder;

use crate::envelope::PackEnvelope;
use crate::store;

/// A block of data compressed on its own, rather than in pack format
pub struct SingleCompressed {
    value: SingleValue,
}

impl SingleCompressed {
    /// Tries to compress the given blob with the given zstd level; will not compress
    /// if the result of compression is an increase in size
    pub fn new(zstd_level: i32, blob: BlobstoreBytes) -> Result<SingleCompressed> {
        let value = blob.into_bytes();
        let mut compressor = Compressor::new(zstd_level)?;
        let compressed = compressor.compress(&value)?;
        let value = if compressed.len() < value.len() {
            SingleValue::Zstd(Bytes::from(compressed))
        } else {
            SingleValue::Raw(value)
        };
        Ok(Self { value })
    }
    /// Always stores the value raw and uncompressed
    pub(crate) fn new_uncompressed(blob: BlobstoreBytes) -> SingleCompressed {
        let value = SingleValue::Raw(blob.into_bytes());
        Self { value }
    }
    /// Gets the size of this blob in compressed form, minus framing overheads
    pub fn get_compressed_size(&self) -> Result<usize> {
        get_value_compressed_size(&self.value)
    }
    pub(crate) fn into_blobstore_bytes(self) -> BlobstoreBytes {
        // Wrap in thrift encoding
        PackEnvelope(StorageEnvelope {
            storage: StorageFormat::Single(self.value),
        })
        .into()
    }
}

/// An empty pack with no data. Cannot be uploaded, takes a dictionary blob
#[derive(Debug)]
pub struct EmptyPack(i32);

/// A pack containing multiple entries, ready to extend or upload
pub struct Pack {
    zstd_level: i32,
    dictionaries: HashMap<String, EncoderDictionary<'static>>,
    entries: Vec<PackedEntry>,
}

impl EmptyPack {
    /// Creates a new EmptyPack
    pub fn new(zstd_level: i32) -> Self {
        EmptyPack(zstd_level)
    }

    /// Adds the first blob to the empty pack
    pub fn add_base_blob(self, key: String, blob: BlobstoreBytes) -> Result<Pack> {
        let zstd_level = self.0;
        let bytes = blob.into_bytes();

        let dictionary = EncoderDictionary::copy(&bytes, zstd_level);

        let cursor = Cursor::new(&bytes);
        let compressed = zstd::encode_all(cursor, zstd_level)?;
        let data = PackedValue::Single(SingleValue::Zstd(Bytes::from(compressed)));

        let mut dictionaries = HashMap::new();
        dictionaries.insert(key.clone(), dictionary);
        let entries = vec![PackedEntry { key, data }];
        Ok(Pack {
            zstd_level,
            dictionaries,
            entries,
        })
    }
}

impl Pack {
    /// Adds another data blob to a pack, delta'd against a previous key
    pub fn add_delta_blob(
        &mut self,
        dict_key: String,
        key: String,
        blob: BlobstoreBytes,
    ) -> Result<()> {
        if self.dictionaries.contains_key(&key) {
            bail!("Key {} cannot appear in the same pack twice", key);
        }
        let zstd = {
            let dictionary = self
                .dictionaries
                .get(&dict_key)
                .ok_or_else(|| format_err!("Cannot find dictionary for blob {}", dict_key))?;

            let mut compressed_blob = BytesMut::with_capacity(blob.len());
            let writer = (&mut compressed_blob).writer();
            let mut encoder = ZstdEncoder::with_prepared_dictionary(writer, dictionary)?;

            encoder.write_all(blob.as_bytes())?;
            encoder.finish()?;
            compressed_blob.freeze()
        };
        // This uses `blob` (raw data) to create a dictionary that improves compression
        // at the expense of requiring the decompressor to find blob before it can
        // decompress the resulting blob
        let dictionary = EncoderDictionary::copy(blob.as_bytes(), self.zstd_level);
        let data = PackedValue::ZstdFromDict(ZstdFromDictValue { dict_key, zstd });
        self.dictionaries.insert(key.clone(), dictionary);
        self.entries.push(PackedEntry { key, data });
        Ok(())
    }

    /// Returns the compressed size of the pack contents, minus framing overheads
    pub fn get_compressed_size(&self) -> Result<usize> {
        let mut size = 0;
        for e in self.entries.iter() {
            size = size + get_entry_compressed_size(e)? + e.key.len();
        }
        Ok(size)
    }

    pub fn entries(&self) -> &[PackedEntry] {
        &self.entries
    }

    /// Converts the pack into something that can go into a blobstore
    pub(crate) fn into_blobstore_bytes(
        self,
        pack_prefix: String,
    ) -> Result<(String, Vec<String>, BlobstoreBytes)> {
        for entry in &self.entries {
            let (prefix, _) = split_key_prefix(&entry.key);
            if !prefix.is_empty() {
                bail!(
                    "Key prefix {} found in packed blob key {}",
                    prefix,
                    entry.key
                );
            }
        }
        let link_keys = {
            let mut link_keys: Vec<String> =
                self.entries.iter().map(|entry| entry.key.clone()).collect();

            // As long as it has the same entries its the same pack.  Sort keys before hashing them
            link_keys.sort_unstable();
            link_keys
        };

        // Build the pack identity
        let mut pack_key = pack_prefix;
        pack_key.push_str(compute_pack_hash(&link_keys).as_str());
        pack_key.push_str(store::ENVELOPE_SUFFIX);

        let pack = PackedFormat {
            key: pack_key.clone(),
            entries: self.entries,
        };

        // Wrap in thrift encoding and return as bytes
        Ok((
            pack_key,
            link_keys,
            PackEnvelope(StorageEnvelope {
                storage: StorageFormat::Packed(pack),
            })
            .into(),
        ))
    }
}

fn get_value_compressed_size(value: &SingleValue) -> Result<usize> {
    match value {
        SingleValue::Raw(bytes) | SingleValue::Zstd(bytes) => Ok(bytes.len()),
        // Can't happen, by construction - this only takes values created by this module
        SingleValue::UnknownField(_) => bail!("Unknown field"),
    }
}

pub fn get_entry_compressed_size(entry: &PackedEntry) -> Result<usize> {
    match &entry.data {
        PackedValue::Single(value) => get_value_compressed_size(value),
        PackedValue::ZstdFromDict(ZstdFromDictValue { zstd, .. }) => Ok(zstd.len()),
        // Can't happen, by construction - this only takes values created by this module
        PackedValue::UnknownField(_) => bail!("Unknown field"),
    }
}

// returns (decoded, unique_compressed_size)
pub(crate) fn decode_independent(v: SingleValue) -> Result<(BlobstoreBytes, u64)> {
    let (compressed_size, decoded) = match v {
        SingleValue::Raw(v) => (v.len() as u64, BlobstoreBytes::from_bytes(v)),
        SingleValue::Zstd(v) => (
            v.len() as u64,
            zstd::decode_all(v.reader()).map(BlobstoreBytes::from_bytes)?,
        ),
        SingleValue::UnknownField(e) => bail!("SingleValue::UnknownField {:?}", e),
    };
    Ok((decoded, compressed_size))
}

fn decode_zstd_from_dict(
    k: &str,
    v: ZstdFromDictValue,
    dicts: &HashMap<String, BlobstoreBytes>,
) -> Result<(BlobstoreBytes, u64), Error> {
    match dicts.get(&v.dict_key) {
        Some(dict) => {
            let uncompressed_size = v.zstd.len() as u64;
            let data = v.zstd.reader();
            let mut decoder = ZstdDecoder::with_dictionary(data, dict.as_bytes())?;
            let mut output_bytes = BytesMut::new();
            let mut writer = (&mut output_bytes).writer();
            io::copy(&mut decoder, &mut writer)?;

            Ok((BlobstoreBytes::from_bytes(output_bytes), uncompressed_size))
        }
        None => Err(format_err!(
            "Dictionary {} not found for key {}",
            v.dict_key,
            k
        )),
    }
}

// Unpack `key` from `packed`
pub(crate) fn decode_pack(
    packed: PackedFormat,
    key: &str,
) -> Result<(BlobstoreBytes, SizeMetadata)> {
    // Strip key prefix, if any
    let (_, key) = split_key_prefix(key);

    let PackedFormat {
        key: pack_key,
        entries: pack_entries,
    } = packed;

    let mut entry_map = HashMap::new();
    for entry in pack_entries {
        entry_map.insert(entry.key, entry.data);
        if entry_map.contains_key(key) {
            // Dictionaries must come before their users, so we don't care about the rest of the pack
            break;
        }
    }
    // Decode time
    let mut decoded_blobs = HashMap::new();
    let mut keys_to_decode = vec![key.to_string()];
    let mut unique_compressed_size = 0;
    let mut relevant_compressed_size = 0;
    let mut relevant_uncompressed_size = 0;
    while let Some(next_key) = keys_to_decode.pop() {
        match entry_map.remove(&next_key) {
            None => {
                if next_key == key {
                    // Handled below
                    break;
                }
                bail!(
                    "Key {} needs dictionary {} but it is not in the pack",
                    key,
                    next_key
                );
            }
            Some(PackedValue::UnknownField(e)) => {
                bail!("PackedValue::UnknownField {:?}", e);
            }
            Some(PackedValue::Single(v)) => {
                let (decoded, compressed_size) = decode_independent(v)?;
                relevant_uncompressed_size += decoded.len() as u64;
                if next_key == key {
                    unique_compressed_size += compressed_size;
                }
                relevant_compressed_size += compressed_size;
                decoded_blobs.insert(next_key, decoded);
            }
            Some(PackedValue::ZstdFromDict(v)) => {
                if decoded_blobs.contains_key(&v.dict_key) {
                    let (decoded, compressed_size) =
                        decode_zstd_from_dict(&next_key, v, &decoded_blobs)?;
                    relevant_uncompressed_size += decoded.len() as u64;
                    if next_key == key {
                        unique_compressed_size += compressed_size;
                    }
                    relevant_compressed_size += compressed_size;
                    decoded_blobs.insert(next_key, decoded);
                } else {
                    // Can't yet decode it - push the keys we need to decode this onto the work queue in order.
                    keys_to_decode.push(next_key.clone());
                    keys_to_decode.push(v.dict_key.clone());
                    // And reinsert this for the next loop
                    entry_map.insert(next_key, PackedValue::ZstdFromDict(v));
                }
            }
        }
    }

    let decoded = decoded_blobs
        .remove(key)
        .ok_or_else(|| format_err!("Key {} not in the pack it is pointing to {}", key, pack_key))?;

    let pack_meta = PackMetadata {
        pack_key,
        relevant_compressed_size,
        relevant_uncompressed_size,
    };
    let sizing = SizeMetadata {
        unique_compressed_size,
        pack_meta: Some(pack_meta),
    };

    Ok((decoded, sizing))
}

// Hash the keys as they themselves are hashes
fn compute_pack_hash(keys: &[String]) -> AsciiString {
    let mut hash_context = HashContext::new(b"pack");
    for key in keys {
        hash_context.update(key.len().to_le_bytes());
        hash_context.update(key.as_bytes());
    }
    hash_context.finish().to_hex()
}

/// Find the key prefix for a given key.  Key prefixes are removed when
/// keys are stored in packs.  Returns the key prefix and the remainder
/// of the key.
fn split_key_prefix(key: &str) -> (&str, &str) {
    if let Some(m) = REPO_PREFIX_REGEX.find(key) {
        key.split_at(m.end())
    } else if let Some(m) = EPH_REPO_PREFIX_REGEX.find(key) {
        key.split_at(m.end())
    } else {
        ("", key)
    }
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;
    use rand::Rng;
    use rand::RngCore;
    use rand::SeedableRng;
    use rand_xorshift::XorShiftRng;

    use super::*;

    #[test]
    fn decode_independent_zstd_test() -> Result<()> {
        // Highly compressable!
        let bytes_in = vec![7u8; 65535];

        // Prepare a compressed blob
        let input = Cursor::new(bytes_in.clone());
        let bytes = Bytes::from(zstd::encode_all(input, 0 /* default */)?);
        assert!(bytes.len() < bytes_in.len());
        let expected_compressed_size = bytes.len() as u64;

        // Test the decoder
        let (decoded, compressed_size) = decode_independent(SingleValue::Zstd(bytes))?;
        assert_eq!(decoded.as_bytes(), &Bytes::from(bytes_in));

        // Check the metadata
        assert_eq!(expected_compressed_size, compressed_size);

        Ok(())
    }

    #[test]
    fn decode_zstd_from_dict_test() -> Result<()> {
        let mut rng = XorShiftRng::seed_from_u64(0); // reproducable Rng

        // Some partially incompressible data as base version
        let mut base_version = vec![7u8; 65535];
        rng.fill_bytes(&mut base_version[0..60000]);

        // Extend it with incompressible new data
        let mut next_version = base_version.clone();
        rng.fill_bytes(&mut next_version[60000..]);

        let diff = {
            let mut compressed_blob = BytesMut::new();
            let writer = (&mut compressed_blob).writer();
            let mut encoder = ZstdEncoder::with_dictionary(writer, 0, &base_version)?;

            encoder.write_all(&next_version)?;
            encoder.finish()?;
            compressed_blob.freeze()
        };
        let expected_compressed_size = diff.len() as u64;

        let base_key = "base".to_string();

        let mut dicts = HashMap::new();
        dicts.insert(base_key.clone(), BlobstoreBytes::from_bytes(base_version));

        // Test the decoder
        let (decoded, compressed_size) = decode_zstd_from_dict(
            "appkey",
            ZstdFromDictValue {
                dict_key: base_key,
                zstd: diff,
            },
            &dicts,
        )?;
        assert_eq!(decoded.as_bytes(), &Bytes::from(next_version));
        assert_eq!(expected_compressed_size, compressed_size);

        Ok(())
    }

    #[test]
    fn pack_zstd_from_dict_chain_test() -> Result<()> {
        let mut rng = XorShiftRng::seed_from_u64(0); // reproducable Rng

        let mut raw_data = vec![];
        let pack = EmptyPack::new(0);

        // Some partially compressible data as base version
        let mut base_version = vec![7u8; 65535];
        rng.fill_bytes(&mut base_version[0..30000]);
        let base_key = "0".to_string();
        raw_data.push(base_version.clone());
        let mut pack =
            pack.add_base_blob(base_key, BlobstoreBytes::from_bytes(base_version.clone()))?;

        // incrementally build a pack from seeded random data
        let mut prev_version = base_version;
        for i in 1..20 {
            let mut this_version = prev_version;
            let start = 30000 + i * 1000;
            let end = start + 1000;
            rng.fill(&mut this_version[start..end]);
            raw_data.push(this_version.clone());
            prev_version = this_version.clone();
            pack.add_delta_blob(
                (i - 1).to_string(),
                i.to_string(),
                BlobstoreBytes::from_bytes(this_version),
            )?;
        }

        let (key, links, blob) = pack.into_blobstore_bytes(String::new())?;

        let mut expected_links: Vec<String> = (0..20_u32).map(|i| i.to_string()).collect();
        expected_links.sort_unstable();
        assert_eq!(expected_links, links);

        let packed = {
            let envelope: PackEnvelope = blob.try_into()?;
            if let StorageFormat::Packed(pack) = envelope.0.storage {
                pack
            } else {
                bail!("Packing resulted in a single value, not a pack");
            }
        };

        assert_eq!(packed.key, key);

        // Check for any change of hashing approach
        assert_eq!(
            "49df56ffd39791720f4d4b5c05d3156dfeab994fcb6b0000b3c2343280b95975.pack",
            packed.key
        );

        // Test reads roundtrip back to the raw form and metadata is populated
        for (raw_data, i) in raw_data.into_iter().zip(0..20) {
            let (value, size_meta) = decode_pack(packed.clone(), &i.to_string())?;
            assert_eq!(value.into_bytes(), Bytes::from(raw_data));
            assert!(size_meta.unique_compressed_size > 0);
            assert!(size_meta.pack_meta.is_some());
            if let Some(pack_meta) = size_meta.pack_meta {
                assert_eq!(&packed.key, &pack_meta.pack_key);
                assert!(pack_meta.relevant_compressed_size > 0);
                assert!(pack_meta.relevant_uncompressed_size > 0);
                assert!(pack_meta.relevant_compressed_size >= size_meta.unique_compressed_size);
            }
        }

        Ok(())
    }

    #[test]
    fn pack_zstd_from_dict_unchained_test() -> Result<()> {
        let mut rng = XorShiftRng::seed_from_u64(0); // reproducable Rng

        let mut raw_data = vec![];
        let pack = EmptyPack::new(0);

        // Some partially compressible data as base version
        let mut base_version = vec![7u8; 65535];
        rng.fill_bytes(&mut base_version[0..30000]);
        let base_key = "0".to_string();
        raw_data.push(base_version.clone());
        let mut pack =
            pack.add_base_blob(base_key, BlobstoreBytes::from_bytes(base_version.clone()))?;

        // incrementally build a pack from seeded random data
        let mut prev_version = base_version;
        for i in 1..20 {
            let mut this_version = prev_version;
            let start = 30000 + i * 1000;
            let end = start + 1000;
            rng.fill(&mut this_version[start..end]);
            raw_data.push(this_version.clone());
            prev_version = this_version.clone();
            pack.add_delta_blob(
                "0".to_string(),
                i.to_string(),
                BlobstoreBytes::from_bytes(this_version),
            )?;
        }

        let (key, links, blob) = pack.into_blobstore_bytes(String::new())?;

        let mut expected_links: Vec<String> = (0..20_u32).map(|i| i.to_string()).collect();
        expected_links.sort_unstable();
        assert_eq!(expected_links, links);

        let packed = {
            let envelope: PackEnvelope = blob.try_into()?;
            if let StorageFormat::Packed(pack) = envelope.0.storage {
                pack
            } else {
                bail!("Packing resulted in a single value, not a pack");
            }
        };

        assert_eq!(packed.key, key);

        // Check for any change of hashing approach
        assert_eq!(
            "49df56ffd39791720f4d4b5c05d3156dfeab994fcb6b0000b3c2343280b95975.pack",
            packed.key
        );

        // Test reads roundtrip back to the raw form and metadata is populated
        for (raw_data, i) in raw_data.into_iter().zip(0..20) {
            let (value, size_meta) = decode_pack(packed.clone(), &i.to_string())?;
            assert_eq!(value.into_bytes(), Bytes::from(raw_data));
            assert!(size_meta.unique_compressed_size > 0);
            assert!(size_meta.pack_meta.is_some());
            if let Some(pack_meta) = size_meta.pack_meta {
                assert_eq!(&packed.key, &pack_meta.pack_key);
                assert!(pack_meta.relevant_compressed_size > 0);
                assert!(pack_meta.relevant_uncompressed_size > 0);
                assert!(pack_meta.relevant_compressed_size >= size_meta.unique_compressed_size);
            }
        }

        Ok(())
    }

    #[test]
    fn pack_size_test() -> Result<()> {
        let mut rng = XorShiftRng::seed_from_u64(0); // reproducable Rng

        let mut raw_data = vec![];
        let pack = EmptyPack::new(19);

        // Some partially compressible data as base version
        let mut base_version = vec![7u8; 65535];
        rng.fill_bytes(&mut base_version[0..30000]);
        let base_key = "0".to_string();
        raw_data.push(base_version.clone());
        let mut pack =
            pack.add_base_blob(base_key, BlobstoreBytes::from_bytes(base_version.clone()))?;

        // Check the compressed size is reasonable
        let base_compressed_size = pack.get_compressed_size()?;
        assert!(
            base_compressed_size > 1024,
            "Compression turned 64 KiB into {} bytes - suspiciously small",
            base_compressed_size
        );
        assert!(
            base_compressed_size < 65535,
            "Compression turned 64 KiB into {} bytes - expansion unexpected",
            base_compressed_size
        );

        // incrementally build a pack from seeded random data
        let mut prev_version = base_version;
        for i in 1..20 {
            let mut this_version = prev_version;
            let start = 30000 + i * 1000;
            let end = start + 1000;
            rng.fill(&mut this_version[start..end]);
            raw_data.push(this_version.clone());
            prev_version = this_version.clone();
            pack.add_delta_blob(
                (i - 1).to_string(),
                i.to_string(),
                BlobstoreBytes::from_bytes(this_version),
            )?;
        }

        // And check it's grown, but not too far given how compressible this should be.
        let compressed_size = pack.get_compressed_size()?;
        assert!(
            compressed_size > base_compressed_size,
            "Pack shrank as it gained data"
        );
        let limit = base_compressed_size + 20 * 1000;
        assert!(
            compressed_size < limit,
            "Pack grew by more than the size of added data. Expected {} < {}",
            compressed_size,
            limit,
        );

        Ok(())
    }
}


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

use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreEnumerationData;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreKeyParam;
use blobstore::BlobstoreKeySource;
use blobstore::BlobstoreMetadata;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use metaconfig_types::PackFormat;
use mononoke_types::BlobstoreBytes;

use crate::envelope::PackEnvelope;
use crate::pack;

#[derive(Clone, Debug, Default)]
pub struct PackOptions {
    // None - user didn't specify
    // Some(xxx) - user wants to override config
    pub override_put_format: Option<PackFormat>,
}

impl PackOptions {
    pub fn new(override_put_format: Option<PackFormat>) -> Self {
        Self {
            override_put_format,
        }
    }
}

/// A layer over an existing blobstore that uses thrift blob wrappers to allow packing and compression
#[derive(Debug)]
pub struct PackBlob<T> {
    inner: T,
    put_format: PackFormat,
}

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

impl<T> PackBlob<T> {
    pub fn new(inner: T, put_format: PackFormat) -> Self {
        Self { inner, put_format }
    }
}

// differentiate keys just in case packblob is run in an existing unpacked store
pub const ENVELOPE_SUFFIX: &str = ".pack";

#[async_trait]
impl<T: Blobstore + BlobstorePutOps> Blobstore for PackBlob<T> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let inner_get_data = {
            let inner_key = &[key, ENVELOPE_SUFFIX].concat();
            self.inner
                .get(ctx, inner_key)
                .await
                .with_context(|| format!("While getting inner data for {:?}", key))?
        };
        let inner_get_data = match inner_get_data {
            Some(inner_get_data) => inner_get_data,
            None => return Ok(None),
        };

        let ctime = inner_get_data.as_meta().ctime();
        let envelope: PackEnvelope = inner_get_data.into_bytes().try_into()?;
        let (decoded, sizing) = envelope.decode(key)?;
        let meta = BlobstoreMetadata::new(ctime, Some(sizing));
        Ok(Some(BlobstoreGetData::new(meta, decoded)))
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.inner
            .is_present(ctx, &[key, ENVELOPE_SUFFIX].concat())
            .await
    }

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

impl<T: BlobstorePutOps> PackBlob<T> {
    async fn put_impl<'a>(
        &'a self,
        ctx: &'a CoreContext,
        mut key: String,
        value: BlobstoreBytes,
        put_behaviour: Option<PutBehaviour>,
    ) -> Result<OverwriteStatus> {
        key.push_str(ENVELOPE_SUFFIX);

        let bytes = match self.put_format {
            PackFormat::ZstdIndividual(zstd_level) => {
                pack::SingleCompressed::new(zstd_level, value)?
            }
            PackFormat::Raw => pack::SingleCompressed::new_uncompressed(value),
        }
        .into_blobstore_bytes();

        // pass through the put after wrapping
        if let Some(put_behaviour) = put_behaviour {
            self.inner
                .put_explicit(ctx, key, bytes, put_behaviour)
                .await
        } else {
            self.inner.put_with_status(ctx, key, bytes).await
        }
    }
}

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

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

#[async_trait]
impl<B: BlobstoreKeySource + BlobstorePutOps> BlobstoreKeySource for PackBlob<B> {
    async fn enumerate<'a>(
        &'a self,
        ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData> {
        let mut enumeration = self.inner.enumerate(ctx, range).await?;
        // Include only keys with the envelope suffix, and remove the suffix
        // from those keys.
        enumeration.keys = enumeration
            .keys
            .into_iter()
            .filter_map(|mut key| {
                if key.ends_with(ENVELOPE_SUFFIX) {
                    let new_len = key.len() - ENVELOPE_SUFFIX.len();
                    key.truncate(new_len);
                    Some(key)
                } else {
                    None
                }
            })
            .collect();
        Ok(enumeration)
    }
}

#[async_trait]
impl<T: BlobstoreUnlinkOps> BlobstoreUnlinkOps for PackBlob<T> {
    async fn unlink<'a>(&'a self, ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        let inner_key = &[key, ENVELOPE_SUFFIX].concat();
        self.inner.unlink(ctx, inner_key).await
    }
}

impl<T: Blobstore + BlobstoreUnlinkOps> PackBlob<T> {
    /// Put packed content, returning the pack's key if successful.
    ///
    /// `key_prefix` is prefixed to all keys within the pack and used to
    /// create links to the pack for each packed key.
    ///
    /// `pack_prefix` is prefixed to the pack key, and is under the control of
    /// the packer, e.g. if packing only filecontent together packer can chose
    /// "repoXXXX.packed.file_content.".  It is used for the temporary pack
    /// file name and stored within the pack itself.
    pub async fn put_packed<'a>(
        &'a self,
        ctx: &'a CoreContext,
        pack: pack::Pack,
        key_prefix: String,
        pack_prefix: String,
    ) -> Result<String> {
        let (pack_key, link_keys, blob) = pack.into_blobstore_bytes(pack_prefix)?;

        // pass through the put after wrapping
        self.inner.put(ctx, pack_key.clone(), blob).await?;

        // add the links
        let links = FuturesUnordered::new();
        for key in link_keys {
            let key = format!("{}{}{}", key_prefix, key, ENVELOPE_SUFFIX);
            links.push(self.inner.copy(ctx, &pack_key, key));
        }
        links.try_collect().await?;

        // remove the pack key, so that only the entries links are keeping it live
        self.inner.unlink(ctx, &pack_key).await?;

        Ok(pack_key)
    }

    pub async fn put_single<'a>(
        &'a self,
        ctx: &'a CoreContext,
        mut key: String,
        value: pack::SingleCompressed,
    ) -> Result<OverwriteStatus> {
        key.push_str(ENVELOPE_SUFFIX);
        self.inner
            .put_explicit(
                ctx,
                key,
                value.into_blobstore_bytes(),
                PutBehaviour::Overwrite,
            )
            .await
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use borrowed::borrowed;
    use bytes::Bytes;
    use fbinit::FacebookInit;
    use memblob::Memblob;
    use rand::RngCore;
    use rand::SeedableRng;
    use rand_xorshift::XorShiftRng;

    use super::*;

    #[fbinit::test]
    async fn simple_roundtrip_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let inner_blobstore = Arc::new(Memblob::default());
        let packblob = PackBlob::new(inner_blobstore.clone(), PackFormat::Raw);

        let outer_key = "repo0000.randomkey";
        let value = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(b"appleveldata"));
        let _ = roundtrip(ctx, inner_blobstore.clone(), &packblob, outer_key, value).await?;
        Ok(())
    }

    #[fbinit::test]
    async fn compressible_roundtrip_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let innerblob = Arc::new(Memblob::default());
        let packblob = PackBlob::new(innerblob.clone(), PackFormat::ZstdIndividual(0));

        let bytes_in = Bytes::from(vec![7u8; 65535]);
        let value = BlobstoreBytes::from_bytes(bytes_in.clone());

        let outer_key = "repo0000.compressible";
        let inner_key = roundtrip(ctx, innerblob.clone(), &packblob, outer_key, value).await?;

        // check inner value is smaller
        let inner_value = innerblob.get(ctx, &inner_key).await?;
        assert!(inner_value.unwrap().into_bytes().len() < bytes_in.len());
        Ok(())
    }

    #[fbinit::test]
    async fn incompressible_roundtrip_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let innerblob = Arc::new(Memblob::default());
        let packblob = PackBlob::new(innerblob.clone(), PackFormat::ZstdIndividual(0));

        let mut rng = XorShiftRng::seed_from_u64(0); // reproducable Rng
        let mut bytes_in = vec![7u8; 65535];
        rng.fill_bytes(&mut bytes_in);
        let bytes_in = Bytes::from(bytes_in);

        let outer_key = "repo0000.incompressible";
        let value = BlobstoreBytes::from_bytes(bytes_in.clone());
        let inner_key = roundtrip(ctx, innerblob.clone(), &packblob, outer_key, value).await?;

        // check inner value is larger (due to being raw plus thrift encoding)
        let inner_value = innerblob.get(ctx, &inner_key).await?;
        assert!(inner_value.unwrap().into_bytes().len() > bytes_in.len());
        Ok(())
    }

    async fn roundtrip(
        ctx: &CoreContext,
        inner_blobstore: Arc<Memblob>,
        packblob: &PackBlob<Arc<Memblob>>,
        outer_key: &str,
        value: BlobstoreBytes,
    ) -> Result<String> {
        // Put, this will apply the thrift envelope and save to the inner store
        packblob
            .put(ctx, outer_key.to_owned(), value.clone())
            .await?;

        // Get, should remove the thrift envelope as it is loaded
        let fetched_value = packblob.get(ctx, outer_key).await?.unwrap();

        // Make sure the thrift wrapper is not still there!
        assert_eq!(value, fetched_value.into_bytes());

        // Make sure that inner blobstore stores has packed value (i.e. not equal to what was written)
        let inner_key = &[outer_key, ENVELOPE_SUFFIX].concat();
        let fetched_value = inner_blobstore.get(ctx, inner_key).await?.unwrap();

        assert_ne!(value, fetched_value.into_bytes());

        // Check is_present matches
        let is_present = inner_blobstore
            .is_present(ctx, inner_key)
            .await?
            .assume_not_found_if_unsure();
        assert!(is_present);

        // Check the key without suffix is not there
        let is_not_present = !inner_blobstore
            .is_present(ctx, outer_key)
            .await?
            .assume_not_found_if_unsure();
        assert!(is_not_present);

        Ok(inner_key.to_owned())
    }

    #[fbinit::test]
    async fn simple_pack_test(fb: FacebookInit) -> Result<()> {
        let mut input_values = vec![];
        let pack = pack::EmptyPack::new(0);

        let base_key = "app_key0".to_string();
        let base_data = BlobstoreBytes::from_bytes(b"app_data0" as &[u8]);
        input_values.push(base_data.clone());

        let mut pack = pack.add_base_blob(base_key.clone(), base_data)?;
        for i in 1..3 {
            let mut app_key = "app_key".to_string();
            app_key.push_str(&i.to_string());

            let app_data = format!("app_data{}", i);
            let app_data = BlobstoreBytes::from_bytes(Bytes::copy_from_slice(app_data.as_bytes()));
            input_values.push(app_data.clone());

            pack.add_delta_blob(base_key.clone(), app_key, app_data)?;
        }

        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let inner_blobstore = Memblob::default();
        let packblob = PackBlob::new(inner_blobstore.clone(), PackFormat::Raw);

        // put_packed, this will apply the thrift envelope and save to the inner store
        let inner_key = packblob
            .put_packed(
                ctx,
                pack,
                "repo0000.".to_string(),
                "repo0000.packed_app_data.".to_string(),
            )
            .await?;

        // Check the inner key is not visible, the pack operation unlinks it
        let is_present = inner_blobstore
            .is_present(ctx, &inner_key)
            .await?
            .assume_not_found_if_unsure();
        assert!(!is_present);

        for (expected, i) in input_values.into_iter().zip(0..3usize) {
            // Get, should remove the thrift envelope as it is loaded
            let fetched_value = packblob.get(ctx, &format!("repo0000.app_key{}", i)).await?;

            assert!(
                fetched_value.is_some(),
                "Failed to fetch repo0000.app_key{}",
                i
            );

            // Make sure the thrift wrapper is not still there
            assert_eq!(expected, fetched_value.unwrap().into_bytes());
        }
        Ok(())
    }

    #[fbinit::test]
    async fn single_precompressed_test(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let innerblob = Arc::new(Memblob::default());
        let packblob = PackBlob::new(innerblob, PackFormat::ZstdIndividual(0));

        let bytes_in = Bytes::from(vec![7u8; 65535]);
        let value = BlobstoreBytes::from_bytes(bytes_in.clone());

        let key = "repo0000.compressible";
        let compressed = pack::SingleCompressed::new(19, value.clone())?;

        assert!(
            compressed.get_compressed_size()? < 65535,
            "Blob grew in compression"
        );

        packblob
            .put_single(ctx, key.to_string(), compressed)
            .await?;

        assert_eq!(
            packblob.get(ctx, key).await?.map(|b| b.into_bytes()),
            Some(value)
        );
        Ok(())
    }
}


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

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreEnumerationData;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreKeyParam;
use blobstore::BlobstoreKeyRange;
use blobstore::BlobstoreKeySource;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use inlinable_string::InlinableString;
use mononoke_types::BlobstoreBytes;

/// A layer over an existing blobstore that prepends a fixed string to each get and put.
#[derive(Clone, Debug)]
pub struct PrefixBlobstore<T> {
    // Try to inline the prefix to ensure copies remain cheap. Most prefixes are short anyway.
    prefix: InlinableString,
    blobstore: T,
}

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

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

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

    pub fn prefix(&self) -> String {
        self.prefix.to_string()
    }
}

impl<T> PrefixBlobstore<T> {
    pub fn new<S: Into<InlinableString>>(blobstore: T, prefix: S) -> Self {
        let prefix = prefix.into();
        Self { prefix, blobstore }
    }

    #[inline]
    pub fn prepend(&self, key: impl AsRef<str>) -> String {
        [&self.prefix, key.as_ref()].concat()
    }

    #[inline]
    pub fn unprepend(&self, key: &str) -> String {
        key[self.prefix.len()..].to_string()
    }
}

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

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

    #[inline]
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.blobstore.is_present(ctx, &self.prepend(key)).await
    }

    async fn copy<'a>(
        &'a self,
        ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        self.blobstore
            .copy(ctx, &self.prepend(old_key), self.prepend(new_key))
            .await
    }
}

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

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

#[async_trait]
impl<T: BlobstoreUnlinkOps> BlobstoreUnlinkOps for PrefixBlobstore<T> {
    async fn unlink<'a>(&'a self, ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        self.blobstore.unlink(ctx, &self.prepend(key)).await
    }
}

#[async_trait]
impl<T: BlobstoreKeySource> BlobstoreKeySource for PrefixBlobstore<T> {
    async fn enumerate<'a>(
        &'a self,
        ctx: &'a CoreContext,
        range: &'a BlobstoreKeyParam,
    ) -> Result<BlobstoreEnumerationData> {
        let new_param = match range {
            BlobstoreKeyParam::Start(range) => BlobstoreKeyParam::Start(BlobstoreKeyRange {
                // Regardless of the value of the begin_key (empty or non-empty), we
                // need to prepend the prefix to begin the search from the first
                // prefix-included entry in the underlying blobstore.
                begin_key: self.prepend(&range.begin_key),
                end_key: if range.end_key.is_empty() {
                    // When the end-key is empty, we need to prepend the prefix to ensure
                    // that the search is limited to prefix-included entries. The \u{10ffff}
                    // (final valid unicode value) is added as a representative end-of-range
                    // character for restricting the search.
                    self.prepend("\u{10ffff}")
                } else {
                    self.prepend(&range.end_key)
                },
            }),
            // No need to prepend Continuation as we don't unprepend it
            p => p.clone(),
        };
        let mut res = self.blobstore.enumerate(ctx, &new_param).await?;
        res.keys = res.keys.into_iter().map(|k| self.unprepend(&k)).collect();
        Ok(res)
    }
}

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

    use super::*;

    #[fbinit::test]
    async fn test_prefix(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let prefixed = PrefixBlobstore::new(base.clone(), "prefix123-");
        let unprefixed_key = "foobar".to_string();
        let prefixed_key = "prefix123-foobar".to_string();

        prefixed
            .put(
                ctx,
                unprefixed_key.clone(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await
            .expect("put should succeed");

        // Test that both the prefixed and the unprefixed stores can access the key.
        assert_eq!(
            prefixed
                .get(ctx, &unprefixed_key)
                .await
                .expect("get should succeed")
                .expect("value should be present")
                .into_raw_bytes(),
            Bytes::from("test foobar"),
        );
        assert_eq!(
            base.get(ctx, &prefixed_key)
                .await
                .expect("get should succeed")
                .expect("value should be present")
                .into_raw_bytes(),
            Bytes::from("test foobar"),
        );

        // Test that is_present works for both the prefixed and unprefixed stores.
        assert!(
            prefixed
                .is_present(ctx, &unprefixed_key)
                .await
                .expect("is_present should succeed")
                .assume_not_found_if_unsure()
        );
        assert!(
            base.is_present(ctx, &prefixed_key)
                .await
                .expect("is_present should succeed")
                .assume_not_found_if_unsure()
        );

        let enumerated = prefixed
            .enumerate(ctx, &BlobstoreKeyParam::from(..))
            .await
            .unwrap();

        assert_eq!(enumerated.keys, hashset! { unprefixed_key.clone() });

        assert!(
            prefixed
                .enumerate(ctx, &BlobstoreKeyParam::from("foobar1".to_string()..))
                .await
                .unwrap()
                .keys
                .is_empty()
        );

        assert!(
            !prefixed
                .enumerate(ctx, &BlobstoreKeyParam::from("fooba".to_string()..))
                .await
                .unwrap()
                .keys
                .is_empty()
        );
    }
}


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

use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Attempt to put to ReadOnlyBlobstore for key {0}")]
    ReadOnlyPut(String),
}