From fa8fab43c3da600d2b4b43b17a7b7c980ce83ab4 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 29 Jan 2025 14:38:30 +0200 Subject: [PATCH 01/47] start downloader engine --- src/downloader2.rs | 220 +++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + 2 files changed, 222 insertions(+) create mode 100644 src/downloader2.rs diff --git a/src/downloader2.rs b/src/downloader2.rs new file mode 100644 index 000000000..0f57d404b --- /dev/null +++ b/src/downloader2.rs @@ -0,0 +1,220 @@ +//! Downloader version that supports range downloads and downloads from multiple sources. + +use std::collections::{BTreeMap, VecDeque}; + +use anyhow::Context; +use bao_tree::ChunkRanges; +use chrono::Duration; +use futures_lite::Stream; +use iroh::NodeId; +use range_collections::range_set::RangeSetRange; +use serde::{Deserialize, Serialize}; +use crate::Hash; + +/// Identifier for a download intent. +#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, derive_more::Display)] +pub struct IntentId(pub u64); + +/// Announce kind +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] +pub enum AnnounceKind { + /// The peer supposedly has some of the data. + Partial = 0, + /// The peer supposedly has the complete data. + Complete, +} + +struct FindPeersOpts { + /// Kind of announce + kind: AnnounceKind, +} + +/// A pluggable content discovery mechanism +trait ContentDiscovery { + /// Find peers that have the given blob. + /// + /// The returned stream is a handle for the discovery task. It should be an + /// infinite stream that only stops when it is dropped. + fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Unpin; +} + +/// Global information about a peer +#[derive(Debug, Default)] +struct PeerState { + /// Executed downloads, to calculate the average download speed. + /// + /// This gets updated as soon as possible when new data has been downloaded. + download_history: VecDeque<(Duration, (u64, u64))>, +} + +/// Information about one blob on one peer +struct PeerBlobState { + /// + subscription_id: u64, + /// chunk ranges this peer reports to have + ranges: ChunkRanges, +} + +struct DownloadRequest { + /// The blob we are downloading + hash: Hash, + /// The ranges we are interested in + ranges: ChunkRanges, +} + +struct DownloadState { + /// The request this state is for + request: DownloadRequest, + +} + +struct DownloaderState { + peers: BTreeMap, + bitmaps: BTreeMap<(NodeId, Hash), PeerBlobState>, + downloads: BTreeMap, + me: NodeId, + next_subscription_id: u64, +} + +impl DownloaderState { + fn new(me: NodeId) -> Self { + Self { + me, + peers: BTreeMap::new(), + downloads: BTreeMap::new(), + bitmaps: BTreeMap::new(), + next_subscription_id: 0, + } + } +} + +enum Command { + /// A request to start a download. + StartDownload { + /// The download request + request: DownloadRequest, + /// The unique id, to be assigned by the caller + id: IntentId, + }, + /// A request to abort a download. + StopDownload { + id: IntentId, + }, + /// A bitmap for a blob and a peer + Bitmap { + /// The peer that sent the bitmap. + peer: NodeId, + /// The blob for which the bitmap is + hash: Hash, + /// The complete bitmap + bitmap: ChunkRanges, + }, + /// An update of a bitmap for a hash + /// + /// This is used both to update the bitmap of remote peers, and to update + /// the local bitmap. + BitmapUpdate { + /// The peer that sent the update. + peer: NodeId, + /// The blob that was updated. + hash: Hash, + /// The ranges that were added + added: ChunkRanges, + /// The ranges that were removed + removed: ChunkRanges, + }, + /// A chunk was downloaded + ChunksDownloaded { + /// Time when the download was received + time: Duration, + /// The peer that sent the chunk + peer: NodeId, + /// The blob that was downloaded + hash: Hash, + /// The ranges that were added locally + added: ChunkRanges, + } +} + +enum Event { + SubscribeBitmap { + peer: NodeId, + hash: Hash, + subscription_id: u64, + }, + UnsubscribeBitmap { + subscription_id: u64, + }, + Error { + message: String, + }, +} + +impl DownloaderState { + + fn apply(&mut self, cmd: Command, events: &mut Vec) { + if let Err(cause) = self.apply0(cmd, events) { + events.push(Event::Error { message: format!("{cause}") }); + } + } + + fn next_subscription_id(&mut self) -> u64 { + let id = self.next_subscription_id; + self.next_subscription_id += 1; + id + } + + fn apply0(&mut self, cmd: Command, events: &mut Vec) -> anyhow::Result<()> { + match cmd { + Command::StartDownload { request, id } => { + if self.downloads.contains_key(&id) { + anyhow::bail!("duplicate download request {id}"); + } + if !self.bitmaps.contains_key(&(self.me, request.hash)) { + let subscription_id = self.next_subscription_id(); + events.push(Event::SubscribeBitmap { + peer: self.me, + hash: request.hash, + subscription_id, + }); + self.bitmaps.insert((self.me, request.hash), PeerBlobState { subscription_id, ranges: ChunkRanges::empty() }); + } + self.downloads.insert(id, DownloadState { request }); + } + Command::StopDownload { id } => { + self.downloads.remove(&id); + } + Command::Bitmap { peer, hash, bitmap } => { + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer} and hash {hash}"))?; + state.ranges = bitmap; + } + Command::BitmapUpdate { peer, hash, added, removed } => { + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer} and hash {hash}"))?; + state.ranges |= added; + state.ranges &= !removed; + } + Command::ChunksDownloaded { time, peer, hash, added } => { + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("chunks downloaded for unknown peer {peer} and hash {hash}"))?; + let total_downloaded = total_chunks(&added).context("open range")?; + let total_before = total_chunks(&state.ranges).context("open range")?; + state.ranges |= added; + let total_after = total_chunks(&state.ranges).context("open range")?; + let useful_downloaded = total_after - total_before; + let peer = self.peers.entry(peer).or_default(); + peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); + } + } + Ok(()) + } +} + +fn total_chunks(chunks: &ChunkRanges) -> Option { + let mut total = 0; + for range in chunks.iter() { + match range { + RangeSetRange::RangeFrom(_range) => return None, + RangeSetRange::Range(range) => total += range.end.0 - range.start.0, + } + } + Some(total) +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 7091ad795..5901bc93b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,6 +32,8 @@ pub mod cli; #[cfg(feature = "downloader")] pub mod downloader; +#[cfg(feature = "downloader")] +pub mod downloader2; pub mod export; pub mod format; pub mod get; From 23817cc0ed291774883f002ce5b212fb27782ab3 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 29 Jan 2025 15:13:57 +0200 Subject: [PATCH 02/47] WIP --- src/downloader2.rs | 111 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 92 insertions(+), 19 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 0f57d404b..6e7092dd0 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -11,9 +11,9 @@ use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use crate::Hash; -/// Identifier for a download intent. -#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, derive_more::Display)] -pub struct IntentId(pub u64); +type DownloadId = u64; +type BitmapSubscriptionId = u64; +type DiscoveryId = u64; /// Announce kind #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] @@ -49,8 +49,10 @@ struct PeerState { /// Information about one blob on one peer struct PeerBlobState { - /// - subscription_id: u64, + /// The subscription id for the subscription + subscription_id: BitmapSubscriptionId, + /// The number of subscriptions this peer has + subscription_count: usize, /// chunk ranges this peer reports to have ranges: ChunkRanges, } @@ -69,11 +71,20 @@ struct DownloadState { } struct DownloaderState { + // my own node id + me: NodeId, + // all peers I am tracking for any download peers: BTreeMap, + // all bitmaps I am tracking, both for myself and for remote peers bitmaps: BTreeMap<(NodeId, Hash), PeerBlobState>, - downloads: BTreeMap, - me: NodeId, - next_subscription_id: u64, + // all active downloads + downloads: BTreeMap, + // discovery tasks + discovery: BTreeMap, + // the next subscription id + next_subscription_id: BitmapSubscriptionId, + // the next discovery id + next_discovery_id: u64, } impl DownloaderState { @@ -83,7 +94,9 @@ impl DownloaderState { peers: BTreeMap::new(), downloads: BTreeMap::new(), bitmaps: BTreeMap::new(), + discovery: BTreeMap::new(), next_subscription_id: 0, + next_discovery_id: 0, } } } @@ -94,11 +107,11 @@ enum Command { /// The download request request: DownloadRequest, /// The unique id, to be assigned by the caller - id: IntentId, + id: u64, }, /// A request to abort a download. StopDownload { - id: IntentId, + id: u64, }, /// A bitmap for a blob and a peer Bitmap { @@ -133,6 +146,10 @@ enum Command { hash: Hash, /// The ranges that were added locally added: ChunkRanges, + }, + /// Stop tracking a peer for all blobs, for whatever reason + DropPeer { + peer: NodeId, } } @@ -140,10 +157,14 @@ enum Event { SubscribeBitmap { peer: NodeId, hash: Hash, - subscription_id: u64, + id: u64, }, UnsubscribeBitmap { - subscription_id: u64, + id: u64, + }, + StartDiscovery { + hash: Hash, + id: u64, }, Error { message: String, @@ -164,25 +185,71 @@ impl DownloaderState { id } + fn next_discovery_id(&mut self) -> u64 { + let id = self.next_discovery_id; + self.next_discovery_id += 1; + id + } + + fn count_providers(&self, hash: Hash) -> usize { + self.bitmaps.iter().filter(|((peer, x), _)| *peer != self.me && *x == hash).count() + } + fn apply0(&mut self, cmd: Command, events: &mut Vec) -> anyhow::Result<()> { match cmd { Command::StartDownload { request, id } => { - if self.downloads.contains_key(&id) { - anyhow::bail!("duplicate download request {id}"); - } - if !self.bitmaps.contains_key(&(self.me, request.hash)) { + // ids must be uniquely assigned by the caller! + anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); + // either we have a subscription for this blob, or we have to create one + if let Some(state) = self.bitmaps.get_mut(&(self.me, request.hash)) { + // just increment the count + state.subscription_count += 1; + } else { + // create a new subscription let subscription_id = self.next_subscription_id(); events.push(Event::SubscribeBitmap { peer: self.me, hash: request.hash, - subscription_id, + id: subscription_id, + }); + self.bitmaps.insert((self.me, request.hash), PeerBlobState { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() }); + } + if !self.discovery.contains_key(&request.hash) { + // start a discovery task + let discovery_id = self.next_discovery_id(); + events.push(Event::StartDiscovery { + hash: request.hash, + id: discovery_id, }); - self.bitmaps.insert((self.me, request.hash), PeerBlobState { subscription_id, ranges: ChunkRanges::empty() }); + self.discovery.insert(request.hash, discovery_id); } self.downloads.insert(id, DownloadState { request }); } Command::StopDownload { id } => { - self.downloads.remove(&id); + let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; + self.bitmaps.retain(|(_peer, hash), state| { + if *hash == removed.request.hash { + state.subscription_count -= 1; + if state.subscription_count == 0 { + events.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + return false; + } + } + true + }); + } + Command::DropPeer { peer } => { + anyhow::ensure!(peer != self.me, "not a remote peer"); + self.bitmaps.retain(|(p, _), state| { + if *p == peer { + // todo: should we emit unsubscribe events here? + events.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + return false; + } else { + return true; + } + }); + self.peers.remove(&peer); } Command::Bitmap { peer, hash, bitmap } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer} and hash {hash}"))?; @@ -217,4 +284,10 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { } } Some(total) +} + +#[cfg(test)] +mod tests { + + } \ No newline at end of file From 090e54be4f36d01c2005684c296e50c5f5b0aef5 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 29 Jan 2025 18:09:42 +0200 Subject: [PATCH 03/47] WIP --- src/downloader2.rs | 196 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 181 insertions(+), 15 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 6e7092dd0..27350f90f 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -4,7 +4,7 @@ use std::collections::{BTreeMap, VecDeque}; use anyhow::Context; use bao_tree::ChunkRanges; -use chrono::Duration; +use std::time::Duration; use futures_lite::Stream; use iroh::NodeId; use range_collections::range_set::RangeSetRange; @@ -57,6 +57,12 @@ struct PeerBlobState { ranges: ChunkRanges, } +impl PeerBlobState { + fn new(subscription_id: BitmapSubscriptionId) -> Self { + Self { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() } + } +} + struct DownloadRequest { /// The blob we are downloading hash: Hash, @@ -67,7 +73,20 @@ struct DownloadRequest { struct DownloadState { /// The request this state is for request: DownloadRequest, + /// Ongoing downloads + downloads: BTreeMap, +} + +impl DownloadState { + + fn new(request: DownloadRequest) -> Self { + Self { request, downloads: BTreeMap::new() } + } +} +struct PeerDownloadState { + id: u64, + ranges: ChunkRanges, } struct DownloaderState { @@ -85,6 +104,8 @@ struct DownloaderState { next_subscription_id: BitmapSubscriptionId, // the next discovery id next_discovery_id: u64, + // the next peer download id + next_peer_download_id: u64, } impl DownloaderState { @@ -97,6 +118,7 @@ impl DownloaderState { discovery: BTreeMap::new(), next_subscription_id: 0, next_discovery_id: 0, + next_peer_download_id: 0, } } } @@ -150,9 +172,15 @@ enum Command { /// Stop tracking a peer for all blobs, for whatever reason DropPeer { peer: NodeId, + }, + /// A peer has been discovered + PeerDiscovered { + peer: NodeId, + hash: Hash, } } +#[derive(Debug, PartialEq, Eq)] enum Event { SubscribeBitmap { peer: NodeId, @@ -166,6 +194,20 @@ enum Event { hash: Hash, id: u64, }, + StopDiscovery { + id: u64, + }, + StartPeerDownload { + id: u64, + peer: NodeId, + ranges: ChunkRanges, + }, + StopPeerDownload { + id: u64, + }, + DownloadComplete { + id: u64, + }, Error { message: String, }, @@ -173,9 +215,15 @@ enum Event { impl DownloaderState { - fn apply(&mut self, cmd: Command, events: &mut Vec) { - if let Err(cause) = self.apply0(cmd, events) { - events.push(Event::Error { message: format!("{cause}") }); + fn apply_and_get_evs(&mut self, cmd: Command) -> Vec { + let mut evs = vec![]; + self.apply(cmd, &mut evs); + evs + } + + fn apply(&mut self, cmd: Command, evs: &mut Vec) { + if let Err(cause) = self.apply0(cmd, evs) { + evs.push(Event::Error { message: format!("{cause}") }); } } @@ -191,11 +239,17 @@ impl DownloaderState { id } + fn next_peer_download_id(&mut self) -> u64 { + let id = self.next_peer_download_id; + self.next_peer_download_id += 1; + id + } + fn count_providers(&self, hash: Hash) -> usize { self.bitmaps.iter().filter(|((peer, x), _)| *peer != self.me && *x == hash).count() } - fn apply0(&mut self, cmd: Command, events: &mut Vec) -> anyhow::Result<()> { + fn apply0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! @@ -207,23 +261,23 @@ impl DownloaderState { } else { // create a new subscription let subscription_id = self.next_subscription_id(); - events.push(Event::SubscribeBitmap { + evs.push(Event::SubscribeBitmap { peer: self.me, hash: request.hash, id: subscription_id, }); - self.bitmaps.insert((self.me, request.hash), PeerBlobState { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() }); + self.bitmaps.insert((self.me, request.hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task let discovery_id = self.next_discovery_id(); - events.push(Event::StartDiscovery { + evs.push(Event::StartDiscovery { hash: request.hash, id: discovery_id, }); self.discovery.insert(request.hash, discovery_id); } - self.downloads.insert(id, DownloadState { request }); + self.downloads.insert(id, DownloadState::new(request)); } Command::StopDownload { id } => { let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; @@ -231,19 +285,40 @@ impl DownloaderState { if *hash == removed.request.hash { state.subscription_count -= 1; if state.subscription_count == 0 { - events.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); return false; } } true }); } + Command::PeerDiscovered { peer, hash } => { + anyhow::ensure!(peer != self.me, "not a remote peer"); + if self.bitmaps.contains_key(&(peer, hash)) { + // we already have a subscription for this peer + return Ok(()); + }; + // check if anybody needs this peer + if !self.downloads.iter().any(|(_id, state)| state.request.hash == hash) { + return Ok(()); + } + // create a peer state if it does not exist + let _state = self.peers.entry(peer).or_default(); + // create a new subscription + let subscription_id = self.next_subscription_id(); + evs.push(Event::SubscribeBitmap { + peer, + hash, + id: subscription_id, + }); + self.bitmaps.insert((peer, hash), PeerBlobState::new(subscription_id)); + }, Command::DropPeer { peer } => { anyhow::ensure!(peer != self.me, "not a remote peer"); self.bitmaps.retain(|(p, _), state| { if *p == peer { - // todo: should we emit unsubscribe events here? - events.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + // todo: should we emit unsubscribe evs here? + evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); return false; } else { return true; @@ -253,26 +328,70 @@ impl DownloaderState { } Command::Bitmap { peer, hash, bitmap } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer} and hash {hash}"))?; + let _chunks = total_chunks(&bitmap).context("open range")?; state.ranges = bitmap; + self.rebalance_downloads(hash, evs)?; } Command::BitmapUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer} and hash {hash}"))?; state.ranges |= added; state.ranges &= !removed; + self.rebalance_downloads(hash, evs)?; } Command::ChunksDownloaded { time, peer, hash, added } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("chunks downloaded for unknown peer {peer} and hash {hash}"))?; + anyhow::ensure!(peer != self.me, "not a remote peer"); + let state = self.bitmaps.get_mut(&(self.me, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; let total_after = total_chunks(&state.ranges).context("open range")?; let useful_downloaded = total_after - total_before; - let peer = self.peers.entry(peer).or_default(); + let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); + self.rebalance_downloads(hash, evs)?; } } Ok(()) } + + fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + let Some(self_state) = self.bitmaps.get(&(self.me, hash)) else { + // 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 mut completed = vec![]; + for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { + let remaining = &download.request.ranges - &self_state.ranges; + if remaining.is_empty() { + // cancel all downloads, if needed + evs.push(Event::DownloadComplete { id: *id }); + completed.push(*id); + continue; + } + let mut candidates = vec![]; + for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *x == hash && *peer != self.me) { + let intersection = &bitmap.ranges & &remaining; + if !intersection.is_empty() { + candidates.push((*peer, intersection)); + } + } + for (_, state) in &download.downloads { + // stop all downloads + evs.push(Event::StopPeerDownload { id: state.id }); + } + download.downloads.clear(); + for (peer, ranges) in candidates { + let id = self.next_peer_download_id; + self.next_peer_download_id += 1; + evs.push(Event::StartPeerDownload { id, peer, ranges: ranges.clone() }); + download.downloads.insert(peer, PeerDownloadState { id, ranges }); + } + }; + for id in completed { + self.downloads.remove(&id); + } + Ok(()) + } } fn total_chunks(chunks: &ChunkRanges) -> Option { @@ -288,6 +407,53 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { #[cfg(test)] mod tests { + use std::ops::Range; + + use super::*; + use bao_tree::ChunkNum; + use testresult::TestResult; + + fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { + let mut res = ChunkRanges::empty(); + for range in ranges.into_iter() { + res |= ChunkRanges::from(ChunkNum(range.start)..ChunkNum(range.end)); + } + res + } + + #[test] + fn smoke() -> TestResult<()> { + let me = "0000000000000000000000000000000000000000000000000000000000000000".parse()?; + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; + let mut state = DownloaderState::new(me); + let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); + assert!(evs.iter().filter(|e| **e == Event::StartDiscovery { hash, id: 0 }).count() == 1, "starting a download should start a discovery task"); + assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: me, hash, id: 0 }).count() == 1, "starting a download should subscribe to the local bitmap"); + println!("{evs:?}"); + let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash, bitmap: ChunkRanges::all() }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); + let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash: unknown_hash, bitmap: ChunkRanges::all() }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); + let evs = state.apply_and_get_evs(super::Command::DropPeer { peer: me }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "dropping self should produce an error!"); + let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); + let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash, bitmap: initial_bitmap.clone() }); + assert!(evs.is_empty()); + assert_eq!(state.bitmaps.get(&(me, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + let evs = state.apply_and_get_evs(super::Command::BitmapUpdate { peer: me, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + assert!(evs.is_empty()); + assert_eq!(state.bitmaps.get(&(me, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); + let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)) }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); + let evs = state.apply_and_get_evs(super::Command::PeerDiscovered { peer: peer_a, hash }); + assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: peer_a, hash, id: 1 }).count() == 1, "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: peer_a, hash, bitmap: chunk_ranges([0..64]) }); + assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); + let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(64)) }); + println!("{evs:?}"); + Ok(()) + } - } \ No newline at end of file From 8caf03498759bc5a7ce593cf25f85463888f7186 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 30 Jan 2025 14:07:28 +0200 Subject: [PATCH 04/47] WIP --- src/downloader2.rs | 463 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 358 insertions(+), 105 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 27350f90f..bea1c3767 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -1,19 +1,37 @@ //! Downloader version that supports range downloads and downloads from multiple sources. - +//! +//! # Structure +//! +//! The [DownloaderState] is a synchronous state machine containing the logic. +//! It gets commands and produces events. It does not do any IO and also does +//! not have any time dependency. So [DownloaderState::apply_and_get_evs] is a +//! pure function of the state and the command and can therefore be tested +//! easily. +//! +//! In several places, requests are identified by a unique id. It is the responsibility +//! of the caller to generate unique ids. We could use random uuids here, but using +//! integers simplifies testing. +//! +//! The [DownloaderDriver] is the asynchronous driver for the state machine. It +//! owns the actual tasks that perform IO. use std::collections::{BTreeMap, VecDeque}; +use crate::{store::Store, Hash}; use anyhow::Context; use bao_tree::ChunkRanges; -use std::time::Duration; use futures_lite::Stream; use iroh::NodeId; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; -use crate::Hash; +use std::time::Duration; +use tokio::task::JoinSet; +use tokio_util::task::AbortOnDropHandle; +/// todo: make newtypes? type DownloadId = u64; type BitmapSubscriptionId = u64; type DiscoveryId = u64; +type PeerDownloadId = u64; /// Announce kind #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] @@ -32,10 +50,11 @@ struct FindPeersOpts { /// A pluggable content discovery mechanism trait ContentDiscovery { /// Find peers that have the given blob. - /// + /// /// The returned stream is a handle for the discovery task. It should be an /// infinite stream that only stops when it is dropped. - fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Unpin; + fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) + -> impl Stream + Unpin; } /// Global information about a peer @@ -59,7 +78,11 @@ struct PeerBlobState { impl PeerBlobState { fn new(subscription_id: BitmapSubscriptionId) -> Self { - Self { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() } + Self { + subscription_id, + subscription_count: 1, + ranges: ChunkRanges::empty(), + } } } @@ -74,32 +97,43 @@ struct DownloadState { /// The request this state is for request: DownloadRequest, /// Ongoing downloads - downloads: BTreeMap, + peer_downloads: BTreeMap, } impl DownloadState { - fn new(request: DownloadRequest) -> Self { - Self { request, downloads: BTreeMap::new() } + Self { + request, + peer_downloads: BTreeMap::new(), + } } } struct PeerDownloadState { - id: u64, + id: PeerDownloadId, ranges: ChunkRanges, } struct DownloaderState { - // my own node id - me: NodeId, // all peers I am tracking for any download peers: BTreeMap, // all bitmaps I am tracking, both for myself and for remote peers - bitmaps: BTreeMap<(NodeId, Hash), PeerBlobState>, + // + // each item here corresponds to an active subscription + bitmaps: BTreeMap<(Option, Hash), PeerBlobState>, // all active downloads + // + // these are user downloads. each user download gets split into one or more + // peer downloads. downloads: BTreeMap, // discovery tasks + // + // there is a discovery task for each blob we are interested in. discovery: BTreeMap, + + // Counters to generate unique ids for various requests. + // We could use uuid here, but using integers simplifies testing. + // // the next subscription id next_subscription_id: BitmapSubscriptionId, // the next discovery id @@ -109,9 +143,8 @@ struct DownloaderState { } impl DownloaderState { - fn new(me: NodeId) -> Self { + fn new() -> Self { Self { - me, peers: BTreeMap::new(), downloads: BTreeMap::new(), bitmaps: BTreeMap::new(), @@ -124,21 +157,19 @@ impl DownloaderState { } enum Command { - /// A request to start a download. + /// A user request to start a download. StartDownload { /// The download request request: DownloadRequest, /// The unique id, to be assigned by the caller id: u64, }, - /// A request to abort a download. - StopDownload { - id: u64, - }, - /// A bitmap for a blob and a peer + /// A user request to abort a download. + StopDownload { id: u64 }, + /// A full bitmap for a blob and a peer Bitmap { /// The peer that sent the bitmap. - peer: NodeId, + peer: Option, /// The blob for which the bitmap is hash: Hash, /// The complete bitmap @@ -150,7 +181,7 @@ enum Command { /// the local bitmap. BitmapUpdate { /// The peer that sent the update. - peer: NodeId, + peer: Option, /// The blob that was updated. hash: Hash, /// The ranges that were added @@ -170,63 +201,51 @@ enum Command { added: ChunkRanges, }, /// Stop tracking a peer for all blobs, for whatever reason - DropPeer { - peer: NodeId, - }, + DropPeer { peer: NodeId }, /// A peer has been discovered - PeerDiscovered { - peer: NodeId, - hash: Hash, - } + PeerDiscovered { peer: NodeId, hash: Hash }, } #[derive(Debug, PartialEq, Eq)] enum Event { SubscribeBitmap { - peer: NodeId, + peer: Option, hash: Hash, + /// The unique id of the subscription id: u64, }, UnsubscribeBitmap { + /// The unique id of the subscription id: u64, }, StartDiscovery { hash: Hash, + /// The unique id of the discovery task id: u64, }, StopDiscovery { + /// The unique id of the discovery task id: u64, }, StartPeerDownload { + /// The unique id of the peer download task id: u64, peer: NodeId, ranges: ChunkRanges, }, StopPeerDownload { + /// The unique id of the peer download task id: u64, }, DownloadComplete { + /// The unique id of the user download id: u64, }, - Error { - message: String, - }, + /// An error that stops processing the command + Error { message: String }, } impl DownloaderState { - - fn apply_and_get_evs(&mut self, cmd: Command) -> Vec { - let mut evs = vec![]; - self.apply(cmd, &mut evs); - evs - } - - fn apply(&mut self, cmd: Command, evs: &mut Vec) { - if let Err(cause) = self.apply0(cmd, evs) { - evs.push(Event::Error { message: format!("{cause}") }); - } - } - fn next_subscription_id(&mut self) -> u64 { let id = self.next_subscription_id; self.next_subscription_id += 1; @@ -246,27 +265,51 @@ impl DownloaderState { } fn count_providers(&self, hash: Hash) -> usize { - self.bitmaps.iter().filter(|((peer, x), _)| *peer != self.me && *x == hash).count() + self.bitmaps + .iter() + .filter(|((peer, x), _)| peer.is_some() && *x == hash) + .count() } + /// Apply a command and return the events that were generated + fn apply_and_get_evs(&mut self, cmd: Command) -> Vec { + let mut evs = vec![]; + self.apply(cmd, &mut evs); + evs + } + + /// Apply a command + fn apply(&mut self, cmd: Command, evs: &mut Vec) { + if let Err(cause) = self.apply0(cmd, evs) { + evs.push(Event::Error { + message: format!("{cause}"), + }); + } + } + + /// Apply a command and bail out on error fn apply0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! - anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); + anyhow::ensure!( + !self.downloads.contains_key(&id), + "duplicate download request {id}" + ); // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitmaps.get_mut(&(self.me, request.hash)) { + if let Some(state) = self.bitmaps.get_mut(&(None, request.hash)) { // just increment the count state.subscription_count += 1; } else { // create a new subscription let subscription_id = self.next_subscription_id(); evs.push(Event::SubscribeBitmap { - peer: self.me, + peer: None, hash: request.hash, id: subscription_id, }); - self.bitmaps.insert((self.me, request.hash), PeerBlobState::new(subscription_id)); + self.bitmaps + .insert((None, request.hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task @@ -280,12 +323,17 @@ impl DownloaderState { self.downloads.insert(id, DownloadState::new(request)); } Command::StopDownload { id } => { - let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; + let removed = self + .downloads + .remove(&id) + .context(format!("removed unknown download {id}"))?; self.bitmaps.retain(|(_peer, hash), state| { if *hash == removed.request.hash { state.subscription_count -= 1; if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitmap { + id: state.subscription_id, + }); return false; } } @@ -293,13 +341,16 @@ impl DownloaderState { }); } Command::PeerDiscovered { peer, hash } => { - anyhow::ensure!(peer != self.me, "not a remote peer"); - if self.bitmaps.contains_key(&(peer, hash)) { + if self.bitmaps.contains_key(&(Some(peer), hash)) { // we already have a subscription for this peer return Ok(()); }; // check if anybody needs this peer - if !self.downloads.iter().any(|(_id, state)| state.request.hash == hash) { + if !self + .downloads + .iter() + .any(|(_id, state)| state.request.hash == hash) + { return Ok(()); } // create a peer state if it does not exist @@ -307,18 +358,20 @@ impl DownloaderState { // create a new subscription let subscription_id = self.next_subscription_id(); evs.push(Event::SubscribeBitmap { - peer, + peer: Some(peer), hash, id: subscription_id, }); - self.bitmaps.insert((peer, hash), PeerBlobState::new(subscription_id)); - }, + self.bitmaps + .insert((Some(peer), hash), PeerBlobState::new(subscription_id)); + } Command::DropPeer { peer } => { - anyhow::ensure!(peer != self.me, "not a remote peer"); self.bitmaps.retain(|(p, _), state| { - if *p == peer { + if *p == Some(peer) { // todo: should we emit unsubscribe evs here? - evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitmap { + id: state.subscription_id, + }); return false; } else { return true; @@ -327,27 +380,46 @@ impl DownloaderState { self.peers.remove(&peer); } Command::Bitmap { peer, hash, bitmap } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer} and hash {hash}"))?; + let state = self + .bitmaps + .get_mut(&(peer, hash)) + .context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; state.ranges = bitmap; self.rebalance_downloads(hash, evs)?; } - Command::BitmapUpdate { peer, hash, added, removed } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer} and hash {hash}"))?; + Command::BitmapUpdate { + peer, + hash, + added, + removed, + } => { + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!( + "bitmap update for unknown peer {peer:?} and hash {hash}" + ))?; state.ranges |= added; state.ranges &= !removed; self.rebalance_downloads(hash, evs)?; } - Command::ChunksDownloaded { time, peer, hash, added } => { - anyhow::ensure!(peer != self.me, "not a remote peer"); - let state = self.bitmaps.get_mut(&(self.me, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; + Command::ChunksDownloaded { + time, + peer, + hash, + added, + } => { + let state = self.bitmaps.get_mut(&(None, hash)).context(format!( + "chunks downloaded before having local bitmap for {hash}" + ))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; let total_after = total_chunks(&state.ranges).context("open range")?; let useful_downloaded = total_after - total_before; - let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; - peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); + let peer = self.peers.get_mut(&peer).context(format!( + "performing download before having peer state for {peer}" + ))?; + peer.download_history + .push_back((time, (total_downloaded, useful_downloaded))); self.rebalance_downloads(hash, evs)?; } } @@ -355,38 +427,59 @@ impl DownloaderState { } fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(self.me, hash)) else { + let Some(self_state) = self.bitmaps.get(&(None, hash)) else { // 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 mut completed = vec![]; - for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { + for (id, download) in self + .downloads + .iter_mut() + .filter(|(_id, download)| download.request.hash == hash) + { let remaining = &download.request.ranges - &self_state.ranges; if remaining.is_empty() { // cancel all downloads, if needed + for (_, peer_download) in &download.peer_downloads { + evs.push(Event::StopPeerDownload { + id: peer_download.id, + }); + } + // notify the user that the download is complete evs.push(Event::DownloadComplete { id: *id }); + // mark the download for later removal completed.push(*id); continue; } let mut candidates = vec![]; - for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *x == hash && *peer != self.me) { + for ((peer, _), bitmap) in self + .bitmaps + .iter() + .filter(|((peer, x), _)| peer.is_some() && *x == hash) + { let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { - candidates.push((*peer, intersection)); + candidates.push((peer.unwrap(), intersection)); } } - for (_, state) in &download.downloads { + for (_, state) in &download.peer_downloads { // stop all downloads evs.push(Event::StopPeerDownload { id: state.id }); } - download.downloads.clear(); + download.peer_downloads.clear(); for (peer, ranges) in candidates { let id = self.next_peer_download_id; self.next_peer_download_id += 1; - evs.push(Event::StartPeerDownload { id, peer, ranges: ranges.clone() }); - download.downloads.insert(peer, PeerDownloadState { id, ranges }); + evs.push(Event::StartPeerDownload { + id, + peer, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); } - }; + } for id in completed { self.downloads.remove(&id); } @@ -405,6 +498,51 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { Some(total) } +struct DownloaderDriver { + state: DownloaderState, + store: S, + discovery: D, + peer_download_tasks: BTreeMap>, + discovery_tasks: BTreeMap>, + bitmap_subscription_tasks: BTreeMap>, + next_download_id: DownloadId, +} + +impl DownloaderDriver { + fn new(me: NodeId, store: S, discovery: D) -> Self { + Self { + state: DownloaderState::new(), + store, + discovery, + peer_download_tasks: BTreeMap::new(), + discovery_tasks: BTreeMap::new(), + bitmap_subscription_tasks: BTreeMap::new(), + next_download_id: 0, + } + } + + fn next_download_id(&mut self) -> DownloadId { + let id = self.next_download_id; + self.next_download_id += 1; + id + } + + fn handle_event(&mut self, ev: Event) { + println!("event: {:?}", ev); + } + + fn download(&mut self, request: DownloadRequest) -> DownloadId { + let id = self.next_download_id(); + let evs = self + .state + .apply_and_get_evs(Command::StartDownload { request, id }); + for ev in evs { + self.handle_event(ev); + } + id + } +} + #[cfg(test)] mod tests { use std::ops::Range; @@ -413,6 +551,7 @@ mod tests { use bao_tree::ChunkNum; use testresult::TestResult; + /// Create chunk ranges from an array of u64 ranges fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { let mut res = ChunkRanges::empty(); for range in ranges.into_iter() { @@ -422,38 +561,152 @@ mod tests { } #[test] - fn smoke() -> TestResult<()> { - let me = "0000000000000000000000000000000000000000000000000000000000000000".parse()?; + fn downloader_state_smoke() -> TestResult<()> { + // let me = "0000000000000000000000000000000000000000000000000000000000000000".parse()?; let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; - let mut state = DownloaderState::new(me); - let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); - assert!(evs.iter().filter(|e| **e == Event::StartDiscovery { hash, id: 0 }).count() == 1, "starting a download should start a discovery task"); - assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: me, hash, id: 0 }).count() == 1, "starting a download should subscribe to the local bitmap"); + let unknown_hash = + "0000000000000000000000000000000000000000000000000000000000000002".parse()?; + let mut state = DownloaderState::new(); + let evs = state.apply_and_get_evs(super::Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: 1, + }); + assert!( + evs.iter() + .filter(|e| **e == Event::StartDiscovery { hash, id: 0 }) + .count() + == 1, + "starting a download should start a discovery task" + ); + assert!( + evs.iter() + .filter(|e| **e + == Event::SubscribeBitmap { + peer: None, + hash, + id: 0 + }) + .count() + == 1, + "starting a download should subscribe to the local bitmap" + ); println!("{evs:?}"); - let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash, bitmap: ChunkRanges::all() }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); - let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash: unknown_hash, bitmap: ChunkRanges::all() }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); - let evs = state.apply_and_get_evs(super::Command::DropPeer { peer: me }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "dropping self should produce an error!"); + let evs = state.apply_and_get_evs(Command::Bitmap { + peer: None, + hash, + bitmap: ChunkRanges::all(), + }); + assert!( + evs.iter().any(|e| matches!(e, Event::Error { .. })), + "adding an open bitmap should produce an error!" + ); + let evs = state.apply_and_get_evs(Command::Bitmap { + peer: None, + hash: unknown_hash, + bitmap: ChunkRanges::all(), + }); + assert!( + evs.iter().any(|e| matches!(e, Event::Error { .. })), + "adding an open bitmap for an unknown hash should produce an error!" + ); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: me, hash, bitmap: initial_bitmap.clone() }); + let evs = state.apply_and_get_evs(Command::Bitmap { + peer: None, + hash, + bitmap: initial_bitmap.clone(), + }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(me, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); - let evs = state.apply_and_get_evs(super::Command::BitmapUpdate { peer: me, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + assert_eq!( + state + .bitmaps + .get(&(None, hash)) + .context("bitmap should be present")? + .ranges, + initial_bitmap, + "bitmap should be set to the initial bitmap" + ); + let evs = state.apply_and_get_evs(Command::BitmapUpdate { + peer: None, + hash, + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(me, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); - let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)) }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); - let evs = state.apply_and_get_evs(super::Command::PeerDiscovered { peer: peer_a, hash }); - assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: peer_a, hash, id: 1 }).count() == 1, "adding a new peer for a hash we are interested in should subscribe to the bitmap"); - let evs = state.apply_and_get_evs(super::Command::Bitmap { peer: peer_a, hash, bitmap: chunk_ranges([0..64]) }); - assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); - let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(64)) }); + assert_eq!( + state + .bitmaps + .get(&(None, hash)) + .context("bitmap should be present")? + .ranges, + ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), + "bitmap should be updated" + ); + let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)), + }); + assert!( + evs.iter().any(|e| matches!(e, Event::Error { .. })), + "download from unknown peer should lead to an error!" + ); + let evs = state.apply_and_get_evs(Command::PeerDiscovered { peer: peer_a, hash }); + assert!( + evs.iter() + .filter(|e| **e + == Event::SubscribeBitmap { + peer: Some(peer_a), + hash, + id: 1 + }) + .count() + == 1, + "adding a new peer for a hash we are interested in should subscribe to the bitmap" + ); + let evs = state.apply_and_get_evs(Command::Bitmap { + peer: Some(peer_a), + hash, + bitmap: chunk_ranges([0..64]), + }); + assert!( + evs.iter() + .filter(|e| **e + == Event::StartPeerDownload { + id: 0, + peer: peer_a, + ranges: chunk_ranges([32..64]) + }) + .count() + == 1, + "bitmap from a peer should start a download" + ); + let evs = state.apply_and_get_evs(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([32..64]), + }); + assert!( + evs.iter() + .any(|e| matches!(e, Event::DownloadComplete { .. })), + "download should be completed by the data" + ); println!("{evs:?}"); Ok(()) } -} \ No newline at end of file + #[tokio::test] + async fn downloader_driver_smoke() -> TestResult<()> { + let store = crate::store::mem::Store::new(); + let endpoint = iroh::Endpoint::builder() + .alpns(vec![crate::protocol::ALPN.to_vec()]) + .bind() + .await?; + Ok(()) + } +} From e945d1c097c319c4c21994e9b38b38a9efdf68da Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 30 Jan 2025 16:02:04 +0200 Subject: [PATCH 05/47] WIP --- src/downloader2.rs | 520 +++++++++++++++++++++++---------------------- 1 file changed, 271 insertions(+), 249 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index bea1c3767..51de06d7e 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -12,20 +12,40 @@ //! of the caller to generate unique ids. We could use random uuids here, but using //! integers simplifies testing. //! +//! Inside the state machine, we use [ChunkRanges] to represent avaiability bitmaps. +//! We treat operations on such bitmaps as very cheap, which is the case as long as +//! the bitmaps are not very fragmented. We can introduce an even more optimized +//! bitmap type, or prevent fragmentation. +//! //! The [DownloaderDriver] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. -use std::collections::{BTreeMap, VecDeque}; - -use crate::{store::Store, Hash}; +use std::{ + collections::{BTreeMap, VecDeque}, + future::Future, + io, + sync::Arc, +}; + +use crate::{ + get::{ + fsm::{AtInitial, BlobContentNext, ConnectedNext, EndBlobNext}, + Stats, + }, + protocol::{GetRequest, RangeSpec, RangeSpecSeq}, + store::Store, + Hash, +}; use anyhow::Context; -use bao_tree::ChunkRanges; -use futures_lite::Stream; -use iroh::NodeId; +use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; +use bytes::Bytes; +use futures_lite::{Stream, StreamExt}; +use iroh::{discovery, Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use std::time::Duration; -use tokio::task::JoinSet; +use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; +use tracing::{debug, error, trace}; /// todo: make newtypes? type DownloadId = u64; @@ -34,27 +54,28 @@ type DiscoveryId = u64; type PeerDownloadId = u64; /// Announce kind -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] pub enum AnnounceKind { /// The peer supposedly has some of the data. Partial = 0, /// The peer supposedly has the complete data. + #[default] Complete, } +#[derive(Default)] struct FindPeersOpts { /// Kind of announce kind: AnnounceKind, } /// A pluggable content discovery mechanism -trait ContentDiscovery { +trait ContentDiscovery: Send + 'static { /// Find peers that have the given blob. /// /// The returned stream is a handle for the discovery task. It should be an /// infinite stream that only stops when it is dropped. - fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) - -> impl Stream + Unpin; + fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Send + Unpin + 'static; } /// Global information about a peer @@ -78,14 +99,11 @@ struct PeerBlobState { impl PeerBlobState { fn new(subscription_id: BitmapSubscriptionId) -> Self { - Self { - subscription_id, - subscription_count: 1, - ranges: ChunkRanges::empty(), - } + Self { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() } } } +#[derive(Debug)] struct DownloadRequest { /// The blob we are downloading hash: Hash, @@ -102,10 +120,7 @@ struct DownloadState { impl DownloadState { fn new(request: DownloadRequest) -> Self { - Self { - request, - peer_downloads: BTreeMap::new(), - } + Self { request, peer_downloads: BTreeMap::new() } } } @@ -231,6 +246,7 @@ enum Event { /// The unique id of the peer download task id: u64, peer: NodeId, + hash: Hash, ranges: ChunkRanges, }, StopPeerDownload { @@ -265,10 +281,7 @@ impl DownloaderState { } fn count_providers(&self, hash: Hash) -> usize { - self.bitmaps - .iter() - .filter(|((peer, x), _)| peer.is_some() && *x == hash) - .count() + self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash).count() } /// Apply a command and return the events that were generated @@ -281,9 +294,7 @@ impl DownloaderState { /// Apply a command fn apply(&mut self, cmd: Command, evs: &mut Vec) { if let Err(cause) = self.apply0(cmd, evs) { - evs.push(Event::Error { - message: format!("{cause}"), - }); + evs.push(Event::Error { message: format!("{cause}") }); } } @@ -292,10 +303,7 @@ impl DownloaderState { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! - anyhow::ensure!( - !self.downloads.contains_key(&id), - "duplicate download request {id}" - ); + anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); // either we have a subscription for this blob, or we have to create one if let Some(state) = self.bitmaps.get_mut(&(None, request.hash)) { // just increment the count @@ -303,37 +311,24 @@ impl DownloaderState { } else { // create a new subscription let subscription_id = self.next_subscription_id(); - evs.push(Event::SubscribeBitmap { - peer: None, - hash: request.hash, - id: subscription_id, - }); - self.bitmaps - .insert((None, request.hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitmap { peer: None, hash: request.hash, id: subscription_id }); + self.bitmaps.insert((None, request.hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task let discovery_id = self.next_discovery_id(); - evs.push(Event::StartDiscovery { - hash: request.hash, - id: discovery_id, - }); + evs.push(Event::StartDiscovery { hash: request.hash, id: discovery_id }); self.discovery.insert(request.hash, discovery_id); } self.downloads.insert(id, DownloadState::new(request)); } Command::StopDownload { id } => { - let removed = self - .downloads - .remove(&id) - .context(format!("removed unknown download {id}"))?; + let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; self.bitmaps.retain(|(_peer, hash), state| { if *hash == removed.request.hash { state.subscription_count -= 1; if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitmap { - id: state.subscription_id, - }); + evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); return false; } } @@ -346,32 +341,21 @@ impl DownloaderState { return Ok(()); }; // check if anybody needs this peer - if !self - .downloads - .iter() - .any(|(_id, state)| state.request.hash == hash) - { + if !self.downloads.iter().any(|(_id, state)| state.request.hash == hash) { return Ok(()); } // create a peer state if it does not exist let _state = self.peers.entry(peer).or_default(); // create a new subscription let subscription_id = self.next_subscription_id(); - evs.push(Event::SubscribeBitmap { - peer: Some(peer), - hash, - id: subscription_id, - }); - self.bitmaps - .insert((Some(peer), hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitmap { peer: Some(peer), hash, id: subscription_id }); + self.bitmaps.insert((Some(peer), hash), PeerBlobState::new(subscription_id)); } Command::DropPeer { peer } => { self.bitmaps.retain(|(p, _), state| { if *p == Some(peer) { // todo: should we emit unsubscribe evs here? - evs.push(Event::UnsubscribeBitmap { - id: state.subscription_id, - }); + evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); return false; } else { return true; @@ -380,46 +364,26 @@ impl DownloaderState { self.peers.remove(&peer); } Command::Bitmap { peer, hash, bitmap } => { - let state = self - .bitmaps - .get_mut(&(peer, hash)) - .context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; state.ranges = bitmap; self.rebalance_downloads(hash, evs)?; } - Command::BitmapUpdate { - peer, - hash, - added, - removed, - } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!( - "bitmap update for unknown peer {peer:?} and hash {hash}" - ))?; + Command::BitmapUpdate { peer, hash, added, removed } => { + let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; state.ranges |= added; state.ranges &= !removed; self.rebalance_downloads(hash, evs)?; } - Command::ChunksDownloaded { - time, - peer, - hash, - added, - } => { - let state = self.bitmaps.get_mut(&(None, hash)).context(format!( - "chunks downloaded before having local bitmap for {hash}" - ))?; + Command::ChunksDownloaded { time, peer, hash, added } => { + let state = self.bitmaps.get_mut(&(None, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; let total_after = total_chunks(&state.ranges).context("open range")?; let useful_downloaded = total_after - total_before; - let peer = self.peers.get_mut(&peer).context(format!( - "performing download before having peer state for {peer}" - ))?; - peer.download_history - .push_back((time, (total_downloaded, useful_downloaded))); + let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; + peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); self.rebalance_downloads(hash, evs)?; } } @@ -432,18 +396,12 @@ impl DownloaderState { return Ok(()); }; let mut completed = vec![]; - for (id, download) in self - .downloads - .iter_mut() - .filter(|(_id, download)| download.request.hash == hash) - { + for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { let remaining = &download.request.ranges - &self_state.ranges; if remaining.is_empty() { // cancel all downloads, if needed for (_, peer_download) in &download.peer_downloads { - evs.push(Event::StopPeerDownload { - id: peer_download.id, - }); + evs.push(Event::StopPeerDownload { id: peer_download.id }); } // notify the user that the download is complete evs.push(Event::DownloadComplete { id: *id }); @@ -452,11 +410,7 @@ impl DownloaderState { continue; } let mut candidates = vec![]; - for ((peer, _), bitmap) in self - .bitmaps - .iter() - .filter(|((peer, x), _)| peer.is_some() && *x == hash) - { + for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash) { let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { candidates.push((peer.unwrap(), intersection)); @@ -470,14 +424,8 @@ impl DownloaderState { for (peer, ranges) in candidates { let id = self.next_peer_download_id; self.next_peer_download_id += 1; - evs.push(Event::StartPeerDownload { - id, - peer, - ranges: ranges.clone(), - }); - download - .peer_downloads - .insert(peer, PeerDownloadState { id, ranges }); + evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); + download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); } } for id in completed { @@ -498,51 +446,220 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { Some(total) } -struct DownloaderDriver { +#[derive(Debug, Clone)] +struct Downloader { + send: mpsc::Sender, + task: Arc>, +} + +impl Downloader { + async fn download(&self, request: DownloadRequest) -> anyhow::Result<()> { + let (send, recv) = tokio::sync::oneshot::channel::<()>(); + self.send.send(UserCommand::Download { request, done: send }).await?; + recv.await?; + Ok(()) + } + + fn new(endpoint: Endpoint, store: S, discovery: D) -> Self { + let actor = DownloaderActor::new(endpoint, store, discovery); + let (send, recv) = tokio::sync::mpsc::channel(256); + let task = Arc::new(spawn(async move { actor.run(recv).await })); + Self { send, task } + } +} + +/// An user-facing command +#[derive(Debug)] +enum UserCommand { + Download { request: DownloadRequest, done: tokio::sync::oneshot::Sender<()> }, +} + +struct DownloaderActor { + endpoint: Endpoint, + command_rx: mpsc::Receiver, + command_tx: mpsc::Sender, state: DownloaderState, store: S, discovery: D, + download_futs: BTreeMap>, peer_download_tasks: BTreeMap>, discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, next_download_id: DownloadId, } -impl DownloaderDriver { - fn new(me: NodeId, store: S, discovery: D) -> Self { +impl DownloaderActor { + fn new(endpoint: Endpoint, store: S, discovery: D) -> Self { + let (send, recv) = mpsc::channel(256); Self { + endpoint, state: DownloaderState::new(), store, discovery, peer_download_tasks: BTreeMap::new(), discovery_tasks: BTreeMap::new(), bitmap_subscription_tasks: BTreeMap::new(), + download_futs: BTreeMap::new(), + command_tx: send, + command_rx: recv, next_download_id: 0, } } + async fn run(mut self, mut channel: mpsc::Receiver) { + loop { + tokio::select! { + biased; + Some(cmd) = self.command_rx.recv() => { + let evs = self.state.apply_and_get_evs(cmd); + for ev in evs { + self.handle_event(ev, 0); + } + }, + Some(cmd) = channel.recv() => { + debug!("user command {cmd:?}"); + match cmd { + UserCommand::Download { + request, done, + } => { + let id = self.next_download_id(); + self.download_futs.insert(id, done); + self.command_tx.send(Command::StartDownload { request, id }).await.ok(); + } + } + }, + } + } + } + fn next_download_id(&mut self) -> DownloadId { let id = self.next_download_id; self.next_download_id += 1; id } - fn handle_event(&mut self, ev: Event) { - println!("event: {:?}", ev); + fn handle_event(&mut self, ev: Event, depth: usize) { + trace!("handle_event {ev:?} {depth}"); + match ev { + Event::SubscribeBitmap { peer, hash, id } => { + let send = self.command_tx.clone(); + let task = spawn(async move { + let cmd = if peer.is_none() { + // we don't have any data, for now + Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::empty() } + } else { + // all peers have all the data, for now + Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)) } + }; + send.send(cmd).await.ok(); + futures_lite::future::pending().await + }); + self.bitmap_subscription_tasks.insert(id, task); + } + Event::StartDiscovery { hash, id } => { + let send = self.command_tx.clone(); + let mut stream = self.discovery.find_peers(hash, Default::default()); + let task = spawn(async move { + // process the infinite discovery stream and send commands + while let Some(peer) = stream.next().await { + println!("peer discovered for hash {hash}: {peer}"); + let res = send.send(Command::PeerDiscovered { peer, hash }).await; + if res.is_err() { + // only reason for this is actor task dropped + break; + } + } + }); + self.discovery_tasks.insert(id, task); + } + Event::StartPeerDownload { id, peer, hash, ranges } => { + let send = self.command_tx.clone(); + let endpoint = self.endpoint.clone(); + let task = spawn(async move { + let conn = endpoint.connect(peer, crate::ALPN).await?; + let spec = RangeSpec::new(ranges); + let initial = crate::get::fsm::start(conn, GetRequest { hash, ranges: RangeSpecSeq::new([spec]) }); + anyhow::Ok(()) + }); + } + Event::UnsubscribeBitmap { id } => { + self.bitmap_subscription_tasks.remove(&id); + } + Event::StopDiscovery { id } => { + self.discovery_tasks.remove(&id); + } + Event::StopPeerDownload { id } => { + self.peer_download_tasks.remove(&id); + } + Event::Error { message } => { + error!("Error during processing event {}", message); + } + _ => { + println!("event: {:?}", ev); + } + } } +} - fn download(&mut self, request: DownloadRequest) -> DownloadId { - let id = self.next_download_id(); - let evs = self - .state - .apply_and_get_evs(Command::StartDownload { request, id }); - for ev in evs { - self.handle_event(ev); - } - id +/// A simple static content discovery mechanism +struct StaticContentDiscovery { + info: BTreeMap>, + default: Vec, +} + +impl ContentDiscovery for StaticContentDiscovery { + fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> impl Stream + Unpin + 'static { + let peers = self.info.get(&hash).unwrap_or(&self.default).clone(); + futures_lite::stream::iter(peers).chain(futures_lite::stream::pending()) } } +async fn stream_blob(initial: AtInitial) -> io::Result { + // connect + let connected = initial.next().await?; + // read the first bytes + let ConnectedNext::StartRoot(start_root) = connected.next().await? else { + return Err(io::Error::new(io::ErrorKind::Other, "expected start root")); + }; + let header = start_root.next(); + + // get the size of the content + let (mut content, _size) = header.next().await?; + // manually loop over the content and yield all data + let done = loop { + match content.next().await { + BlobContentNext::More((next, data)) => { + if let BaoContentItem::Leaf(leaf) = data? { + // yield the data + // co.yield_(Ok(leaf.data)).await; + } + content = next; + } + BlobContentNext::Done(done) => { + // we are done with the root blob + break done; + } + } + }; + // close the connection even if there is more data + let closing = match done.next() { + EndBlobNext::Closing(closing) => closing, + EndBlobNext::MoreChildren(more) => more.finish(), + }; + // close the connection + let stats = closing.next().await?; + Ok(stats) +} + +fn spawn(f: F) -> AbortOnDropHandle +where + F: Future + Send + 'static, + T: Send + 'static, +{ + let task = tokio::spawn(f); + AbortOnDropHandle::new(task) +} + #[cfg(test)] mod tests { use std::ops::Range; @@ -550,6 +667,7 @@ mod tests { use super::*; use bao_tree::ChunkNum; use testresult::TestResult; + use tracing::info; /// Create chunk ranges from an array of u64 ranges fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { @@ -562,151 +680,55 @@ mod tests { #[test] fn downloader_state_smoke() -> TestResult<()> { + let _ = tracing_subscriber::fmt::try_init(); // let me = "0000000000000000000000000000000000000000000000000000000000000000".parse()?; let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let unknown_hash = - "0000000000000000000000000000000000000000000000000000000000000002".parse()?; + let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(); - let evs = state.apply_and_get_evs(super::Command::StartDownload { - request: DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }, - id: 1, - }); - assert!( - evs.iter() - .filter(|e| **e == Event::StartDiscovery { hash, id: 0 }) - .count() - == 1, - "starting a download should start a discovery task" - ); - assert!( - evs.iter() - .filter(|e| **e - == Event::SubscribeBitmap { - peer: None, - hash, - id: 0 - }) - .count() - == 1, - "starting a download should subscribe to the local bitmap" - ); + let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); + assert!(evs.iter().filter(|e| **e == Event::StartDiscovery { hash, id: 0 }).count() == 1, "starting a download should start a discovery task"); + assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: None, hash, id: 0 }).count() == 1, "starting a download should subscribe to the local bitmap"); println!("{evs:?}"); - let evs = state.apply_and_get_evs(Command::Bitmap { - peer: None, - hash, - bitmap: ChunkRanges::all(), - }); - assert!( - evs.iter().any(|e| matches!(e, Event::Error { .. })), - "adding an open bitmap should produce an error!" - ); - let evs = state.apply_and_get_evs(Command::Bitmap { - peer: None, - hash: unknown_hash, - bitmap: ChunkRanges::all(), - }); - assert!( - evs.iter().any(|e| matches!(e, Event::Error { .. })), - "adding an open bitmap for an unknown hash should produce an error!" - ); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::all() }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash: unknown_hash, bitmap: ChunkRanges::all() }); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply_and_get_evs(Command::Bitmap { - peer: None, - hash, - bitmap: initial_bitmap.clone(), - }); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); - assert_eq!( - state - .bitmaps - .get(&(None, hash)) - .context("bitmap should be present")? - .ranges, - initial_bitmap, - "bitmap should be set to the initial bitmap" - ); - let evs = state.apply_and_get_evs(Command::BitmapUpdate { - peer: None, - hash, - added: chunk_ranges([16..32]), - removed: ChunkRanges::empty(), - }); + assert_eq!(state.bitmaps.get(&(None, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: None, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); - assert_eq!( - state - .bitmaps - .get(&(None, hash)) - .context("bitmap should be present")? - .ranges, - ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), - "bitmap should be updated" - ); + assert_eq!(state.bitmaps.get(&(None, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)), }); - assert!( - evs.iter().any(|e| matches!(e, Event::Error { .. })), - "download from unknown peer should lead to an error!" - ); + assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); let evs = state.apply_and_get_evs(Command::PeerDiscovered { peer: peer_a, hash }); - assert!( - evs.iter() - .filter(|e| **e - == Event::SubscribeBitmap { - peer: Some(peer_a), - hash, - id: 1 - }) - .count() - == 1, - "adding a new peer for a hash we are interested in should subscribe to the bitmap" - ); - let evs = state.apply_and_get_evs(Command::Bitmap { - peer: Some(peer_a), - hash, - bitmap: chunk_ranges([0..64]), - }); - assert!( - evs.iter() - .filter(|e| **e - == Event::StartPeerDownload { - id: 0, - peer: peer_a, - ranges: chunk_ranges([32..64]) - }) - .count() - == 1, - "bitmap from a peer should start a download" - ); - let evs = state.apply_and_get_evs(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([32..64]), - }); - assert!( - evs.iter() - .any(|e| matches!(e, Event::DownloadComplete { .. })), - "download should be completed by the data" - ); + assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: Some(peer_a), hash, id: 1 }).count() == 1, "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: Some(peer_a), hash, bitmap: chunk_ranges([0..64]) }); + assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); + let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); + assert!(evs.iter().any(|e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); println!("{evs:?}"); Ok(()) } #[tokio::test] async fn downloader_driver_smoke() -> TestResult<()> { + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let store = crate::store::mem::Store::new(); - let endpoint = iroh::Endpoint::builder() - .alpns(vec![crate::protocol::ALPN.to_vec()]) - .bind() - .await?; + let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).bind().await?; + let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer_a] }; + let downloader = Downloader::new(endpoint, store, discovery); + let fut = downloader.download(DownloadRequest { hash, ranges: ChunkRanges::all() }); + fut.await?; Ok(()) } } From 7f3f22f53b9e75d99b3825bb98e28bcb65e34a50 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 30 Jan 2025 16:45:17 +0200 Subject: [PATCH 06/47] WIP --- src/downloader2.rs | 50 ++++++++++++++++++++++++++++++---------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 51de06d7e..d6775040b 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -30,22 +30,20 @@ use crate::{ get::{ fsm::{AtInitial, BlobContentNext, ConnectedNext, EndBlobNext}, Stats, - }, - protocol::{GetRequest, RangeSpec, RangeSpecSeq}, - store::Store, - Hash, + }, protocol::{GetRequest, RangeSpec, RangeSpecSeq}, store::{BaoBatchWriter, MapEntryMut, Store}, util::local_pool::{self, LocalPool}, Hash }; use anyhow::Context; use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; use bytes::Bytes; use futures_lite::{Stream, StreamExt}; use iroh::{discovery, Endpoint, NodeId}; +use quinn::Chunk; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use std::time::Duration; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; -use tracing::{debug, error, trace}; +use tracing::{debug, error, info, trace}; /// todo: make newtypes? type DownloadId = u64; @@ -460,8 +458,8 @@ impl Downloader { Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D) -> Self { - let actor = DownloaderActor::new(endpoint, store, discovery); + fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool) -> Self { + let actor = DownloaderActor::new(endpoint, store, discovery, local_pool); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); Self { send, task } @@ -475,6 +473,7 @@ enum UserCommand { } struct DownloaderActor { + local_pool: LocalPool, endpoint: Endpoint, command_rx: mpsc::Receiver, command_tx: mpsc::Sender, @@ -482,16 +481,17 @@ struct DownloaderActor { store: S, discovery: D, download_futs: BTreeMap>, - peer_download_tasks: BTreeMap>, + peer_download_tasks: BTreeMap>>, discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, next_download_id: DownloadId, } impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: D) -> Self { + fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool) -> Self { let (send, recv) = mpsc::channel(256); Self { + local_pool, endpoint, state: DownloaderState::new(), store, @@ -575,12 +575,16 @@ impl DownloaderActor { Event::StartPeerDownload { id, peer, hash, ranges } => { let send = self.command_tx.clone(); let endpoint = self.endpoint.clone(); - let task = spawn(async move { + let store = self.store.clone(); + let task = self.local_pool.spawn(move || async move { + info!("Connecting to peer {peer}"); let conn = endpoint.connect(peer, crate::ALPN).await?; let spec = RangeSpec::new(ranges); let initial = crate::get::fsm::start(conn, GetRequest { hash, ranges: RangeSpecSeq::new([spec]) }); + stream_to_db(initial, store, hash, peer, send).await?; anyhow::Ok(()) }); + self.peer_download_tasks.insert(id, task); } Event::UnsubscribeBitmap { id } => { self.bitmap_subscription_tasks.remove(&id); @@ -614,7 +618,7 @@ impl ContentDiscovery for StaticContentDiscovery { } } -async fn stream_blob(initial: AtInitial) -> io::Result { +async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: NodeId, sender: mpsc::Sender) -> io::Result { // connect let connected = initial.next().await?; // read the first bytes @@ -624,14 +628,25 @@ async fn stream_blob(initial: AtInitial) -> io::Result { let header = start_root.next(); // get the size of the content - let (mut content, _size) = header.next().await?; + let (mut content, size) = header.next().await?; + let entry = store.get_or_create(hash, size).await?; + let mut writer = entry.batch_writer().await?; + let mut batch = Vec::new(); // manually loop over the content and yield all data let done = loop { match content.next().await { BlobContentNext::More((next, data)) => { - if let BaoContentItem::Leaf(leaf) = data? { - // yield the data - // co.yield_(Ok(leaf.data)).await; + match data? { + BaoContentItem::Parent(parent) => { + batch.push(parent.into()); + }, + BaoContentItem::Leaf(leaf) => { + let added = ChunkRanges::from(ChunkNum(leaf.offset / 1024)..); + sender.send(Command::ChunksDownloaded { time: Duration::ZERO, peer, hash, added: added.clone() }).await.ok(); + batch.push(leaf.into()); + writer.write_batch(size, std::mem::take(&mut batch)).await?; + sender.send(Command::BitmapUpdate { peer: None, hash, added, removed: ChunkRanges::empty() }).await.ok(); + } } content = next; } @@ -713,7 +728,7 @@ mod tests { let evs = state.apply_and_get_evs(Command::Bitmap { peer: Some(peer_a), hash, bitmap: chunk_ranges([0..64]) }); assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); - assert!(evs.iter().any(|e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); + assert!(evs.iter().filter(|e| matches!(e, Event::DownloadComplete { .. })).count() == 1, "download should be completed by the data"); println!("{evs:?}"); Ok(()) } @@ -726,7 +741,8 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer_a] }; - let downloader = Downloader::new(endpoint, store, discovery); + let local_pool = LocalPool::single(); + let downloader = Downloader::new(endpoint, store, discovery, local_pool); let fut = downloader.download(DownloadRequest { hash, ranges: ChunkRanges::all() }); fut.await?; Ok(()) From 9af5698aadd3d3b8a878b818a87df178468cfab5 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 30 Jan 2025 17:13:48 +0200 Subject: [PATCH 07/47] wip --- src/downloader2.rs | 36 +++++++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index d6775040b..6d25222c6 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -202,7 +202,9 @@ enum Command { /// The ranges that were removed removed: ChunkRanges, }, - /// A chunk was downloaded + /// A chunk was downloaded, but not yet stored + /// + /// This can only be used for updating peer stats, not for completing downloads. ChunksDownloaded { /// Time when the download was received time: Duration, @@ -578,7 +580,10 @@ impl DownloaderActor { let store = self.store.clone(); let task = self.local_pool.spawn(move || async move { info!("Connecting to peer {peer}"); - let conn = endpoint.connect(peer, crate::ALPN).await?; + let conn = endpoint.connect(peer, crate::ALPN).await; + info!("Got connection to peer {peer} {}", conn.is_err()); + println!("{conn:?}"); + let conn = conn?; let spec = RangeSpec::new(ranges); let initial = crate::get::fsm::start(conn, GetRequest { hash, ranges: RangeSpecSeq::new([spec]) }); stream_to_db(initial, store, hash, peer, send).await?; @@ -679,10 +684,26 @@ where mod tests { use std::ops::Range; + use crate::net_protocol::Blobs; + use super::*; use bao_tree::ChunkNum; + use iroh::protocol::Router; use testresult::TestResult; - use tracing::info; + + #[cfg(feature = "rpc")] + async fn make_test_node(data: &[u8]) -> anyhow::Result<(Router, NodeId, Hash)> { + let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; + let node_id = endpoint.node_id(); + let store = crate::store::mem::Store::new(); + let blobs = Blobs::builder(store) + .build(&endpoint); + let hash = blobs.client().add_bytes(bytes::Bytes::copy_from_slice(data)).await?.hash; + let router = iroh::protocol::Router::builder(endpoint) + .accept(crate::ALPN, blobs) + .spawn().await?; + Ok((router, node_id, hash)) + } /// Create chunk ranges from an array of u64 ranges fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { @@ -734,15 +755,16 @@ mod tests { } #[tokio::test] + #[cfg(feature = "rpc")] async fn downloader_driver_smoke() -> TestResult<()> { let _ = tracing_subscriber::fmt::try_init(); - let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let (router, peer, hash) = make_test_node(b"test").await?; let store = crate::store::mem::Store::new(); - let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).bind().await?; - let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer_a] }; + let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; + let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; let local_pool = LocalPool::single(); let downloader = Downloader::new(endpoint, store, discovery, local_pool); + tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: ChunkRanges::all() }); fut.await?; Ok(()) From ebc845476235237e241e452cd7a2e42ff73d02c7 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 30 Jan 2025 19:08:12 +0200 Subject: [PATCH 08/47] Add large test --- src/downloader2.rs | 76 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 58 insertions(+), 18 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 6d25222c6..42e5d9d14 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -30,7 +30,11 @@ use crate::{ get::{ fsm::{AtInitial, BlobContentNext, ConnectedNext, EndBlobNext}, Stats, - }, protocol::{GetRequest, RangeSpec, RangeSpecSeq}, store::{BaoBatchWriter, MapEntryMut, Store}, util::local_pool::{self, LocalPool}, Hash + }, + protocol::{GetRequest, RangeSpec, RangeSpecSeq}, + store::{BaoBatchWriter, MapEntryMut, Store}, + util::local_pool::{self, LocalPool}, + Hash, }; use anyhow::Context; use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; @@ -203,7 +207,7 @@ enum Command { removed: ChunkRanges, }, /// A chunk was downloaded, but not yet stored - /// + /// /// This can only be used for updating peer stats, not for completing downloads. ChunksDownloaded { /// Time when the download was received @@ -384,7 +388,6 @@ impl DownloaderState { let useful_downloaded = total_after - total_before; let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); - self.rebalance_downloads(hash, evs)?; } } Ok(()) @@ -398,6 +401,7 @@ impl DownloaderState { let mut completed = vec![]; for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { let remaining = &download.request.ranges - &self_state.ranges; + info!("Remaining chunks for download {id}: {remaining:?}"); if remaining.is_empty() { // cancel all downloads, if needed for (_, peer_download) in &download.peer_downloads { @@ -416,12 +420,15 @@ impl DownloaderState { candidates.push((peer.unwrap(), intersection)); } } + info!("Stopping {} old peer downloads", download.peer_downloads.len()); for (_, state) in &download.peer_downloads { // stop all downloads evs.push(Event::StopPeerDownload { id: state.id }); } + info!("Creating {} new peer downloads", candidates.len()); download.peer_downloads.clear(); for (peer, ranges) in candidates { + info!(" Starting download from {peer} for {hash} {ranges:?}"); let id = self.next_peer_download_id; self.next_peer_download_id += 1; evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); @@ -551,7 +558,7 @@ impl DownloaderActor { Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::empty() } } else { // all peers have all the data, for now - Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)) } + Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } }; send.send(cmd).await.ok(); futures_lite::future::pending().await @@ -580,12 +587,13 @@ impl DownloaderActor { let store = self.store.clone(); let task = self.local_pool.spawn(move || async move { info!("Connecting to peer {peer}"); - let conn = endpoint.connect(peer, crate::ALPN).await; - info!("Got connection to peer {peer} {}", conn.is_err()); - println!("{conn:?}"); - let conn = conn?; + let conn = endpoint.connect(peer, crate::ALPN).await?; + info!("Got connection to peer {peer}"); let spec = RangeSpec::new(ranges); - let initial = crate::get::fsm::start(conn, GetRequest { hash, ranges: RangeSpecSeq::new([spec]) }); + let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); + info!("starting download from {peer} for {hash} {ranges:?}"); + let request = GetRequest::new(hash, ranges); + let initial = crate::get::fsm::start(conn, request); stream_to_db(initial, store, hash, peer, send).await?; anyhow::Ok(()) }); @@ -600,6 +608,11 @@ impl DownloaderActor { Event::StopPeerDownload { id } => { self.peer_download_tasks.remove(&id); } + Event::DownloadComplete { id } => { + if let Some(done) = self.download_futs.remove(&id) { + done.send(()).ok(); + } + } Event::Error { message } => { error!("Error during processing event {}", message); } @@ -644,9 +657,10 @@ async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: match data? { BaoContentItem::Parent(parent) => { batch.push(parent.into()); - }, + } BaoContentItem::Leaf(leaf) => { - let added = ChunkRanges::from(ChunkNum(leaf.offset / 1024)..); + let start_chunk = leaf.offset / 1024; + let added = ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 1)); sender.send(Command::ChunksDownloaded { time: Duration::ZERO, peer, hash, added: added.clone() }).await.ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; @@ -696,12 +710,9 @@ mod tests { let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; let node_id = endpoint.node_id(); let store = crate::store::mem::Store::new(); - let blobs = Blobs::builder(store) - .build(&endpoint); + let blobs = Blobs::builder(store).build(&endpoint); let hash = blobs.client().add_bytes(bytes::Bytes::copy_from_slice(data)).await?.hash; - let router = iroh::protocol::Router::builder(endpoint) - .accept(crate::ALPN, blobs) - .spawn().await?; + let router = iroh::protocol::Router::builder(endpoint).accept(crate::ALPN, blobs).spawn().await?; Ok((router, node_id, hash)) } @@ -748,7 +759,11 @@ mod tests { assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: Some(peer_a), hash, id: 1 }).count() == 1, "adding a new peer for a hash we are interested in should subscribe to the bitmap"); let evs = state.apply_and_get_evs(Command::Bitmap { peer: Some(peer_a), hash, bitmap: chunk_ranges([0..64]) }); assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); + // ChunksDownloaded just updates the peer stats let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); + assert!(evs.is_empty()); + // Bitmap update for the local bitmap should complete the download + let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: None, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); assert!(evs.iter().filter(|e| matches!(e, Event::DownloadComplete { .. })).count() == 1, "download should be completed by the data"); println!("{evs:?}"); Ok(()) @@ -758,14 +773,39 @@ mod tests { #[cfg(feature = "rpc")] async fn downloader_driver_smoke() -> TestResult<()> { let _ = tracing_subscriber::fmt::try_init(); - let (router, peer, hash) = make_test_node(b"test").await?; + let (_router1, peer, hash) = make_test_node(b"test").await?; let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; let local_pool = LocalPool::single(); let downloader = Downloader::new(endpoint, store, discovery, local_pool); tokio::time::sleep(Duration::from_secs(2)).await; - let fut = downloader.download(DownloadRequest { hash, ranges: ChunkRanges::all() }); + let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); + fut.await?; + Ok(()) + } + + #[tokio::test] + #[cfg(feature = "rpc")] + async fn downloader_driver_large() -> TestResult<()> { + use std::collections::BTreeSet; + + let _ = tracing_subscriber::fmt::try_init(); + let data = vec![0u8; 1024 * 1024]; + let mut nodes = vec![]; + for _i in 0..10 { + nodes.push(make_test_node(&data).await?); + } + let peers = nodes.iter().map(|(_, peer, _)| *peer).collect::>(); + let hashes = nodes.iter().map(|(_, _, hash)| *hash).collect::>(); + let hash = *hashes.iter().next().unwrap(); + let store = crate::store::mem::Store::new(); + let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; + let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; + let local_pool = LocalPool::single(); + let downloader = Downloader::new(endpoint, store, discovery, local_pool); + tokio::time::sleep(Duration::from_secs(2)).await; + let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; Ok(()) } From 002503fd1b2d44cf93f4006e7378f4ac89d5c202 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 15:46:34 +0200 Subject: [PATCH 09/47] add planner --- src/downloader2.rs | 441 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 404 insertions(+), 37 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 42e5d9d14..2fdb3944e 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -20,7 +20,7 @@ //! The [DownloaderDriver] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. use std::{ - collections::{BTreeMap, VecDeque}, + collections::{BTreeMap, BTreeSet, VecDeque}, future::Future, io, sync::Arc, @@ -126,6 +126,119 @@ impl DownloadState { } } +/// Trait for a download planner. +/// +/// A download planner has the option to be stateful and keep track of plans +/// depending on the hash, but many planners will be stateless. +/// +/// Planners can do whatever they want with the chunk ranges. Usually, they +/// want to deduplicate the ranges, but they could also do other things, like +/// eliminate gaps or even extend ranges. The only thing they should not do is +/// to add new peers to the list of options. +trait DownloadPlanner: Send + 'static { + /// Make a download plan for a hash, by reducing or eliminating the overlap of chunk ranges + fn plan(&mut self, hash: Hash, options: &mut BTreeMap); +} + +type BoxedDownloadPlanner = Box; + +/// A download planner that just leaves everything as is. +/// +/// Data will be downloaded from all peers wherever multiple peers have the same data. +struct NoopPlanner; + +impl DownloadPlanner for NoopPlanner { + fn plan(&mut self, _hash: Hash, _options: &mut BTreeMap) {} +} + +/// A download planner that fully removes overlap between peers. +/// +/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, +/// and for each stripe decides on a single peer to download from, based on the +/// peer id and a random seed. +struct StripePlanner { + /// seed for the score function. This can be set to 0 for testing for + /// maximum determinism, but can be set to a random value for production + /// to avoid multiple downloaders coming up with the same plan. + seed: u64, + /// The log of the stripe size in chunks. This planner is relatively + /// dumb and does not try to come up with continuous ranges, but you can + /// just set this to a large value to avoid fragmentation. + /// + /// In the very common case where you have small downloads, this will + /// frequently just choose a single peer for the entire download. + /// + /// This is a feature, not a bug. For small downloads, it is not worth + /// the effort to come up with a more sophisticated plan. + stripe_size_log: u8, +} + +impl StripePlanner { + pub fn new(seed: u64, stripe_size_log: u8) -> Self { + Self { seed, stripe_size_log } + } + + /// The score function to decide which peer to download from. + fn score(peer: &NodeId, seed: u64, stripe: u64) -> u64 { + // todo: use fnv? blake3 is a bit overkill + let mut data = [0u8; 32 + 8 + 8]; + data[..32].copy_from_slice(peer.as_bytes()); + data[32..40].copy_from_slice(&stripe.to_be_bytes()); + data[40..48].copy_from_slice(&seed.to_be_bytes()); + let hash = blake3::hash(&data); + u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) + } +} + +impl DownloadPlanner for StripePlanner { + fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { + options.retain(|_, x| !x.is_empty()); + if options.len() <= 1 { + return; + } + assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); + let mut ranges = Vec::new(); + for x in options.values() { + ranges.extend(x.boundaries().iter().map(|x| x.0)); + } + let min = ranges.iter().next().copied().unwrap(); + let max = ranges.iter().next_back().copied().unwrap(); + // add stripe subdividers + for i in (min >> self.stripe_size_log)..(max >> self.stripe_size_log) { + let x = i << self.stripe_size_log; + if x > min && x < max { + ranges.push(x); + } + } + ranges.sort(); + ranges.dedup(); + for range in ranges.windows(2) { + let start = ChunkNum(range[0]); + let end = ChunkNum(range[1]); + let curr = ChunkRanges::from(start..end); + let stripe = range[0] >> self.stripe_size_log; + let mut best_peer = None; + let mut best_score = 0; + let mut matching = vec![]; + for (peer, peer_ranges) in options.iter_mut() { + if peer_ranges.contains(&start) { + let score = Self::score(peer, self.seed, stripe); + if score > best_score && peer_ranges.contains(&start) { + best_peer = Some(*peer); + best_score = score; + } + matching.push((peer, peer_ranges)); + } + } + for (peer, peer_ranges) in matching { + if *peer != best_peer.unwrap() { + peer_ranges.difference_with(&curr); + } + } + } + } +} + struct PeerDownloadState { id: PeerDownloadId, ranges: ChunkRanges, @@ -157,10 +270,12 @@ struct DownloaderState { next_discovery_id: u64, // the next peer download id next_peer_download_id: u64, + // the download planner + planner: Box, } impl DownloaderState { - fn new() -> Self { + fn new(planner: Box) -> Self { Self { peers: BTreeMap::new(), downloads: BTreeMap::new(), @@ -169,6 +284,7 @@ impl DownloaderState { next_subscription_id: 0, next_discovery_id: 0, next_peer_download_id: 0, + planner, } } } @@ -302,6 +418,42 @@ impl DownloaderState { } } + /// Stop a download and clean up + /// + /// This is called both for stopping a download before completion, and for + /// cleaning up after a successful download. + /// + /// Cleanup involves emitting events for + /// - stopping all peer downloads + /// - unsubscribing from bitmaps if needed + /// - stopping the discovery task if needed + fn stop_download(&mut self, id: u64, evs: &mut Vec) -> anyhow::Result<()> { + let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; + let removed_hash = removed.request.hash; + // stop associated peer downloads + for peer_download in removed.peer_downloads.values() { + evs.push(Event::StopPeerDownload { id: peer_download.id }); + } + // unsubscribe from bitmaps that have no more subscriptions + self.bitmaps.retain(|(_peer, hash), state| { + if *hash == removed_hash { + state.subscription_count -= 1; + if state.subscription_count == 0 { + evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + return false; + } + } + true + }); + let hash_interest = self.downloads.values().filter(|x| x.request.hash == removed.request.hash).count(); + if hash_interest == 0 { + // stop the discovery task if we were the last one interested in the hash + let discovery_id = self.discovery.remove(&removed.request.hash).context(format!("removed unknown discovery task for {}", removed.request.hash))?; + evs.push(Event::StopDiscovery { id: discovery_id }); + } + Ok(()) + } + /// Apply a command and bail out on error fn apply0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { match cmd { @@ -327,17 +479,7 @@ impl DownloaderState { self.downloads.insert(id, DownloadState::new(request)); } Command::StopDownload { id } => { - let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; - self.bitmaps.retain(|(_peer, hash), state| { - if *hash == removed.request.hash { - state.subscription_count -= 1; - if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); - return false; - } - } - true - }); + self.stop_download(id, evs)?; } Command::PeerDiscovered { peer, hash } => { if self.bitmaps.contains_key(&(Some(peer), hash)) { @@ -371,12 +513,18 @@ impl DownloaderState { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; state.ranges = bitmap; + if peer.is_none() { + self.check_completion(hash, evs)?; + } self.rebalance_downloads(hash, evs)?; } Command::BitmapUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; state.ranges |= added; state.ranges &= !removed; + if peer.is_none() { + self.check_completion(hash, evs)?; + } self.rebalance_downloads(hash, evs)?; } Command::ChunksDownloaded { time, peer, hash, added } => { @@ -393,26 +541,73 @@ impl DownloaderState { Ok(()) } - fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + /// Check for completion of a download or of an individual peer download + /// + /// This must be called after each change of the local bitmap for a hash + fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get(&(None, hash)) else { // 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 mut completed = vec![]; for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { - let remaining = &download.request.ranges - &self_state.ranges; - info!("Remaining chunks for download {id}: {remaining:?}"); - if remaining.is_empty() { - // cancel all downloads, if needed - for (_, peer_download) in &download.peer_downloads { - evs.push(Event::StopPeerDownload { id: peer_download.id }); - } + // 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 evs.push(Event::DownloadComplete { id: *id }); - // mark the download for later removal + // remember id for later cleanup completed.push(*id); + // no need to look at individual peer downloads in this case continue; } + // check if any peer download is complete, and remove it. + let mut available = vec![]; + download.peer_downloads.retain(|peer, peer_download| { + if self_state.ranges.is_superset(&peer_download.ranges) { + // stop this peer download. + // + // Might be a noop if the cause for this local change was the same peer download, but we don't know. + evs.push(Event::StopPeerDownload { id: peer_download.id }); + // mark this peer as available + available.push(*peer); + false + } else { + true + } + }); + // reassign the newly available peers without doing a full rebalance + if !available.is_empty() { + // check if any of the available peers can provide something of the remaining data + let remaining = &download.request.ranges - &self_state.ranges; + // see what the new peers can do for us + let mut candidates= BTreeMap::new(); + for peer in available { + let Some(peer_state) = self.bitmaps.get(&(Some(peer), hash)) else { + // weird. we should have a bitmap for this peer since it just completed a download + continue; + }; + let intersection = &peer_state.ranges & &remaining; + if !intersection.is_empty() { + candidates.insert(peer, intersection); + } + } + // todo: make a plan and create new peer downloads if possible + } + } + // cleanup completed downloads, has to happen later to avoid double mutable borrow + for id in completed { + self.stop_download(id, evs)?; + } + Ok(()) + } + + fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + let Some(self_state) = self.bitmaps.get(&(None, hash)) else { + // we don't have the self state yet, so we can't really decide if we need to download anything at all + return Ok(()); + }; + for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { + let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = vec![]; for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash) { let intersection = &bitmap.ranges & &remaining; @@ -435,9 +630,6 @@ impl DownloaderState { download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); } } - for id in completed { - self.downloads.remove(&id); - } Ok(()) } } @@ -467,8 +659,8 @@ impl Downloader { Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool) -> Self { - let actor = DownloaderActor::new(endpoint, store, discovery, local_pool); + fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool, planner: Box) -> Self { + let actor = DownloaderActor::new(endpoint, store, discovery, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); Self { send, task } @@ -497,12 +689,12 @@ struct DownloaderActor { } impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool) -> Self { + fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool, planner: Box) -> Self { let (send, recv) = mpsc::channel(256); Self { local_pool, endpoint, - state: DownloaderState::new(), + state: DownloaderState::new(planner), store, discovery, peer_download_tasks: BTreeMap::new(), @@ -694,6 +886,122 @@ where AbortOnDropHandle::new(task) } +fn print_bitmap(iter: impl IntoIterator) -> String { + let mut chars = String::new(); + for x in iter { + chars.push(if x { '█' } else { ' ' }); + } + chars +} + +fn print_bitmap_compact(iter: impl IntoIterator) -> String { + let mut chars = String::new(); + let mut iter = iter.into_iter(); + + while let (Some(left), Some(right)) = (iter.next(), iter.next()) { + let c = match (left, right) { + (true, true) => '█', // Both pixels are "on" + (true, false) => '▌', // Left pixel is "on" + (false, true) => '▐', // Right pixel is "on" + (false, false) => ' ', // Both are "off" + }; + chars.push(c); + } + + // If there's an odd pixel at the end, print only a left block. + if let Some(left) = iter.next() { + chars.push(if left { '▌' } else { ' ' }); + } + + chars +} + +fn as_bool_iter(x: ChunkRanges, max: u64) -> impl Iterator { + let max = x.iter().last().map(|x| match x { + RangeSetRange::RangeFrom(_) => max, + RangeSetRange::Range(x) => x.end.0, + }).unwrap_or_default(); + (0..max).map(move |i| x.contains(&ChunkNum(i))) +} + +/// Given a set of ranges, make them non-overlapping according to some rules. +fn select_ranges(ranges: &[ChunkRanges], continuity_bonus: u64) -> Vec> { + let mut total = vec![0u64; ranges.len()]; + let mut boundaries = BTreeSet::new(); + assert!(ranges.iter().all(|x| x.boundaries().len() % 2 == 0)); + for range in ranges.iter() { + for x in range.boundaries() { + boundaries.insert(x.0); + } + } + let max = boundaries.iter().max().copied().unwrap_or(0); + let mut last_selected = None; + let mut res = vec![]; + for i in 0..max { + let mut lowest_score = u64::MAX; + let mut selected = None; + for j in 0..ranges.len() { + if ranges[j].contains(&ChunkNum(i)) { + let consecutive = last_selected == Some(j); + let score = if consecutive { + total[j].saturating_sub(continuity_bonus) + } else { + total[j] + }; + if score < lowest_score { + lowest_score = score; + selected = Some(j); + } + } + } + res.push(selected); + if let Some(selected) = selected { + total[selected] += 1; + } + last_selected = selected; + } + res +} + +fn create_ranges(indexes: impl IntoIterator>) -> Vec { + let mut res = vec![]; + for (i, n) in indexes.into_iter().enumerate() { + let x = i as u64; + if let Some(n) = n { + while res.len() <= n { + res.push(ChunkRanges::empty()); + } + res[n] |= ChunkRanges::from(ChunkNum(x)..ChunkNum(x + 1)); + } + } + res +} + +fn print_colored_bitmap(data: &[u8], colors: &[u8]) -> String { + let mut chars = String::new(); + let mut iter = data.iter(); + + while let Some(&left) = iter.next() { + let right = iter.next(); // Try to fetch the next element + + let left_color = colors.get(left as usize).map(|x| *x).unwrap_or_default(); + let right_char = match right { + Some(&right) => { + let right_color = colors.get(right as usize).map(|x| *x).unwrap_or_default(); + // Use ANSI escape codes to color left and right halves of `█` + format!( + "\x1b[38;5;{}m\x1b[48;5;{}m▌\x1b[0m", + left_color, right_color + ) + } + None => format!("\x1b[38;5;{}m▌\x1b[0m", left_color), // Handle odd-length case + }; + + chars.push_str(&right_char); + } + chars +} + #[cfg(test)] mod tests { use std::ops::Range; @@ -705,6 +1013,53 @@ mod tests { use iroh::protocol::Router; use testresult::TestResult; + #[test] + fn print_chunk_range() { + let x = chunk_ranges([0..3, 4..30, 40..50]); + let s = print_bitmap_compact(as_bool_iter(x, 50)); + println!("{}", s); + } + + #[test] + fn test_select_ranges() { + let ranges = [ + chunk_ranges([0..90]), + chunk_ranges([0..100]), + chunk_ranges([0..80]), + ]; + let iter = select_ranges(ranges.as_slice(), 8); + for (i, range) in ranges.iter().enumerate() { + let bools = as_bool_iter(range.clone(), 100).collect::>(); + println!("{i:4}{}", print_bitmap(bools)); + } + print!(" "); + for x in &iter { + print!("{}", x.map(|x| x.to_string()).unwrap_or(" ".to_string())); + } + println!(); + let ranges = create_ranges(iter); + for (i, range) in ranges.iter().enumerate() { + let bools = as_bool_iter(range.clone(), 100).collect::>(); + println!("{i:4}{}", print_bitmap(bools)); + } + } + + #[test] + fn test_is_superset() { + let local = ChunkRanges::from(ChunkNum(0)..ChunkNum(100)); + let request = ChunkRanges::from(ChunkNum(0)..ChunkNum(50)); + assert!(local.is_superset(&request)); + } + + #[test] + fn test_print_colored_bitmap() { + let bitmap = vec![ + 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, + ]; // Notice: odd-length input + + println!("{}", print_colored_bitmap(&bitmap, &[1,2,3,4])); + } + #[cfg(feature = "rpc")] async fn make_test_node(data: &[u8]) -> anyhow::Result<(Router, NodeId, Hash)> { let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; @@ -725,6 +1080,19 @@ mod tests { res } + fn noop_planner() -> BoxedDownloadPlanner { + Box::new(NoopPlanner) + } + + /// Checks if an exact event is present exactly once in a list of events + fn has_one_event(evs: &[Event], ev: &Event) -> bool { + evs.iter().filter(|e| *e == ev).count() == 1 + } + + fn has_one_event_matching(evs: &[Event], f: impl Fn(&Event) -> bool) -> bool { + evs.iter().filter(|e| f(e)).count() == 1 + } + #[test] fn downloader_state_smoke() -> TestResult<()> { let _ = tracing_subscriber::fmt::try_init(); @@ -732,13 +1100,12 @@ mod tests { let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; - let mut state = DownloaderState::new(); + let mut state = DownloaderState::new(noop_planner()); let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); - assert!(evs.iter().filter(|e| **e == Event::StartDiscovery { hash, id: 0 }).count() == 1, "starting a download should start a discovery task"); - assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: None, hash, id: 0 }).count() == 1, "starting a download should subscribe to the local bitmap"); - println!("{evs:?}"); + assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: 0 }), "starting a download should start a discovery task"); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: None, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::all() }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash: unknown_hash, bitmap: ChunkRanges::all() }); assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); @@ -765,7 +1132,7 @@ mod tests { // Bitmap update for the local bitmap should complete the download let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: None, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); assert!(evs.iter().filter(|e| matches!(e, Event::DownloadComplete { .. })).count() == 1, "download should be completed by the data"); - println!("{evs:?}"); + println!("XXX {evs:?}"); Ok(()) } @@ -778,7 +1145,7 @@ mod tests { let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, local_pool); + let downloader = Downloader::new(endpoint, store, discovery, local_pool, noop_planner()); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); fut.await?; @@ -803,7 +1170,7 @@ mod tests { let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, local_pool); + let downloader = Downloader::new(endpoint, store, discovery, local_pool, noop_planner()); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; From f20bc2e2a83193e16936b311a5a41a06799aeeb7 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 17:31:54 +0200 Subject: [PATCH 10/47] Add fast download start --- src/downloader2.rs | 290 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 226 insertions(+), 64 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 2fdb3944e..0a1cecf9e 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -38,10 +38,8 @@ use crate::{ }; use anyhow::Context; use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; -use bytes::Bytes; use futures_lite::{Stream, StreamExt}; -use iroh::{discovery, Endpoint, NodeId}; -use quinn::Chunk; +use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use std::time::Duration; @@ -126,6 +124,20 @@ impl DownloadState { } } +#[derive(Debug, Default)] +struct IdGenerator { + next_id: u64, +} + +impl IdGenerator { + + fn next(&mut self) -> u64 { + let id = self.next_id; + self.next_id += 1; + id + } +} + /// Trait for a download planner. /// /// A download planner has the option to be stateful and keep track of plans @@ -192,26 +204,12 @@ impl StripePlanner { impl DownloadPlanner for StripePlanner { fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { + assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); options.retain(|_, x| !x.is_empty()); if options.len() <= 1 { return; } - assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); - let mut ranges = Vec::new(); - for x in options.values() { - ranges.extend(x.boundaries().iter().map(|x| x.0)); - } - let min = ranges.iter().next().copied().unwrap(); - let max = ranges.iter().next_back().copied().unwrap(); - // add stripe subdividers - for i in (min >> self.stripe_size_log)..(max >> self.stripe_size_log) { - let x = i << self.stripe_size_log; - if x > min && x < max { - ranges.push(x); - } - } - ranges.sort(); - ranges.dedup(); + let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); for range in ranges.windows(2) { let start = ChunkNum(range[0]); let end = ChunkNum(range[1]); @@ -236,9 +234,110 @@ impl DownloadPlanner for StripePlanner { } } } + options.retain(|_, x| !x.is_empty()); + } +} + +fn get_continuous_ranges(options: &mut BTreeMap, stripe_size_log: u8) -> Option> { + let mut ranges = BTreeSet::new(); + for x in options.values() { + ranges.extend(x.boundaries().iter().map(|x| x.0)); + } + let min = ranges.iter().next().copied()?; + let max = ranges.iter().next_back().copied()?; + // add stripe subdividers + for i in (min >> stripe_size_log)..(max >> stripe_size_log) { + let x = i << stripe_size_log; + if x > min && x < max { + ranges.insert(x); + } + } + let ranges = ranges.into_iter().collect::>(); + Some(ranges) +} + + +/// A download planner that fully removes overlap between peers. +/// +/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, +/// and for each stripe decides on a single peer to download from, based on the +/// peer id and a random seed. +struct StripePlanner2 { + /// seed for the score function. This can be set to 0 for testing for + /// maximum determinism, but can be set to a random value for production + /// to avoid multiple downloaders coming up with the same plan. + seed: u64, + /// The log of the stripe size in chunks. This planner is relatively + /// dumb and does not try to come up with continuous ranges, but you can + /// just set this to a large value to avoid fragmentation. + /// + /// In the very common case where you have small downloads, this will + /// frequently just choose a single peer for the entire download. + /// + /// This is a feature, not a bug. For small downloads, it is not worth + /// the effort to come up with a more sophisticated plan. + stripe_size_log: u8, +} + +impl StripePlanner2 { + pub fn new(seed: u64, stripe_size_log: u8) -> Self { + Self { seed, stripe_size_log } + } + + /// The score function to decide which peer to download from. + fn score(peer: &NodeId, seed: u64) -> u64 { + // todo: use fnv? blake3 is a bit overkill + let mut data = [0u8; 32 + 8]; + data[..32].copy_from_slice(peer.as_bytes()); + data[32..40].copy_from_slice(&seed.to_be_bytes()); + let hash = blake3::hash(&data); + u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) + } +} + +impl DownloadPlanner for StripePlanner2 { + fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { + assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); + options.retain(|_, x| !x.is_empty()); + if options.len() <= 1 { + return; + } + let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); + for range in ranges.windows(2) { + let start = ChunkNum(range[0]); + let end = ChunkNum(range[1]); + let curr = ChunkRanges::from(start..end); + let stripe = range[0] >> self.stripe_size_log; + let mut best_peer = None; + let mut best_score = None; + let mut matching = vec![]; + for (peer, peer_ranges) in options.iter_mut() { + if peer_ranges.contains(&start) { + matching.push((peer, peer_ranges)); + } + } + let mut peer_and_score = matching.iter().map(|(peer, _)| (Self::score(peer, self.seed), peer)).collect::>(); + peer_and_score.sort(); + let peer_to_rank = peer_and_score.into_iter().enumerate().map(|(i, (_, peer))| (*peer, i as u64)).collect::>(); + let n = matching.len() as u64; + for (peer, _) in matching.iter() { + let score = Some((peer_to_rank[*peer] + stripe) % n); + if score > best_score { + best_peer = Some(**peer); + best_score = score; + } + } + for (peer, peer_ranges) in matching { + if *peer != best_peer.unwrap() { + peer_ranges.difference_with(&curr); + } + } + } + options.retain(|_, x| !x.is_empty()); } } + struct PeerDownloadState { id: PeerDownloadId, ranges: ChunkRanges, @@ -265,11 +364,11 @@ struct DownloaderState { // We could use uuid here, but using integers simplifies testing. // // the next subscription id - next_subscription_id: BitmapSubscriptionId, + subscription_id_gen: IdGenerator, // the next discovery id - next_discovery_id: u64, + discovery_id_gen: IdGenerator, // the next peer download id - next_peer_download_id: u64, + peer_download_id_gen: IdGenerator, // the download planner planner: Box, } @@ -281,9 +380,9 @@ impl DownloaderState { downloads: BTreeMap::new(), bitmaps: BTreeMap::new(), discovery: BTreeMap::new(), - next_subscription_id: 0, - next_discovery_id: 0, - next_peer_download_id: 0, + subscription_id_gen: Default::default(), + discovery_id_gen: Default::default(), + peer_download_id_gen: Default::default(), planner, } } @@ -382,23 +481,6 @@ enum Event { } impl DownloaderState { - fn next_subscription_id(&mut self) -> u64 { - let id = self.next_subscription_id; - self.next_subscription_id += 1; - id - } - - fn next_discovery_id(&mut self) -> u64 { - let id = self.next_discovery_id; - self.next_discovery_id += 1; - id - } - - fn next_peer_download_id(&mut self) -> u64 { - let id = self.next_peer_download_id; - self.next_peer_download_id += 1; - id - } fn count_providers(&self, hash: Hash) -> usize { self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash).count() @@ -466,13 +548,13 @@ impl DownloaderState { state.subscription_count += 1; } else { // create a new subscription - let subscription_id = self.next_subscription_id(); + let subscription_id = self.subscription_id_gen.next(); evs.push(Event::SubscribeBitmap { peer: None, hash: request.hash, id: subscription_id }); self.bitmaps.insert((None, request.hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task - let discovery_id = self.next_discovery_id(); + let discovery_id = self.discovery_id_gen.next(); evs.push(Event::StartDiscovery { hash: request.hash, id: discovery_id }); self.discovery.insert(request.hash, discovery_id); } @@ -493,7 +575,7 @@ impl DownloaderState { // create a peer state if it does not exist let _state = self.peers.entry(peer).or_default(); // create a new subscription - let subscription_id = self.next_subscription_id(); + let subscription_id = self.subscription_id_gen.next(); evs.push(Event::SubscribeBitmap { peer: Some(peer), hash, id: subscription_id }); self.bitmaps.insert((Some(peer), hash), PeerBlobState::new(subscription_id)); } @@ -544,6 +626,8 @@ impl DownloaderState { /// Check for completion of a download or of an individual peer download /// /// This must be called after each change of the local bitmap for a hash + /// + /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get(&(None, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all @@ -578,7 +662,11 @@ impl DownloaderState { // reassign the newly available peers without doing a full rebalance if !available.is_empty() { // check if any of the available peers can provide something of the remaining data - let remaining = &download.request.ranges - &self_state.ranges; + let mut remaining = &download.request.ranges - &self_state.ranges; + // subtract the ranges that are already being taken care of by remaining peer downloads + for peer_download in download.peer_downloads.values() { + remaining.difference_with(&peer_download.ranges); + } // see what the new peers can do for us let mut candidates= BTreeMap::new(); for peer in available { @@ -591,7 +679,14 @@ impl DownloaderState { candidates.insert(peer, intersection); } } - // todo: make a plan and create new peer downloads if possible + // deduplicate the ranges + self.planner.plan(hash, &mut candidates); + // start new downloads + for (peer, ranges) in candidates { + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); + download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + } } } // cleanup completed downloads, has to happen later to avoid double mutable borrow @@ -601,6 +696,32 @@ impl DownloaderState { Ok(()) } + /// Look at all downloads for a hash and see start peer downloads for those that do not have any yet + fn start_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + let Some(self_state) = self.bitmaps.get(&(None, hash)) else { + // we don't have the self state yet, so we can't really decide if we need to download anything at all + return Ok(()); + }; + for (_id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash && download.peer_downloads.is_empty()) { + let remaining = &download.request.ranges - &self_state.ranges; + let mut candidates = BTreeMap::new(); + for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash) { + let intersection = &bitmap.ranges & &remaining; + if !intersection.is_empty() { + candidates.insert(peer.unwrap(), intersection); + } + } + self.planner.plan(hash, &mut candidates); + for (peer, ranges) in candidates { + info!(" Starting download from {peer} for {hash} {ranges:?}"); + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); + download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + } + } + Ok(()) + } + fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get(&(None, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all @@ -624,8 +745,7 @@ impl DownloaderState { download.peer_downloads.clear(); for (peer, ranges) in candidates { info!(" Starting download from {peer} for {hash} {ranges:?}"); - let id = self.next_peer_download_id; - self.next_peer_download_id += 1; + let id = self.peer_download_id_gen.next(); evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); } @@ -685,7 +805,7 @@ struct DownloaderActor { peer_download_tasks: BTreeMap>>, discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, - next_download_id: DownloadId, + download_id_gen: IdGenerator, } impl DownloaderActor { @@ -703,7 +823,7 @@ impl DownloaderActor { download_futs: BTreeMap::new(), command_tx: send, command_rx: recv, - next_download_id: 0, + download_id_gen: Default::default(), } } @@ -723,7 +843,7 @@ impl DownloaderActor { UserCommand::Download { request, done, } => { - let id = self.next_download_id(); + let id = self.download_id_gen.next(); self.download_futs.insert(id, done); self.command_tx.send(Command::StartDownload { request, id }).await.ok(); } @@ -733,12 +853,6 @@ impl DownloaderActor { } } - fn next_download_id(&mut self) -> DownloadId { - let id = self.next_download_id; - self.next_download_id += 1; - id - } - fn handle_event(&mut self, ev: Event, depth: usize) { trace!("handle_event {ev:?} {depth}"); match ev { @@ -916,12 +1030,13 @@ fn print_bitmap_compact(iter: impl IntoIterator) -> String { chars } -fn as_bool_iter(x: ChunkRanges, max: u64) -> impl Iterator { +fn as_bool_iter(x: &ChunkRanges, max: u64) -> impl Iterator { let max = x.iter().last().map(|x| match x { RangeSetRange::RangeFrom(_) => max, RangeSetRange::Range(x) => x.end.0, }).unwrap_or_default(); - (0..max).map(move |i| x.contains(&ChunkNum(i))) + let res = (0..max).map(move |i| x.contains(&ChunkNum(i))).collect::>(); + res.into_iter() } /// Given a set of ranges, make them non-overlapping according to some rules. @@ -1010,16 +1125,63 @@ mod tests { use super::*; use bao_tree::ChunkNum; - use iroh::protocol::Router; + use iroh::{protocol::Router, SecretKey}; use testresult::TestResult; #[test] fn print_chunk_range() { let x = chunk_ranges([0..3, 4..30, 40..50]); - let s = print_bitmap_compact(as_bool_iter(x, 50)); + let s = print_bitmap_compact(as_bool_iter(&x, 50)); println!("{}", s); } + fn peer(id: u8) -> NodeId { + let mut secret = [0; 32]; + secret[0] = id; + SecretKey::from(secret).public() + } + + #[test] + fn test_planner() { + let mut planner = StripePlanner2::new(0, 4); + let hash = Hash::new(b"test"); + let mut ranges = make_range_map(&[ + chunk_ranges([0..100]), + chunk_ranges([0..110]), + chunk_ranges([0..120]), + ]); + print_range_map(&ranges); + println!("planning"); + planner.plan(hash, &mut ranges); + print_range_map(&ranges); + println!("---"); + let mut ranges = make_range_map(&[ + chunk_ranges([0..100]), + chunk_ranges([0..110]), + chunk_ranges([0..120]), + chunk_ranges([0..50]), + ]); + print_range_map(&ranges); + println!("planning"); + planner.plan(hash, &mut ranges); + print_range_map(&ranges); + } + + fn make_range_map(ranges: &[ChunkRanges]) -> BTreeMap { + let mut res = BTreeMap::new(); + for (i, range) in ranges.iter().enumerate() { + res.insert(peer(i as u8), range.clone()); + } + res + } + + fn print_range_map(ranges: &BTreeMap) { + for (peer, ranges) in ranges { + let x = print_bitmap(as_bool_iter(ranges, 100)); + println!("{peer}: {x}"); + } + } + #[test] fn test_select_ranges() { let ranges = [ @@ -1029,7 +1191,7 @@ mod tests { ]; let iter = select_ranges(ranges.as_slice(), 8); for (i, range) in ranges.iter().enumerate() { - let bools = as_bool_iter(range.clone(), 100).collect::>(); + let bools = as_bool_iter(&range, 100).collect::>(); println!("{i:4}{}", print_bitmap(bools)); } print!(" "); @@ -1039,7 +1201,7 @@ mod tests { println!(); let ranges = create_ranges(iter); for (i, range) in ranges.iter().enumerate() { - let bools = as_bool_iter(range.clone(), 100).collect::>(); + let bools = as_bool_iter(&range, 100).collect::>(); println!("{i:4}{}", print_bitmap(bools)); } } From 38859542fbcc2a01a8cfb81a59451a5263239daf Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 18:00:46 +0200 Subject: [PATCH 11/47] DRY --- src/downloader2.rs | 106 ++++++++++++++++++++++++++------------------- 1 file changed, 61 insertions(+), 45 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 0a1cecf9e..e320856d0 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -349,7 +349,7 @@ struct DownloaderState { // all bitmaps I am tracking, both for myself and for remote peers // // each item here corresponds to an active subscription - bitmaps: BTreeMap<(Option, Hash), PeerBlobState>, + bitmaps: BTreeMap<(BitmapPeer, Hash), PeerBlobState>, // all active downloads // // these are user downloads. each user download gets split into one or more @@ -388,6 +388,12 @@ impl DownloaderState { } } +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +enum BitmapPeer { + Local, + Remote(NodeId), +} + enum Command { /// A user request to start a download. StartDownload { @@ -401,7 +407,7 @@ enum Command { /// A full bitmap for a blob and a peer Bitmap { /// The peer that sent the bitmap. - peer: Option, + peer: BitmapPeer, /// The blob for which the bitmap is hash: Hash, /// The complete bitmap @@ -413,7 +419,7 @@ enum Command { /// the local bitmap. BitmapUpdate { /// The peer that sent the update. - peer: Option, + peer: BitmapPeer, /// The blob that was updated. hash: Hash, /// The ranges that were added @@ -443,7 +449,7 @@ enum Command { #[derive(Debug, PartialEq, Eq)] enum Event { SubscribeBitmap { - peer: Option, + peer: BitmapPeer, hash: Hash, /// The unique id of the subscription id: u64, @@ -483,7 +489,7 @@ enum Event { impl DownloaderState { fn count_providers(&self, hash: Hash) -> usize { - self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash).count() + self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash).count() } /// Apply a command and return the events that were generated @@ -543,14 +549,14 @@ impl DownloaderState { // ids must be uniquely assigned by the caller! anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitmaps.get_mut(&(None, request.hash)) { + if let Some(state) = self.bitmaps.get_mut(&(BitmapPeer::Local, request.hash)) { // just increment the count state.subscription_count += 1; } else { // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitmap { peer: None, hash: request.hash, id: subscription_id }); - self.bitmaps.insert((None, request.hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Local, hash: request.hash, id: subscription_id }); + self.bitmaps.insert((BitmapPeer::Local, request.hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task @@ -564,7 +570,7 @@ impl DownloaderState { self.stop_download(id, evs)?; } Command::PeerDiscovered { peer, hash } => { - if self.bitmaps.contains_key(&(Some(peer), hash)) { + if self.bitmaps.contains_key(&(BitmapPeer::Remote(peer), hash)) { // we already have a subscription for this peer return Ok(()); }; @@ -576,12 +582,12 @@ impl DownloaderState { let _state = self.peers.entry(peer).or_default(); // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitmap { peer: Some(peer), hash, id: subscription_id }); - self.bitmaps.insert((Some(peer), hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Remote(peer), hash, id: subscription_id }); + self.bitmaps.insert((BitmapPeer::Remote(peer), hash), PeerBlobState::new(subscription_id)); } Command::DropPeer { peer } => { self.bitmaps.retain(|(p, _), state| { - if *p == Some(peer) { + if *p == BitmapPeer::Remote(peer) { // todo: should we emit unsubscribe evs here? evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); return false; @@ -595,22 +601,25 @@ impl DownloaderState { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; state.ranges = bitmap; - if peer.is_none() { + if peer == BitmapPeer::Local { self.check_completion(hash, evs)?; } - self.rebalance_downloads(hash, evs)?; + // we have to call start_downloads even if the local bitmap set, since we don't know in which order local and remote bitmaps arrive + self.start_downloads(hash, evs)?; } Command::BitmapUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; state.ranges |= added; state.ranges &= !removed; - if peer.is_none() { + if peer == BitmapPeer::Local { self.check_completion(hash, evs)?; + } else { + // a local bitmap update does not make more data available, so we don't need to start downloads + self.start_downloads(hash, evs)?; } - self.rebalance_downloads(hash, evs)?; } Command::ChunksDownloaded { time, peer, hash, added } => { - let state = self.bitmaps.get_mut(&(None, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; + let state = self.bitmaps.get_mut(&(BitmapPeer::Local, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; @@ -629,7 +638,7 @@ impl DownloaderState { /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(None, hash)) else { + let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; @@ -670,7 +679,7 @@ impl DownloaderState { // see what the new peers can do for us let mut candidates= BTreeMap::new(); for peer in available { - let Some(peer_state) = self.bitmaps.get(&(Some(peer), hash)) else { + let Some(peer_state) = self.bitmaps.get(&(BitmapPeer::Remote(peer), hash)) else { // weird. we should have a bitmap for this peer since it just completed a download continue; }; @@ -698,17 +707,18 @@ impl DownloaderState { /// Look at all downloads for a hash and see start peer downloads for those that do not have any yet fn start_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(None, hash)) else { + let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; for (_id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash && download.peer_downloads.is_empty()) { let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = BTreeMap::new(); - for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash) { + for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash) { + let BitmapPeer::Remote(peer) = peer else { panic!() }; let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { - candidates.insert(peer.unwrap(), intersection); + candidates.insert(*peer, intersection); } } self.planner.plan(hash, &mut candidates); @@ -723,17 +733,18 @@ impl DownloaderState { } fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(None, hash)) else { + let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = vec![]; - for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| peer.is_some() && *x == hash) { + for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash) { + let BitmapPeer::Remote(peer) = peer else { panic!(); }; let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { - candidates.push((peer.unwrap(), intersection)); + candidates.push((*peer, intersection)); } } info!("Stopping {} old peer downloads", download.peer_downloads.len()); @@ -859,9 +870,9 @@ impl DownloaderActor { Event::SubscribeBitmap { peer, hash, id } => { let send = self.command_tx.clone(); let task = spawn(async move { - let cmd = if peer.is_none() { + let cmd = if peer == BitmapPeer::Local { // we don't have any data, for now - Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::empty() } + Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: ChunkRanges::empty() } } else { // all peers have all the data, for now Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } @@ -970,7 +981,7 @@ async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: sender.send(Command::ChunksDownloaded { time: Duration::ZERO, peer, hash, added: added.clone() }).await.ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; - sender.send(Command::BitmapUpdate { peer: None, hash, added, removed: ChunkRanges::empty() }).await.ok(); + sender.send(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); } } content = next; @@ -1265,36 +1276,41 @@ mod tests { let mut state = DownloaderState::new(noop_planner()); let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: 0 }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: None, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash, bitmap: ChunkRanges::all() }); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: BitmapPeer::Local, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: ChunkRanges::all() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash: unknown_hash, bitmap: ChunkRanges::all() }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: None, hash, bitmap: initial_bitmap.clone() }); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(None, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); - let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: None, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + assert_eq!(state.bitmaps.get(&(BitmapPeer::Local, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(None, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); + assert_eq!(state.bitmaps.get(&(BitmapPeer::Local, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)), }); - assert!(evs.iter().any(|e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); let evs = state.apply_and_get_evs(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(evs.iter().filter(|e| **e == Event::SubscribeBitmap { peer: Some(peer_a), hash, id: 1 }).count() == 1, "adding a new peer for a hash we are interested in should subscribe to the bitmap"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: Some(peer_a), hash, bitmap: chunk_ranges([0..64]) }); - assert!(evs.iter().filter(|e| **e == Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }).count() == 1, "bitmap from a peer should start a download"); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: BitmapPeer::Remote(peer_a), hash, id: 1 }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); + // ChunksDownloaded just updates the peer stats + let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); + assert!(evs.is_empty()); + // Bitmap update does not yet complete the download + let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); + assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats - let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); + let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); assert!(evs.is_empty()); - // Bitmap update for the local bitmap should complete the download - let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: None, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); - assert!(evs.iter().filter(|e| matches!(e, Event::DownloadComplete { .. })).count() == 1, "download should be completed by the data"); - println!("XXX {evs:?}"); + // Final bitmap update for the local bitmap should complete the download + let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); Ok(()) } From ae5f05cc252b4d280408c7771c05e3122bde6201 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 18:51:35 +0200 Subject: [PATCH 12/47] incremental scenario --- src/downloader2.rs | 200 ++++++++++++++++++++++++++------------------- 1 file changed, 115 insertions(+), 85 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index e320856d0..af45879b2 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -130,7 +130,6 @@ struct IdGenerator { } impl IdGenerator { - fn next(&mut self) -> u64 { let id = self.next_id; self.next_id += 1; @@ -142,7 +141,7 @@ impl IdGenerator { /// /// A download planner has the option to be stateful and keep track of plans /// depending on the hash, but many planners will be stateless. -/// +/// /// Planners can do whatever they want with the chunk ranges. Usually, they /// want to deduplicate the ranges, but they could also do other things, like /// eliminate gaps or even extend ranges. The only thing they should not do is @@ -164,7 +163,7 @@ impl DownloadPlanner for NoopPlanner { } /// A download planner that fully removes overlap between peers. -/// +/// /// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, /// and for each stripe decides on a single peer to download from, based on the /// peer id and a random seed. @@ -176,7 +175,7 @@ struct StripePlanner { /// The log of the stripe size in chunks. This planner is relatively /// dumb and does not try to come up with continuous ranges, but you can /// just set this to a large value to avoid fragmentation. - /// + /// /// In the very common case where you have small downloads, this will /// frequently just choose a single peer for the entire download. /// @@ -256,9 +255,8 @@ fn get_continuous_ranges(options: &mut BTreeMap, stripe_siz Some(ranges) } - /// A download planner that fully removes overlap between peers. -/// +/// /// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, /// and for each stripe decides on a single peer to download from, based on the /// peer id and a random seed. @@ -270,7 +268,7 @@ struct StripePlanner2 { /// The log of the stripe size in chunks. This planner is relatively /// dumb and does not try to come up with continuous ranges, but you can /// just set this to a large value to avoid fragmentation. - /// + /// /// In the very common case where you have small downloads, this will /// frequently just choose a single peer for the entire download. /// @@ -337,7 +335,6 @@ impl DownloadPlanner for StripePlanner2 { } } - struct PeerDownloadState { id: PeerDownloadId, ranges: ChunkRanges, @@ -487,30 +484,29 @@ enum Event { } impl DownloaderState { - fn count_providers(&self, hash: Hash) -> usize { self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash).count() } /// Apply a command and return the events that were generated - fn apply_and_get_evs(&mut self, cmd: Command) -> Vec { + fn apply(&mut self, cmd: Command) -> Vec { let mut evs = vec![]; - self.apply(cmd, &mut evs); + self.apply_mut(cmd, &mut evs); evs } - /// Apply a command - fn apply(&mut self, cmd: Command, evs: &mut Vec) { - if let Err(cause) = self.apply0(cmd, evs) { + /// Apply a command, using a mutable reference to the events + fn apply_mut(&mut self, cmd: Command, evs: &mut Vec) { + if let Err(cause) = self.apply_mut_0(cmd, evs) { evs.push(Event::Error { message: format!("{cause}") }); } } /// Stop a download and clean up - /// + /// /// This is called both for stopping a download before completion, and for /// cleaning up after a successful download. - /// + /// /// Cleanup involves emitting events for /// - stopping all peer downloads /// - unsubscribing from bitmaps if needed @@ -543,7 +539,7 @@ impl DownloaderState { } /// Apply a command and bail out on error - fn apply0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { + fn apply_mut_0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! @@ -635,7 +631,7 @@ impl DownloaderState { /// Check for completion of a download or of an individual peer download /// /// This must be called after each change of the local bitmap for a hash - /// + /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { @@ -677,7 +673,7 @@ impl DownloaderState { remaining.difference_with(&peer_download.ranges); } // see what the new peers can do for us - let mut candidates= BTreeMap::new(); + let mut candidates = BTreeMap::new(); for peer in available { let Some(peer_state) = self.bitmaps.get(&(BitmapPeer::Remote(peer), hash)) else { // weird. we should have a bitmap for this peer since it just completed a download @@ -705,7 +701,7 @@ impl DownloaderState { Ok(()) } - /// Look at all downloads for a hash and see start peer downloads for those that do not have any yet + /// Look at all downloads for a hash and start peer downloads for those that do not have any yet fn start_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all @@ -741,7 +737,9 @@ impl DownloaderState { let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = vec![]; for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash) { - let BitmapPeer::Remote(peer) = peer else { panic!(); }; + let BitmapPeer::Remote(peer) = peer else { + panic!(); + }; let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { candidates.push((*peer, intersection)); @@ -843,7 +841,7 @@ impl DownloaderActor { tokio::select! { biased; Some(cmd) = self.command_rx.recv() => { - let evs = self.state.apply_and_get_evs(cmd); + let evs = self.state.apply(cmd); for ev in evs { self.handle_event(ev, 0); } @@ -1025,9 +1023,9 @@ fn print_bitmap_compact(iter: impl IntoIterator) -> String { while let (Some(left), Some(right)) = (iter.next(), iter.next()) { let c = match (left, right) { - (true, true) => '█', // Both pixels are "on" - (true, false) => '▌', // Left pixel is "on" - (false, true) => '▐', // Right pixel is "on" + (true, true) => '█', // Both pixels are "on" + (true, false) => '▌', // Left pixel is "on" + (false, true) => '▐', // Right pixel is "on" (false, false) => ' ', // Both are "off" }; chars.push(c); @@ -1042,10 +1040,14 @@ fn print_bitmap_compact(iter: impl IntoIterator) -> String { } fn as_bool_iter(x: &ChunkRanges, max: u64) -> impl Iterator { - let max = x.iter().last().map(|x| match x { - RangeSetRange::RangeFrom(_) => max, - RangeSetRange::Range(x) => x.end.0, - }).unwrap_or_default(); + let max = x + .iter() + .last() + .map(|x| match x { + RangeSetRange::RangeFrom(_) => max, + RangeSetRange::Range(x) => x.end.0, + }) + .unwrap_or_default(); let res = (0..max).map(move |i| x.contains(&ChunkNum(i))).collect::>(); res.into_iter() } @@ -1069,11 +1071,7 @@ fn select_ranges(ranges: &[ChunkRanges], continuity_bonus: u64) -> Vec String { while let Some(&left) = iter.next() { let right = iter.next(); // Try to fetch the next element - + let left_color = colors.get(left as usize).map(|x| *x).unwrap_or_default(); let right_char = match right { Some(&right) => { let right_color = colors.get(right as usize).map(|x| *x).unwrap_or_default(); // Use ANSI escape codes to color left and right halves of `█` - format!( - "\x1b[38;5;{}m\x1b[48;5;{}m▌\x1b[0m", - left_color, right_color - ) + format!("\x1b[38;5;{}m\x1b[48;5;{}m▌\x1b[0m", left_color, right_color) } None => format!("\x1b[38;5;{}m▌\x1b[0m", left_color), // Handle odd-length case }; @@ -1156,22 +1151,13 @@ mod tests { fn test_planner() { let mut planner = StripePlanner2::new(0, 4); let hash = Hash::new(b"test"); - let mut ranges = make_range_map(&[ - chunk_ranges([0..100]), - chunk_ranges([0..110]), - chunk_ranges([0..120]), - ]); + let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..110]), chunk_ranges([0..120])]); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); print_range_map(&ranges); println!("---"); - let mut ranges = make_range_map(&[ - chunk_ranges([0..100]), - chunk_ranges([0..110]), - chunk_ranges([0..120]), - chunk_ranges([0..50]), - ]); + let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..110]), chunk_ranges([0..120]), chunk_ranges([0..50])]); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); @@ -1195,11 +1181,7 @@ mod tests { #[test] fn test_select_ranges() { - let ranges = [ - chunk_ranges([0..90]), - chunk_ranges([0..100]), - chunk_ranges([0..80]), - ]; + let ranges = [chunk_ranges([0..90]), chunk_ranges([0..100]), chunk_ranges([0..80])]; let iter = select_ranges(ranges.as_slice(), 8); for (i, range) in ranges.iter().enumerate() { let bools = as_bool_iter(&range, 100).collect::>(); @@ -1226,11 +1208,9 @@ mod tests { #[test] fn test_print_colored_bitmap() { - let bitmap = vec![ - 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, - ]; // Notice: odd-length input - - println!("{}", print_colored_bitmap(&bitmap, &[1,2,3,4])); + let bitmap = vec![0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]; // Notice: odd-length input + + println!("{}", print_colored_bitmap(&bitmap, &[1, 2, 3, 4])); } #[cfg(feature = "rpc")] @@ -1262,58 +1242,108 @@ mod tests { evs.iter().filter(|e| *e == ev).count() == 1 } + fn has_all_events(evs: &[Event], evs2: &[&Event]) -> bool { + evs2.iter().all(|ev| has_one_event(evs, ev)) + } + fn has_one_event_matching(evs: &[Event], f: impl Fn(&Event) -> bool) -> bool { evs.iter().filter(|e| f(e)).count() == 1 } + /// Test various things that should produce errors #[test] - fn downloader_state_smoke() -> TestResult<()> { + fn downloader_state_errors() -> TestResult<()> { + use BitmapPeer::*; let _ = tracing_subscriber::fmt::try_init(); - // let me = "0000000000000000000000000000000000000000000000000000000000000000".parse()?; let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply_and_get_evs(super::Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 1 }); - assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: 0 }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: BitmapPeer::Local, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: ChunkRanges::all() }); + let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: ChunkRanges::all() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); + let evs = state.apply(Command::Bitmap { peer: Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); + let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); + Ok(()) + } + + /// Test a simple scenario where a download is started and completed + #[test] + fn downloader_state_smoke() -> TestResult<()> { + use BitmapPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0 }); + assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: 0 }), "starting a download should start a discovery task"); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Local, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: initial_bitmap.clone() }); + let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(BitmapPeer::Local, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); - let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + assert_eq!(state.bitmaps.get(&(Local, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(BitmapPeer::Local, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); - let evs = state.apply_and_get_evs(super::Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: ChunkRanges::from(ChunkNum(0)..ChunkNum(16)), - }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); - let evs = state.apply_and_get_evs(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: BitmapPeer::Remote(peer_a), hash, id: 1 }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); - let evs = state.apply_and_get_evs(Command::Bitmap { peer: BitmapPeer::Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); + assert_eq!(state.bitmaps.get(&(Local, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); + let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Remote(peer_a), hash, id: 1 }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); // ChunksDownloaded just updates the peer stats - let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); + let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); assert!(evs.is_empty()); // Bitmap update does not yet complete the download - let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats - let evs = state.apply_and_get_evs(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); + let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); assert!(evs.is_empty()); // Final bitmap update for the local bitmap should complete the download - let evs = state.apply_and_get_evs(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); Ok(()) } + /// Test a scenario where more data becomes available at the remote peer as the download progresses + #[test] + fn downloader_state_incremental() -> TestResult<()> { + use BitmapPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + // Start a download + state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0 }); + // Initially, we have nothing + state.apply(Command::Bitmap { peer: Local, hash, bitmap: ChunkRanges::empty() }); + // We have a peer for the hash + state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + // We have a bitmap from the peer + let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..32]) }); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitmap from a peer should start a download"); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); + // Bitmap update does not yet complete the download + state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([0..16]), removed: ChunkRanges::empty() }); + // The peer now has more data + state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([32..64]) }); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([16..32]) }); + // Complete the first part of the download + let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + // This triggers cancellation of the first peer download and starting a new one for the remaining data + assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0 }), "first peer download should be stopped"); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); + // Final bitmap update for the local bitmap should complete the download + let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); + assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1 }, &Event::DownloadComplete { id: 0 }, &Event::UnsubscribeBitmap { id: 0 }, &Event::StopDiscovery { id: 0 },]), "download should be completed by the data"); + println!("{evs:?}"); + Ok(()) + } + #[tokio::test] #[cfg(feature = "rpc")] async fn downloader_driver_smoke() -> TestResult<()> { From e46182e37953103cafd5f62220d434e1443fb864 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 22:25:18 +0200 Subject: [PATCH 13/47] delayed rebalance --- src/downloader2.rs | 282 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 220 insertions(+), 62 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index af45879b2..decf808ce 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -24,6 +24,7 @@ use std::{ future::Future, io, sync::Arc, + time::Instant, }; use crate::{ @@ -116,11 +117,13 @@ struct DownloadState { request: DownloadRequest, /// Ongoing downloads peer_downloads: BTreeMap, + /// Set to true if the download needs rebalancing + needs_rebalancing: bool, } impl DownloadState { fn new(request: DownloadRequest) -> Self { - Self { request, peer_downloads: BTreeMap::new() } + Self { request, peer_downloads: BTreeMap::new(), needs_rebalancing: false } } } @@ -137,6 +140,101 @@ impl IdGenerator { } } +/// Wrapper for the downloads map +/// +/// This is so we can later optimize access by fields other than id, such as hash. +#[derive(Default)] +struct Downloads { + by_id: BTreeMap, +} + +impl Downloads { + fn remove(&mut self, id: &DownloadId) -> Option { + self.by_id.remove(id) + } + + fn contains_key(&self, id: &DownloadId) -> bool { + self.by_id.contains_key(id) + } + + fn insert(&mut self, id: DownloadId, state: DownloadState) { + self.by_id.insert(id, state); + } + + fn iter_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { + self.by_id.iter_mut().filter(move |x| x.1.request.hash == hash) + } + + fn iter(&mut self) -> impl Iterator { + self.by_id.iter() + } + + /// Iterate over all downloads for a given hash + fn values_for_hash(&self, hash: Hash) -> impl Iterator { + self.by_id.values().filter(move |x| x.request.hash == hash) + } + + fn values_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { + self.by_id.values_mut().filter(move |x| x.request.hash == hash) + } + + fn by_id_mut(&mut self, id: DownloadId) -> Option<&mut DownloadState> { + self.by_id.get_mut(&id) + } +} + +#[derive(Default)] +struct Bitmaps { + by_peer_and_hash: BTreeMap<(BitmapPeer, Hash), PeerBlobState>, +} + +impl Bitmaps { + fn retain(&mut self, mut f: F) + where + F: FnMut(&(BitmapPeer, Hash), &mut PeerBlobState) -> bool, + { + self.by_peer_and_hash.retain(|k, v| f(k, v)); + } + + fn get(&self, key: &(BitmapPeer, Hash)) -> Option<&PeerBlobState> { + self.by_peer_and_hash.get(key) + } + + fn get_local(&self, hash: Hash) -> Option<&PeerBlobState> { + self.by_peer_and_hash.get(&(BitmapPeer::Local, hash)) + } + + fn get_mut(&mut self, key: &(BitmapPeer, Hash)) -> Option<&mut PeerBlobState> { + self.by_peer_and_hash.get_mut(key) + } + + fn get_local_mut(&mut self, hash: Hash) -> Option<&mut PeerBlobState> { + self.by_peer_and_hash.get_mut(&(BitmapPeer::Local, hash)) + } + + fn insert(&mut self, key: (BitmapPeer, Hash), value: PeerBlobState) { + self.by_peer_and_hash.insert(key, value); + } + + fn contains_key(&self, key: &(BitmapPeer, Hash)) -> bool { + self.by_peer_and_hash.contains_key(key) + } + + fn remote_for_hash(&self, hash: Hash) -> impl Iterator { + self.by_peer_and_hash.iter().filter_map(move |((peer, h), state)| { + if let BitmapPeer::Remote(peer) = peer { + if *h == hash { + Some((peer, state)) + } else { + None + } + } else { + None + } + }) + } +} + /// Trait for a download planner. /// /// A download planner has the option to be stateful and keep track of plans @@ -346,17 +444,16 @@ struct DownloaderState { // all bitmaps I am tracking, both for myself and for remote peers // // each item here corresponds to an active subscription - bitmaps: BTreeMap<(BitmapPeer, Hash), PeerBlobState>, + bitmaps: Bitmaps, // all active downloads // // these are user downloads. each user download gets split into one or more // peer downloads. - downloads: BTreeMap, + downloads: Downloads, // discovery tasks // // there is a discovery task for each blob we are interested in. discovery: BTreeMap, - // Counters to generate unique ids for various requests. // We could use uuid here, but using integers simplifies testing. // @@ -374,8 +471,8 @@ impl DownloaderState { fn new(planner: Box) -> Self { Self { peers: BTreeMap::new(), - downloads: BTreeMap::new(), - bitmaps: BTreeMap::new(), + downloads: Default::default(), + bitmaps: Default::default(), discovery: BTreeMap::new(), subscription_id_gen: Default::default(), discovery_id_gen: Default::default(), @@ -391,6 +488,7 @@ enum BitmapPeer { Remote(NodeId), } +#[derive(Debug)] enum Command { /// A user request to start a download. StartDownload { @@ -441,6 +539,9 @@ enum Command { DropPeer { peer: NodeId }, /// A peer has been discovered PeerDiscovered { peer: NodeId, hash: Hash }, + /// A tick from the driver, for rebalancing + Tick { time: Duration }, + // todo: peerdownloadaborted } #[derive(Debug, PartialEq, Eq)] @@ -484,10 +585,6 @@ enum Event { } impl DownloaderState { - fn count_providers(&self, hash: Hash) -> usize { - self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash).count() - } - /// Apply a command and return the events that were generated fn apply(&mut self, cmd: Command) -> Vec { let mut evs = vec![]; @@ -529,7 +626,7 @@ impl DownloaderState { } true }); - let hash_interest = self.downloads.values().filter(|x| x.request.hash == removed.request.hash).count(); + let hash_interest = self.downloads.values_for_hash(removed.request.hash).count(); if hash_interest == 0 { // stop the discovery task if we were the last one interested in the hash let discovery_id = self.discovery.remove(&removed.request.hash).context(format!("removed unknown discovery task for {}", removed.request.hash))?; @@ -540,12 +637,13 @@ impl DownloaderState { /// Apply a command and bail out on error fn apply_mut_0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { + trace!("handle_command {cmd:?}"); match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitmaps.get_mut(&(BitmapPeer::Local, request.hash)) { + if let Some(state) = self.bitmaps.get_local_mut(request.hash) { // just increment the count state.subscription_count += 1; } else { @@ -571,7 +669,7 @@ impl DownloaderState { return Ok(()); }; // check if anybody needs this peer - if !self.downloads.iter().any(|(_id, state)| state.request.hash == hash) { + if self.downloads.values_for_hash(hash).next().is_none() { return Ok(()); } // create a peer state if it does not exist @@ -596,26 +694,43 @@ impl DownloaderState { Command::Bitmap { peer, hash, bitmap } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; - state.ranges = bitmap; if peer == BitmapPeer::Local { + state.ranges = bitmap; self.check_completion(hash, evs)?; + } else { + // We got an entirely new peer, mark all affected downloads for rebalancing + for download in self.downloads.values_mut_for_hash(hash) { + if bitmap.intersects(&download.request.ranges) { + download.needs_rebalancing = true; + } + } + state.ranges = bitmap; } // we have to call start_downloads even if the local bitmap set, since we don't know in which order local and remote bitmaps arrive self.start_downloads(hash, evs)?; } Command::BitmapUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; - state.ranges |= added; - state.ranges &= !removed; if peer == BitmapPeer::Local { + state.ranges |= added; + state.ranges &= !removed; self.check_completion(hash, evs)?; } else { + // We got more data for this hash, mark all affected downloads for rebalancing + for download in self.downloads.values_mut_for_hash(hash) { + // if removed is non-empty, that is so weird that we just rebalance in any case + if !removed.is_empty() || added.intersects(&download.request.ranges) { + download.needs_rebalancing = true; + } + } + state.ranges |= added; + state.ranges &= !removed; // a local bitmap update does not make more data available, so we don't need to start downloads self.start_downloads(hash, evs)?; } } Command::ChunksDownloaded { time, peer, hash, added } => { - let state = self.bitmaps.get_mut(&(BitmapPeer::Local, hash)).context(format!("chunks downloaded before having local bitmap for {hash}"))?; + let state = self.bitmaps.get_local_mut(hash).context(format!("chunks downloaded before having local bitmap for {hash}"))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; @@ -624,6 +739,40 @@ impl DownloaderState { let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); } + Command::Tick { time } => { + let window = 10; + let horizon = time.saturating_sub(Duration::from_secs(window)); + // clean up download history + let mut to_rebalance = vec![]; + for (peer, state) in self.peers.iter_mut() { + state.download_history.retain(|(duration, _)| *duration > horizon); + let mut sum_total = 0; + let mut sum_useful = 0; + for (_, (total, useful)) in state.download_history.iter() { + sum_total += total; + sum_useful += useful; + } + let speed_useful = (sum_useful as f64) / (window as f64); + let speed_total = (sum_total as f64) / (window as f64); + trace!("peer {peer} download speed {speed_total} cps total, {speed_useful} cps useful"); + } + + for (id, download) in self.downloads.iter() { + if !download.needs_rebalancing { + // nothing has changed that affects this download + continue; + } + let n_peers = self.bitmaps.remote_for_hash(download.request.hash).count(); + if download.peer_downloads.len() >= n_peers { + // we are already downloading from all peers for this hash + continue; + } + to_rebalance.push(*id); + } + for id in to_rebalance { + self.rebalance_download(id, evs)?; + } + } } Ok(()) } @@ -634,12 +783,12 @@ impl DownloaderState { /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { + let Some(self_state) = self.bitmaps.get_local(hash) else { // 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 mut completed = vec![]; - for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { + for (id, download) in self.downloads.iter_mut_for_hash(hash) { // 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 @@ -703,15 +852,14 @@ impl DownloaderState { /// Look at all downloads for a hash and start peer downloads for those that do not have any yet fn start_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { + let Some(self_state) = self.bitmaps.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; - for (_id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash && download.peer_downloads.is_empty()) { + for download in self.downloads.values_mut_for_hash(hash).filter(|download| download.peer_downloads.is_empty()) { let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = BTreeMap::new(); - for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash) { - let BitmapPeer::Remote(peer) = peer else { panic!() }; + for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { candidates.insert(*peer, intersection); @@ -728,37 +876,37 @@ impl DownloaderState { Ok(()) } - fn rebalance_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get(&(BitmapPeer::Local, hash)) else { + /// rebalance a single download + fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { + let download = self.downloads.by_id_mut(id).context(format!("rebalancing unknown download {id}"))?; + download.needs_rebalancing = false; + tracing::error!("Rebalancing download {id} {:?}", download.request); + let hash = download.request.hash; + let Some(self_state) = self.bitmaps.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; - for (id, download) in self.downloads.iter_mut().filter(|(_id, download)| download.request.hash == hash) { - let remaining = &download.request.ranges - &self_state.ranges; - let mut candidates = vec![]; - for ((peer, _), bitmap) in self.bitmaps.iter().filter(|((peer, x), _)| *peer != BitmapPeer::Local && *x == hash) { - let BitmapPeer::Remote(peer) = peer else { - panic!(); - }; - let intersection = &bitmap.ranges & &remaining; - if !intersection.is_empty() { - candidates.push((*peer, intersection)); - } - } - info!("Stopping {} old peer downloads", download.peer_downloads.len()); - for (_, state) in &download.peer_downloads { - // stop all downloads - evs.push(Event::StopPeerDownload { id: state.id }); - } - info!("Creating {} new peer downloads", candidates.len()); - download.peer_downloads.clear(); - for (peer, ranges) in candidates { - info!(" Starting download from {peer} for {hash} {ranges:?}"); - let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); - download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + let remaining = &download.request.ranges - &self_state.ranges; + let mut candidates = vec![]; + for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { + let intersection = &bitmap.ranges & &remaining; + if !intersection.is_empty() { + candidates.push((*peer, intersection)); } } + info!("Stopping {} old peer downloads", download.peer_downloads.len()); + for (_, state) in &download.peer_downloads { + // stop all downloads + evs.push(Event::StopPeerDownload { id: state.id }); + } + info!("Creating {} new peer downloads", candidates.len()); + download.peer_downloads.clear(); + for (peer, ranges) in candidates { + info!(" Starting download from {peer} for {hash} {ranges:?}"); + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); + download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + } Ok(()) } } @@ -814,7 +962,10 @@ struct DownloaderActor { peer_download_tasks: BTreeMap>>, discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, + /// Id generator for download ids download_id_gen: IdGenerator, + /// The time when the actor was started, serves as the epoch for time messages to the state machine + start: Instant, } impl DownloaderActor { @@ -833,19 +984,15 @@ impl DownloaderActor { command_tx: send, command_rx: recv, download_id_gen: Default::default(), + start: Instant::now(), } } async fn run(mut self, mut channel: mpsc::Receiver) { + let mut ticks = tokio::time::interval(Duration::from_millis(50)); loop { tokio::select! { biased; - Some(cmd) = self.command_rx.recv() => { - let evs = self.state.apply(cmd); - for ev in evs { - self.handle_event(ev, 0); - } - }, Some(cmd) = channel.recv() => { debug!("user command {cmd:?}"); match cmd { @@ -858,6 +1005,16 @@ impl DownloaderActor { } } }, + Some(cmd) = self.command_rx.recv() => { + let evs = self.state.apply(cmd); + for ev in evs { + self.handle_event(ev, 0); + } + }, + _ = ticks.tick() => { + let time = self.start.elapsed(); + self.command_tx.send(Command::Tick { time }).await.ok(); + }, } } } @@ -900,6 +1057,7 @@ impl DownloaderActor { let send = self.command_tx.clone(); let endpoint = self.endpoint.clone(); let store = self.store.clone(); + let start = self.start; let task = self.local_pool.spawn(move || async move { info!("Connecting to peer {peer}"); let conn = endpoint.connect(peer, crate::ALPN).await?; @@ -909,7 +1067,7 @@ impl DownloaderActor { info!("starting download from {peer} for {hash} {ranges:?}"); let request = GetRequest::new(hash, ranges); let initial = crate::get::fsm::start(conn, request); - stream_to_db(initial, store, hash, peer, send).await?; + stream_to_db(initial, store, hash, peer, send, start).await?; anyhow::Ok(()) }); self.peer_download_tasks.insert(id, task); @@ -951,7 +1109,7 @@ impl ContentDiscovery for StaticContentDiscovery { } } -async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: NodeId, sender: mpsc::Sender) -> io::Result { +async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: NodeId, sender: mpsc::Sender, start: Instant) -> io::Result { // connect let connected = initial.next().await?; // read the first bytes @@ -975,8 +1133,8 @@ async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: } BaoContentItem::Leaf(leaf) => { let start_chunk = leaf.offset / 1024; - let added = ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 1)); - sender.send(Command::ChunksDownloaded { time: Duration::ZERO, peer, hash, added: added.clone() }).await.ok(); + let added = ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16)); + sender.send(Command::ChunksDownloaded { time: start.elapsed(), peer, hash, added: added.clone() }).await.ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; sender.send(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); @@ -1282,10 +1440,10 @@ mod tests { let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(Local, hash)).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get(&(Local, hash)).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); + assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Remote(peer_a), hash, id: 1 }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); @@ -1368,7 +1526,7 @@ mod tests { let _ = tracing_subscriber::fmt::try_init(); let data = vec![0u8; 1024 * 1024]; let mut nodes = vec![]; - for _i in 0..10 { + for _i in 0..2 { nodes.push(make_test_node(&data).await?); } let peers = nodes.iter().map(|(_, peer, _)| *peer).collect::>(); From a14ab1f1a8b282abc6995bd8ef0e7a944cad6340 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 31 Jan 2025 22:36:39 +0200 Subject: [PATCH 14/47] plug in planner --- src/downloader2.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index decf808ce..cabe543fe 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -887,13 +887,14 @@ impl DownloaderState { return Ok(()); }; let remaining = &download.request.ranges - &self_state.ranges; - let mut candidates = vec![]; + let mut candidates = BTreeMap::new(); for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { let intersection = &bitmap.ranges & &remaining; if !intersection.is_empty() { - candidates.push((*peer, intersection)); + candidates.insert(*peer, intersection); } } + self.planner.plan(hash, &mut candidates); info!("Stopping {} old peer downloads", download.peer_downloads.len()); for (_, state) in &download.peer_downloads { // stop all downloads @@ -1536,7 +1537,7 @@ mod tests { let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, local_pool, noop_planner()); + let downloader = Downloader::new(endpoint, store, discovery, local_pool, Box::new(StripePlanner2::new(0, 8))); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; From 3bc8b3500565624a232f1b7f8fd9316b7097f51c Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 3 Feb 2025 11:01:10 +0200 Subject: [PATCH 15/47] newtypes --- src/downloader2.rs | 113 +++++++++++++++++++++++++-------------------- 1 file changed, 64 insertions(+), 49 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index cabe543fe..13adf08e3 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -20,11 +20,7 @@ //! The [DownloaderDriver] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, - future::Future, - io, - sync::Arc, - time::Instant, + collections::{BTreeMap, BTreeSet, VecDeque}, future::Future, io, marker::PhantomData, sync::Arc, time::Instant }; use crate::{ @@ -48,11 +44,17 @@ use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; -/// todo: make newtypes? -type DownloadId = u64; -type BitmapSubscriptionId = u64; -type DiscoveryId = u64; -type PeerDownloadId = u64; +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +struct DownloadId(u64); + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +struct DiscoveryId(u64); + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +struct PeerDownloadId(u64); + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +struct BitmapSubscriptionId(u64); /// Announce kind #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] @@ -127,16 +129,26 @@ impl DownloadState { } } -#[derive(Debug, Default)] -struct IdGenerator { +#[derive(Debug)] +struct IdGenerator { next_id: u64, + _p: PhantomData, } -impl IdGenerator { - fn next(&mut self) -> u64 { +impl Default for IdGenerator { + fn default() -> Self { + Self { next_id: 0, _p: PhantomData } + } +} + +impl IdGenerator +where + T: From + Copy, +{ + fn next(&mut self) -> T { let id = self.next_id; self.next_id += 1; - id + T::from(id) } } @@ -458,11 +470,11 @@ struct DownloaderState { // We could use uuid here, but using integers simplifies testing. // // the next subscription id - subscription_id_gen: IdGenerator, + subscription_id_gen: IdGenerator, // the next discovery id - discovery_id_gen: IdGenerator, + discovery_id_gen: IdGenerator, // the next peer download id - peer_download_id_gen: IdGenerator, + peer_download_id_gen: IdGenerator, // the download planner planner: Box, } @@ -495,10 +507,10 @@ enum Command { /// The download request request: DownloadRequest, /// The unique id, to be assigned by the caller - id: u64, + id: DownloadId, }, /// A user request to abort a download. - StopDownload { id: u64 }, + StopDownload { id: DownloadId }, /// A full bitmap for a blob and a peer Bitmap { /// The peer that sent the bitmap. @@ -550,35 +562,35 @@ enum Event { peer: BitmapPeer, hash: Hash, /// The unique id of the subscription - id: u64, + id: BitmapSubscriptionId, }, UnsubscribeBitmap { /// The unique id of the subscription - id: u64, + id: BitmapSubscriptionId, }, StartDiscovery { hash: Hash, /// The unique id of the discovery task - id: u64, + id: DiscoveryId, }, StopDiscovery { /// The unique id of the discovery task - id: u64, + id: DiscoveryId, }, StartPeerDownload { /// The unique id of the peer download task - id: u64, + id: PeerDownloadId, peer: NodeId, hash: Hash, ranges: ChunkRanges, }, StopPeerDownload { /// The unique id of the peer download task - id: u64, + id: PeerDownloadId, }, DownloadComplete { /// The unique id of the user download - id: u64, + id: DownloadId, }, /// An error that stops processing the command Error { message: String }, @@ -608,8 +620,8 @@ impl DownloaderState { /// - stopping all peer downloads /// - unsubscribing from bitmaps if needed /// - stopping the discovery task if needed - fn stop_download(&mut self, id: u64, evs: &mut Vec) -> anyhow::Result<()> { - let removed = self.downloads.remove(&id).context(format!("removed unknown download {id}"))?; + fn stop_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { + let removed = self.downloads.remove(&id).context(format!("removed unknown download {id:?}"))?; let removed_hash = removed.request.hash; // stop associated peer downloads for peer_download in removed.peer_downloads.values() { @@ -641,24 +653,27 @@ impl DownloaderState { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! - anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id}"); + anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id:?}"); + let hash = request.hash; // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitmaps.get_local_mut(request.hash) { + if let Some(state) = self.bitmaps.get_local_mut(hash) { // just increment the count state.subscription_count += 1; } else { // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Local, hash: request.hash, id: subscription_id }); - self.bitmaps.insert((BitmapPeer::Local, request.hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Local, hash, id: subscription_id }); + self.bitmaps.insert((BitmapPeer::Local, hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task - let discovery_id = self.discovery_id_gen.next(); - evs.push(Event::StartDiscovery { hash: request.hash, id: discovery_id }); - self.discovery.insert(request.hash, discovery_id); + let id = self.discovery_id_gen.next(); + evs.push(Event::StartDiscovery { hash, id }); + self.discovery.insert(request.hash, id ); } self.downloads.insert(id, DownloadState::new(request)); + self.check_completion(hash, evs)?; + self.start_downloads(hash, evs)?; } Command::StopDownload { id } => { self.stop_download(id, evs)?; @@ -878,9 +893,9 @@ impl DownloaderState { /// rebalance a single download fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { - let download = self.downloads.by_id_mut(id).context(format!("rebalancing unknown download {id}"))?; + let download = self.downloads.by_id_mut(id).context(format!("rebalancing unknown download {id:?}"))?; download.needs_rebalancing = false; - tracing::error!("Rebalancing download {id} {:?}", download.request); + tracing::error!("Rebalancing download {id:?} {:?}", download.request); let hash = download.request.hash; let Some(self_state) = self.bitmaps.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all @@ -964,7 +979,7 @@ struct DownloaderActor { discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, /// Id generator for download ids - download_id_gen: IdGenerator, + download_id_gen: IdGenerator, /// The time when the actor was started, serves as the epoch for time messages to the state machine start: Instant, } @@ -1435,9 +1450,9 @@ mod tests { let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0 }); - assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: 0 }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Local, hash, id: 0 }), "starting a download should subscribe to the local bitmap"); + let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); + assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "starting a download should start a discovery task"); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Local, hash, id: BitmapSubscriptionId(0) }), "starting a download should subscribe to the local bitmap"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); @@ -1446,9 +1461,9 @@ mod tests { assert!(evs.is_empty()); assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Remote(peer_a), hash, id: 1 }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: PeerDownloadId(0), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); // ChunksDownloaded just updates the peer stats let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); assert!(evs.is_empty()); @@ -1473,14 +1488,14 @@ mod tests { let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); // Start a download - state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0 }); + state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); // Initially, we have nothing state.apply(Command::Bitmap { peer: Local, hash, bitmap: ChunkRanges::empty() }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitmap from the peer let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..32]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0, peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitmap from a peer should start a download"); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitmap from a peer should start a download"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); // Bitmap update does not yet complete the download @@ -1492,13 +1507,13 @@ mod tests { // Complete the first part of the download let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); // This triggers cancellation of the first peer download and starting a new one for the remaining data - assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0 }), "first peer download should be stopped"); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1, peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); + assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), "first peer download should be stopped"); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1.into(), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); // Final bitmap update for the local bitmap should complete the download let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); - assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1 }, &Event::DownloadComplete { id: 0 }, &Event::UnsubscribeBitmap { id: 0 }, &Event::StopDiscovery { id: 0 },]), "download should be completed by the data"); + assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1.into() }, &Event::DownloadComplete { id: 0.into() }, &Event::UnsubscribeBitmap { id: 0.into() }, &Event::StopDiscovery { id: 0.into() },]), "download should be completed by the data"); println!("{evs:?}"); Ok(()) } From a94d8002f2447e0fcd0c36942dca264a78658755 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 3 Feb 2025 14:21:31 +0200 Subject: [PATCH 16/47] Add and handle PeerDownloadComplete This is mostly for the case where a peer download aborts with an error, since if it just downloads the required data we will make progress when the local bitmap changes. --- src/downloader2.rs | 121 +++++++++++++++++++++++++++++++++------------ 1 file changed, 89 insertions(+), 32 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 13adf08e3..b16ece7bd 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -25,7 +25,7 @@ use std::{ use crate::{ get::{ - fsm::{AtInitial, BlobContentNext, ConnectedNext, EndBlobNext}, + fsm::{BlobContentNext, ConnectedNext, EndBlobNext}, Stats, }, protocol::{GetRequest, RangeSpec, RangeSpecSeq}, @@ -81,6 +81,32 @@ trait ContentDiscovery: Send + 'static { fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Send + Unpin + 'static; } +trait BitmapSubscription: Send + 'static { + /// Subscribe to a bitmap + fn subscribe(&mut self, peer: BitmapPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; +} + +enum BitmapSubscriptionEvent { + Bitmap { ranges: ChunkRanges }, + BitmapUpdate { added: ChunkRanges, removed: ChunkRanges}, +} + +struct TestBitmapSubscription; + +impl BitmapSubscription for TestBitmapSubscription { + fn subscribe(&mut self, peer: BitmapPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { + let bitmap = match peer { + BitmapPeer::Local => { + ChunkRanges::empty() + } + BitmapPeer::Remote(_) => { + ChunkRanges::all() + } + }; + futures_lite::stream::once(BitmapSubscriptionEvent::Bitmap { ranges: bitmap }).chain(futures_lite::stream::pending()) + } +} + /// Global information about a peer #[derive(Debug, Default)] struct PeerState { @@ -193,6 +219,10 @@ impl Downloads { fn by_id_mut(&mut self, id: DownloadId) -> Option<&mut DownloadState> { self.by_id.get_mut(&id) } + + fn by_peer_download_id_mut(&mut self, id: PeerDownloadId) -> Option<(&DownloadId, &mut DownloadState)> { + self.by_id.iter_mut().filter(|(k, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)).next() + } } #[derive(Default)] @@ -547,6 +577,11 @@ enum Command { /// The ranges that were added locally added: ChunkRanges, }, + /// A peer download has completed + PeerDownloadComplete { + id: PeerDownloadId, + result: anyhow::Result, + }, /// Stop tracking a peer for all blobs, for whatever reason DropPeer { peer: NodeId }, /// A peer has been discovered @@ -672,8 +707,18 @@ impl DownloaderState { self.discovery.insert(request.hash, id ); } self.downloads.insert(id, DownloadState::new(request)); - self.check_completion(hash, evs)?; - self.start_downloads(hash, evs)?; + self.check_completion(hash, Some(id), evs)?; + self.start_downloads(hash, Some(id), evs)?; + } + Command::PeerDownloadComplete { id, .. } => { + let Some((download_id, download)) = self.downloads.by_peer_download_id_mut(id) else { + // the download was already removed + return Ok(()); + }; + let download_id = *download_id; + let hash = download.request.hash; + download.peer_downloads.retain(|_, v| v.id != id); + self.start_downloads(hash, Some(download_id), evs)?; } Command::StopDownload { id } => { self.stop_download(id, evs)?; @@ -711,7 +756,7 @@ impl DownloaderState { let _chunks = total_chunks(&bitmap).context("open range")?; if peer == BitmapPeer::Local { state.ranges = bitmap; - self.check_completion(hash, evs)?; + self.check_completion(hash, None, evs)?; } else { // We got an entirely new peer, mark all affected downloads for rebalancing for download in self.downloads.values_mut_for_hash(hash) { @@ -722,14 +767,14 @@ impl DownloaderState { state.ranges = bitmap; } // we have to call start_downloads even if the local bitmap set, since we don't know in which order local and remote bitmaps arrive - self.start_downloads(hash, evs)?; + self.start_downloads(hash, None, evs)?; } Command::BitmapUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; if peer == BitmapPeer::Local { state.ranges |= added; state.ranges &= !removed; - self.check_completion(hash, evs)?; + self.check_completion(hash, None, evs)?; } else { // We got more data for this hash, mark all affected downloads for rebalancing for download in self.downloads.values_mut_for_hash(hash) { @@ -741,7 +786,7 @@ impl DownloaderState { state.ranges |= added; state.ranges &= !removed; // a local bitmap update does not make more data available, so we don't need to start downloads - self.start_downloads(hash, evs)?; + self.start_downloads(hash, None, evs)?; } } Command::ChunksDownloaded { time, peer, hash, added } => { @@ -797,13 +842,16 @@ impl DownloaderState { /// This must be called after each change of the local bitmap for a hash /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. - fn check_completion(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + fn check_completion(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get_local(hash) else { // 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 mut completed = vec![]; for (id, download) in self.downloads.iter_mut_for_hash(hash) { + if just_id.is_some() && just_id != Some(*id) { + continue; + } // 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 @@ -866,12 +914,15 @@ impl DownloaderState { } /// Look at all downloads for a hash and start peer downloads for those that do not have any yet - fn start_downloads(&mut self, hash: Hash, evs: &mut Vec) -> anyhow::Result<()> { + fn start_downloads(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { let Some(self_state) = self.bitmaps.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; - for download in self.downloads.values_mut_for_hash(hash).filter(|download| download.peer_downloads.is_empty()) { + for (id, download) in self.downloads.iter_mut_for_hash(hash).filter(|(_, download)| download.peer_downloads.is_empty()) { + if just_id.is_some() && just_id != Some(*id) { + continue; + } let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = BTreeMap::new(); for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { @@ -952,8 +1003,8 @@ impl Downloader { Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool, planner: Box) -> Self { - let actor = DownloaderActor::new(endpoint, store, discovery, local_pool, planner); + fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { + let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitmap, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); Self { send, task } @@ -966,7 +1017,7 @@ enum UserCommand { Download { request: DownloadRequest, done: tokio::sync::oneshot::Sender<()> }, } -struct DownloaderActor { +struct DownloaderActor { local_pool: LocalPool, endpoint: Endpoint, command_rx: mpsc::Receiver, @@ -974,8 +1025,9 @@ struct DownloaderActor { state: DownloaderState, store: S, discovery: D, + subscribe_bitmap: B, download_futs: BTreeMap>, - peer_download_tasks: BTreeMap>>, + peer_download_tasks: BTreeMap>, discovery_tasks: BTreeMap>, bitmap_subscription_tasks: BTreeMap>, /// Id generator for download ids @@ -984,8 +1036,8 @@ struct DownloaderActor { start: Instant, } -impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: D, local_pool: LocalPool, planner: Box) -> Self { +impl DownloaderActor { + fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { let (send, recv) = mpsc::channel(256); Self { local_pool, @@ -993,6 +1045,7 @@ impl DownloaderActor { state: DownloaderState::new(planner), store, discovery, + subscribe_bitmap, peer_download_tasks: BTreeMap::new(), discovery_tasks: BTreeMap::new(), bitmap_subscription_tasks: BTreeMap::new(), @@ -1074,18 +1127,7 @@ impl DownloaderActor { let endpoint = self.endpoint.clone(); let store = self.store.clone(); let start = self.start; - let task = self.local_pool.spawn(move || async move { - info!("Connecting to peer {peer}"); - let conn = endpoint.connect(peer, crate::ALPN).await?; - info!("Got connection to peer {peer}"); - let spec = RangeSpec::new(ranges); - let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); - info!("starting download from {peer} for {hash} {ranges:?}"); - let request = GetRequest::new(hash, ranges); - let initial = crate::get::fsm::start(conn, request); - stream_to_db(initial, store, hash, peer, send, start).await?; - anyhow::Ok(()) - }); + let task = self.local_pool.spawn(move || peer_download_task(id, endpoint, store, hash, peer, ranges, send, start)); self.peer_download_tasks.insert(id, task); } Event::UnsubscribeBitmap { id } => { @@ -1125,12 +1167,25 @@ impl ContentDiscovery for StaticContentDiscovery { } } -async fn stream_to_db(initial: AtInitial, store: S, hash: Hash, peer: NodeId, sender: mpsc::Sender, start: Instant) -> io::Result { +async fn peer_download_task(id: PeerDownloadId, endpoint: Endpoint, store: S, hash: Hash, peer: NodeId, ranges: ChunkRanges, sender: mpsc::Sender, start: Instant) { + let result = peer_download(endpoint, store, hash, peer, ranges, &sender, start).await; + sender.send(Command::PeerDownloadComplete { id, result }).await.ok(); +} + +async fn peer_download(endpoint: Endpoint, store: S, hash: Hash, peer: NodeId, ranges: ChunkRanges, sender: &mpsc::Sender, start: Instant) -> anyhow::Result { + info!("Connecting to peer {peer}"); + let conn = endpoint.connect(peer, crate::ALPN).await?; + info!("Got connection to peer {peer}"); + let spec = RangeSpec::new(ranges); + let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); + info!("starting download from {peer} for {hash} {ranges:?}"); + let request = GetRequest::new(hash, ranges); + let initial = crate::get::fsm::start(conn, request); // connect let connected = initial.next().await?; // read the first bytes let ConnectedNext::StartRoot(start_root) = connected.next().await? else { - return Err(io::Error::new(io::ErrorKind::Other, "expected start root")); + return Err(io::Error::new(io::ErrorKind::Other, "expected start root").into()); }; let header = start_root.next(); @@ -1526,8 +1581,9 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; + let bitmap_subscription = TestBitmapSubscription; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, local_pool, noop_planner()); + let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, noop_planner()); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); fut.await?; @@ -1551,8 +1607,9 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; + let bitmap_subscription = TestBitmapSubscription; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, local_pool, Box::new(StripePlanner2::new(0, 8))); + let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; From 871b3437b8db5d669e4668234a99bcf9b14f423f Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 3 Feb 2025 15:51:51 +0200 Subject: [PATCH 17/47] WIP --- src/downloader2.rs | 71 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 53 insertions(+), 18 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index b16ece7bd..32da7521a 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -4,7 +4,7 @@ //! //! The [DownloaderState] is a synchronous state machine containing the logic. //! It gets commands and produces events. It does not do any IO and also does -//! not have any time dependency. So [DownloaderState::apply_and_get_evs] is a +//! not have any time dependency. So [DownloaderState::apply] is a //! pure function of the state and the command and can therefore be tested //! easily. //! @@ -17,7 +17,7 @@ //! the bitmaps are not very fragmented. We can introduce an even more optimized //! bitmap type, or prevent fragmentation. //! -//! The [DownloaderDriver] is the asynchronous driver for the state machine. It +//! The [DownloaderActor] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, future::Future, io, marker::PhantomData, sync::Arc, time::Instant @@ -100,7 +100,7 @@ impl BitmapSubscription for TestBitmapSubscription { ChunkRanges::empty() } BitmapPeer::Remote(_) => { - ChunkRanges::all() + ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } }; futures_lite::stream::once(BitmapSubscriptionEvent::Bitmap { ranges: bitmap }).chain(futures_lite::stream::pending()) @@ -588,7 +588,6 @@ enum Command { PeerDiscovered { peer: NodeId, hash: Hash }, /// A tick from the driver, for rebalancing Tick { time: Duration }, - // todo: peerdownloadaborted } #[derive(Debug, PartialEq, Eq)] @@ -1058,7 +1057,7 @@ impl DownloaderActor) { - let mut ticks = tokio::time::interval(Duration::from_millis(50)); + let mut ticks = tokio::time::interval(Duration::from_millis(100)); loop { tokio::select! { biased; @@ -1093,16 +1092,28 @@ impl DownloaderActor { let send = self.command_tx.clone(); + let mut stream = self.subscribe_bitmap.subscribe(peer, hash); let task = spawn(async move { - let cmd = if peer == BitmapPeer::Local { - // we don't have any data, for now - Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: ChunkRanges::empty() } - } else { - // all peers have all the data, for now - Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } - }; - send.send(cmd).await.ok(); - futures_lite::future::pending().await + while let Some(ev) = stream.next().await { + let cmd = match ev { + BitmapSubscriptionEvent::Bitmap { ranges } => { + Command::Bitmap { peer, hash, bitmap: ranges } + } + BitmapSubscriptionEvent::BitmapUpdate { added, removed } => { + Command::BitmapUpdate { peer, hash, added, removed } + } + }; + send.send(cmd).await.ok(); + } + // let cmd = if peer == BitmapPeer::Local { + // // we don't have any data, for now + // Command::Bitmap { peer: BitmapPeer::Local, hash, bitmap: ChunkRanges::empty() } + // } else { + // // all peers have all the data, for now + // Command::Bitmap { peer, hash, bitmap: ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } + // }; + // send.send(cmd).await.ok(); + // futures_lite::future::pending().await }); self.bitmap_subscription_tasks.insert(id, task); } @@ -1377,16 +1388,35 @@ mod tests { } #[test] - fn test_planner() { + fn test_planner_1() { let mut planner = StripePlanner2::new(0, 4); let hash = Hash::new(b"test"); - let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..110]), chunk_ranges([0..120])]); + let mut ranges = make_range_map(&[chunk_ranges([0..50]), chunk_ranges([50..100])]); + println!(""); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); print_range_map(&ranges); - println!("---"); + } + + #[test] + fn test_planner_2() { + let mut planner = StripePlanner2::new(0, 4); + let hash = Hash::new(b"test"); + let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..100]), chunk_ranges([0..100])]); + println!(""); + print_range_map(&ranges); + println!("planning"); + planner.plan(hash, &mut ranges); + print_range_map(&ranges); + } + + #[test] + fn test_planner_3() { + let mut planner = StripePlanner2::new(0, 4); + let hash = Hash::new(b"test"); let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..110]), chunk_ranges([0..120]), chunk_ranges([0..50])]); + println!(""); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); @@ -1444,6 +1474,11 @@ mod tests { #[cfg(feature = "rpc")] async fn make_test_node(data: &[u8]) -> anyhow::Result<(Router, NodeId, Hash)> { + // let noop_subscriber = tracing_subscriber::fmt::Subscriber::builder() + // .with_writer(io::sink) // all output is discarded + // .with_max_level(tracing::level_filters::LevelFilter::OFF) // effectively disable logging + // .finish(); + // let noop_dispatch = tracing::Dispatch::new(noop_subscriber); let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; let node_id = endpoint.node_id(); let store = crate::store::mem::Store::new(); @@ -1598,7 +1633,7 @@ mod tests { let _ = tracing_subscriber::fmt::try_init(); let data = vec![0u8; 1024 * 1024]; let mut nodes = vec![]; - for _i in 0..2 { + for _i in 0..4 { nodes.push(make_test_node(&data).await?); } let peers = nodes.iter().map(|(_, peer, _)| *peer).collect::>(); From 414715fc82b0f4815f65d5248cc79ad2b34f56fb Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 4 Feb 2025 12:40:03 +0200 Subject: [PATCH 18/47] remame bitmap to bitfield part 1 --- src/downloader2.rs | 117 +++++++++++++++++++++++---------------------- 1 file changed, 60 insertions(+), 57 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 32da7521a..2630f1ef8 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -12,10 +12,10 @@ //! of the caller to generate unique ids. We could use random uuids here, but using //! integers simplifies testing. //! -//! Inside the state machine, we use [ChunkRanges] to represent avaiability bitmaps. -//! We treat operations on such bitmaps as very cheap, which is the case as long as -//! the bitmaps are not very fragmented. We can introduce an even more optimized -//! bitmap type, or prevent fragmentation. +//! Inside the state machine, we use [ChunkRanges] to represent avaiability bitfields. +//! We treat operations on such bitfields as very cheap, which is the case as long as +//! the bitfields are not very fragmented. We can introduce an even more optimized +//! bitfield type, or prevent fragmentation. //! //! The [DownloaderActor] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. @@ -54,7 +54,7 @@ struct DiscoveryId(u64); struct PeerDownloadId(u64); #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] -struct BitmapSubscriptionId(u64); +struct BitfieldSubscriptionId(u64); /// Announce kind #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] @@ -81,20 +81,23 @@ trait ContentDiscovery: Send + 'static { fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Send + Unpin + 'static; } -trait BitmapSubscription: Send + 'static { +trait BitfieldSubscription: Send + 'static { /// Subscribe to a bitmap - fn subscribe(&mut self, peer: BitmapPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; + fn subscribe(&mut self, peer: BitmapPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; } -enum BitmapSubscriptionEvent { - Bitmap { ranges: ChunkRanges }, - BitmapUpdate { added: ChunkRanges, removed: ChunkRanges}, +enum BitfieldSubscriptionEvent { + Bitfield { ranges: ChunkRanges }, + BitfieldUpdate { added: ChunkRanges, removed: ChunkRanges}, } -struct TestBitmapSubscription; +/// A bitmap subscription that just returns nothing for local and everything(*) for remote +/// +/// * Still need to figure out how to deal with open ended chunk ranges. +struct TestBitfieldSubscription; -impl BitmapSubscription for TestBitmapSubscription { - fn subscribe(&mut self, peer: BitmapPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { +impl BitfieldSubscription for TestBitfieldSubscription { + fn subscribe(&mut self, peer: BitmapPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { let bitmap = match peer { BitmapPeer::Local => { ChunkRanges::empty() @@ -103,7 +106,7 @@ impl BitmapSubscription for TestBitmapSubscription { ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } }; - futures_lite::stream::once(BitmapSubscriptionEvent::Bitmap { ranges: bitmap }).chain(futures_lite::stream::pending()) + futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges: bitmap }).chain(futures_lite::stream::pending()) } } @@ -119,7 +122,7 @@ struct PeerState { /// Information about one blob on one peer struct PeerBlobState { /// The subscription id for the subscription - subscription_id: BitmapSubscriptionId, + subscription_id: BitfieldSubscriptionId, /// The number of subscriptions this peer has subscription_count: usize, /// chunk ranges this peer reports to have @@ -127,7 +130,7 @@ struct PeerBlobState { } impl PeerBlobState { - fn new(subscription_id: BitmapSubscriptionId) -> Self { + fn new(subscription_id: BitfieldSubscriptionId) -> Self { Self { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() } } } @@ -500,7 +503,7 @@ struct DownloaderState { // We could use uuid here, but using integers simplifies testing. // // the next subscription id - subscription_id_gen: IdGenerator, + subscription_id_gen: IdGenerator, // the next discovery id discovery_id_gen: IdGenerator, // the next peer download id @@ -542,7 +545,7 @@ enum Command { /// A user request to abort a download. StopDownload { id: DownloadId }, /// A full bitmap for a blob and a peer - Bitmap { + Bitfield { /// The peer that sent the bitmap. peer: BitmapPeer, /// The blob for which the bitmap is @@ -554,7 +557,7 @@ enum Command { /// /// This is used both to update the bitmap of remote peers, and to update /// the local bitmap. - BitmapUpdate { + BitfieldUpdate { /// The peer that sent the update. peer: BitmapPeer, /// The blob that was updated. @@ -592,15 +595,15 @@ enum Command { #[derive(Debug, PartialEq, Eq)] enum Event { - SubscribeBitmap { + SubscribeBitfield { peer: BitmapPeer, hash: Hash, /// The unique id of the subscription - id: BitmapSubscriptionId, + id: BitfieldSubscriptionId, }, - UnsubscribeBitmap { + UnsubscribeBitfield { /// The unique id of the subscription - id: BitmapSubscriptionId, + id: BitfieldSubscriptionId, }, StartDiscovery { hash: Hash, @@ -666,7 +669,7 @@ impl DownloaderState { if *hash == removed_hash { state.subscription_count -= 1; if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitfield { id: state.subscription_id }); return false; } } @@ -696,7 +699,7 @@ impl DownloaderState { } else { // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Local, hash, id: subscription_id }); + evs.push(Event::SubscribeBitfield { peer: BitmapPeer::Local, hash, id: subscription_id }); self.bitmaps.insert((BitmapPeer::Local, hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { @@ -735,14 +738,14 @@ impl DownloaderState { let _state = self.peers.entry(peer).or_default(); // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitmap { peer: BitmapPeer::Remote(peer), hash, id: subscription_id }); + evs.push(Event::SubscribeBitfield { peer: BitmapPeer::Remote(peer), hash, id: subscription_id }); self.bitmaps.insert((BitmapPeer::Remote(peer), hash), PeerBlobState::new(subscription_id)); } Command::DropPeer { peer } => { self.bitmaps.retain(|(p, _), state| { if *p == BitmapPeer::Remote(peer) { // todo: should we emit unsubscribe evs here? - evs.push(Event::UnsubscribeBitmap { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitfield { id: state.subscription_id }); return false; } else { return true; @@ -750,7 +753,7 @@ impl DownloaderState { }); self.peers.remove(&peer); } - Command::Bitmap { peer, hash, bitmap } => { + Command::Bitfield { peer, hash, bitmap } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; let _chunks = total_chunks(&bitmap).context("open range")?; if peer == BitmapPeer::Local { @@ -768,7 +771,7 @@ impl DownloaderState { // we have to call start_downloads even if the local bitmap set, since we don't know in which order local and remote bitmaps arrive self.start_downloads(hash, None, evs)?; } - Command::BitmapUpdate { peer, hash, added, removed } => { + Command::BitfieldUpdate { peer, hash, added, removed } => { let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; if peer == BitmapPeer::Local { state.ranges |= added; @@ -1002,7 +1005,7 @@ impl Downloader { Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { + fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitmap, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); @@ -1028,14 +1031,14 @@ struct DownloaderActor { download_futs: BTreeMap>, peer_download_tasks: BTreeMap>, discovery_tasks: BTreeMap>, - bitmap_subscription_tasks: BTreeMap>, + bitmap_subscription_tasks: BTreeMap>, /// Id generator for download ids download_id_gen: IdGenerator, /// The time when the actor was started, serves as the epoch for time messages to the state machine start: Instant, } -impl DownloaderActor { +impl DownloaderActor { fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { let (send, recv) = mpsc::channel(256); Self { @@ -1090,17 +1093,17 @@ impl DownloaderActor { + Event::SubscribeBitfield { peer, hash, id } => { let send = self.command_tx.clone(); let mut stream = self.subscribe_bitmap.subscribe(peer, hash); let task = spawn(async move { while let Some(ev) = stream.next().await { let cmd = match ev { - BitmapSubscriptionEvent::Bitmap { ranges } => { - Command::Bitmap { peer, hash, bitmap: ranges } + BitfieldSubscriptionEvent::Bitfield { ranges } => { + Command::Bitfield { peer, hash, bitmap: ranges } } - BitmapSubscriptionEvent::BitmapUpdate { added, removed } => { - Command::BitmapUpdate { peer, hash, added, removed } + BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { + Command::BitfieldUpdate { peer, hash, added, removed } } }; send.send(cmd).await.ok(); @@ -1141,7 +1144,7 @@ impl DownloaderActor { + Event::UnsubscribeBitfield { id } => { self.bitmap_subscription_tasks.remove(&id); } Event::StopDiscovery { id } => { @@ -1219,7 +1222,7 @@ async fn peer_download(endpoint: Endpoint, store: S, hash: Hash, peer: sender.send(Command::ChunksDownloaded { time: start.elapsed(), peer, hash, added: added.clone() }).await.ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; - sender.send(Command::BitmapUpdate { peer: BitmapPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); + sender.send(Command::BitfieldUpdate { peer: BitmapPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); } } content = next; @@ -1523,9 +1526,9 @@ mod tests { let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: ChunkRanges::all() }); + let evs = state.apply(Command::Bitfield { peer: Local, hash, bitmap: ChunkRanges::all() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); - let evs = state.apply(Command::Bitmap { peer: Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); + let evs = state.apply(Command::Bitfield { peer: Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); @@ -1542,29 +1545,29 @@ mod tests { let mut state = DownloaderState::new(noop_planner()); let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Local, hash, id: BitmapSubscriptionId(0) }), "starting a download should subscribe to the local bitmap"); + assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "starting a download should subscribe to the local bitmap"); let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply(Command::Bitmap { peer: Local, hash, bitmap: initial_bitmap.clone() }); + let evs = state.apply(Command::Bitfield { peer: Local, hash, bitmap: initial_bitmap.clone() }); assert!(evs.is_empty()); assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); - let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(has_one_event(&evs, &Event::SubscribeBitmap { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); - let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); + assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); + let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); assert!(has_one_event(&evs, &Event::StartPeerDownload { id: PeerDownloadId(0), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); // ChunksDownloaded just updates the peer stats let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); assert!(evs.is_empty()); // Bitmap update does not yet complete the download - let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); assert!(evs.is_empty()); // Final bitmap update for the local bitmap should complete the download - let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); Ok(()) } @@ -1580,30 +1583,30 @@ mod tests { // Start a download state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); // Initially, we have nothing - state.apply(Command::Bitmap { peer: Local, hash, bitmap: ChunkRanges::empty() }); + state.apply(Command::Bitfield { peer: Local, hash, bitmap: ChunkRanges::empty() }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitmap from the peer - let evs = state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..32]) }); + let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..32]) }); assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitmap from a peer should start a download"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); // Bitmap update does not yet complete the download - state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([0..16]), removed: ChunkRanges::empty() }); + state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([0..16]), removed: ChunkRanges::empty() }); // The peer now has more data - state.apply(Command::Bitmap { peer: Remote(peer_a), hash, bitmap: chunk_ranges([32..64]) }); + state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([32..64]) }); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([16..32]) }); // Complete the first part of the download - let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); // This triggers cancellation of the first peer download and starting a new one for the remaining data assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), "first peer download should be stopped"); assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1.into(), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); // Final bitmap update for the local bitmap should complete the download - let evs = state.apply(Command::BitmapUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); - assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1.into() }, &Event::DownloadComplete { id: 0.into() }, &Event::UnsubscribeBitmap { id: 0.into() }, &Event::StopDiscovery { id: 0.into() },]), "download should be completed by the data"); + let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); + assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1.into() }, &Event::DownloadComplete { id: 0.into() }, &Event::UnsubscribeBitfield { id: 0.into() }, &Event::StopDiscovery { id: 0.into() },]), "download should be completed by the data"); println!("{evs:?}"); Ok(()) } @@ -1616,7 +1619,7 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; - let bitmap_subscription = TestBitmapSubscription; + let bitmap_subscription = TestBitfieldSubscription; let local_pool = LocalPool::single(); let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, noop_planner()); tokio::time::sleep(Duration::from_secs(2)).await; @@ -1642,7 +1645,7 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; - let bitmap_subscription = TestBitmapSubscription; + let bitmap_subscription = TestBitfieldSubscription; let local_pool = LocalPool::single(); let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); tokio::time::sleep(Duration::from_secs(2)).await; From b39d1f9a4c86105c8ee21d75baf312e724d8e43b Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 4 Feb 2025 12:55:35 +0200 Subject: [PATCH 19/47] rename bitmap to bitfield bitmap evokes the thought of 2d bitmaps like images --- src/downloader2.rs | 245 ++++++++++++++++++++++----------------------- 1 file changed, 118 insertions(+), 127 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 2630f1ef8..8227bd19f 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -82,8 +82,8 @@ trait ContentDiscovery: Send + 'static { } trait BitfieldSubscription: Send + 'static { - /// Subscribe to a bitmap - fn subscribe(&mut self, peer: BitmapPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; + /// Subscribe to a bitfield + fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; } enum BitfieldSubscriptionEvent { @@ -91,22 +91,22 @@ enum BitfieldSubscriptionEvent { BitfieldUpdate { added: ChunkRanges, removed: ChunkRanges}, } -/// A bitmap subscription that just returns nothing for local and everything(*) for remote +/// A bitfield subscription that just returns nothing for local and everything(*) for remote /// /// * Still need to figure out how to deal with open ended chunk ranges. struct TestBitfieldSubscription; impl BitfieldSubscription for TestBitfieldSubscription { - fn subscribe(&mut self, peer: BitmapPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { - let bitmap = match peer { - BitmapPeer::Local => { + fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { + let ranges = match peer { + BitfieldPeer::Local => { ChunkRanges::empty() } - BitmapPeer::Remote(_) => { + BitfieldPeer::Remote(_) => { ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) } }; - futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges: bitmap }).chain(futures_lite::stream::pending()) + futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }).chain(futures_lite::stream::pending()) } } @@ -229,45 +229,45 @@ impl Downloads { } #[derive(Default)] -struct Bitmaps { - by_peer_and_hash: BTreeMap<(BitmapPeer, Hash), PeerBlobState>, +struct Bitfields { + by_peer_and_hash: BTreeMap<(BitfieldPeer, Hash), PeerBlobState>, } -impl Bitmaps { +impl Bitfields { fn retain(&mut self, mut f: F) where - F: FnMut(&(BitmapPeer, Hash), &mut PeerBlobState) -> bool, + F: FnMut(&(BitfieldPeer, Hash), &mut PeerBlobState) -> bool, { self.by_peer_and_hash.retain(|k, v| f(k, v)); } - fn get(&self, key: &(BitmapPeer, Hash)) -> Option<&PeerBlobState> { + fn get(&self, key: &(BitfieldPeer, Hash)) -> Option<&PeerBlobState> { self.by_peer_and_hash.get(key) } fn get_local(&self, hash: Hash) -> Option<&PeerBlobState> { - self.by_peer_and_hash.get(&(BitmapPeer::Local, hash)) + self.by_peer_and_hash.get(&(BitfieldPeer::Local, hash)) } - fn get_mut(&mut self, key: &(BitmapPeer, Hash)) -> Option<&mut PeerBlobState> { + fn get_mut(&mut self, key: &(BitfieldPeer, Hash)) -> Option<&mut PeerBlobState> { self.by_peer_and_hash.get_mut(key) } fn get_local_mut(&mut self, hash: Hash) -> Option<&mut PeerBlobState> { - self.by_peer_and_hash.get_mut(&(BitmapPeer::Local, hash)) + self.by_peer_and_hash.get_mut(&(BitfieldPeer::Local, hash)) } - fn insert(&mut self, key: (BitmapPeer, Hash), value: PeerBlobState) { + fn insert(&mut self, key: (BitfieldPeer, Hash), value: PeerBlobState) { self.by_peer_and_hash.insert(key, value); } - fn contains_key(&self, key: &(BitmapPeer, Hash)) -> bool { + fn contains_key(&self, key: &(BitfieldPeer, Hash)) -> bool { self.by_peer_and_hash.contains_key(key) } fn remote_for_hash(&self, hash: Hash) -> impl Iterator { self.by_peer_and_hash.iter().filter_map(move |((peer, h), state)| { - if let BitmapPeer::Remote(peer) = peer { + if let BitfieldPeer::Remote(peer) = peer { if *h == hash { Some((peer, state)) } else { @@ -486,10 +486,10 @@ struct PeerDownloadState { struct DownloaderState { // all peers I am tracking for any download peers: BTreeMap, - // all bitmaps I am tracking, both for myself and for remote peers + // all bitfields I am tracking, both for myself and for remote peers // // each item here corresponds to an active subscription - bitmaps: Bitmaps, + bitfields: Bitfields, // all active downloads // // these are user downloads. each user download gets split into one or more @@ -517,7 +517,7 @@ impl DownloaderState { Self { peers: BTreeMap::new(), downloads: Default::default(), - bitmaps: Default::default(), + bitfields: Default::default(), discovery: BTreeMap::new(), subscription_id_gen: Default::default(), discovery_id_gen: Default::default(), @@ -528,7 +528,7 @@ impl DownloaderState { } #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] -enum BitmapPeer { +enum BitfieldPeer { Local, Remote(NodeId), } @@ -544,22 +544,22 @@ enum Command { }, /// A user request to abort a download. StopDownload { id: DownloadId }, - /// A full bitmap for a blob and a peer + /// A full bitfield for a blob and a peer Bitfield { - /// The peer that sent the bitmap. - peer: BitmapPeer, - /// The blob for which the bitmap is + /// The peer that sent the bitfield. + peer: BitfieldPeer, + /// The blob for which the bitfield is hash: Hash, - /// The complete bitmap - bitmap: ChunkRanges, + /// The complete bitfield + ranges: ChunkRanges, }, - /// An update of a bitmap for a hash + /// An update of a bitfield for a hash /// - /// This is used both to update the bitmap of remote peers, and to update - /// the local bitmap. + /// This is used both to update the bitfield of remote peers, and to update + /// the local bitfield. BitfieldUpdate { /// The peer that sent the update. - peer: BitmapPeer, + peer: BitfieldPeer, /// The blob that was updated. hash: Hash, /// The ranges that were added @@ -596,7 +596,7 @@ enum Command { #[derive(Debug, PartialEq, Eq)] enum Event { SubscribeBitfield { - peer: BitmapPeer, + peer: BitfieldPeer, hash: Hash, /// The unique id of the subscription id: BitfieldSubscriptionId, @@ -655,7 +655,7 @@ impl DownloaderState { /// /// Cleanup involves emitting events for /// - stopping all peer downloads - /// - unsubscribing from bitmaps if needed + /// - unsubscribing from bitfields if needed /// - stopping the discovery task if needed fn stop_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { let removed = self.downloads.remove(&id).context(format!("removed unknown download {id:?}"))?; @@ -664,8 +664,8 @@ impl DownloaderState { for peer_download in removed.peer_downloads.values() { evs.push(Event::StopPeerDownload { id: peer_download.id }); } - // unsubscribe from bitmaps that have no more subscriptions - self.bitmaps.retain(|(_peer, hash), state| { + // unsubscribe from bitfields that have no more subscriptions + self.bitfields.retain(|(_peer, hash), state| { if *hash == removed_hash { state.subscription_count -= 1; if state.subscription_count == 0 { @@ -693,14 +693,14 @@ impl DownloaderState { anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id:?}"); let hash = request.hash; // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitmaps.get_local_mut(hash) { + if let Some(state) = self.bitfields.get_local_mut(hash) { // just increment the count state.subscription_count += 1; } else { // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { peer: BitmapPeer::Local, hash, id: subscription_id }); - self.bitmaps.insert((BitmapPeer::Local, hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitfield { peer: BitfieldPeer::Local, hash, id: subscription_id }); + self.bitfields.insert((BitfieldPeer::Local, hash), PeerBlobState::new(subscription_id)); } if !self.discovery.contains_key(&request.hash) { // start a discovery task @@ -726,7 +726,7 @@ impl DownloaderState { self.stop_download(id, evs)?; } Command::PeerDiscovered { peer, hash } => { - if self.bitmaps.contains_key(&(BitmapPeer::Remote(peer), hash)) { + if self.bitfields.contains_key(&(BitfieldPeer::Remote(peer), hash)) { // we already have a subscription for this peer return Ok(()); }; @@ -738,12 +738,12 @@ impl DownloaderState { let _state = self.peers.entry(peer).or_default(); // create a new subscription let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { peer: BitmapPeer::Remote(peer), hash, id: subscription_id }); - self.bitmaps.insert((BitmapPeer::Remote(peer), hash), PeerBlobState::new(subscription_id)); + evs.push(Event::SubscribeBitfield { peer: BitfieldPeer::Remote(peer), hash, id: subscription_id }); + self.bitfields.insert((BitfieldPeer::Remote(peer), hash), PeerBlobState::new(subscription_id)); } Command::DropPeer { peer } => { - self.bitmaps.retain(|(p, _), state| { - if *p == BitmapPeer::Remote(peer) { + self.bitfields.retain(|(p, _), state| { + if *p == BitfieldPeer::Remote(peer) { // todo: should we emit unsubscribe evs here? evs.push(Event::UnsubscribeBitfield { id: state.subscription_id }); return false; @@ -753,27 +753,27 @@ impl DownloaderState { }); self.peers.remove(&peer); } - Command::Bitfield { peer, hash, bitmap } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap for unknown peer {peer:?} and hash {hash}"))?; - let _chunks = total_chunks(&bitmap).context("open range")?; - if peer == BitmapPeer::Local { - state.ranges = bitmap; + Command::Bitfield { peer, hash, ranges } => { + let state = self.bitfields.get_mut(&(peer, hash)).context(format!("bitfields for unknown peer {peer:?} and hash {hash}"))?; + let _chunks = total_chunks(&ranges).context("open range")?; + if peer == BitfieldPeer::Local { + state.ranges = ranges; self.check_completion(hash, None, evs)?; } else { // We got an entirely new peer, mark all affected downloads for rebalancing for download in self.downloads.values_mut_for_hash(hash) { - if bitmap.intersects(&download.request.ranges) { + if ranges.intersects(&download.request.ranges) { download.needs_rebalancing = true; } } - state.ranges = bitmap; + state.ranges = ranges; } - // we have to call start_downloads even if the local bitmap set, since we don't know in which order local and remote bitmaps arrive + // 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::BitfieldUpdate { peer, hash, added, removed } => { - let state = self.bitmaps.get_mut(&(peer, hash)).context(format!("bitmap update for unknown peer {peer:?} and hash {hash}"))?; - if peer == BitmapPeer::Local { + let state = self.bitfields.get_mut(&(peer, hash)).context(format!("bitfield update for unknown peer {peer:?} and hash {hash}"))?; + if peer == BitfieldPeer::Local { state.ranges |= added; state.ranges &= !removed; self.check_completion(hash, None, evs)?; @@ -787,12 +787,12 @@ impl DownloaderState { } state.ranges |= added; state.ranges &= !removed; - // a local bitmap update does not make more data available, so we don't need to start downloads + // a local bitfield update does not make more data available, so we don't need to start downloads self.start_downloads(hash, None, evs)?; } } Command::ChunksDownloaded { time, peer, hash, added } => { - let state = self.bitmaps.get_local_mut(hash).context(format!("chunks downloaded before having local bitmap for {hash}"))?; + let state = self.bitfields.get_local_mut(hash).context(format!("chunks downloaded before having local bitfield for {hash}"))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; @@ -824,7 +824,7 @@ impl DownloaderState { // nothing has changed that affects this download continue; } - let n_peers = self.bitmaps.remote_for_hash(download.request.hash).count(); + let n_peers = self.bitfields.remote_for_hash(download.request.hash).count(); if download.peer_downloads.len() >= n_peers { // we are already downloading from all peers for this hash continue; @@ -841,11 +841,11 @@ impl DownloaderState { /// Check for completion of a download or of an individual peer download /// - /// This must be called after each change of the local bitmap for a hash + /// This must be called after each change of the local bitfield for a hash /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. fn check_completion(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get_local(hash) else { + let Some(self_state) = self.bitfields.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; @@ -889,8 +889,8 @@ impl DownloaderState { // see what the new peers can do for us let mut candidates = BTreeMap::new(); for peer in available { - let Some(peer_state) = self.bitmaps.get(&(BitmapPeer::Remote(peer), hash)) else { - // weird. we should have a bitmap for this peer since it just completed a download + let Some(peer_state) = self.bitfields.get(&(BitfieldPeer::Remote(peer), hash)) else { + // weird. we should have a bitfield for this peer since it just completed a download continue; }; let intersection = &peer_state.ranges & &remaining; @@ -917,7 +917,7 @@ impl DownloaderState { /// Look at all downloads for a hash and start peer downloads for those that do not have any yet fn start_downloads(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { - let Some(self_state) = self.bitmaps.get_local(hash) else { + let Some(self_state) = self.bitfields.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; @@ -927,8 +927,8 @@ impl DownloaderState { } let remaining = &download.request.ranges - &self_state.ranges; let mut candidates = BTreeMap::new(); - for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { - let intersection = &bitmap.ranges & &remaining; + for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { + let intersection = &bitfield.ranges & &remaining; if !intersection.is_empty() { candidates.insert(*peer, intersection); } @@ -950,14 +950,14 @@ impl DownloaderState { download.needs_rebalancing = false; tracing::error!("Rebalancing download {id:?} {:?}", download.request); let hash = download.request.hash; - let Some(self_state) = self.bitmaps.get_local(hash) else { + let Some(self_state) = self.bitfields.get_local(hash) else { // 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 remaining = &download.request.ranges - &self_state.ranges; let mut candidates = BTreeMap::new(); - for (peer, bitmap) in self.bitmaps.remote_for_hash(hash) { - let intersection = &bitmap.ranges & &remaining; + for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { + let intersection = &bitfield.ranges & &remaining; if !intersection.is_empty() { candidates.insert(*peer, intersection); } @@ -1005,8 +1005,8 @@ impl Downloader { Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { - let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitmap, local_pool, planner); + fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { + let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitfield, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); Self { send, task } @@ -1027,11 +1027,11 @@ struct DownloaderActor { state: DownloaderState, store: S, discovery: D, - subscribe_bitmap: B, + subscribe_bitfield: B, download_futs: BTreeMap>, peer_download_tasks: BTreeMap>, discovery_tasks: BTreeMap>, - bitmap_subscription_tasks: BTreeMap>, + bitfield_subscription_tasks: BTreeMap>, /// Id generator for download ids download_id_gen: IdGenerator, /// The time when the actor was started, serves as the epoch for time messages to the state machine @@ -1039,7 +1039,7 @@ struct DownloaderActor { } impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitmap: B, local_pool: LocalPool, planner: Box) -> Self { + fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { let (send, recv) = mpsc::channel(256); Self { local_pool, @@ -1047,10 +1047,10 @@ impl DownloaderActor DownloaderActor { let send = self.command_tx.clone(); - let mut stream = self.subscribe_bitmap.subscribe(peer, hash); + 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 { BitfieldSubscriptionEvent::Bitfield { ranges } => { - Command::Bitfield { peer, hash, bitmap: ranges } + Command::Bitfield { peer, hash, ranges } } BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { Command::BitfieldUpdate { peer, hash, added, removed } @@ -1108,17 +1108,8 @@ impl DownloaderActor { let send = self.command_tx.clone(); @@ -1145,7 +1136,7 @@ impl DownloaderActor { - self.bitmap_subscription_tasks.remove(&id); + self.bitfield_subscription_tasks.remove(&id); } Event::StopDiscovery { id } => { self.discovery_tasks.remove(&id); @@ -1222,7 +1213,7 @@ async fn peer_download(endpoint: Endpoint, store: S, hash: Hash, peer: sender.send(Command::ChunksDownloaded { time: start.elapsed(), peer, hash, added: added.clone() }).await.ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; - sender.send(Command::BitfieldUpdate { peer: BitmapPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); + sender.send(Command::BitfieldUpdate { peer: BitfieldPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); } } content = next; @@ -1252,7 +1243,7 @@ where AbortOnDropHandle::new(task) } -fn print_bitmap(iter: impl IntoIterator) -> String { +fn print_bitfield(iter: impl IntoIterator) -> String { let mut chars = String::new(); for x in iter { chars.push(if x { '█' } else { ' ' }); @@ -1260,7 +1251,7 @@ fn print_bitmap(iter: impl IntoIterator) -> String { chars } -fn print_bitmap_compact(iter: impl IntoIterator) -> String { +fn print_bitfield_compact(iter: impl IntoIterator) -> String { let mut chars = String::new(); let mut iter = iter.into_iter(); @@ -1344,7 +1335,7 @@ fn create_ranges(indexes: impl IntoIterator>) -> Vec String { +fn print_colored_bitfield(data: &[u8], colors: &[u8]) -> String { let mut chars = String::new(); let mut iter = data.iter(); @@ -1380,7 +1371,7 @@ mod tests { #[test] fn print_chunk_range() { let x = chunk_ranges([0..3, 4..30, 40..50]); - let s = print_bitmap_compact(as_bool_iter(&x, 50)); + let s = print_bitfield_compact(as_bool_iter(&x, 50)); println!("{}", s); } @@ -1436,7 +1427,7 @@ mod tests { fn print_range_map(ranges: &BTreeMap) { for (peer, ranges) in ranges { - let x = print_bitmap(as_bool_iter(ranges, 100)); + let x = print_bitfield(as_bool_iter(ranges, 100)); println!("{peer}: {x}"); } } @@ -1447,7 +1438,7 @@ mod tests { let iter = select_ranges(ranges.as_slice(), 8); for (i, range) in ranges.iter().enumerate() { let bools = as_bool_iter(&range, 100).collect::>(); - println!("{i:4}{}", print_bitmap(bools)); + println!("{i:4}{}", print_bitfield(bools)); } print!(" "); for x in &iter { @@ -1457,7 +1448,7 @@ mod tests { let ranges = create_ranges(iter); for (i, range) in ranges.iter().enumerate() { let bools = as_bool_iter(&range, 100).collect::>(); - println!("{i:4}{}", print_bitmap(bools)); + println!("{i:4}{}", print_bitfield(bools)); } } @@ -1469,10 +1460,10 @@ mod tests { } #[test] - fn test_print_colored_bitmap() { - let bitmap = vec![0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]; // Notice: odd-length input + fn test_print_colored_bitfield() { + let bitfield = vec![0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]; // Notice: odd-length input - println!("{}", print_colored_bitmap(&bitmap, &[1, 2, 3, 4])); + println!("{}", print_colored_bitfield(&bitfield, &[1, 2, 3, 4])); } #[cfg(feature = "rpc")] @@ -1520,16 +1511,16 @@ mod tests { /// Test various things that should produce errors #[test] fn downloader_state_errors() -> TestResult<()> { - use BitmapPeer::*; + use BitfieldPeer::*; let _ = tracing_subscriber::fmt::try_init(); let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::Bitfield { peer: Local, hash, bitmap: ChunkRanges::all() }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap should produce an error!"); - let evs = state.apply(Command::Bitfield { peer: Local, hash: unknown_hash, bitmap: ChunkRanges::all() }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitmap for an unknown hash should produce an error!"); + let evs = state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::all() }); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitfield should produce an error!"); + let evs = state.apply(Command::Bitfield { peer: Local, hash: unknown_hash, ranges: ChunkRanges::all() }); + assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitfield for an unknown hash should produce an error!"); let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); Ok(()) @@ -1538,35 +1529,35 @@ mod tests { /// Test a simple scenario where a download is started and completed #[test] fn downloader_state_smoke() -> TestResult<()> { - use BitmapPeer::*; + use BitfieldPeer::*; let _ = tracing_subscriber::fmt::try_init(); let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "starting a download should subscribe to the local bitmap"); - let initial_bitmap = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply(Command::Bitfield { peer: Local, hash, bitmap: initial_bitmap.clone() }); + assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "starting a download should subscribe to the local bitfield"); + let initial_bitfield = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); + let evs = state.apply(Command::Bitfield { peer: Local, hash, ranges: initial_bitfield.clone() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, initial_bitmap, "bitmap should be set to the initial bitmap"); + assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, initial_bitfield, "bitfield should be set to the initial bitfield"); let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); - assert_eq!(state.bitmaps.get_local(hash).context("bitmap should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitmap should be updated"); + assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitfield should be updated"); let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitmap"); - let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..64]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: PeerDownloadId(0), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitmap from a peer should start a download"); + assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitfield"); + let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..64]) }); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: PeerDownloadId(0), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitfield from a peer should start a download"); // ChunksDownloaded just updates the peer stats let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); assert!(evs.is_empty()); - // Bitmap update does not yet complete the download + // Bitfield update does not yet complete the download let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); assert!(evs.is_empty()); - // Final bitmap update for the local bitmap should complete the download + // Final bitfield update for the local bitfield should complete the download let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); Ok(()) @@ -1575,7 +1566,7 @@ mod tests { /// Test a scenario where more data becomes available at the remote peer as the download progresses #[test] fn downloader_state_incremental() -> TestResult<()> { - use BitmapPeer::*; + use BitfieldPeer::*; let _ = tracing_subscriber::fmt::try_init(); let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; @@ -1583,18 +1574,18 @@ mod tests { // Start a download state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); // Initially, we have nothing - state.apply(Command::Bitfield { peer: Local, hash, bitmap: ChunkRanges::empty() }); + state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::empty() }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - // We have a bitmap from the peer - let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([0..32]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitmap from a peer should start a download"); + // We have a bitfield from the peer + let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..32]) }); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitfield from a peer should start a download"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); - // Bitmap update does not yet complete the download + // Bitfield update does not yet complete the download state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([0..16]), removed: ChunkRanges::empty() }); // The peer now has more data - state.apply(Command::Bitfield { peer: Remote(peer_a), hash, bitmap: chunk_ranges([32..64]) }); + state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([32..64]) }); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([16..32]) }); // Complete the first part of the download @@ -1604,7 +1595,7 @@ mod tests { assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1.into(), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); - // Final bitmap update for the local bitmap should complete the download + // Final bitfield update for the local bitfield should complete the download let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1.into() }, &Event::DownloadComplete { id: 0.into() }, &Event::UnsubscribeBitfield { id: 0.into() }, &Event::StopDiscovery { id: 0.into() },]), "download should be completed by the data"); println!("{evs:?}"); @@ -1619,9 +1610,9 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; - let bitmap_subscription = TestBitfieldSubscription; + let bitfield_subscription = TestBitfieldSubscription; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, noop_planner()); + let downloader = Downloader::new(endpoint, store, discovery, bitfield_subscription, local_pool, noop_planner()); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); fut.await?; @@ -1645,9 +1636,9 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; - let bitmap_subscription = TestBitfieldSubscription; + let bitfield_subscription = TestBitfieldSubscription; let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, bitmap_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); + let downloader = Downloader::new(endpoint, store, discovery, bitfield_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; From 94eaf5feb8cc267c922dc7e68ab669ed0b1c3ecc Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 4 Feb 2025 14:35:18 +0200 Subject: [PATCH 20/47] Properly handle dropping downloads --- src/downloader2.rs | 41 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 8227bd19f..c462fc3ec 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -40,7 +40,7 @@ use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use std::time::Duration; -use tokio::sync::mpsc; +use tokio::{sync::mpsc, task::JoinSet}; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; @@ -1062,6 +1062,7 @@ impl DownloaderActor) { let mut ticks = tokio::time::interval(Duration::from_millis(100)); loop { + trace!("downloader actor tick"); tokio::select! { biased; Some(cmd) = channel.recv() => { @@ -1085,6 +1086,15 @@ impl DownloaderActor { let time = self.start.elapsed(); self.command_tx.send(Command::Tick { time }).await.ok(); + // clean up dropped futures + // + // todo: is there a better mechanism than periodic checks? + // I don't want some cancellation token rube goldberg machine. + for (id, fut) in self.download_futs.iter() { + if fut.is_closed() { + self.command_tx.send(Command::StopDownload { id: *id }).await.ok(); + } + } }, } } @@ -1602,6 +1612,33 @@ mod tests { Ok(()) } + /// Test a scenario where more data becomes available at the remote peer as the download progresses + #[test] + fn downloader_state_drop() -> TestResult<()> { + use BitfieldPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + // Start a download + state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0.into() }); + // Initially, we have nothing + state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::empty() }); + // We have a peer for the hash + state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + // We have a bitfield from the peer + let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..32]) }); + assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitfield from a peer should start a download"); + // Sending StopDownload should stop the download and all associated tasks + // This is what happens (delayed) when the user drops the download future + let evs = state.apply(Command::StopDownload { id: 0.into() }); + assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() })); + assert!(has_one_event(&evs, &Event::UnsubscribeBitfield { id: 0.into() })); + assert!(has_one_event(&evs, &Event::UnsubscribeBitfield { id: 1.into() })); + assert!(has_one_event(&evs, &Event::StopDiscovery { id: 0.into() })); + Ok(()) + } + #[tokio::test] #[cfg(feature = "rpc")] async fn downloader_driver_smoke() -> TestResult<()> { @@ -1639,7 +1676,7 @@ mod tests { let bitfield_subscription = TestBitfieldSubscription; let local_pool = LocalPool::single(); let downloader = Downloader::new(endpoint, store, discovery, bitfield_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(1)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; Ok(()) From 6073d6c104167830efc95475b3b6131541268b04 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 7 Feb 2025 13:19:49 +0200 Subject: [PATCH 21/47] more tests --- src/downloader2.rs | 73 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/src/downloader2.rs b/src/downloader2.rs index c462fc3ec..fcd878dc0 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -1551,6 +1551,7 @@ mod tests { let evs = state.apply(Command::Bitfield { peer: Local, hash, ranges: initial_bitfield.clone() }); assert!(evs.is_empty()); assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, initial_bitfield, "bitfield should be set to the initial bitfield"); + assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.subscription_count, 1, "we have one download interested in the bitfield"); let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); assert!(evs.is_empty()); assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitfield should be updated"); @@ -1570,6 +1571,10 @@ mod tests { // Final bitfield update for the local bitfield should complete the download let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); + // quick check that everything got cleaned up + assert!(state.downloads.by_id.is_empty()); + assert!(state.bitfields.by_peer_and_hash.is_empty()); + assert!(state.discovery.is_empty()); Ok(()) } @@ -1612,6 +1617,74 @@ mod tests { Ok(()) } + #[test] + fn downloader_state_multiple_downloads() -> testresult::TestResult<()> { + use BitfieldPeer::*; + // Use a constant hash (the same style as used in other tests). + let hash = "0000000000000000000000000000000000000000000000000000000000000001" + .parse()?; + // Create a downloader state with a no‐op planner. + let mut state = DownloaderState::new(noop_planner()); + + // --- Start the first (ongoing) download. + // Request a range from 0..64. + let download0 = DownloadId(0); + let req0 = DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }; + let evs0 = state.apply(Command::StartDownload { request: req0, id: download0 }); + // When starting the download, we expect a discovery task to be started + // and a subscription to the local bitfield to be requested. + assert!( + has_one_event(&evs0, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), + "download0 should start discovery" + ); + assert!( + has_one_event(&evs0, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), + "download0 should subscribe to the local bitfield" + ); + + // --- Simulate some progress for the first download. + // Let’s say only chunks 0..32 are available locally. + let evs1 = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: chunk_ranges([0..32]), + }); + // No completion event should be generated for download0 because its full range 0..64 is not yet met. + assert!(evs1.is_empty(), "Partial bitfield update should not complete download0"); + + // --- Start a second download for the same hash. + // This new download only requires chunks 0..32 which are already available. + let download1 = DownloadId(1); + let req1 = DownloadRequest { + hash, + ranges: chunk_ranges([0..32]), + }; + let evs2 = state.apply(Command::StartDownload { request: req1, id: download1 }); + // Because the local bitfield (0..32) is already a superset of the new download’s request, + // a DownloadComplete event for download1 should be generated immediately. + assert!( + has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), + "New download should complete immediately" + ); + + // --- Verify state: + // The ongoing download (download0) should still be present in the state, + // while the newly completed download (download1) is removed. + assert!( + state.downloads.contains_key(&download0), + "download0 should still be active" + ); + assert!( + !state.downloads.contains_key(&download1), + "download1 should have been cleaned up after completion" + ); + + Ok(()) + } + /// Test a scenario where more data becomes available at the remote peer as the download progresses #[test] fn downloader_state_drop() -> TestResult<()> { From 04712e3c3acb5469eb030b940721aa1065a02d75 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 7 Feb 2025 15:11:50 +0200 Subject: [PATCH 22/47] Fix warnings --- src/downloader2.rs | 370 +++++++++++++++------------------------------ 1 file changed, 125 insertions(+), 245 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index fcd878dc0..8d3893a65 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -20,7 +20,12 @@ //! The [DownloaderActor] is the asynchronous driver for the state machine. It //! owns the actual tasks that perform IO. use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, future::Future, io, marker::PhantomData, sync::Arc, time::Instant + collections::{BTreeMap, BTreeSet, VecDeque}, + future::Future, + io, + marker::PhantomData, + sync::Arc, + time::Instant, }; use crate::{ @@ -40,7 +45,7 @@ use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; use std::time::Duration; -use tokio::{sync::mpsc, task::JoinSet}; +use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; @@ -66,14 +71,15 @@ pub enum AnnounceKind { Complete, } -#[derive(Default)] -struct FindPeersOpts { +/// Options for finding peers +#[derive(Debug, Default)] +pub struct FindPeersOpts { /// Kind of announce - kind: AnnounceKind, + pub kind: AnnounceKind, } /// A pluggable content discovery mechanism -trait ContentDiscovery: Send + 'static { +pub trait ContentDiscovery: Send + 'static { /// Find peers that have the given blob. /// /// The returned stream is a handle for the discovery task. It should be an @@ -81,33 +87,27 @@ trait ContentDiscovery: Send + 'static { fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Send + Unpin + 'static; } -trait BitfieldSubscription: Send + 'static { +/// A pluggable bitfield subscription mechanism +pub trait BitfieldSubscription: Send + 'static { /// Subscribe to a bitfield fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; } -enum BitfieldSubscriptionEvent { - Bitfield { ranges: ChunkRanges }, - BitfieldUpdate { added: ChunkRanges, removed: ChunkRanges}, -} - -/// A bitfield subscription that just returns nothing for local and everything(*) for remote -/// -/// * Still need to figure out how to deal with open ended chunk ranges. -struct TestBitfieldSubscription; - -impl BitfieldSubscription for TestBitfieldSubscription { - fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { - let ranges = match peer { - BitfieldPeer::Local => { - ChunkRanges::empty() - } - BitfieldPeer::Remote(_) => { - ChunkRanges::from(ChunkNum(0)..ChunkNum(1024)) - } - }; - futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }).chain(futures_lite::stream::pending()) - } +/// 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, + }, } /// Global information about a peer @@ -135,12 +135,13 @@ impl PeerBlobState { } } +/// A download request #[derive(Debug)] -struct DownloadRequest { +pub struct DownloadRequest { /// The blob we are downloading - hash: Hash, + pub hash: Hash, /// The ranges we are interested in - ranges: ChunkRanges, + pub ranges: ChunkRanges, } struct DownloadState { @@ -224,7 +225,7 @@ impl Downloads { } fn by_peer_download_id_mut(&mut self, id: PeerDownloadId) -> Option<(&DownloadId, &mut DownloadState)> { - self.by_id.iter_mut().filter(|(k, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)).next() + self.by_id.iter_mut().filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)).next() } } @@ -289,17 +290,19 @@ impl Bitfields { /// want to deduplicate the ranges, but they could also do other things, like /// eliminate gaps or even extend ranges. The only thing they should not do is /// to add new peers to the list of options. -trait DownloadPlanner: Send + 'static { +pub trait DownloadPlanner: Send + 'static { /// Make a download plan for a hash, by reducing or eliminating the overlap of chunk ranges fn plan(&mut self, hash: Hash, options: &mut BTreeMap); } -type BoxedDownloadPlanner = Box; +/// A boxed download planner +pub type BoxedDownloadPlanner = Box; /// A download planner that just leaves everything as is. /// /// Data will be downloaded from all peers wherever multiple peers have the same data. -struct NoopPlanner; +#[derive(Debug, Clone, Copy)] +pub struct NoopPlanner; impl DownloadPlanner for NoopPlanner { fn plan(&mut self, _hash: Hash, _options: &mut BTreeMap) {} @@ -310,7 +313,8 @@ impl DownloadPlanner for NoopPlanner { /// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, /// and for each stripe decides on a single peer to download from, based on the /// peer id and a random seed. -struct StripePlanner { +#[derive(Debug)] +pub struct StripePlanner { /// seed for the score function. This can be set to 0 for testing for /// maximum determinism, but can be set to a random value for production /// to avoid multiple downloaders coming up with the same plan. @@ -328,6 +332,7 @@ struct StripePlanner { } impl StripePlanner { + /// Create a new planner with the given seed and stripe size. pub fn new(seed: u64, stripe_size_log: u8) -> Self { Self { seed, stripe_size_log } } @@ -403,7 +408,8 @@ fn get_continuous_ranges(options: &mut BTreeMap, stripe_siz /// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, /// and for each stripe decides on a single peer to download from, based on the /// peer id and a random seed. -struct StripePlanner2 { +#[derive(Debug)] +pub struct StripePlanner2 { /// seed for the score function. This can be set to 0 for testing for /// maximum determinism, but can be set to a random value for production /// to avoid multiple downloaders coming up with the same plan. @@ -421,6 +427,7 @@ struct StripePlanner2 { } impl StripePlanner2 { + /// Create a new planner with the given seed and stripe size. pub fn new(seed: u64, stripe_size_log: u8) -> Self { Self { seed, stripe_size_log } } @@ -527,9 +534,12 @@ impl DownloaderState { } } +/// Peer for a bitfield subscription #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] -enum BitfieldPeer { +pub enum BitfieldPeer { + /// The local bitfield Local, + /// A bitfield from a remote peer Remote(NodeId), } @@ -583,9 +593,11 @@ enum Command { /// A peer download has completed PeerDownloadComplete { id: PeerDownloadId, + #[allow(dead_code)] result: anyhow::Result, }, /// Stop tracking a peer for all blobs, for whatever reason + #[allow(dead_code)] DropPeer { peer: NodeId }, /// A peer has been discovered PeerDiscovered { peer: NodeId, hash: Hash }, @@ -706,7 +718,7 @@ impl DownloaderState { // start a discovery task let id = self.discovery_id_gen.next(); evs.push(Event::StartDiscovery { hash, id }); - self.discovery.insert(request.hash, id ); + self.discovery.insert(request.hash, id); } self.downloads.insert(id, DownloadState::new(request)); self.check_completion(hash, Some(id), evs)?; @@ -948,7 +960,7 @@ impl DownloaderState { fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { let download = self.downloads.by_id_mut(id).context(format!("rebalancing unknown download {id:?}"))?; download.needs_rebalancing = false; - tracing::error!("Rebalancing download {id:?} {:?}", download.request); + tracing::info!("Rebalancing download {id:?} {:?}", download.request); let hash = download.request.hash; let Some(self_state) = self.bitfields.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all @@ -991,25 +1003,30 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { Some(total) } +/// A downloader that allows range downloads and downloads from multiple peers. #[derive(Debug, Clone)] -struct Downloader { +pub struct Downloader { send: mpsc::Sender, - task: Arc>, + _task: Arc>, } impl Downloader { - async fn download(&self, request: DownloadRequest) -> anyhow::Result<()> { + /// Create a new download + /// + /// The download will be cancelled if the returned future is dropped. + pub async fn download(&self, request: DownloadRequest) -> anyhow::Result<()> { let (send, recv) = tokio::sync::oneshot::channel::<()>(); self.send.send(UserCommand::Download { request, done: send }).await?; recv.await?; Ok(()) } - fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { + /// Create a new downloader + pub fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitfield, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); - Self { send, task } + Self { send, _task: task } } } @@ -1080,7 +1097,7 @@ impl DownloaderActor { let evs = self.state.apply(cmd); for ev in evs { - self.handle_event(ev, 0); + self.handle_event(ev); } }, _ = ticks.tick() => { @@ -1100,8 +1117,8 @@ impl DownloaderActor { let send = self.command_tx.clone(); @@ -1109,12 +1126,8 @@ impl DownloaderActor { - Command::Bitfield { peer, hash, ranges } - } - BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { - Command::BitfieldUpdate { peer, hash, added, removed } - } + BitfieldSubscriptionEvent::Bitfield { ranges } => Command::Bitfield { peer, hash, ranges }, + BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => Command::BitfieldUpdate { peer, hash, added, removed }, }; send.send(cmd).await.ok(); } @@ -1162,19 +1175,24 @@ impl DownloaderActor { error!("Error during processing event {}", message); } - _ => { - println!("event: {:?}", ev); - } } } } /// A simple static content discovery mechanism -struct StaticContentDiscovery { +#[derive(Debug)] +pub struct StaticContentDiscovery { info: BTreeMap>, default: Vec, } +impl StaticContentDiscovery { + /// Create a new static content discovery mechanism + pub fn new(info: BTreeMap>, default: Vec) -> Self { + Self { info, default } + } +} + impl ContentDiscovery for StaticContentDiscovery { fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> impl Stream + Unpin + 'static { let peers = self.info.get(&hash).unwrap_or(&self.default).clone(); @@ -1253,120 +1271,6 @@ where AbortOnDropHandle::new(task) } -fn print_bitfield(iter: impl IntoIterator) -> String { - let mut chars = String::new(); - for x in iter { - chars.push(if x { '█' } else { ' ' }); - } - chars -} - -fn print_bitfield_compact(iter: impl IntoIterator) -> String { - let mut chars = String::new(); - let mut iter = iter.into_iter(); - - while let (Some(left), Some(right)) = (iter.next(), iter.next()) { - let c = match (left, right) { - (true, true) => '█', // Both pixels are "on" - (true, false) => '▌', // Left pixel is "on" - (false, true) => '▐', // Right pixel is "on" - (false, false) => ' ', // Both are "off" - }; - chars.push(c); - } - - // If there's an odd pixel at the end, print only a left block. - if let Some(left) = iter.next() { - chars.push(if left { '▌' } else { ' ' }); - } - - chars -} - -fn as_bool_iter(x: &ChunkRanges, max: u64) -> impl Iterator { - let max = x - .iter() - .last() - .map(|x| match x { - RangeSetRange::RangeFrom(_) => max, - RangeSetRange::Range(x) => x.end.0, - }) - .unwrap_or_default(); - let res = (0..max).map(move |i| x.contains(&ChunkNum(i))).collect::>(); - res.into_iter() -} - -/// Given a set of ranges, make them non-overlapping according to some rules. -fn select_ranges(ranges: &[ChunkRanges], continuity_bonus: u64) -> Vec> { - let mut total = vec![0u64; ranges.len()]; - let mut boundaries = BTreeSet::new(); - assert!(ranges.iter().all(|x| x.boundaries().len() % 2 == 0)); - for range in ranges.iter() { - for x in range.boundaries() { - boundaries.insert(x.0); - } - } - let max = boundaries.iter().max().copied().unwrap_or(0); - let mut last_selected = None; - let mut res = vec![]; - for i in 0..max { - let mut lowest_score = u64::MAX; - let mut selected = None; - for j in 0..ranges.len() { - if ranges[j].contains(&ChunkNum(i)) { - let consecutive = last_selected == Some(j); - let score = if consecutive { total[j].saturating_sub(continuity_bonus) } else { total[j] }; - if score < lowest_score { - lowest_score = score; - selected = Some(j); - } - } - } - res.push(selected); - if let Some(selected) = selected { - total[selected] += 1; - } - last_selected = selected; - } - res -} - -fn create_ranges(indexes: impl IntoIterator>) -> Vec { - let mut res = vec![]; - for (i, n) in indexes.into_iter().enumerate() { - let x = i as u64; - if let Some(n) = n { - while res.len() <= n { - res.push(ChunkRanges::empty()); - } - res[n] |= ChunkRanges::from(ChunkNum(x)..ChunkNum(x + 1)); - } - } - res -} - -fn print_colored_bitfield(data: &[u8], colors: &[u8]) -> String { - let mut chars = String::new(); - let mut iter = data.iter(); - - while let Some(&left) = iter.next() { - let right = iter.next(); // Try to fetch the next element - - let left_color = colors.get(left as usize).map(|x| *x).unwrap_or_default(); - let right_char = match right { - Some(&right) => { - let right_color = colors.get(right as usize).map(|x| *x).unwrap_or_default(); - // Use ANSI escape codes to color left and right halves of `█` - format!("\x1b[38;5;{}m\x1b[48;5;{}m▌\x1b[0m", left_color, right_color) - } - None => format!("\x1b[38;5;{}m▌\x1b[0m", left_color), // Handle odd-length case - }; - - chars.push_str(&right_char); - } - chars -} - #[cfg(test)] mod tests { use std::ops::Range; @@ -1378,11 +1282,40 @@ mod tests { use iroh::{protocol::Router, SecretKey}; use testresult::TestResult; - #[test] - fn print_chunk_range() { - let x = chunk_ranges([0..3, 4..30, 40..50]); - let s = print_bitfield_compact(as_bool_iter(&x, 50)); - println!("{}", s); + fn print_bitfield(iter: impl IntoIterator) -> String { + let mut chars = String::new(); + for x in iter { + chars.push(if x { '█' } else { ' ' }); + } + chars + } + + fn as_bool_iter(x: &ChunkRanges, max: u64) -> impl Iterator { + let max = x + .iter() + .last() + .map(|x| match x { + RangeSetRange::RangeFrom(_) => max, + RangeSetRange::Range(x) => x.end.0, + }) + .unwrap_or_default(); + let res = (0..max).map(move |i| x.contains(&ChunkNum(i))).collect::>(); + res.into_iter() + } + + /// A bitfield subscription that just returns nothing for local and everything(*) for remote + /// + /// * Still need to figure out how to deal with open ended chunk ranges. + struct TestBitfieldSubscription; + + impl BitfieldSubscription for TestBitfieldSubscription { + fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { + let ranges = match peer { + BitfieldPeer::Local => ChunkRanges::empty(), + BitfieldPeer::Remote(_) => ChunkRanges::from(ChunkNum(0)..ChunkNum(1024 * 1024 * 1024 * 1024)), + }; + futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }).chain(futures_lite::stream::pending()) + } } fn peer(id: u8) -> NodeId { @@ -1442,26 +1375,6 @@ mod tests { } } - #[test] - fn test_select_ranges() { - let ranges = [chunk_ranges([0..90]), chunk_ranges([0..100]), chunk_ranges([0..80])]; - let iter = select_ranges(ranges.as_slice(), 8); - for (i, range) in ranges.iter().enumerate() { - let bools = as_bool_iter(&range, 100).collect::>(); - println!("{i:4}{}", print_bitfield(bools)); - } - print!(" "); - for x in &iter { - print!("{}", x.map(|x| x.to_string()).unwrap_or(" ".to_string())); - } - println!(); - let ranges = create_ranges(iter); - for (i, range) in ranges.iter().enumerate() { - let bools = as_bool_iter(&range, 100).collect::>(); - println!("{i:4}{}", print_bitfield(bools)); - } - } - #[test] fn test_is_superset() { let local = ChunkRanges::from(ChunkNum(0)..ChunkNum(100)); @@ -1469,13 +1382,6 @@ mod tests { assert!(local.is_superset(&request)); } - #[test] - fn test_print_colored_bitfield() { - let bitfield = vec![0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]; // Notice: odd-length input - - println!("{}", print_colored_bitfield(&bitfield, &[1, 2, 3, 4])); - } - #[cfg(feature = "rpc")] async fn make_test_node(data: &[u8]) -> anyhow::Result<(Router, NodeId, Hash)> { // let noop_subscriber = tracing_subscriber::fmt::Subscriber::builder() @@ -1621,67 +1527,41 @@ mod tests { fn downloader_state_multiple_downloads() -> testresult::TestResult<()> { use BitfieldPeer::*; // Use a constant hash (the same style as used in other tests). - let hash = "0000000000000000000000000000000000000000000000000000000000000001" - .parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; // Create a downloader state with a no‐op planner. let mut state = DownloaderState::new(noop_planner()); - + // --- Start the first (ongoing) download. // Request a range from 0..64. let download0 = DownloadId(0); - let req0 = DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }; + let req0 = DownloadRequest { hash, ranges: chunk_ranges([0..64]) }; let evs0 = state.apply(Command::StartDownload { request: req0, id: download0 }); // When starting the download, we expect a discovery task to be started // and a subscription to the local bitfield to be requested. - assert!( - has_one_event(&evs0, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), - "download0 should start discovery" - ); - assert!( - has_one_event(&evs0, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), - "download0 should subscribe to the local bitfield" - ); - + assert!(has_one_event(&evs0, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "download0 should start discovery"); + assert!(has_one_event(&evs0, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "download0 should subscribe to the local bitfield"); + // --- Simulate some progress for the first download. // Let’s say only chunks 0..32 are available locally. - let evs1 = state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: chunk_ranges([0..32]), - }); + let evs1 = state.apply(Command::Bitfield { peer: Local, hash, ranges: chunk_ranges([0..32]) }); // No completion event should be generated for download0 because its full range 0..64 is not yet met. assert!(evs1.is_empty(), "Partial bitfield update should not complete download0"); - + // --- Start a second download for the same hash. // This new download only requires chunks 0..32 which are already available. let download1 = DownloadId(1); - let req1 = DownloadRequest { - hash, - ranges: chunk_ranges([0..32]), - }; + let req1 = DownloadRequest { hash, ranges: chunk_ranges([0..32]) }; let evs2 = state.apply(Command::StartDownload { request: req1, id: download1 }); // Because the local bitfield (0..32) is already a superset of the new download’s request, // a DownloadComplete event for download1 should be generated immediately. - assert!( - has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), - "New download should complete immediately" - ); - + assert!(has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), "New download should complete immediately"); + // --- Verify state: // The ongoing download (download0) should still be present in the state, // while the newly completed download (download1) is removed. - assert!( - state.downloads.contains_key(&download0), - "download0 should still be active" - ); - assert!( - !state.downloads.contains_key(&download1), - "download1 should have been cleaned up after completion" - ); - + assert!(state.downloads.contains_key(&download0), "download0 should still be active"); + assert!(!state.downloads.contains_key(&download1), "download1 should have been cleaned up after completion"); + Ok(()) } From ff029a6d79ce416c6becb6a6553abf53b0d358c0 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 7 Feb 2025 16:56:38 +0200 Subject: [PATCH 23/47] Add multiprovider example --- Cargo.toml | 4 + examples/multiprovider.rs | 76 +++++++++++++++++ src/downloader2.rs | 166 +++++++++++++++++++++++++------------- 3 files changed, 192 insertions(+), 54 deletions(-) create mode 100644 examples/multiprovider.rs diff --git a/Cargo.toml b/Cargo.toml index 9d5d12298..18a1e6983 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -156,6 +156,10 @@ required-features = ["example-iroh"] name = "custom-protocol" required-features = ["example-iroh"] +[[example]] +name = "multiprovider" +required-features = ["example-iroh"] + [lints.rust] missing_debug_implementations = "warn" diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs new file mode 100644 index 000000000..94f767187 --- /dev/null +++ b/examples/multiprovider.rs @@ -0,0 +1,76 @@ +use std::path::PathBuf; + +use bao_tree::ChunkRanges; +use iroh::NodeId; +use iroh_blobs::{downloader2::{DownloadRequest, Downloader, StaticContentDiscovery}, store::Store, Hash}; +use clap::Parser; + +#[derive(Debug, Parser)] +struct Args { + #[clap(subcommand)] + subcommand: Subcommand, +} + +#[derive(Debug, Parser)] +enum Subcommand { + Download(DownloadArgs), + Provide(ProvideArgs), +} + +#[derive(Debug, Parser)] +struct DownloadArgs { + #[clap(help = "hash to download")] + hash: Hash, + + providers: Vec, +} + +#[derive(Debug, Parser)] +struct ProvideArgs { + #[clap(help = "path to provide")] + path: Vec, +} + +async fn provide(args: ProvideArgs) -> anyhow::Result<()> { + let store = iroh_blobs::store::mem::Store::new(); + let mut tags = Vec::new(); + for path in args.path { + let data = std::fs::read(&path)?; + let tag = store.import_bytes(data.into(), iroh_blobs::BlobFormat::Raw).await?; + println!("added {} as {}", path.display(), tag.hash()); + tags.push((path, tag)); + } + let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; + let id = endpoint.node_id(); + let blobs = iroh_blobs::net_protocol::Blobs::builder(store).build(&endpoint); + let router = iroh::protocol::Router::builder(endpoint) + .accept(iroh_blobs::ALPN, blobs) + .spawn().await?; + println!("listening on {}", id); + tokio::signal::ctrl_c().await?; + router.shutdown().await?; + Ok(()) +} + +async fn download(args: DownloadArgs) -> anyhow::Result<()> { + let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; + let store = iroh_blobs::store::mem::Store::new(); + let discovery = StaticContentDiscovery::new(Default::default(), args.providers); + let downloader = Downloader::builder(endpoint, store).discovery(discovery).build(); + let request = DownloadRequest { + hash: args.hash, + ranges: ChunkRanges::all(), + }; + downloader.download(request).await?; + Ok(()) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let args = Args::parse(); + match args.subcommand { + Subcommand::Download(args) => download(args).await?, + Subcommand::Provide(args) => provide(args).await?, + } + Ok(()) +} \ No newline at end of file diff --git a/src/downloader2.rs b/src/downloader2.rs index 8d3893a65..c1a09fb2e 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -1,24 +1,16 @@ //! Downloader version that supports range downloads and downloads from multiple sources. //! -//! # Structure -//! -//! The [DownloaderState] is a synchronous state machine containing the logic. -//! It gets commands and produces events. It does not do any IO and also does -//! not have any time dependency. So [DownloaderState::apply] is a -//! pure function of the state and the command and can therefore be tested -//! easily. -//! -//! In several places, requests are identified by a unique id. It is the responsibility -//! of the caller to generate unique ids. We could use random uuids here, but using -//! integers simplifies testing. -//! -//! Inside the state machine, we use [ChunkRanges] to represent avaiability bitfields. -//! We treat operations on such bitfields as very cheap, which is the case as long as -//! the bitfields are not very fragmented. We can introduce an even more optimized -//! bitfield type, or prevent fragmentation. -//! -//! The [DownloaderActor] is the asynchronous driver for the state machine. It -//! owns the actual tasks that perform IO. +//! The entry point is the [Downloader::builder] function, which creates a downloader +//! builder. The downloader is highly configurable. +//! +//! Content discovery is configurable via the [ContentDiscovery] trait. +//! Bitfield subscriptions are configurable via the [BitfieldSubscription] trait. +//! Download planning is configurable via the [DownloadPlanner] trait. +//! +//! After creating a downloader, you can schedule downloads using the +//! [Downloader::download] function. The function returns a future that +//! resolves once the download is complete. The download can be cancelled by +//! dropping the future. use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, future::Future, @@ -40,7 +32,8 @@ use crate::{ }; use anyhow::Context; use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; -use futures_lite::{Stream, StreamExt}; +use futures_lite::StreamExt; +use futures_util::stream::BoxStream; use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; @@ -79,20 +72,26 @@ pub struct FindPeersOpts { } /// A pluggable content discovery mechanism -pub trait ContentDiscovery: Send + 'static { +pub trait ContentDiscovery: std::fmt::Debug + Send + 'static { /// Find peers that have the given blob. /// /// The returned stream is a handle for the discovery task. It should be an /// infinite stream that only stops when it is dropped. - fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> impl Stream + Send + Unpin + 'static; + fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> BoxStream<'static, NodeId>; } +/// A boxed content discovery +pub type BoxedContentDiscovery = Box; + /// A pluggable bitfield subscription mechanism -pub trait BitfieldSubscription: Send + 'static { +pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static { /// Subscribe to a bitfield - fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> impl Stream + Send + Unpin + 'static; + fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> BoxStream<'static, BitfieldSubscriptionEvent>; } +/// A boxed bitfield subscription +pub type BoxedBitfieldSubscription = Box; + /// An event from a bitfield subscription #[derive(Debug)] pub enum BitfieldSubscriptionEvent { @@ -290,7 +289,7 @@ impl Bitfields { /// want to deduplicate the ranges, but they could also do other things, like /// eliminate gaps or even extend ranges. The only thing they should not do is /// to add new peers to the list of options. -pub trait DownloadPlanner: Send + 'static { +pub trait DownloadPlanner: Send + std::fmt::Debug + 'static { /// Make a download plan for a hash, by reducing or eliminating the overlap of chunk ranges fn plan(&mut self, hash: Hash, options: &mut BTreeMap); } @@ -1010,6 +1009,56 @@ pub struct Downloader { _task: Arc>, } +/// A builder for a downloader +#[derive(Debug)] +pub struct DownloaderBuilder { + endpoint: Endpoint, + store: S, + discovery: Option, + subscribe_bitfield: Option, + local_pool: Option, + planner: Option, +} + +impl DownloaderBuilder { + + /// Set the content discovery + pub fn discovery(self, discovery: D) -> Self { + Self { discovery: Some(Box::new(discovery)), ..self } + } + + /// Set the bitfield subscription + pub fn bitfield_subscription(self, value: B) -> Self { + Self { + subscribe_bitfield: Some(Box::new(value)), + ..self + } + } + + /// Set the local pool + pub fn local_pool(self, local_pool: LocalPool) -> Self { + Self { local_pool: Some(local_pool), ..self } + } + + /// Set the download planner + pub fn planner(self, planner: P) -> Self { + Self { planner: Some(Box::new(planner)), ..self } + } + + /// Build the downloader + pub fn build(self) -> Downloader + where + S: Store, + { + let store = self.store; + let discovery = self.discovery.expect("discovery not set"); + let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| Box::new(TestBitfieldSubscription)); + let local_pool = self.local_pool.unwrap_or_else(|| LocalPool::single()); + let planner = self.planner.unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); + Downloader::new(self.endpoint, store, discovery, subscribe_bitfield, local_pool, planner) + } +} + impl Downloader { /// Create a new download /// @@ -1021,8 +1070,13 @@ impl Downloader { Ok(()) } + /// Create a new downloader builder + pub fn builder(endpoint: Endpoint, store: S) -> DownloaderBuilder { + DownloaderBuilder { endpoint, store, discovery: None, subscribe_bitfield: None, local_pool: None, planner: None } + } + /// Create a new downloader - pub fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { + fn new(endpoint: Endpoint, store: S, discovery: BoxedContentDiscovery, subscribe_bitfield: BoxedBitfieldSubscription, local_pool: LocalPool, planner: Box) -> Self { let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitfield, local_pool, planner); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); @@ -1036,15 +1090,15 @@ enum UserCommand { Download { request: DownloadRequest, done: tokio::sync::oneshot::Sender<()> }, } -struct DownloaderActor { +struct DownloaderActor { local_pool: LocalPool, endpoint: Endpoint, command_rx: mpsc::Receiver, command_tx: mpsc::Sender, state: DownloaderState, store: S, - discovery: D, - subscribe_bitfield: B, + discovery: BoxedContentDiscovery, + subscribe_bitfield: BoxedBitfieldSubscription, download_futs: BTreeMap>, peer_download_tasks: BTreeMap>, discovery_tasks: BTreeMap>, @@ -1055,8 +1109,8 @@ struct DownloaderActor { start: Instant, } -impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: D, subscribe_bitfield: B, local_pool: LocalPool, planner: Box) -> Self { +impl DownloaderActor { + fn new(endpoint: Endpoint, store: S, discovery: BoxedContentDiscovery, subscribe_bitfield: BoxedBitfieldSubscription, local_pool: LocalPool, planner: Box) -> Self { let (send, recv) = mpsc::channel(256); Self { local_pool, @@ -1188,15 +1242,21 @@ pub struct StaticContentDiscovery { impl StaticContentDiscovery { /// Create a new static content discovery mechanism - pub fn new(info: BTreeMap>, default: Vec) -> Self { + pub fn new(mut info: BTreeMap>, mut default: Vec) -> Self { + default.sort(); + default.dedup(); + for (_, peers) in info.iter_mut() { + peers.sort(); + peers.dedup(); + } Self { info, default } } } impl ContentDiscovery for StaticContentDiscovery { - fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> impl Stream + Unpin + 'static { + fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> BoxStream<'static, NodeId> { let peers = self.info.get(&hash).unwrap_or(&self.default).clone(); - futures_lite::stream::iter(peers).chain(futures_lite::stream::pending()) + Box::pin(futures_lite::stream::iter(peers).chain(futures_lite::stream::pending())) } } @@ -1271,6 +1331,22 @@ where AbortOnDropHandle::new(task) } +/// A bitfield subscription that just returns nothing for local and everything(*) for remote +/// +/// * Still need to figure out how to deal with open ended chunk ranges. +#[derive(Debug)] +struct TestBitfieldSubscription; + +impl BitfieldSubscription for TestBitfieldSubscription { + fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> BoxStream<'static, BitfieldSubscriptionEvent> { + 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 }).chain(futures_lite::stream::pending())) + } +} + #[cfg(test)] mod tests { use std::ops::Range; @@ -1303,21 +1379,6 @@ mod tests { res.into_iter() } - /// A bitfield subscription that just returns nothing for local and everything(*) for remote - /// - /// * Still need to figure out how to deal with open ended chunk ranges. - struct TestBitfieldSubscription; - - impl BitfieldSubscription for TestBitfieldSubscription { - fn subscribe(&mut self, peer: BitfieldPeer, _hash: Hash) -> impl Stream + Send + Unpin + 'static { - let ranges = match peer { - BitfieldPeer::Local => ChunkRanges::empty(), - BitfieldPeer::Remote(_) => ChunkRanges::from(ChunkNum(0)..ChunkNum(1024 * 1024 * 1024 * 1024)), - }; - futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }).chain(futures_lite::stream::pending()) - } - } - fn peer(id: u8) -> NodeId { let mut secret = [0; 32]; secret[0] = id; @@ -1601,8 +1662,7 @@ mod tests { let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; let bitfield_subscription = TestBitfieldSubscription; - let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, bitfield_subscription, local_pool, noop_planner()); + let downloader = Downloader::builder(endpoint, store).discovery(discovery).bitfield_subscription(bitfield_subscription).build(); tokio::time::sleep(Duration::from_secs(2)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); fut.await?; @@ -1626,9 +1686,7 @@ mod tests { let store = crate::store::mem::Store::new(); let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; - let bitfield_subscription = TestBitfieldSubscription; - let local_pool = LocalPool::single(); - let downloader = Downloader::new(endpoint, store, discovery, bitfield_subscription, local_pool, Box::new(StripePlanner2::new(0, 8))); + let downloader = Downloader::builder(endpoint, store).discovery(discovery).planner(StripePlanner2::new(0, 8)).build(); tokio::time::sleep(Duration::from_secs(1)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); fut.await?; From 64d7ea837c9f94a7629c4b830b6f0f17b471470b Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 7 Feb 2025 17:09:15 +0200 Subject: [PATCH 24/47] Add debug output --- examples/multiprovider.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 94f767187..b2265b28a 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -36,8 +36,9 @@ async fn provide(args: ProvideArgs) -> anyhow::Result<()> { let mut tags = Vec::new(); for path in args.path { let data = std::fs::read(&path)?; + let len = data.len(); let tag = store.import_bytes(data.into(), iroh_blobs::BlobFormat::Raw).await?; - println!("added {} as {}", path.display(), tag.hash()); + println!("added {} as {}, {} bytes, {} chunks", path.display(), tag.hash(), len, (len + 1023) / 1024); tags.push((path, tag)); } let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; @@ -67,6 +68,7 @@ async fn download(args: DownloadArgs) -> anyhow::Result<()> { #[tokio::main] async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt::init(); let args = Args::parse(); match args.subcommand { Subcommand::Download(args) => download(args).await?, From 730bcf62b50dc22065e1d74e412c0fb9ed9a0b22 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 7 Feb 2025 19:10:52 +0200 Subject: [PATCH 25/47] Working multiprovider example --- examples/multiprovider.rs | 4 +- src/downloader2.rs | 78 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index b2265b28a..e3d44f047 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use bao_tree::ChunkRanges; +use bao_tree::{ChunkNum, ChunkRanges}; use iroh::NodeId; use iroh_blobs::{downloader2::{DownloadRequest, Downloader, StaticContentDiscovery}, store::Store, Hash}; use clap::Parser; @@ -60,7 +60,7 @@ async fn download(args: DownloadArgs) -> anyhow::Result<()> { let downloader = Downloader::builder(endpoint, store).discovery(discovery).build(); let request = DownloadRequest { hash: args.hash, - ranges: ChunkRanges::all(), + ranges: ChunkRanges::from(ChunkNum(0)..ChunkNum(25421)), }; downloader.download(request).await?; Ok(()) diff --git a/src/downloader2.rs b/src/downloader2.rs index c1a09fb2e..3d969e26a 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -27,13 +27,13 @@ use crate::{ }, protocol::{GetRequest, RangeSpec, RangeSpecSeq}, store::{BaoBatchWriter, MapEntryMut, Store}, - util::local_pool::{self, LocalPool}, + util::local_pool::{self, LocalPool, LocalPoolHandle}, Hash, }; use anyhow::Context; use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; use futures_lite::StreamExt; -use futures_util::stream::BoxStream; +use futures_util::{stream::BoxStream, FutureExt}; use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; @@ -1052,9 +1052,9 @@ impl DownloaderBuilder { { let store = self.store; let discovery = self.discovery.expect("discovery not set"); - let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| Box::new(TestBitfieldSubscription)); let local_pool = self.local_pool.unwrap_or_else(|| LocalPool::single()); let planner = self.planner.unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); + let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| Box::new(SimpleBitfieldSubscription::new(self.endpoint.clone(), store.clone(), local_pool.handle().clone()))); Downloader::new(self.endpoint, store, discovery, subscribe_bitfield, local_pool, planner) } } @@ -1334,6 +1334,7 @@ where /// A bitfield subscription that just returns nothing for local and everything(*) for remote /// /// * Still need to figure out how to deal with open ended chunk ranges. +#[allow(dead_code)] #[derive(Debug)] struct TestBitfieldSubscription; @@ -1347,6 +1348,77 @@ impl BitfieldSubscription for TestBitfieldSubscription { } } +/// A simple bitfield subscription that gets the valid ranges from a remote node, and the bitmap from a local store +#[derive(Debug)] +pub struct SimpleBitfieldSubscription { + endpoint: Endpoint, + store: S, + local_pool: LocalPoolHandle, +} + +impl SimpleBitfieldSubscription { + /// Create a new bitfield subscription + pub fn new(endpoint: Endpoint, store: S, local_pool: LocalPoolHandle) -> Self { + Self { endpoint, store, local_pool } + } +} + +async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { + if let Some(entry) = store.get_mut(&hash).await? { + crate::get::db::valid_ranges::(&entry).await + } else { + Ok(ChunkRanges::empty()) + } +} + +async fn get_valid_ranges_remote(endpoint: &Endpoint, id: NodeId, hash: &Hash) -> anyhow::Result { + 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))) +} + +impl BitfieldSubscription for SimpleBitfieldSubscription { + fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> BoxStream<'static, BitfieldSubscriptionEvent> { + let (send, recv) = tokio::sync::oneshot::channel(); + match peer { + BitfieldPeer::Local => { + let store = self.store.clone(); + self.local_pool.spawn_detached(move || async move { + match get_valid_ranges_local(&hash, store).await { + Ok(ranges) => { + send.send(ranges).ok(); + }, + Err(e) => { + tracing::error!("error getting bitfield: {e}"); + }, + }; + }); + } + BitfieldPeer::Remote(id) => { + let endpoint = self.endpoint.clone(); + tokio::spawn(async move { + match get_valid_ranges_remote(&endpoint, id, &hash).await { + Ok(ranges) => { + send.send(ranges).ok(); + }, + Err(cause) => { + tracing::error!("error getting bitfield: {cause}"); + } + } + }); + } + } + Box::pin(async move { + let ranges = match recv.await { + Ok(ev) => ev, + Err(_) => ChunkRanges::empty(), + }; + BitfieldSubscriptionEvent::Bitfield { ranges } + }.into_stream()) + } +} + #[cfg(test)] mod tests { use std::ops::Range; From 5c91fd4036bff28fc8e5cc7017c5530a8f396042 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 11:27:24 +0200 Subject: [PATCH 26/47] Add way to keep node id stable in multiprovider example provide --- examples/multiprovider.rs | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index e3d44f047..4fbc67e6e 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -1,7 +1,7 @@ -use std::path::PathBuf; +use std::{env::VarError, path::PathBuf, str::FromStr}; use bao_tree::{ChunkNum, ChunkRanges}; -use iroh::NodeId; +use iroh::{NodeId, SecretKey}; use iroh_blobs::{downloader2::{DownloadRequest, Downloader, StaticContentDiscovery}, store::Store, Hash}; use clap::Parser; @@ -31,6 +31,27 @@ struct ProvideArgs { path: Vec, } +fn load_secret_key() -> anyhow::Result> { + match std::env::var("IROH_SECRET") { + Ok(secret) => Ok(Some(SecretKey::from_str(&secret)?)), + Err(VarError::NotPresent) => Ok(None), + Err(x) => Err(x.into()), + } +} + +fn get_or_create_secret_key() -> iroh::SecretKey { + match load_secret_key() { + Ok(Some(secret)) => return secret, + Ok(None) => {} + Err(cause) => { + println!("failed to load secret key: {}", cause); + } + }; + let secret = SecretKey::generate(rand::thread_rng()); + println!("Using secret key {secret}. Set IROH_SECRET env var to use the same key next time."); + secret +} + async fn provide(args: ProvideArgs) -> anyhow::Result<()> { let store = iroh_blobs::store::mem::Store::new(); let mut tags = Vec::new(); @@ -41,7 +62,9 @@ async fn provide(args: ProvideArgs) -> anyhow::Result<()> { println!("added {} as {}, {} bytes, {} chunks", path.display(), tag.hash(), len, (len + 1023) / 1024); tags.push((path, tag)); } - let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; + let secret_key = get_or_create_secret_key(); + let endpoint = iroh::Endpoint::builder().discovery_n0() + .secret_key(secret_key).bind().await?; let id = endpoint.node_id(); let blobs = iroh_blobs::net_protocol::Blobs::builder(store).build(&endpoint); let router = iroh::protocol::Router::builder(endpoint) From 5a658f687c0aa34e32659754f7ffbe1c01e99e6a Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 15:27:16 +0200 Subject: [PATCH 27/47] Show some progress bars --- examples/multiprovider.rs | 172 +++++- src/downloader2.rs | 1210 +++++++++++++++++++++++++++++++------ src/util.rs | 2 +- 3 files changed, 1189 insertions(+), 195 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 4fbc67e6e..53aca8ed2 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -1,9 +1,12 @@ use std::{env::VarError, path::PathBuf, str::FromStr}; use bao_tree::{ChunkNum, ChunkRanges}; -use iroh::{NodeId, SecretKey}; -use iroh_blobs::{downloader2::{DownloadRequest, Downloader, StaticContentDiscovery}, store::Store, Hash}; use clap::Parser; +use console::Term; +use iroh::{NodeId, SecretKey}; +use iroh_blobs::{ + downloader2::{DownloadRequest, Downloader, ObserveEvent, ObserveRequest, StaticContentDiscovery}, store::Store, util::total_bytes, Hash +}; #[derive(Debug, Parser)] struct Args { @@ -31,7 +34,7 @@ struct ProvideArgs { path: Vec, } -fn load_secret_key() -> anyhow::Result> { +fn load_secret_key() -> anyhow::Result> { match std::env::var("IROH_SECRET") { Ok(secret) => Ok(Some(SecretKey::from_str(&secret)?)), Err(VarError::NotPresent) => Ok(None), @@ -58,33 +61,182 @@ async fn provide(args: ProvideArgs) -> anyhow::Result<()> { for path in args.path { let data = std::fs::read(&path)?; let len = data.len(); - let tag = store.import_bytes(data.into(), iroh_blobs::BlobFormat::Raw).await?; - println!("added {} as {}, {} bytes, {} chunks", path.display(), tag.hash(), len, (len + 1023) / 1024); + let tag = store + .import_bytes(data.into(), iroh_blobs::BlobFormat::Raw) + .await?; + println!( + "added {} as {}, {} bytes, {} chunks", + path.display(), + tag.hash(), + len, + (len + 1023) / 1024 + ); tags.push((path, tag)); } let secret_key = get_or_create_secret_key(); - let endpoint = iroh::Endpoint::builder().discovery_n0() - .secret_key(secret_key).bind().await?; + let endpoint = iroh::Endpoint::builder() + .discovery_n0() + .secret_key(secret_key) + .bind() + .await?; let id = endpoint.node_id(); let blobs = iroh_blobs::net_protocol::Blobs::builder(store).build(&endpoint); let router = iroh::protocol::Router::builder(endpoint) .accept(iroh_blobs::ALPN, blobs) - .spawn().await?; + .spawn() + .await?; println!("listening on {}", id); tokio::signal::ctrl_c().await?; router.shutdown().await?; Ok(()) } +/// Progress for a single download +struct BlobDownloadProgress { + request: DownloadRequest, + current: ChunkRanges, +} + +impl BlobDownloadProgress { + fn new(request: DownloadRequest) -> Self { + Self { + request, + current: ChunkRanges::empty(), + } + } + + fn update(&mut self, ev: ObserveEvent) { + match ev { + ObserveEvent::Bitfield { ranges } => { + self.current = ranges; + } + ObserveEvent::BitfieldUpdate { added, removed } => { + self.current |= added; + self.current -= 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); + (downloaded, total) + } + + #[allow(dead_code)] + fn get_bitmap(&self) -> String { + format!("{:?}", self.current) + } + + fn is_done(&self) -> bool { + self.current == self.request.ranges + } +} + +fn bitmap(current: &[ChunkNum], requested: &[ChunkNum], n: usize) -> String { + // If n is 0, return an empty string. + if n == 0 { + return String::new(); + } + + // Determine the overall bitfield size. + // Since the ranges are sorted, we take the last element as the total size. + let total = if let Some(&last) = requested.last() { + last.0 + } else { + // If there are no ranges, we assume the bitfield is empty. + 0 + }; + + // If total is 0, output n spaces. + if total == 0 { + return " ".repeat(n); + } + + let mut result = String::with_capacity(n); + + // For each of the n output buckets: + for bucket in 0..n { + // Calculate the bucket's start and end in the overall bitfield. + let bucket_start = bucket as u64 * total / n as u64; + let bucket_end = (bucket as u64 + 1) * total / n as u64; + let bucket_size = bucket_end.saturating_sub(bucket_start); + + // Sum the number of bits that are set in this bucket. + let mut set_bits = 0u64; + for pair in current.chunks_exact(2) { + let start = pair[0]; + let end = pair[1]; + // Determine the overlap between the bucket and the current range. + let overlap_start = start.0.max(bucket_start); + let overlap_end = end.0.min(bucket_end); + if overlap_start < overlap_end { + set_bits += overlap_end - overlap_start; + } + } + + // Calculate the fraction of the bucket that is set. + let fraction = if bucket_size > 0 { + set_bits as f64 / bucket_size as f64 + } else { + 0.0 + }; + + // Map the fraction to a grayscale character. + let ch = if fraction == 0.0 { + ' ' // completely empty + } else if fraction == 1.0 { + '█' // completely full + } else if fraction < 0.25 { + '░' + } else if fraction < 0.5 { + '▒' + } else { + '▓' + }; + + result.push(ch); + } + + result +} + async fn download(args: DownloadArgs) -> anyhow::Result<()> { let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; let store = iroh_blobs::store::mem::Store::new(); let discovery = StaticContentDiscovery::new(Default::default(), args.providers); - let downloader = Downloader::builder(endpoint, store).discovery(discovery).build(); + let downloader = Downloader::builder(endpoint, store) + .discovery(discovery) + .build(); let request = DownloadRequest { hash: args.hash, ranges: ChunkRanges::from(ChunkNum(0)..ChunkNum(25421)), }; + 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)), + 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 requested = progress.request.ranges.boundaries(); + let bitmap = bitmap(current, requested, rows as usize); + print!("\r{bitmap}"); + if progress.is_done() { + println!(""); + break; + } + } + anyhow::Ok(()) + }); downloader.download(request).await?; Ok(()) } @@ -98,4 +250,4 @@ async fn main() -> anyhow::Result<()> { Subcommand::Provide(args) => provide(args).await?, } Ok(()) -} \ No newline at end of file +} diff --git a/src/downloader2.rs b/src/downloader2.rs index 3d969e26a..2dc3496e0 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -2,11 +2,11 @@ //! //! The entry point is the [Downloader::builder] function, which creates a downloader //! builder. The downloader is highly configurable. -//! +//! //! Content discovery is configurable via the [ContentDiscovery] trait. //! Bitfield subscriptions are configurable via the [BitfieldSubscription] trait. //! Download planning is configurable via the [DownloadPlanner] trait. -//! +//! //! After creating a downloader, you can schedule downloads using the //! [Downloader::download] function. The function returns a future that //! resolves once the download is complete. The download can be cancelled by @@ -42,16 +42,29 @@ use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, +)] struct DownloadId(u64); -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, +)] +struct ObserveId(u64); + +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, +)] struct DiscoveryId(u64); -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, +)] struct PeerDownloadId(u64); -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From)] +#[derive( + Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, +)] struct BitfieldSubscriptionId(u64); /// Announce kind @@ -86,7 +99,11 @@ pub type BoxedContentDiscovery = Box; /// 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, BitfieldSubscriptionEvent>; } /// A boxed bitfield subscription @@ -109,6 +126,23 @@ pub enum BitfieldSubscriptionEvent { }, } +/// Events from observing a local bitfield +#[derive(Debug)] +pub enum ObserveEvent { + /// 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, + }, +} + /// Global information about a peer #[derive(Debug, Default)] struct PeerState { @@ -130,12 +164,16 @@ struct PeerBlobState { impl PeerBlobState { fn new(subscription_id: BitfieldSubscriptionId) -> Self { - Self { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty() } + Self { + subscription_id, + subscription_count: 1, + ranges: ChunkRanges::empty(), + } } } /// A download request -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct DownloadRequest { /// The blob we are downloading pub hash: Hash, @@ -143,6 +181,17 @@ pub struct DownloadRequest { pub ranges: ChunkRanges, } +/// A request to observe the local bitmap for a blob +#[derive(Debug, Clone)] +pub struct ObserveRequest { + /// The blob we are observing + pub hash: Hash, + /// The ranges we are interested in + pub ranges: ChunkRanges, + /// Buffer size + pub buffer: usize, +} + struct DownloadState { /// The request this state is for request: DownloadRequest, @@ -154,7 +203,11 @@ struct DownloadState { impl DownloadState { fn new(request: DownloadRequest) -> Self { - Self { request, peer_downloads: BTreeMap::new(), needs_rebalancing: false } + Self { + request, + peer_downloads: BTreeMap::new(), + needs_rebalancing: false, + } } } @@ -166,7 +219,10 @@ struct IdGenerator { impl Default for IdGenerator { fn default() -> Self { - Self { next_id: 0, _p: PhantomData } + Self { + next_id: 0, + _p: PhantomData, + } } } @@ -202,8 +258,13 @@ impl Downloads { self.by_id.insert(id, state); } - fn iter_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { - self.by_id.iter_mut().filter(move |x| x.1.request.hash == hash) + fn iter_mut_for_hash( + &mut self, + hash: Hash, + ) -> impl Iterator { + self.by_id + .iter_mut() + .filter(move |x| x.1.request.hash == hash) } fn iter(&mut self) -> impl Iterator { @@ -216,20 +277,33 @@ impl Downloads { } fn values_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { - self.by_id.values_mut().filter(move |x| x.request.hash == hash) + self.by_id + .values_mut() + .filter(move |x| x.request.hash == hash) } fn by_id_mut(&mut self, id: DownloadId) -> Option<&mut DownloadState> { self.by_id.get_mut(&id) } - fn by_peer_download_id_mut(&mut self, id: PeerDownloadId) -> Option<(&DownloadId, &mut DownloadState)> { - self.by_id.iter_mut().filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)).next() + fn by_peer_download_id_mut( + &mut self, + id: PeerDownloadId, + ) -> Option<(&DownloadId, &mut DownloadState)> { + self.by_id + .iter_mut() + .filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) + .next() } } #[derive(Default)] struct Bitfields { + // Counters to generate unique ids for various requests. + // We could use uuid here, but using integers simplifies testing. + // + // the next subscription id + subscription_id_gen: IdGenerator, by_peer_and_hash: BTreeMap<(BitfieldPeer, Hash), PeerBlobState>, } @@ -266,17 +340,19 @@ impl Bitfields { } fn remote_for_hash(&self, hash: Hash) -> impl Iterator { - self.by_peer_and_hash.iter().filter_map(move |((peer, h), state)| { - if let BitfieldPeer::Remote(peer) = peer { - if *h == hash { - Some((peer, state)) + self.by_peer_and_hash + .iter() + .filter_map(move |((peer, h), state)| { + if let BitfieldPeer::Remote(peer) = peer { + if *h == hash { + Some((peer, state)) + } else { + None + } } else { None } - } else { - None - } - }) + }) } } @@ -333,7 +409,10 @@ pub struct StripePlanner { impl StripePlanner { /// Create a new planner with the given seed and stripe size. pub fn new(seed: u64, stripe_size_log: u8) -> Self { - Self { seed, stripe_size_log } + Self { + seed, + stripe_size_log, + } } /// The score function to decide which peer to download from. @@ -350,7 +429,10 @@ impl StripePlanner { impl DownloadPlanner for StripePlanner { fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { - assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); + assert!( + options.values().all(|x| x.boundaries().len() % 2 == 0), + "open ranges not supported" + ); options.retain(|_, x| !x.is_empty()); if options.len() <= 1 { return; @@ -384,7 +466,10 @@ impl DownloadPlanner for StripePlanner { } } -fn get_continuous_ranges(options: &mut BTreeMap, stripe_size_log: u8) -> Option> { +fn get_continuous_ranges( + options: &mut BTreeMap, + stripe_size_log: u8, +) -> Option> { let mut ranges = BTreeSet::new(); for x in options.values() { ranges.extend(x.boundaries().iter().map(|x| x.0)); @@ -428,7 +513,10 @@ pub struct StripePlanner2 { impl StripePlanner2 { /// Create a new planner with the given seed and stripe size. pub fn new(seed: u64, stripe_size_log: u8) -> Self { - Self { seed, stripe_size_log } + Self { + seed, + stripe_size_log, + } } /// The score function to decide which peer to download from. @@ -444,7 +532,10 @@ impl StripePlanner2 { impl DownloadPlanner for StripePlanner2 { fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { - assert!(options.values().all(|x| x.boundaries().len() % 2 == 0), "open ranges not supported"); + assert!( + options.values().all(|x| x.boundaries().len() % 2 == 0), + "open ranges not supported" + ); options.retain(|_, x| !x.is_empty()); if options.len() <= 1 { return; @@ -463,9 +554,16 @@ impl DownloadPlanner for StripePlanner2 { matching.push((peer, peer_ranges)); } } - let mut peer_and_score = matching.iter().map(|(peer, _)| (Self::score(peer, self.seed), peer)).collect::>(); + let mut peer_and_score = matching + .iter() + .map(|(peer, _)| (Self::score(peer, self.seed), peer)) + .collect::>(); peer_and_score.sort(); - let peer_to_rank = peer_and_score.into_iter().enumerate().map(|(i, (_, peer))| (*peer, i as u64)).collect::>(); + let peer_to_rank = peer_and_score + .into_iter() + .enumerate() + .map(|(i, (_, peer))| (*peer, i as u64)) + .collect::>(); let n = matching.len() as u64; for (peer, _) in matching.iter() { let score = Some((peer_to_rank[*peer] + stripe) % n); @@ -496,6 +594,8 @@ struct DownloaderState { // // each item here corresponds to an active subscription bitfields: Bitfields, + /// Observers for local bitfields + observers: Observers, // all active downloads // // these are user downloads. each user download gets split into one or more @@ -505,11 +605,6 @@ struct DownloaderState { // // there is a discovery task for each blob we are interested in. discovery: BTreeMap, - // Counters to generate unique ids for various requests. - // We could use uuid here, but using integers simplifies testing. - // - // the next subscription id - subscription_id_gen: IdGenerator, // the next discovery id discovery_id_gen: IdGenerator, // the next peer download id @@ -521,11 +616,11 @@ struct DownloaderState { impl DownloaderState { fn new(planner: Box) -> Self { Self { - peers: BTreeMap::new(), + peers: Default::default(), downloads: Default::default(), bitfields: Default::default(), - discovery: BTreeMap::new(), - subscription_id_gen: Default::default(), + discovery: Default::default(), + observers: Default::default(), discovery_id_gen: Default::default(), peer_download_id_gen: Default::default(), planner, @@ -600,6 +695,10 @@ enum Command { DropPeer { peer: NodeId }, /// A peer has been discovered PeerDiscovered { peer: NodeId, hash: Hash }, + /// + ObserveLocal { id: ObserveId, hash: Hash, ranges: ChunkRanges }, + /// + StopObserveLocal { id: ObserveId }, /// A tick from the driver, for rebalancing Tick { time: Duration }, } @@ -616,6 +715,15 @@ enum Event { /// The unique id of the subscription id: BitfieldSubscriptionId, }, + LocalBitfield { + id: ObserveId, + ranges: ChunkRanges, + }, + LocalBitfieldUpdate { + id: ObserveId, + added: ChunkRanges, + removed: ChunkRanges, + }, StartDiscovery { hash: Hash, /// The unique id of the discovery task @@ -655,7 +763,9 @@ impl DownloaderState { /// Apply a command, using a mutable reference to the events fn apply_mut(&mut self, cmd: Command, evs: &mut Vec) { if let Err(cause) = self.apply_mut_0(cmd, evs) { - evs.push(Event::Error { message: format!("{cause}") }); + evs.push(Event::Error { + message: format!("{cause}"), + }); } } @@ -669,18 +779,25 @@ impl DownloaderState { /// - unsubscribing from bitfields if needed /// - stopping the discovery task if needed fn stop_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { - let removed = self.downloads.remove(&id).context(format!("removed unknown download {id:?}"))?; + let removed = self + .downloads + .remove(&id) + .context(format!("removed unknown download {id:?}"))?; let removed_hash = removed.request.hash; // stop associated peer downloads for peer_download in removed.peer_downloads.values() { - evs.push(Event::StopPeerDownload { id: peer_download.id }); + evs.push(Event::StopPeerDownload { + id: peer_download.id, + }); } // unsubscribe from bitfields that have no more subscriptions self.bitfields.retain(|(_peer, hash), state| { if *hash == removed_hash { state.subscription_count -= 1; if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitfield { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); return false; } } @@ -689,7 +806,13 @@ impl DownloaderState { let hash_interest = self.downloads.values_for_hash(removed.request.hash).count(); if hash_interest == 0 { // stop the discovery task if we were the last one interested in the hash - let discovery_id = self.discovery.remove(&removed.request.hash).context(format!("removed unknown discovery task for {}", removed.request.hash))?; + let discovery_id = self + .discovery + .remove(&removed.request.hash) + .context(format!( + "removed unknown discovery task for {}", + removed.request.hash + ))?; evs.push(Event::StopDiscovery { id: discovery_id }); } Ok(()) @@ -701,7 +824,10 @@ impl DownloaderState { match cmd { Command::StartDownload { request, id } => { // ids must be uniquely assigned by the caller! - anyhow::ensure!(!self.downloads.contains_key(&id), "duplicate download request {id:?}"); + anyhow::ensure!( + !self.downloads.contains_key(&id), + "duplicate download request {id:?}" + ); let hash = request.hash; // either we have a subscription for this blob, or we have to create one if let Some(state) = self.bitfields.get_local_mut(hash) { @@ -709,9 +835,16 @@ impl DownloaderState { state.subscription_count += 1; } else { // create a new subscription - let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { peer: BitfieldPeer::Local, hash, id: subscription_id }); - self.bitfields.insert((BitfieldPeer::Local, hash), PeerBlobState::new(subscription_id)); + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Local, + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Local, hash), + PeerBlobState::new(subscription_id), + ); } if !self.discovery.contains_key(&request.hash) { // start a discovery task @@ -724,7 +857,8 @@ impl DownloaderState { self.start_downloads(hash, Some(id), evs)?; } Command::PeerDownloadComplete { id, .. } => { - let Some((download_id, download)) = self.downloads.by_peer_download_id_mut(id) else { + let Some((download_id, download)) = self.downloads.by_peer_download_id_mut(id) + else { // the download was already removed return Ok(()); }; @@ -737,7 +871,10 @@ impl DownloaderState { self.stop_download(id, evs)?; } Command::PeerDiscovered { peer, hash } => { - if self.bitfields.contains_key(&(BitfieldPeer::Remote(peer), hash)) { + if self + .bitfields + .contains_key(&(BitfieldPeer::Remote(peer), hash)) + { // we already have a subscription for this peer return Ok(()); }; @@ -748,15 +885,24 @@ impl DownloaderState { // create a peer state if it does not exist let _state = self.peers.entry(peer).or_default(); // create a new subscription - let subscription_id = self.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { peer: BitfieldPeer::Remote(peer), hash, id: subscription_id }); - self.bitfields.insert((BitfieldPeer::Remote(peer), hash), PeerBlobState::new(subscription_id)); + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Remote(peer), + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Remote(peer), hash), + PeerBlobState::new(subscription_id), + ); } Command::DropPeer { peer } => { self.bitfields.retain(|(p, _), state| { if *p == BitfieldPeer::Remote(peer) { // todo: should we emit unsubscribe evs here? - evs.push(Event::UnsubscribeBitfield { id: state.subscription_id }); + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); return false; } else { return true; @@ -765,9 +911,20 @@ impl DownloaderState { self.peers.remove(&peer); } Command::Bitfield { peer, hash, ranges } => { - let state = self.bitfields.get_mut(&(peer, hash)).context(format!("bitfields for unknown peer {peer:?} and hash {hash}"))?; + let state = self.bitfields.get_mut(&(peer, hash)).context(format!( + "bitfields for unknown peer {peer:?} and hash {hash}" + ))?; let _chunks = total_chunks(&ranges).context("open range")?; if peer == BitfieldPeer::Local { + if let Some(observers) = self.observers.get_by_hash(&hash) { + for (id, request) in observers { + let ranges = &ranges & &request.ranges; + evs.push(Event::LocalBitfield { + id: *id, + ranges, + }); + } + } state.ranges = ranges; self.check_completion(hash, None, evs)?; } else { @@ -782,9 +939,29 @@ impl DownloaderState { // 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::BitfieldUpdate { peer, hash, added, removed } => { - let state = self.bitfields.get_mut(&(peer, hash)).context(format!("bitfield update for unknown peer {peer:?} and hash {hash}"))?; + Command::BitfieldUpdate { + peer, + hash, + added, + removed, + } => { + let state = self.bitfields.get_mut(&(peer, hash)).context(format!( + "bitfield update for unknown peer {peer:?} and hash {hash}" + ))?; if peer == BitfieldPeer::Local { + if let Some(observers) = self.observers.get_by_hash(&hash) { + for (id, request) in observers { + let added = &added & &request.ranges; + let removed = &removed & &request.ranges; + if !added.is_empty() || !removed.is_empty() { + evs.push(Event::LocalBitfieldUpdate { + id: *id, + added: &added & &request.ranges, + removed: &removed & &request.ranges, + }); + } + } + } state.ranges |= added; state.ranges &= !removed; self.check_completion(hash, None, evs)?; @@ -802,15 +979,25 @@ impl DownloaderState { self.start_downloads(hash, None, evs)?; } } - Command::ChunksDownloaded { time, peer, hash, added } => { - let state = self.bitfields.get_local_mut(hash).context(format!("chunks downloaded before having local bitfield for {hash}"))?; + Command::ChunksDownloaded { + time, + peer, + hash, + added, + } => { + let state = self.bitfields.get_local_mut(hash).context(format!( + "chunks downloaded before having local bitfield for {hash}" + ))?; let total_downloaded = total_chunks(&added).context("open range")?; let total_before = total_chunks(&state.ranges).context("open range")?; state.ranges |= added; let total_after = total_chunks(&state.ranges).context("open range")?; let useful_downloaded = total_after - total_before; - let peer = self.peers.get_mut(&peer).context(format!("performing download before having peer state for {peer}"))?; - peer.download_history.push_back((time, (total_downloaded, useful_downloaded))); + let peer = self.peers.get_mut(&peer).context(format!( + "performing download before having peer state for {peer}" + ))?; + peer.download_history + .push_back((time, (total_downloaded, useful_downloaded))); } Command::Tick { time } => { let window = 10; @@ -818,7 +1005,9 @@ impl DownloaderState { // clean up download history let mut to_rebalance = vec![]; for (peer, state) in self.peers.iter_mut() { - state.download_history.retain(|(duration, _)| *duration > horizon); + state + .download_history + .retain(|(duration, _)| *duration > horizon); let mut sum_total = 0; let mut sum_useful = 0; for (_, (total, useful)) in state.download_history.iter() { @@ -835,7 +1024,10 @@ impl DownloaderState { // nothing has changed that affects this download continue; } - let n_peers = self.bitfields.remote_for_hash(download.request.hash).count(); + let n_peers = self + .bitfields + .remote_for_hash(download.request.hash) + .count(); if download.peer_downloads.len() >= n_peers { // we are already downloading from all peers for this hash continue; @@ -846,6 +1038,47 @@ impl DownloaderState { self.rebalance_download(id, evs)?; } } + Command::ObserveLocal { id, hash, ranges } => { + // either we have a subscription for this blob, or we have to create one + if let Some(state) = self.bitfields.get_local_mut(hash) { + // just increment the count + state.subscription_count += 1; + // emit the current bitfield + evs.push(Event::LocalBitfield { id, ranges: state.ranges.clone() }); + } else { + // create a new subscription + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Local, + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Local, hash), + PeerBlobState::new(subscription_id), + ); + } + self.observers.insert(id, ObserveRequest { hash, ranges, buffer: 0 }); + } + Command::StopObserveLocal { id } => { + let request = self.observers.remove(&id).context(format!( + "stop observing unknown local bitfield {id:?}" + ))?; + let removed_hash = request.hash; + // unsubscribe from bitfields that have no more subscriptions + self.bitfields.retain(|(_peer, hash), state| { + if *hash == removed_hash { + state.subscription_count -= 1; + if state.subscription_count == 0 { + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); + return false; + } + } + true + }); + } } Ok(()) } @@ -855,7 +1088,12 @@ impl DownloaderState { /// This must be called after each change of the local bitfield for a hash /// /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. - fn check_completion(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { + fn check_completion( + &mut self, + hash: Hash, + just_id: Option, + evs: &mut Vec, + ) -> anyhow::Result<()> { let Some(self_state) = self.bitfields.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); @@ -881,7 +1119,9 @@ impl DownloaderState { // stop this peer download. // // Might be a noop if the cause for this local change was the same peer download, but we don't know. - evs.push(Event::StopPeerDownload { id: peer_download.id }); + evs.push(Event::StopPeerDownload { + id: peer_download.id, + }); // mark this peer as available available.push(*peer); false @@ -900,7 +1140,8 @@ impl DownloaderState { // see what the new peers can do for us let mut candidates = BTreeMap::new(); for peer in available { - let Some(peer_state) = self.bitfields.get(&(BitfieldPeer::Remote(peer), hash)) else { + let Some(peer_state) = self.bitfields.get(&(BitfieldPeer::Remote(peer), hash)) + else { // weird. we should have a bitfield for this peer since it just completed a download continue; }; @@ -914,8 +1155,15 @@ impl DownloaderState { // start new downloads for (peer, ranges) in candidates { let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); - download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); } } } @@ -927,12 +1175,21 @@ impl DownloaderState { } /// Look at all downloads for a hash and start peer downloads for those that do not have any yet - fn start_downloads(&mut self, hash: Hash, just_id: Option, evs: &mut Vec) -> anyhow::Result<()> { + fn start_downloads( + &mut self, + hash: Hash, + just_id: Option, + evs: &mut Vec, + ) -> anyhow::Result<()> { let Some(self_state) = self.bitfields.get_local(hash) else { // we don't have the self state yet, so we can't really decide if we need to download anything at all return Ok(()); }; - for (id, download) in self.downloads.iter_mut_for_hash(hash).filter(|(_, download)| download.peer_downloads.is_empty()) { + for (id, download) in self + .downloads + .iter_mut_for_hash(hash) + .filter(|(_, download)| download.peer_downloads.is_empty()) + { if just_id.is_some() && just_id != Some(*id) { continue; } @@ -948,8 +1205,15 @@ impl DownloaderState { for (peer, ranges) in candidates { info!(" Starting download from {peer} for {hash} {ranges:?}"); let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); - download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); } } Ok(()) @@ -957,7 +1221,10 @@ impl DownloaderState { /// rebalance a single download fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { - let download = self.downloads.by_id_mut(id).context(format!("rebalancing unknown download {id:?}"))?; + let download = self + .downloads + .by_id_mut(id) + .context(format!("rebalancing unknown download {id:?}"))?; download.needs_rebalancing = false; tracing::info!("Rebalancing download {id:?} {:?}", download.request); let hash = download.request.hash; @@ -974,7 +1241,10 @@ impl DownloaderState { } } self.planner.plan(hash, &mut candidates); - info!("Stopping {} old peer downloads", download.peer_downloads.len()); + info!( + "Stopping {} old peer downloads", + download.peer_downloads.len() + ); for (_, state) in &download.peer_downloads { // stop all downloads evs.push(Event::StopPeerDownload { id: state.id }); @@ -984,8 +1254,15 @@ impl DownloaderState { for (peer, ranges) in candidates { info!(" Starting download from {peer} for {hash} {ranges:?}"); let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { id, peer, hash, ranges: ranges.clone() }); - download.peer_downloads.insert(peer, PeerDownloadState { id, ranges }); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); } Ok(()) } @@ -1021,10 +1298,12 @@ pub struct DownloaderBuilder { } impl DownloaderBuilder { - /// Set the content discovery pub fn discovery(self, discovery: D) -> Self { - Self { discovery: Some(Box::new(discovery)), ..self } + Self { + discovery: Some(Box::new(discovery)), + ..self + } } /// Set the bitfield subscription @@ -1037,12 +1316,18 @@ impl DownloaderBuilder { /// Set the local pool pub fn local_pool(self, local_pool: LocalPool) -> Self { - Self { local_pool: Some(local_pool), ..self } + Self { + local_pool: Some(local_pool), + ..self + } } /// Set the download planner pub fn planner(self, planner: P) -> Self { - Self { planner: Some(Box::new(planner)), ..self } + Self { + planner: Some(Box::new(planner)), + ..self + } } /// Build the downloader @@ -1053,9 +1338,24 @@ impl DownloaderBuilder { let store = self.store; let discovery = self.discovery.expect("discovery not set"); let local_pool = self.local_pool.unwrap_or_else(|| LocalPool::single()); - let planner = self.planner.unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); - let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| Box::new(SimpleBitfieldSubscription::new(self.endpoint.clone(), store.clone(), local_pool.handle().clone()))); - Downloader::new(self.endpoint, store, discovery, subscribe_bitfield, local_pool, planner) + let planner = self + .planner + .unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); + let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| { + Box::new(SimpleBitfieldSubscription::new( + self.endpoint.clone(), + store.clone(), + local_pool.handle().clone(), + )) + }); + Downloader::new( + self.endpoint, + store, + discovery, + subscribe_bitfield, + local_pool, + planner, + ) } } @@ -1065,29 +1365,101 @@ impl Downloader { /// The download will be cancelled if the returned future is dropped. pub async fn download(&self, request: DownloadRequest) -> anyhow::Result<()> { let (send, recv) = tokio::sync::oneshot::channel::<()>(); - self.send.send(UserCommand::Download { request, done: send }).await?; + self.send + .send(UserCommand::Download { + request, + done: send, + }) + .await?; recv.await?; Ok(()) } + /// Observe a local bitmap + pub async fn observe( + &self, + request: ObserveRequest, + ) -> anyhow::Result> { + let (send, recv) = tokio::sync::mpsc::channel(request.buffer); + self.send + .send(UserCommand::Observe { request, send }) + .await?; + Ok(recv) + } + /// Create a new downloader builder pub fn builder(endpoint: Endpoint, store: S) -> DownloaderBuilder { - DownloaderBuilder { endpoint, store, discovery: None, subscribe_bitfield: None, local_pool: None, planner: None } + DownloaderBuilder { + endpoint, + store, + discovery: None, + subscribe_bitfield: None, + local_pool: None, + planner: None, + } } /// Create a new downloader - fn new(endpoint: Endpoint, store: S, discovery: BoxedContentDiscovery, subscribe_bitfield: BoxedBitfieldSubscription, local_pool: LocalPool, planner: Box) -> Self { - let actor = DownloaderActor::new(endpoint, store, discovery, subscribe_bitfield, local_pool, planner); + fn new( + endpoint: Endpoint, + store: S, + discovery: BoxedContentDiscovery, + subscribe_bitfield: BoxedBitfieldSubscription, + local_pool: LocalPool, + planner: Box, + ) -> Self { + let actor = DownloaderActor::new( + endpoint, + store, + discovery, + subscribe_bitfield, + local_pool, + planner, + ); let (send, recv) = tokio::sync::mpsc::channel(256); let task = Arc::new(spawn(async move { actor.run(recv).await })); Self { send, _task: task } } } +#[derive(Debug, Default)] +struct Observers { + by_hash_and_id: BTreeMap>, +} + +impl Observers { + fn insert(&mut self, id: ObserveId, request: ObserveRequest) { + self.by_hash_and_id + .entry(request.hash) + .or_default() + .insert(id, request); + } + + fn remove(&mut self, id: &ObserveId) -> Option { + for requests in self.by_hash_and_id.values_mut() { + if let Some(request) = requests.remove(id) { + return Some(request); + } + } + None + } + + fn get_by_hash(&self, hash: &Hash) -> Option<&BTreeMap> { + self.by_hash_and_id.get(hash) + } +} + /// An user-facing command #[derive(Debug)] enum UserCommand { - Download { request: DownloadRequest, done: tokio::sync::oneshot::Sender<()> }, + Download { + request: DownloadRequest, + done: tokio::sync::oneshot::Sender<()>, + }, + Observe { + request: ObserveRequest, + send: tokio::sync::mpsc::Sender, + }, } struct DownloaderActor { @@ -1105,12 +1477,23 @@ struct DownloaderActor { bitfield_subscription_tasks: BTreeMap>, /// Id generator for download ids download_id_gen: IdGenerator, + /// Id generator for observe ids + observe_id_gen: IdGenerator, + /// Observers + observers: BTreeMap>, /// The time when the actor was started, serves as the epoch for time messages to the state machine start: Instant, } impl DownloaderActor { - fn new(endpoint: Endpoint, store: S, discovery: BoxedContentDiscovery, subscribe_bitfield: BoxedBitfieldSubscription, local_pool: LocalPool, planner: Box) -> Self { + fn new( + endpoint: Endpoint, + store: S, + discovery: BoxedContentDiscovery, + subscribe_bitfield: BoxedBitfieldSubscription, + local_pool: LocalPool, + planner: Box, + ) -> Self { let (send, recv) = mpsc::channel(256); Self { local_pool, @@ -1126,6 +1509,8 @@ impl DownloaderActor { command_tx: send, command_rx: recv, download_id_gen: Default::default(), + observe_id_gen: Default::default(), + observers: Default::default(), start: Instant::now(), } } @@ -1146,6 +1531,11 @@ impl DownloaderActor { self.download_futs.insert(id, done); self.command_tx.send(Command::StartDownload { request, id }).await.ok(); } + UserCommand::Observe { request, send } => { + let id = self.observe_id_gen.next(); + self.command_tx.send(Command::ObserveLocal { id, hash: request.hash, ranges: request.ranges }).await.ok(); + self.observers.insert(id, send); + } } }, Some(cmd) = self.command_rx.recv() => { @@ -1161,11 +1551,27 @@ impl DownloaderActor { // // todo: is there a better mechanism than periodic checks? // I don't want some cancellation token rube goldberg machine. + let mut to_delete = vec![]; for (id, fut) in self.download_futs.iter() { if fut.is_closed() { + to_delete.push(*id); self.command_tx.send(Command::StopDownload { id: *id }).await.ok(); } } + for id in to_delete { + self.download_futs.remove(&id); + } + // clean up dropped observers + let mut to_delete = vec![]; + for (id, sender) in self.observers.iter() { + if sender.is_closed() { + to_delete.push(*id); + self.command_tx.send(Command::StopObserveLocal { id: *id }).await.ok(); + } + } + for id in to_delete { + self.observers.remove(&id); + } }, } } @@ -1180,8 +1586,17 @@ impl DownloaderActor { let task = spawn(async move { while let Some(ev) = stream.next().await { let cmd = match ev { - BitfieldSubscriptionEvent::Bitfield { ranges } => Command::Bitfield { peer, hash, ranges }, - BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => Command::BitfieldUpdate { peer, hash, added, removed }, + BitfieldSubscriptionEvent::Bitfield { ranges } => { + Command::Bitfield { peer, hash, ranges } + } + BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { + Command::BitfieldUpdate { + peer, + hash, + added, + removed, + } + } }; send.send(cmd).await.ok(); } @@ -1204,12 +1619,19 @@ impl DownloaderActor { }); self.discovery_tasks.insert(id, task); } - Event::StartPeerDownload { id, peer, hash, ranges } => { + Event::StartPeerDownload { + id, + peer, + hash, + ranges, + } => { let send = self.command_tx.clone(); let endpoint = self.endpoint.clone(); let store = self.store.clone(); let start = self.start; - let task = self.local_pool.spawn(move || peer_download_task(id, endpoint, store, hash, peer, ranges, send, start)); + let task = self.local_pool.spawn(move || { + peer_download_task(id, endpoint, store, hash, peer, ranges, send, start) + }); self.peer_download_tasks.insert(id, task); } Event::UnsubscribeBitfield { id } => { @@ -1226,6 +1648,24 @@ impl DownloaderActor { done.send(()).ok(); } } + Event::LocalBitfield { id, ranges } => { + let Some(sender) = self.observers.get(&id) else { + return; + }; + if sender.try_send(ObserveEvent::Bitfield { ranges }).is_err() { + // the observer has been dropped + self.observers.remove(&id); + } + } + Event::LocalBitfieldUpdate { id, added, removed } => { + let Some(sender) = self.observers.get(&id) else { + return; + }; + if sender.try_send(ObserveEvent::BitfieldUpdate { added, removed }).is_err() { + // the observer has been dropped + self.observers.remove(&id); + } + } Event::Error { message } => { error!("Error during processing event {}", message); } @@ -1260,12 +1700,32 @@ impl ContentDiscovery for StaticContentDiscovery { } } -async fn peer_download_task(id: PeerDownloadId, endpoint: Endpoint, store: S, hash: Hash, peer: NodeId, ranges: ChunkRanges, sender: mpsc::Sender, start: Instant) { +async fn peer_download_task( + id: PeerDownloadId, + endpoint: Endpoint, + store: S, + hash: Hash, + peer: NodeId, + ranges: ChunkRanges, + sender: mpsc::Sender, + start: Instant, +) { let result = peer_download(endpoint, store, hash, peer, ranges, &sender, start).await; - sender.send(Command::PeerDownloadComplete { id, result }).await.ok(); + sender + .send(Command::PeerDownloadComplete { id, result }) + .await + .ok(); } -async fn peer_download(endpoint: Endpoint, store: S, hash: Hash, peer: NodeId, ranges: ChunkRanges, sender: &mpsc::Sender, start: Instant) -> anyhow::Result { +async fn peer_download( + endpoint: Endpoint, + store: S, + hash: Hash, + peer: NodeId, + ranges: ChunkRanges, + sender: &mpsc::Sender, + start: Instant, +) -> anyhow::Result { info!("Connecting to peer {peer}"); let conn = endpoint.connect(peer, crate::ALPN).await?; info!("Got connection to peer {peer}"); @@ -1297,11 +1757,28 @@ async fn peer_download(endpoint: Endpoint, store: S, hash: Hash, peer: } BaoContentItem::Leaf(leaf) => { let start_chunk = leaf.offset / 1024; - let added = ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16)); - sender.send(Command::ChunksDownloaded { time: start.elapsed(), peer, hash, added: added.clone() }).await.ok(); + let added = + ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16)); + sender + .send(Command::ChunksDownloaded { + time: start.elapsed(), + peer, + hash, + added: added.clone(), + }) + .await + .ok(); batch.push(leaf.into()); writer.write_batch(size, std::mem::take(&mut batch)).await?; - sender.send(Command::BitfieldUpdate { peer: BitfieldPeer::Local, hash, added, removed: ChunkRanges::empty() }).await.ok(); + sender + .send(Command::BitfieldUpdate { + peer: BitfieldPeer::Local, + hash, + added, + removed: ChunkRanges::empty(), + }) + .await + .ok(); } } content = next; @@ -1339,12 +1816,21 @@ where 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, BitfieldSubscriptionEvent> { let ranges = match peer { BitfieldPeer::Local => ChunkRanges::empty(), - BitfieldPeer::Remote(_) => ChunkRanges::from(ChunkNum(0)..ChunkNum(1024 * 1024 * 1024 * 1024)), + BitfieldPeer::Remote(_) => { + ChunkRanges::from(ChunkNum(0)..ChunkNum(1024 * 1024 * 1024 * 1024)) + } }; - Box::pin(futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }).chain(futures_lite::stream::pending())) + Box::pin( + futures_lite::stream::once(BitfieldSubscriptionEvent::Bitfield { ranges }) + .chain(futures_lite::stream::pending()), + ) } } @@ -1359,7 +1845,11 @@ pub struct SimpleBitfieldSubscription { impl SimpleBitfieldSubscription { /// Create a new bitfield subscription pub fn new(endpoint: Endpoint, store: S, local_pool: LocalPoolHandle) -> Self { - Self { endpoint, store, local_pool } + Self { + endpoint, + store, + local_pool, + } } } @@ -1371,7 +1861,11 @@ async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Resu } } -async fn get_valid_ranges_remote(endpoint: &Endpoint, id: NodeId, hash: &Hash) -> anyhow::Result { +async fn get_valid_ranges_remote( + endpoint: &Endpoint, + id: NodeId, + hash: &Hash, +) -> anyhow::Result { let conn = endpoint.connect(id, crate::ALPN).await?; let (size, _) = crate::get::request::get_verified_size(&conn, &hash).await?; let chunks = (size + 1023) / 1024; @@ -1379,7 +1873,11 @@ async fn get_valid_ranges_remote(endpoint: &Endpoint, id: NodeId, hash: &Hash) - } impl BitfieldSubscription for SimpleBitfieldSubscription { - fn subscribe(&mut self, peer: BitfieldPeer, hash: Hash) -> BoxStream<'static, BitfieldSubscriptionEvent> { + fn subscribe( + &mut self, + peer: BitfieldPeer, + hash: Hash, + ) -> BoxStream<'static, BitfieldSubscriptionEvent> { let (send, recv) = tokio::sync::oneshot::channel(); match peer { BitfieldPeer::Local => { @@ -1388,10 +1886,10 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { match get_valid_ranges_local(&hash, store).await { Ok(ranges) => { send.send(ranges).ok(); - }, + } Err(e) => { tracing::error!("error getting bitfield: {e}"); - }, + } }; }); } @@ -1401,7 +1899,7 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { match get_valid_ranges_remote(&endpoint, id, &hash).await { Ok(ranges) => { send.send(ranges).ok(); - }, + } Err(cause) => { tracing::error!("error getting bitfield: {cause}"); } @@ -1409,13 +1907,16 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { }); } } - Box::pin(async move { - let ranges = match recv.await { - Ok(ev) => ev, - Err(_) => ChunkRanges::empty(), - }; - BitfieldSubscriptionEvent::Bitfield { ranges } - }.into_stream()) + Box::pin( + async move { + let ranges = match recv.await { + Ok(ev) => ev, + Err(_) => ChunkRanges::empty(), + }; + BitfieldSubscriptionEvent::Bitfield { ranges } + } + .into_stream(), + ) } } @@ -1447,7 +1948,9 @@ mod tests { RangeSetRange::Range(x) => x.end.0, }) .unwrap_or_default(); - let res = (0..max).map(move |i| x.contains(&ChunkNum(i))).collect::>(); + let res = (0..max) + .map(move |i| x.contains(&ChunkNum(i))) + .collect::>(); res.into_iter() } @@ -1473,7 +1976,11 @@ mod tests { fn test_planner_2() { let mut planner = StripePlanner2::new(0, 4); let hash = Hash::new(b"test"); - let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..100]), chunk_ranges([0..100])]); + let mut ranges = make_range_map(&[ + chunk_ranges([0..100]), + chunk_ranges([0..100]), + chunk_ranges([0..100]), + ]); println!(""); print_range_map(&ranges); println!("planning"); @@ -1485,7 +1992,12 @@ mod tests { fn test_planner_3() { let mut planner = StripePlanner2::new(0, 4); let hash = Hash::new(b"test"); - let mut ranges = make_range_map(&[chunk_ranges([0..100]), chunk_ranges([0..110]), chunk_ranges([0..120]), chunk_ranges([0..50])]); + let mut ranges = make_range_map(&[ + chunk_ranges([0..100]), + chunk_ranges([0..110]), + chunk_ranges([0..120]), + chunk_ranges([0..50]), + ]); println!(""); print_range_map(&ranges); println!("planning"); @@ -1526,8 +2038,15 @@ mod tests { let node_id = endpoint.node_id(); let store = crate::store::mem::Store::new(); let blobs = Blobs::builder(store).build(&endpoint); - let hash = blobs.client().add_bytes(bytes::Bytes::copy_from_slice(data)).await?.hash; - let router = iroh::protocol::Router::builder(endpoint).accept(crate::ALPN, blobs).spawn().await?; + let hash = blobs + .client() + .add_bytes(bytes::Bytes::copy_from_slice(data)) + .await? + .hash; + let router = iroh::protocol::Router::builder(endpoint) + .accept(crate::ALPN, blobs) + .spawn() + .await?; Ok((router, node_id, hash)) } @@ -1564,14 +2083,37 @@ mod tests { let _ = tracing_subscriber::fmt::try_init(); let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; + let unknown_hash = + "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::all() }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitfield should produce an error!"); - let evs = state.apply(Command::Bitfield { peer: Local, hash: unknown_hash, ranges: ChunkRanges::all() }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitfield for an unknown hash should produce an error!"); - let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "download from unknown peer should lead to an error!"); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::all(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "adding an open bitfield should produce an error!" + ); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash: unknown_hash, + ranges: ChunkRanges::all(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "adding an open bitfield for an unknown hash should produce an error!" + ); + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([0..16]), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "download from unknown peer should lead to an error!" + ); Ok(()) } @@ -1583,33 +2125,139 @@ mod tests { let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); - assert!(has_one_event(&evs, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "starting a download should start a discovery task"); - assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "starting a download should subscribe to the local bitfield"); + let evs = state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: DownloadId(0), + }); + assert!( + has_one_event( + &evs, + &Event::StartDiscovery { + hash, + id: DiscoveryId(0) + } + ), + "starting a download should start a discovery task" + ); + assert!( + has_one_event( + &evs, + &Event::SubscribeBitfield { + peer: Local, + hash, + id: BitfieldSubscriptionId(0) + } + ), + "starting a download should subscribe to the local bitfield" + ); let initial_bitfield = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply(Command::Bitfield { peer: Local, hash, ranges: initial_bitfield.clone() }); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: initial_bitfield.clone(), + }); assert!(evs.is_empty()); - assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, initial_bitfield, "bitfield should be set to the initial bitfield"); - assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.subscription_count, 1, "we have one download interested in the bitfield"); - let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .ranges, + initial_bitfield, + "bitfield should be set to the initial bitfield" + ); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .subscription_count, + 1, + "we have one download interested in the bitfield" + ); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }); assert!(evs.is_empty()); - assert_eq!(state.bitfields.get_local(hash).context("bitfield should be present")?.ranges, ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), "bitfield should be updated"); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .ranges, + ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), + "bitfield should be updated" + ); let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - assert!(has_one_event(&evs, &Event::SubscribeBitfield { peer: Remote(peer_a), hash, id: 1.into() }), "adding a new peer for a hash we are interested in should subscribe to the bitfield"); - let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..64]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: PeerDownloadId(0), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "bitfield from a peer should start a download"); + assert!( + has_one_event( + &evs, + &Event::SubscribeBitfield { + peer: Remote(peer_a), + hash, + id: 1.into() + } + ), + "adding a new peer for a hash we are interested in should subscribe to the bitfield" + ); + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..64]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: PeerDownloadId(0), + peer: peer_a, + hash, + ranges: chunk_ranges([32..64]) + } + ), + "bitfield from a peer should start a download" + ); // ChunksDownloaded just updates the peer stats - let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..48]) }); + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([32..48]), + }); assert!(evs.is_empty()); // Bitfield update does not yet complete the download - let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..48]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([32..48]), + removed: ChunkRanges::empty(), + }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats - let evs = state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([48..64]) }); + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([48..64]), + }); assert!(evs.is_empty()); // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([48..64]), removed: ChunkRanges::empty() }); - assert!(has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), "download should be completed by the data"); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([48..64]), + removed: ChunkRanges::empty(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), + "download should be completed by the data" + ); // quick check that everything got cleaned up assert!(state.downloads.by_id.is_empty()); assert!(state.bitfields.by_peer_and_hash.is_empty()); @@ -1626,32 +2274,116 @@ mod tests { let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); // Start a download - state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: DownloadId(0) }); + state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: DownloadId(0), + }); // Initially, we have nothing - state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::empty() }); + state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::empty(), + }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..32]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitfield from a peer should start a download"); + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..32]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 0.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([0..32]) + } + ), + "bitfield from a peer should start a download" + ); // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([0..16]) }); + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([0..16]), + }); // Bitfield update does not yet complete the download - state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([0..16]), removed: ChunkRanges::empty() }); + state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([0..16]), + removed: ChunkRanges::empty(), + }); // The peer now has more data - state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([32..64]) }); + state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([32..64]), + }); // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([16..32]) }); + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([16..32]), + }); // Complete the first part of the download - let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([16..32]), removed: ChunkRanges::empty() }); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }); // This triggers cancellation of the first peer download and starting a new one for the remaining data - assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), "first peer download should be stopped"); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 1.into(), peer: peer_a, hash, ranges: chunk_ranges([32..64]) }), "second peer download should be started"); + assert!( + has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), + "first peer download should be stopped" + ); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 1.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([32..64]) + } + ), + "second peer download should be started" + ); // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { time: Duration::ZERO, peer: peer_a, hash, added: chunk_ranges([32..64]) }); + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([32..64]), + }); // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { peer: Local, hash, added: chunk_ranges([32..64]), removed: ChunkRanges::empty() }); - assert!(has_all_events(&evs, &[&Event::StopPeerDownload { id: 1.into() }, &Event::DownloadComplete { id: 0.into() }, &Event::UnsubscribeBitfield { id: 0.into() }, &Event::StopDiscovery { id: 0.into() },]), "download should be completed by the data"); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([32..64]), + removed: ChunkRanges::empty(), + }); + assert!( + has_all_events( + &evs, + &[ + &Event::StopPeerDownload { id: 1.into() }, + &Event::DownloadComplete { id: 0.into() }, + &Event::UnsubscribeBitfield { id: 0.into() }, + &Event::StopDiscovery { id: 0.into() }, + ] + ), + "download should be completed by the data" + ); println!("{evs:?}"); Ok(()) } @@ -1667,33 +2399,80 @@ mod tests { // --- Start the first (ongoing) download. // Request a range from 0..64. let download0 = DownloadId(0); - let req0 = DownloadRequest { hash, ranges: chunk_ranges([0..64]) }; - let evs0 = state.apply(Command::StartDownload { request: req0, id: download0 }); + let req0 = DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }; + let evs0 = state.apply(Command::StartDownload { + request: req0, + id: download0, + }); // When starting the download, we expect a discovery task to be started // and a subscription to the local bitfield to be requested. - assert!(has_one_event(&evs0, &Event::StartDiscovery { hash, id: DiscoveryId(0) }), "download0 should start discovery"); - assert!(has_one_event(&evs0, &Event::SubscribeBitfield { peer: Local, hash, id: BitfieldSubscriptionId(0) }), "download0 should subscribe to the local bitfield"); + assert!( + has_one_event( + &evs0, + &Event::StartDiscovery { + hash, + id: DiscoveryId(0) + } + ), + "download0 should start discovery" + ); + assert!( + has_one_event( + &evs0, + &Event::SubscribeBitfield { + peer: Local, + hash, + id: BitfieldSubscriptionId(0) + } + ), + "download0 should subscribe to the local bitfield" + ); // --- Simulate some progress for the first download. // Let’s say only chunks 0..32 are available locally. - let evs1 = state.apply(Command::Bitfield { peer: Local, hash, ranges: chunk_ranges([0..32]) }); + let evs1 = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: chunk_ranges([0..32]), + }); // No completion event should be generated for download0 because its full range 0..64 is not yet met. - assert!(evs1.is_empty(), "Partial bitfield update should not complete download0"); + assert!( + evs1.is_empty(), + "Partial bitfield update should not complete download0" + ); // --- Start a second download for the same hash. // This new download only requires chunks 0..32 which are already available. let download1 = DownloadId(1); - let req1 = DownloadRequest { hash, ranges: chunk_ranges([0..32]) }; - let evs2 = state.apply(Command::StartDownload { request: req1, id: download1 }); + let req1 = DownloadRequest { + hash, + ranges: chunk_ranges([0..32]), + }; + let evs2 = state.apply(Command::StartDownload { + request: req1, + id: download1, + }); // Because the local bitfield (0..32) is already a superset of the new download’s request, // a DownloadComplete event for download1 should be generated immediately. - assert!(has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), "New download should complete immediately"); + assert!( + has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), + "New download should complete immediately" + ); // --- Verify state: // The ongoing download (download0) should still be present in the state, // while the newly completed download (download1) is removed. - assert!(state.downloads.contains_key(&download0), "download0 should still be active"); - assert!(!state.downloads.contains_key(&download1), "download1 should have been cleaned up after completion"); + assert!( + state.downloads.contains_key(&download0), + "download0 should still be active" + ); + assert!( + !state.downloads.contains_key(&download1), + "download1 should have been cleaned up after completion" + ); Ok(()) } @@ -1707,20 +2486,54 @@ mod tests { let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; let mut state = DownloaderState::new(noop_planner()); // Start a download - state.apply(Command::StartDownload { request: DownloadRequest { hash, ranges: chunk_ranges([0..64]) }, id: 0.into() }); + state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: 0.into(), + }); // Initially, we have nothing - state.apply(Command::Bitfield { peer: Local, hash, ranges: ChunkRanges::empty() }); + state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::empty(), + }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { peer: Remote(peer_a), hash, ranges: chunk_ranges([0..32]) }); - assert!(has_one_event(&evs, &Event::StartPeerDownload { id: 0.into(), peer: peer_a, hash, ranges: chunk_ranges([0..32]) }), "bitfield from a peer should start a download"); + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..32]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 0.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([0..32]) + } + ), + "bitfield from a peer should start a download" + ); // Sending StopDownload should stop the download and all associated tasks // This is what happens (delayed) when the user drops the download future let evs = state.apply(Command::StopDownload { id: 0.into() }); - assert!(has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() })); - assert!(has_one_event(&evs, &Event::UnsubscribeBitfield { id: 0.into() })); - assert!(has_one_event(&evs, &Event::UnsubscribeBitfield { id: 1.into() })); + assert!(has_one_event( + &evs, + &Event::StopPeerDownload { id: 0.into() } + )); + assert!(has_one_event( + &evs, + &Event::UnsubscribeBitfield { id: 0.into() } + )); + assert!(has_one_event( + &evs, + &Event::UnsubscribeBitfield { id: 1.into() } + )); assert!(has_one_event(&evs, &Event::StopDiscovery { id: 0.into() })); Ok(()) } @@ -1731,12 +2544,25 @@ mod tests { let _ = tracing_subscriber::fmt::try_init(); let (_router1, peer, hash) = make_test_node(b"test").await?; let store = crate::store::mem::Store::new(); - let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; - let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: vec![peer] }; + let endpoint = iroh::Endpoint::builder() + .alpns(vec![crate::protocol::ALPN.to_vec()]) + .discovery_n0() + .bind() + .await?; + let discovery = StaticContentDiscovery { + info: BTreeMap::new(), + default: vec![peer], + }; let bitfield_subscription = TestBitfieldSubscription; - let downloader = Downloader::builder(endpoint, store).discovery(discovery).bitfield_subscription(bitfield_subscription).build(); + let downloader = Downloader::builder(endpoint, store) + .discovery(discovery) + .bitfield_subscription(bitfield_subscription) + .build(); tokio::time::sleep(Duration::from_secs(2)).await; - let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]) }); + let fut = downloader.download(DownloadRequest { + hash, + ranges: chunk_ranges([0..1]), + }); fut.await?; Ok(()) } @@ -1753,14 +2579,30 @@ mod tests { nodes.push(make_test_node(&data).await?); } let peers = nodes.iter().map(|(_, peer, _)| *peer).collect::>(); - let hashes = nodes.iter().map(|(_, _, hash)| *hash).collect::>(); + let hashes = nodes + .iter() + .map(|(_, _, hash)| *hash) + .collect::>(); let hash = *hashes.iter().next().unwrap(); let store = crate::store::mem::Store::new(); - let endpoint = iroh::Endpoint::builder().alpns(vec![crate::protocol::ALPN.to_vec()]).discovery_n0().bind().await?; - let discovery = StaticContentDiscovery { info: BTreeMap::new(), default: peers }; - let downloader = Downloader::builder(endpoint, store).discovery(discovery).planner(StripePlanner2::new(0, 8)).build(); + let endpoint = iroh::Endpoint::builder() + .alpns(vec![crate::protocol::ALPN.to_vec()]) + .discovery_n0() + .bind() + .await?; + let discovery = StaticContentDiscovery { + info: BTreeMap::new(), + default: peers, + }; + let downloader = Downloader::builder(endpoint, store) + .discovery(discovery) + .planner(StripePlanner2::new(0, 8)) + .build(); tokio::time::sleep(Duration::from_secs(1)).await; - let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1024]) }); + let fut = downloader.download(DownloadRequest { + hash, + ranges: chunk_ranges([0..1024]), + }); fut.await?; Ok(()) } diff --git a/src/util.rs b/src/util.rs index 735a9feb1..6e3d29b62 100644 --- a/src/util.rs +++ b/src/util.rs @@ -238,7 +238,7 @@ impl Drop for TempTag { /// Get the number of bytes given a set of chunk ranges and the total size. /// /// If some ranges are out of bounds, they will be clamped to the size. -pub fn total_bytes(ranges: ChunkRanges, size: u64) -> u64 { +pub fn total_bytes(ranges: &ChunkRanges, size: u64) -> u64 { ranges .iter() .map(|range| { From 9ba84c300a6231849137965676faa9821d850811 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 15:55:35 +0200 Subject: [PATCH 28/47] Don't use DownloadRequest for internal state --- examples/multiprovider.rs | 13 +++++--- src/downloader2.rs | 52 ++++++++++++++++++++--------- src/downloader2/downloader_state.rs | 0 3 files changed, 45 insertions(+), 20 deletions(-) create mode 100644 src/downloader2/downloader_state.rs diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 53aca8ed2..e00c2a037 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -5,7 +5,12 @@ use clap::Parser; use console::Term; use iroh::{NodeId, SecretKey}; use iroh_blobs::{ - downloader2::{DownloadRequest, Downloader, ObserveEvent, ObserveRequest, StaticContentDiscovery}, store::Store, util::total_bytes, Hash + downloader2::{ + DownloadRequest, Downloader, ObserveEvent, ObserveRequest, StaticContentDiscovery, + }, + store::Store, + util::total_bytes, + Hash, }; #[derive(Debug, Parser)] @@ -132,7 +137,7 @@ impl BlobDownloadProgress { fn is_done(&self) -> bool { self.current == self.request.ranges } -} +} fn bitmap(current: &[ChunkNum], requested: &[ChunkNum], n: usize) -> String { // If n is 0, return an empty string. @@ -185,9 +190,9 @@ fn bitmap(current: &[ChunkNum], requested: &[ChunkNum], n: usize) -> String { // Map the fraction to a grayscale character. let ch = if fraction == 0.0 { - ' ' // completely empty + ' ' // completely empty } else if fraction == 1.0 { - '█' // completely full + '█' // completely full } else if fraction < 0.25 { '░' } else if fraction < 0.5 { diff --git a/src/downloader2.rs b/src/downloader2.rs index 2dc3496e0..021d69f9e 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -696,7 +696,11 @@ enum Command { /// A peer has been discovered PeerDiscovered { peer: NodeId, hash: Hash }, /// - ObserveLocal { id: ObserveId, hash: Hash, ranges: ChunkRanges }, + ObserveLocal { + id: ObserveId, + hash: Hash, + ranges: ChunkRanges, + }, /// StopObserveLocal { id: ObserveId }, /// A tick from the driver, for rebalancing @@ -749,7 +753,9 @@ enum Event { id: DownloadId, }, /// An error that stops processing the command - Error { message: String }, + Error { + message: String, + }, } impl DownloaderState { @@ -916,13 +922,12 @@ impl DownloaderState { ))?; let _chunks = total_chunks(&ranges).context("open range")?; 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 if let Some(observers) = self.observers.get_by_hash(&hash) { for (id, request) in observers { let ranges = &ranges & &request.ranges; - evs.push(Event::LocalBitfield { - id: *id, - ranges, - }); + evs.push(Event::LocalBitfield { id: *id, ranges }); } } state.ranges = ranges; @@ -949,6 +954,8 @@ impl DownloaderState { "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 if let Some(observers) = self.observers.get_by_hash(&hash) { for (id, request) in observers { let added = &added & &request.ranges; @@ -1044,7 +1051,10 @@ impl DownloaderState { // just increment the count state.subscription_count += 1; // emit the current bitfield - evs.push(Event::LocalBitfield { id, ranges: state.ranges.clone() }); + evs.push(Event::LocalBitfield { + id, + ranges: state.ranges.clone(), + }); } else { // create a new subscription let subscription_id = self.bitfields.subscription_id_gen.next(); @@ -1058,12 +1068,13 @@ impl DownloaderState { PeerBlobState::new(subscription_id), ); } - self.observers.insert(id, ObserveRequest { hash, ranges, buffer: 0 }); + self.observers.insert(id, ObserveInfo { hash, ranges }); } Command::StopObserveLocal { id } => { - let request = self.observers.remove(&id).context(format!( - "stop observing unknown local bitfield {id:?}" - ))?; + let request = self + .observers + .remove(&id) + .context(format!("stop observing unknown local bitfield {id:?}"))?; let removed_hash = request.hash; // unsubscribe from bitfields that have no more subscriptions self.bitfields.retain(|(_peer, hash), state| { @@ -1422,20 +1433,26 @@ impl Downloader { } } +#[derive(Debug)] +struct ObserveInfo { + hash: Hash, + ranges: ChunkRanges, +} + #[derive(Debug, Default)] struct Observers { - by_hash_and_id: BTreeMap>, + by_hash_and_id: BTreeMap>, } impl Observers { - fn insert(&mut self, id: ObserveId, request: ObserveRequest) { + fn insert(&mut self, id: ObserveId, request: ObserveInfo) { self.by_hash_and_id .entry(request.hash) .or_default() .insert(id, request); } - fn remove(&mut self, id: &ObserveId) -> Option { + fn remove(&mut self, id: &ObserveId) -> Option { for requests in self.by_hash_and_id.values_mut() { if let Some(request) = requests.remove(id) { return Some(request); @@ -1444,7 +1461,7 @@ impl Observers { None } - fn get_by_hash(&self, hash: &Hash) -> Option<&BTreeMap> { + fn get_by_hash(&self, hash: &Hash) -> Option<&BTreeMap> { self.by_hash_and_id.get(hash) } } @@ -1661,7 +1678,10 @@ impl DownloaderActor { let Some(sender) = self.observers.get(&id) else { return; }; - if sender.try_send(ObserveEvent::BitfieldUpdate { added, removed }).is_err() { + if sender + .try_send(ObserveEvent::BitfieldUpdate { added, removed }) + .is_err() + { // the observer has been dropped self.observers.remove(&id); } diff --git a/src/downloader2/downloader_state.rs b/src/downloader2/downloader_state.rs new file mode 100644 index 000000000..e69de29bb From 87938af5189193dbfb2d3664ca955c6f885f6a9c Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 15:59:50 +0200 Subject: [PATCH 29/47] start moving stuff into separate mods --- src/downloader2.rs | 231 +---------------- .../{downloader_state.rs => actor.rs} | 0 src/downloader2/planners.rs | 234 ++++++++++++++++++ src/downloader2/state.rs | 0 4 files changed, 238 insertions(+), 227 deletions(-) rename src/downloader2/{downloader_state.rs => actor.rs} (100%) create mode 100644 src/downloader2/planners.rs create mode 100644 src/downloader2/state.rs diff --git a/src/downloader2.rs b/src/downloader2.rs index 021d69f9e..7f69050d5 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -12,7 +12,7 @@ //! resolves once the download is complete. The download can be cancelled by //! dropping the future. use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, + collections::{BTreeMap, VecDeque}, future::Future, io, marker::PhantomData, @@ -42,6 +42,9 @@ use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; +mod planners; +pub(crate) use planners::*; + #[derive( Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, )] @@ -356,232 +359,6 @@ impl Bitfields { } } -/// Trait for a download planner. -/// -/// A download planner has the option to be stateful and keep track of plans -/// depending on the hash, but many planners will be stateless. -/// -/// Planners can do whatever they want with the chunk ranges. Usually, they -/// want to deduplicate the ranges, but they could also do other things, like -/// eliminate gaps or even extend ranges. The only thing they should not do is -/// to add new peers to the list of options. -pub trait DownloadPlanner: Send + std::fmt::Debug + 'static { - /// Make a download plan for a hash, by reducing or eliminating the overlap of chunk ranges - fn plan(&mut self, hash: Hash, options: &mut BTreeMap); -} - -/// A boxed download planner -pub type BoxedDownloadPlanner = Box; - -/// A download planner that just leaves everything as is. -/// -/// Data will be downloaded from all peers wherever multiple peers have the same data. -#[derive(Debug, Clone, Copy)] -pub struct NoopPlanner; - -impl DownloadPlanner for NoopPlanner { - fn plan(&mut self, _hash: Hash, _options: &mut BTreeMap) {} -} - -/// A download planner that fully removes overlap between peers. -/// -/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, -/// and for each stripe decides on a single peer to download from, based on the -/// peer id and a random seed. -#[derive(Debug)] -pub struct StripePlanner { - /// seed for the score function. This can be set to 0 for testing for - /// maximum determinism, but can be set to a random value for production - /// to avoid multiple downloaders coming up with the same plan. - seed: u64, - /// The log of the stripe size in chunks. This planner is relatively - /// dumb and does not try to come up with continuous ranges, but you can - /// just set this to a large value to avoid fragmentation. - /// - /// In the very common case where you have small downloads, this will - /// frequently just choose a single peer for the entire download. - /// - /// This is a feature, not a bug. For small downloads, it is not worth - /// the effort to come up with a more sophisticated plan. - stripe_size_log: u8, -} - -impl StripePlanner { - /// Create a new planner with the given seed and stripe size. - pub fn new(seed: u64, stripe_size_log: u8) -> Self { - Self { - seed, - stripe_size_log, - } - } - - /// The score function to decide which peer to download from. - fn score(peer: &NodeId, seed: u64, stripe: u64) -> u64 { - // todo: use fnv? blake3 is a bit overkill - let mut data = [0u8; 32 + 8 + 8]; - data[..32].copy_from_slice(peer.as_bytes()); - data[32..40].copy_from_slice(&stripe.to_be_bytes()); - data[40..48].copy_from_slice(&seed.to_be_bytes()); - let hash = blake3::hash(&data); - u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) - } -} - -impl DownloadPlanner for StripePlanner { - fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { - assert!( - options.values().all(|x| x.boundaries().len() % 2 == 0), - "open ranges not supported" - ); - options.retain(|_, x| !x.is_empty()); - if options.len() <= 1 { - return; - } - let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); - for range in ranges.windows(2) { - let start = ChunkNum(range[0]); - let end = ChunkNum(range[1]); - let curr = ChunkRanges::from(start..end); - let stripe = range[0] >> self.stripe_size_log; - let mut best_peer = None; - let mut best_score = 0; - let mut matching = vec![]; - for (peer, peer_ranges) in options.iter_mut() { - if peer_ranges.contains(&start) { - let score = Self::score(peer, self.seed, stripe); - if score > best_score && peer_ranges.contains(&start) { - best_peer = Some(*peer); - best_score = score; - } - matching.push((peer, peer_ranges)); - } - } - for (peer, peer_ranges) in matching { - if *peer != best_peer.unwrap() { - peer_ranges.difference_with(&curr); - } - } - } - options.retain(|_, x| !x.is_empty()); - } -} - -fn get_continuous_ranges( - options: &mut BTreeMap, - stripe_size_log: u8, -) -> Option> { - let mut ranges = BTreeSet::new(); - for x in options.values() { - ranges.extend(x.boundaries().iter().map(|x| x.0)); - } - let min = ranges.iter().next().copied()?; - let max = ranges.iter().next_back().copied()?; - // add stripe subdividers - for i in (min >> stripe_size_log)..(max >> stripe_size_log) { - let x = i << stripe_size_log; - if x > min && x < max { - ranges.insert(x); - } - } - let ranges = ranges.into_iter().collect::>(); - Some(ranges) -} - -/// A download planner that fully removes overlap between peers. -/// -/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, -/// and for each stripe decides on a single peer to download from, based on the -/// peer id and a random seed. -#[derive(Debug)] -pub struct StripePlanner2 { - /// seed for the score function. This can be set to 0 for testing for - /// maximum determinism, but can be set to a random value for production - /// to avoid multiple downloaders coming up with the same plan. - seed: u64, - /// The log of the stripe size in chunks. This planner is relatively - /// dumb and does not try to come up with continuous ranges, but you can - /// just set this to a large value to avoid fragmentation. - /// - /// In the very common case where you have small downloads, this will - /// frequently just choose a single peer for the entire download. - /// - /// This is a feature, not a bug. For small downloads, it is not worth - /// the effort to come up with a more sophisticated plan. - stripe_size_log: u8, -} - -impl StripePlanner2 { - /// Create a new planner with the given seed and stripe size. - pub fn new(seed: u64, stripe_size_log: u8) -> Self { - Self { - seed, - stripe_size_log, - } - } - - /// The score function to decide which peer to download from. - fn score(peer: &NodeId, seed: u64) -> u64 { - // todo: use fnv? blake3 is a bit overkill - let mut data = [0u8; 32 + 8]; - data[..32].copy_from_slice(peer.as_bytes()); - data[32..40].copy_from_slice(&seed.to_be_bytes()); - let hash = blake3::hash(&data); - u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) - } -} - -impl DownloadPlanner for StripePlanner2 { - fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { - assert!( - options.values().all(|x| x.boundaries().len() % 2 == 0), - "open ranges not supported" - ); - options.retain(|_, x| !x.is_empty()); - if options.len() <= 1 { - return; - } - let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); - for range in ranges.windows(2) { - let start = ChunkNum(range[0]); - let end = ChunkNum(range[1]); - let curr = ChunkRanges::from(start..end); - let stripe = range[0] >> self.stripe_size_log; - let mut best_peer = None; - let mut best_score = None; - let mut matching = vec![]; - for (peer, peer_ranges) in options.iter_mut() { - if peer_ranges.contains(&start) { - matching.push((peer, peer_ranges)); - } - } - let mut peer_and_score = matching - .iter() - .map(|(peer, _)| (Self::score(peer, self.seed), peer)) - .collect::>(); - peer_and_score.sort(); - let peer_to_rank = peer_and_score - .into_iter() - .enumerate() - .map(|(i, (_, peer))| (*peer, i as u64)) - .collect::>(); - let n = matching.len() as u64; - for (peer, _) in matching.iter() { - let score = Some((peer_to_rank[*peer] + stripe) % n); - if score > best_score { - best_peer = Some(**peer); - best_score = score; - } - } - for (peer, peer_ranges) in matching { - if *peer != best_peer.unwrap() { - peer_ranges.difference_with(&curr); - } - } - } - options.retain(|_, x| !x.is_empty()); - } -} - struct PeerDownloadState { id: PeerDownloadId, ranges: ChunkRanges, diff --git a/src/downloader2/downloader_state.rs b/src/downloader2/actor.rs similarity index 100% rename from src/downloader2/downloader_state.rs rename to src/downloader2/actor.rs diff --git a/src/downloader2/planners.rs b/src/downloader2/planners.rs new file mode 100644 index 000000000..5d87ad73c --- /dev/null +++ b/src/downloader2/planners.rs @@ -0,0 +1,234 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use bao_tree::{ChunkNum, ChunkRanges}; +use iroh::NodeId; + +use crate::Hash; + +/// Trait for a download planner. +/// +/// A download planner has the option to be stateful and keep track of plans +/// depending on the hash, but many planners will be stateless. +/// +/// Planners can do whatever they want with the chunk ranges. Usually, they +/// want to deduplicate the ranges, but they could also do other things, like +/// eliminate gaps or even extend ranges. The only thing they should not do is +/// to add new peers to the list of options. +pub trait DownloadPlanner: Send + std::fmt::Debug + 'static { + /// Make a download plan for a hash, by reducing or eliminating the overlap of chunk ranges + fn plan(&mut self, hash: Hash, options: &mut BTreeMap); +} + +/// A boxed download planner +pub type BoxedDownloadPlanner = Box; + + +/// A download planner that just leaves everything as is. +/// +/// Data will be downloaded from all peers wherever multiple peers have the same data. +#[derive(Debug, Clone, Copy)] +pub struct NoopPlanner; + +impl DownloadPlanner for NoopPlanner { + fn plan(&mut self, _hash: Hash, _options: &mut BTreeMap) {} +} + +/// A download planner that fully removes overlap between peers. +/// +/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, +/// and for each stripe decides on a single peer to download from, based on the +/// peer id and a random seed. +#[derive(Debug)] +pub struct StripePlanner { + /// seed for the score function. This can be set to 0 for testing for + /// maximum determinism, but can be set to a random value for production + /// to avoid multiple downloaders coming up with the same plan. + seed: u64, + /// The log of the stripe size in chunks. This planner is relatively + /// dumb and does not try to come up with continuous ranges, but you can + /// just set this to a large value to avoid fragmentation. + /// + /// In the very common case where you have small downloads, this will + /// frequently just choose a single peer for the entire download. + /// + /// This is a feature, not a bug. For small downloads, it is not worth + /// the effort to come up with a more sophisticated plan. + stripe_size_log: u8, +} + +impl StripePlanner { + /// Create a new planner with the given seed and stripe size. + #[allow(dead_code)] + pub fn new(seed: u64, stripe_size_log: u8) -> Self { + Self { + seed, + stripe_size_log, + } + } + + /// The score function to decide which peer to download from. + fn score(peer: &NodeId, seed: u64, stripe: u64) -> u64 { + // todo: use fnv? blake3 is a bit overkill + let mut data = [0u8; 32 + 8 + 8]; + data[..32].copy_from_slice(peer.as_bytes()); + data[32..40].copy_from_slice(&stripe.to_be_bytes()); + data[40..48].copy_from_slice(&seed.to_be_bytes()); + let hash = blake3::hash(&data); + u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) + } +} + +impl DownloadPlanner for StripePlanner { + fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { + assert!( + options.values().all(|x| x.boundaries().len() % 2 == 0), + "open ranges not supported" + ); + options.retain(|_, x| !x.is_empty()); + if options.len() <= 1 { + return; + } + let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); + for range in ranges.windows(2) { + let start = ChunkNum(range[0]); + let end = ChunkNum(range[1]); + let curr = ChunkRanges::from(start..end); + let stripe = range[0] >> self.stripe_size_log; + let mut best_peer = None; + let mut best_score = 0; + let mut matching = vec![]; + for (peer, peer_ranges) in options.iter_mut() { + if peer_ranges.contains(&start) { + let score = Self::score(peer, self.seed, stripe); + if score > best_score && peer_ranges.contains(&start) { + best_peer = Some(*peer); + best_score = score; + } + matching.push((peer, peer_ranges)); + } + } + for (peer, peer_ranges) in matching { + if *peer != best_peer.unwrap() { + peer_ranges.difference_with(&curr); + } + } + } + options.retain(|_, x| !x.is_empty()); + } +} + +fn get_continuous_ranges( + options: &mut BTreeMap, + stripe_size_log: u8, +) -> Option> { + let mut ranges = BTreeSet::new(); + for x in options.values() { + ranges.extend(x.boundaries().iter().map(|x| x.0)); + } + let min = ranges.iter().next().copied()?; + let max = ranges.iter().next_back().copied()?; + // add stripe subdividers + for i in (min >> stripe_size_log)..(max >> stripe_size_log) { + let x = i << stripe_size_log; + if x > min && x < max { + ranges.insert(x); + } + } + let ranges = ranges.into_iter().collect::>(); + Some(ranges) +} + +/// A download planner that fully removes overlap between peers. +/// +/// It divides files into stripes of a fixed size `1 << stripe_size_log` chunks, +/// and for each stripe decides on a single peer to download from, based on the +/// peer id and a random seed. +#[derive(Debug)] +pub struct StripePlanner2 { + /// seed for the score function. This can be set to 0 for testing for + /// maximum determinism, but can be set to a random value for production + /// to avoid multiple downloaders coming up with the same plan. + seed: u64, + /// The log of the stripe size in chunks. This planner is relatively + /// dumb and does not try to come up with continuous ranges, but you can + /// just set this to a large value to avoid fragmentation. + /// + /// In the very common case where you have small downloads, this will + /// frequently just choose a single peer for the entire download. + /// + /// This is a feature, not a bug. For small downloads, it is not worth + /// the effort to come up with a more sophisticated plan. + stripe_size_log: u8, +} + +impl StripePlanner2 { + /// Create a new planner with the given seed and stripe size. + pub fn new(seed: u64, stripe_size_log: u8) -> Self { + Self { + seed, + stripe_size_log, + } + } + + /// The score function to decide which peer to download from. + fn score(peer: &NodeId, seed: u64) -> u64 { + // todo: use fnv? blake3 is a bit overkill + let mut data = [0u8; 32 + 8]; + data[..32].copy_from_slice(peer.as_bytes()); + data[32..40].copy_from_slice(&seed.to_be_bytes()); + let hash = blake3::hash(&data); + u64::from_be_bytes(hash.as_bytes()[..8].try_into().unwrap()) + } +} + +impl DownloadPlanner for StripePlanner2 { + fn plan(&mut self, _hash: Hash, options: &mut BTreeMap) { + assert!( + options.values().all(|x| x.boundaries().len() % 2 == 0), + "open ranges not supported" + ); + options.retain(|_, x| !x.is_empty()); + if options.len() <= 1 { + return; + } + let ranges = get_continuous_ranges(options, self.stripe_size_log).unwrap(); + for range in ranges.windows(2) { + let start = ChunkNum(range[0]); + let end = ChunkNum(range[1]); + let curr = ChunkRanges::from(start..end); + let stripe = range[0] >> self.stripe_size_log; + let mut best_peer = None; + let mut best_score = None; + let mut matching = vec![]; + for (peer, peer_ranges) in options.iter_mut() { + if peer_ranges.contains(&start) { + matching.push((peer, peer_ranges)); + } + } + let mut peer_and_score = matching + .iter() + .map(|(peer, _)| (Self::score(peer, self.seed), peer)) + .collect::>(); + peer_and_score.sort(); + let peer_to_rank = peer_and_score + .into_iter() + .enumerate() + .map(|(i, (_, peer))| (*peer, i as u64)) + .collect::>(); + let n = matching.len() as u64; + for (peer, _) in matching.iter() { + let score = Some((peer_to_rank[*peer] + stripe) % n); + if score > best_score { + best_peer = Some(**peer); + best_score = score; + } + } + for (peer, peer_ranges) in matching { + if *peer != best_peer.unwrap() { + peer_ranges.difference_with(&curr); + } + } + } + options.retain(|_, x| !x.is_empty()); + } +} diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs new file mode 100644 index 000000000..e69de29bb From e3ff68ff07a5ecb91bd896237a745426e40ee396 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 17:24:50 +0200 Subject: [PATCH 30/47] Move things into modules --- src/downloader2.rs | 1695 +---------------------------------- src/downloader2/actor.rs | 323 +++++++ src/downloader2/planners.rs | 1 - src/downloader2/state.rs | 1361 ++++++++++++++++++++++++++++ 4 files changed, 1698 insertions(+), 1682 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index 7f69050d5..fe5361c7c 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -43,7 +43,13 @@ use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; mod planners; -pub(crate) use planners::*; +use planners::*; + +mod state; +use state::*; + +mod actor; +use actor::*; #[derive( Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, @@ -146,35 +152,6 @@ pub enum ObserveEvent { }, } -/// Global information about a peer -#[derive(Debug, Default)] -struct PeerState { - /// Executed downloads, to calculate the average download speed. - /// - /// This gets updated as soon as possible when new data has been downloaded. - download_history: VecDeque<(Duration, (u64, u64))>, -} - -/// Information about one blob on one peer -struct PeerBlobState { - /// The subscription id for the subscription - subscription_id: BitfieldSubscriptionId, - /// The number of subscriptions this peer has - subscription_count: usize, - /// chunk ranges this peer reports to have - ranges: ChunkRanges, -} - -impl PeerBlobState { - fn new(subscription_id: BitfieldSubscriptionId) -> Self { - Self { - subscription_id, - subscription_count: 1, - ranges: ChunkRanges::empty(), - } - } -} - /// A download request #[derive(Debug, Clone)] pub struct DownloadRequest { @@ -195,25 +172,6 @@ pub struct ObserveRequest { pub buffer: usize, } -struct DownloadState { - /// The request this state is for - request: DownloadRequest, - /// Ongoing downloads - peer_downloads: BTreeMap, - /// Set to true if the download needs rebalancing - needs_rebalancing: bool, -} - -impl DownloadState { - fn new(request: DownloadRequest) -> Self { - Self { - request, - peer_downloads: BTreeMap::new(), - needs_rebalancing: false, - } - } -} - #[derive(Debug)] struct IdGenerator { next_id: u64, @@ -240,171 +198,6 @@ where } } -/// Wrapper for the downloads map -/// -/// This is so we can later optimize access by fields other than id, such as hash. -#[derive(Default)] -struct Downloads { - by_id: BTreeMap, -} - -impl Downloads { - fn remove(&mut self, id: &DownloadId) -> Option { - self.by_id.remove(id) - } - - fn contains_key(&self, id: &DownloadId) -> bool { - self.by_id.contains_key(id) - } - - fn insert(&mut self, id: DownloadId, state: DownloadState) { - self.by_id.insert(id, state); - } - - fn iter_mut_for_hash( - &mut self, - hash: Hash, - ) -> impl Iterator { - self.by_id - .iter_mut() - .filter(move |x| x.1.request.hash == hash) - } - - fn iter(&mut self) -> impl Iterator { - self.by_id.iter() - } - - /// Iterate over all downloads for a given hash - fn values_for_hash(&self, hash: Hash) -> impl Iterator { - self.by_id.values().filter(move |x| x.request.hash == hash) - } - - fn values_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { - self.by_id - .values_mut() - .filter(move |x| x.request.hash == hash) - } - - fn by_id_mut(&mut self, id: DownloadId) -> Option<&mut DownloadState> { - self.by_id.get_mut(&id) - } - - fn by_peer_download_id_mut( - &mut self, - id: PeerDownloadId, - ) -> Option<(&DownloadId, &mut DownloadState)> { - self.by_id - .iter_mut() - .filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) - .next() - } -} - -#[derive(Default)] -struct Bitfields { - // Counters to generate unique ids for various requests. - // We could use uuid here, but using integers simplifies testing. - // - // the next subscription id - subscription_id_gen: IdGenerator, - by_peer_and_hash: BTreeMap<(BitfieldPeer, Hash), PeerBlobState>, -} - -impl Bitfields { - fn retain(&mut self, mut f: F) - where - F: FnMut(&(BitfieldPeer, Hash), &mut PeerBlobState) -> bool, - { - self.by_peer_and_hash.retain(|k, v| f(k, v)); - } - - fn get(&self, key: &(BitfieldPeer, Hash)) -> Option<&PeerBlobState> { - self.by_peer_and_hash.get(key) - } - - fn get_local(&self, hash: Hash) -> Option<&PeerBlobState> { - self.by_peer_and_hash.get(&(BitfieldPeer::Local, hash)) - } - - fn get_mut(&mut self, key: &(BitfieldPeer, Hash)) -> Option<&mut PeerBlobState> { - self.by_peer_and_hash.get_mut(key) - } - - fn get_local_mut(&mut self, hash: Hash) -> Option<&mut PeerBlobState> { - self.by_peer_and_hash.get_mut(&(BitfieldPeer::Local, hash)) - } - - fn insert(&mut self, key: (BitfieldPeer, Hash), value: PeerBlobState) { - self.by_peer_and_hash.insert(key, value); - } - - fn contains_key(&self, key: &(BitfieldPeer, Hash)) -> bool { - self.by_peer_and_hash.contains_key(key) - } - - fn remote_for_hash(&self, hash: Hash) -> impl Iterator { - self.by_peer_and_hash - .iter() - .filter_map(move |((peer, h), state)| { - if let BitfieldPeer::Remote(peer) = peer { - if *h == hash { - Some((peer, state)) - } else { - None - } - } else { - None - } - }) - } -} - -struct PeerDownloadState { - id: PeerDownloadId, - ranges: ChunkRanges, -} - -struct DownloaderState { - // all peers I am tracking for any download - peers: BTreeMap, - // all bitfields I am tracking, both for myself and for remote peers - // - // each item here corresponds to an active subscription - bitfields: Bitfields, - /// Observers for local bitfields - observers: Observers, - // all active downloads - // - // these are user downloads. each user download gets split into one or more - // peer downloads. - downloads: Downloads, - // discovery tasks - // - // there is a discovery task for each blob we are interested in. - discovery: BTreeMap, - // the next discovery id - discovery_id_gen: IdGenerator, - // the next peer download id - peer_download_id_gen: IdGenerator, - // the download planner - planner: Box, -} - -impl DownloaderState { - fn new(planner: Box) -> Self { - Self { - peers: Default::default(), - downloads: Default::default(), - bitfields: Default::default(), - discovery: Default::default(), - observers: Default::default(), - discovery_id_gen: Default::default(), - peer_download_id_gen: Default::default(), - planner, - } - } -} - /// Peer for a bitfield subscription #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] pub enum BitfieldPeer { @@ -414,648 +207,6 @@ pub enum BitfieldPeer { Remote(NodeId), } -#[derive(Debug)] -enum Command { - /// A user request to start a download. - StartDownload { - /// The download request - request: DownloadRequest, - /// The unique id, to be assigned by the caller - id: DownloadId, - }, - /// A user request to abort a download. - StopDownload { id: DownloadId }, - /// A full bitfield for a blob and a peer - Bitfield { - /// The peer that sent the bitfield. - peer: BitfieldPeer, - /// The blob for which the bitfield is - hash: Hash, - /// The complete bitfield - ranges: ChunkRanges, - }, - /// An update of a bitfield for a hash - /// - /// This is used both to update the bitfield of remote peers, and to update - /// the local bitfield. - BitfieldUpdate { - /// The peer that sent the update. - peer: BitfieldPeer, - /// The blob that was updated. - hash: Hash, - /// The ranges that were added - added: ChunkRanges, - /// The ranges that were removed - removed: ChunkRanges, - }, - /// A chunk was downloaded, but not yet stored - /// - /// This can only be used for updating peer stats, not for completing downloads. - ChunksDownloaded { - /// Time when the download was received - time: Duration, - /// The peer that sent the chunk - peer: NodeId, - /// The blob that was downloaded - hash: Hash, - /// The ranges that were added locally - added: ChunkRanges, - }, - /// A peer download has completed - PeerDownloadComplete { - id: PeerDownloadId, - #[allow(dead_code)] - result: anyhow::Result, - }, - /// Stop tracking a peer for all blobs, for whatever reason - #[allow(dead_code)] - DropPeer { peer: NodeId }, - /// A peer has been discovered - PeerDiscovered { peer: NodeId, hash: Hash }, - /// - ObserveLocal { - id: ObserveId, - hash: Hash, - ranges: ChunkRanges, - }, - /// - StopObserveLocal { id: ObserveId }, - /// A tick from the driver, for rebalancing - Tick { time: Duration }, -} - -#[derive(Debug, PartialEq, Eq)] -enum Event { - SubscribeBitfield { - peer: BitfieldPeer, - hash: Hash, - /// The unique id of the subscription - id: BitfieldSubscriptionId, - }, - UnsubscribeBitfield { - /// The unique id of the subscription - id: BitfieldSubscriptionId, - }, - LocalBitfield { - id: ObserveId, - ranges: ChunkRanges, - }, - LocalBitfieldUpdate { - id: ObserveId, - added: ChunkRanges, - removed: ChunkRanges, - }, - StartDiscovery { - hash: Hash, - /// The unique id of the discovery task - id: DiscoveryId, - }, - StopDiscovery { - /// The unique id of the discovery task - id: DiscoveryId, - }, - StartPeerDownload { - /// The unique id of the peer download task - id: PeerDownloadId, - peer: NodeId, - hash: Hash, - ranges: ChunkRanges, - }, - StopPeerDownload { - /// The unique id of the peer download task - id: PeerDownloadId, - }, - DownloadComplete { - /// The unique id of the user download - id: DownloadId, - }, - /// An error that stops processing the command - Error { - message: String, - }, -} - -impl DownloaderState { - /// Apply a command and return the events that were generated - fn apply(&mut self, cmd: Command) -> Vec { - let mut evs = vec![]; - self.apply_mut(cmd, &mut evs); - evs - } - - /// Apply a command, using a mutable reference to the events - fn apply_mut(&mut self, cmd: Command, evs: &mut Vec) { - if let Err(cause) = self.apply_mut_0(cmd, evs) { - evs.push(Event::Error { - message: format!("{cause}"), - }); - } - } - - /// Stop a download and clean up - /// - /// This is called both for stopping a download before completion, and for - /// cleaning up after a successful download. - /// - /// Cleanup involves emitting events for - /// - stopping all peer downloads - /// - unsubscribing from bitfields if needed - /// - stopping the discovery task if needed - fn stop_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { - let removed = self - .downloads - .remove(&id) - .context(format!("removed unknown download {id:?}"))?; - let removed_hash = removed.request.hash; - // stop associated peer downloads - for peer_download in removed.peer_downloads.values() { - evs.push(Event::StopPeerDownload { - id: peer_download.id, - }); - } - // unsubscribe from bitfields that have no more subscriptions - self.bitfields.retain(|(_peer, hash), state| { - if *hash == removed_hash { - state.subscription_count -= 1; - if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitfield { - id: state.subscription_id, - }); - return false; - } - } - true - }); - let hash_interest = self.downloads.values_for_hash(removed.request.hash).count(); - if hash_interest == 0 { - // stop the discovery task if we were the last one interested in the hash - let discovery_id = self - .discovery - .remove(&removed.request.hash) - .context(format!( - "removed unknown discovery task for {}", - removed.request.hash - ))?; - evs.push(Event::StopDiscovery { id: discovery_id }); - } - Ok(()) - } - - /// Apply a command and bail out on error - fn apply_mut_0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { - trace!("handle_command {cmd:?}"); - match cmd { - Command::StartDownload { request, id } => { - // ids must be uniquely assigned by the caller! - anyhow::ensure!( - !self.downloads.contains_key(&id), - "duplicate download request {id:?}" - ); - let hash = request.hash; - // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitfields.get_local_mut(hash) { - // just increment the count - state.subscription_count += 1; - } else { - // create a new subscription - let subscription_id = self.bitfields.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { - peer: BitfieldPeer::Local, - hash, - id: subscription_id, - }); - self.bitfields.insert( - (BitfieldPeer::Local, hash), - PeerBlobState::new(subscription_id), - ); - } - if !self.discovery.contains_key(&request.hash) { - // start a discovery task - let id = self.discovery_id_gen.next(); - evs.push(Event::StartDiscovery { hash, id }); - self.discovery.insert(request.hash, id); - } - self.downloads.insert(id, DownloadState::new(request)); - self.check_completion(hash, Some(id), evs)?; - self.start_downloads(hash, Some(id), evs)?; - } - Command::PeerDownloadComplete { id, .. } => { - let Some((download_id, download)) = self.downloads.by_peer_download_id_mut(id) - else { - // the download was already removed - return Ok(()); - }; - let download_id = *download_id; - let hash = download.request.hash; - download.peer_downloads.retain(|_, v| v.id != id); - self.start_downloads(hash, Some(download_id), evs)?; - } - Command::StopDownload { id } => { - self.stop_download(id, evs)?; - } - Command::PeerDiscovered { peer, hash } => { - if self - .bitfields - .contains_key(&(BitfieldPeer::Remote(peer), hash)) - { - // we already have a subscription for this peer - return Ok(()); - }; - // check if anybody needs this peer - if self.downloads.values_for_hash(hash).next().is_none() { - return Ok(()); - } - // create a peer state if it does not exist - let _state = self.peers.entry(peer).or_default(); - // create a new subscription - let subscription_id = self.bitfields.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { - peer: BitfieldPeer::Remote(peer), - hash, - id: subscription_id, - }); - self.bitfields.insert( - (BitfieldPeer::Remote(peer), hash), - PeerBlobState::new(subscription_id), - ); - } - Command::DropPeer { peer } => { - self.bitfields.retain(|(p, _), state| { - if *p == BitfieldPeer::Remote(peer) { - // todo: should we emit unsubscribe evs here? - evs.push(Event::UnsubscribeBitfield { - id: state.subscription_id, - }); - return false; - } else { - return true; - } - }); - self.peers.remove(&peer); - } - Command::Bitfield { peer, hash, ranges } => { - let state = self.bitfields.get_mut(&(peer, hash)).context(format!( - "bitfields for unknown peer {peer:?} and hash {hash}" - ))?; - let _chunks = total_chunks(&ranges).context("open range")?; - 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 - if let Some(observers) = self.observers.get_by_hash(&hash) { - for (id, request) in observers { - let ranges = &ranges & &request.ranges; - evs.push(Event::LocalBitfield { id: *id, ranges }); - } - } - state.ranges = ranges; - self.check_completion(hash, None, evs)?; - } else { - // We got an entirely new peer, mark all affected downloads for rebalancing - for download in self.downloads.values_mut_for_hash(hash) { - if ranges.intersects(&download.request.ranges) { - download.needs_rebalancing = true; - } - } - state.ranges = ranges; - } - // 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::BitfieldUpdate { - peer, - hash, - added, - removed, - } => { - 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 - if let Some(observers) = self.observers.get_by_hash(&hash) { - for (id, request) in observers { - let added = &added & &request.ranges; - let removed = &removed & &request.ranges; - if !added.is_empty() || !removed.is_empty() { - evs.push(Event::LocalBitfieldUpdate { - id: *id, - added: &added & &request.ranges, - removed: &removed & &request.ranges, - }); - } - } - } - state.ranges |= added; - state.ranges &= !removed; - self.check_completion(hash, None, evs)?; - } else { - // We got more data for this hash, mark all affected downloads for rebalancing - for download in self.downloads.values_mut_for_hash(hash) { - // if removed is non-empty, that is so weird that we just rebalance in any case - if !removed.is_empty() || added.intersects(&download.request.ranges) { - download.needs_rebalancing = true; - } - } - state.ranges |= added; - state.ranges &= !removed; - // a local bitfield update does not make more data available, so we don't need to start downloads - self.start_downloads(hash, None, evs)?; - } - } - Command::ChunksDownloaded { - time, - peer, - hash, - added, - } => { - let state = self.bitfields.get_local_mut(hash).context(format!( - "chunks downloaded before having local bitfield for {hash}" - ))?; - let total_downloaded = total_chunks(&added).context("open range")?; - let total_before = total_chunks(&state.ranges).context("open range")?; - state.ranges |= added; - let total_after = total_chunks(&state.ranges).context("open range")?; - let useful_downloaded = total_after - total_before; - let peer = self.peers.get_mut(&peer).context(format!( - "performing download before having peer state for {peer}" - ))?; - peer.download_history - .push_back((time, (total_downloaded, useful_downloaded))); - } - Command::Tick { time } => { - let window = 10; - let horizon = time.saturating_sub(Duration::from_secs(window)); - // clean up download history - let mut to_rebalance = vec![]; - for (peer, state) in self.peers.iter_mut() { - state - .download_history - .retain(|(duration, _)| *duration > horizon); - let mut sum_total = 0; - let mut sum_useful = 0; - for (_, (total, useful)) in state.download_history.iter() { - sum_total += total; - sum_useful += useful; - } - let speed_useful = (sum_useful as f64) / (window as f64); - let speed_total = (sum_total as f64) / (window as f64); - trace!("peer {peer} download speed {speed_total} cps total, {speed_useful} cps useful"); - } - - for (id, download) in self.downloads.iter() { - if !download.needs_rebalancing { - // nothing has changed that affects this download - continue; - } - let n_peers = self - .bitfields - .remote_for_hash(download.request.hash) - .count(); - if download.peer_downloads.len() >= n_peers { - // we are already downloading from all peers for this hash - continue; - } - to_rebalance.push(*id); - } - for id in to_rebalance { - self.rebalance_download(id, evs)?; - } - } - Command::ObserveLocal { id, hash, ranges } => { - // either we have a subscription for this blob, or we have to create one - if let Some(state) = self.bitfields.get_local_mut(hash) { - // just increment the count - state.subscription_count += 1; - // emit the current bitfield - evs.push(Event::LocalBitfield { - id, - ranges: state.ranges.clone(), - }); - } else { - // create a new subscription - let subscription_id = self.bitfields.subscription_id_gen.next(); - evs.push(Event::SubscribeBitfield { - peer: BitfieldPeer::Local, - hash, - id: subscription_id, - }); - self.bitfields.insert( - (BitfieldPeer::Local, hash), - PeerBlobState::new(subscription_id), - ); - } - self.observers.insert(id, ObserveInfo { hash, ranges }); - } - Command::StopObserveLocal { id } => { - let request = self - .observers - .remove(&id) - .context(format!("stop observing unknown local bitfield {id:?}"))?; - let removed_hash = request.hash; - // unsubscribe from bitfields that have no more subscriptions - self.bitfields.retain(|(_peer, hash), state| { - if *hash == removed_hash { - state.subscription_count -= 1; - if state.subscription_count == 0 { - evs.push(Event::UnsubscribeBitfield { - id: state.subscription_id, - }); - return false; - } - } - true - }); - } - } - Ok(()) - } - - /// Check for completion of a download or of an individual peer download - /// - /// This must be called after each change of the local bitfield for a hash - /// - /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. - fn check_completion( - &mut self, - hash: Hash, - just_id: Option, - evs: &mut Vec, - ) -> anyhow::Result<()> { - let Some(self_state) = self.bitfields.get_local(hash) else { - // 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 mut completed = vec![]; - for (id, download) in self.downloads.iter_mut_for_hash(hash) { - if just_id.is_some() && just_id != Some(*id) { - continue; - } - // 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 - evs.push(Event::DownloadComplete { id: *id }); - // remember id for later cleanup - completed.push(*id); - // no need to look at individual peer downloads in this case - continue; - } - // check if any peer download is complete, and remove it. - let mut available = vec![]; - download.peer_downloads.retain(|peer, peer_download| { - if self_state.ranges.is_superset(&peer_download.ranges) { - // stop this peer download. - // - // Might be a noop if the cause for this local change was the same peer download, but we don't know. - evs.push(Event::StopPeerDownload { - id: peer_download.id, - }); - // mark this peer as available - available.push(*peer); - false - } else { - true - } - }); - // reassign the newly available peers without doing a full rebalance - if !available.is_empty() { - // check if any of the available peers can provide something of the remaining data - let mut remaining = &download.request.ranges - &self_state.ranges; - // subtract the ranges that are already being taken care of by remaining peer downloads - for peer_download in download.peer_downloads.values() { - remaining.difference_with(&peer_download.ranges); - } - // see what the new peers can do for us - let mut candidates = BTreeMap::new(); - for peer in available { - let Some(peer_state) = self.bitfields.get(&(BitfieldPeer::Remote(peer), hash)) - else { - // weird. we should have a bitfield for this peer since it just completed a download - continue; - }; - let intersection = &peer_state.ranges & &remaining; - if !intersection.is_empty() { - candidates.insert(peer, intersection); - } - } - // deduplicate the ranges - self.planner.plan(hash, &mut candidates); - // start new downloads - for (peer, ranges) in candidates { - let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { - id, - peer, - hash, - ranges: ranges.clone(), - }); - download - .peer_downloads - .insert(peer, PeerDownloadState { id, ranges }); - } - } - } - // cleanup completed downloads, has to happen later to avoid double mutable borrow - for id in completed { - self.stop_download(id, evs)?; - } - Ok(()) - } - - /// Look at all downloads for a hash and start peer downloads for those that do not have any yet - fn start_downloads( - &mut self, - hash: Hash, - just_id: Option, - evs: &mut Vec, - ) -> anyhow::Result<()> { - let Some(self_state) = self.bitfields.get_local(hash) else { - // we don't have the self state yet, so we can't really decide if we need to download anything at all - return Ok(()); - }; - for (id, download) in self - .downloads - .iter_mut_for_hash(hash) - .filter(|(_, download)| download.peer_downloads.is_empty()) - { - if just_id.is_some() && just_id != Some(*id) { - continue; - } - let remaining = &download.request.ranges - &self_state.ranges; - let mut candidates = BTreeMap::new(); - for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { - let intersection = &bitfield.ranges & &remaining; - if !intersection.is_empty() { - candidates.insert(*peer, intersection); - } - } - self.planner.plan(hash, &mut candidates); - for (peer, ranges) in candidates { - info!(" Starting download from {peer} for {hash} {ranges:?}"); - let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { - id, - peer, - hash, - ranges: ranges.clone(), - }); - download - .peer_downloads - .insert(peer, PeerDownloadState { id, ranges }); - } - } - Ok(()) - } - - /// rebalance a single download - fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { - let download = self - .downloads - .by_id_mut(id) - .context(format!("rebalancing unknown download {id:?}"))?; - download.needs_rebalancing = false; - tracing::info!("Rebalancing download {id:?} {:?}", download.request); - let hash = download.request.hash; - let Some(self_state) = self.bitfields.get_local(hash) else { - // 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 remaining = &download.request.ranges - &self_state.ranges; - let mut candidates = BTreeMap::new(); - for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { - let intersection = &bitfield.ranges & &remaining; - if !intersection.is_empty() { - candidates.insert(*peer, intersection); - } - } - self.planner.plan(hash, &mut candidates); - info!( - "Stopping {} old peer downloads", - download.peer_downloads.len() - ); - for (_, state) in &download.peer_downloads { - // stop all downloads - evs.push(Event::StopPeerDownload { id: state.id }); - } - info!("Creating {} new peer downloads", candidates.len()); - download.peer_downloads.clear(); - for (peer, ranges) in candidates { - info!(" Starting download from {peer} for {hash} {ranges:?}"); - let id = self.peer_download_id_gen.next(); - evs.push(Event::StartPeerDownload { - id, - peer, - hash, - ranges: ranges.clone(), - }); - download - .peer_downloads - .insert(peer, PeerDownloadState { id, ranges }); - } - Ok(()) - } -} - fn total_chunks(chunks: &ChunkRanges) -> Option { let mut total = 0; for range in chunks.iter() { @@ -1210,42 +361,9 @@ impl Downloader { } } -#[derive(Debug)] -struct ObserveInfo { - hash: Hash, - ranges: ChunkRanges, -} - -#[derive(Debug, Default)] -struct Observers { - by_hash_and_id: BTreeMap>, -} - -impl Observers { - fn insert(&mut self, id: ObserveId, request: ObserveInfo) { - self.by_hash_and_id - .entry(request.hash) - .or_default() - .insert(id, request); - } - - fn remove(&mut self, id: &ObserveId) -> Option { - for requests in self.by_hash_and_id.values_mut() { - if let Some(request) = requests.remove(id) { - return Some(request); - } - } - None - } - - fn get_by_hash(&self, hash: &Hash) -> Option<&BTreeMap> { - self.by_hash_and_id.get(hash) - } -} - /// An user-facing command #[derive(Debug)] -enum UserCommand { +pub(super) enum UserCommand { Download { request: DownloadRequest, done: tokio::sync::oneshot::Sender<()>, @@ -1255,221 +373,6 @@ enum UserCommand { send: tokio::sync::mpsc::Sender, }, } - -struct DownloaderActor { - local_pool: LocalPool, - endpoint: Endpoint, - command_rx: mpsc::Receiver, - command_tx: mpsc::Sender, - state: DownloaderState, - store: S, - discovery: BoxedContentDiscovery, - subscribe_bitfield: BoxedBitfieldSubscription, - download_futs: BTreeMap>, - peer_download_tasks: BTreeMap>, - discovery_tasks: BTreeMap>, - bitfield_subscription_tasks: BTreeMap>, - /// Id generator for download ids - download_id_gen: IdGenerator, - /// Id generator for observe ids - observe_id_gen: IdGenerator, - /// Observers - observers: BTreeMap>, - /// The time when the actor was started, serves as the epoch for time messages to the state machine - start: Instant, -} - -impl DownloaderActor { - fn new( - endpoint: Endpoint, - store: S, - discovery: BoxedContentDiscovery, - subscribe_bitfield: BoxedBitfieldSubscription, - local_pool: LocalPool, - planner: Box, - ) -> Self { - let (send, recv) = mpsc::channel(256); - Self { - local_pool, - endpoint, - state: DownloaderState::new(planner), - store, - discovery, - subscribe_bitfield, - peer_download_tasks: BTreeMap::new(), - discovery_tasks: BTreeMap::new(), - bitfield_subscription_tasks: BTreeMap::new(), - download_futs: BTreeMap::new(), - command_tx: send, - command_rx: recv, - download_id_gen: Default::default(), - observe_id_gen: Default::default(), - observers: Default::default(), - start: Instant::now(), - } - } - - async fn run(mut self, mut channel: mpsc::Receiver) { - let mut ticks = tokio::time::interval(Duration::from_millis(100)); - loop { - trace!("downloader actor tick"); - tokio::select! { - biased; - Some(cmd) = channel.recv() => { - debug!("user command {cmd:?}"); - match cmd { - UserCommand::Download { - request, done, - } => { - let id = self.download_id_gen.next(); - self.download_futs.insert(id, done); - self.command_tx.send(Command::StartDownload { request, id }).await.ok(); - } - UserCommand::Observe { request, send } => { - let id = self.observe_id_gen.next(); - self.command_tx.send(Command::ObserveLocal { id, hash: request.hash, ranges: request.ranges }).await.ok(); - self.observers.insert(id, send); - } - } - }, - Some(cmd) = self.command_rx.recv() => { - let evs = self.state.apply(cmd); - for ev in evs { - self.handle_event(ev); - } - }, - _ = ticks.tick() => { - let time = self.start.elapsed(); - self.command_tx.send(Command::Tick { time }).await.ok(); - // clean up dropped futures - // - // todo: is there a better mechanism than periodic checks? - // I don't want some cancellation token rube goldberg machine. - let mut to_delete = vec![]; - for (id, fut) in self.download_futs.iter() { - if fut.is_closed() { - to_delete.push(*id); - self.command_tx.send(Command::StopDownload { id: *id }).await.ok(); - } - } - for id in to_delete { - self.download_futs.remove(&id); - } - // clean up dropped observers - let mut to_delete = vec![]; - for (id, sender) in self.observers.iter() { - if sender.is_closed() { - to_delete.push(*id); - self.command_tx.send(Command::StopObserveLocal { id: *id }).await.ok(); - } - } - for id in to_delete { - self.observers.remove(&id); - } - }, - } - } - } - - fn handle_event(&mut self, ev: Event) { - trace!("handle_event {ev:?}"); - match ev { - Event::SubscribeBitfield { peer, hash, id } => { - 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 { - BitfieldSubscriptionEvent::Bitfield { ranges } => { - Command::Bitfield { peer, hash, ranges } - } - BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { - Command::BitfieldUpdate { - peer, - hash, - added, - removed, - } - } - }; - send.send(cmd).await.ok(); - } - }); - self.bitfield_subscription_tasks.insert(id, task); - } - Event::StartDiscovery { hash, id } => { - let send = self.command_tx.clone(); - let mut stream = self.discovery.find_peers(hash, Default::default()); - let task = spawn(async move { - // process the infinite discovery stream and send commands - while let Some(peer) = stream.next().await { - println!("peer discovered for hash {hash}: {peer}"); - let res = send.send(Command::PeerDiscovered { peer, hash }).await; - if res.is_err() { - // only reason for this is actor task dropped - break; - } - } - }); - self.discovery_tasks.insert(id, task); - } - Event::StartPeerDownload { - id, - peer, - hash, - ranges, - } => { - let send = self.command_tx.clone(); - let endpoint = self.endpoint.clone(); - let store = self.store.clone(); - let start = self.start; - let task = self.local_pool.spawn(move || { - peer_download_task(id, endpoint, store, hash, peer, ranges, send, start) - }); - self.peer_download_tasks.insert(id, task); - } - Event::UnsubscribeBitfield { id } => { - self.bitfield_subscription_tasks.remove(&id); - } - Event::StopDiscovery { id } => { - self.discovery_tasks.remove(&id); - } - Event::StopPeerDownload { id } => { - self.peer_download_tasks.remove(&id); - } - Event::DownloadComplete { id } => { - if let Some(done) = self.download_futs.remove(&id) { - done.send(()).ok(); - } - } - Event::LocalBitfield { id, ranges } => { - let Some(sender) = self.observers.get(&id) else { - return; - }; - if sender.try_send(ObserveEvent::Bitfield { ranges }).is_err() { - // the observer has been dropped - self.observers.remove(&id); - } - } - Event::LocalBitfieldUpdate { id, added, removed } => { - let Some(sender) = self.observers.get(&id) else { - return; - }; - if sender - .try_send(ObserveEvent::BitfieldUpdate { added, removed }) - .is_err() - { - // the observer has been dropped - self.observers.remove(&id); - } - } - Event::Error { message } => { - error!("Error during processing event {}", message); - } - } - } -} - /// A simple static content discovery mechanism #[derive(Debug)] pub struct StaticContentDiscovery { @@ -1497,114 +400,6 @@ impl ContentDiscovery for StaticContentDiscovery { } } -async fn peer_download_task( - id: PeerDownloadId, - endpoint: Endpoint, - store: S, - hash: Hash, - peer: NodeId, - ranges: ChunkRanges, - sender: mpsc::Sender, - start: Instant, -) { - let result = peer_download(endpoint, store, hash, peer, ranges, &sender, start).await; - sender - .send(Command::PeerDownloadComplete { id, result }) - .await - .ok(); -} - -async fn peer_download( - endpoint: Endpoint, - store: S, - hash: Hash, - peer: NodeId, - ranges: ChunkRanges, - sender: &mpsc::Sender, - start: Instant, -) -> anyhow::Result { - info!("Connecting to peer {peer}"); - let conn = endpoint.connect(peer, crate::ALPN).await?; - info!("Got connection to peer {peer}"); - let spec = RangeSpec::new(ranges); - let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); - info!("starting download from {peer} for {hash} {ranges:?}"); - let request = GetRequest::new(hash, ranges); - let initial = crate::get::fsm::start(conn, request); - // connect - let connected = initial.next().await?; - // read the first bytes - let ConnectedNext::StartRoot(start_root) = connected.next().await? else { - return Err(io::Error::new(io::ErrorKind::Other, "expected start root").into()); - }; - let header = start_root.next(); - - // get the size of the content - let (mut content, size) = header.next().await?; - let entry = store.get_or_create(hash, size).await?; - let mut writer = entry.batch_writer().await?; - let mut batch = Vec::new(); - // manually loop over the content and yield all data - let done = loop { - match content.next().await { - BlobContentNext::More((next, data)) => { - match data? { - BaoContentItem::Parent(parent) => { - batch.push(parent.into()); - } - BaoContentItem::Leaf(leaf) => { - let start_chunk = leaf.offset / 1024; - let added = - ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16)); - sender - .send(Command::ChunksDownloaded { - time: start.elapsed(), - peer, - hash, - added: added.clone(), - }) - .await - .ok(); - batch.push(leaf.into()); - writer.write_batch(size, std::mem::take(&mut batch)).await?; - sender - .send(Command::BitfieldUpdate { - peer: BitfieldPeer::Local, - hash, - added, - removed: ChunkRanges::empty(), - }) - .await - .ok(); - } - } - content = next; - } - BlobContentNext::Done(done) => { - // we are done with the root blob - break done; - } - } - }; - // close the connection even if there is more data - let closing = match done.next() { - EndBlobNext::Closing(closing) => closing, - EndBlobNext::MoreChildren(more) => more.finish(), - }; - // close the connection - let stats = closing.next().await?; - Ok(stats) -} - -fn spawn(f: F) -> AbortOnDropHandle -where - F: Future + Send + 'static, - T: Send + 'static, -{ - let task = tokio::spawn(f); - AbortOnDropHandle::new(task) -} - /// A bitfield subscription that just returns nothing for local and everything(*) for remote /// /// * Still need to figure out how to deal with open ended chunk ranges. @@ -1848,7 +643,7 @@ mod tests { } /// Create chunk ranges from an array of u64 ranges - fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { + pub fn chunk_ranges(ranges: impl IntoIterator>) -> ChunkRanges { let mut res = ChunkRanges::empty(); for range in ranges.into_iter() { res |= ChunkRanges::from(ChunkNum(range.start)..ChunkNum(range.end)); @@ -1856,485 +651,23 @@ mod tests { res } - fn noop_planner() -> BoxedDownloadPlanner { + pub fn noop_planner() -> BoxedDownloadPlanner { Box::new(NoopPlanner) } /// Checks if an exact event is present exactly once in a list of events - fn has_one_event(evs: &[Event], ev: &Event) -> bool { + pub fn has_one_event(evs: &[Event], ev: &Event) -> bool { evs.iter().filter(|e| *e == ev).count() == 1 } - fn has_all_events(evs: &[Event], evs2: &[&Event]) -> bool { + pub fn has_all_events(evs: &[Event], evs2: &[&Event]) -> bool { evs2.iter().all(|ev| has_one_event(evs, ev)) } - fn has_one_event_matching(evs: &[Event], f: impl Fn(&Event) -> bool) -> bool { + pub fn has_one_event_matching(evs: &[Event], f: impl Fn(&Event) -> bool) -> bool { evs.iter().filter(|e| f(e)).count() == 1 } - /// Test various things that should produce errors - #[test] - fn downloader_state_errors() -> TestResult<()> { - use BitfieldPeer::*; - let _ = tracing_subscriber::fmt::try_init(); - let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let unknown_hash = - "0000000000000000000000000000000000000000000000000000000000000002".parse()?; - let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: ChunkRanges::all(), - }); - assert!( - has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), - "adding an open bitfield should produce an error!" - ); - let evs = state.apply(Command::Bitfield { - peer: Local, - hash: unknown_hash, - ranges: ChunkRanges::all(), - }); - assert!( - has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), - "adding an open bitfield for an unknown hash should produce an error!" - ); - let evs = state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([0..16]), - }); - assert!( - has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), - "download from unknown peer should lead to an error!" - ); - Ok(()) - } - - /// Test a simple scenario where a download is started and completed - #[test] - fn downloader_state_smoke() -> TestResult<()> { - use BitfieldPeer::*; - let _ = tracing_subscriber::fmt::try_init(); - let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::StartDownload { - request: DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }, - id: DownloadId(0), - }); - assert!( - has_one_event( - &evs, - &Event::StartDiscovery { - hash, - id: DiscoveryId(0) - } - ), - "starting a download should start a discovery task" - ); - assert!( - has_one_event( - &evs, - &Event::SubscribeBitfield { - peer: Local, - hash, - id: BitfieldSubscriptionId(0) - } - ), - "starting a download should subscribe to the local bitfield" - ); - let initial_bitfield = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: initial_bitfield.clone(), - }); - assert!(evs.is_empty()); - assert_eq!( - state - .bitfields - .get_local(hash) - .context("bitfield should be present")? - .ranges, - initial_bitfield, - "bitfield should be set to the initial bitfield" - ); - assert_eq!( - state - .bitfields - .get_local(hash) - .context("bitfield should be present")? - .subscription_count, - 1, - "we have one download interested in the bitfield" - ); - let evs = state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([16..32]), - removed: ChunkRanges::empty(), - }); - assert!(evs.is_empty()); - assert_eq!( - state - .bitfields - .get_local(hash) - .context("bitfield should be present")? - .ranges, - ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), - "bitfield should be updated" - ); - let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - assert!( - has_one_event( - &evs, - &Event::SubscribeBitfield { - peer: Remote(peer_a), - hash, - id: 1.into() - } - ), - "adding a new peer for a hash we are interested in should subscribe to the bitfield" - ); - let evs = state.apply(Command::Bitfield { - peer: Remote(peer_a), - hash, - ranges: chunk_ranges([0..64]), - }); - assert!( - has_one_event( - &evs, - &Event::StartPeerDownload { - id: PeerDownloadId(0), - peer: peer_a, - hash, - ranges: chunk_ranges([32..64]) - } - ), - "bitfield from a peer should start a download" - ); - // ChunksDownloaded just updates the peer stats - let evs = state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([32..48]), - }); - assert!(evs.is_empty()); - // Bitfield update does not yet complete the download - let evs = state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([32..48]), - removed: ChunkRanges::empty(), - }); - assert!(evs.is_empty()); - // ChunksDownloaded just updates the peer stats - let evs = state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([48..64]), - }); - assert!(evs.is_empty()); - // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([48..64]), - removed: ChunkRanges::empty(), - }); - assert!( - has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), - "download should be completed by the data" - ); - // quick check that everything got cleaned up - assert!(state.downloads.by_id.is_empty()); - assert!(state.bitfields.by_peer_and_hash.is_empty()); - assert!(state.discovery.is_empty()); - Ok(()) - } - - /// Test a scenario where more data becomes available at the remote peer as the download progresses - #[test] - fn downloader_state_incremental() -> TestResult<()> { - use BitfieldPeer::*; - let _ = tracing_subscriber::fmt::try_init(); - let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let mut state = DownloaderState::new(noop_planner()); - // Start a download - state.apply(Command::StartDownload { - request: DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }, - id: DownloadId(0), - }); - // Initially, we have nothing - state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: ChunkRanges::empty(), - }); - // We have a peer for the hash - state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { - peer: Remote(peer_a), - hash, - ranges: chunk_ranges([0..32]), - }); - assert!( - has_one_event( - &evs, - &Event::StartPeerDownload { - id: 0.into(), - peer: peer_a, - hash, - ranges: chunk_ranges([0..32]) - } - ), - "bitfield from a peer should start a download" - ); - // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([0..16]), - }); - // Bitfield update does not yet complete the download - state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([0..16]), - removed: ChunkRanges::empty(), - }); - // The peer now has more data - state.apply(Command::Bitfield { - peer: Remote(peer_a), - hash, - ranges: chunk_ranges([32..64]), - }); - // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([16..32]), - }); - // Complete the first part of the download - let evs = state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([16..32]), - removed: ChunkRanges::empty(), - }); - // This triggers cancellation of the first peer download and starting a new one for the remaining data - assert!( - has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), - "first peer download should be stopped" - ); - assert!( - has_one_event( - &evs, - &Event::StartPeerDownload { - id: 1.into(), - peer: peer_a, - hash, - ranges: chunk_ranges([32..64]) - } - ), - "second peer download should be started" - ); - // ChunksDownloaded just updates the peer stats - state.apply(Command::ChunksDownloaded { - time: Duration::ZERO, - peer: peer_a, - hash, - added: chunk_ranges([32..64]), - }); - // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { - peer: Local, - hash, - added: chunk_ranges([32..64]), - removed: ChunkRanges::empty(), - }); - assert!( - has_all_events( - &evs, - &[ - &Event::StopPeerDownload { id: 1.into() }, - &Event::DownloadComplete { id: 0.into() }, - &Event::UnsubscribeBitfield { id: 0.into() }, - &Event::StopDiscovery { id: 0.into() }, - ] - ), - "download should be completed by the data" - ); - println!("{evs:?}"); - Ok(()) - } - - #[test] - fn downloader_state_multiple_downloads() -> testresult::TestResult<()> { - use BitfieldPeer::*; - // Use a constant hash (the same style as used in other tests). - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - // Create a downloader state with a no‐op planner. - let mut state = DownloaderState::new(noop_planner()); - - // --- Start the first (ongoing) download. - // Request a range from 0..64. - let download0 = DownloadId(0); - let req0 = DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }; - let evs0 = state.apply(Command::StartDownload { - request: req0, - id: download0, - }); - // When starting the download, we expect a discovery task to be started - // and a subscription to the local bitfield to be requested. - assert!( - has_one_event( - &evs0, - &Event::StartDiscovery { - hash, - id: DiscoveryId(0) - } - ), - "download0 should start discovery" - ); - assert!( - has_one_event( - &evs0, - &Event::SubscribeBitfield { - peer: Local, - hash, - id: BitfieldSubscriptionId(0) - } - ), - "download0 should subscribe to the local bitfield" - ); - - // --- Simulate some progress for the first download. - // Let’s say only chunks 0..32 are available locally. - let evs1 = state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: chunk_ranges([0..32]), - }); - // No completion event should be generated for download0 because its full range 0..64 is not yet met. - assert!( - evs1.is_empty(), - "Partial bitfield update should not complete download0" - ); - - // --- Start a second download for the same hash. - // This new download only requires chunks 0..32 which are already available. - let download1 = DownloadId(1); - let req1 = DownloadRequest { - hash, - ranges: chunk_ranges([0..32]), - }; - let evs2 = state.apply(Command::StartDownload { - request: req1, - id: download1, - }); - // Because the local bitfield (0..32) is already a superset of the new download’s request, - // a DownloadComplete event for download1 should be generated immediately. - assert!( - has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), - "New download should complete immediately" - ); - - // --- Verify state: - // The ongoing download (download0) should still be present in the state, - // while the newly completed download (download1) is removed. - assert!( - state.downloads.contains_key(&download0), - "download0 should still be active" - ); - assert!( - !state.downloads.contains_key(&download1), - "download1 should have been cleaned up after completion" - ); - - Ok(()) - } - - /// Test a scenario where more data becomes available at the remote peer as the download progresses - #[test] - fn downloader_state_drop() -> TestResult<()> { - use BitfieldPeer::*; - let _ = tracing_subscriber::fmt::try_init(); - let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; - let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; - let mut state = DownloaderState::new(noop_planner()); - // Start a download - state.apply(Command::StartDownload { - request: DownloadRequest { - hash, - ranges: chunk_ranges([0..64]), - }, - id: 0.into(), - }); - // Initially, we have nothing - state.apply(Command::Bitfield { - peer: Local, - hash, - ranges: ChunkRanges::empty(), - }); - // We have a peer for the hash - state.apply(Command::PeerDiscovered { peer: peer_a, hash }); - // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { - peer: Remote(peer_a), - hash, - ranges: chunk_ranges([0..32]), - }); - assert!( - has_one_event( - &evs, - &Event::StartPeerDownload { - id: 0.into(), - peer: peer_a, - hash, - ranges: chunk_ranges([0..32]) - } - ), - "bitfield from a peer should start a download" - ); - // Sending StopDownload should stop the download and all associated tasks - // This is what happens (delayed) when the user drops the download future - let evs = state.apply(Command::StopDownload { id: 0.into() }); - assert!(has_one_event( - &evs, - &Event::StopPeerDownload { id: 0.into() } - )); - assert!(has_one_event( - &evs, - &Event::UnsubscribeBitfield { id: 0.into() } - )); - assert!(has_one_event( - &evs, - &Event::UnsubscribeBitfield { id: 1.into() } - )); - assert!(has_one_event(&evs, &Event::StopDiscovery { id: 0.into() })); - Ok(()) - } - #[tokio::test] #[cfg(feature = "rpc")] async fn downloader_driver_smoke() -> TestResult<()> { @@ -2355,7 +688,7 @@ mod tests { .discovery(discovery) .bitfield_subscription(bitfield_subscription) .build(); - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(1)).await; let fut = downloader.download(DownloadRequest { hash, ranges: chunk_ranges([0..1]), diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index e69de29bb..f0b08fb26 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -0,0 +1,323 @@ +use super::*; + +pub(super) struct DownloaderActor { + local_pool: LocalPool, + endpoint: Endpoint, + command_rx: mpsc::Receiver, + command_tx: mpsc::Sender, + state: DownloaderState, + store: S, + discovery: BoxedContentDiscovery, + subscribe_bitfield: BoxedBitfieldSubscription, + download_futs: BTreeMap>, + peer_download_tasks: BTreeMap>, + discovery_tasks: BTreeMap>, + bitfield_subscription_tasks: BTreeMap>, + /// Id generator for download ids + download_id_gen: IdGenerator, + /// Id generator for observe ids + observe_id_gen: IdGenerator, + /// Observers + observers: BTreeMap>, + /// The time when the actor was started, serves as the epoch for time messages to the state machine + start: Instant, +} + +impl DownloaderActor { + pub(super) fn new( + endpoint: Endpoint, + store: S, + discovery: BoxedContentDiscovery, + subscribe_bitfield: BoxedBitfieldSubscription, + local_pool: LocalPool, + planner: Box, + ) -> Self { + let (send, recv) = mpsc::channel(256); + Self { + local_pool, + endpoint, + state: DownloaderState::new(planner), + store, + discovery, + subscribe_bitfield, + peer_download_tasks: BTreeMap::new(), + discovery_tasks: BTreeMap::new(), + bitfield_subscription_tasks: BTreeMap::new(), + download_futs: BTreeMap::new(), + command_tx: send, + command_rx: recv, + download_id_gen: Default::default(), + observe_id_gen: Default::default(), + observers: Default::default(), + start: Instant::now(), + } + } + + pub(super) async fn run(mut self, mut channel: mpsc::Receiver) { + let mut ticks = tokio::time::interval(Duration::from_millis(100)); + loop { + trace!("downloader actor tick"); + tokio::select! { + biased; + Some(cmd) = channel.recv() => { + debug!("user command {cmd:?}"); + match cmd { + UserCommand::Download { + request, done, + } => { + let id = self.download_id_gen.next(); + self.download_futs.insert(id, done); + self.command_tx.send(Command::StartDownload { request, id }).await.ok(); + } + UserCommand::Observe { request, send } => { + let id = self.observe_id_gen.next(); + self.command_tx.send(Command::ObserveLocal { id, hash: request.hash, ranges: request.ranges }).await.ok(); + self.observers.insert(id, send); + } + } + }, + Some(cmd) = self.command_rx.recv() => { + let evs = self.state.apply(cmd); + for ev in evs { + self.handle_event(ev); + } + }, + _ = ticks.tick() => { + let time = self.start.elapsed(); + self.command_tx.send(Command::Tick { time }).await.ok(); + // clean up dropped futures + // + // todo: is there a better mechanism than periodic checks? + // I don't want some cancellation token rube goldberg machine. + let mut to_delete = vec![]; + for (id, fut) in self.download_futs.iter() { + if fut.is_closed() { + to_delete.push(*id); + self.command_tx.send(Command::StopDownload { id: *id }).await.ok(); + } + } + for id in to_delete { + self.download_futs.remove(&id); + } + // clean up dropped observers + let mut to_delete = vec![]; + for (id, sender) in self.observers.iter() { + if sender.is_closed() { + to_delete.push(*id); + self.command_tx.send(Command::StopObserveLocal { id: *id }).await.ok(); + } + } + for id in to_delete { + self.observers.remove(&id); + } + }, + } + } + } + + fn handle_event(&mut self, ev: Event) { + trace!("handle_event {ev:?}"); + match ev { + Event::SubscribeBitfield { peer, hash, id } => { + 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 { + BitfieldSubscriptionEvent::Bitfield { ranges } => { + Command::Bitfield { peer, hash, ranges } + } + BitfieldSubscriptionEvent::BitfieldUpdate { added, removed } => { + Command::BitfieldUpdate { + peer, + hash, + added, + removed, + } + } + }; + send.send(cmd).await.ok(); + } + }); + self.bitfield_subscription_tasks.insert(id, task); + } + Event::StartDiscovery { hash, id } => { + let send = self.command_tx.clone(); + let mut stream = self.discovery.find_peers(hash, Default::default()); + let task = spawn(async move { + // process the infinite discovery stream and send commands + while let Some(peer) = stream.next().await { + println!("peer discovered for hash {hash}: {peer}"); + let res = send.send(Command::PeerDiscovered { peer, hash }).await; + if res.is_err() { + // only reason for this is actor task dropped + break; + } + } + }); + self.discovery_tasks.insert(id, task); + } + Event::StartPeerDownload { + id, + peer, + hash, + ranges, + } => { + let send = self.command_tx.clone(); + let endpoint = self.endpoint.clone(); + let store = self.store.clone(); + let start = self.start; + let task = self.local_pool.spawn(move || { + peer_download_task(id, endpoint, store, hash, peer, ranges, send, start) + }); + self.peer_download_tasks.insert(id, task); + } + Event::UnsubscribeBitfield { id } => { + self.bitfield_subscription_tasks.remove(&id); + } + Event::StopDiscovery { id } => { + self.discovery_tasks.remove(&id); + } + Event::StopPeerDownload { id } => { + self.peer_download_tasks.remove(&id); + } + Event::DownloadComplete { id } => { + if let Some(done) = self.download_futs.remove(&id) { + done.send(()).ok(); + } + } + Event::LocalBitfield { id, ranges } => { + let Some(sender) = self.observers.get(&id) else { + return; + }; + if sender.try_send(ObserveEvent::Bitfield { ranges }).is_err() { + // the observer has been dropped + self.observers.remove(&id); + } + } + Event::LocalBitfieldUpdate { id, added, removed } => { + let Some(sender) = self.observers.get(&id) else { + return; + }; + if sender + .try_send(ObserveEvent::BitfieldUpdate { added, removed }) + .is_err() + { + // the observer has been dropped + self.observers.remove(&id); + } + } + Event::Error { message } => { + error!("Error during processing event {}", message); + } + } + } +} + +async fn peer_download_task( + id: PeerDownloadId, + endpoint: Endpoint, + store: S, + hash: Hash, + peer: NodeId, + ranges: ChunkRanges, + sender: mpsc::Sender, + start: Instant, +) { + let result = peer_download(endpoint, store, hash, peer, ranges, &sender, start).await; + sender + .send(Command::PeerDownloadComplete { id, result }) + .await + .ok(); +} + +async fn peer_download( + endpoint: Endpoint, + store: S, + hash: Hash, + peer: NodeId, + ranges: ChunkRanges, + sender: &mpsc::Sender, + start: Instant, +) -> anyhow::Result { + info!("Connecting to peer {peer}"); + let conn = endpoint.connect(peer, crate::ALPN).await?; + info!("Got connection to peer {peer}"); + let spec = RangeSpec::new(ranges); + let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); + info!("starting download from {peer} for {hash} {ranges:?}"); + let request = GetRequest::new(hash, ranges); + let initial = crate::get::fsm::start(conn, request); + // connect + let connected = initial.next().await?; + // read the first bytes + let ConnectedNext::StartRoot(start_root) = connected.next().await? else { + return Err(io::Error::new(io::ErrorKind::Other, "expected start root").into()); + }; + let header = start_root.next(); + + // get the size of the content + let (mut content, size) = header.next().await?; + let entry = store.get_or_create(hash, size).await?; + let mut writer = entry.batch_writer().await?; + let mut batch = Vec::new(); + // manually loop over the content and yield all data + let done = loop { + match content.next().await { + BlobContentNext::More((next, data)) => { + match data? { + BaoContentItem::Parent(parent) => { + batch.push(parent.into()); + } + BaoContentItem::Leaf(leaf) => { + let start_chunk = leaf.offset / 1024; + let added = + ChunkRanges::from(ChunkNum(start_chunk)..ChunkNum(start_chunk + 16)); + sender + .send(Command::ChunksDownloaded { + time: start.elapsed(), + peer, + hash, + added: added.clone(), + }) + .await + .ok(); + batch.push(leaf.into()); + writer.write_batch(size, std::mem::take(&mut batch)).await?; + sender + .send(Command::BitfieldUpdate { + peer: BitfieldPeer::Local, + hash, + added, + removed: ChunkRanges::empty(), + }) + .await + .ok(); + } + } + content = next; + } + BlobContentNext::Done(done) => { + // we are done with the root blob + break done; + } + } + }; + // close the connection even if there is more data + let closing = match done.next() { + EndBlobNext::Closing(closing) => closing, + EndBlobNext::MoreChildren(more) => more.finish(), + }; + // close the connection + let stats = closing.next().await?; + Ok(stats) +} + +pub(super) fn spawn(f: F) -> AbortOnDropHandle +where + F: Future + Send + 'static, + T: Send + 'static, +{ + let task = tokio::spawn(f); + AbortOnDropHandle::new(task) +} diff --git a/src/downloader2/planners.rs b/src/downloader2/planners.rs index 5d87ad73c..6825fad6d 100644 --- a/src/downloader2/planners.rs +++ b/src/downloader2/planners.rs @@ -22,7 +22,6 @@ pub trait DownloadPlanner: Send + std::fmt::Debug + 'static { /// A boxed download planner pub type BoxedDownloadPlanner = Box; - /// A download planner that just leaves everything as is. /// /// Data will be downloaded from all peers wherever multiple peers have the same data. diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index e69de29bb..eea84b12d 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -0,0 +1,1361 @@ +//! The state machine for the downloader +use super::*; + +#[derive(Debug)] +pub(super) enum Command { + /// A user request to start a download. + StartDownload { + /// The download request + request: DownloadRequest, + /// The unique id, to be assigned by the caller + id: DownloadId, + }, + /// A user request to abort a download. + StopDownload { id: DownloadId }, + /// A full bitfield for a blob and a peer + Bitfield { + /// The peer that sent the bitfield. + peer: BitfieldPeer, + /// The blob for which the bitfield is + hash: Hash, + /// The complete bitfield + ranges: ChunkRanges, + }, + /// An update of a bitfield for a hash + /// + /// This is used both to update the bitfield of remote peers, and to update + /// the local bitfield. + BitfieldUpdate { + /// The peer that sent the update. + peer: BitfieldPeer, + /// The blob that was updated. + hash: Hash, + /// The ranges that were added + added: ChunkRanges, + /// The ranges that were removed + removed: ChunkRanges, + }, + /// A chunk was downloaded, but not yet stored + /// + /// This can only be used for updating peer stats, not for completing downloads. + ChunksDownloaded { + /// Time when the download was received + time: Duration, + /// The peer that sent the chunk + peer: NodeId, + /// The blob that was downloaded + hash: Hash, + /// The ranges that were added locally + added: ChunkRanges, + }, + /// A peer download has completed + PeerDownloadComplete { + id: PeerDownloadId, + #[allow(dead_code)] + result: anyhow::Result, + }, + /// Stop tracking a peer for all blobs, for whatever reason + #[allow(dead_code)] + DropPeer { peer: NodeId }, + /// A peer has been discovered + PeerDiscovered { peer: NodeId, hash: Hash }, + /// + ObserveLocal { + id: ObserveId, + hash: Hash, + ranges: ChunkRanges, + }, + /// + StopObserveLocal { id: ObserveId }, + /// A tick from the driver, for rebalancing + Tick { time: Duration }, +} + +#[derive(Debug, PartialEq, Eq)] +pub(super) enum Event { + SubscribeBitfield { + peer: BitfieldPeer, + hash: Hash, + /// The unique id of the subscription + id: BitfieldSubscriptionId, + }, + UnsubscribeBitfield { + /// The unique id of the subscription + id: BitfieldSubscriptionId, + }, + LocalBitfield { + id: ObserveId, + ranges: ChunkRanges, + }, + LocalBitfieldUpdate { + id: ObserveId, + added: ChunkRanges, + removed: ChunkRanges, + }, + StartDiscovery { + hash: Hash, + /// The unique id of the discovery task + id: DiscoveryId, + }, + StopDiscovery { + /// The unique id of the discovery task + id: DiscoveryId, + }, + StartPeerDownload { + /// The unique id of the peer download task + id: PeerDownloadId, + peer: NodeId, + hash: Hash, + ranges: ChunkRanges, + }, + StopPeerDownload { + /// The unique id of the peer download task + id: PeerDownloadId, + }, + DownloadComplete { + /// The unique id of the user download + id: DownloadId, + }, + /// An error that stops processing the command + Error { + message: String, + }, +} + +pub struct DownloaderState { + // all peers I am tracking for any download + peers: BTreeMap, + // all bitfields I am tracking, both for myself and for remote peers + // + // each item here corresponds to an active subscription + bitfields: Bitfields, + /// Observers for local bitfields + observers: Observers, + // all active downloads + // + // these are user downloads. each user download gets split into one or more + // peer downloads. + downloads: Downloads, + // discovery tasks + // + // there is a discovery task for each blob we are interested in. + discovery: BTreeMap, + // the next discovery id + discovery_id_gen: IdGenerator, + // the next peer download id + peer_download_id_gen: IdGenerator, + // the download planner + planner: Box, +} + +impl DownloaderState { + pub fn new(planner: Box) -> Self { + Self { + peers: Default::default(), + downloads: Default::default(), + bitfields: Default::default(), + discovery: Default::default(), + observers: Default::default(), + discovery_id_gen: Default::default(), + peer_download_id_gen: Default::default(), + planner, + } + } +} + +impl DownloaderState { + /// Apply a command and return the events that were generated + pub(super) fn apply(&mut self, cmd: Command) -> Vec { + let mut evs = vec![]; + self.apply_mut(cmd, &mut evs); + evs + } + + /// Apply a command, using a mutable reference to the events + fn apply_mut(&mut self, cmd: Command, evs: &mut Vec) { + if let Err(cause) = self.apply_mut_0(cmd, evs) { + evs.push(Event::Error { + message: format!("{cause}"), + }); + } + } + + /// Stop a download and clean up + /// + /// This is called both for stopping a download before completion, and for + /// cleaning up after a successful download. + /// + /// Cleanup involves emitting events for + /// - stopping all peer downloads + /// - unsubscribing from bitfields if needed + /// - stopping the discovery task if needed + fn stop_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { + let removed = self + .downloads + .remove(&id) + .context(format!("removed unknown download {id:?}"))?; + let removed_hash = removed.request.hash; + // stop associated peer downloads + for peer_download in removed.peer_downloads.values() { + evs.push(Event::StopPeerDownload { + id: peer_download.id, + }); + } + // unsubscribe from bitfields that have no more subscriptions + self.bitfields.retain(|(_peer, hash), state| { + if *hash == removed_hash { + state.subscription_count -= 1; + if state.subscription_count == 0 { + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); + return false; + } + } + true + }); + let hash_interest = self.downloads.values_for_hash(removed.request.hash).count(); + if hash_interest == 0 { + // stop the discovery task if we were the last one interested in the hash + let discovery_id = self + .discovery + .remove(&removed.request.hash) + .context(format!( + "removed unknown discovery task for {}", + removed.request.hash + ))?; + evs.push(Event::StopDiscovery { id: discovery_id }); + } + Ok(()) + } + + /// Apply a command and bail out on error + fn apply_mut_0(&mut self, cmd: Command, evs: &mut Vec) -> anyhow::Result<()> { + trace!("handle_command {cmd:?}"); + match cmd { + Command::StartDownload { request, id } => { + // ids must be uniquely assigned by the caller! + anyhow::ensure!( + !self.downloads.contains_key(&id), + "duplicate download request {id:?}" + ); + let hash = request.hash; + // either we have a subscription for this blob, or we have to create one + if let Some(state) = self.bitfields.get_local_mut(hash) { + // just increment the count + state.subscription_count += 1; + } else { + // create a new subscription + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Local, + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Local, hash), + PeerBlobState::new(subscription_id), + ); + } + if !self.discovery.contains_key(&request.hash) { + // start a discovery task + let id = self.discovery_id_gen.next(); + evs.push(Event::StartDiscovery { hash, id }); + self.discovery.insert(request.hash, id); + } + self.downloads.insert(id, DownloadState::new(request)); + self.check_completion(hash, Some(id), evs)?; + self.start_downloads(hash, Some(id), evs)?; + } + Command::PeerDownloadComplete { id, .. } => { + let Some((download_id, download)) = self.downloads.by_peer_download_id_mut(id) + else { + // the download was already removed + return Ok(()); + }; + let download_id = *download_id; + let hash = download.request.hash; + download.peer_downloads.retain(|_, v| v.id != id); + self.start_downloads(hash, Some(download_id), evs)?; + } + Command::StopDownload { id } => { + self.stop_download(id, evs)?; + } + Command::PeerDiscovered { peer, hash } => { + if self + .bitfields + .contains_key(&(BitfieldPeer::Remote(peer), hash)) + { + // we already have a subscription for this peer + return Ok(()); + }; + // check if anybody needs this peer + if self.downloads.values_for_hash(hash).next().is_none() { + return Ok(()); + } + // create a peer state if it does not exist + let _state = self.peers.entry(peer).or_default(); + // create a new subscription + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Remote(peer), + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Remote(peer), hash), + PeerBlobState::new(subscription_id), + ); + } + Command::DropPeer { peer } => { + self.bitfields.retain(|(p, _), state| { + if *p == BitfieldPeer::Remote(peer) { + // todo: should we emit unsubscribe evs here? + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); + return false; + } else { + return true; + } + }); + self.peers.remove(&peer); + } + Command::Bitfield { peer, hash, ranges } => { + let state = self.bitfields.get_mut(&(peer, hash)).context(format!( + "bitfields for unknown peer {peer:?} and hash {hash}" + ))?; + let _chunks = total_chunks(&ranges).context("open range")?; + 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 + if let Some(observers) = self.observers.get_by_hash(&hash) { + for (id, request) in observers { + let ranges = &ranges & &request.ranges; + evs.push(Event::LocalBitfield { id: *id, ranges }); + } + } + state.ranges = ranges; + self.check_completion(hash, None, evs)?; + } else { + // We got an entirely new peer, mark all affected downloads for rebalancing + for download in self.downloads.values_mut_for_hash(hash) { + if ranges.intersects(&download.request.ranges) { + download.needs_rebalancing = true; + } + } + state.ranges = ranges; + } + // 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::BitfieldUpdate { + peer, + hash, + added, + removed, + } => { + 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 + if let Some(observers) = self.observers.get_by_hash(&hash) { + for (id, request) in observers { + let added = &added & &request.ranges; + let removed = &removed & &request.ranges; + if !added.is_empty() || !removed.is_empty() { + evs.push(Event::LocalBitfieldUpdate { + id: *id, + added: &added & &request.ranges, + removed: &removed & &request.ranges, + }); + } + } + } + state.ranges |= added; + state.ranges &= !removed; + self.check_completion(hash, None, evs)?; + } else { + // We got more data for this hash, mark all affected downloads for rebalancing + for download in self.downloads.values_mut_for_hash(hash) { + // if removed is non-empty, that is so weird that we just rebalance in any case + if !removed.is_empty() || added.intersects(&download.request.ranges) { + download.needs_rebalancing = true; + } + } + state.ranges |= added; + state.ranges &= !removed; + // a local bitfield update does not make more data available, so we don't need to start downloads + self.start_downloads(hash, None, evs)?; + } + } + Command::ChunksDownloaded { + time, + peer, + hash, + added, + } => { + let state = self.bitfields.get_local_mut(hash).context(format!( + "chunks downloaded before having local bitfield for {hash}" + ))?; + let total_downloaded = total_chunks(&added).context("open range")?; + let total_before = total_chunks(&state.ranges).context("open range")?; + state.ranges |= added; + let total_after = total_chunks(&state.ranges).context("open range")?; + let useful_downloaded = total_after - total_before; + let peer = self.peers.get_mut(&peer).context(format!( + "performing download before having peer state for {peer}" + ))?; + peer.download_history + .push_back((time, (total_downloaded, useful_downloaded))); + } + Command::Tick { time } => { + let window = 10; + let horizon = time.saturating_sub(Duration::from_secs(window)); + // clean up download history + let mut to_rebalance = vec![]; + for (peer, state) in self.peers.iter_mut() { + state + .download_history + .retain(|(duration, _)| *duration > horizon); + let mut sum_total = 0; + let mut sum_useful = 0; + for (_, (total, useful)) in state.download_history.iter() { + sum_total += total; + sum_useful += useful; + } + let speed_useful = (sum_useful as f64) / (window as f64); + let speed_total = (sum_total as f64) / (window as f64); + trace!("peer {peer} download speed {speed_total} cps total, {speed_useful} cps useful"); + } + + for (id, download) in self.downloads.iter() { + if !download.needs_rebalancing { + // nothing has changed that affects this download + continue; + } + let n_peers = self + .bitfields + .remote_for_hash(download.request.hash) + .count(); + if download.peer_downloads.len() >= n_peers { + // we are already downloading from all peers for this hash + continue; + } + to_rebalance.push(*id); + } + for id in to_rebalance { + self.rebalance_download(id, evs)?; + } + } + Command::ObserveLocal { id, hash, ranges } => { + // either we have a subscription for this blob, or we have to create one + if let Some(state) = self.bitfields.get_local_mut(hash) { + // just increment the count + state.subscription_count += 1; + // emit the current bitfield + evs.push(Event::LocalBitfield { + id, + ranges: state.ranges.clone(), + }); + } else { + // create a new subscription + let subscription_id = self.bitfields.subscription_id_gen.next(); + evs.push(Event::SubscribeBitfield { + peer: BitfieldPeer::Local, + hash, + id: subscription_id, + }); + self.bitfields.insert( + (BitfieldPeer::Local, hash), + PeerBlobState::new(subscription_id), + ); + } + self.observers.insert(id, ObserveInfo { hash, ranges }); + } + Command::StopObserveLocal { id } => { + let request = self + .observers + .remove(&id) + .context(format!("stop observing unknown local bitfield {id:?}"))?; + let removed_hash = request.hash; + // unsubscribe from bitfields that have no more subscriptions + self.bitfields.retain(|(_peer, hash), state| { + if *hash == removed_hash { + state.subscription_count -= 1; + if state.subscription_count == 0 { + evs.push(Event::UnsubscribeBitfield { + id: state.subscription_id, + }); + return false; + } + } + true + }); + } + } + Ok(()) + } + + /// Check for completion of a download or of an individual peer download + /// + /// This must be called after each change of the local bitfield for a hash + /// + /// In addition to checking for completion, this also create new peer downloads if a peer download is complete and there is more data available for that peer. + fn check_completion( + &mut self, + hash: Hash, + just_id: Option, + evs: &mut Vec, + ) -> anyhow::Result<()> { + let Some(self_state) = self.bitfields.get_local(hash) else { + // 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 mut completed = vec![]; + for (id, download) in self.downloads.iter_mut_for_hash(hash) { + if just_id.is_some() && just_id != Some(*id) { + continue; + } + // 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 + evs.push(Event::DownloadComplete { id: *id }); + // remember id for later cleanup + completed.push(*id); + // no need to look at individual peer downloads in this case + continue; + } + // check if any peer download is complete, and remove it. + let mut available = vec![]; + download.peer_downloads.retain(|peer, peer_download| { + if self_state.ranges.is_superset(&peer_download.ranges) { + // stop this peer download. + // + // Might be a noop if the cause for this local change was the same peer download, but we don't know. + evs.push(Event::StopPeerDownload { + id: peer_download.id, + }); + // mark this peer as available + available.push(*peer); + false + } else { + true + } + }); + // reassign the newly available peers without doing a full rebalance + if !available.is_empty() { + // check if any of the available peers can provide something of the remaining data + let mut remaining = &download.request.ranges - &self_state.ranges; + // subtract the ranges that are already being taken care of by remaining peer downloads + for peer_download in download.peer_downloads.values() { + remaining.difference_with(&peer_download.ranges); + } + // see what the new peers can do for us + let mut candidates = BTreeMap::new(); + for peer in available { + let Some(peer_state) = self.bitfields.get(&(BitfieldPeer::Remote(peer), hash)) + else { + // weird. we should have a bitfield for this peer since it just completed a download + continue; + }; + let intersection = &peer_state.ranges & &remaining; + if !intersection.is_empty() { + candidates.insert(peer, intersection); + } + } + // deduplicate the ranges + self.planner.plan(hash, &mut candidates); + // start new downloads + for (peer, ranges) in candidates { + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); + } + } + } + // cleanup completed downloads, has to happen later to avoid double mutable borrow + for id in completed { + self.stop_download(id, evs)?; + } + Ok(()) + } + + /// Look at all downloads for a hash and start peer downloads for those that do not have any yet + fn start_downloads( + &mut self, + hash: Hash, + just_id: Option, + evs: &mut Vec, + ) -> anyhow::Result<()> { + let Some(self_state) = self.bitfields.get_local(hash) else { + // we don't have the self state yet, so we can't really decide if we need to download anything at all + return Ok(()); + }; + for (id, download) in self + .downloads + .iter_mut_for_hash(hash) + .filter(|(_, download)| download.peer_downloads.is_empty()) + { + if just_id.is_some() && just_id != Some(*id) { + continue; + } + let remaining = &download.request.ranges - &self_state.ranges; + let mut candidates = BTreeMap::new(); + for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { + let intersection = &bitfield.ranges & &remaining; + if !intersection.is_empty() { + candidates.insert(*peer, intersection); + } + } + self.planner.plan(hash, &mut candidates); + for (peer, ranges) in candidates { + info!(" Starting download from {peer} for {hash} {ranges:?}"); + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); + } + } + Ok(()) + } + + /// rebalance a single download + fn rebalance_download(&mut self, id: DownloadId, evs: &mut Vec) -> anyhow::Result<()> { + let download = self + .downloads + .by_id_mut(id) + .context(format!("rebalancing unknown download {id:?}"))?; + download.needs_rebalancing = false; + tracing::info!("Rebalancing download {id:?} {:?}", download.request); + let hash = download.request.hash; + let Some(self_state) = self.bitfields.get_local(hash) else { + // 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 remaining = &download.request.ranges - &self_state.ranges; + let mut candidates = BTreeMap::new(); + for (peer, bitfield) in self.bitfields.remote_for_hash(hash) { + let intersection = &bitfield.ranges & &remaining; + if !intersection.is_empty() { + candidates.insert(*peer, intersection); + } + } + self.planner.plan(hash, &mut candidates); + info!( + "Stopping {} old peer downloads", + download.peer_downloads.len() + ); + for (_, state) in &download.peer_downloads { + // stop all downloads + evs.push(Event::StopPeerDownload { id: state.id }); + } + info!("Creating {} new peer downloads", candidates.len()); + download.peer_downloads.clear(); + for (peer, ranges) in candidates { + info!(" Starting download from {peer} for {hash} {ranges:?}"); + let id = self.peer_download_id_gen.next(); + evs.push(Event::StartPeerDownload { + id, + peer, + hash, + ranges: ranges.clone(), + }); + download + .peer_downloads + .insert(peer, PeerDownloadState { id, ranges }); + } + Ok(()) + } +} + +#[derive(Default)] +struct Bitfields { + // Counters to generate unique ids for various requests. + // We could use uuid here, but using integers simplifies testing. + // + // the next subscription id + subscription_id_gen: IdGenerator, + by_peer_and_hash: BTreeMap<(BitfieldPeer, Hash), PeerBlobState>, +} + +impl Bitfields { + fn retain(&mut self, mut f: F) + where + F: FnMut(&(BitfieldPeer, Hash), &mut PeerBlobState) -> bool, + { + self.by_peer_and_hash.retain(|k, v| f(k, v)); + } + + fn get(&self, key: &(BitfieldPeer, Hash)) -> Option<&PeerBlobState> { + self.by_peer_and_hash.get(key) + } + + fn get_local(&self, hash: Hash) -> Option<&PeerBlobState> { + self.by_peer_and_hash.get(&(BitfieldPeer::Local, hash)) + } + + fn get_mut(&mut self, key: &(BitfieldPeer, Hash)) -> Option<&mut PeerBlobState> { + self.by_peer_and_hash.get_mut(key) + } + + fn get_local_mut(&mut self, hash: Hash) -> Option<&mut PeerBlobState> { + self.by_peer_and_hash.get_mut(&(BitfieldPeer::Local, hash)) + } + + fn insert(&mut self, key: (BitfieldPeer, Hash), value: PeerBlobState) { + self.by_peer_and_hash.insert(key, value); + } + + fn contains_key(&self, key: &(BitfieldPeer, Hash)) -> bool { + self.by_peer_and_hash.contains_key(key) + } + + fn remote_for_hash(&self, hash: Hash) -> impl Iterator { + self.by_peer_and_hash + .iter() + .filter_map(move |((peer, h), state)| { + if let BitfieldPeer::Remote(peer) = peer { + if *h == hash { + Some((peer, state)) + } else { + None + } + } else { + None + } + }) + } +} + +struct PeerDownloadState { + id: PeerDownloadId, + ranges: ChunkRanges, +} + +struct DownloadState { + /// The request this state is for + request: DownloadRequest, + /// Ongoing downloads + peer_downloads: BTreeMap, + /// Set to true if the download needs rebalancing + needs_rebalancing: bool, +} + +impl DownloadState { + fn new(request: DownloadRequest) -> Self { + Self { + request, + peer_downloads: BTreeMap::new(), + needs_rebalancing: false, + } + } +} +/// Wrapper for the downloads map +/// +/// This is so we can later optimize access by fields other than id, such as hash. +#[derive(Default)] +struct Downloads { + by_id: BTreeMap, +} + +impl Downloads { + fn remove(&mut self, id: &DownloadId) -> Option { + self.by_id.remove(id) + } + + fn contains_key(&self, id: &DownloadId) -> bool { + self.by_id.contains_key(id) + } + + fn insert(&mut self, id: DownloadId, state: DownloadState) { + self.by_id.insert(id, state); + } + + fn iter_mut_for_hash( + &mut self, + hash: Hash, + ) -> impl Iterator { + self.by_id + .iter_mut() + .filter(move |x| x.1.request.hash == hash) + } + + fn iter(&mut self) -> impl Iterator { + self.by_id.iter() + } + + /// Iterate over all downloads for a given hash + fn values_for_hash(&self, hash: Hash) -> impl Iterator { + self.by_id.values().filter(move |x| x.request.hash == hash) + } + + fn values_mut_for_hash(&mut self, hash: Hash) -> impl Iterator { + self.by_id + .values_mut() + .filter(move |x| x.request.hash == hash) + } + + fn by_id_mut(&mut self, id: DownloadId) -> Option<&mut DownloadState> { + self.by_id.get_mut(&id) + } + + fn by_peer_download_id_mut( + &mut self, + id: PeerDownloadId, + ) -> Option<(&DownloadId, &mut DownloadState)> { + self.by_id + .iter_mut() + .filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) + .next() + } +} + +#[derive(Debug)] +struct ObserveInfo { + hash: Hash, + ranges: ChunkRanges, +} + +#[derive(Debug, Default)] +struct Observers { + by_hash_and_id: BTreeMap>, +} + +impl Observers { + fn insert(&mut self, id: ObserveId, request: ObserveInfo) { + self.by_hash_and_id + .entry(request.hash) + .or_default() + .insert(id, request); + } + + fn remove(&mut self, id: &ObserveId) -> Option { + for requests in self.by_hash_and_id.values_mut() { + if let Some(request) = requests.remove(id) { + return Some(request); + } + } + None + } + + fn get_by_hash(&self, hash: &Hash) -> Option<&BTreeMap> { + self.by_hash_and_id.get(hash) + } +} + +/// Global information about a peer +#[derive(Debug, Default)] +struct PeerState { + /// Executed downloads, to calculate the average download speed. + /// + /// This gets updated as soon as possible when new data has been downloaded. + download_history: VecDeque<(Duration, (u64, u64))>, +} + +/// Information about one blob on one peer +struct PeerBlobState { + /// The subscription id for the subscription + subscription_id: BitfieldSubscriptionId, + /// The number of subscriptions this peer has + subscription_count: usize, + /// chunk ranges this peer reports to have + ranges: ChunkRanges, +} + +impl PeerBlobState { + fn new(subscription_id: BitfieldSubscriptionId) -> Self { + Self { + subscription_id, + subscription_count: 1, + ranges: ChunkRanges::empty(), + } + } +} + +#[cfg(test)] +mod tests { + + use super::super::tests::{ + chunk_ranges, has_all_events, has_one_event, has_one_event_matching, noop_planner, + }; + use super::*; + use testresult::TestResult; + + /// Test a simple scenario where a download is started and completed + #[test] + fn downloader_state_smoke() -> TestResult<()> { + use BitfieldPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + let evs = state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: DownloadId(0), + }); + assert!( + has_one_event( + &evs, + &Event::StartDiscovery { + hash, + id: DiscoveryId(0) + } + ), + "starting a download should start a discovery task" + ); + assert!( + has_one_event( + &evs, + &Event::SubscribeBitfield { + peer: Local, + hash, + id: BitfieldSubscriptionId(0) + } + ), + "starting a download should subscribe to the local bitfield" + ); + let initial_bitfield = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: initial_bitfield.clone(), + }); + assert!(evs.is_empty()); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .ranges, + initial_bitfield, + "bitfield should be set to the initial bitfield" + ); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .subscription_count, + 1, + "we have one download interested in the bitfield" + ); + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }); + assert!(evs.is_empty()); + assert_eq!( + state + .bitfields + .get_local(hash) + .context("bitfield should be present")? + .ranges, + ChunkRanges::from(ChunkNum(0)..ChunkNum(32)), + "bitfield should be updated" + ); + let evs = state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + assert!( + has_one_event( + &evs, + &Event::SubscribeBitfield { + peer: Remote(peer_a), + hash, + id: 1.into() + } + ), + "adding a new peer for a hash we are interested in should subscribe to the bitfield" + ); + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..64]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: PeerDownloadId(0), + peer: peer_a, + hash, + ranges: chunk_ranges([32..64]) + } + ), + "bitfield from a peer should start a download" + ); + // ChunksDownloaded just updates the peer stats + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([32..48]), + }); + assert!(evs.is_empty()); + // Bitfield update does not yet complete the download + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([32..48]), + removed: ChunkRanges::empty(), + }); + assert!(evs.is_empty()); + // ChunksDownloaded just updates the peer stats + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([48..64]), + }); + assert!(evs.is_empty()); + // Final bitfield update for the local bitfield should complete the download + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([48..64]), + removed: ChunkRanges::empty(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), + "download should be completed by the data" + ); + // quick check that everything got cleaned up + assert!(state.downloads.by_id.is_empty()); + assert!(state.bitfields.by_peer_and_hash.is_empty()); + assert!(state.discovery.is_empty()); + Ok(()) + } + + /// Test a scenario where more data becomes available at the remote peer as the download progresses + #[test] + fn downloader_state_incremental() -> TestResult<()> { + use BitfieldPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + // Start a download + state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: DownloadId(0), + }); + // Initially, we have nothing + state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::empty(), + }); + // We have a peer for the hash + state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + // We have a bitfield from the peer + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..32]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 0.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([0..32]) + } + ), + "bitfield from a peer should start a download" + ); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([0..16]), + }); + // Bitfield update does not yet complete the download + state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([0..16]), + removed: ChunkRanges::empty(), + }); + // The peer now has more data + state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([32..64]), + }); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([16..32]), + }); + // Complete the first part of the download + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }); + // This triggers cancellation of the first peer download and starting a new one for the remaining data + assert!( + has_one_event(&evs, &Event::StopPeerDownload { id: 0.into() }), + "first peer download should be stopped" + ); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 1.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([32..64]) + } + ), + "second peer download should be started" + ); + // ChunksDownloaded just updates the peer stats + state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([32..64]), + }); + // Final bitfield update for the local bitfield should complete the download + let evs = state.apply(Command::BitfieldUpdate { + peer: Local, + hash, + added: chunk_ranges([32..64]), + removed: ChunkRanges::empty(), + }); + assert!( + has_all_events( + &evs, + &[ + &Event::StopPeerDownload { id: 1.into() }, + &Event::DownloadComplete { id: 0.into() }, + &Event::UnsubscribeBitfield { id: 0.into() }, + &Event::StopDiscovery { id: 0.into() }, + ] + ), + "download should be completed by the data" + ); + println!("{evs:?}"); + Ok(()) + } + + #[test] + fn downloader_state_multiple_downloads() -> testresult::TestResult<()> { + use BitfieldPeer::*; + // Use a constant hash (the same style as used in other tests). + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + // Create a downloader state with a no‐op planner. + let mut state = DownloaderState::new(noop_planner()); + + // --- Start the first (ongoing) download. + // Request a range from 0..64. + let download0 = DownloadId(0); + let req0 = DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }; + let evs0 = state.apply(Command::StartDownload { + request: req0, + id: download0, + }); + // When starting the download, we expect a discovery task to be started + // and a subscription to the local bitfield to be requested. + assert!( + has_one_event( + &evs0, + &Event::StartDiscovery { + hash, + id: DiscoveryId(0) + } + ), + "download0 should start discovery" + ); + assert!( + has_one_event( + &evs0, + &Event::SubscribeBitfield { + peer: Local, + hash, + id: BitfieldSubscriptionId(0) + } + ), + "download0 should subscribe to the local bitfield" + ); + + // --- Simulate some progress for the first download. + // Let’s say only chunks 0..32 are available locally. + let evs1 = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: chunk_ranges([0..32]), + }); + // No completion event should be generated for download0 because its full range 0..64 is not yet met. + assert!( + evs1.is_empty(), + "Partial bitfield update should not complete download0" + ); + + // --- Start a second download for the same hash. + // This new download only requires chunks 0..32 which are already available. + let download1 = DownloadId(1); + let req1 = DownloadRequest { + hash, + ranges: chunk_ranges([0..32]), + }; + let evs2 = state.apply(Command::StartDownload { + request: req1, + id: download1, + }); + // Because the local bitfield (0..32) is already a superset of the new download’s request, + // a DownloadComplete event for download1 should be generated immediately. + assert!( + has_one_event(&evs2, &Event::DownloadComplete { id: download1 }), + "New download should complete immediately" + ); + + // --- Verify state: + // The ongoing download (download0) should still be present in the state, + // while the newly completed download (download1) is removed. + assert!( + state.downloads.contains_key(&download0), + "download0 should still be active" + ); + assert!( + !state.downloads.contains_key(&download1), + "download1 should have been cleaned up after completion" + ); + + Ok(()) + } + + /// Test a scenario where more data becomes available at the remote peer as the download progresses + #[test] + fn downloader_state_drop() -> TestResult<()> { + use BitfieldPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let mut state = DownloaderState::new(noop_planner()); + // Start a download + state.apply(Command::StartDownload { + request: DownloadRequest { + hash, + ranges: chunk_ranges([0..64]), + }, + id: 0.into(), + }); + // Initially, we have nothing + state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::empty(), + }); + // We have a peer for the hash + state.apply(Command::PeerDiscovered { peer: peer_a, hash }); + // We have a bitfield from the peer + let evs = state.apply(Command::Bitfield { + peer: Remote(peer_a), + hash, + ranges: chunk_ranges([0..32]), + }); + assert!( + has_one_event( + &evs, + &Event::StartPeerDownload { + id: 0.into(), + peer: peer_a, + hash, + ranges: chunk_ranges([0..32]) + } + ), + "bitfield from a peer should start a download" + ); + // Sending StopDownload should stop the download and all associated tasks + // This is what happens (delayed) when the user drops the download future + let evs = state.apply(Command::StopDownload { id: 0.into() }); + assert!(has_one_event( + &evs, + &Event::StopPeerDownload { id: 0.into() } + )); + assert!(has_one_event( + &evs, + &Event::UnsubscribeBitfield { id: 0.into() } + )); + assert!(has_one_event( + &evs, + &Event::UnsubscribeBitfield { id: 1.into() } + )); + assert!(has_one_event(&evs, &Event::StopDiscovery { id: 0.into() })); + Ok(()) + } + + /// Test various things that should produce errors + #[test] + fn downloader_state_errors() -> TestResult<()> { + use BitfieldPeer::*; + let _ = tracing_subscriber::fmt::try_init(); + let peer_a = "1000000000000000000000000000000000000000000000000000000000000000".parse()?; + let hash = "0000000000000000000000000000000000000000000000000000000000000001".parse()?; + let unknown_hash = + "0000000000000000000000000000000000000000000000000000000000000002".parse()?; + let mut state = DownloaderState::new(noop_planner()); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash, + ranges: ChunkRanges::all(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "adding an open bitfield should produce an error!" + ); + let evs = state.apply(Command::Bitfield { + peer: Local, + hash: unknown_hash, + ranges: ChunkRanges::all(), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "adding an open bitfield for an unknown hash should produce an error!" + ); + let evs = state.apply(Command::ChunksDownloaded { + time: Duration::ZERO, + peer: peer_a, + hash, + added: chunk_ranges([0..16]), + }); + assert!( + has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), + "download from unknown peer should lead to an error!" + ); + Ok(()) + } +} From 070609840901069eb0e722173a0fbcb035400bc0 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Mon, 10 Feb 2025 18:05:45 +0200 Subject: [PATCH 31/47] more moving stuff around --- src/downloader2.rs | 23 ----------------------- src/downloader2/actor.rs | 14 ++++++++++++++ src/downloader2/state.rs | 18 +++++++++++++++++- 3 files changed, 31 insertions(+), 24 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index fe5361c7c..4b87f1bea 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -207,17 +207,6 @@ pub enum BitfieldPeer { Remote(NodeId), } -fn total_chunks(chunks: &ChunkRanges) -> Option { - let mut total = 0; - for range in chunks.iter() { - match range { - RangeSetRange::RangeFrom(_range) => return None, - RangeSetRange::Range(range) => total += range.end.0 - range.start.0, - } - } - Some(total) -} - /// A downloader that allows range downloads and downloads from multiple peers. #[derive(Debug, Clone)] pub struct Downloader { @@ -361,18 +350,6 @@ impl Downloader { } } -/// An user-facing command -#[derive(Debug)] -pub(super) enum UserCommand { - Download { - request: DownloadRequest, - done: tokio::sync::oneshot::Sender<()>, - }, - Observe { - request: ObserveRequest, - send: tokio::sync::mpsc::Sender, - }, -} /// A simple static content discovery mechanism #[derive(Debug)] pub struct StaticContentDiscovery { diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index f0b08fb26..242d5b66e 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -1,5 +1,19 @@ +//! The actor for the downloader use super::*; +/// An user-facing command +#[derive(Debug)] +pub(super) enum UserCommand { + Download { + request: DownloadRequest, + done: tokio::sync::oneshot::Sender<()>, + }, + Observe { + request: ObserveRequest, + send: tokio::sync::mpsc::Sender, + }, +} + pub(super) struct DownloaderActor { local_pool: LocalPool, endpoint: Endpoint, diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index eea84b12d..37e890ec6 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -1,4 +1,9 @@ -//! The state machine for the downloader +//! The state machine for the downloader, as well as the commands and events. +//! +//! In addition to the state machine, there are also some structs encapsulating +//! a part of the state. These are at this point just wrappers around a single +//! map, but can be made more efficient later if needed without breaking the +//! interface. use super::*; #[derive(Debug)] @@ -888,6 +893,17 @@ impl PeerBlobState { } } +fn total_chunks(chunks: &ChunkRanges) -> Option { + let mut total = 0; + for range in chunks.iter() { + match range { + RangeSetRange::RangeFrom(_range) => return None, + RangeSetRange::Range(range) => total += range.end.0 - range.start.0, + } + } + Some(total) +} + #[cfg(test)] mod tests { From 84fbe80ffdd19357aa1167aa1d1fea5756e0c4da Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 11 Feb 2025 11:05:06 +0200 Subject: [PATCH 32/47] clippy --- examples/multiprovider.rs | 2 +- src/downloader2.rs | 82 ++++------------------------ src/downloader2/actor.rs | 14 ++++- src/downloader2/content_discovery.rs | 63 +++++++++++++++++++++ src/downloader2/state.rs | 11 ++-- 5 files changed, 93 insertions(+), 79 deletions(-) create mode 100644 src/downloader2/content_discovery.rs diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index e00c2a037..9bec4ce7b 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -236,7 +236,7 @@ async fn download(args: DownloadArgs) -> anyhow::Result<()> { let bitmap = bitmap(current, requested, rows as usize); print!("\r{bitmap}"); if progress.is_done() { - println!(""); + println!(); break; } } diff --git a/src/downloader2.rs b/src/downloader2.rs index 4b87f1bea..28ca7c1bb 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -51,6 +51,9 @@ use state::*; mod actor; use actor::*; +mod content_discovery; +pub use content_discovery::*; + #[derive( Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, derive_more::From, )] @@ -76,35 +79,6 @@ struct PeerDownloadId(u64); )] struct BitfieldSubscriptionId(u64); -/// Announce kind -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] -pub enum AnnounceKind { - /// The peer supposedly has some of the data. - Partial = 0, - /// The peer supposedly has the complete data. - #[default] - Complete, -} - -/// Options for finding peers -#[derive(Debug, Default)] -pub struct FindPeersOpts { - /// Kind of announce - pub kind: AnnounceKind, -} - -/// A pluggable content discovery mechanism -pub trait ContentDiscovery: std::fmt::Debug + Send + 'static { - /// Find peers that have the given blob. - /// - /// The returned stream is a handle for the discovery task. It should be an - /// infinite stream that only stops when it is dropped. - fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> BoxStream<'static, NodeId>; -} - -/// A boxed content discovery -pub type BoxedContentDiscovery = Box; - /// A pluggable bitfield subscription mechanism pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static { /// Subscribe to a bitfield @@ -265,7 +239,7 @@ impl DownloaderBuilder { { let store = self.store; let discovery = self.discovery.expect("discovery not set"); - let local_pool = self.local_pool.unwrap_or_else(|| LocalPool::single()); + let local_pool = self.local_pool.unwrap_or_else(LocalPool::single); let planner = self .planner .unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); @@ -350,33 +324,6 @@ impl Downloader { } } -/// A simple static content discovery mechanism -#[derive(Debug)] -pub struct StaticContentDiscovery { - info: BTreeMap>, - default: Vec, -} - -impl StaticContentDiscovery { - /// Create a new static content discovery mechanism - pub fn new(mut info: BTreeMap>, mut default: Vec) -> Self { - default.sort(); - default.dedup(); - for (_, peers) in info.iter_mut() { - peers.sort(); - peers.dedup(); - } - Self { info, default } - } -} - -impl ContentDiscovery for StaticContentDiscovery { - fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> BoxStream<'static, NodeId> { - let peers = self.info.get(&hash).unwrap_or(&self.default).clone(); - Box::pin(futures_lite::stream::iter(peers).chain(futures_lite::stream::pending())) - } -} - /// A bitfield subscription that just returns nothing for local and everything(*) for remote /// /// * Still need to figure out how to deal with open ended chunk ranges. @@ -423,7 +370,7 @@ impl SimpleBitfieldSubscription { } async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { - if let Some(entry) = store.get_mut(&hash).await? { + if let Some(entry) = store.get_mut(hash).await? { crate::get::db::valid_ranges::(&entry).await } else { Ok(ChunkRanges::empty()) @@ -436,7 +383,7 @@ async fn get_valid_ranges_remote( hash: &Hash, ) -> anyhow::Result { let conn = endpoint.connect(id, crate::ALPN).await?; - let (size, _) = crate::get::request::get_verified_size(&conn, &hash).await?; + let (size, _) = crate::get::request::get_verified_size(&conn, hash).await?; let chunks = (size + 1023) / 1024; Ok(ChunkRanges::from(ChunkNum(0)..ChunkNum(chunks))) } @@ -491,6 +438,7 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { #[cfg(test)] mod tests { + #![allow(clippy::single_range_in_vec_init)] use std::ops::Range; use crate::net_protocol::Blobs; @@ -534,7 +482,7 @@ mod tests { let mut planner = StripePlanner2::new(0, 4); let hash = Hash::new(b"test"); let mut ranges = make_range_map(&[chunk_ranges([0..50]), chunk_ranges([50..100])]); - println!(""); + println!(); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); @@ -550,7 +498,7 @@ mod tests { chunk_ranges([0..100]), chunk_ranges([0..100]), ]); - println!(""); + println!(); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); @@ -567,7 +515,7 @@ mod tests { chunk_ranges([0..120]), chunk_ranges([0..50]), ]); - println!(""); + println!(); print_range_map(&ranges); println!("planning"); planner.plan(hash, &mut ranges); @@ -656,10 +604,7 @@ mod tests { .discovery_n0() .bind() .await?; - let discovery = StaticContentDiscovery { - info: BTreeMap::new(), - default: vec![peer], - }; + let discovery = StaticContentDiscovery::new(BTreeMap::new(), vec![peer]); let bitfield_subscription = TestBitfieldSubscription; let downloader = Downloader::builder(endpoint, store) .discovery(discovery) @@ -697,10 +642,7 @@ mod tests { .discovery_n0() .bind() .await?; - let discovery = StaticContentDiscovery { - info: BTreeMap::new(), - default: peers, - }; + let discovery = StaticContentDiscovery::new(BTreeMap::new(), peers); let downloader = Downloader::builder(endpoint, store) .discovery(discovery) .planner(StripePlanner2::new(0, 8)) diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index 242d5b66e..3b376a090 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -15,17 +15,23 @@ pub(super) enum UserCommand { } pub(super) struct DownloaderActor { + state: DownloaderState, local_pool: LocalPool, endpoint: Endpoint, command_rx: mpsc::Receiver, command_tx: mpsc::Sender, - state: DownloaderState, store: S, + /// Content discovery discovery: BoxedContentDiscovery, + /// Bitfield subscription subscribe_bitfield: BoxedBitfieldSubscription, + /// Futures for downloads download_futs: BTreeMap>, + /// Tasks for peer downloads peer_download_tasks: BTreeMap>, + /// Tasks for discovery discovery_tasks: BTreeMap>, + /// Tasks for bitfield subscriptions bitfield_subscription_tasks: BTreeMap>, /// Id generator for download ids download_id_gen: IdGenerator, @@ -67,13 +73,13 @@ impl DownloaderActor { } } - pub(super) async fn run(mut self, mut channel: mpsc::Receiver) { + pub(super) async fn run(mut self, mut user_commands: mpsc::Receiver) { let mut ticks = tokio::time::interval(Duration::from_millis(100)); loop { trace!("downloader actor tick"); tokio::select! { biased; - Some(cmd) = channel.recv() => { + Some(cmd) = user_commands.recv() => { debug!("user command {cmd:?}"); match cmd { UserCommand::Download { @@ -228,6 +234,7 @@ impl DownloaderActor { } } +#[allow(clippy::too_many_arguments)] async fn peer_download_task( id: PeerDownloadId, endpoint: Endpoint, @@ -327,6 +334,7 @@ async fn peer_download( Ok(stats) } +/// Spawn a future and wrap it in a [`AbortOnDropHandle`] pub(super) fn spawn(f: F) -> AbortOnDropHandle where F: Future + Send + 'static, diff --git a/src/downloader2/content_discovery.rs b/src/downloader2/content_discovery.rs new file mode 100644 index 000000000..a09f04ea9 --- /dev/null +++ b/src/downloader2/content_discovery.rs @@ -0,0 +1,63 @@ +use std::collections::BTreeMap; + +use crate::Hash; +use futures_lite::stream::StreamExt; +use futures_util::stream::BoxStream; +use iroh::NodeId; +use serde::{Deserialize, Serialize}; + +/// Announce kind +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] +pub enum AnnounceKind { + /// The peer supposedly has some of the data. + Partial = 0, + /// The peer supposedly has the complete data. + #[default] + Complete, +} + +/// Options for finding peers +#[derive(Debug, Default)] +pub struct FindPeersOpts { + /// Kind of announce + pub kind: AnnounceKind, +} + +/// A pluggable content discovery mechanism +pub trait ContentDiscovery: std::fmt::Debug + Send + 'static { + /// Find peers that have the given blob. + /// + /// The returned stream is a handle for the discovery task. It should be an + /// infinite stream that only stops when it is dropped. + fn find_peers(&mut self, hash: Hash, opts: FindPeersOpts) -> BoxStream<'static, NodeId>; +} + +/// A boxed content discovery +pub type BoxedContentDiscovery = Box; + +/// A simple static content discovery mechanism +#[derive(Debug)] +pub struct StaticContentDiscovery { + info: BTreeMap>, + default: Vec, +} + +impl StaticContentDiscovery { + /// Create a new static content discovery mechanism + pub fn new(mut info: BTreeMap>, mut default: Vec) -> Self { + default.sort(); + default.dedup(); + for (_, peers) in info.iter_mut() { + peers.sort(); + peers.dedup(); + } + Self { info, default } + } +} + +impl ContentDiscovery for StaticContentDiscovery { + fn find_peers(&mut self, hash: Hash, _opts: FindPeersOpts) -> BoxStream<'static, NodeId> { + let peers = self.info.get(&hash).unwrap_or(&self.default).clone(); + Box::pin(futures_lite::stream::iter(peers).chain(futures_lite::stream::pending())) + } +} diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index 37e890ec6..bdb4595da 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -64,13 +64,13 @@ pub(super) enum Command { DropPeer { peer: NodeId }, /// A peer has been discovered PeerDiscovered { peer: NodeId, hash: Hash }, - /// + /// Start observing a local bitfield ObserveLocal { id: ObserveId, hash: Hash, ranges: ChunkRanges, }, - /// + /// Stop observing a local bitfield StopObserveLocal { id: ObserveId }, /// A tick from the driver, for rebalancing Tick { time: Duration }, @@ -319,9 +319,9 @@ impl DownloaderState { evs.push(Event::UnsubscribeBitfield { id: state.subscription_id, }); - return false; + false } else { - return true; + true } }); self.peers.remove(&peer); @@ -666,7 +666,7 @@ impl DownloaderState { "Stopping {} old peer downloads", download.peer_downloads.len() ); - for (_, state) in &download.peer_downloads { + for state in download.peer_downloads.values() { // stop all downloads evs.push(Event::StopPeerDownload { id: state.id }); } @@ -906,6 +906,7 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { #[cfg(test)] mod tests { + #![allow(clippy::single_range_in_vec_init)] use super::super::tests::{ chunk_ranges, has_all_events, has_one_event, has_one_event_matching, noop_planner, From dbf20de600aa21f215cd42423d44f6ed12da4901 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 11 Feb 2025 12:42:19 +0200 Subject: [PATCH 33/47] more clippy --- src/downloader2/state.rs | 8 +++----- src/store/traits.rs | 2 +- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index bdb4595da..ac98b2f78 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -262,11 +262,11 @@ impl DownloaderState { PeerBlobState::new(subscription_id), ); } - if !self.discovery.contains_key(&request.hash) { + if let std::collections::btree_map::Entry::Vacant(e) = self.discovery.entry(request.hash) { // start a discovery task let id = self.discovery_id_gen.next(); evs.push(Event::StartDiscovery { hash, id }); - self.discovery.insert(request.hash, id); + e.insert(id); } self.downloads.insert(id, DownloadState::new(request)); self.check_completion(hash, Some(id), evs)?; @@ -825,9 +825,7 @@ impl Downloads { id: PeerDownloadId, ) -> Option<(&DownloadId, &mut DownloadState)> { self.by_id - .iter_mut() - .filter(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) - .next() + .iter_mut().find(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) } } diff --git a/src/store/traits.rs b/src/store/traits.rs index 01c48229d..6afeb0c19 100644 --- a/src/store/traits.rs +++ b/src/store/traits.rs @@ -52,7 +52,7 @@ pub enum BaoBlobSize { } impl BaoBlobSize { - /// Create a new `BaoFileSize` with the given size and verification status. + /// Create a new `BaoBlobSize` with the given size and verification status. pub fn new(size: u64, verified: bool) -> Self { if verified { BaoBlobSize::Verified(size) From a8015d55156f3a868f551ef3d7c1ee7dd472a469 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 11 Feb 2025 13:59:58 +0200 Subject: [PATCH 34/47] Extend multiprovider example with persistent local storage --- examples/multiprovider.rs | 92 ++++++++++----------------------------- src/downloader2.rs | 69 +++++++++++++++++++++++++++++ src/downloader2/state.rs | 7 ++- 3 files changed, 98 insertions(+), 70 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 9bec4ce7b..10afc0fb3 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -6,7 +6,8 @@ use console::Term; use iroh::{NodeId, SecretKey}; use iroh_blobs::{ downloader2::{ - DownloadRequest, Downloader, ObserveEvent, ObserveRequest, StaticContentDiscovery, + print_bitmap, DownloadRequest, Downloader, ObserveEvent, ObserveRequest, + StaticContentDiscovery, }, store::Store, util::total_bytes, @@ -30,7 +31,11 @@ struct DownloadArgs { #[clap(help = "hash to download")] hash: Hash, + #[clap(help = "providers to download from")] providers: Vec, + + #[clap(long, help = "path to save to")] + path: Option, } #[derive(Debug, Parser)] @@ -139,77 +144,28 @@ impl BlobDownloadProgress { } } -fn bitmap(current: &[ChunkNum], requested: &[ChunkNum], n: usize) -> String { - // If n is 0, return an empty string. - if n == 0 { - return String::new(); - } - - // Determine the overall bitfield size. - // Since the ranges are sorted, we take the last element as the total size. - let total = if let Some(&last) = requested.last() { - last.0 - } else { - // If there are no ranges, we assume the bitfield is empty. - 0 - }; - - // If total is 0, output n spaces. - if total == 0 { - return " ".repeat(n); - } - - let mut result = String::with_capacity(n); - - // For each of the n output buckets: - for bucket in 0..n { - // Calculate the bucket's start and end in the overall bitfield. - let bucket_start = bucket as u64 * total / n as u64; - let bucket_end = (bucket as u64 + 1) * total / n as u64; - let bucket_size = bucket_end.saturating_sub(bucket_start); - - // Sum the number of bits that are set in this bucket. - let mut set_bits = 0u64; - for pair in current.chunks_exact(2) { - let start = pair[0]; - let end = pair[1]; - // Determine the overlap between the bucket and the current range. - let overlap_start = start.0.max(bucket_start); - let overlap_end = end.0.min(bucket_end); - if overlap_start < overlap_end { - set_bits += overlap_end - overlap_start; - } +async fn download(args: DownloadArgs) -> anyhow::Result<()> { + match &args.path { + Some(path) => { + tokio::fs::create_dir_all(path).await?; + let store = iroh_blobs::store::fs::Store::load(path).await?; + // make sure we properly shut down the store on ctrl-c + let res = tokio::select! { + x = download_impl(args, store.clone()) => x, + _ = tokio::signal::ctrl_c() => Ok(()), + }; + store.shutdown().await; + res + } + None => { + let store = iroh_blobs::store::mem::Store::new(); + download_impl(args, store).await } - - // Calculate the fraction of the bucket that is set. - let fraction = if bucket_size > 0 { - set_bits as f64 / bucket_size as f64 - } else { - 0.0 - }; - - // Map the fraction to a grayscale character. - let ch = if fraction == 0.0 { - ' ' // completely empty - } else if fraction == 1.0 { - '█' // completely full - } else if fraction < 0.25 { - '░' - } else if fraction < 0.5 { - '▒' - } else { - '▓' - }; - - result.push(ch); } - - result } -async fn download(args: DownloadArgs) -> anyhow::Result<()> { +async fn download_impl(args: DownloadArgs, store: S) -> anyhow::Result<()> { let endpoint = iroh::Endpoint::builder().discovery_n0().bind().await?; - let store = iroh_blobs::store::mem::Store::new(); let discovery = StaticContentDiscovery::new(Default::default(), args.providers); let downloader = Downloader::builder(endpoint, store) .discovery(discovery) @@ -233,7 +189,7 @@ async fn download(args: DownloadArgs) -> anyhow::Result<()> { progress.update(chunk); let current = progress.current.boundaries(); let requested = progress.request.ranges.boundaries(); - let bitmap = bitmap(current, requested, rows as usize); + let bitmap = print_bitmap(current, requested, rows as usize); print!("\r{bitmap}"); if progress.is_done() { println!(); diff --git a/src/downloader2.rs b/src/downloader2.rs index 28ca7c1bb..d01e5059d 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -436,6 +436,75 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { } } +/// Print a bitmap +pub fn print_bitmap(current: &[ChunkNum], requested: &[ChunkNum], n: usize) -> String { + // If n is 0, return an empty string. + if n == 0 { + return String::new(); + } + + // Determine the overall bitfield size. + // Since the ranges are sorted, we take the last element as the total size. + let total = if let Some(&last) = requested.last() { + last.0 + } else { + // If there are no ranges, we assume the bitfield is empty. + 0 + }; + + // If total is 0, output n spaces. + if total == 0 { + return " ".repeat(n); + } + + let mut result = String::with_capacity(n); + + // For each of the n output buckets: + for bucket in 0..n { + // Calculate the bucket's start and end in the overall bitfield. + let bucket_start = bucket as u64 * total / n as u64; + let bucket_end = (bucket as u64 + 1) * total / n as u64; + let bucket_size = bucket_end.saturating_sub(bucket_start); + + // Sum the number of bits that are set in this bucket. + let mut set_bits = 0u64; + for pair in current.chunks_exact(2) { + let start = pair[0]; + let end = pair[1]; + // Determine the overlap between the bucket and the current range. + let overlap_start = start.0.max(bucket_start); + let overlap_end = end.0.min(bucket_end); + if overlap_start < overlap_end { + set_bits += overlap_end - overlap_start; + } + } + + // Calculate the fraction of the bucket that is set. + let fraction = if bucket_size > 0 { + set_bits as f64 / bucket_size as f64 + } else { + 0.0 + }; + + // Map the fraction to a grayscale character. + let ch = if fraction == 0.0 { + ' ' // completely empty + } else if fraction == 1.0 { + '█' // completely full + } else if fraction < 0.25 { + '░' + } else if fraction < 0.5 { + '▒' + } else { + '▓' + }; + + result.push(ch); + } + + result +} + #[cfg(test)] mod tests { #![allow(clippy::single_range_in_vec_init)] diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index ac98b2f78..e5343af77 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -262,7 +262,9 @@ impl DownloaderState { PeerBlobState::new(subscription_id), ); } - if let std::collections::btree_map::Entry::Vacant(e) = self.discovery.entry(request.hash) { + if let std::collections::btree_map::Entry::Vacant(e) = + self.discovery.entry(request.hash) + { // start a discovery task let id = self.discovery_id_gen.next(); evs.push(Event::StartDiscovery { hash, id }); @@ -825,7 +827,8 @@ impl Downloads { id: PeerDownloadId, ) -> Option<(&DownloadId, &mut DownloadState)> { self.by_id - .iter_mut().find(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) + .iter_mut() + .find(|(_, v)| v.peer_downloads.iter().any(|(_, state)| state.id == id)) } } From d5112867ac92436ad81d174fd20a935a662696db Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Tue, 11 Feb 2025 14:35:48 +0200 Subject: [PATCH 35/47] fix off by one bug in the quick hacky valid_ranges --- src/downloader2.rs | 12 +++++++++++- src/get/db.rs | 2 +- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index d01e5059d..a408966d6 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -510,7 +510,7 @@ mod tests { #![allow(clippy::single_range_in_vec_init)] use std::ops::Range; - use crate::net_protocol::Blobs; + use crate::{net_protocol::Blobs, store::MapMut}; use super::*; use bao_tree::ChunkNum; @@ -546,6 +546,16 @@ mod tests { SecretKey::from(secret).public() } + #[tokio::test] + async fn test_chunk_ranges() -> TestResult<()> { + let store = crate::store::mem::Store::new(); + 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::(&entry).await?; + println!("{valid:?}"); + Ok(()) + } + #[test] fn test_planner_1() { let mut planner = StripePlanner2::new(0, 4); diff --git a/src/get/db.rs b/src/get/db.rs index 783bbabc5..d76945a2c 100644 --- a/src/get/db.rs +++ b/src/get/db.rs @@ -241,7 +241,7 @@ pub async fn valid_ranges(entry: &D::EntryMut) -> anyhow::Result Date: Wed, 12 Feb 2025 10:37:05 +0200 Subject: [PATCH 36/47] Make the test a real test --- src/downloader2.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/downloader2.rs b/src/downloader2.rs index a408966d6..21fa674a1 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -547,12 +547,12 @@ mod tests { } #[tokio::test] - async fn test_chunk_ranges() -> TestResult<()> { + 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 entry = store.get_mut(tt.hash()).await?.unwrap(); let valid = crate::get::db::valid_ranges::(&entry).await?; - println!("{valid:?}"); + assert!(valid == ChunkRanges::from(ChunkNum(0)..ChunkNum(2))); Ok(()) } From ea3595a698f22c247397d69a697f6581bbc05bf5 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 12 Feb 2025 10:43:30 +0200 Subject: [PATCH 37/47] Unify BitfieldEvent and BitfieldSubscriptionEvent --- examples/multiprovider.rs | 8 +++--- src/downloader2.rs | 55 ++++++++++----------------------------- src/downloader2/actor.rs | 24 ++++++++--------- 3 files changed, 29 insertions(+), 58 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 10afc0fb3..2f9dda72c 100644 --- a/examples/multiprovider.rs +++ b/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; } diff --git a/src/downloader2.rs b/src/downloader2.rs index 21fa674a1..d116520c9 100644 --- a/src/downloader2.rs +++ b/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; -/// 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> { + ) -> anyhow::Result> { let (send, recv) = tokio::sync::mpsc::channel(request.buffer); self.send .send(UserCommand::Observe { request, send }) @@ -332,11 +311,7 @@ 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(_) => { @@ -344,7 +319,7 @@ impl BitfieldSubscription for TestBitfieldSubscription { } }; 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 BitfieldSubscription for SimpleBitfieldSubscription { - 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 BitfieldSubscription for SimpleBitfieldSubscription { 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::(&entry).await?; assert!(valid == ChunkRanges::from(ChunkNum(0)..ChunkNum(2))); diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index 3b376a090..b4b84a7c0 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -10,7 +10,7 @@ pub(super) enum UserCommand { }, Observe { request: ObserveRequest, - send: tokio::sync::mpsc::Sender, + send: tokio::sync::mpsc::Sender, }, } @@ -38,7 +38,7 @@ pub(super) struct DownloaderActor { /// Id generator for observe ids observe_id_gen: IdGenerator, /// Observers - observers: BTreeMap>, + observers: BTreeMap>, /// 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 DownloaderActor { 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 DownloaderActor { 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 DownloaderActor { return; }; if sender - .try_send(ObserveEvent::BitfieldUpdate { added, removed }) + .try_send(BitfieldEvent::Update { added, removed }) .is_err() { // the observer has been dropped From 704c7eba9b6c0fc87ab1cb4b6a00097c78cfad42 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 12 Feb 2025 11:19:16 +0200 Subject: [PATCH 38/47] Unify BitfieldEvents and start adding size (not used yet) --- examples/multiprovider.rs | 2 +- src/downloader2.rs | 10 ++- src/downloader2/actor.rs | 38 ++------- src/downloader2/state.rs | 172 +++++++++++++++++++++++--------------- 4 files changed, 119 insertions(+), 103 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 2f9dda72c..6fe3c63cf 100644 --- a/examples/multiprovider.rs +++ b/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 } => { diff --git a/src/downloader2.rs b/src/downloader2.rs index d116520c9..21401835f 100644 --- a/src/downloader2.rs +++ b/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; /// 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 BitfieldSubscription for SimpleBitfieldSubscription { Ok(ev) => ev, Err(_) => ChunkRanges::empty(), }; - BitfieldEvent::State { ranges } + BitfieldEvent::State { ranges, size: u64::MAX } } .into_stream(), ) diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index b4b84a7c0..47a69a4d9 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -142,18 +142,8 @@ impl DownloaderActor { 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 DownloaderActor { 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( 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(); diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index e5343af77..c0ec9ff29 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -17,28 +17,17 @@ pub(super) enum Command { }, /// A user request to abort a download. StopDownload { id: DownloadId }, - /// A full bitfield for a blob and a peer - Bitfield { - /// The peer that sent the bitfield. - peer: BitfieldPeer, - /// The blob for which the bitfield is - hash: Hash, - /// The complete bitfield - ranges: ChunkRanges, - }, /// An update of a bitfield for a hash /// /// This is used both to update the bitfield of remote peers, and to update /// the local bitfield. - BitfieldUpdate { + BitfieldInfo { /// The peer that sent the update. peer: BitfieldPeer, /// The blob that was updated. hash: Hash, - /// The ranges that were added - added: ChunkRanges, - /// The ranges that were removed - removed: ChunkRanges, + /// The state or update event + event: BitfieldEvent, }, /// A chunk was downloaded, but not yet stored /// @@ -88,14 +77,9 @@ pub(super) enum Event { /// The unique id of the subscription id: BitfieldSubscriptionId, }, - LocalBitfield { + LocalBitfieldInfo { id: ObserveId, - ranges: ChunkRanges, - }, - LocalBitfieldUpdate { - id: ObserveId, - added: ChunkRanges, - removed: ChunkRanges, + event: BitfieldEvent, }, StartDiscovery { hash: Hash, @@ -328,7 +312,11 @@ impl DownloaderState { }); self.peers.remove(&peer); } - Command::Bitfield { peer, hash, ranges } => { + Command::BitfieldInfo { + peer, + hash, + event: BitfieldEvent::State { ranges, size }, + } => { let state = self.bitfields.get_mut(&(peer, hash)).context(format!( "bitfields for unknown peer {peer:?} and hash {hash}" ))?; @@ -339,7 +327,7 @@ impl DownloaderState { if let Some(observers) = self.observers.get_by_hash(&hash) { for (id, request) in observers { let ranges = &ranges & &request.ranges; - evs.push(Event::LocalBitfield { id: *id, ranges }); + evs.push(Event::LocalBitfieldInfo { id: *id, event: BitfieldEvent::State { ranges: ranges.clone(), size } }); } } state.ranges = ranges; @@ -356,11 +344,10 @@ impl DownloaderState { // 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::BitfieldUpdate { + Command::BitfieldInfo { peer, hash, - added, - removed, + event: BitfieldEvent::Update { added, removed }, } => { let state = self.bitfields.get_mut(&(peer, hash)).context(format!( "bitfield update for unknown peer {peer:?} and hash {hash}" @@ -373,10 +360,12 @@ impl DownloaderState { let added = &added & &request.ranges; let removed = &removed & &request.ranges; if !added.is_empty() || !removed.is_empty() { - evs.push(Event::LocalBitfieldUpdate { + evs.push(Event::LocalBitfieldInfo { id: *id, - added: &added & &request.ranges, - removed: &removed & &request.ranges, + event: BitfieldEvent::Update { + added: &added & &request.ranges, + removed: &removed & &request.ranges, + } }); } } @@ -463,9 +452,12 @@ impl DownloaderState { // just increment the count state.subscription_count += 1; // emit the current bitfield - evs.push(Event::LocalBitfield { + evs.push(Event::LocalBitfieldInfo { id, - ranges: state.ranges.clone(), + event: BitfieldEvent::State { + ranges: state.ranges.clone(), + size: u64::MAX, + } }); } else { // create a new subscription @@ -952,10 +944,13 @@ mod tests { "starting a download should subscribe to the local bitfield" ); let initial_bitfield = ChunkRanges::from(ChunkNum(0)..ChunkNum(16)); - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - ranges: initial_bitfield.clone(), + event: BitfieldEvent::State { + ranges: initial_bitfield.clone(), + size: u64::MAX, + }, }); assert!(evs.is_empty()); assert_eq!( @@ -976,11 +971,13 @@ mod tests { 1, "we have one download interested in the bitfield" ); - let evs = state.apply(Command::BitfieldUpdate { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([16..32]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }, }); assert!(evs.is_empty()); assert_eq!( @@ -1004,10 +1001,13 @@ mod tests { ), "adding a new peer for a hash we are interested in should subscribe to the bitfield" ); - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - ranges: chunk_ranges([0..64]), + event: BitfieldEvent::State { + ranges: chunk_ranges([0..64]), + size: u64::MAX, + }, }); assert!( has_one_event( @@ -1030,11 +1030,13 @@ mod tests { }); assert!(evs.is_empty()); // Bitfield update does not yet complete the download - let evs = state.apply(Command::BitfieldUpdate { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([32..48]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([32..48]), + removed: ChunkRanges::empty(), + }, }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats @@ -1046,11 +1048,13 @@ mod tests { }); assert!(evs.is_empty()); // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([48..64]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([48..64]), + removed: ChunkRanges::empty(), + }, }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), @@ -1080,18 +1084,24 @@ mod tests { id: DownloadId(0), }); // Initially, we have nothing - state.apply(Command::Bitfield { + state.apply(Command::BitfieldInfo { peer: Local, hash, - ranges: ChunkRanges::empty(), + event: BitfieldEvent::State { + ranges: ChunkRanges::empty(), + size: u64::MAX, + }, }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - ranges: chunk_ranges([0..32]), + event: BitfieldEvent::State { + ranges: chunk_ranges([0..32]), + size: u64::MAX, + }, }); assert!( has_one_event( @@ -1113,17 +1123,22 @@ mod tests { added: chunk_ranges([0..16]), }); // Bitfield update does not yet complete the download - state.apply(Command::BitfieldUpdate { + state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([0..16]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([0..16]), + removed: ChunkRanges::empty(), + }, }); // The peer now has more data - state.apply(Command::Bitfield { + state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - ranges: chunk_ranges([32..64]), + event: BitfieldEvent::State { + ranges: chunk_ranges([32..64]), + size: u64::MAX, + }, }); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { @@ -1133,11 +1148,13 @@ mod tests { added: chunk_ranges([16..32]), }); // Complete the first part of the download - let evs = state.apply(Command::BitfieldUpdate { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([16..32]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([16..32]), + removed: ChunkRanges::empty(), + }, }); // This triggers cancellation of the first peer download and starting a new one for the remaining data assert!( @@ -1164,11 +1181,13 @@ mod tests { added: chunk_ranges([32..64]), }); // Final bitfield update for the local bitfield should complete the download - let evs = state.apply(Command::BitfieldUpdate { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - added: chunk_ranges([32..64]), - removed: ChunkRanges::empty(), + event: BitfieldEvent::Update { + added: chunk_ranges([32..64]), + removed: ChunkRanges::empty(), + }, }); assert!( has_all_events( @@ -1231,10 +1250,13 @@ mod tests { // --- Simulate some progress for the first download. // Let’s say only chunks 0..32 are available locally. - let evs1 = state.apply(Command::Bitfield { + let evs1 = state.apply(Command::BitfieldInfo { peer: Local, hash, - ranges: chunk_ranges([0..32]), + event: BitfieldEvent::State { + ranges: chunk_ranges([0..32]), + size: u64::MAX, + }, }); // No completion event should be generated for download0 because its full range 0..64 is not yet met. assert!( @@ -1292,18 +1314,24 @@ mod tests { id: 0.into(), }); // Initially, we have nothing - state.apply(Command::Bitfield { + state.apply(Command::BitfieldInfo { peer: Local, hash, - ranges: ChunkRanges::empty(), + event: BitfieldEvent::State { + ranges: ChunkRanges::empty(), + size: u64::MAX, + }, }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); // We have a bitfield from the peer - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - ranges: chunk_ranges([0..32]), + event: BitfieldEvent::State { + ranges: chunk_ranges([0..32]), + size: u64::MAX, + }, }); assert!( has_one_event( @@ -1346,19 +1374,25 @@ mod tests { let unknown_hash = "0000000000000000000000000000000000000000000000000000000000000002".parse()?; let mut state = DownloaderState::new(noop_planner()); - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - ranges: ChunkRanges::all(), + event: BitfieldEvent::State { + ranges: ChunkRanges::all(), + size: u64::MAX, + }, }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), "adding an open bitfield should produce an error!" ); - let evs = state.apply(Command::Bitfield { + let evs = state.apply(Command::BitfieldInfo { peer: Local, hash: unknown_hash, - ranges: ChunkRanges::all(), + event: BitfieldEvent::State { + ranges: ChunkRanges::all(), + size: u64::MAX, + }, }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), From ced9b956738728974470bdba59a92a06652d8c01 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 12 Feb 2025 12:03:58 +0200 Subject: [PATCH 39/47] Make use of the size --- examples/multiprovider.rs | 2 +- src/downloader2.rs | 27 ++++++++++++++++++--------- src/downloader2/actor.rs | 1 + src/downloader2/state.rs | 16 +++++++++++++++- src/get/db.rs | 8 +++++++- 5 files changed, 42 insertions(+), 12 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 6fe3c63cf..12df101bd 100644 --- a/examples/multiprovider.rs +++ b/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; } diff --git a/src/downloader2.rs b/src/downloader2.rs index 21401835f..656baf312 100644 --- a/src/downloader2.rs +++ b/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,11 +350,12 @@ impl SimpleBitfieldSubscription { } } -async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { +async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { if let Some(entry) = store.get_mut(hash).await? { - crate::get::db::valid_ranges::(&entry).await + let (ranges, size) = crate::get::db::valid_ranges_and_size::(&entry).await?; + Ok(BitfieldEvent::State { ranges, size }) } else { - Ok(ChunkRanges::empty()) + Ok(BitfieldEvent::State { ranges: ChunkRanges::empty(), size: u64::MAX }) } } @@ -358,11 +363,12 @@ async fn get_valid_ranges_remote( endpoint: &Endpoint, id: NodeId, hash: &Hash, -) -> anyhow::Result { +) -> anyhow::Result { 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 BitfieldSubscription for SimpleBitfieldSubscription { @@ -398,11 +404,14 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { } 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(), ) diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index 47a69a4d9..8dd4ca00a 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -288,6 +288,7 @@ async fn peer_download( event: BitfieldEvent::Update { added, removed: ChunkRanges::empty(), + size, }, }) .await diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index c0ec9ff29..596fefacf 100644 --- a/src/downloader2/state.rs +++ b/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,6 +341,7 @@ 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)?; @@ -347,7 +349,7 @@ impl DownloaderState { 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}" @@ -355,6 +357,7 @@ impl DownloaderState { 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!( diff --git a/src/get/db.rs b/src/get/db.rs index d76945a2c..b165ff503 100644 --- a/src/get/db.rs +++ b/src/get/db.rs @@ -237,6 +237,12 @@ async fn get_blob( /// Given a partial entry, get the valid ranges. pub async fn valid_ranges(entry: &D::EntryMut) -> anyhow::Result { + let (ranges, _) = valid_ranges_and_size::(entry).await?; + Ok(ranges) +} + +/// Given a partial entry, get the valid ranges. +pub async fn valid_ranges_and_size(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(entry: &D::EntryMut) -> anyhow::Result Date: Wed, 12 Feb 2025 12:21:10 +0200 Subject: [PATCH 40/47] introduce BitfieldUpdate and BitfieldState structs enum is just a wrapper around them now --- examples/multiprovider.rs | 8 +-- src/downloader2.rs | 68 ++++++++++++++++--------- src/downloader2/actor.rs | 5 +- src/downloader2/state.rs | 104 ++++++++++++++++++++++++-------------- src/get/db.rs | 4 +- 5 files changed, 120 insertions(+), 69 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 12df101bd..e3f54ad17 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider.rs @@ -6,8 +6,8 @@ use console::Term; use iroh::{NodeId, SecretKey}; use iroh_blobs::{ downloader2::{ - print_bitmap, BitfieldEvent, DownloadRequest, Downloader, ObserveRequest, - StaticContentDiscovery, + print_bitmap, BitfieldEvent, BitfieldState, BitfieldUpdate, DownloadRequest, Downloader, + ObserveRequest, StaticContentDiscovery, }, store::Store, util::total_bytes, @@ -117,10 +117,10 @@ impl BlobDownloadProgress { fn update(&mut self, ev: BitfieldEvent) { match ev { - BitfieldEvent::State { ranges, .. } => { + BitfieldEvent::State(BitfieldState { ranges, .. }) => { self.current = ranges; } - BitfieldEvent::Update { added, removed, .. } => { + BitfieldEvent::Update(BitfieldUpdate { added, removed, .. }) => { self.current |= added; self.current -= removed; } diff --git a/src/downloader2.rs b/src/downloader2.rs index 656baf312..bbc791f95 100644 --- a/src/downloader2.rs +++ b/src/downloader2.rs @@ -17,7 +17,8 @@ use std::{ io, marker::PhantomData, sync::Arc, - time::Instant, u64, + time::Instant, + u64, }; use crate::{ @@ -89,26 +90,32 @@ pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static { pub type BoxedBitfieldSubscription = Box; /// Events from observing a local bitfield -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, derive_more::From)] pub enum BitfieldEvent { /// The full state of the bitfield - State { - /// The entire bitfield - ranges: ChunkRanges, - /// The most precise known size of the blob. - /// - /// If I know nothing about the blob, this is u64::MAX. - size: u64, - }, + State(BitfieldState), /// An update to the bitfield - Update { - /// The ranges that were added - added: ChunkRanges, - /// The ranges that were removed - removed: ChunkRanges, - /// A refinement of the size of the blob. - size: u64, - }, + Update(BitfieldUpdate), +} + +/// 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, } /// A download request @@ -325,8 +332,14 @@ impl BitfieldSubscription for TestBitfieldSubscription { } }; Box::pin( - futures_lite::stream::once(BitfieldEvent::State { ranges, size: 1024 * 1024 * 1024 * 1024 * 1024 }) - .chain(futures_lite::stream::pending()), + futures_lite::stream::once( + BitfieldState { + ranges, + size: 1024 * 1024 * 1024 * 1024 * 1024, + } + .into(), + ) + .chain(futures_lite::stream::pending()), ) } } @@ -353,9 +366,13 @@ impl SimpleBitfieldSubscription { async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { if let Some(entry) = store.get_mut(hash).await? { let (ranges, size) = crate::get::db::valid_ranges_and_size::(&entry).await?; - Ok(BitfieldEvent::State { ranges, size }) + Ok(BitfieldState { ranges, size }.into()) } else { - Ok(BitfieldEvent::State { ranges: ChunkRanges::empty(), size: u64::MAX }) + Ok(BitfieldState { + ranges: ChunkRanges::empty(), + size: u64::MAX, + } + .into()) } } @@ -368,7 +385,7 @@ 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(BitfieldEvent::State { ranges, size }) + Ok(BitfieldState { ranges, size }.into()) } impl BitfieldSubscription for SimpleBitfieldSubscription { @@ -406,10 +423,11 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { async move { let event = match recv.await { Ok(ev) => ev, - Err(_) => BitfieldEvent::State { + Err(_) => BitfieldState { ranges: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }; event } diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index 8dd4ca00a..3420d0aec 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -285,11 +285,12 @@ async fn peer_download( .send(Command::BitfieldInfo { peer: BitfieldPeer::Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added, removed: ChunkRanges::empty(), size, - }, + } + .into(), }) .await .ok(); diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index 596fefacf..55ff352ca 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -315,7 +315,7 @@ impl DownloaderState { Command::BitfieldInfo { peer, hash, - event: BitfieldEvent::State { ranges, size }, + event: BitfieldEvent::State(BitfieldState { ranges, size }), } => { let state = self.bitfields.get_mut(&(peer, hash)).context(format!( "bitfields for unknown peer {peer:?} and hash {hash}" @@ -327,7 +327,14 @@ impl DownloaderState { 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: BitfieldEvent::State { ranges: ranges.clone(), size } }); + evs.push(Event::LocalBitfieldInfo { + id: *id, + event: BitfieldState { + ranges: ranges.clone(), + size, + } + .into(), + }); } } state.ranges = ranges; @@ -349,7 +356,12 @@ impl DownloaderState { Command::BitfieldInfo { peer, hash, - event: BitfieldEvent::Update { added, removed, size }, + event: + BitfieldEvent::Update(BitfieldUpdate { + added, + removed, + size, + }), } => { let state = self.bitfields.get_mut(&(peer, hash)).context(format!( "bitfield update for unknown peer {peer:?} and hash {hash}" @@ -365,11 +377,12 @@ impl DownloaderState { if !added.is_empty() || !removed.is_empty() { evs.push(Event::LocalBitfieldInfo { id: *id, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: &added & &request.ranges, removed: &removed & &request.ranges, size: state.size, } + .into(), }); } } @@ -459,10 +472,11 @@ impl DownloaderState { // emit the current bitfield evs.push(Event::LocalBitfieldInfo { id, - event: BitfieldEvent::State { + event: BitfieldState { ranges: state.ranges.clone(), size: u64::MAX, } + .into(), }); } else { // create a new subscription @@ -955,10 +969,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: initial_bitfield.clone(), size: u64::MAX, - }, + } + .into(), }); assert!(evs.is_empty()); assert_eq!( @@ -982,11 +997,12 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([16..32]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); assert!(evs.is_empty()); assert_eq!( @@ -1013,10 +1029,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: chunk_ranges([0..64]), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event( @@ -1042,11 +1059,12 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([32..48]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); assert!(evs.is_empty()); // ChunksDownloaded just updates the peer stats @@ -1061,11 +1079,12 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([48..64]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::DownloadComplete { .. })), @@ -1098,10 +1117,11 @@ mod tests { state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); @@ -1109,10 +1129,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: chunk_ranges([0..32]), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event( @@ -1137,20 +1158,22 @@ mod tests { state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([0..16]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); // The peer now has more data state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: chunk_ranges([32..64]), size: u64::MAX, - }, + } + .into(), }); // ChunksDownloaded just updates the peer stats state.apply(Command::ChunksDownloaded { @@ -1163,11 +1186,12 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([16..32]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); // This triggers cancellation of the first peer download and starting a new one for the remaining data assert!( @@ -1197,11 +1221,12 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::Update { + event: BitfieldUpdate { added: chunk_ranges([32..64]), removed: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); assert!( has_all_events( @@ -1267,10 +1292,11 @@ mod tests { let evs1 = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: chunk_ranges([0..32]), size: u64::MAX, - }, + } + .into(), }); // No completion event should be generated for download0 because its full range 0..64 is not yet met. assert!( @@ -1331,10 +1357,11 @@ mod tests { state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: ChunkRanges::empty(), size: u64::MAX, - }, + } + .into(), }); // We have a peer for the hash state.apply(Command::PeerDiscovered { peer: peer_a, hash }); @@ -1342,10 +1369,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Remote(peer_a), hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: chunk_ranges([0..32]), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event( @@ -1391,10 +1419,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: ChunkRanges::all(), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), @@ -1403,10 +1432,11 @@ mod tests { let evs = state.apply(Command::BitfieldInfo { peer: Local, hash: unknown_hash, - event: BitfieldEvent::State { + event: BitfieldState { ranges: ChunkRanges::all(), size: u64::MAX, - }, + } + .into(), }); assert!( has_one_event_matching(&evs, |e| matches!(e, Event::Error { .. })), diff --git a/src/get/db.rs b/src/get/db.rs index b165ff503..1a80bba7b 100644 --- a/src/get/db.rs +++ b/src/get/db.rs @@ -242,7 +242,9 @@ pub async fn valid_ranges(entry: &D::EntryMut) -> anyhow::Result(entry: &D::EntryMut) -> anyhow::Result<(ChunkRanges, u64)> { +pub async fn valid_ranges_and_size( + 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?; From 18d76884903ae1f3627d43286c3eb80d8681143f Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Wed, 12 Feb 2025 13:55:35 +0200 Subject: [PATCH 41/47] First somewhat working size handling approach --- examples/multiprovider.rs | 35 ++++++++++++++++++++++------------- src/downloader2.rs | 22 ++++++++++++---------- src/downloader2/state.rs | 21 +++++++++------------ 3 files changed, 43 insertions(+), 35 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index e3f54ad17..380c6c3e7 100644 --- a/examples/multiprovider.rs +++ b/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,14 +180,14 @@ async fn download_impl(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?; @@ -187,8 +195,9 @@ async fn download_impl(args: DownloadArgs, store: S) -> anyhow::Result 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() { diff --git a/src/downloader2.rs b/src/downloader2.rs index bbc791f95..fbb4b675c 100644 --- a/src/downloader2.rs +++ b/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(hash: &Hash, store: S) -> anyhow::Resu let (ranges, size) = crate::get::db::valid_ranges_and_size::(&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 BitfieldSubscription for SimpleBitfieldSubscription { async move { let event = match recv.await { Ok(ev) => ev, - Err(_) => BitfieldState { - ranges: ChunkRanges::empty(), - size: u64::MAX, - } - .into(), + Err(_) => BitfieldState::unknown().into(), }; event } diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index 55ff352ca..bf0237ff5 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -324,6 +324,7 @@ 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; @@ -331,14 +332,13 @@ impl DownloaderState { 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 }); From 54ec3e331daed54a1c245caa2d5fc0fda0fffbe1 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 13 Feb 2025 11:05:44 +0200 Subject: [PATCH 42/47] Even more size refactoring --- examples/multiprovider.rs | 24 ++++----- src/downloader2.rs | 111 ++++++++++++++++++++++++++++++++------ src/downloader2/actor.rs | 18 +++++-- src/downloader2/state.rs | 47 ++++++++-------- src/get/db.rs | 12 +---- src/util.rs | 3 +- 6 files changed, 150 insertions(+), 65 deletions(-) diff --git a/examples/multiprovider.rs b/examples/multiprovider.rs index 380c6c3e7..a861f6566 100644 --- a/examples/multiprovider.rs +++ b/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(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() { diff --git a/src/downloader2.rs b/src/downloader2.rs index fbb4b675c..6ae1dd1a0 100644 --- a/src/downloader2.rs +++ b/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; +/// 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 { + 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 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,6 +156,15 @@ 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 { @@ -105,17 +172,8 @@ pub struct BitfieldUpdate { 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 { @@ -123,7 +181,7 @@ impl BitfieldState { 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 SimpleBitfieldSubscription { async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { if let Some(entry) = store.get_mut(hash).await? { - let (ranges, size) = crate::get::db::valid_ranges_and_size::(&entry).await?; + let ranges = crate::get::db::valid_ranges::(&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 BitfieldSubscription for SimpleBitfieldSubscription { diff --git a/src/downloader2/actor.rs b/src/downloader2/actor.rs index 3420d0aec..c4c9828fd 100644 --- a/src/downloader2/actor.rs +++ b/src/downloader2/actor.rs @@ -267,9 +267,16 @@ async fn peer_download( 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( .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, diff --git a/src/downloader2/state.rs b/src/downloader2/state.rs index bf0237ff5..41c1770dc 100644 --- a/src/downloader2/state.rs +++ b/src/downloader2/state.rs @@ -324,7 +324,7 @@ 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); + state.size.update(size)?; if let Some(observers) = self.observers.get_by_hash(&hash) { for (id, request) in observers { let ranges = &ranges & &request.ranges; @@ -348,7 +348,7 @@ impl DownloaderState { } } state.ranges = ranges; - state.size = state.size.min(size); + state.size.update(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)?; @@ -369,7 +369,7 @@ impl DownloaderState { 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); + state.size.update(size)?; if let Some(observers) = self.observers.get_by_hash(&hash) { for (id, request) in observers { let added = &added & &request.ranges; @@ -400,7 +400,7 @@ impl DownloaderState { } state.ranges |= added; state.ranges &= !removed; - state.size = state.size.min(size); + state.size.update(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)?; } @@ -474,7 +474,7 @@ impl DownloaderState { id, event: BitfieldState { ranges: state.ranges.clone(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -532,7 +532,10 @@ 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 mask = match self_state.size { + BaoBlobSizeOpt::Verified(size) => ChunkRanges::from(..ChunkNum::chunks(size)), + _ => ChunkRanges::all(), + }; let mut completed = vec![]; for (id, download) in self.downloads.iter_mut_for_hash(hash) { if just_id.is_some() && just_id != Some(*id) { @@ -899,7 +902,7 @@ struct PeerBlobState { /// chunk ranges this peer reports to have ranges: ChunkRanges, /// The minimum reported size of the blob - size: u64, + size: BaoBlobSizeOpt, } impl PeerBlobState { @@ -908,7 +911,7 @@ impl PeerBlobState { subscription_id, subscription_count: 1, ranges: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } } } @@ -976,7 +979,7 @@ mod tests { hash, event: BitfieldState { ranges: initial_bitfield.clone(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1005,7 +1008,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([16..32]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1036,7 +1039,7 @@ mod tests { hash, event: BitfieldState { ranges: chunk_ranges([0..64]), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1067,7 +1070,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([32..48]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1087,7 +1090,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([48..64]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1132,7 +1135,7 @@ mod tests { hash, event: BitfieldState { ranges: chunk_ranges([0..32]), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1162,7 +1165,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([0..16]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1172,7 +1175,7 @@ mod tests { hash, event: BitfieldState { ranges: chunk_ranges([32..64]), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1190,7 +1193,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([16..32]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1225,7 +1228,7 @@ mod tests { event: BitfieldUpdate { added: chunk_ranges([32..64]), removed: ChunkRanges::empty(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1295,7 +1298,7 @@ mod tests { hash, event: BitfieldState { ranges: chunk_ranges([0..32]), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1368,7 +1371,7 @@ mod tests { hash, event: BitfieldState { ranges: chunk_ranges([0..32]), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1418,7 +1421,7 @@ mod tests { hash, event: BitfieldState { ranges: ChunkRanges::all(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); @@ -1431,7 +1434,7 @@ mod tests { hash: unknown_hash, event: BitfieldState { ranges: ChunkRanges::all(), - size: u64::MAX, + size: BaoBlobSizeOpt::Unknown, } .into(), }); diff --git a/src/get/db.rs b/src/get/db.rs index 1a80bba7b..414ad9ff9 100644 --- a/src/get/db.rs +++ b/src/get/db.rs @@ -237,14 +237,6 @@ async fn get_blob( /// Given a partial entry, get the valid ranges. pub async fn valid_ranges(entry: &D::EntryMut) -> anyhow::Result { - let (ranges, _) = valid_ranges_and_size::(entry).await?; - Ok(ranges) -} - -/// Given a partial entry, get the valid ranges. -pub async fn valid_ranges_and_size( - 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?; @@ -260,8 +252,8 @@ pub async fn valid_ranges_and_size( } 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, data_size)) + log!("valid_from_outboard: {:?}", valid_from_outboard); + Ok(valid) } /// Get a blob that was requested completely. diff --git a/src/util.rs b/src/util.rs index 6e3d29b62..db2edc2a7 100644 --- a/src/util.rs +++ b/src/util.rs @@ -238,7 +238,8 @@ impl Drop for TempTag { /// Get the number of bytes given a set of chunk ranges and the total size. /// /// If some ranges are out of bounds, they will be clamped to the size. -pub fn total_bytes(ranges: &ChunkRanges, size: u64) -> u64 { +pub fn total_bytes(ranges: &ChunkRanges, size: Option) -> u64 { + let size = size.unwrap_or(u64::MAX); ranges .iter() .map(|range| { From 419f8d8b5a851f4ae0859185b5ca2192eacc5c4e Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 21 Feb 2025 11:00:55 +0200 Subject: [PATCH 43/47] Move entire multiprovider downloader into an example. This is a bit unfortunate, but there are major refactorings coming for the store, and the current code has to fake a part of the functionality that is needed from the store. This way people can take a look at the code structure and also try it out. Once the main store changes are merged, the multiprovider downloader will move into the main crate. --- {src => examples/multiprovider}/downloader2.rs | 8 ++++---- {src => examples/multiprovider}/downloader2/actor.rs | 4 ++-- .../multiprovider}/downloader2/content_discovery.rs | 0 {src => examples/multiprovider}/downloader2/planners.rs | 0 {src => examples/multiprovider}/downloader2/state.rs | 0 examples/{multiprovider.rs => multiprovider/main.rs} | 9 +++++---- src/lib.rs | 2 -- 7 files changed, 11 insertions(+), 12 deletions(-) rename {src => examples/multiprovider}/downloader2.rs (99%) rename {src => examples/multiprovider}/downloader2/actor.rs (99%) rename {src => examples/multiprovider}/downloader2/content_discovery.rs (100%) rename {src => examples/multiprovider}/downloader2/planners.rs (100%) rename {src => examples/multiprovider}/downloader2/state.rs (100%) rename examples/{multiprovider.rs => multiprovider/main.rs} (97%) diff --git a/src/downloader2.rs b/examples/multiprovider/downloader2.rs similarity index 99% rename from src/downloader2.rs rename to examples/multiprovider/downloader2.rs index 6ae1dd1a0..0a7373b10 100644 --- a/src/downloader2.rs +++ b/examples/multiprovider/downloader2.rs @@ -21,7 +21,7 @@ use std::{ u64, }; -use crate::{ +use iroh_blobs::{ get::{ fsm::{BlobContentNext, ConnectedNext, EndBlobNext}, Stats, @@ -433,7 +433,7 @@ impl SimpleBitfieldSubscription { async fn get_valid_ranges_local(hash: &Hash, store: S) -> anyhow::Result { if let Some(entry) = store.get_mut(hash).await? { - let ranges = crate::get::db::valid_ranges::(&entry).await?; + let ranges = iroh_blobs::get::db::valid_ranges::(&entry).await?; let size = entry.size(); let size = match size { size @ BaoBlobSize::Unverified(value) => { @@ -462,8 +462,8 @@ async fn get_valid_ranges_remote( id: NodeId, hash: &Hash, ) -> anyhow::Result { - let conn = endpoint.connect(id, crate::ALPN).await?; - let (size, _) = crate::get::request::get_verified_size(&conn, hash).await?; + let conn = endpoint.connect(id, iroh_blobs::ALPN).await?; + let (size, _) = iroh_blobs::get::request::get_verified_size(&conn, hash).await?; let chunks = (size + 1023) / 1024; let ranges = ChunkRanges::from(ChunkNum(0)..ChunkNum(chunks)); Ok(BitfieldState { diff --git a/src/downloader2/actor.rs b/examples/multiprovider/downloader2/actor.rs similarity index 99% rename from src/downloader2/actor.rs rename to examples/multiprovider/downloader2/actor.rs index c4c9828fd..f457abf76 100644 --- a/src/downloader2/actor.rs +++ b/examples/multiprovider/downloader2/actor.rs @@ -238,13 +238,13 @@ async fn peer_download( start: Instant, ) -> anyhow::Result { info!("Connecting to peer {peer}"); - let conn = endpoint.connect(peer, crate::ALPN).await?; + let conn = endpoint.connect(peer, iroh_blobs::ALPN).await?; info!("Got connection to peer {peer}"); let spec = RangeSpec::new(ranges); let ranges = RangeSpecSeq::new([spec, RangeSpec::EMPTY]); info!("starting download from {peer} for {hash} {ranges:?}"); let request = GetRequest::new(hash, ranges); - let initial = crate::get::fsm::start(conn, request); + let initial = iroh_blobs::get::fsm::start(conn, request); // connect let connected = initial.next().await?; // read the first bytes diff --git a/src/downloader2/content_discovery.rs b/examples/multiprovider/downloader2/content_discovery.rs similarity index 100% rename from src/downloader2/content_discovery.rs rename to examples/multiprovider/downloader2/content_discovery.rs diff --git a/src/downloader2/planners.rs b/examples/multiprovider/downloader2/planners.rs similarity index 100% rename from src/downloader2/planners.rs rename to examples/multiprovider/downloader2/planners.rs diff --git a/src/downloader2/state.rs b/examples/multiprovider/downloader2/state.rs similarity index 100% rename from src/downloader2/state.rs rename to examples/multiprovider/downloader2/state.rs diff --git a/examples/multiprovider.rs b/examples/multiprovider/main.rs similarity index 97% rename from examples/multiprovider.rs rename to examples/multiprovider/main.rs index a861f6566..65c073bb0 100644 --- a/examples/multiprovider.rs +++ b/examples/multiprovider/main.rs @@ -4,11 +4,12 @@ use bao_tree::{ChunkNum, ChunkRanges}; use clap::Parser; use console::Term; use iroh::{NodeId, SecretKey}; +mod downloader2; +use downloader2::{ + print_bitmap, BitfieldEvent, BitfieldState, BitfieldUpdate, DownloadRequest, Downloader, + ObserveRequest, StaticContentDiscovery, +}; use iroh_blobs::{ - downloader2::{ - print_bitmap, BitfieldEvent, BitfieldState, BitfieldUpdate, DownloadRequest, Downloader, - ObserveRequest, StaticContentDiscovery, - }, store::Store, Hash, }; diff --git a/src/lib.rs b/src/lib.rs index 5901bc93b..7091ad795 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,8 +32,6 @@ pub mod cli; #[cfg(feature = "downloader")] pub mod downloader; -#[cfg(feature = "downloader")] -pub mod downloader2; pub mod export; pub mod format; pub mod get; From 19050894b6b8f98b2b6af50d62e7c6c0a04cc87a Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 21 Feb 2025 11:45:21 +0200 Subject: [PATCH 44/47] fmt & clippy --- examples/multiprovider/downloader2.rs | 26 +++++++++---------- .../downloader2/content_discovery.rs | 4 ++- examples/multiprovider/downloader2/state.rs | 11 +++++--- examples/multiprovider/main.rs | 5 +--- 4 files changed, 23 insertions(+), 23 deletions(-) diff --git a/examples/multiprovider/downloader2.rs b/examples/multiprovider/downloader2.rs index 0a7373b10..49dc90f98 100644 --- a/examples/multiprovider/downloader2.rs +++ b/examples/multiprovider/downloader2.rs @@ -17,10 +17,14 @@ use std::{ io, marker::PhantomData, sync::Arc, - time::Instant, - u64, + time::{Duration, Instant}, }; +use anyhow::Context; +use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; +use futures_lite::StreamExt; +use futures_util::{stream::BoxStream, FutureExt}; +use iroh::{Endpoint, NodeId}; use iroh_blobs::{ get::{ fsm::{BlobContentNext, ConnectedNext, EndBlobNext}, @@ -31,14 +35,8 @@ use iroh_blobs::{ util::local_pool::{self, LocalPool, LocalPoolHandle}, Hash, }; -use anyhow::Context; -use bao_tree::{io::BaoContentItem, ChunkNum, ChunkRanges}; -use futures_lite::StreamExt; -use futures_util::{stream::BoxStream, FutureExt}; -use iroh::{Endpoint, NodeId}; use range_collections::range_set::RangeSetRange; use serde::{Deserialize, Serialize}; -use std::time::Duration; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, error, info, trace}; @@ -259,6 +257,7 @@ pub struct DownloaderBuilder { planner: Option, } +#[allow(dead_code)] impl DownloaderBuilder { /// Set the content discovery pub fn discovery(self, discovery: D) -> Self { @@ -506,11 +505,10 @@ impl BitfieldSubscription for SimpleBitfieldSubscription { } Box::pin( async move { - let event = match recv.await { + match recv.await { Ok(ev) => ev, Err(_) => BitfieldState::unknown().into(), - }; - event + } } .into_stream(), ) @@ -591,13 +589,13 @@ mod tests { #![allow(clippy::single_range_in_vec_init)] use std::ops::Range; - use crate::{net_protocol::Blobs, store::MapMut}; - - use super::*; use bao_tree::ChunkNum; use iroh::{protocol::Router, SecretKey}; use testresult::TestResult; + use super::*; + use crate::{net_protocol::Blobs, store::MapMut}; + fn print_bitfield(iter: impl IntoIterator) -> String { let mut chars = String::new(); for x in iter { diff --git a/examples/multiprovider/downloader2/content_discovery.rs b/examples/multiprovider/downloader2/content_discovery.rs index a09f04ea9..1be36349d 100644 --- a/examples/multiprovider/downloader2/content_discovery.rs +++ b/examples/multiprovider/downloader2/content_discovery.rs @@ -1,11 +1,12 @@ use std::collections::BTreeMap; -use crate::Hash; use futures_lite::stream::StreamExt; use futures_util::stream::BoxStream; use iroh::NodeId; use serde::{Deserialize, Serialize}; +use crate::Hash; + /// Announce kind #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)] pub enum AnnounceKind { @@ -20,6 +21,7 @@ pub enum AnnounceKind { #[derive(Debug, Default)] pub struct FindPeersOpts { /// Kind of announce + #[allow(dead_code)] pub kind: AnnounceKind, } diff --git a/examples/multiprovider/downloader2/state.rs b/examples/multiprovider/downloader2/state.rs index 41c1770dc..1da59c2b6 100644 --- a/examples/multiprovider/downloader2/state.rs +++ b/examples/multiprovider/downloader2/state.rs @@ -931,12 +931,15 @@ fn total_chunks(chunks: &ChunkRanges) -> Option { mod tests { #![allow(clippy::single_range_in_vec_init)] - use super::super::tests::{ - chunk_ranges, has_all_events, has_one_event, has_one_event_matching, noop_planner, - }; - use super::*; use testresult::TestResult; + use super::{ + super::tests::{ + chunk_ranges, has_all_events, has_one_event, has_one_event_matching, noop_planner, + }, + *, + }; + /// Test a simple scenario where a download is started and completed #[test] fn downloader_state_smoke() -> TestResult<()> { diff --git a/examples/multiprovider/main.rs b/examples/multiprovider/main.rs index 65c073bb0..387da4dfb 100644 --- a/examples/multiprovider/main.rs +++ b/examples/multiprovider/main.rs @@ -9,10 +9,7 @@ use downloader2::{ print_bitmap, BitfieldEvent, BitfieldState, BitfieldUpdate, DownloadRequest, Downloader, ObserveRequest, StaticContentDiscovery, }; -use iroh_blobs::{ - store::Store, - Hash, -}; +use iroh_blobs::{store::Store, Hash}; #[derive(Debug, Parser)] struct Args { From c9146e95fe36c005b0a7ab3eff8f282534b3b556 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 21 Feb 2025 12:00:07 +0200 Subject: [PATCH 45/47] codespell --- examples/multiprovider/downloader2.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/multiprovider/downloader2.rs b/examples/multiprovider/downloader2.rs index 49dc90f98..5df1b2cb9 100644 --- a/examples/multiprovider/downloader2.rs +++ b/examples/multiprovider/downloader2.rs @@ -87,7 +87,7 @@ pub trait BitfieldSubscription: std::fmt::Debug + Send + 'static { /// A boxed bitfield subscription pub type BoxedBitfieldSubscription = Box; -/// Knowlege about the size of a blob +/// Knowledge 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 @@ -113,7 +113,7 @@ impl BaoBlobSizeOpt { /// Update the size information /// - /// Unkown sizes are always updated + /// Unknown 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<()> { From 177dfcef2c878f75643a8f7f8877a44b857babb7 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 21 Feb 2025 13:00:46 +0200 Subject: [PATCH 46/47] Increase stripe size and add readme --- examples/multiprovider/README.md | 92 +++++++++++++++++++++++++++ examples/multiprovider/downloader2.rs | 2 +- 2 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 examples/multiprovider/README.md diff --git a/examples/multiprovider/README.md b/examples/multiprovider/README.md new file mode 100644 index 000000000..a3d9cc5a2 --- /dev/null +++ b/examples/multiprovider/README.md @@ -0,0 +1,92 @@ +# Multiprovider + +This example shows how to use iroh-blobs to download concurrently from multiple +providers. As of now, the infrastructure to do this is included in the example. +It will move into the main crate soon. + +## Usage + +This example requires the `rpc` feature, so is is easiest to run it with +`--all-features`. Also, if you want try it out with large blobs such as ML +models, it is best to run in release mode. + +There are two subcommands, `provide` and `download`. + +### Provide + +Provide provides a *single* blob, printing the blob hash and the node id. + +### Download + +Download downloads a *single* hash from any number of node ids. + +In the long +term we are going to have content discovery based on trackers or other mechanisms, +but for this example you just have to provide the node ids in the command line. + +To have a stable node id, it is +possible to provide the iroh node secret in an environment variable. + +**This is fine for an example, but don't do it in production** + +## Trying it out + +Multiprovider downloads are mostly relevant for very large downloads, so let's +use a large file, a ~4GB ML model. + +Terminal 1: + +``` +> IROH_SECRET= \ + cargo run --release --all-features --example multiprovider \ + provide ~/.ollama/models/blobs/sha256-96c415656d377afbff962f6cdb2394ab092ccbcbaab4b82525bc4ca800fe8a49 +added /Users/rklaehn/.ollama/models/blobs/sha256-96c415656d377afbff962f6cdb2394ab092ccbcbaab4b82525bc4ca800fe8a49 as e5njueepdum3ks2usqdxw3ofztj63jgedtnfak34smgvw5b6cr3a, 4683073184 bytes, 4573314 chunks +listening on 28300fcb69830c3e094c68f383ffd568dd9aa9126a6aa537c3dcfec077b60af9 +``` + +Terminal 2: + +``` +❯ IROH_SECRET= \ + cargo run --release --all-features --example multiprovider \ + provide ~/.ollama/models/blobs/sha256-96c415656d377afbff962f6cdb2394ab092ccbcbaab4b82525bc4ca800fe8a49 +added /Users/rklaehn/.ollama/models/blobs/sha256-96c415656d377afbff962f6cdb2394ab092ccbcbaab4b82525bc4ca800fe8a49 as e5njueepdum3ks2usqdxw3ofztj63jgedtnfak34smgvw5b6cr3a, 4683073184 bytes, 4573314 chunks +listening on 77d81595422c0a757b9e3f739f9a67eab9646f13d941654e9074982c5c800a5a +``` + +So now we got 2 node ids, +`77d81595422c0a757b9e3f739f9a67eab9646f13d941654e9074982c5c800a5a` and +`28300fcb69830c3e094c68f383ffd568dd9aa9126a6aa537c3dcfec077b60af9`, providing +the data. + +Note that the provide side is not in any way special. It is just using the +existing iroh-blobs protocol, so any other iroh node could be used as well. + +For downloading, we don't need a stable node id, so we don't need to bother with +setting IROH_SECRET. + +``` +> cargo run --release --all-features --example multiprovider \ + download e5njueepdum3ks2usqdxw3ofztj63jgedtnfak34smgvw5b6cr3a \ + 28300fcb69830c3e094c68f383ffd568dd9aa9126a6aa537c3dcfec077b60af9 \ + 77d81595422c0a757b9e3f739f9a67eab9646f13d941654e9074982c5c800a5a + +peer discovered for hash e5njueepdum3ks2usqdxw3ofztj63jgedtnfak34smgvw5b6cr3a: 28300fcb69830c3e094c68f383ffd568dd9aa9126a6aa537c3dcfec077b60af9 +peer discovered for hash e5njueepdum3ks2usqdxw3ofztj63jgedtnfak34smgvw5b6cr3a: 77d81595422c0a757b9e3f739f9a67eab9646f13d941654e9074982c5c800a5a +█████████▓ ░█████████░ +``` + +The download side will initially download from the first peer, then quickly +rebalance the download as a new peer becomes available. It will currently +download from each peer in "stripes". + +When running without `--path` argument it will download into a memory store. +When providing a `--path` argument it will download into a persitent store at the +given path, and the download will resume if you interrupt the download process. + +## Notes on the current state + +The current state of the downloader is highly experimental. While peers that don't +respond at all are handled properly, peers that are slow or become slow over time +are not properly punished. Also, there is not yet a mechanism to limit the number +of peers to download from. \ No newline at end of file diff --git a/examples/multiprovider/downloader2.rs b/examples/multiprovider/downloader2.rs index 5df1b2cb9..11be870e1 100644 --- a/examples/multiprovider/downloader2.rs +++ b/examples/multiprovider/downloader2.rs @@ -301,7 +301,7 @@ impl DownloaderBuilder { let local_pool = self.local_pool.unwrap_or_else(LocalPool::single); let planner = self .planner - .unwrap_or_else(|| Box::new(StripePlanner2::new(0, 10))); + .unwrap_or_else(|| Box::new(StripePlanner2::new(0, 18))); let subscribe_bitfield = self.subscribe_bitfield.unwrap_or_else(|| { Box::new(SimpleBitfieldSubscription::new( self.endpoint.clone(), From e5d231a4d99cb7f683e92ee7d83951bc933360e9 Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Fri, 21 Feb 2025 13:38:59 +0200 Subject: [PATCH 47/47] more limitations --- examples/multiprovider/README.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/examples/multiprovider/README.md b/examples/multiprovider/README.md index a3d9cc5a2..daa6036ac 100644 --- a/examples/multiprovider/README.md +++ b/examples/multiprovider/README.md @@ -81,7 +81,7 @@ rebalance the download as a new peer becomes available. It will currently download from each peer in "stripes". When running without `--path` argument it will download into a memory store. -When providing a `--path` argument it will download into a persitent store at the +When providing a `--path` argument it will download into a persistent store at the given path, and the download will resume if you interrupt the download process. ## Notes on the current state @@ -89,4 +89,8 @@ given path, and the download will resume if you interrupt the download process. The current state of the downloader is highly experimental. While peers that don't respond at all are handled properly, peers that are slow or become slow over time are not properly punished. Also, there is not yet a mechanism to limit the number -of peers to download from. \ No newline at end of file +of peers to download from. + +In addition, the current blob store does not have the ability to watch a bitfield +of available chunks for a blob. The current multiprovider downloader just fakes +this by assuming that all remote stores have the full file. \ No newline at end of file