[6/n] simplify FetchArtifactImpl interface (#8040)

sunshowers · web-flow · commit 780981103c17 · 2025-04-26T00:02:07.000Z
diff --git a/installinator/src/dispatch.rs b/installinator/src/dispatch.rs
@@ -91,15 +91,15 @@ struct DebugDiscoverOpts {
 
 impl DebugDiscoverOpts {
     async fn exec(self, log: &slog::Logger) -> Result<()> {
-        let peers = FetchArtifactBackend::new(
+        let backend = FetchArtifactBackend::new(
             log,
             Box::new(HttpFetchBackend::new(
                 &log,
                 self.opts.mechanism.discover_peers(&log).await?,
             )),
             Duration::from_secs(10),
         );
-        println!("discovered peers: {}", peers.display());
+        println!("discovered peers: {}", backend.peers().display());
         Ok(())
     }
 }
diff --git a/installinator/src/fetch.rs b/installinator/src/fetch.rs
@@ -15,15 +15,14 @@ use installinator_client::ClientError;
 use installinator_common::{
     InstallinatorProgressMetadata, StepContext, StepProgress,
 };
-use itertools::Itertools;
 use tokio::{sync::mpsc, time::Instant};
 use tufaceous_artifact::ArtifactHashId;
 use update_engine::events::ProgressUnits;
 
 use crate::{
     artifact::ArtifactClient,
     errors::{ArtifactFetchError, DiscoverPeersError, HttpError},
-    peers::PeerAddress,
+    peers::{PeerAddress, PeerAddresses},
 };
 
 /// A fetched artifact.
@@ -55,7 +54,7 @@ impl FetchedArtifact {
         let mut attempt = 0;
         loop {
             attempt += 1;
-            let peers = match discover_fn().await {
+            let fetch_backend = match discover_fn().await {
                 Ok(peers) => peers,
                 Err(DiscoverPeersError::Retry(error)) => {
                     slog::warn!(
@@ -78,13 +77,14 @@ impl FetchedArtifact {
                 }
             };
 
+            let peers = fetch_backend.peers();
             slog::info!(
                 log,
                 "discovered {} peers: [{}]",
-                peers.peer_count(),
+                peers.len(),
                 peers.display(),
             );
-            match peers.fetch_artifact(&cx, artifact_hash_id).await {
+            match fetch_backend.fetch_artifact(&cx, artifact_hash_id).await {
                 Some((peer, artifact)) => {
                     return Ok(Self { attempt, peer, artifact });
                 }
@@ -95,7 +95,7 @@ impl FetchedArtifact {
                     );
                     cx.send_progress(StepProgress::retry(format!(
                         "unable to fetch artifact from any of {} peers, retrying",
-                        peers.peer_count(),
+                        peers.len(),
                     )))
                     .await;
                     tokio::time::sleep(RETRY_DELAY).await;
@@ -146,15 +146,15 @@ impl FetchArtifactBackend {
     ) -> Option<(PeerAddress, BufList)> {
         // TODO: do we want a check phase that happens before the download?
         let peers = self.peers();
-        let mut remaining_peers = self.peer_count();
+        let mut remaining_peers = peers.len();
 
         let log = self.log.new(
             slog::o!("artifact_hash_id" => format!("{artifact_hash_id:?}")),
         );
 
         slog::debug!(log, "start fetch from peers"; "remaining_peers" => remaining_peers);
 
-        for peer in peers {
+        for &peer in peers.peers() {
             remaining_peers -= 1;
 
             slog::debug!(
@@ -188,18 +188,10 @@ impl FetchArtifactBackend {
         None
     }
 
-    pub(crate) fn peers(&self) -> impl Iterator<Item = PeerAddress> + '_ {
+    pub(crate) fn peers(&self) -> &PeerAddresses {
         self.imp.peers()
     }
 
-    pub(crate) fn peer_count(&self) -> usize {
-        self.imp.peer_count()
-    }
-
-    pub(crate) fn display(&self) -> impl fmt::Display {
-        self.peers().join(", ")
-    }
-
     async fn fetch_from_peer(
         &self,
         cx: &StepContext,
@@ -307,10 +299,16 @@ impl FetchArtifactBackend {
     }
 }
 
+/// Backend implementation for fetching artifacts.
+///
+/// Note: While [`crate::reporter::ReportProgressImpl`] is a persistent
+/// structure, a new `FetchArtifactImpl` is generated separately each time
+/// discovery occurs. We should align this with `ReportProgressImpl` in the
+/// future, though we'd need some way of looking up delay information in
+/// `mock_peers`.
 #[async_trait]
 pub(crate) trait FetchArtifactImpl: fmt::Debug + Send + Sync {
-    fn peers(&self) -> Box<dyn Iterator<Item = PeerAddress> + Send + '_>;
-    fn peer_count(&self) -> usize;
+    fn peers(&self) -> &PeerAddresses;
 
     /// Returns (size, receiver) on success, and an error on failure.
     async fn fetch_from_peer_impl(
@@ -329,24 +327,20 @@ pub(crate) type FetchReceiver = mpsc::Receiver<Result<Bytes, ClientError>>;
 #[derive(Clone, Debug)]
 pub(crate) struct HttpFetchBackend {
     log: slog::Logger,
-    peers: Vec<PeerAddress>,
+    peers: PeerAddresses,
 }
 
 impl HttpFetchBackend {
-    pub(crate) fn new(log: &slog::Logger, peers: Vec<PeerAddress>) -> Self {
+    pub(crate) fn new(log: &slog::Logger, peers: PeerAddresses) -> Self {
         let log = log.new(slog::o!("component" => "HttpPeers"));
         Self { log, peers }
     }
 }
 
 #[async_trait]
 impl FetchArtifactImpl for HttpFetchBackend {
-    fn peers(&self) -> Box<dyn Iterator<Item = PeerAddress> + Send + '_> {
-        Box::new(self.peers.iter().copied())
-    }
-
-    fn peer_count(&self) -> usize {
-        self.peers.len()
+    fn peers(&self) -> &PeerAddresses {
+        &self.peers
     }
 
     async fn fetch_from_peer_impl(
diff --git a/installinator/src/mock_peers.rs b/installinator/src/mock_peers.rs
@@ -28,7 +28,7 @@ use uuid::Uuid;
 use crate::{
     errors::{DiscoverPeersError, HttpError},
     fetch::{FetchArtifactImpl, FetchReceiver},
-    peers::PeerAddress,
+    peers::{PeerAddress, PeerAddresses},
     reporter::ReportProgressImpl,
 };
 
@@ -143,10 +143,10 @@ impl MockPeersUniverse {
                                     .then(|| (*addr, peer.clone()))
                             })
                             .collect();
-                        Ok(MockFetchBackend {
-                            artifact: self.artifact.clone(),
+                        Ok(MockFetchBackend::new(
+                            self.artifact.clone(),
                             selected_peers,
-                        })
+                        ))
                     }
                     AttemptBitmap::Failure => {
                         bail!(
@@ -177,20 +177,27 @@ struct MockFetchBackend {
     artifact: Bytes,
     // Peers within the universe that have been selected
     selected_peers: BTreeMap<PeerAddress, MockPeer>,
+    // selected_peers keys stored in a suitable form for the
+    // FetchArtifactImpl trait
+    peer_addresses: PeerAddresses,
 }
 
 impl MockFetchBackend {
-    fn get(&self, peer: PeerAddress) -> Option<&MockPeer> {
-        self.selected_peers.get(&peer)
+    fn new(
+        artifact: Bytes,
+        selected_peers: BTreeMap<PeerAddress, MockPeer>,
+    ) -> Self {
+        let peer_addresses = selected_peers.keys().copied().collect();
+        Self { artifact, selected_peers, peer_addresses }
     }
 
-    fn peers(&self) -> impl Iterator<Item = (&PeerAddress, &MockPeer)> + '_ {
-        self.selected_peers.iter()
+    fn get(&self, peer: PeerAddress) -> Option<&MockPeer> {
+        self.selected_peers.get(&peer)
     }
 
     /// Returns the peer that can return the entire dataset within the timeout.
     fn successful_peer(&self, timeout: Duration) -> Option<PeerAddress> {
-        self.peers()
+        self.selected_peers.iter()
             .filter_map(|(addr, peer)| {
                 if peer.artifact != self.artifact {
                     // We don't handle the case where the peer returns the wrong artifact yet.
@@ -235,12 +242,8 @@ impl MockFetchBackend {
 
 #[async_trait]
 impl FetchArtifactImpl for MockFetchBackend {
-    fn peers(&self) -> Box<dyn Iterator<Item = PeerAddress> + Send + '_> {
-        Box::new(self.selected_peers.keys().copied())
-    }
-
-    fn peer_count(&self) -> usize {
-        self.selected_peers.len()
+    fn peers(&self) -> &PeerAddresses {
+        &self.peer_addresses
     }
 
     async fn fetch_from_peer_impl(
@@ -478,8 +481,10 @@ impl MockProgressBackend {
 impl ReportProgressImpl for MockProgressBackend {
     async fn discover_peers(
         &self,
-    ) -> Result<Vec<PeerAddress>, DiscoverPeersError> {
-        Ok(vec![Self::VALID_PEER, Self::INVALID_PEER, Self::UNRESPONSIVE_PEER])
+    ) -> Result<PeerAddresses, DiscoverPeersError> {
+        Ok([Self::VALID_PEER, Self::INVALID_PEER, Self::UNRESPONSIVE_PEER]
+            .into_iter()
+            .collect())
     }
 
     async fn report_progress_impl(
diff --git a/installinator/src/peers.rs b/installinator/src/peers.rs
@@ -3,6 +3,7 @@
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
 use std::{
+    collections::BTreeSet,
     fmt,
     net::{AddrParseError, IpAddr, SocketAddr},
     str::FromStr,
@@ -31,7 +32,7 @@ impl DiscoveryMechanism {
     pub(crate) async fn discover_peers(
         &self,
         log: &slog::Logger,
-    ) -> Result<Vec<PeerAddress>, DiscoverPeersError> {
+    ) -> Result<PeerAddresses, DiscoverPeersError> {
         let peers = match self {
             Self::Bootstrap => {
                 // Note: we do not abort this process and instead keep retrying
@@ -61,7 +62,7 @@ impl DiscoveryMechanism {
                     })
                     .collect()
             }
-            Self::List(peers) => peers.clone(),
+            Self::List(peers) => peers.iter().copied().collect(),
         };
 
         Ok(peers)
@@ -100,6 +101,32 @@ impl FromStr for DiscoveryMechanism {
     }
 }
 
+#[derive(Clone, Debug)]
+pub(crate) struct PeerAddresses {
+    peers: BTreeSet<PeerAddress>,
+}
+
+impl PeerAddresses {
+    pub(crate) fn peers(&self) -> &BTreeSet<PeerAddress> {
+        &self.peers
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.peers.len()
+    }
+
+    pub(crate) fn display(&self) -> impl fmt::Display {
+        self.peers().iter().join(", ")
+    }
+}
+
+impl FromIterator<PeerAddress> for PeerAddresses {
+    fn from_iter<I: IntoIterator<Item = PeerAddress>>(iter: I) -> Self {
+        let peers = iter.into_iter().collect::<BTreeSet<_>>();
+        Self { peers }
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq)]
 #[cfg_attr(test, derive(test_strategy::Arbitrary))]
 pub(crate) struct PeerAddress {
diff --git a/installinator/src/reporter.rs b/installinator/src/reporter.rs
@@ -19,7 +19,7 @@ use uuid::Uuid;
 use crate::{
     artifact::ArtifactClient,
     errors::DiscoverPeersError,
-    peers::{DiscoveryMechanism, PeerAddress},
+    peers::{DiscoveryMechanism, PeerAddress, PeerAddresses},
 };
 
 #[derive(Debug)]
@@ -129,17 +129,18 @@ impl ProgressReporter {
             // that if two servers both say, only one of them will
             // deterministically get the update. Need to decide post-PVT1.
             let last_reported = report.last_seen;
-            let results: Vec<_> = futures::stream::iter(peers)
-                .map(|peer| {
-                    report_backend.send_report_to_peer(
-                        peer,
-                        update_id,
-                        report.clone(),
-                    )
-                })
-                .buffer_unordered(8)
-                .collect()
-                .await;
+            let results: Vec<_> =
+                futures::stream::iter(peers.peers().iter().copied())
+                    .map(|peer| {
+                        report_backend.send_report_to_peer(
+                            peer,
+                            update_id,
+                            report.clone(),
+                        )
+                    })
+                    .buffer_unordered(8)
+                    .collect()
+                    .await;
 
             if results.iter().any(|res| res.is_ok()) {
                 Some(last_reported)
@@ -188,7 +189,7 @@ impl ReportProgressBackend {
 
     pub(crate) async fn discover_peers(
         &self,
-    ) -> Result<Vec<PeerAddress>, DiscoverPeersError> {
+    ) -> Result<PeerAddresses, DiscoverPeersError> {
         let log = self.log.new(slog::o!("task" => "discover_peers"));
         slog::debug!(log, "discovering peers");
 
@@ -255,9 +256,8 @@ pub(crate) enum SendReportStatus {
 
 #[async_trait]
 pub(crate) trait ReportProgressImpl: fmt::Debug + Send + Sync {
-    async fn discover_peers(
-        &self,
-    ) -> Result<Vec<PeerAddress>, DiscoverPeersError>;
+    async fn discover_peers(&self)
+    -> Result<PeerAddresses, DiscoverPeersError>;
 
     async fn report_progress_impl(
         &self,
@@ -287,7 +287,7 @@ impl HttpProgressBackend {
 impl ReportProgressImpl for HttpProgressBackend {
     async fn discover_peers(
         &self,
-    ) -> Result<Vec<PeerAddress>, DiscoverPeersError> {
+    ) -> Result<PeerAddresses, DiscoverPeersError> {
         self.discovery.discover_peers(&self.log).await
     }
 
