Make use of the size

Author: rklaehn

examples/multiprovider.rs

@@ -120,7 +120,7 @@ impl BlobDownloadProgress {
             BitfieldEvent::State { ranges, .. } => {
                 self.current = ranges;
             }
-            BitfieldEvent::Update { added, removed } => {
+            BitfieldEvent::Update { added, removed, .. } => {
                 self.current |= added;
                 self.current -= removed;
             }

src/downloader2.rs

@@ -95,7 +95,9 @@ pub enum BitfieldEvent {
     State {
         /// The entire bitfield
         ranges: ChunkRanges,
-        /// The most precise known size of the blob
+        /// The most precise known size of the blob.
+        ///
+        /// If I know nothing about the blob, this is u64::MAX.
         size: u64,
     },
     /// An update to the bitfield
@@ -104,6 +106,8 @@ pub enum BitfieldEvent {
         added: ChunkRanges,
         /// The ranges that were removed
         removed: ChunkRanges,
+        /// A refinement of the size of the blob.
+        size: u64,
     },
 }
 
@@ -346,23 +350,25 @@ impl<S> SimpleBitfieldSubscription<S> {
     }
 }
 
-async fn get_valid_ranges_local<S: Store>(hash: &Hash, store: S) -> anyhow::Result<ChunkRanges> {
+async fn get_valid_ranges_local<S: Store>(hash: &Hash, store: S) -> anyhow::Result<BitfieldEvent> {
     if let Some(entry) = store.get_mut(hash).await? {
-        crate::get::db::valid_ranges::<S>(&entry).await
+        let (ranges, size) = crate::get::db::valid_ranges_and_size::<S>(&entry).await?;
+        Ok(BitfieldEvent::State { ranges, size })
     } else {
-        Ok(ChunkRanges::empty())
+        Ok(BitfieldEvent::State { ranges: ChunkRanges::empty(), size: u64::MAX })
     }
 }
 
 async fn get_valid_ranges_remote(
     endpoint: &Endpoint,
     id: NodeId,
     hash: &Hash,
-) -> anyhow::Result<ChunkRanges> {
+) -> anyhow::Result<BitfieldEvent> {
     let conn = endpoint.connect(id, crate::ALPN).await?;
     let (size, _) = crate::get::request::get_verified_size(&conn, hash).await?;
     let chunks = (size + 1023) / 1024;
-    Ok(ChunkRanges::from(ChunkNum(0)..ChunkNum(chunks)))
+    let ranges = ChunkRanges::from(ChunkNum(0)..ChunkNum(chunks));
+    Ok(BitfieldEvent::State { ranges, size })
 }
 
 impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
@@ -398,11 +404,14 @@ impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
         }
         Box::pin(
             async move {
-                let ranges = match recv.await {
+                let event = match recv.await {
                     Ok(ev) => ev,
-                    Err(_) => ChunkRanges::empty(),
+                    Err(_) => BitfieldEvent::State {
+                        ranges: ChunkRanges::empty(),
+                        size: u64::MAX,
+                    },
                 };
-                BitfieldEvent::State { ranges, size: u64::MAX }
+                event
             }
             .into_stream(),
         )

src/downloader2/actor.rs

@@ -288,6 +288,7 @@ async fn peer_download<S: Store>(
                                 event: BitfieldEvent::Update {
                                     added,
                                     removed: ChunkRanges::empty(),
+                                    size,
                                 },
                             })
                             .await

src/downloader2/state.rs

@@ -331,6 +331,7 @@ impl DownloaderState {
                         }
                     }
                     state.ranges = ranges;
+                    state.size = state.size.min(size);
                     self.check_completion(hash, None, evs)?;
                 } else {
                     // We got an entirely new peer, mark all affected downloads for rebalancing
@@ -340,21 +341,23 @@ impl DownloaderState {
                         }
                     }
                     state.ranges = ranges;
+                    state.size = state.size.min(size);
                 }
                 // we have to call start_downloads even if the local bitfield set, since we don't know in which order local and remote bitfields arrive
                 self.start_downloads(hash, None, evs)?;
             }
             Command::BitfieldInfo {
                 peer,
                 hash,
-                event: BitfieldEvent::Update { added, removed },
+                event: BitfieldEvent::Update { added, removed, size },
             } => {
                 let state = self.bitfields.get_mut(&(peer, hash)).context(format!(
                     "bitfield update for unknown peer {peer:?} and hash {hash}"
                 ))?;
                 if peer == BitfieldPeer::Local {
                     // we got a local bitfield update, notify local observers
                     // for updates we can just notify the observers that have a non-empty intersection with the update
+                    state.size = state.size.min(size);
                     if let Some(observers) = self.observers.get_by_hash(&hash) {
                         for (id, request) in observers {
                             let added = &added & &request.ranges;
@@ -365,6 +368,7 @@ impl DownloaderState {
                                     event: BitfieldEvent::Update {
                                         added: &added & &request.ranges,
                                         removed: &removed & &request.ranges,
+                                        size: state.size,
                                     }
                                 });
                             }
@@ -383,6 +387,7 @@ impl DownloaderState {
                     }
                     state.ranges |= added;
                     state.ranges &= !removed;
+                    state.size = state.size.min(size);
                     // a local bitfield update does not make more data available, so we don't need to start downloads
                     self.start_downloads(hash, None, evs)?;
                 }
@@ -874,6 +879,8 @@ struct PeerBlobState {
     subscription_count: usize,
     /// chunk ranges this peer reports to have
     ranges: ChunkRanges,
+    /// The minimum reported size of the blob
+    size: u64,
 }
 
 impl PeerBlobState {
@@ -882,6 +889,7 @@ impl PeerBlobState {
             subscription_id,
             subscription_count: 1,
             ranges: ChunkRanges::empty(),
+            size: u64::MAX,
         }
     }
 }
@@ -977,6 +985,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([16..32]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         assert!(evs.is_empty());
@@ -1036,6 +1045,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([32..48]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         assert!(evs.is_empty());
@@ -1054,6 +1064,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([48..64]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         assert!(
@@ -1129,6 +1140,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([0..16]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         // The peer now has more data
@@ -1154,6 +1166,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([16..32]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         // This triggers cancellation of the first peer download and starting a new one for the remaining data
@@ -1187,6 +1200,7 @@ mod tests {
             event: BitfieldEvent::Update {
                 added: chunk_ranges([32..64]),
                 removed: ChunkRanges::empty(),
+                size: u64::MAX,
             },
         });
         assert!(

src/get/db.rs

@@ -237,6 +237,12 @@ async fn get_blob<D: BaoStore>(
 
 /// Given a partial entry, get the valid ranges.
 pub async fn valid_ranges<D: MapMut>(entry: &D::EntryMut) -> anyhow::Result<ChunkRanges> {
+    let (ranges, _) = valid_ranges_and_size::<D>(entry).await?;
+    Ok(ranges)
+}
+
+/// Given a partial entry, get the valid ranges.
+pub async fn valid_ranges_and_size<D: MapMut>(entry: &D::EntryMut) -> anyhow::Result<(ChunkRanges, u64)> {
     use tracing::trace as log;
     // compute the valid range from just looking at the data file
     let mut data_reader = entry.data_reader().await?;
@@ -253,7 +259,7 @@ pub async fn valid_ranges<D: MapMut>(entry: &D::EntryMut) -> anyhow::Result<Chun
     let valid: ChunkRanges = valid_from_data.intersection(&valid_from_outboard);
     log!("valid_from_data: {:?}", valid_from_data);
     log!("valid_from_outboard: {:?}", valid_from_data);
-    Ok(valid)
+    Ok((valid, data_size))
 }
 
 /// Get a blob that was requested completely.