pulled coalesce_ranges conditionally in-tree, feature-gated object_store and reqwest

Author: feefladder

Cargo.toml

@@ -15,14 +15,19 @@ flate2 = "1.0.20"
 futures = "0.3.31"
 jpeg = { package = "jpeg-decoder", version = "0.3.0", default-features = false }
 num_enum = "0.7.3"
-object_store = "0.12"
+object_store = { version = "0.12", optional = true }
 # In the future we could make this feature-flagged, but for now we depend on
 # object_store which uses reqwest.
-reqwest = { version = "0.12", default-features = false }
+reqwest = { version = "0.12", default-features = false, optional = true }
 thiserror = "1"
 tokio = { version = "1.43.0", optional = true }
 weezl = "0.1.0"
 
 [dev-dependencies]
 tiff = "0.9.1"
 tokio = { version = "1.9", features = ["macros", "fs", "rt-multi-thread"] }
+
+[features]
+default = ["object_store"]
+reqwest = ["dep:reqwest"]
+object_store = ["dep:object_store", "reqwest"]

src/cog.rs

@@ -56,6 +56,7 @@ impl TIFF {
     }
 }
 
+#[cfg(feature = "object_store")]
 #[cfg(test)]
 mod test {
     use std::io::BufReader;

src/error.rs

@@ -24,6 +24,7 @@ pub enum AsyncTiffError {
     JPEGDecodingError(#[from] jpeg::Error),
 
     /// Error while fetching data using object store.
+    #[cfg(feature = "object_store")]
     #[error(transparent)]
     ObjectStore(#[from] object_store::Error),
 
@@ -32,6 +33,7 @@ pub enum AsyncTiffError {
     InternalTIFFError(#[from] crate::tiff::TiffError),
 
     /// Reqwest error
+    #[cfg(feature = "reqwest")] // see https://www.reddit.com/r/rust/comments/xyik51/comment/irhei39/
     #[error(transparent)]
     ReqwestError(#[from] reqwest::Error),
 

src/lib.rs

@@ -11,6 +11,13 @@ mod ifd;
 pub mod tiff;
 mod tile;
 
+#[cfg(not(feature = "object_store"))]
+mod object_store;
+#[cfg(feature = "object_store")]
+pub use object_store::coalesce_ranges;
+#[cfg(not(feature = "object_store"))]
+pub use object_store::util::coalesce_ranges;
+
 pub use cog::TIFF;
 pub use ifd::{ImageFileDirectories, ImageFileDirectory};
 pub use tile::Tile;

src/object_store/mod.rs

@@ -0,0 +1 @@
+pub mod util;

src/object_store/util.rs

@@ -0,0 +1,84 @@
+use bytes::Bytes;
+use futures::stream::{StreamExt, TryStreamExt};
+use std::ops::Range;
+
+const COALESCE_PARALLEL: usize = 10;
+
+/// Takes a function `fetch` that can fetch a range of bytes and uses this to
+/// fetch the provided byte `ranges`
+///
+/// To improve performance it will:
+///
+/// * Combine ranges less than `coalesce` bytes apart into a single call to `fetch`
+/// * Make multiple `fetch` requests in parallel (up to maximum of 10)
+///
+pub async fn coalesce_ranges<F, E, Fut>(
+    ranges: &[Range<u64>],
+    fetch: F,
+    coalesce: u64,
+) -> Result<Vec<Bytes>, E>
+where
+    F: Send + FnMut(Range<u64>) -> Fut,
+    E: Send,
+    Fut: std::future::Future<Output = Result<Bytes, E>> + Send,
+{
+    let fetch_ranges = merge_ranges(ranges, coalesce);
+
+    let fetched: Vec<_> = futures::stream::iter(fetch_ranges.iter().cloned())
+        .map(fetch)
+        .buffered(COALESCE_PARALLEL)
+        .try_collect()
+        .await?;
+
+    Ok(ranges
+        .iter()
+        .map(|range| {
+            let idx = fetch_ranges.partition_point(|v| v.start <= range.start) - 1;
+            let fetch_range = &fetch_ranges[idx];
+            let fetch_bytes = &fetched[idx];
+
+            let start = range.start - fetch_range.start;
+            let end = range.end - fetch_range.start;
+            let range = (start as usize)..(end as usize).min(fetch_bytes.len());
+            fetch_bytes.slice(range)
+        })
+        .collect())
+}
+
+/// Returns a sorted list of ranges that cover `ranges`
+fn merge_ranges(ranges: &[Range<u64>], coalesce: u64) -> Vec<Range<u64>> {
+    if ranges.is_empty() {
+        return vec![];
+    }
+
+    let mut ranges = ranges.to_vec();
+    ranges.sort_unstable_by_key(|range| range.start);
+
+    let mut ret = Vec::with_capacity(ranges.len());
+    let mut start_idx = 0;
+    let mut end_idx = 1;
+
+    while start_idx != ranges.len() {
+        let mut range_end = ranges[start_idx].end;
+
+        while end_idx != ranges.len()
+            && ranges[end_idx]
+                .start
+                .checked_sub(range_end)
+                .map(|delta| delta <= coalesce)
+                .unwrap_or(true)
+        {
+            range_end = range_end.max(ranges[end_idx].end);
+            end_idx += 1;
+        }
+
+        let start = ranges[start_idx].start;
+        let end = range_end;
+        ret.push(start..end);
+
+        start_idx = end_idx;
+        end_idx += 1;
+    }
+
+    ret
+}

src/reader.rs

@@ -9,8 +9,11 @@ use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
 use bytes::buf::Reader;
 use bytes::{Buf, Bytes};
 use futures::future::{BoxFuture, FutureExt, TryFutureExt};
+#[cfg(feature = "object_store")]
 use object_store::ObjectStore;
 
+use crate::coalesce_ranges;
+
 use crate::error::{AsyncTiffError, AsyncTiffResult};
 
 /// The asynchronous interface used to read COG files
@@ -33,17 +36,19 @@ pub trait AsyncFileReader: Debug + Send + Sync {
     /// Retrieve the bytes in `range`
     fn get_bytes(&self, range: Range<u64>) -> BoxFuture<'_, AsyncTiffResult<Bytes>>;
 
-    /// Retrieve multiple byte ranges. The default implementation will call `get_bytes`
-    /// sequentially
+    /// Retrieve multiple byte ranges. The default implementation will
+    /// coalesce ranges with:
+    /// - less than 1024*1024=1MB space in between
+    /// - 10 parallel requests
     fn get_byte_ranges(
         &self,
         ranges: Vec<Range<u64>>,
     ) -> BoxFuture<'_, AsyncTiffResult<Vec<Bytes>>> {
         async move {
             let mut result = Vec::with_capacity(ranges.len());
 
-            for range in ranges.into_iter() {
-                let data = self.get_bytes(range).await?;
+            for data in coalesce_ranges(&ranges, |range| self.get_bytes(range), 1024 * 1024).await?
+            {
                 result.push(data);
             }
 
@@ -91,12 +96,13 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
 // }
 
 /// An AsyncFileReader that reads from an [`ObjectStore`] instance.
+#[cfg(feature = "object_store")]
 #[derive(Clone, Debug)]
 pub struct ObjectReader {
     store: Arc<dyn ObjectStore>,
     path: object_store::path::Path,
 }
-
+#[cfg(feature = "object_store")]
 impl ObjectReader {
     /// Creates a new [`ObjectReader`] for the provided [`ObjectStore`] and path
     ///
@@ -105,7 +111,7 @@ impl ObjectReader {
         Self { store, path }
     }
 }
-
+#[cfg(feature = "object_store")]
 impl AsyncFileReader for ObjectReader {
     fn get_bytes(&self, range: Range<u64>) -> BoxFuture<'_, AsyncTiffResult<Bytes>> {
         let range = range.start as _..range.end as _;
@@ -134,19 +140,20 @@ impl AsyncFileReader for ObjectReader {
 }
 
 /// An AsyncFileReader that reads from a URL using reqwest.
+#[cfg(feature = "reqwest")]
 #[derive(Debug, Clone)]
 pub struct ReqwestReader {
     client: reqwest::Client,
     url: reqwest::Url,
 }
-
+#[cfg(feature = "reqwest")]
 impl ReqwestReader {
     /// Construct a new ReqwestReader from a reqwest client and URL.
     pub fn new(client: reqwest::Client, url: reqwest::Url) -> Self {
         Self { client, url }
     }
 }
-
+#[cfg(feature = "reqwest")]
 impl AsyncFileReader for ReqwestReader {
     fn get_bytes(&self, range: Range<u64>) -> BoxFuture<'_, AsyncTiffResult<Bytes>> {
         let url = self.url.clone();