Unify BitfieldEvents and start adding size (not used yet)

Author: rklaehn

examples/multiprovider.rs

@@ -117,7 +117,7 @@ impl BlobDownloadProgress {
 
     fn update(&mut self, ev: BitfieldEvent) {
         match ev {
-            BitfieldEvent::State { ranges } => {
+            BitfieldEvent::State { ranges, .. } => {
                 self.current = ranges;
             }
             BitfieldEvent::Update { added, removed } => {

src/downloader2.rs

@@ -17,7 +17,7 @@ use std::{
     io,
     marker::PhantomData,
     sync::Arc,
-    time::Instant,
+    time::Instant, u64,
 };
 
 use crate::{
@@ -89,12 +89,14 @@ pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static {
 pub type BoxedBitfieldSubscription = Box<dyn BitfieldSubscription>;
 
 /// Events from observing a local bitfield
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum BitfieldEvent {
     /// The full state of the bitfield
     State {
         /// The entire bitfield
         ranges: ChunkRanges,
+        /// The most precise known size of the blob
+        size: u64,
     },
     /// An update to the bitfield
     Update {
@@ -319,7 +321,7 @@ impl BitfieldSubscription for TestBitfieldSubscription {
             }
         };
         Box::pin(
-            futures_lite::stream::once(BitfieldEvent::State { ranges })
+            futures_lite::stream::once(BitfieldEvent::State { ranges, size: 1024 * 1024 * 1024 * 1024 * 1024 })
                 .chain(futures_lite::stream::pending()),
         )
     }
@@ -400,7 +402,7 @@ impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
                     Ok(ev) => ev,
                     Err(_) => ChunkRanges::empty(),
                 };
-                BitfieldEvent::State { ranges }
+                BitfieldEvent::State { ranges, size: u64::MAX }
             }
             .into_stream(),
         )

src/downloader2/actor.rs

@@ -142,18 +142,8 @@ impl<S: Store> DownloaderActor<S> {
                 let send = self.command_tx.clone();
                 let mut stream = self.subscribe_bitfield.subscribe(peer, hash);
                 let task = spawn(async move {
-                    while let Some(ev) = stream.next().await {
-                        let cmd = match ev {
-                            BitfieldEvent::State { ranges } => {
-                                Command::Bitfield { peer, hash, ranges }
-                            }
-                            BitfieldEvent::Update { added, removed } => Command::BitfieldUpdate {
-                                peer,
-                                hash,
-                                added,
-                                removed,
-                            },
-                        };
+                    while let Some(event) = stream.next().await {
+                        let cmd = Command::BitfieldInfo { peer, hash, event };
                         send.send(cmd).await.ok();
                     }
                 });
@@ -204,23 +194,11 @@ impl<S: Store> DownloaderActor<S> {
                     done.send(()).ok();
                 }
             }
-            Event::LocalBitfield { id, ranges } => {
-                let Some(sender) = self.observers.get(&id) else {
-                    return;
-                };
-                if sender.try_send(BitfieldEvent::State { ranges }).is_err() {
-                    // the observer has been dropped
-                    self.observers.remove(&id);
-                }
-            }
-            Event::LocalBitfieldUpdate { id, added, removed } => {
+            Event::LocalBitfieldInfo { id, event } => {
                 let Some(sender) = self.observers.get(&id) else {
                     return;
                 };
-                if sender
-                    .try_send(BitfieldEvent::Update { added, removed })
-                    .is_err()
-                {
+                if sender.try_send(event).is_err() {
                     // the observer has been dropped
                     self.observers.remove(&id);
                 }
@@ -304,11 +282,13 @@ async fn peer_download<S: Store>(
                         batch.push(leaf.into());
                         writer.write_batch(size, std::mem::take(&mut batch)).await?;
                         sender
-                            .send(Command::BitfieldUpdate {
+                            .send(Command::BitfieldInfo {
                                 peer: BitfieldPeer::Local,
                                 hash,
-                                added,
-                                removed: ChunkRanges::empty(),
+                                event: BitfieldEvent::Update {
+                                    added,
+                                    removed: ChunkRanges::empty(),
+                                },
                             })
                             .await
                             .ok();