Move the ConnectionPool to util and restructure a bunch of stuff so we can

make util public.

The ConnPool will probably move somewhere else in the longer term.

Author: rklaehn

src/api/blobs/reader.rs

@@ -221,14 +221,14 @@ mod tests {
 
     use super::*;
     use crate::{
+        protocol::ChunkRangesExt,
         store::{
             fs::{
                 tests::{create_n0_bao, test_data, INTERESTING_SIZES},
                 FsStore,
             },
             mem::MemStore,
         },
-        util::ChunkRangesExt,
     };
 
     async fn reader_smoke(blobs: &Blobs) -> TestResult<()> {

src/api/downloader.rs

@@ -442,7 +442,7 @@ async fn execute_get(
                 request: request.clone(),
             })
             .await?;
-        let conn = pool.connect(provider);
+        let conn = pool.get_or_connect(provider);
         let local = remote.local_for_request(request.clone()).await?;
         if local.is_complete() {
             return Ok(());

src/api/remote.rs

@@ -1067,7 +1067,7 @@ mod tests {
 
     use crate::{
         api::blobs::Blobs,
-        protocol::{ChunkRangesSeq, GetRequest},
+        protocol::{ChunkRangesExt, ChunkRangesSeq, GetRequest},
         store::{
             fs::{
                 tests::{create_n0_bao, test_data, INTERESTING_SIZES},
@@ -1076,7 +1076,6 @@ mod tests {
             mem::MemStore,
         },
         tests::{add_test_hash_seq, add_test_hash_seq_incomplete},
-        util::ChunkRangesExt,
     };
 
     #[tokio::test]

src/get/request.rs

@@ -27,8 +27,7 @@ use super::{fsm, GetError, GetResult, Stats};
 use crate::{
     get::error::{BadRequestSnafu, LocalFailureSnafu},
     hashseq::HashSeq,
-    protocol::{ChunkRangesSeq, GetRequest},
-    util::ChunkRangesExt,
+    protocol::{ChunkRangesExt, ChunkRangesSeq, GetRequest},
     Hash, HashAndFormat,
 };
 

src/lib.rs

@@ -43,7 +43,7 @@ pub mod ticket;
 
 #[doc(hidden)]
 pub mod test;
-mod util;
+pub mod util;
 
 #[cfg(test)]
 mod tests;

src/protocol.rs

@@ -373,21 +373,25 @@
 //! a large existing system that has demonstrated performance issues.
 //!
 //! If in doubt, just use multiple requests and multiple connections.
-use std::io;
+use std::{
+    io,
+    ops::{Bound, RangeBounds},
+};
 
+use bao_tree::{io::round_up_to_chunks, ChunkNum};
 use builder::GetRequestBuilder;
 use derive_more::From;
 use iroh::endpoint::VarInt;
 use irpc::util::AsyncReadVarintExt;
 use postcard::experimental::max_size::MaxSize;
+use range_collections::{range_set::RangeSetEntry, RangeSet2};
 use serde::{Deserialize, Serialize};
 mod range_spec;
 pub use bao_tree::ChunkRanges;
 pub use range_spec::{ChunkRangesSeq, NonEmptyRequestRangeSpecIter, RangeSpec};
 use snafu::{GenerateImplicitData, Snafu};
 use tokio::io::AsyncReadExt;
 
-pub use crate::util::ChunkRangesExt;
 use crate::{api::blobs::Bitfield, provider::CountingReader, BlobFormat, Hash, HashAndFormat};
 
 /// Maximum message size is limited to 100MiB for now.
@@ -714,6 +718,73 @@ impl TryFrom<VarInt> for Closed {
     }
 }
 
+pub trait ChunkRangesExt {
+    fn last_chunk() -> Self;
+    fn chunk(offset: u64) -> Self;
+    fn bytes(ranges: impl RangeBounds<u64>) -> Self;
+    fn chunks(ranges: impl RangeBounds<u64>) -> Self;
+    fn offset(offset: u64) -> Self;
+}
+
+impl ChunkRangesExt for ChunkRanges {
+    fn last_chunk() -> Self {
+        ChunkRanges::from(ChunkNum(u64::MAX)..)
+    }
+
+    /// Create a chunk range that contains a single chunk.
+    fn chunk(offset: u64) -> Self {
+        ChunkRanges::from(ChunkNum(offset)..ChunkNum(offset + 1))
+    }
+
+    /// Create a range of chunks that contains the given byte ranges.
+    /// The byte ranges are rounded up to the nearest chunk size.
+    fn bytes(ranges: impl RangeBounds<u64>) -> Self {
+        round_up_to_chunks(&bounds_from_range(ranges, |v| v))
+    }
+
+    /// Create a range of chunks from u64 chunk bounds.
+    ///
+    /// This is equivalent but more convenient than using the ChunkNum newtype.
+    fn chunks(ranges: impl RangeBounds<u64>) -> Self {
+        bounds_from_range(ranges, ChunkNum)
+    }
+
+    /// Create a chunk range that contains a single byte offset.
+    fn offset(offset: u64) -> Self {
+        Self::bytes(offset..offset + 1)
+    }
+}
+
+// todo: move to range_collections
+pub(crate) fn bounds_from_range<R, T, F>(range: R, f: F) -> RangeSet2<T>
+where
+    R: RangeBounds<u64>,
+    T: RangeSetEntry,
+    F: Fn(u64) -> T,
+{
+    let from = match range.start_bound() {
+        Bound::Included(start) => Some(*start),
+        Bound::Excluded(start) => {
+            let Some(start) = start.checked_add(1) else {
+                return RangeSet2::empty();
+            };
+            Some(start)
+        }
+        Bound::Unbounded => None,
+    };
+    let to = match range.end_bound() {
+        Bound::Included(end) => end.checked_add(1),
+        Bound::Excluded(end) => Some(*end),
+        Bound::Unbounded => None,
+    };
+    match (from, to) {
+        (Some(from), Some(to)) => RangeSet2::from(f(from)..f(to)),
+        (Some(from), None) => RangeSet2::from(f(from)..),
+        (None, Some(to)) => RangeSet2::from(..f(to)),
+        (None, None) => RangeSet2::all(),
+    }
+}
+
 pub mod builder {
     use std::collections::BTreeMap;
 
@@ -863,7 +934,7 @@ pub mod builder {
         use bao_tree::ChunkNum;
 
         use super::*;
-        use crate::{protocol::GetManyRequest, util::ChunkRangesExt};
+        use crate::protocol::{ChunkRangesExt, GetManyRequest};
 
         #[test]
         fn chunk_ranges_ext() {

src/protocol/range_spec.rs

@@ -12,7 +12,7 @@ use bao_tree::{ChunkNum, ChunkRangesRef};
 use serde::{Deserialize, Serialize};
 use smallvec::{smallvec, SmallVec};
 
-pub use crate::util::ChunkRangesExt;
+use crate::protocol::ChunkRangesExt;
 
 static CHUNK_RANGES_EMPTY: OnceLock<ChunkRanges> = OnceLock::new();
 
@@ -511,7 +511,7 @@ mod tests {
     use proptest::prelude::*;
 
     use super::*;
-    use crate::util::ChunkRangesExt;
+    use crate::protocol::ChunkRangesExt;
 
     fn ranges(value_range: Range<u64>) -> impl Strategy<Value = ChunkRanges> {
         prop::collection::vec((value_range.clone(), value_range), 0..16).prop_map(|v| {

src/store/fs.rs

@@ -111,6 +111,7 @@ use crate::{
         },
         ApiClient,
     },
+    protocol::ChunkRangesExt,
     store::{
         fs::{
             bao_file::{
@@ -125,7 +126,6 @@ use crate::{
     util::{
         channel::oneshot,
         temp_tag::{TagDrop, TempTag, TempTagScope, TempTags},
-        ChunkRangesExt,
     },
 };
 mod bao_file;

src/store/mem.rs

@@ -56,14 +56,12 @@ use crate::{
         tags::TagInfo,
         ApiClient,
     },
+    protocol::ChunkRangesExt,
     store::{
         util::{SizeInfo, SparseMemFile, Tag},
         HashAndFormat, IROH_BLOCK_SIZE,
     },
-    util::{
-        temp_tag::{TagDrop, TempTagScope, TempTags},
-        ChunkRangesExt,
-    },
+    util::temp_tag::{TagDrop, TempTagScope, TempTags},
     BlobFormat, Hash,
 };
 

src/store/readonly_mem.rs

@@ -41,8 +41,8 @@ use crate::{
         },
         ApiClient, TempTag,
     },
+    protocol::ChunkRangesExt,
     store::{mem::CompleteStorage, IROH_BLOCK_SIZE},
-    util::ChunkRangesExt,
     Hash,
 };