Unify BitfieldEvent and BitfieldSubscriptionEvent

Author: rklaehn

examples/multiprovider.rs

@@ -6,7 +6,7 @@ use console::Term;
 use iroh::{NodeId, SecretKey};
 use iroh_blobs::{
     downloader2::{
-        print_bitmap, DownloadRequest, Downloader, ObserveEvent, ObserveRequest,
+        print_bitmap, BitfieldEvent, DownloadRequest, Downloader, ObserveRequest,
         StaticContentDiscovery,
     },
     store::Store,
@@ -115,12 +115,12 @@ impl BlobDownloadProgress {
         }
     }
 
-    fn update(&mut self, ev: ObserveEvent) {
+    fn update(&mut self, ev: BitfieldEvent) {
         match ev {
-            ObserveEvent::Bitfield { ranges } => {
+            BitfieldEvent::State { ranges } => {
                 self.current = ranges;
             }
-            ObserveEvent::BitfieldUpdate { added, removed } => {
+            BitfieldEvent::Update { added, removed } => {
                 self.current |= added;
                 self.current -= removed;
             }

src/downloader2.rs

@@ -82,43 +82,22 @@ struct BitfieldSubscriptionId(u64);
 /// A pluggable bitfield subscription mechanism
 pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static {
     /// Subscribe to a bitfield
-    fn subscribe(
-        &mut self,
-        peer: BitfieldPeer,
-        hash: Hash,
-    ) -> BoxStream<'static, BitfieldSubscriptionEvent>;
+    fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> BoxStream<'static, BitfieldEvent>;
 }
 
 /// A boxed bitfield subscription
 pub type BoxedBitfieldSubscription = Box<dyn BitfieldSubscription>;
 
-/// An event from a bitfield subscription
-#[derive(Debug)]
-pub enum BitfieldSubscriptionEvent {
-    /// Set the bitfield to the given ranges
-    Bitfield {
-        /// The entire bitfield
-        ranges: ChunkRanges,
-    },
-    /// Update the bitfield with the given ranges
-    BitfieldUpdate {
-        /// The ranges that were added
-        added: ChunkRanges,
-        /// The ranges that were removed
-        removed: ChunkRanges,
-    },
-}
-
 /// Events from observing a local bitfield
 #[derive(Debug)]
-pub enum ObserveEvent {
-    /// Set the bitfield to the given ranges
-    Bitfield {
+pub enum BitfieldEvent {
+    /// The full state of the bitfield
+    State {
         /// The entire bitfield
         ranges: ChunkRanges,
     },
-    /// Update the bitfield with the given ranges
-    BitfieldUpdate {
+    /// An update to the bitfield
+    Update {
         /// The ranges that were added
         added: ChunkRanges,
         /// The ranges that were removed
@@ -281,7 +260,7 @@ impl Downloader {
     pub async fn observe(
         &self,
         request: ObserveRequest,
-    ) -> anyhow::Result<tokio::sync::mpsc::Receiver<ObserveEvent>> {
+    ) -> anyhow::Result<tokio::sync::mpsc::Receiver<BitfieldEvent>> {
         let (send, recv) = tokio::sync::mpsc::channel(request.buffer);
         self.send
             .send(UserCommand::Observe { request, send })
@@ -332,19 +311,15 @@ impl Downloader {
 struct TestBitfieldSubscription;
 
 impl BitfieldSubscription for TestBitfieldSubscription {
-    fn subscribe(
-        &mut self,
-        peer: BitfieldPeer,
-        _hash: Hash,
-    ) -> BoxStream<'static, BitfieldSubscriptionEvent> {
+    fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> BoxStream<'static, BitfieldEvent> {
         let ranges = match peer {
             BitfieldPeer::Local => ChunkRanges::empty(),
             BitfieldPeer::Remote(_) => {
                 ChunkRanges::from(ChunkNum(0)..ChunkNum(1024 * 1024 * 1024 * 1024))
             }
         };
         Box::pin(
-            futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges })
+            futures_lite::stream::once(BitfieldEvent::State { ranges })
                 .chain(futures_lite::stream::pending()),
         )
     }
@@ -389,11 +364,7 @@ async fn get_valid_ranges_remote(
 }
 
 impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
-    fn subscribe(
-        &mut self,
-        peer: BitfieldPeer,
-        hash: Hash,
-    ) -> BoxStream<'static, BitfieldSubscriptionEvent> {
+    fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> BoxStream<'static, BitfieldEvent> {
         let (send, recv) = tokio::sync::oneshot::channel();
         match peer {
             BitfieldPeer::Local => {
@@ -429,7 +400,7 @@ impl<S: Store> BitfieldSubscription for SimpleBitfieldSubscription<S> {
                     Ok(ev) => ev,
                     Err(_) => ChunkRanges::empty(),
                 };
-                BitfieldSubscriptionEvent::Bitfield { ranges }
+                BitfieldEvent::State { ranges }
             }
             .into_stream(),
         )
@@ -549,7 +520,9 @@ mod tests {
     #[tokio::test]
     async fn test_valid_ranges() -> TestResult<()> {
         let store = crate::store::mem::Store::new();
-        let tt = store.import_bytes(vec![0u8;1025].into(), crate::BlobFormat::Raw).await?;
+        let tt = store
+            .import_bytes(vec![0u8; 1025].into(), crate::BlobFormat::Raw)
+            .await?;
         let entry = store.get_mut(tt.hash()).await?.unwrap();
         let valid = crate::get::db::valid_ranges::<crate::store::mem::Store>(&entry).await?;
         assert!(valid == ChunkRanges::from(ChunkNum(0)..ChunkNum(2)));

src/downloader2/actor.rs

@@ -10,7 +10,7 @@ pub(super) enum UserCommand {
     },
     Observe {
         request: ObserveRequest,
-        send: tokio::sync::mpsc::Sender<ObserveEvent>,
+        send: tokio::sync::mpsc::Sender<BitfieldEvent>,
     },
 }
 
@@ -38,7 +38,7 @@ pub(super) struct DownloaderActor<S> {
     /// Id generator for observe ids
     observe_id_gen: IdGenerator<ObserveId>,
     /// Observers
-    observers: BTreeMap<ObserveId, tokio::sync::mpsc::Sender<ObserveEvent>>,
+    observers: BTreeMap<ObserveId, tokio::sync::mpsc::Sender<BitfieldEvent>>,
     /// The time when the actor was started, serves as the epoch for time messages to the state machine
     start: Instant,
 }
@@ -144,17 +144,15 @@ impl<S: Store> DownloaderActor<S> {
                 let task = spawn(async move {
                     while let Some(ev) = stream.next().await {
                         let cmd = match ev {
-                            BitfieldSubscriptionEvent::Bitfield { ranges } => {
+                            BitfieldEvent::State { ranges } => {
                                 Command::Bitfield { peer, hash, ranges }
                             }
-                            BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => {
-                                Command::BitfieldUpdate {
-                                    peer,
-                                    hash,
-                                    added,
-                                    removed,
-                                }
-                            }
+                            BitfieldEvent::Update { added, removed } => Command::BitfieldUpdate {
+                                peer,
+                                hash,
+                                added,
+                                removed,
+                            },
                         };
                         send.send(cmd).await.ok();
                     }
@@ -210,7 +208,7 @@ impl<S: Store> DownloaderActor<S> {
                 let Some(sender) = self.observers.get(&id) else {
                     return;
                 };
-                if sender.try_send(ObserveEvent::Bitfield { ranges }).is_err() {
+                if sender.try_send(BitfieldEvent::State { ranges }).is_err() {
                     // the observer has been dropped
                     self.observers.remove(&id);
                 }
@@ -220,7 +218,7 @@ impl<S: Store> DownloaderActor<S> {
                     return;
                 };
                 if sender
-                    .try_send(ObserveEvent::BitfieldUpdate { added, removed })
+                    .try_send(BitfieldEvent::Update { added, removed })
                     .is_err()
                 {
                     // the observer has been dropped