First somewhat working size handling approach

Author: rklaehn

examples/multiprovider.rs

@@ -104,33 +104,41 @@ async fn provide(args: ProvideArgs) -> anyhow::Result<()> {
 /// Progress for a single download
 struct BlobDownloadProgress {
     request: DownloadRequest,
-    current: ChunkRanges,
+    current: BitfieldState,
 }
 
 impl BlobDownloadProgress {
     fn new(request: DownloadRequest) -> Self {
         Self {
             request,
-            current: ChunkRanges::empty(),
+            current: BitfieldState::unknown(),
         }
     }
 
     fn update(&mut self, ev: BitfieldEvent) {
         match ev {
-            BitfieldEvent::State(BitfieldState { ranges, .. }) => {
-                self.current = ranges;
+            BitfieldEvent::State(BitfieldState { ranges, size }) => {
+                self.current.size = self.current.size.min(size);
+                self.current.ranges = ranges;
+                self.request.ranges &= ChunkRanges::from(..ChunkNum::chunks(self.current.size));
             }
-            BitfieldEvent::Update(BitfieldUpdate { added, removed, .. }) => {
-                self.current |= added;
-                self.current -= removed;
+            BitfieldEvent::Update(BitfieldUpdate {
+                added,
+                removed,
+                size,
+            }) => {
+                self.current.size = self.current.size.min(size);
+                self.request.ranges &= ChunkRanges::from(..ChunkNum::chunks(self.current.size));
+                self.current.ranges |= added;
+                self.current.ranges -= removed;
             }
         }
     }
 
     #[allow(dead_code)]
     fn get_stats(&self) -> (u64, u64) {
-        let total = total_bytes(&self.request.ranges, u64::MAX);
-        let downloaded = total_bytes(&self.current, u64::MAX);
+        let total = total_bytes(&self.request.ranges, self.current.size);
+        let downloaded = total_bytes(&self.current.ranges, self.current.size);
         (downloaded, total)
     }
 
@@ -140,7 +148,7 @@ impl BlobDownloadProgress {
     }
 
     fn is_done(&self) -> bool {
-        self.current == self.request.ranges
+        self.current.ranges == self.request.ranges
     }
 }
 
@@ -172,23 +180,24 @@ async fn download_impl<S: Store>(args: DownloadArgs, store: S) -> anyhow::Result
         .build();
     let request = DownloadRequest {
         hash: args.hash,
-        ranges: ChunkRanges::from(ChunkNum(0)..ChunkNum(25421)),
+        ranges: ChunkRanges::all(),
     };
     let downloader2 = downloader.clone();
     let mut progress = BlobDownloadProgress::new(request.clone());
     tokio::spawn(async move {
         let request = ObserveRequest {
             hash: args.hash,
-            ranges: ChunkRanges::from(ChunkNum(0)..ChunkNum(25421)),
+            ranges: ChunkRanges::all(),
             buffer: 1024,
         };
         let mut observe = downloader2.observe(request).await?;
         let term = Term::stdout();
         let (_, rows) = term.size();
         while let Some(chunk) = observe.recv().await {
             progress.update(chunk);
-            let current = progress.current.boundaries();
+            let current = progress.current.ranges.boundaries();
             let requested = progress.request.ranges.boundaries();
+            println!("observe print_bitmap {:?} {:?}", current, requested);
             let bitmap = print_bitmap(current, requested, rows as usize);
             print!("\r{bitmap}");
             if progress.is_done() {

src/downloader2.rs

@@ -118,6 +118,16 @@ pub struct BitfieldState {
     pub size: u64,
 }
 
+impl BitfieldState {
+    /// State for a completely unknown bitfield
+    pub fn unknown() -> Self {
+        Self {
+            ranges: ChunkRanges::empty(),
+            size: u64::MAX,
+        }
+    }
+}
+
 /// A download request
 #[derive(Debug, Clone)]
 pub struct DownloadRequest {
@@ -368,11 +378,7 @@ async fn get_valid_ranges_local<S: Store>(hash: &Hash, store: S) -> anyhow::Resu
         let (ranges, size) = crate::get::db::valid_ranges_and_size::<S>(&entry).await?;
         Ok(BitfieldState { ranges, size }.into())
     } else {
-        Ok(BitfieldState {
-            ranges: ChunkRanges::empty(),
-            size: u64::MAX,
-        }
-        .into())
+        Ok(BitfieldState::unknown().into())
     }
 }
 
@@ -423,11 +429,7 @@ impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
             async move {
                 let event = match recv.await {
                     Ok(ev) => ev,
-                    Err(_) => BitfieldState {
-                        ranges: ChunkRanges::empty(),
-                        size: u64::MAX,
-                    }
-                    .into(),
+                    Err(_) => BitfieldState::unknown().into(),
                 };
                 event
             }

src/downloader2/state.rs

@@ -324,21 +324,21 @@ impl DownloaderState {
                 if peer == BitfieldPeer::Local {
                     // we got a new local bitmap, notify local observers
                     // we must notify all local observers, even if the bitmap is empty
+                    state.size = state.size.min(size);
                     if let Some(observers) = self.observers.get_by_hash(&hash) {
                         for (id, request) in observers {
                             let ranges = &ranges & &request.ranges;
                             evs.push(Event::LocalBitfieldInfo {
                                 id: *id,
                                 event: BitfieldState {
                                     ranges: ranges.clone(),
-                                    size,
+                                    size: state.size,
                                 }
                                 .into(),
                             });
                         }
                     }
                     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
@@ -532,11 +532,13 @@ impl DownloaderState {
             // we don't have the self state yet, so we can't really decide if we need to download anything at all
             return Ok(());
         };
+        let mask = ChunkRanges::from(ChunkNum(0)..ChunkNum::chunks(self_state.size));
         let mut completed = vec![];
         for (id, download) in self.downloads.iter_mut_for_hash(hash) {
             if just_id.is_some() && just_id != Some(*id) {
                 continue;
             }
+            download.request.ranges &= mask.clone();
             // check if the entire download is complete. If this is the case, peer downloads will be cleaned up later
             if self_state.ranges.is_superset(&download.request.ranges) {
                 // notify the user that the download is complete
@@ -886,6 +888,9 @@ struct PeerState {
 }
 
 /// Information about one blob on one peer
+///
+/// Note that for remote peers we can't really trust this information.
+/// They could lie about the size, and the ranges could be either wrong or outdated.
 struct PeerBlobState {
     /// The subscription id for the subscription
     subscription_id: BitfieldSubscriptionId,
@@ -1117,11 +1122,7 @@ mod tests {
         state.apply(Command::BitfieldInfo {
             peer: Local,
             hash,
-            event: BitfieldState {
-                ranges: ChunkRanges::empty(),
-                size: u64::MAX,
-            }
-            .into(),
+            event: BitfieldState::unknown().into(),
         });
         // We have a peer for the hash
         state.apply(Command::PeerDiscovered { peer: peer_a, hash });
@@ -1357,11 +1358,7 @@ mod tests {
         state.apply(Command::BitfieldInfo {
             peer: Local,
             hash,
-            event: BitfieldState {
-                ranges: ChunkRanges::empty(),
-                size: u64::MAX,
-            }
-            .into(),
+            event: BitfieldState::unknown().into(),
         });
         // We have a peer for the hash
         state.apply(Command::PeerDiscovered { peer: peer_a, hash });