azure_storage_blob 0.11.0

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

use std::{
    cmp::min,
    collections::VecDeque,
    ops::Range,
    pin::pin,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};

use async_trait::async_trait;
use azure_core::{
    async_runtime::{get_async_runtime, SpawnedTask},
    error::ErrorKind,
    http::{AsyncRawResponse, StatusCode},
    Error,
};
use bytes::Bytes;
use futures::{
    channel::mpsc::{self, UnboundedReceiver, UnboundedSender},
    future::{self, Either},
    stream, SinkExt, StreamExt,
};

use crate::models::{
    drains::{SequentialBoundedDrain, UnorderedFuturesDrain},
    http_ranges::ContentRange,
    response_ext::AsyncResponseBodyExt,
    slice::SendSlice,
};

use super::*;

#[async_trait]
pub(crate) trait PartitionedDownloadBehavior {
    async fn transfer_range(&self, range: Option<Range<usize>>) -> AzureResult<AsyncRawResponse>;
}

/// Returns a stream that runs up to parallel-many ranged downloads at a time.
///
/// Downloads are stored in-order. The returned stream will produce an item only when the next
/// download in the sequence has been buffered, regardless of the state of any other downloads.
/// This means completed ranged downloads may sit for a while while earlier ones complete.
///
/// This implementation makes an initial download request to gauge the actual size of the remote
/// resource while not wasting a roundtrip just for a HEAD request. It then determines the
/// correct set of additional ranges to download and queues them up. The returned `Stream`
/// executes these downloads, maintaining limits for parallel downloads and buffer count.
pub(crate) async fn download<Behavior>(
    range: Option<Range<usize>>,
    parallel: NonZero<usize>,
    partition_size: NonZero<usize>,
    client: Arc<Behavior>,
) -> AzureResult<AsyncRawResponse>
where
    Behavior: PartitionedDownloadBehavior + Send + Sync + 'static,
{
    let parallel = parallel.get();
    let max_buffers = parallel * 2;
    let partition_size = partition_size.get();

    let (initial_response, stats) =
        get_initial_response_and_analyze(range, partition_size, client.clone()).await?;

    let status = initial_response.status();
    let headers = initial_response.headers().clone();

    let mut remaining_ranges = stats
        .map(|s| s.remaining_download_ranges)
        .unwrap_or_default();
    if remaining_ranges.is_empty() {
        return Ok(AsyncRawResponse::new(
            status,
            headers,
            Box::pin(initial_response.into_body()),
        ));
    }
    let total_chunks = remaining_ranges.len() + 1;

    // channel for download workers to send results to their coordinator.
    let (tx, mut rx) = mpsc::unbounded();

    // start with one initial download task at index 0
    let active_tasks_counter = Arc::new(AtomicUsize::new(1));
    let mut next_task_index = 1;
    let mut task_bucket = vec![start_initial_download_task(
        initial_response,
        tx.clone(),
        active_tasks_counter.clone(),
        partition_size,
    )];

    let mut drain = SequentialBoundedDrain::new(max_buffers);
    let mut tx_opt = Some(tx);

    // This stream maintains up to parallel-many active client downloads at a time while maintaining
    // up to max_buffers-many buffers of length partition_size.
    // It re-sequences these buffers, only yielding the sequentially next buffer when it is ready,
    // regardless of the state of subsequent buffers.
    // Drain serves double duty of holding completed buffers as well as tracking position of the
    // download, indexed by chunk.
    let stream = async_stream::try_stream! {
        while drain.position() < total_chunks {
            // while there is room in the buffer drain and not at max connections, start new range downloads
            while drain.currently_accepting().contains(&next_task_index) && active_tasks_counter.load(Ordering::Relaxed) < parallel {
                match remaining_ranges.pop_front() {
                    Some(range) => {
                        let i = next_task_index;
                        next_task_index += 1;
                        active_tasks_counter.fetch_add(1, Ordering::Relaxed);

                        let t = tx_opt.as_ref().ok_or_else(||Error::with_message(ErrorKind::Other, "Channel closed unexpectedly."))?.clone();
                        task_bucket.push(start_download_task(client.clone(), range, t, active_tasks_counter.clone(), i));
                    }
                    None => {
                        // if ranges are finished, we'll never need to clone the transmitter again.
                        // drop this transmitter to ensure channel closes when expected
                        tx_opt = None;
                        break;
                    }
                }
            }

            // return next readied bytes, if any
            while let Some(bytes) = drain.pop() {
                yield bytes;
            }

            // early break if finished
            if drain.position() >= total_chunks {
                break;
            }

            // max tasks are spawned and sequential ready bytes already returned
            // this will not change until either:
            //   1. a task sends a message through this channel
            //   2. a task fails
            let channel_message;
            (channel_message, task_bucket) = await_message_while_joining_workers(&mut rx, task_bucket).await?;
            let (idx, bytes) = channel_message?;
            drain.push(idx, bytes)?;
        }
    };

    Ok(AsyncRawResponse::new(status, headers, Box::pin(stream)))
}

pub(crate) async fn download_into<Behavior>(
    mut buffer: &mut [u8],
    range: Option<Range<usize>>,
    parallel: NonZero<usize>,
    partition_size: NonZero<usize>,
    client: Arc<Behavior>,
) -> AzureResult<usize>
where
    Behavior: PartitionedDownloadBehavior + Send + Sync + 'static,
{
    let parallel = parallel.get();
    let partition_size = partition_size.get();
    let mut tasks = UnorderedFuturesDrain::with_capacity(parallel);

    // SAFETY: This function includes unsafe code that sends slices of `buffer` to other worker threads.
    // Those threads MUST ALL terminate before this function returns.
    // We accomplish this by emulating a try-finally block. The ENTIRE implementation is wrapped in a
    // closure and it's awaited result stored, allowing idiomatic `?` error returns to be assigned to
    // that result, rather than exit this whole function. Before returning that result, ensure ALL
    // tasks in `tasks` are joined.
    let overall_result = async {
        let (initial_response, stats_opt) =
            get_initial_response_and_analyze(range, partition_size, client.clone()).await?;

        let mut stats = match stats_opt {
            Some(s) => s,
            // If no information on ranges, no subsequent ranges to get.
            // Write whatever we have and return.
            None => return initial_response.into_body().collect_into(buffer).await,
        };

        // fail fast if we know we'll receive too much data
        if stats.overall_download_range.len() > buffer.len() {
            return Err(Error::with_message(
                ErrorKind::Other,
                format!(
                    "Buffer with length {} cannot fit payload with length {}.",
                    buffer.len(),
                    stats.overall_download_range.len()
                ),
            ));
        }

        let total_ranges = stats.remaining_download_ranges.len() + 1;

        // if no real parallelism, take the simple option of executing downloads sequentially.
        // no worker spawning.
        if parallel == 1 || stats.remaining_download_ranges.is_empty() {
            // sequence the initial response with a stream that fetches responses for each subsequent range
            let mut all_responses = pin!(stream::once(future::ready(Ok(initial_response))).chain(
                stream::iter(stats.remaining_download_ranges).then(|range| {
                    let client = client.clone();
                    async move { client.transfer_range(Some(range)).await }
                }),
            ));

            let mut total_written = 0;
            while let Some(response) = all_responses.try_next().await? {
                let written = response.into_body().collect_into(buffer).await?;
                buffer = &mut buffer[written..];
                total_written += written;
            }
            return Ok(total_written);
        }

        // channel for download workers to send results to their coordinator.
        let (tx, mut rx) = mpsc::unbounded();

        // Start collecting the body of the initial download.
        // Special-cased since we already have the response for it.

        // SAFETY: Ensure mutable reference safety for the SendSlice created from `buffer`.
        // This memory address is being sent to a separate worker thread. That worker thread CANNOT
        // live longer than the lifetime of `buffer`, and therefore MUST be cleaned up before
        // this function returns. Storing those worker references in `tasks` allows this function
        // to guarantee that cleanup at the end.
        unsafe {
            let init_dst;
            (init_dst, buffer) = buffer.split_at_mut(stats.initial_download_range.len());
            let init_dst = SendSlice::from_raw(init_dst.as_mut_ptr(), init_dst.len());
            tasks.push(start_initial_download_into_task(
                initial_response,
                init_dst,
                tx.clone(),
            ));
        }

        // tracking of completed work
        let mut received_result_count = 0;
        let mut bytes_copied = 0;
        let mut consume_available_messages = || {
            while received_result_count < total_ranges {
                match rx.try_recv() {
                    Ok(worker_result) => {
                        bytes_copied += worker_result?;
                        received_result_count += 1;
                    }
                    Err(mpsc::TryRecvError::Empty) => break,
                    Err(mpsc::TryRecvError::Closed) => Err(Error::with_message(
                        ErrorKind::Other,
                        "Download incomplete. Premature channel close.",
                    ))?,
                }
            }
            Ok::<(), Error>(())
        };

        // Download all remaining ranges, parallel-many at a time
        while let Some(range) = stats.remaining_download_ranges.pop_front() {
            // If at parallel limit, wait for a slot to open
            while tasks.len() >= parallel {
                if let Some(result) = tasks.next().await {
                    result.map_err(map_spawned_task_error)?;
                } else {
                    break;
                }
            }
            consume_available_messages()?;

            // SAFETY: Ensure mutable reference safety for the SendSlice created from `buffer`.
            // This memory address is being sent to a separate worker thread. That worker thread CANNOT
            // live longer than the lifetime of `buffer`, and therefore MUST be cleaned up before
            // this function returns. Storing those worker references in `tasks` allows this function
            // to guarantee that cleanup at the end.
            unsafe {
                let dst;
                (dst, buffer) = buffer.split_at_mut(range.len());
                let dst = SendSlice::from_raw(dst.as_mut_ptr(), dst.len());
                tasks.push(start_download_into_task(
                    client.clone(),
                    dst,
                    range,
                    tx.clone(),
                ));
            }
        }

        // wait on all remaining tasks, failing fast on error
        while let Some(join_result) = tasks.next().await {
            join_result.map_err(map_spawned_task_error)?;
            consume_available_messages()?;
        }

        // final validation of work done
        if tasks.total_completed() != received_result_count {
            Err(Error::with_message(
                ErrorKind::Other,
                format!(
                    "download completion count mismatch: tasks_completed={} messages_received={}",
                    tasks.total_completed(),
                    received_result_count
                ),
            ))
        } else {
            Ok(bytes_copied)
        }
    }
    .await;

    // SAFETY: This is where we prevent any tasks from living beyond this method return,
    // guaranteeing no references to `buffer` exist.
    // TODO: Implement aborting spawned runtime workers.
    while let Some(worker_result) = tasks.next().await {
        worker_result.map_err(map_spawned_task_error)?;
    }

    overall_result
}

async fn get_initial_response_and_analyze<Behavior>(
    range: Option<Range<usize>>,
    partition_size: usize,
    client: Arc<Behavior>,
) -> AzureResult<(AsyncRawResponse, Option<InitialResponseAnalysis>)>
where
    Behavior: PartitionedDownloadBehavior + Send + Sync + 'static,
{
    // Outer bound estimate of the resource range that will be downloaded. The actual download
    // range will never exceed these bounds, but it may be smaller, based on the actual size
    // of the remote resource.
    let max_download_range = range.unwrap_or(0..usize::MAX);
    if max_download_range.is_empty() {
        return Err(Error::with_message(
            ErrorKind::Other,
            "Provided range must have length > 0.",
        ));
    }

    let initial_response = download_with_empty_blob_safety(
        client.as_ref(),
        max_download_range.start
            ..min(
                max_download_range.end,
                max_download_range.start.saturating_add(partition_size),
            ),
    )
    .await?;

    let stats =
        analyze_initial_response(&initial_response, partition_size, max_download_range.end)?;

    Ok((initial_response, stats))
}

/// Race awaiting a message vs checking if tasks have completed successfully,
/// until either message is received or a task failure is found.
///
/// # Returns
///
/// - Returns Ok with received message and remaining un-joined tasks.
/// - Returns Err if channel closed before a message received.
/// - Returns Err if joined task closed with an error.
async fn await_message_while_joining_workers<T>(
    receiver: &mut UnboundedReceiver<T>,
    mut task_bucket: Vec<SpawnedTask>,
) -> AzureResult<(T, Vec<SpawnedTask>)> {
    let on_recv_err = |_| {
        Error::with_message(
            ErrorKind::Other,
            "Download incomplete. Premature channel close.",
        )
    };

    let mut message_fut = receiver.recv();
    // `task_bucket`` may be empty. `select_all` cannot handle that.
    while !task_bucket.is_empty() {
        match future::select(message_fut, future::select_all(task_bucket)).await {
            Either::Left((message, task_select)) => {
                return Ok((message.map_err(on_recv_err)?, task_select.into_inner()));
            }
            Either::Right(((completed_task, _, remaining_tasks), m_fut)) => {
                completed_task.map_err(map_spawned_task_error)?;
                task_bucket = remaining_tasks;
                message_fut = m_fut;
            }
        }
    }

    Ok((message_fut.await.map_err(on_recv_err)?, task_bucket))
}

/// Spawns a worker to take the given raw response and stream it into a buffer.
/// That buffer result is then sent through sender with chunk index 0.
fn start_initial_download_task(
    initial_response: AsyncRawResponse,
    mut sender: UnboundedSender<Result<(usize, Bytes), Error>>,
    active_tasks_counter: Arc<AtomicUsize>,
    partition_size: usize,
) -> SpawnedTask {
    get_async_runtime().spawn(Box::pin(async move {
        let mut dst = vec![0u8; partition_size];
        let res = initial_response
            .into_body()
            .collect_into(&mut dst)
            .await
            // this is the initial download task, it's chunk index is 0
            .map(|_| (0usize, dst.into()));
        active_tasks_counter.fetch_sub(1, Ordering::Relaxed);
        let _send_res = sender.send(res).await;
    }))
}

/// Spawns a worker to take the given raw response and stream it into the given buffer.
/// A result of the number of bytes written is then sent through sender.
///
/// If the body does not *exactly* fill the buffer, an error will be sent instead.
///
/// # Safety
/// Caller must ensure the slice represented by `buffer` is unused elsewhere until the returned future completes.
unsafe fn start_initial_download_into_task(
    initial_response: AsyncRawResponse,
    mut buffer: SendSlice<u8>,
    mut sender: UnboundedSender<Result<usize, Error>>,
) -> SpawnedTask {
    get_async_runtime().spawn(Box::pin(async move {
        let res = initial_response
            .into_body()
            .collect_into_exact(buffer.as_mut_slice())
            .await
            .map(|_| buffer.len());
        let _send_res = sender.send(res).await;
    }))
}

/// Spawns a worker to request the given range and stream it into a buffer.
/// That buffer result is then sent through sender with the given chunk index.
fn start_download_task<Behavior: PartitionedDownloadBehavior + Send + Sync + 'static>(
    client: Arc<Behavior>,
    range: Range<usize>,
    mut sender: UnboundedSender<Result<(usize, Bytes), Error>>,
    active_tasks_counter: Arc<AtomicUsize>,
    chunk_idx: usize,
) -> SpawnedTask {
    get_async_runtime().spawn(Box::pin(async move {
        let mut dst = vec![0u8; range.len()];
        let res = async {
            client
                .transfer_range(Some(range))
                .await?
                .into_body()
                .collect_into(&mut dst)
                .await
        }
        .await;
        if let Ok(count) = res {
            dst.truncate(count);
        }
        active_tasks_counter.fetch_sub(1, Ordering::Relaxed);
        let _send_res = sender.send(res.map(|_| (chunk_idx, dst.into()))).await;
    }))
}

/// Spawns a worker to request the given range and stream it into the given buffer.
/// The success result is then sent through sender.
///
/// If the body does not *exactly* fill the buffer, an error will be sent instead.
///
/// # Safety
/// Caller must ensure the slice represented by `buffer` is unused elsewhere until the returned future completes.
unsafe fn start_download_into_task<
    Behavior: PartitionedDownloadBehavior + Send + Sync + 'static,
>(
    client: Arc<Behavior>,
    mut buffer: SendSlice<u8>,
    range: Range<usize>,
    mut sender: UnboundedSender<Result<usize, Error>>,
) -> SpawnedTask {
    get_async_runtime().spawn(Box::pin(async move {
        let res = async {
            client
                .transfer_range(Some(range))
                .await?
                .into_body()
                .collect_into_exact(buffer.as_mut_slice())
                .await
        }
        .await
        .map(|_| buffer.len());
        let _send_res = sender.send(res).await;
    }))
}

/// Performs a `transfer_range()` call with the given range. If this results in a
/// RequestedRangeNotSatisfiable error, and if the requested range begins at the
/// start of the blob, retries the operation without a range argument.
/// This handles the service's edge case where a ranged get on an empty blob
/// always fails. Retrying with an empty range gives the correct empty blob data
/// as well as all the header information we expect.
async fn download_with_empty_blob_safety<Behavior>(
    client: &Behavior,
    range: Range<usize>,
) -> AzureResult<AsyncRawResponse>
where
    Behavior: PartitionedDownloadBehavior + Send + Sync + 'static,
{
    let range_start = range.start;
    match client.transfer_range(Some(range)).await {
        Ok(response) => Ok(response),
        Err(err) => match (err.http_status(), range_start) {
            (Some(StatusCode::RequestedRangeNotSatisfiable), 0) => {
                client.transfer_range(None).await
            }
            _ => Err(err),
        },
    }
}

struct InitialResponseAnalysis {
    overall_download_range: Range<usize>,
    initial_download_range: Range<usize>,
    remaining_download_ranges: VecDeque<Range<usize>>,
}
/// Reads over the response headers of the initial download response and compiles all relevant
/// information to perform the remaining downloads and arrange all resulting bytes.
///
/// # Returns
///
/// Ok(Some(analysis)) if the appropriate information was available.
///
/// Ok(None) if the appropriate information was not available.
///
/// Err(error) if there was an error parsing the appropriate information.
fn analyze_initial_response(
    initial_response: &AsyncRawResponse,
    partition_len: usize,
    max_range_end: usize,
) -> AzureResult<Option<InitialResponseAnalysis>> {
    if let Some(content_range) = initial_response
        .headers()
        .get_optional_as::<ContentRange, _>(&"content-range".into())?
    {
        if let (Some(received_range), Some(resource_len)) =
            (content_range.range, content_range.total_len)
        {
            let remainder_start = received_range.1;
            let remainder_end = min(max_range_end, resource_len);
            return Ok(Some(InitialResponseAnalysis {
                overall_download_range: received_range.0..remainder_end,
                initial_download_range: received_range.0..received_range.1,
                remaining_download_ranges: (remainder_start..remainder_end)
                    .step_by(partition_len)
                    .map(|i| i..min(i.saturating_add(partition_len), remainder_end))
                    .collect(),
            }));
        }
    }
    Ok(None)
}

fn map_spawned_task_error(err: Box<dyn std::error::Error + Send>) -> Error {
    Error::with_message(ErrorKind::Other, err.to_string())
}

trait DownloadRangeFuture: Future + Send {}
impl<T: Future + Send> DownloadRangeFuture for T {}

#[cfg(test)]
mod tests {
    use std::cmp::min;

    use azure_core::{
        http::{
            headers::{Header, Headers},
            StatusCode,
        },
        stream::BytesStream,
    };

    use azure_core_test::ErrorKind;
    use tokio::{
        sync::Mutex,
        time::{sleep, Duration},
    };

    use super::*;

    pub const KB: usize = 1024;
    pub const MB: usize = KB * 1024;
    pub const GB: usize = MB * 1024;

    #[derive(Debug)]
    enum MockPartitionedDownloadBehaviorInvocation {
        TransferRange(Option<Range<usize>>),
    }

    struct MockPartitionedDownloadBehavior {
        pub invocations: Mutex<Vec<MockPartitionedDownloadBehaviorInvocation>>,
        pub data: Bytes,
        pub delay_millis: Option<Range<u64>>,
    }

    impl MockPartitionedDownloadBehavior {
        pub fn new(data: impl Into<Bytes>, delay_millis: Option<Range<u64>>) -> Self {
            Self {
                invocations: Mutex::new(vec![]),
                data: data.into(),
                delay_millis,
            }
        }
    }

    #[async_trait::async_trait]
    impl PartitionedDownloadBehavior for MockPartitionedDownloadBehavior {
        async fn transfer_range(
            &self,
            requested_range: Option<Range<usize>>,
        ) -> AzureResult<AsyncRawResponse> {
            {
                self.invocations.lock().await.push(
                    MockPartitionedDownloadBehaviorInvocation::TransferRange(
                        requested_range.clone(),
                    ),
                );
            }

            if let Some(delay_millis_range) = self.delay_millis.clone() {
                let millis = rand::random_range(delay_millis_range);
                sleep(Duration::from_millis(millis)).await
            }

            struct ContentLength(usize);
            impl Header for ContentLength {
                fn name(&self) -> azure_core::http::headers::HeaderName {
                    "content-length".into()
                }
                fn value(&self) -> azure_core::http::headers::HeaderValue {
                    self.0.to_string().into()
                }
            }
            let mut headers = Headers::new();
            match (requested_range, self.data.len()) {
                (Some(range), data_len) => {
                    if range.start >= data_len {
                        return Err(ErrorKind::HttpResponse {
                            status: StatusCode::RequestedRangeNotSatisfiable,
                            error_code: Some("InvalidRange".into()),
                            raw_response: None,
                        }
                        .into_error());
                    }
                    let range = range.start..min(range.end, data_len);
                    if !range.is_empty() {
                        headers.add(ContentRange {
                            range: Some((range.start, range.end - 1)),
                            total_len: Some(self.data.len()),
                        })?
                    };
                    headers.add(ContentLength(range.len()))?;
                    let range = range.start..range.end;
                    Ok(AsyncRawResponse::new(
                        StatusCode::PartialContent,
                        headers,
                        Box::pin(BytesStream::from(self.data.slice(range))),
                    ))
                }
                (None, 0) => {
                    headers.add(ContentRange {
                        range: None,
                        total_len: None,
                    })?;
                    headers.add(ContentLength(0))?;
                    Ok(AsyncRawResponse::new(
                        StatusCode::Ok,
                        headers,
                        Box::pin(BytesStream::new_empty()),
                    ))
                }
                (None, data_len) => {
                    headers.add(ContentRange {
                        range: Some((0, data_len - 1)),
                        total_len: Some(data_len),
                    })?;
                    headers.add(ContentLength(data_len))?;
                    Ok(AsyncRawResponse::new(
                        StatusCode::Ok,
                        headers,
                        Box::pin(BytesStream::from(self.data.clone())),
                    ))
                }
            }
        }
    }

    struct SingleRangeArgSet {
        partition_len: usize,
        download_range: Option<(usize, usize)>,
    }
    fn single_range_args(data_len: usize) -> impl IntoIterator<Item = SingleRangeArgSet> {
        // trait not implemented for usize
        let part_len = data_len / 5;
        let extra = data_len / 5;
        let offset = data_len / 5;

        let start_range = (0, part_len);
        let mid_range = (offset, offset + part_len);
        let end_range = (data_len - part_len, data_len);
        [
            // exact len
            SingleRangeArgSet {
                partition_len: data_len,
                download_range: None,
            },
            // oversize len
            SingleRangeArgSet {
                partition_len: data_len + extra,
                download_range: None,
            },
            // exact range len (start)
            SingleRangeArgSet {
                partition_len: part_len,
                download_range: Some(start_range),
            },
            // oversize range len (start)
            SingleRangeArgSet {
                partition_len: part_len + extra,
                download_range: Some(start_range),
            },
            // exact range len (mid))
            SingleRangeArgSet {
                partition_len: part_len,
                download_range: Some(mid_range),
            },
            // oversize range len (mid))
            SingleRangeArgSet {
                partition_len: part_len + extra,
                download_range: Some(mid_range),
            },
            // exact range len (end)
            SingleRangeArgSet {
                partition_len: part_len,
                download_range: Some(end_range),
            },
            // oversize range len (end)
            SingleRangeArgSet {
                partition_len: part_len + extra,
                download_range: Some(end_range),
            },
        ]
    }

    #[tokio::test]
    async fn download_single_range() -> AzureResult<()> {
        const DATA_LEN: usize = 1024;
        const PARALLEL: usize = 2;

        let data = get_random_data(DATA_LEN);

        for args in single_range_args(DATA_LEN) {
            let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));

            let mut body = download(
                args.download_range.map(|r| r.0..r.1),
                PARALLEL.try_into().unwrap(),
                args.partition_len.try_into().unwrap(),
                mock.clone(),
            )
            .await?
            .into_body();
            let downloaded_data = body.buffer_all().await?;

            assert_eq!(
                &downloaded_data[..],
                match args.download_range {
                    Some(r) => &data[r.0..r.1],
                    None => &data[..],
                }
            );
            assert_eq!(mock.invocations.lock().await.len(), 1);
        }

        Ok(())
    }

    #[tokio::test]
    async fn download_into_single_range() -> AzureResult<()> {
        const DATA_LEN: usize = 1024;
        const PARALLEL: usize = 2;

        let data = get_random_data(DATA_LEN);

        for args in single_range_args(DATA_LEN) {
            let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));
            let mut destination = vec![
                0u8;
                if let Some((start, end)) = args.download_range {
                    end - start
                } else {
                    DATA_LEN
                }
            ];

            let written = download_into(
                &mut destination,
                args.download_range.map(|r| r.0..r.1),
                PARALLEL.try_into().unwrap(),
                args.partition_len.try_into().unwrap(),
                mock.clone(),
            )
            .await?;

            assert_eq!(destination.len(), written);
            assert_eq!(
                &destination,
                match args.download_range {
                    Some(r) => &data[r.0..r.1],
                    None => &data[..],
                }
            );
            assert_eq!(mock.invocations.lock().await.len(), 1);
        }

        Ok(())
    }

    struct MultiRangeArgSet {
        pub parallel: usize,
        pub partition_len: usize,
        pub download_range: Option<(usize, usize)>,
        pub expected_parts: usize,
    }
    fn multi_range_args(data_len: usize) -> impl IntoIterator<Item = MultiRangeArgSet> {
        let offset = data_len / 9;
        let range_len = data_len / 9;

        let mut combos = Vec::new();
        for parallel in [1, 4] {
            for blob_range in [
                (0, range_len),
                (offset, offset + range_len),
                (data_len - range_len, data_len),
            ] {
                for (partition_len, download_range) in [
                    (data_len - 1, None),              // barely smaller
                    (data_len / 2, None),              // half size
                    (data_len / 41, None),             // oddball size
                    (range_len - 1, Some(blob_range)), // barely smaller, range
                    (range_len / 2, Some(blob_range)), // half size, range
                    (data_len / 41, Some(blob_range)), // oddball size, range
                ] {
                    let expected_parts = match download_range {
                        Some((start, end)) => (end - start).div_ceil(partition_len),
                        None => data_len.div_ceil(partition_len),
                    };
                    combos.push(MultiRangeArgSet {
                        parallel,
                        partition_len,
                        download_range,
                        expected_parts,
                    });
                }
            }
        }

        combos
    }

    #[tokio::test]
    async fn download_multi_range() -> AzureResult<()> {
        const DATA_LEN: usize = 4096;

        let data = get_random_data(DATA_LEN);

        for args in multi_range_args(DATA_LEN) {
            let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));

            let mut body = download(
                args.download_range.map(|r| r.0..r.1),
                args.parallel.try_into().unwrap(),
                args.partition_len.try_into().unwrap(),
                mock.clone(),
            )
            .await?
            .into_body();
            let downloaded_data = body.buffer_all().await?;

            assert_eq!(
                downloaded_data.len(),
                args.download_range
                    .map_or(DATA_LEN, |range| range.1 - range.0),
                "Data mismatch. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts
            );
            assert_eq!(
                &downloaded_data[..],
                match args.download_range {
                    Some(r) => &data[r.0..r.1],
                    None => &data[..],
                },
                "Data mismatch. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts
            );
            assert_eq!(
                mock.invocations.lock().await.len(),
                args.expected_parts,
                "Unexpected invocation count. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts);
        }

        Ok(())
    }

    #[tokio::test]
    async fn download_into_multi_range() -> AzureResult<()> {
        const DATA_LEN: usize = 4096;

        let data = get_random_data(DATA_LEN);

        for args in multi_range_args(DATA_LEN) {
            let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));
            let mut downloaded_data = vec![
                0u8;
                if let Some((start, end)) = args.download_range {
                    end - start
                } else {
                    DATA_LEN
                }
            ];

            let written = download_into(
                &mut downloaded_data,
                args.download_range.map(|r| r.0..r.1),
                args.parallel.try_into().unwrap(),
                args.partition_len.try_into().unwrap(),
                mock.clone(),
            )
            .await?;

            assert_eq!(
                written,
                args.download_range
                    .map_or(DATA_LEN, |range| range.1 - range.0),
                "Data mismatch. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts
            );
            assert_eq!(
                &downloaded_data,
                match args.download_range {
                    Some(r) => &data[r.0..r.1],
                    None => &data,
                },
                "Data mismatch. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts
            );
            assert_eq!(
                mock.invocations.lock().await.len(),
                args.expected_parts,
                "Unexpected invocation count. partition_len={}. download_range={:?}, expected_parts={}",
                args.partition_len,
                args.download_range,
                args.expected_parts);
        }

        Ok(())
    }

    #[tokio::test]
    async fn download_ranges_parallel_maintain_order() -> AzureResult<()> {
        let segments: usize = 20;
        let partition_size = NonZero::new(3).unwrap();
        let parallel = NonZero::new(16).unwrap();
        let data_size: usize = partition_size.get() * segments;

        let data = get_random_data(data_size);
        let mock = Arc::new(MockPartitionedDownloadBehavior::new(
            data.clone(),
            Some(1..5),
        ));

        let mut body = download(None, parallel, partition_size, mock.clone())
            .await?
            .into_body();
        let downloaded_data = body.buffer_all().await?;

        assert_eq!(downloaded_data[..], data[..]);
        assert_eq!(mock.invocations.lock().await.len(), segments);

        Ok(())
    }

    #[tokio::test]
    async fn download_into_ranges_parallel_maintain_order() -> AzureResult<()> {
        let segments: usize = 20;
        let partition_size = NonZero::new(3).unwrap();
        let parallel = NonZero::new(16).unwrap();
        let data_size: usize = partition_size.get() * segments;

        let data = get_random_data(data_size);
        let mock = Arc::new(MockPartitionedDownloadBehavior::new(
            data.clone(),
            Some(1..5),
        ));

        let mut destination = vec![0u8; data.len()];
        let written = download_into(
            &mut destination,
            None,
            parallel,
            partition_size,
            mock.clone(),
        )
        .await?;

        assert_eq!(written, data_size);
        assert_eq!(&destination, &data);
        assert_eq!(mock.invocations.lock().await.len(), segments);

        Ok(())
    }

    #[tokio::test]
    async fn download_empty_resource() -> AzureResult<()> {
        let parallel = NonZero::new(1).unwrap();
        let partition_len = NonZero::new(MB).unwrap();
        let data = get_random_data(0);
        let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));

        let mut body = download(None, parallel, partition_len, mock.clone())
            .await?
            .into_body();
        let downloaded_data = body.buffer_all().await?;

        assert_eq!(downloaded_data.len(), 0);

        Ok(())
    }

    #[tokio::test]
    async fn download_into_empty_resource() -> AzureResult<()> {
        let parallel = NonZero::new(1).unwrap();
        let partition_len = NonZero::new(MB).unwrap();
        let data = get_random_data(0);
        let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));

        let mut destination = Vec::new();
        let written = download_into(
            &mut destination,
            None,
            parallel,
            partition_len,
            mock.clone(),
        )
        .await?;

        assert_eq!(written, 0);

        Ok(())
    }

    #[tokio::test]
    async fn download_into_insufficient_buffer() -> AzureResult<()> {
        let partition_len = NonZero::new(3).unwrap();
        let parallel = NonZero::new(2).unwrap();
        let data_len = 1024;

        let data = get_random_data(data_len);
        let mock = Arc::new(MockPartitionedDownloadBehavior::new(data.clone(), None));

        for buffer_len in [0, data_len - 1, data_len / 2] {
            let mut buffer = vec![0; buffer_len];
            assert!(
                download_into(&mut buffer, None, parallel, partition_len, mock.clone())
                    .await
                    .is_err()
            );
        }

        for range_len in [1, data_len, data_len / 2] {
            let mut buffer = vec![0; range_len - 1];
            assert!(download_into(
                &mut buffer,
                Some(0..range_len),
                parallel,
                partition_len,
                mock.clone()
            )
            .await
            .is_err());
        }

        Ok(())
    }

    trait BytesTryStreamExt {
        async fn buffer_all(&mut self) -> AzureResult<Vec<u8>>;
    }
    impl<S> BytesTryStreamExt for S
    where
        S: ?Sized + Stream<Item = AzureResult<Bytes>> + Unpin,
    {
        async fn buffer_all(&mut self) -> AzureResult<Vec<u8>> {
            let mut buffer = Vec::<u8>::new();
            while let Some(bytes) = self.try_next().await? {
                buffer.extend_from_slice(&bytes);
            }

            Ok(buffer)
        }
    }

    fn get_random_data(len: usize) -> Vec<u8> {
        let mut data: Vec<u8> = vec![0; len];
        rand::fill(&mut data[..]);
        data
    }
}