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

use anyhow::format_err;
use anyhow::Error;
use futures::channel::mpsc;
use futures::future;
use futures::future::Future;
use futures::future::TryFutureExt;
use futures::sink::Sink;
use futures::sink::SinkExt;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use tokio::task::JoinHandle;

// NOTE: This buffer size is used by the Multiplexer to let the overall multiplexer make progress
// even if one part is being a little slow. The multiplexer typically has individual tasks running
// that are hashing data, so in practice, this buffer is only useful to smooth out one task hitting
// a bit of latency for some reason. It doesn't need to be very large (in fact, empirically, I've
// observed that it needs to be very small), and it shouldn't be mistaken with the Filestore's
// upload concurrency, which is how many upload tasks will be running concurrently (those are tasks
// that perform I/O); those tasks are spawned by prepare.rs, and their number is controlled by the
// Filestore's upload concurrency level.
const BUFFER_SIZE: usize = 5;

type InnerSink<T> = dyn Sink<T, Error = mpsc::SendError> + Send + std::marker::Unpin + 'static;

pub struct Multiplexer<T: Send + Sync> {
    buffer_size: usize,
    sink: Option<Box<InnerSink<T>>>,
}

#[derive(Debug)]
pub enum MultiplexerError<E> {
    /// InputError indicates that the input stream hit an error E.
    InputError(E),

    /// Cancelled indicates that one of the consumers add()'ed to this Multiplexer stopped
    /// accepting input (likely because it hit an error).
    Cancelled,
}

impl From<MultiplexerError<Error>> for Error {
    fn from(me: MultiplexerError<Error>) -> Error {
        use MultiplexerError::*;

        match me {
            InputError(e) => e,
            e @ Cancelled => format_err!("MultiplexerError: {:?}", e),
        }
    }
}

/// Multiplexer lets you attach a set of readers on a Stream, but run them as separate Tokio tasks.
/// In the Filestore, this is used to do hashing of the data Stream on separate tasks.
impl<T: Send + Sync + Clone + 'static> Multiplexer<T> {
    pub fn new() -> Self {
        Self {
            buffer_size: BUFFER_SIZE,
            sink: None,
        }
    }

    /// Add a new consumer on the Multiplexer. To do so, pass a builder that expects a Stream of
    /// data as input (the same data you'll pass in later when draining the Multiplexer), and
    /// returns a Future. This returns a Future with the same result, and an Error type of
    /// SpawnError, which is a thin wrapper over your Future's Error type.
    ///
    /// NOTE: The builder added to the multiplexer should consume the ENTIRE stream that gets
    /// drained. Terminating halfway through the stream through the use of short-circuiting
    /// functions like `any`, `all`, etc. results in `Multiplexer Cancelled` Error.
    pub fn add<I, F, B>(&mut self, builder: B) -> JoinHandle<I>
    where
        I: Send + 'static,
        F: Future<Output = I> + Send + 'static,
        B: FnOnce(mpsc::Receiver<T>) -> F,
    {
        let (sender, receiver) = mpsc::channel::<T>(self.buffer_size);

        let sink = match self.sink.take() {
            Some(sink) => Box::new(sink.fanout(sender)) as Box<InnerSink<T>>,
            None => Box::new(sender) as Box<InnerSink<T>>,
        };

        self.sink = Some(sink);

        // NOTE: We need to start the built future here to make sure it makes progress even if the
        // receiving channel we return here is not polled. This ensures that consumers can't
        // deadlock their Multiplexer.
        // TODO: Pass through an executor?
        tokio::task::spawn(builder(receiver))
    }

    /// Drain a Stream into the multiplexer.
    pub async fn drain<'a, S, E>(self, stream: S) -> Result<(), MultiplexerError<E>>
    where
        S: Stream<Item = Result<T, E>> + Send,
    {
        let Self { sink, .. } = self;

        match sink {
            // If we have a Sink, then forward the Stream's data into it, and differentiate between
            // errors originating from the Stream and errors originating from the Sink.
            Some(sink) => {
                let mut sink = sink.sink_map_err(|_| MultiplexerError::Cancelled);
                // NOTE: I'd like to use futures::pin_mut! here, but:
                // https://github.com/rust-lang/rust/issues/64552
                let mut stream = stream.map_err(MultiplexerError::InputError).boxed();
                sink.send_all(&mut stream).await
            }
            // If we have no Sink, then consume the Stream regardless.
            None => {
                stream
                    .try_for_each(|_| future::ready(Ok(())))
                    .map_err(MultiplexerError::InputError)
                    .await
            }
        }
    }
}


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

use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use futures::future;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::join;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::stream::{self};
use futures::task::Poll;
use mononoke_types::content_chunk::new_blob_and_pointer;
use mononoke_types::content_metadata_v2::ends_in_newline;
use mononoke_types::content_metadata_v2::first_line;
use mononoke_types::content_metadata_v2::is_ascii;
use mononoke_types::content_metadata_v2::is_binary;
use mononoke_types::content_metadata_v2::is_generated;
use mononoke_types::content_metadata_v2::is_partially_generated;
use mononoke_types::content_metadata_v2::is_utf8;
use mononoke_types::content_metadata_v2::newline_count;
use mononoke_types::content_metadata_v2::PartialMetadata;
use mononoke_types::hash;
use mononoke_types::BlobstoreKey;
use mononoke_types::ChunkedFileContents;
use mononoke_types::FileContents;

use crate::alias::add_aliases_to_multiplexer;
use crate::expected_size::ExpectedSize;
use crate::incremental_hash::hash_bytes;
use crate::incremental_hash::Blake3IncrementalHasher;
use crate::incremental_hash::ContentIdIncrementalHasher;
use crate::incremental_hash::GitSha1IncrementalHasher;
use crate::incremental_hash::Sha1IncrementalHasher;
use crate::incremental_hash::Sha256IncrementalHasher;
use crate::multiplexer::Multiplexer;
use crate::multiplexer::MultiplexerError;
use crate::streamhash::hash_stream;

#[derive(Debug, Clone)]
pub struct Prepared {
    pub sha1: hash::Sha1,
    pub sha256: hash::Sha256,
    pub git_sha1: hash::RichGitSha1,
    pub seeded_blake3: hash::Blake3,
    pub is_binary: bool,
    pub is_ascii: bool,
    pub is_utf8: bool,
    pub ends_in_newline: bool,
    pub newline_count: u64,
    pub first_line: Option<String>,
    pub is_generated: bool,
    pub is_partially_generated: bool,
    pub contents: FileContents,
}

pub async fn prepare_bytes(bytes: Bytes) -> Prepared {
    let sha1 = hash_bytes(Sha1IncrementalHasher::new(), &bytes);
    let sha256 = hash_bytes(Sha256IncrementalHasher::new(), &bytes);
    let git_sha1 = hash_bytes(GitSha1IncrementalHasher::new(&bytes), &bytes);
    let seeded_blake3 = hash_bytes(Blake3IncrementalHasher::new_seeded(), &bytes);
    let is_binary = is_binary(stream::once(future::ready(bytes.clone())));
    let is_ascii = is_ascii(stream::once(future::ready(bytes.clone())));
    let is_utf8 = is_utf8(stream::once(future::ready(bytes.clone())));
    let ends_in_newline = ends_in_newline(stream::once(future::ready(bytes.clone())));
    let newline_count = newline_count(stream::once(future::ready(bytes.clone())));
    let first_line = first_line(stream::once(future::ready(bytes.clone())));
    let is_generated = is_generated(stream::once(future::ready(bytes.clone())));
    let is_partially_generated = is_partially_generated(stream::once(future::ready(bytes.clone())));
    let (
        is_binary,
        is_ascii,
        is_utf8,
        ends_in_newline,
        newline_count,
        first_line,
        is_generated,
        is_partially_generated,
    ) = join!(
        is_binary,
        is_ascii,
        is_utf8,
        ends_in_newline,
        newline_count,
        first_line,
        is_generated,
        is_partially_generated
    );
    let contents = FileContents::Bytes(bytes);

    Prepared {
        sha1,
        sha256,
        git_sha1,
        seeded_blake3,
        contents,
        is_binary,
        is_ascii,
        is_utf8,
        ends_in_newline,
        newline_count,
        first_line,
        is_generated,
        is_partially_generated,
    }
}

/// Prepare a stream of bytes for upload. This will return a Prepared struct that can be used to
/// finalize the upload. The hashes we compute may depend on the size hint.
pub async fn prepare_chunked<B: Blobstore + Clone + 'static, S>(
    ctx: CoreContext,
    blobstore: B,
    expected_size: ExpectedSize,
    chunks: S,
    concurrency: usize,
) -> Result<Prepared, Error>
where
    S: Stream<Item = Result<Bytes, Error>> + Send,
{
    // NOTE: The Multiplexer makes clones of the Bytes we pass in. It's worth noting that Bytes
    // actually behaves like an Arc with an inner reference-counted handle to data, so those
    // clones are actually fairly cheap
    let mut multiplexer = Multiplexer::<Bytes>::new();

    // Spawn a stream for each hash we need to produce.
    let content_id =
        multiplexer.add(|stream| hash_stream(ContentIdIncrementalHasher::new(), stream));
    let aliases = add_aliases_to_multiplexer(&mut multiplexer, expected_size);
    let metadata = add_partial_metadata_to_multiplexer(&mut multiplexer);
    // For the file's contents, spawn new tasks for each individual chunk. This ensures that
    // each chunk is hashed and uploaded separately, and potentially on a different CPU core.
    // We allow up to concurrency uploads to progress at the same time, which creates
    // backpressure into the chunks Stream.
    let contents = multiplexer
        .add(move |stream| {
            stream
                .map(move |bytes| {
                    // NOTE: This is lazy to allow the hash computation for this chunk's ID to
                    // happen on a separate core.
                    let fut = {
                        cloned!(blobstore, ctx);
                        async move {
                            let (blob, pointer) = new_blob_and_pointer(bytes);

                            // TODO: Convert this along with other store calls to impl Storable for
                            // MononokeId.
                            blobstore
                                .put(&ctx, blob.id().blobstore_key(), blob.into())
                                .await?;

                            Result::<_, Error>::Ok(pointer)
                        }
                    };

                    async move { tokio::task::spawn(fut).await? }
                })
                .buffered(concurrency)
                .try_fold(vec![], |mut chunks, chunk| async move {
                    chunks.push(chunk);
                    Result::<_, Error>::Ok(chunks)
                })
        })
        .map(|res| match res {
            Ok(Ok(r)) => Ok(r),
            Ok(Err(e)) => Err(e),
            Err(e) => Err(e.into()),
        });

    let res = multiplexer.drain(chunks).await;

    // Coerce the Error value for all our futures to Error.
    let content_id = content_id.map_err(Error::from);
    let aliases = aliases.map_err(Error::from);
    let contents = contents.map_err(Error::from);
    let metadata = metadata.map_err(Error::from);

    let futs = future::try_join4(content_id, aliases, contents, metadata);

    match res {
        // All is well - get the results when our futures complete.
        Ok(_) => {
            let (content_id, aliases, chunks, metadata) = futs.await?;
            let contents = FileContents::Chunked(ChunkedFileContents::new(content_id, chunks));

            let (sha1, sha256, git_sha1, seeded_blake3) = aliases.redeem(contents.size())?;

            let prepared = Prepared {
                sha1,
                sha256,
                git_sha1,
                seeded_blake3,
                contents,
                is_ascii: metadata.is_ascii,
                is_binary: metadata.is_binary,
                is_utf8: metadata.is_utf8,
                ends_in_newline: metadata.ends_in_newline,
                newline_count: metadata.newline_count,
                first_line: metadata.first_line,
                is_generated: metadata.is_generated,
                is_partially_generated: metadata.is_partially_generated,
            };

            Ok(prepared)
        }
        // If the Multiplexer hit an error, then it's worth handling the Cancelled case
        // separately: Cancelled means our Multiplexer noted that one of its readers
        // stopped reading. Usually, this will be because one of the readers failed. So,
        // let's just poll the readers once to see if they have an error value ready, and
        // if so, let's return that Error (because it'll be a more usable one). If not,
        // we'll passthrough the cancellation (but, we do have a unit test to make sure we
        // hit the happy path that prettifies the error).
        Err(m @ MultiplexerError::Cancelled) => match futures::poll!(futs) {
            Poll::Ready(Err(e)) => Err(e),
            _ => Err(m.into()),
        },

        Err(m @ MultiplexerError::InputError(..)) => Err(m.into()),
    }
}

pub fn add_partial_metadata_to_multiplexer(
    multiplexer: &mut Multiplexer<Bytes>,
) -> impl Future<Output = Result<PartialMetadata, Error>> + std::marker::Unpin {
    let is_ascii = multiplexer.add(is_ascii).map_err(Error::from);
    let is_utf8 = multiplexer.add(is_utf8).map_err(Error::from);
    let is_binary = multiplexer.add(is_binary).map_err(Error::from);
    let ends_in_newline = multiplexer.add(ends_in_newline).map_err(Error::from);
    let newline_count = multiplexer.add(newline_count).map_err(Error::from);
    let first_line = multiplexer.add(first_line).map_err(Error::from);
    let is_generated = multiplexer.add(is_generated).map_err(Error::from);
    let is_partially_generated = multiplexer.add(is_partially_generated).map_err(Error::from);

    let fut1 = future::try_join4(is_ascii, is_utf8, is_binary, is_generated);
    let fut2 = future::try_join4(
        ends_in_newline,
        newline_count,
        first_line,
        is_partially_generated,
    );
    future::try_join(fut1, fut2)
        .map_ok(
            |(
                (is_ascii, is_utf8, is_binary, is_generated),
                (ends_in_newline, newline_count, first_line, is_partially_generated),
            )| {
                PartialMetadata {
                    is_binary,
                    is_ascii,
                    is_utf8,
                    ends_in_newline,
                    newline_count,
                    first_line,
                    is_generated,
                    is_partially_generated,
                }
            },
        )
        .map_err(Error::from)
}


/*
 * 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::cmp::Ordering;

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use context::CoreContext;
use futures::future::TryFutureExt;
use mononoke_types::ChunkedFileContents;
use mononoke_types::ContentId;
use mononoke_types::ContentMetadataV2;
use mononoke_types::FileContents;
use slog::debug;
use thiserror::Error;

use crate::fetch;
use crate::get_metadata;
use crate::store;
use crate::FetchKey;
use crate::FilestoreConfig;
use crate::StoreRequest;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Content not found: {0:?}")]
    ContentNotFound(ContentId),
}

/// Fetch a file from the blobstore and reupload it in a chunked form.
/// NOTE: This could actually unchunk a file if the chunk size threshold
/// is increased after the file is written.
pub async fn force_rechunk<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    config: FilestoreConfig,
    ctx: &CoreContext,
    content_id: ContentId,
) -> Result<ContentMetadataV2, Error> {
    let file_contents: FileContents = content_id
        .load(ctx, blobstore)
        .map_err(move |err| match err {
            LoadableError::Error(err) => err,
            LoadableError::Missing(_) => ErrorKind::ContentNotFound(content_id).into(),
        })
        .await?;
    do_rechunk_file_contents(blobstore, config, ctx, file_contents, content_id).await
}

/// Fetch a file from the blobstore and reupload it in a chunked form
/// only if it is chunked using a larger chunk size (or unchunked)
/// Note that this fn is not suitable for unchunking a file,
/// as if existing file uses smaller-than-requested chunk size,
/// this fn won't do anything.
/// Returns a future, resolving to the `ContentMetadataV2` of the
/// processed `ContentId` and whether it was *actually* rechunked
pub async fn rechunk<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    filestore_config: FilestoreConfig,
    ctx: &CoreContext,
    content_id: ContentId,
) -> Result<(ContentMetadataV2, bool), Error> {
    let fetch_key = FetchKey::Canonical(content_id.clone());
    let chunk_size = filestore_config.chunk_size;
    let metadata = get_metadata(blobstore, ctx, &fetch_key).await?;
    let content_metadata: ContentMetadataV2 = match metadata {
        Some(content_metadata) => content_metadata,
        None => return Err(ErrorKind::ContentNotFound(content_id).into()),
    };

    match chunk_size {
        Some(chunk_size) if content_metadata.total_size > chunk_size => {
            let r: Result<(ContentMetadataV2, bool), Error> = rechunk_if_uses_larger_chunk_size(
                blobstore,
                chunk_size,
                filestore_config.concurrency,
                ctx,
                content_metadata,
            )
            .await;

            r
        }
        _ => Ok((content_metadata, false)),
    }
}

/// Return true if stored `chunked_file_contents` uses chunks larger
/// than `expected_chunk_size`
fn uses_larger_chunks(
    ctx: &CoreContext,
    chunked_file_contents: &ChunkedFileContents,
    expected_chunk_size: u64,
    content_id: &ContentId,
) -> bool {
    let mut all_smaller_or_equal = true;
    let mut num_smaller_chunks = 0;
    let num_chunks = chunked_file_contents.num_chunks();
    for (idx, content_chunk_pointer) in chunked_file_contents.iter_chunks().enumerate() {
        let real_chunk_size = content_chunk_pointer.size();

        match real_chunk_size.cmp(&expected_chunk_size) {
            Ordering::Greater => {
                all_smaller_or_equal = false;
                break;
            }
            Ordering::Less => {
                if idx == num_chunks - 1 {
                    // last pointer, it's ok if it is smaller
                    continue;
                }

                num_smaller_chunks += 1;
            }
            Ordering::Equal => {}
        }
    }

    if num_smaller_chunks > 0 {
        debug!(
            ctx.logger(),
            "{} chunks of {} have size smaller than what we want: {}. No action will be taken",
            num_smaller_chunks,
            content_id,
            expected_chunk_size
        );
    }

    !all_smaller_or_equal
}

/// For content, represented by `content_metadata`, rechunk it
/// if it is unchunked or uses larger chunk sizes
/// Note: this fn expects `expected_chunk_size` and `concurrency`
/// instead of `FilestoreConfig` to emphasize that it can only be
/// called, if the filestore's chunk size is not `None`
async fn rechunk_if_uses_larger_chunk_size<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    expected_chunk_size: u64,
    concurrency: usize,
    ctx: &CoreContext,
    content_metadata: ContentMetadataV2,
) -> Result<(ContentMetadataV2, bool), Error> {
    let content_id = content_metadata.content_id.clone();

    let file_contents: FileContents = content_id
        .load(ctx, blobstore)
        .map_err(move |err| match err {
            LoadableError::Error(err) => err,
            LoadableError::Missing(_) => ErrorKind::ContentNotFound(content_id).into(),
        })
        .await?;

    let should_rechunk = match file_contents {
        FileContents::Bytes(_) => true,
        FileContents::Chunked(ref chunked_file_contents) => {
            uses_larger_chunks(ctx, chunked_file_contents, expected_chunk_size, &content_id)
        }
    };

    if should_rechunk {
        let filestore_config = FilestoreConfig {
            chunk_size: Some(expected_chunk_size),
            concurrency,
        };

        let content_metadata: ContentMetadataV2 =
            do_rechunk_file_contents(blobstore, filestore_config, ctx, file_contents, content_id)
                .await?;

        Ok((content_metadata, true))
    } else {
        Ok((content_metadata, false))
    }
}

/// Unconditionally rechunk `file_contents` using the `filestore_config`
/// NOTE: This could actually unchunk a file if the chunk size threshold
/// is increased after the file is written.
async fn do_rechunk_file_contents<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    filestore_config: FilestoreConfig,
    ctx: &CoreContext,
    file_contents: FileContents,
    content_id: ContentId,
) -> Result<ContentMetadataV2, Error> {
    let req = StoreRequest::with_canonical(file_contents.size(), content_id);
    let file_stream = fetch::stream_file_bytes(blobstore, ctx, file_contents, fetch::Range::all())?;

    store(blobstore, filestore_config, ctx, &req, file_stream).await
}


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

use futures::future;
use futures::future::FutureExt;
use futures::stream::Stream;
use futures::stream::StreamExt;

use crate::incremental_hash::Hasher;

pub async fn hash_stream<H, I, S>(hasher: impl Hasher<H>, stream: S) -> H
where
    I: AsRef<[u8]>,
    S: Stream<Item = I>,
{
    stream
        .fold(hasher, |mut hasher, bytes| {
            hasher.update(bytes);
            future::ready(hasher)
        })
        .map(|hasher| hasher.finish())
        .await
}

#[cfg(test)]
mod test {
    use bytes::Bytes;
    use futures::stream;
    use mononoke_types::hash::Blake3;
    use mononoke_types::hash::RichGitSha1;
    use mononoke_types::hash::Sha1;
    use mononoke_types::hash::Sha256;
    use mononoke_types::ContentId;

    use super::*;
    use crate::expected_size::ExpectedSize;
    use crate::incremental_hash::Blake3IncrementalHasher;
    use crate::incremental_hash::ContentIdIncrementalHasher;
    use crate::incremental_hash::GitSha1IncrementalHasher;
    use crate::incremental_hash::Sha1IncrementalHasher;
    use crate::incremental_hash::Sha256IncrementalHasher;

    #[tokio::test]
    async fn sha1_simple() {
        let data = Bytes::from(&b"hello, world"[..]); // b7e23ec29af22b0b4e41da31e868d57226121c84
        let s = stream::once(future::ready(data));

        let res: Sha1 = hash_stream(Sha1IncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            Sha1::from_bytes([
                0xb7, 0xe2, 0x3e, 0xc2, 0x9a, 0xf2, 0x2b, 0x0b, 0x4e, 0x41, 0xda, 0x31, 0xe8, 0x68,
                0xd5, 0x72, 0x26, 0x12, 0x1c, 0x84
            ])
            .unwrap()
        );
    }

    #[tokio::test]
    async fn sha1_chunks() {
        let data = vec![&b"hello"[..], &b", "[..], &b"world"[..]] // b7e23ec29af22b0b4e41da31e868d57226121c84
            .into_iter()
            .map(Bytes::from);
        let s = stream::iter(data);

        let res: Sha1 = hash_stream(Sha1IncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            Sha1::from_bytes([
                0xb7, 0xe2, 0x3e, 0xc2, 0x9a, 0xf2, 0x2b, 0x0b, 0x4e, 0x41, 0xda, 0x31, 0xe8, 0x68,
                0xd5, 0x72, 0x26, 0x12, 0x1c, 0x84
            ])
            .unwrap()
        );
    }

    #[tokio::test]
    async fn git_sha1_simple() {
        let data = Bytes::from(&b"hello, world"[..]); // 8c01d89ae06311834ee4b1fab2f0414d35f01102
        let s = stream::once(future::ready(data));

        let res: RichGitSha1 =
            hash_stream(GitSha1IncrementalHasher::new(ExpectedSize::new(12)), s).await;

        assert_eq!(
            res,
            RichGitSha1::from_bytes(
                [
                    0x8c, 0x01, 0xd8, 0x9a, 0xe0, 0x63, 0x11, 0x83, 0x4e, 0xe4, 0xb1, 0xfa, 0xb2,
                    0xf0, 0x41, 0x4d, 0x35, 0xf0, 0x11, 0x02
                ],
                "blob",
                12
            )
            .unwrap()
        );
    }

    #[tokio::test]
    async fn git_sha1_chunks() {
        let data = vec![&b"hello"[..], &b", "[..], &b"world"[..]] // 8c01d89ae06311834ee4b1fab2f0414d35f01102
            .into_iter()
            .map(Bytes::from);
        let s = stream::iter(data);

        let res: RichGitSha1 =
            hash_stream(GitSha1IncrementalHasher::new(ExpectedSize::new(12)), s).await;

        assert_eq!(
            res,
            RichGitSha1::from_bytes(
                [
                    0x8c, 0x01, 0xd8, 0x9a, 0xe0, 0x63, 0x11, 0x83, 0x4e, 0xe4, 0xb1, 0xfa, 0xb2,
                    0xf0, 0x41, 0x4d, 0x35, 0xf0, 0x11, 0x02
                ],
                "blob",
                12
            )
            .unwrap()
        );
    }

    #[tokio::test]
    async fn sha256_simple() {
        let data = Bytes::from(&b"hello, world"[..]); // 09ca7e4eaa6e8ae9c7d261167129184883644d07dfba7cbfbc4c8a2e08360d5b
        let s = stream::once(future::ready(data));

        let res: Sha256 = hash_stream(Sha256IncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            Sha256::from_bytes([
                0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7, 0xd2, 0x61, 0x16, 0x71, 0x29,
                0x18, 0x48, 0x83, 0x64, 0x4d, 0x07, 0xdf, 0xba, 0x7c, 0xbf, 0xbc, 0x4c, 0x8a, 0x2e,
                0x08, 0x36, 0x0d, 0x5b,
            ],)
            .unwrap()
        );
    }

    #[tokio::test]
    async fn sha256_chunks() {
        let data = vec![&b"hello"[..], &b", "[..], &b"world"[..]] // 09ca7e4eaa6e8ae9c7d261167129184883644d07dfba7cbfbc4c8a2e08360d5b
            .into_iter()
            .map(Bytes::from);
        let s = stream::iter(data);

        let res: Sha256 = hash_stream(Sha256IncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            Sha256::from_bytes([
                0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7, 0xd2, 0x61, 0x16, 0x71, 0x29,
                0x18, 0x48, 0x83, 0x64, 0x4d, 0x07, 0xdf, 0xba, 0x7c, 0xbf, 0xbc, 0x4c, 0x8a, 0x2e,
                0x08, 0x36, 0x0d, 0x5b,
            ],)
            .unwrap()
        );
    }

    #[tokio::test]
    async fn content_id_simple() {
        let data = Bytes::from(&b"hello, world"[..]); // 19d95f338fa0f547f773c12353eb6dacb1a7ce9b0402515484e8276055774b35
        let s = stream::once(future::ready(data));

        let res: ContentId = hash_stream(ContentIdIncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            ContentId::from_bytes([
                0x19, 0xd9, 0x5f, 0x33, 0x8f, 0xa0, 0xf5, 0x47, 0xf7, 0x73, 0xc1, 0x23, 0x53, 0xeb,
                0x6d, 0xac, 0xb1, 0xa7, 0xce, 0x9b, 0x04, 0x02, 0x51, 0x54, 0x84, 0xe8, 0x27, 0x60,
                0x55, 0x77, 0x4b, 0x35,
            ],)
            .unwrap()
        );
    }

    #[tokio::test]
    async fn content_id_chunks() {
        let data = vec![&b"hello"[..], &b", "[..], &b"world"[..]] // 19d95f338fa0f547f773c12353eb6dacb1a7ce9b0402515484e8276055774b35
            .into_iter()
            .map(Bytes::from);
        let s = stream::iter(data);

        let res: ContentId = hash_stream(ContentIdIncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            ContentId::from_bytes([
                0x19, 0xd9, 0x5f, 0x33, 0x8f, 0xa0, 0xf5, 0x47, 0xf7, 0x73, 0xc1, 0x23, 0x53, 0xeb,
                0x6d, 0xac, 0xb1, 0xa7, 0xce, 0x9b, 0x04, 0x02, 0x51, 0x54, 0x84, 0xe8, 0x27, 0x60,
                0x55, 0x77, 0x4b, 0x35,
            ],)
            .unwrap()
        );
    }

    #[tokio::test]
    async fn sha1_empty() {
        let s = stream::empty::<Vec<u8>>();

        let res: Sha1 = hash_stream(Sha1IncrementalHasher::new(), s).await;

        assert_eq!(
            res,
            Sha1::from_bytes([
                0xda, 0x39, 0xa3, 0xee, 0x5e, 0x6b, 0x4b, 0x0d, 0x32, 0x55, 0xbf, 0xef, 0x95, 0x60,
                0x18, 0x90, 0xaf, 0xd8, 0x07, 0x09
            ])
            .unwrap()
        );
    }

    #[tokio::test]
    async fn seeded_blake3_simple() {
        let data = Bytes::from(&b"hello, world"[..]); // 8537f0e61caaa8f1aca4eac89bb9f27ba2562dbdd1b1a2d987999c37f159bf18
        let s = stream::once(future::ready(data));

        let res: Blake3 = hash_stream(Blake3IncrementalHasher::new_seeded(), s).await;

        let hello_world_hash = Blake3::from_bytes([
            0x85, 0x37, 0xf0, 0xe6, 0x1c, 0xaa, 0xa8, 0xf1, 0xac, 0xa4, 0xea, 0xc8, 0x9b, 0xb9,
            0xf2, 0x7b, 0xa2, 0x56, 0x2d, 0xbd, 0xd1, 0xb1, 0xa2, 0xd9, 0x87, 0x99, 0x9c, 0x37,
            0xf1, 0x59, 0xbf, 0x18,
        ])
        .unwrap();
        assert_eq!(res, hello_world_hash);
    }

    #[tokio::test]
    async fn seeded_blake3_chunks() {
        let data = vec![&b"hello"[..], &b", "[..], &b"world"[..]] // 8537f0e61caaa8f1aca4eac89bb9f27ba2562dbdd1b1a2d987999c37f159bf18
            .into_iter()
            .map(Bytes::from);
        let s = stream::iter(data);

        let res: Blake3 = hash_stream(Blake3IncrementalHasher::new_seeded(), s).await;
        let hello_world_hash = Blake3::from_bytes([
            0x85, 0x37, 0xf0, 0xe6, 0x1c, 0xaa, 0xa8, 0xf1, 0xac, 0xa4, 0xea, 0xc8, 0x9b, 0xb9,
            0xf2, 0x7b, 0xa2, 0x56, 0x2d, 0xbd, 0xd1, 0xb1, 0xa2, 0xd9, 0x87, 0x99, 0x9c, 0x37,
            0xf1, 0x59, 0xbf, 0x18,
        ])
        .unwrap();
        assert_eq!(res, hello_world_hash);
    }
}


/*
 * 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::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use rand::thread_rng;
use rand::Rng;
use thiserror::Error;

#[derive(Debug, Error)]
#[error("Failing Blobstore Error")]
pub struct FailingBlobstoreError;

#[derive(Debug, Clone)]
pub struct FailingBlobstore<B> {
    inner: B,
    read_success_probability: f64,
    write_success_probability: f64,
}

impl<B> FailingBlobstore<B> {
    pub fn new(inner: B, read_success_probability: f64, write_success_probability: f64) -> Self {
        Self {
            inner,
            read_success_probability,
            write_success_probability,
        }
    }
}

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

#[async_trait]
impl<B: Blobstore> Blobstore for FailingBlobstore<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        if thread_rng().gen_bool(self.read_success_probability) {
            self.inner.get(ctx, key).await
        } else {
            Err(FailingBlobstoreError.into())
        }
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        if thread_rng().gen_bool(self.write_success_probability) {
            self.inner.put(ctx, key, value).await
        } else {
            Err(FailingBlobstoreError.into())
        }
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        if thread_rng().gen_bool(self.read_success_probability) {
            self.inner.is_present(ctx, key).await
        } else {
            Err(FailingBlobstoreError.into())
        }
    }
}


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

use mononoke_types::typed_hash;
use mononoke_types::ContentChunkId;
use mononoke_types::ContentId;

use super::StoreRequest;

mod failing_blobstore;
mod test_api;
mod test_invariants;

fn request(data: impl AsRef<[u8]>) -> StoreRequest {
    StoreRequest::new(data.as_ref().len() as u64)
}

fn canonical(data: impl AsRef<[u8]>) -> ContentId {
    let mut ctx = typed_hash::ContentIdContext::new();
    ctx.update(data.as_ref());
    ctx.finish()
}

fn chunk(data: impl AsRef<[u8]>) -> ContentChunkId {
    let mut ctx = typed_hash::ContentChunkIdContext::new();
    ctx.update(data.as_ref());
    ctx.finish()
}


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

use anyhow::Error;
use anyhow::Result;
use assert_matches::assert_matches;
use blobstore::Blobstore;
use blobstore::PutBehaviour;
use borrowed::borrowed;
use bytes::Bytes;
use bytes::BytesMut;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::future;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::TryStreamExt;
use lazy_static::lazy_static;
use mononoke_types::hash;
use mononoke_types::typed_hash::BlobstoreKey;
use mononoke_types::ContentId;
use mononoke_types::ContentMetadataV2;
use mononoke_types::ContentMetadataV2Id;
use mononoke_types_mocks::contentid::ONES_CTID;

use super::canonical;
use super::chunk;
use super::failing_blobstore::FailingBlobstore;
use super::failing_blobstore::FailingBlobstoreError;
use super::request;
use crate as filestore;
use crate::errors;
use crate::Alias;
use crate::FetchKey;
use crate::FilestoreConfig;
use crate::StoreRequest;

const HELLO_WORLD: &[u8] = b"hello, world";
const HELLO_WORLD_LENGTH: u64 = 12;
const DEFAULT_CONFIG: FilestoreConfig = FilestoreConfig {
    chunk_size: None,
    concurrency: 1,
};

lazy_static! {
    static ref HELLO_WORLD_SHA1: hash::Sha1 = hash::Sha1::from_bytes([
        0xb7, 0xe2, 0x3e, 0xc2, 0x9a, 0xf2, 0x2b, 0x0b, 0x4e, 0x41, 0xda, 0x31, 0xe8, 0x68, 0xd5,
        0x72, 0x26, 0x12, 0x1c, 0x84
    ])
    .unwrap();
    static ref HELLO_WORLD_GIT_SHA1: hash::RichGitSha1 = hash::RichGitSha1::from_bytes(
        [
            0x8c, 0x01, 0xd8, 0x9a, 0xe0, 0x63, 0x11, 0x83, 0x4e, 0xe4, 0xb1, 0xfa, 0xb2, 0xf0,
            0x41, 0x4d, 0x35, 0xf0, 0x11, 0x02
        ],
        "blob",
        HELLO_WORLD_LENGTH
    )
    .unwrap();
    static ref HELLO_WORLD_SHA256: hash::Sha256 = hash::Sha256::from_bytes([
        0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7, 0xd2, 0x61, 0x16, 0x71, 0x29, 0x18,
        0x48, 0x83, 0x64, 0x4d, 0x07, 0xdf, 0xba, 0x7c, 0xbf, 0xbc, 0x4c, 0x8a, 0x2e, 0x08, 0x36,
        0x0d, 0x5b,
    ])
    .unwrap();
    static ref HELLO_WORLD_SEEDED_BLAKE3: hash::Blake3 = hash::Blake3::from_bytes([
        0x85, 0x37, 0xf0, 0xe6, 0x1c, 0xaa, 0xa8, 0xf1, 0xac, 0xa4, 0xea, 0xc8, 0x9b, 0xb9, 0xf2,
        0x7b, 0xa2, 0x56, 0x2d, 0xbd, 0xd1, 0xb1, 0xa2, 0xd9, 0x87, 0x99, 0x9c, 0x37, 0xf1, 0x59,
        0xbf, 0x18
    ])
    .unwrap();
}

#[fbinit::test]
async fn filestore_put_alias(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;
    let res = filestore::get_metadata(blob, ctx, &FetchKey::Canonical(content_id)).await?;

    println!("res = {:#?}", res);

    assert_eq!(
        res,
        Some(ContentMetadataV2 {
            content_id,
            total_size: HELLO_WORLD_LENGTH,
            sha1: *HELLO_WORLD_SHA1,
            sha256: *HELLO_WORLD_SHA256,
            git_sha1: *HELLO_WORLD_GIT_SHA1,
            seeded_blake3: *HELLO_WORLD_SEEDED_BLAKE3,
            is_binary: false,
            is_ascii: true,
            is_utf8: true,
            ends_in_newline: false,
            newline_count: 0,
            first_line: String::from_utf8(HELLO_WORLD.to_vec()).ok(),
            is_generated: false,
            is_partially_generated: false,
        })
    );

    Ok(())
}

#[fbinit::test]
async fn filestore_put_get_canon(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(blob, ctx, &FetchKey::Canonical(content_id)).await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));

    Ok(())
}

#[fbinit::test]
async fn filestore_put_get_sha1(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::Sha1(*HELLO_WORLD_SHA1)),
    )
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_put_get_git_sha1(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::GitSha1(HELLO_WORLD_GIT_SHA1.sha1())),
    )
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_put_get_sha256(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::Sha256(*HELLO_WORLD_SHA256)),
    )
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_put_get_seeded_blake3(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::SeededBlake3(*HELLO_WORLD_SEEDED_BLAKE3)),
    )
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_chunked_put_get(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };

    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_concat_opt(blob, ctx, &FetchKey::Canonical(content_id)).await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_chunked_put_get_nested(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();

    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let large = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);

    let part_data = &b"foo"[..];
    let part_key = request(part_data);
    borrowed!(ctx, blob, full_key, part_key);

    // Store in 3-byte chunks
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    // Now, go and split up one chunk into 1-byte parts.
    filestore::store(
        blob,
        small,
        ctx,
        part_key,
        stream::once(future::ready(Ok(Bytes::from(part_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foo", "bar"]).await?;
    Ok(())
}

#[fbinit::test]
async fn filestore_content_not_found(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob);

    // Missing content shouldn't throw an error

    let data = &b"foobar"[..];
    let content_id = canonical(data);

    // Verify that we can still read the full thing.
    let res = filestore::fetch_concat_opt(blob, ctx, &FetchKey::Canonical(content_id)).await;

    println!("res = {:#?}", res);
    assert_eq!(res?, None);
    Ok(())
}

#[fbinit::test]
async fn filestore_chunk_not_found(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let data = &b"foobar"[..];
    let req = request(data);
    let content_id = canonical(data);
    borrowed!(ctx, blob, req);

    let part = &b"foo"[..];
    let part_id = chunk(part);

    filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(data)))),
    )
    .await?;

    assert!(
        blob.unlink(part_id.blobstore_key())
            .await
            .unwrap()
            .is_some()
    );

    // This should fail
    let res = filestore::fetch_concat_opt(&blob, ctx, &FetchKey::Canonical(content_id)).await;

    println!("res = {:#?}", res);
    assert!(res.is_err());
    Ok(())
}

#[fbinit::test]
async fn filestore_put_invalid_size(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let data = &b"foobar"[..];
    let req = StoreRequest::new(123);
    borrowed!(ctx, blob, req);

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(data)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidSize(..))
    );
    Ok(())
}

#[fbinit::test]
async fn filestore_put_content_id(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    // Bad Content Id should fail
    let req = StoreRequest::with_canonical(HELLO_WORLD_LENGTH, ONES_CTID);
    borrowed!(ctx, blob, req);

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidContentId(..))
    );

    // Correct content Id should succeed
    let req = StoreRequest::with_canonical(HELLO_WORLD_LENGTH, canonical(HELLO_WORLD));
    let res = filestore::store(
        blob,
        config,
        ctx,
        &req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert!(res.is_ok());

    Ok(())
}

#[fbinit::test]
async fn filestore_put_sha1(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    // Bad Content Id should fail
    let req = StoreRequest::with_sha1(HELLO_WORLD_LENGTH, hash::Sha1::from_byte_array([0x00; 20]));
    borrowed!(ctx, blob, req);
    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidSha1(..))
    );

    // Correct content Id should succeed
    let req = StoreRequest::with_sha1(HELLO_WORLD_LENGTH, *HELLO_WORLD_SHA1);
    let res = filestore::store(
        blob,
        config,
        ctx,
        &req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert!(res.is_ok());

    Ok(())
}

#[fbinit::test]
async fn filestore_put_git_sha1(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    // Bad Content Id should fail
    let req = StoreRequest::with_git_sha1(
        HELLO_WORLD_LENGTH,
        hash::RichGitSha1::from_byte_array([0x00; 20], "blob", HELLO_WORLD_LENGTH),
    );
    borrowed!(ctx, blob, req);

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidGitSha1(..))
    );

    // Correct content Id should succeed
    let req = StoreRequest::with_git_sha1(HELLO_WORLD_LENGTH, *HELLO_WORLD_GIT_SHA1);
    let res = filestore::store(
        blob,
        config,
        ctx,
        &req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert!(res.is_ok());

    Ok(())
}

#[fbinit::test]
async fn filestore_put_sha256(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    // Bad Content Id should fail
    let req = StoreRequest::with_sha256(
        HELLO_WORLD_LENGTH,
        hash::Sha256::from_byte_array([0x00; 32]),
    );
    borrowed!(ctx, blob, req);

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidSha256(..))
    );

    // Correct content Id should succeed
    let req = StoreRequest::with_sha256(HELLO_WORLD_LENGTH, *HELLO_WORLD_SHA256);
    let res = filestore::store(
        blob,
        config,
        ctx,
        &req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert!(res.is_ok());

    Ok(())
}

#[fbinit::test]
async fn filestore_put_seeded_blake3(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    // Bad Content Id should fail
    let req = StoreRequest::with_seeded_blake3(
        HELLO_WORLD_LENGTH,
        hash::Blake3::from_byte_array([0x00; 32]),
    );
    borrowed!(ctx, blob, req);

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<errors::ErrorKind>(),
        Ok(errors::ErrorKind::InvalidBlake3(..))
    );

    // Correct content Id should succeed
    let req = StoreRequest::with_seeded_blake3(HELLO_WORLD_LENGTH, *HELLO_WORLD_SEEDED_BLAKE3);
    let res = filestore::store(
        blob,
        config,
        ctx,
        &req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;
    println!("res = {:#?}", res);
    assert!(res.is_ok());

    Ok(())
}

#[fbinit::test]
async fn filestore_get_range(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = async {
        let stream = filestore::fetch_range_with_size(
            &blob,
            ctx,
            &FetchKey::Canonical(content_id),
            filestore::Range::sized(7, 5),
        )
        .await?
        .ok_or_else(|| Error::msg("Object does not exist"))?
        .0;

        let bytes = stream
            .try_fold(BytesMut::new(), |mut buff, chunk| async move {
                buff.extend_from_slice(&chunk);
                Result::<_, Error>::Ok(buff)
            })
            .await?
            .freeze();

        Result::<_, Error>::Ok(bytes)
    }
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Bytes::from(&HELLO_WORLD[7..]));

    Ok(())
}

#[fbinit::test]
async fn filestore_get_invalid_range(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_range_with_size(
        &blob,
        ctx,
        &FetchKey::Canonical(content_id),
        filestore::Range::sized(0, 40),
    )
    .await;

    assert!(res.is_ok());

    let res = filestore::fetch_range_with_size(
        &blob,
        ctx,
        &FetchKey::Canonical(content_id),
        filestore::Range::sized(0, HELLO_WORLD.len() as u64).strict(),
    )
    .await;

    assert!(res.is_ok());

    let res = filestore::fetch_range_with_size(
        &blob,
        ctx,
        &FetchKey::Canonical(content_id),
        filestore::Range::sized(0, 40).strict(),
    )
    .await;

    assert!(res.is_err());

    Ok(())
}

#[fbinit::test]
async fn filestore_get_chunked_range(fb: FacebookInit) -> Result<()> {
    let small = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobarbazquxxyz"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Store in 3-byte chunks
    filestore::store(
        blob,
        small,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    // Check that we get the data we expect (6 bytes starting from the one at offset 4, i.e. the
    // 5th one).
    let res = async {
        let stream = filestore::fetch_range_with_size(
            blob,
            ctx,
            &FetchKey::Canonical(full_id),
            filestore::Range::sized(4, 6),
        )
        .await?
        .ok_or_else(|| Error::msg("Object does not exist"))?
        .0;

        let bytes = stream
            .try_fold(BytesMut::new(), |mut buff, chunk| async move {
                buff.extend_from_slice(&chunk);
                Result::<_, Error>::Ok(buff)
            })
            .await?
            .freeze();

        Result::<_, Error>::Ok(bytes)
    }
    .await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Bytes::from(&b"arbazq"[..]));

    // Check that we don't fetch things we do not need (extra chunks to the left).
    let res = async {
        let stream = filestore::fetch_range_with_size(
            blob,
            ctx,
            &FetchKey::Canonical(full_id),
            filestore::Range::sized(3, 2),
        )
        .await?
        .ok_or_else(|| Error::msg("Object does not exist"))?
        .0;

        stream.try_collect::<Vec<_>>().await
    }
    .await;

    println!("res = {:#?}", res);

    let res = res?;
    assert_eq!(res.len(), 1);
    assert_eq!(res[0], Bytes::from(&b"ba"[..]));

    // Check that we don't fetch things we do not need (extra chunks to the right).
    let res = async {
        let stream = filestore::fetch_range_with_size(
            blob,
            ctx,
            &FetchKey::Canonical(full_id),
            filestore::Range::sized(0, 3),
        )
        .await?
        .ok_or_else(|| Error::msg("Object does not exist"))?
        .0;

        stream.try_collect::<Vec<_>>().await
    }
    .await;

    println!("res = {:#?}", res);

    let res = res?;
    assert_eq!(res.len(), 1);
    assert_eq!(res[0], Bytes::from(&b"foo"[..]));

    Ok(())
}

#[fbinit::test]
async fn filestore_rebuild_metadata(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);
    let metadata: ContentMetadataV2Id = content_id.clone().into();

    let expected = Some(ContentMetadataV2 {
        total_size: HELLO_WORLD_LENGTH,
        content_id,
        sha1: *HELLO_WORLD_SHA1,
        git_sha1: *HELLO_WORLD_GIT_SHA1,
        sha256: *HELLO_WORLD_SHA256,
        seeded_blake3: *HELLO_WORLD_SEEDED_BLAKE3,
        is_binary: false,
        is_ascii: true,
        is_utf8: true,
        ends_in_newline: false,
        newline_count: 0,
        first_line: String::from_utf8(HELLO_WORLD.to_vec()).ok(),
        is_generated: false,
        is_partially_generated: false,
    });

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    // Remove the metadata
    assert!(
        blob.unlink(metadata.blobstore_key())
            .await
            .unwrap()
            .is_some()
    );

    // Getting the metadata should cause it to get recomputed
    let res = filestore::get_metadata(blob, ctx, &FetchKey::Canonical(content_id)).await?;
    println!("res = {:#?}", res);
    assert_eq!(res, expected);

    // Now, delete the content (this shouldn't normally happen, but we're injecting failure here).
    assert!(
        blob.unlink(content_id.blobstore_key())
            .await
            .unwrap()
            .is_some()
    );

    // Query the metadata again. It should succeed because it's saved.
    let res = filestore::get_metadata(blob, ctx, &FetchKey::Canonical(content_id)).await?;
    println!("res = {:#?}", res);
    assert_eq!(res, expected);

    // Delete the metadata now.
    assert!(
        blob.unlink(metadata.blobstore_key())
            .await
            .unwrap()
            .is_some()
    );

    // And then, query it again. This should now return None, because the metadata isn't there,
    // and we can't recreate it.
    let res = filestore::get_metadata(blob, ctx, &FetchKey::Canonical(content_id)).await?;
    println!("res = {:#?}", res);
    assert_eq!(res, None);

    Ok(())
}

#[fbinit::test]
async fn filestore_test_missing_metadata(fb: FacebookInit) -> Result<()> {
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob);

    // No matter the Fetchkey, querying the metadata should return None.

    let res = filestore::get_metadata(blob, ctx, &FetchKey::Canonical(content_id)).await;
    println!("res = {:#?}", res);
    assert_eq!(res?, None);

    let res = filestore::get_metadata(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::Sha1(*HELLO_WORLD_SHA1)),
    )
    .await;
    println!("res = {:#?}", res);
    assert_eq!(res?, None);

    let res = filestore::get_metadata(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::Sha256(*HELLO_WORLD_SHA256)),
    )
    .await;
    println!("res = {:#?}", res);
    assert_eq!(res?, None);

    let res = filestore::get_metadata(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::SeededBlake3(*HELLO_WORLD_SEEDED_BLAKE3)),
    )
    .await;
    println!("res = {:#?}", res);
    assert_eq!(res?, None);

    let res = filestore::get_metadata(
        blob,
        ctx,
        &FetchKey::Aliased(Alias::GitSha1(HELLO_WORLD_GIT_SHA1.sha1())),
    )
    .await;
    println!("res = {:#?}", res);
    assert_eq!(res?, None);

    Ok(())
}

#[fbinit::test]
async fn filestore_test_peek(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::peek(blob, ctx, &FetchKey::Canonical(content_id), 3).await;
    println!("res = {:#?}", res);

    let expected: &[u8] = b"hel";
    assert_eq!(res?, Some(Bytes::from(expected)));

    Ok(())
}

#[fbinit::test]
async fn filestore_test_chunked_peek(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        small,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::peek(blob, ctx, &FetchKey::Canonical(content_id), 3).await;
    println!("res = {:#?}", res);

    let expected: &[u8] = b"hel";
    assert_eq!(res?, Some(Bytes::from(expected)));

    Ok(())
}

#[fbinit::test]
async fn filestore_test_short_peek(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::peek(blob, ctx, &FetchKey::Canonical(content_id), 128).await;
    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));

    Ok(())
}

#[fbinit::test]
async fn filestore_test_empty_peek(fb: FacebookInit) -> Result<()> {
    let bytes = Bytes::new();

    let req = request(&bytes);
    let content_id = canonical(&bytes);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(bytes.clone()))),
    )
    .await?;

    let res = filestore::peek(blob, ctx, &FetchKey::Canonical(content_id), 128).await;
    println!("res = {:#?}", res);

    assert_eq!(res?, Some(bytes.clone()));

    Ok(())
}

#[fbinit::test]
async fn filestore_store_bytes(fb: FacebookInit) -> Result<()> {
    let blob: memblob::Memblob = memblob::Memblob::default();

    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob);
    let ((content_id, _size), fut) =
        filestore::store_bytes(blob, DEFAULT_CONFIG, ctx, Bytes::from(HELLO_WORLD));
    assert_eq!(content_id, canonical(HELLO_WORLD));

    fut.await?;

    let res = filestore::fetch_concat_opt(blob, ctx, &FetchKey::Canonical(content_id)).await;

    println!("res = {:#?}", res);

    assert_eq!(res?, Some(Bytes::from(HELLO_WORLD)));
    Ok(())
}

#[fbinit::test]
async fn filestore_store_error(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();
    let blob = FailingBlobstore::new(blob, 1.0, 0.0); // Blobstore you can't write to.
    let ctx = CoreContext::test_mock(fb);
    let req = request(HELLO_WORLD);
    borrowed!(ctx, blob, req);

    let config = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };

    let res = filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await;

    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<FailingBlobstoreError>(),
        Ok(FailingBlobstoreError)
    );
    Ok(())
}

#[fbinit::test]
async fn filestore_test_rechunk(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::new(PutBehaviour::Overwrite);

    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let large = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Store in 3-byte chunks
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foo", "bar"]).await?;

    // Rechunk the file into 1 byte sections
    filestore::rechunk::force_rechunk(blob, small, ctx, full_id).await?;

    assert_fetches_as(ctx, blob, full_id, vec!["f", "o", "o", "b", "a", "r"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_larger(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::new(PutBehaviour::Overwrite);

    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let large = FilestoreConfig {
        chunk_size: Some(3),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Store in 1 byte chunks
    filestore::store(
        blob,
        small,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["f", "o", "o", "b", "a", "r"]).await?;

    // Rechunk the file into 3 byte sections
    filestore::rechunk::force_rechunk(blob, large, ctx, full_id).await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foo", "bar"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_unchunked(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::new(PutBehaviour::Overwrite);

    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    // This is large enough that the data we upload won't be chunked.
    let large = FilestoreConfig {
        chunk_size: Some(100),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Don't chunk
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    // Rechunk the file into 1 byte sections
    filestore::rechunk::force_rechunk(blob, small, ctx, full_id).await?;

    assert_fetches_as(ctx, blob, full_id, vec!["f", "o", "o", "b", "a", "r"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_missing_content(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();

    let conf = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob);

    let full_data = &b"foobar"[..];
    let full_id = canonical(full_data);

    // Attempt to rechunk the file into 1 byte sections
    let res = filestore::rechunk::force_rechunk(blob, conf, ctx, full_id).await;

    println!("res = {:#?}", res);
    assert_matches!(
        res.unwrap_err().downcast::<filestore::rechunk::ErrorKind>(),
        Ok(filestore::rechunk::ErrorKind::ContentNotFound(..))
    );

    Ok(())
}

#[fbinit::test]
async fn filestore_chunked_put_get_with_size(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);

    let blob = memblob::Memblob::default();
    let config = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };

    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    filestore::store(
        blob,
        config,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    let res = filestore::fetch_with_size(blob, ctx, &FetchKey::Canonical(content_id)).await;

    let (stream, size) = res?.unwrap();

    let fut = stream
        .try_fold(BytesMut::new(), |mut buff, chunk| async move {
            buff.extend_from_slice(&chunk);
            Result::<_, Error>::Ok(buff)
        })
        .map_ok(BytesMut::freeze);

    let bytes = fut.await;

    println!("{:?}", bytes);

    assert_eq!(bytes?, Bytes::from(HELLO_WORLD));
    assert_eq!(size, HELLO_WORLD_LENGTH);
    Ok(())
}

#[fbinit::test]
/// Test a case, where both old and new filestore config do not require
/// chunking of the file (e.g. file size is smaller than a single chunk)
async fn filestore_test_rechunk_if_needed_tiny_unchunked_file(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();

    let large1 = FilestoreConfig {
        chunk_size: Some(100),
        concurrency: 5,
    };
    let large2 = FilestoreConfig {
        chunk_size: Some(200),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Don't chunk
    filestore::store(
        blob,
        large1,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foobar"]).await?;

    // We expect that rechunk is not needed
    filestore::rechunk::rechunk(
        &FailingBlobstore::new(blob.clone(), 1.0, 0.0),
        large2,
        ctx,
        full_id,
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foobar"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_if_needed_large_unchunked_file(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::new(PutBehaviour::Overwrite);

    let large = FilestoreConfig {
        chunk_size: Some(100),
        concurrency: 5,
    };
    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Don't chunk
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["foobar"]).await?;

    // We expect the rechunk is needed
    let (_, rechunked) = filestore::rechunk::rechunk(blob, small, ctx, full_id).await?;
    assert!(rechunked);

    assert_fetches_as(ctx, blob, full_id, vec!["f", "o", "o", "b", "a", "r"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_if_needed_large_chunks(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::new(PutBehaviour::Overwrite);

    let large = FilestoreConfig {
        chunk_size: Some(5),
        concurrency: 5,
    };
    let small = FilestoreConfig {
        chunk_size: Some(1),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Chunk with larger chunks
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;

    assert_fetches_as(ctx, blob, full_id, vec!["fooba", "r"]).await?;

    // We expect the rechunk is needed
    let (_, rechunked) = filestore::rechunk::rechunk(blob, small, ctx, full_id).await?;
    assert!(rechunked);

    assert_fetches_as(ctx, blob, full_id, vec!["f", "o", "o", "b", "a", "r"]).await
}

#[fbinit::test]
async fn filestore_test_rechunk_if_needed_tiny_chunks(fb: FacebookInit) -> Result<()> {
    let blob = memblob::Memblob::default();

    let large = FilestoreConfig {
        chunk_size: Some(4),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);

    let full_data = &b"foobar"[..];
    let full_key = request(full_data);
    let full_id = canonical(full_data);
    borrowed!(ctx, blob, full_key);

    // Chunk
    filestore::store(
        blob,
        large,
        ctx,
        full_key,
        stream::once(future::ready(Ok(Bytes::from(full_data)))),
    )
    .await?;
    assert_fetches_as(ctx, blob, full_id, vec!["foob", "ar"]).await?;

    // We expect the rechunk is not needed
    let (_, rechunked) = filestore::rechunk::rechunk(
        &FailingBlobstore::new(blob.clone(), 1.0, 0.0),
        large,
        ctx,
        full_id,
    )
    .await?;
    assert!(!rechunked);
    assert_fetches_as(ctx, blob, full_id, vec!["foob", "ar"]).await
}

async fn assert_fetches_as<B: Blobstore, S: Into<Bytes>>(
    ctx: &CoreContext,
    blobstore: &B,
    content_id: ContentId,
    expected: Vec<S>,
) -> Result<()> {
    let expected = expected.into_iter().map(|s| s.into()).collect();
    let key = FetchKey::Canonical(content_id);
    let maybe_stream = filestore::fetch(blobstore, ctx, &key).await?;
    let res = match maybe_stream {
        Some(stream) => Some(stream.try_collect::<Vec<_>>().await?),
        None => None,
    };

    println!("res = {:#?}", res);
    assert_eq!(res, Some(expected));
    Ok(())
}

#[fbinit::test]
async fn filestore_exists(fb: FacebookInit) -> Result<()> {
    let req = request(HELLO_WORLD);
    let content_id = canonical(HELLO_WORLD);
    let alias = Alias::Sha1(*HELLO_WORLD_SHA1);

    let blob = memblob::Memblob::default();
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, req);

    assert!(!filestore::exists(blob, ctx, &content_id.into()).await?);
    assert!(!filestore::exists(blob, ctx, &alias.into()).await?);

    filestore::store(
        blob,
        DEFAULT_CONFIG,
        ctx,
        req,
        stream::once(future::ready(Ok(Bytes::from(HELLO_WORLD)))),
    )
    .await?;

    assert!(filestore::exists(blob, ctx, &content_id.into()).await?);
    assert!(filestore::exists(blob, ctx, &alias.into()).await?);

    Ok(())
}


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

use std::collections::HashSet;
use std::sync::Arc;

use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use borrowed::borrowed;
use bytes::Bytes;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::future;
use futures::future::TryFutureExt;
use futures::stream;
use quickcheck::Arbitrary;
use quickcheck::Gen;

use super::failing_blobstore::FailingBlobstore;
use super::request;
use crate as filestore;
use crate::incremental_hash::hash_bytes;
use crate::incremental_hash::ContentIdIncrementalHasher;
use crate::incremental_hash::GitSha1IncrementalHasher;
use crate::incremental_hash::Sha1IncrementalHasher;
use crate::incremental_hash::Sha256IncrementalHasher;
use crate::Alias;
use crate::FetchKey;
use crate::FilestoreConfig;

/// Fetching through any alias should return the same outcome.
async fn check_consistency<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    bytes: &Bytes,
) -> Result<bool, Error> {
    let content_id = hash_bytes(ContentIdIncrementalHasher::new(), bytes);
    let sha1 = hash_bytes(Sha1IncrementalHasher::new(), bytes);
    let sha256 = hash_bytes(Sha256IncrementalHasher::new(), bytes);
    let git_sha1 = hash_bytes(GitSha1IncrementalHasher::new(bytes), bytes);

    let content_id = FetchKey::Canonical(content_id);
    let sha1 = FetchKey::Aliased(Alias::Sha1(sha1));
    let sha256 = FetchKey::Aliased(Alias::Sha256(sha256));
    let git_sha1 = FetchKey::Aliased(Alias::GitSha1(git_sha1.sha1()));

    let futs = vec![
        filestore::fetch(blobstore, ctx, &content_id),
        filestore::fetch(blobstore, ctx, &sha1),
        filestore::fetch(blobstore, ctx, &sha256),
        filestore::fetch(blobstore, ctx, &git_sha1),
    ];

    let futs: Vec<_> = futs
        .into_iter()
        .map(|f| f.map_ok(|r| r.is_some()))
        .collect();

    let outcomes = future::try_join_all(futs).await?;

    // Either all should exist, or none should exist.
    let h: HashSet<_> = outcomes.iter().collect();
    if h.len() == 1 {
        Ok(*h.into_iter().next().unwrap())
    } else {
        Err(format_err!("Inconsistent fetch results: {:?}", outcomes))
    }
}

async fn check_metadata<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    bytes: &Bytes,
) -> Result<bool, Error> {
    let content_id = hash_bytes(ContentIdIncrementalHasher::new(), bytes);

    filestore::get_metadata(blobstore, ctx, &FetchKey::Canonical(content_id))
        .await
        .map(|r| r.is_some())
}

#[fbinit::test]
fn test_invariants(fb: FacebookInit) -> Result<()> {
    // NOTE: We make calls to our Blobstore succeed with 75% probability below. This might seem
    // high, but this actually makes most store() calls fail, since there is a lot that needs to go
    // right for a store() call to succeed (all the chunks need to be saved, then we need to write
    // 3 aliases, and then the content).
    let rt = tokio::runtime::Runtime::new()?;
    let mut gen = Gen::new(128);

    let memblob = Arc::new(memblob::Memblob::default());
    let blob = FailingBlobstore::new(memblob.clone(), 0.75, 0.75);
    let config = FilestoreConfig {
        chunk_size: Some(16),
        concurrency: 5,
    };
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, blob, memblob: &Arc<_>);

    for _ in 0..1000 {
        let bytes = Bytes::from(Vec::arbitrary(&mut gen));
        let req = request(&bytes);

        // Try to store with a broken blobstore. It doesn't matter if we succeed or not.
        let res = rt.block_on(filestore::store(
            blob,
            config,
            ctx,
            &req,
            stream::once(future::ready(Ok(bytes.clone()))),
        ));
        println!("store: {:?}", res);

        // Try to read with a functional blobstore. All results should be consistent.
        let content_ok = rt.block_on(check_consistency(memblob, ctx, &bytes))?;
        println!("content_ok: {:?}", content_ok);

        // If we can read the content metadata, then we should also be able to read a metadata.
        let metadata_ok = rt.block_on(check_metadata(memblob, ctx, &bytes))?;
        println!("metadata_ok: {:?}", metadata_ok);
        assert_eq!(content_ok, metadata_ok)
    }

    Ok(())
}

#[fbinit::test]
async fn test_store_bytes_consistency(fb: FacebookInit) -> Result<(), Error> {
    let mut gen = Gen::new(128);

    let memblob = Arc::new(memblob::Memblob::default());
    let ctx = CoreContext::test_mock(fb);
    borrowed!(ctx, memblob: &Arc<_>);

    for _ in 0..100usize {
        let bytes = Bytes::from(Vec::arbitrary(&mut gen));

        let no_chunking = FilestoreConfig {
            chunk_size: None,
            concurrency: 1,
        };

        let chunked = FilestoreConfig {
            chunk_size: Some(std::cmp::max(1, (bytes.len() as u64) / 2)),
            concurrency: 1,
        };

        let too_small_to_chunk = FilestoreConfig {
            chunk_size: Some(std::cmp::max(1, (bytes.len() as u64) * 2)),
            concurrency: 1,
        };

        let ((id1, len1), fut1) = filestore::store_bytes(memblob, no_chunking, ctx, bytes.clone());
        fut1.await?;

        assert_eq!(bytes, filestore::fetch_concat(memblob, ctx, id1).await?);

        let ((id2, len2), fut2) = filestore::store_bytes(memblob, chunked, ctx, bytes.clone());
        fut2.await?;

        assert_eq!(bytes, filestore::fetch_concat(memblob, ctx, id2).await?);

        let ((id3, len3), fut3) =
            filestore::store_bytes(memblob, too_small_to_chunk, ctx, bytes.clone());
        fut3.await?;

        assert_eq!(bytes, filestore::fetch_concat(memblob, ctx, id3).await?);

        let meta = filestore::store(
            memblob,
            no_chunking,
            ctx,
            &request(&bytes),
            stream::once(future::ready(Ok(bytes.clone()))),
        )
        .await?;

        assert_eq!(meta.content_id, id1);
        assert_eq!(meta.content_id, id2);
        assert_eq!(meta.content_id, id3);

        assert_eq!(meta.total_size, len1);
        assert_eq!(meta.total_size, len2);
        assert_eq!(meta.total_size, len3);
    }

    Result::<_, Error>::Ok(())
}


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

use std::collections::HashMap;
use std::str::FromStr;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Result;
use git2::ObjectType;
use git2::Oid;
use git2::Repository;
use git2::Tree;
use git2::TreeEntry;
use mononoke_types::hash;
use mononoke_types::FileType;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use mononoke_types::RepoPath;
use tokio::sync::mpsc;

use crate::CheckEntry;
use crate::CheckNode;

fn get_sha256(git_lfs: bool, contents: &[u8]) -> hash::Sha256 {
    if git_lfs
        && (contents.starts_with(b"version https://git-lfs.github.com/spec/v1\n")
            || contents.starts_with(b"version https://hawser.github.com/spec/v1\n"))
    {
        if let Ok(pointer) = std::str::from_utf8(contents) {
            // This is a plausible pointer - it's UTF-8, it begins with a version line
            // Skip the version, and look for the embedded hash
            for line in pointer.lines().skip(1) {
                if line > "oid sha256:g" {
                    // Lines apart from version must be in ASCII order. If we get past the last
                    // possible place for our OID, then we should break out.
                    // That way, a file that starts with a known version line but then has random
                    // junk is not going to consume our time
                    break;
                }
                if let Some(hash) = line.strip_prefix("oid sha256:") {
                    // Possible SHA256 for content. Extract if possible
                    if let Ok(hash) = hash::Sha256::from_str(hash) {
                        return hash;
                    } else {
                        // Not a valid SHA256, so not actually a pointer file
                        break;
                    }
                }
            }
        }
    }
    use sha2::Digest;
    use sha2::Sha256;
    let mut hasher = Sha256::new();
    hasher.update(contents);
    hash::Sha256::from_byte_array(hasher.finalize().into())
}

fn process_entry(git_lfs: bool, repo: &Repository, entry: TreeEntry<'_>) -> Result<CheckEntry> {
    let filemode = entry.filemode();

    match entry.kind() {
        Some(ObjectType::Blob) => {
            let obj = entry.to_object(repo)?;
            let blob = obj.peel_to_blob()?;
            let hash = get_sha256(git_lfs, blob.content());
            let filetype = match filemode {
                0o100644 => FileType::Regular,
                0o100755 => FileType::Executable,
                0o120000 => FileType::Symlink,
                _ => bail!("Unknown git filemode {}", filemode),
            };

            Ok(CheckEntry::File(filetype, hash))
        }
        Some(ObjectType::Tree) => Ok(CheckEntry::Directory),
        Some(ObjectType::Commit) => {
            // A commit in a tree is a submodule. Treat as empty directory - this is what
            // `git checkout` does when submodules aren't enabled
            Ok(CheckEntry::Directory)
        }
        kind => Err(anyhow!("Object is of unexpected kind {:?}", kind)),
    }
}

fn walk_tree(
    git_lfs: bool,
    path: RepoPath,
    tx: &mpsc::Sender<CheckNode>,
    repo: &Repository,
    tree: Tree<'_>,
) -> Result<()> {
    let contents: HashMap<MPathElement, CheckEntry> = tree
        .iter()
        .map(|entry| {
            let path_element = MPathElement::new(entry.name_bytes().to_vec())?;
            let check_entry = process_entry(git_lfs, repo, entry)?;
            Ok((path_element, check_entry))
        })
        .collect::<Result<HashMap<_, _>>>()?;

    // Send this entry down for checking, then recurse into children. The checking task depends
    // on getting parents before children in order to avoid doing excess fsnode fetches
    tx.blocking_send(CheckNode {
        path: path.clone(),
        contents,
    })?;

    for entry in tree.iter() {
        if entry.kind() == Some(ObjectType::Tree) {
            let path_element = MPathElement::new(entry.name_bytes().to_vec())?;
            let obj = entry.to_object(repo)?;
            let subtree_path = RepoPath::dir(MPath::join_opt_element(path.mpath(), &path_element))?;
            let tree = obj.peel_to_tree()?;

            walk_tree(git_lfs, subtree_path, tx, repo, tree)?;
        }
    }

    Ok(())
}

/// Spawn a thread that reads the commit, and outputs a chain of `CheckNode`s to a channel
pub(crate) fn thread(
    repo: Repository,
    commit: String,
    git_lfs: bool,
    tx: mpsc::Sender<CheckNode>,
) -> Result<()> {
    let git_commit = repo.find_commit(Oid::from_str(&commit)?)?;

    let root = git_commit.tree()?;

    walk_tree(git_lfs, RepoPath::root(), &tx, &repo, root)
}

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

    // The content used in the LFS pointer
    const CONTENT: &[u8] = b"hello\n";
    // This is a real LFS pointer from `git lfs pointer`
    const LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v1\noid sha256:5891b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03\nsize 6\n";
    // Damaged the hash - it's not a genuine SHA-256 hash
    const BROKEN_HASH_LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v1\noid sha256:5MOO91b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03\nsize 6\n";
    // No hash in the LFS pointer
    const NO_HASH_LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v1\n\nsize 6\n";
    // Not UTF-8
    const NOT_UTF8_LFS_POINTER: &[u8] =
        b"version https://git-lfs.github.com/spec/v1\n\x91a\nsize 6\n";
    // Bad LFS version
    const BAD_VERSION_LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v2\noid sha256:5MOO91b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03\nsize 6\n";
    // Unknown hash type
    const BAD_HASH_TYPE_LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v1\noid sha2-256:5891b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03\nsize 6\n";

    #[test]
    fn test_lfs_pointer() {
        // No LFS, hash of the pointer
        assert_eq!(
            get_sha256(false, LFS_POINTER),
            hash::Sha256::from_str(
                "94cb9a4fb124ed218aeeaefa7927680d5a261652f400f9d4f6a4e729c995d088",
            )
            .unwrap(),
        );
        // LFS, hash extracted from the pointer
        assert_eq!(
            get_sha256(true, LFS_POINTER),
            hash::Sha256::from_str(
                "5891b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03",
            )
            .unwrap(),
        );
        // And check the LFS hash matches the real content
        assert_eq!(get_sha256(true, LFS_POINTER), get_sha256(true, CONTENT));
    }

    #[test]
    fn test_bad_pointers() {
        // All of these are bad pointers of some form, so we should get their hash, not extract a hash
        // from the pointer
        assert_eq!(
            get_sha256(true, BROKEN_HASH_LFS_POINTER),
            hash::Sha256::from_str(
                "8de25c66642bdcfed8e6b5f84dff6eda5ae6054a36c909a52c0ba8208d50ea09",
            )
            .unwrap(),
        );
        assert_ne!(
            get_sha256(true, BROKEN_HASH_LFS_POINTER),
            get_sha256(true, CONTENT)
        );
        assert_eq!(
            get_sha256(true, NO_HASH_LFS_POINTER),
            hash::Sha256::from_str(
                "e326ffde029b1b9101bbb00ca3da8f6e526685b5b1296e46b446a779e7c847a8",
            )
            .unwrap(),
        );
        assert_ne!(
            get_sha256(true, NO_HASH_LFS_POINTER),
            get_sha256(true, CONTENT)
        );
        assert_eq!(
            get_sha256(true, NOT_UTF8_LFS_POINTER),
            hash::Sha256::from_str(
                "328395baea3e9b372ac8040b66f2607ca208668623561d5da35a124b538a1e03",
            )
            .unwrap(),
        );
        assert_ne!(
            get_sha256(true, NOT_UTF8_LFS_POINTER),
            get_sha256(true, CONTENT)
        );
        assert_eq!(
            get_sha256(true, BAD_VERSION_LFS_POINTER),
            hash::Sha256::from_str(
                "3a2d968718c0ba448b90f2288985cbe158128e8d360505bf3f39eefd1f46b382",
            )
            .unwrap(),
        );
        assert_ne!(
            get_sha256(true, BAD_VERSION_LFS_POINTER),
            get_sha256(true, CONTENT)
        );
        assert_eq!(
            get_sha256(true, BAD_HASH_TYPE_LFS_POINTER),
            hash::Sha256::from_str(
                "37fb1ff2a48362f5c4d64c1b9121875a245a96f2787db5dd2f5c3273aabc4cb3",
            )
            .unwrap(),
        );
        assert_ne!(
            get_sha256(true, BAD_HASH_TYPE_LFS_POINTER),
            get_sha256(true, CONTENT)
        );
    }
}


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

#![feature(trait_alias)]

use std::collections::HashMap;
use std::sync::Mutex;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Result;
use blobstore::Loadable;
use context::CoreContext;
use derived_data::BonsaiDerived;
use fsnodes::RootFsnodeId;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use git2::Repository;
use mononoke_types::fsnode::FsnodeEntry;
use mononoke_types::hash;
use mononoke_types::ChangesetId;
use mononoke_types::FileType;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use mononoke_types::RepoPath;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataRef;
use sorted_vector_map::SortedVectorMap;
use tokio::sync::mpsc;

mod git_walker;

pub trait HgRepo = RepoBlobstoreRef + RepoDerivedDataRef + Send + Sync;

#[derive(Debug)]
pub(crate) enum CheckEntry {
    Directory,
    File(FileType, hash::Sha256),
}

#[derive(Debug)]
pub(crate) struct CheckNode {
    pub path: RepoPath,
    pub contents: HashMap<MPathElement, CheckEntry>,
}

async fn check_node(
    node: CheckNode,
    entries: SortedVectorMap<MPathElement, FsnodeEntry>,
) -> Result<()> {
    let CheckNode { path, mut contents } = node;

    // Check that each Mononoke entry is the same as its matching git entry
    // and remove it from contents once checked
    for (filename, fsnode_entry) in entries {
        let git_entry = contents
            .remove(&filename)
            .ok_or_else(|| anyhow!("File {}/{} in Bonsai but not git", path, filename))?;
        match git_entry {
            CheckEntry::Directory => match fsnode_entry {
                FsnodeEntry::File(_) => {
                    let entry_path =
                        RepoPath::dir(MPath::join_opt_element(path.mpath(), &filename))?;
                    bail!("{} is a file in Mononoke", entry_path);
                }
                FsnodeEntry::Directory(_) => {}
            },
            CheckEntry::File(filetype, sha256) => match fsnode_entry {
                FsnodeEntry::File(fsnode_file) => {
                    let entry_path =
                        RepoPath::file(MPath::join_opt_element(path.mpath(), &filename))?;
                    if *fsnode_file.file_type() != filetype {
                        bail!(
                            "{} is type {} in git and {} in Mononoke",
                            path,
                            filetype,
                            *fsnode_file.file_type()
                        );
                    }
                    if *fsnode_file.content_sha256() != sha256 {
                        bail!(
                            "{} has hash {} in git and {} in Mononoke",
                            entry_path,
                            sha256,
                            *fsnode_file.content_sha256()
                        );
                    }
                }
                FsnodeEntry::Directory(_) => {
                    let entry_path =
                        RepoPath::file(MPath::join_opt_element(path.mpath(), &filename))?;
                    bail!("{} is a directory in Mononoke", entry_path);
                }
            },
        }
    }

    // By this point, all the Mononoke entries have been checked
    if let Some((filename, _)) = contents.drain().next() {
        let entry_path = RepoPath::file(MPath::join_opt_element(path.mpath(), &filename))?;
        bail!("{} in git but not Bonsai", entry_path);
    }

    Ok(())
}

async fn check_receiver(
    ctx: &CoreContext,
    hg_repo: &impl HgRepo,
    cs: ChangesetId,
    rx: mpsc::Receiver<CheckNode>,
    scheduled_max: usize,
) -> Result<()> {
    let root_fsnode = RootFsnodeId::derive(ctx, hg_repo, cs)
        .await?
        .into_fsnode_id()
        .load(ctx, hg_repo.repo_blobstore())
        .await?;

    let path_to_fsnode = Mutex::new(HashMap::new());
    let path_to_fsnode = &path_to_fsnode;
    path_to_fsnode
        .lock()
        .expect("lock poisoned")
        .insert(None, Some(root_fsnode));
    let rx = tokio_stream::wrappers::ReceiverStream::new(rx);

    rx.map(Result::Ok)
        .and_then(move |node| async move {
            // This relies on git_walker doing a top-down traversal and only visiting each directory once
            let path = &node.path;
            let this_fsnode = path_to_fsnode
                .lock()
                .expect("lock poisoned")
                .remove(&path.mpath().cloned())
                .ok_or_else(|| anyhow!("{} not found in Mononoke", path))?
                .ok_or_else(|| anyhow!("{} in git is a file in Mononoke", path))?;
            // Fetch all children in parallel, storing them in the map
            stream::iter(this_fsnode.list().map(Result::Ok))
                .try_for_each_concurrent(scheduled_max, move |(element, entry)| async move {
                    let old_entry = {
                        let fsnode = match entry {
                            FsnodeEntry::File(_) => None,
                            FsnodeEntry::Directory(dir) => {
                                Some(dir.id().load(ctx, hg_repo.repo_blobstore()).await?)
                            }
                        };
                        path_to_fsnode
                            .lock()
                            .expect("lock poisoned")
                            .insert(Some(MPath::join_opt_element(path.mpath(), element)), fsnode)
                    };
                    if old_entry.is_some() {
                        Err(anyhow!("Two different routes to the same path?!?"))
                    } else {
                        Ok(())
                    }
                })
                .await?;

            // Pass on this node for comparison to git
            Ok((node, this_fsnode.into_subentries()))
        })
        .map_ok(|(node, fsnode)| check_node(node, fsnode))
        .try_for_each(|f| async move { tokio::spawn(f).await? })
        .await
}

pub async fn check_git_wc(
    ctx: &CoreContext,
    hg_repo: &impl HgRepo,
    cs: ChangesetId,
    git_repo: Repository,
    git_commit: String,
    git_lfs: bool,
    scheduled_max: usize,
) -> Result<()> {
    let (tx, rx) = mpsc::channel(100);
    let git_handle =
        tokio::task::spawn_blocking(move || git_walker::thread(git_repo, git_commit, git_lfs, tx));

    let res = check_receiver(ctx, hg_repo, cs, rx, scheduled_max).await;
    let git_res = git_handle.await?;
    match (res, git_res) {
        (Ok(()), Ok(())) => Ok(()),
        (Err(e), Ok(())) | (Ok(()), Err(e)) => Err(e),
        (Err(e1), Err(e2)) => Err(anyhow!(
            "Both git and Mononoke threads failed.\ngit: {}\nMononoke: {}\n",
            e2,
            e1
        )),
    }
}


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

#![feature(never_type)]
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Error;
use git2::Error as Git2Error;
use git2::Repository;
use r2d2::ManageConnection;
use r2d2::Pool;
use tokio::sync::Semaphore;

const POOL_CONNECTION_TIMEOUT_SEC: u64 = 3600;

#[derive(Debug, Clone)]
pub struct GitPool {
    pool: Pool<GitConnectionManager>,
    /// Semaphore to make sure we are not using to many "blocking" tokio tasks.
    /// Instead we wait in the calling task context before spawning.
    sem: Arc<Semaphore>,
}

impl GitPool {
    pub fn new(path: PathBuf) -> Result<Self, Error> {
        let simultaneous_open_repos = num_cpus::get();
        let manager = GitConnectionManager { path };
        let pool = Pool::builder()
            .max_size(simultaneous_open_repos as u32)
            .connection_timeout(Duration::from_secs(POOL_CONNECTION_TIMEOUT_SEC))
            .build(manager)?;
        Ok(Self {
            pool,
            sem: Arc::new(Semaphore::new(simultaneous_open_repos)),
        })
    }

    pub async fn with<F, T, E>(&self, f: F) -> Result<T, Error>
    where
        F: FnOnce(&Repository) -> Result<T, E> + Send + 'static,
        T: Send + Sync + 'static,
        E: Into<Error> + Send + Sync + 'static,
    {
        let sem = self.sem.clone();
        let pool = self.pool.clone();
        // Note - this tokio::spawn() is an attempt to fix deadlock D31541432.
        let ret = tokio::spawn(async move {
            let permit = sem.acquire_owned().await?;
            let ret = tokio_shim::task::spawn_blocking(move || {
                let result_repo = pool.get()?;
                let repo = match &*result_repo {
                    Ok(repo) => repo,
                    Err(err) => {
                        return Err(anyhow!("error while opening repo: {}", err));
                    }
                };
                let ret = f(repo).map_err(|e| e.into())?;
                drop(result_repo);
                drop(permit);
                Result::<_, Error>::Ok(ret)
            })
            .await??;
            Result::<_, Error>::Ok(ret)
        })
        .await??;

        Ok(ret)
    }
}

#[derive(Debug)]
struct GitConnectionManager {
    path: PathBuf,
}

impl ManageConnection for GitConnectionManager {
    type Connection = Result<Repository, Git2Error>;
    type Error = !;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        let repo = Repository::open(&self.path);
        Ok(repo)
    }

    fn is_valid(&self, _conn: &mut Self::Connection) -> Result<(), Self::Error> {
        Ok(())
    }

    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
        conn.is_err()
    }
}