Even more size refactoring

Author: rklaehn

examples/multiprovider.rs

@@ -10,7 +10,6 @@ use iroh_blobs::{
         ObserveRequest, StaticContentDiscovery,
     },
     store::Store,
-    util::total_bytes,
     Hash,
 };
 
@@ -118,30 +117,30 @@ impl BlobDownloadProgress {
     fn update(&mut self, ev: BitfieldEvent) {
         match ev {
             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));
+                self.current
+                    .size
+                    .update(size)
+                    .expect("verified size changed");
             }
             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;
+                self.current
+                    .size
+                    .update(size)
+                    .expect("verified size changed");
+                if let Some(size) = self.current.size.value() {
+                    self.request.ranges &= ChunkRanges::from(..ChunkNum::chunks(size));
+                }
             }
         }
     }
 
-    #[allow(dead_code)]
-    fn get_stats(&self) -> (u64, u64) {
-        let total = total_bytes(&self.request.ranges, self.current.size);
-        let downloaded = total_bytes(&self.current.ranges, self.current.size);
-        (downloaded, total)
-    }
-
     #[allow(dead_code)]
     fn get_bitmap(&self) -> String {
         format!("{:?}", self.current)
@@ -197,7 +196,6 @@ async fn download_impl<S: Store>(args: DownloadArgs, store: S) -> anyhow::Result
             progress.update(chunk);
             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

@@ -27,7 +27,7 @@ use crate::{
         Stats,
     },
     protocol::{GetRequest, RangeSpec, RangeSpecSeq},
-    store::{BaoBatchWriter, MapEntryMut, Store},
+    store::{BaoBatchWriter, BaoBlobSize, MapEntry, MapEntryMut, Store},
     util::local_pool::{self, LocalPool, LocalPoolHandle},
     Hash,
 };
@@ -89,6 +89,64 @@ pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static {
 /// A boxed bitfield subscription
 pub type BoxedBitfieldSubscription = Box<dyn BitfieldSubscription>;
 
+/// Knowlege about the size of a blob
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BaoBlobSizeOpt {
+    /// We have a size that a peer told us about, but we don't know if it is correct
+    /// It can be off at most by a factor of 2, so it is OK for things like showing
+    /// a progress bar or even for an allocation size
+    Unverified(u64),
+    /// We know the size, and it is verified
+    /// either by having the last chunk locally or by receiving a size proof from a peer
+    Verified(u64),
+    /// We know nothing, e.g. we have never heard of the blob
+    Unknown,
+}
+
+impl BaoBlobSizeOpt {
+    /// Get the value of the size, if known
+    pub fn value(self) -> Option<u64> {
+        match self {
+            BaoBlobSizeOpt::Unverified(x) => Some(x),
+            BaoBlobSizeOpt::Verified(x) => Some(x),
+            BaoBlobSizeOpt::Unknown => None,
+        }
+    }
+
+    /// Update the size information
+    ///
+    /// Unkown sizes are always updated
+    /// Unverified sizes are updated if the new size is verified
+    /// Verified sizes must never change
+    pub fn update(&mut self, size: BaoBlobSizeOpt) -> anyhow::Result<()> {
+        match self {
+            BaoBlobSizeOpt::Verified(old) => {
+                if let BaoBlobSizeOpt::Verified(new) = size {
+                    if *old != new {
+                        anyhow::bail!("mismatched verified sizes: {old} != {new}");
+                    }
+                }
+            }
+            BaoBlobSizeOpt::Unverified(_) => {
+                if let BaoBlobSizeOpt::Verified(new) = size {
+                    *self = BaoBlobSizeOpt::Verified(new);
+                }
+            }
+            BaoBlobSizeOpt::Unknown => *self = size,
+        };
+        Ok(())
+    }
+}
+
+impl From<BaoBlobSize> for BaoBlobSizeOpt {
+    fn from(size: BaoBlobSize) -> Self {
+        match size {
+            BaoBlobSize::Unverified(x) => Self::Unverified(x),
+            BaoBlobSize::Verified(x) => Self::Verified(x),
+        }
+    }
+}
+
 /// Events from observing a local bitfield
 #[derive(Debug, PartialEq, Eq, derive_more::From)]
 pub enum BitfieldEvent {
@@ -98,32 +156,32 @@ pub enum BitfieldEvent {
     Update(BitfieldUpdate),
 }
 
+/// The state of a bitfield
+#[derive(Debug, PartialEq, Eq)]
+pub struct BitfieldState {
+    /// The ranges that are set
+    pub ranges: ChunkRanges,
+    /// Whatever size information is available
+    pub size: BaoBlobSizeOpt,
+}
+
 /// An update to a bitfield
 #[derive(Debug, PartialEq, Eq)]
 pub struct BitfieldUpdate {
     /// The ranges that were added
     pub added: ChunkRanges,
     /// The ranges that were removed
     pub removed: ChunkRanges,
-    /// The total size of the bitfield in bytes
-    pub size: u64,
-}
-
-/// The state of a bitfield
-#[derive(Debug, PartialEq, Eq)]
-pub struct BitfieldState {
-    /// The ranges that are set
-    pub ranges: ChunkRanges,
-    /// The total size of the bitfield in bytes
-    pub size: u64,
+    /// Possible update to the size information
+    pub size: BaoBlobSizeOpt,
 }
 
 impl BitfieldState {
     /// State for a completely unknown bitfield
     pub fn unknown() -> Self {
         Self {
             ranges: ChunkRanges::empty(),
-            size: u64::MAX,
+            size: BaoBlobSizeOpt::Unknown,
         }
     }
 }
@@ -345,7 +403,7 @@ impl BitfieldSubscription for TestBitfieldSubscription {
             futures_lite::stream::once(
                 BitfieldState {
                     ranges,
-                    size: 1024 * 1024 * 1024 * 1024 * 1024,
+                    size: BaoBlobSizeOpt::Unknown,
                 }
                 .into(),
             )
@@ -375,7 +433,24 @@ impl<S> SimpleBitfieldSubscription<S> {
 
 async fn get_valid_ranges_local<S: Store>(hash: &Hash, store: S) -> anyhow::Result<BitfieldEvent> {
     if let Some(entry) = store.get_mut(hash).await? {
-        let (ranges, size) = crate::get::db::valid_ranges_and_size::<S>(&entry).await?;
+        let ranges = crate::get::db::valid_ranges::<S>(&entry).await?;
+        let size = entry.size();
+        let size = match size {
+            size @ BaoBlobSize::Unverified(value) => {
+                if let Some(last_chunk) = ChunkNum::chunks(value).0.checked_sub(1).map(ChunkNum) {
+                    if ranges.contains(&last_chunk) {
+                        BaoBlobSizeOpt::Verified(value)
+                    } else {
+                        size.into()
+                    }
+                } else {
+                    // this branch is just for size == 0
+                    // todo: return BaoBlobSize::Verified(0) if the hash is the hash of the empty blob
+                    BaoBlobSizeOpt::Unknown
+                }
+            }
+            size => size.into(),
+        };
         Ok(BitfieldState { ranges, size }.into())
     } else {
         Ok(BitfieldState::unknown().into())
@@ -391,7 +466,11 @@ async fn get_valid_ranges_remote(
     let (size, _) = crate::get::request::get_verified_size(&conn, hash).await?;
     let chunks = (size + 1023) / 1024;
     let ranges = ChunkRanges::from(ChunkNum(0)..ChunkNum(chunks));
-    Ok(BitfieldState { ranges, size }.into())
+    Ok(BitfieldState {
+        ranges,
+        size: BaoBlobSizeOpt::Verified(size),
+    }
+    .into())
 }
 
 impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {

src/downloader2/actor.rs

@@ -267,9 +267,16 @@ async fn peer_download<S: Store>(
                         batch.push(parent.into());
                     }
                     BaoContentItem::Leaf(leaf) => {
-                        let start_chunk = leaf.offset / 1024;
-                        let added =
-                            ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16));
+                        let size_chunks = ChunkNum::chunks(size);
+                        let start_chunk = ChunkNum::full_chunks(leaf.offset);
+                        let end_chunk =
+                            ChunkNum::full_chunks(leaf.offset + 16 * 1024).min(size_chunks);
+                        let last_chunk = size_chunks
+                            .0
+                            .checked_sub(1)
+                            .map(ChunkNum)
+                            .expect("Size must not be 0");
+                        let added = ChunkRanges::from(start_chunk..end_chunk);
                         sender
                             .send(Command::ChunksDownloaded {
                                 time: start.elapsed(),
@@ -281,6 +288,11 @@ async fn peer_download<S: Store>(
                             .ok();
                         batch.push(leaf.into());
                         writer.write_batch(size, std::mem::take(&mut batch)).await?;
+                        let size = if added.contains(&last_chunk) {
+                            BaoBlobSizeOpt::Verified(size)
+                        } else {
+                            BaoBlobSizeOpt::Unverified(size)
+                        };
                         sender
                             .send(Command::BitfieldInfo {
                                 peer: BitfieldPeer::Local,