[1/n] [installinator] add PeerAddress wrapper (#8035)

Author: sunshowers

installinator/src/dispatch.rs

@@ -24,7 +24,7 @@ use update_engine::StepResult;
 
 use crate::{
     artifact::ArtifactIdOpts,
-    peers::{DiscoveryMechanism, FetchedArtifact, Peers},
+    peers::{DiscoveryMechanism, FetchedArtifact, HttpPeers, Peers},
     reporter::ProgressReporter,
     write::{ArtifactWriter, WriteDestination},
 };
@@ -92,7 +92,10 @@ impl DebugDiscoverOpts {
     async fn exec(self, log: &slog::Logger) -> Result<()> {
         let peers = Peers::new(
             log,
-            self.opts.mechanism.discover_peers(log).await?,
+            Box::new(HttpPeers::new(
+                log,
+                self.opts.mechanism.discover_peers(log).await?,
+            )),
             Duration::from_secs(10),
         );
         println!("discovered peers: {}", peers.display());
@@ -190,7 +193,10 @@ impl InstallOpts {
                 async move {
                     Ok(Peers::new(
                         &log,
-                        discovery.discover_peers(&log).await?,
+                        Box::new(HttpPeers::new(
+                            &log,
+                            discovery.discover_peers(&log).await?,
+                        )),
                         Duration::from_secs(10),
                     ))
                 }
@@ -234,7 +240,7 @@ impl InstallOpts {
                     )
                     .await?;
 
-                    let address = host_phase_2_artifact.addr;
+                    let address = host_phase_2_artifact.peer.address();
 
                     StepSuccess::new(host_phase_2_artifact)
                         .with_metadata(
@@ -273,7 +279,7 @@ impl InstallOpts {
                     )
                     .await?;
 
-                    let address = control_plane_artifact.addr;
+                    let address = control_plane_artifact.peer.address();
 
                     StepSuccess::new(control_plane_artifact)
                         .with_metadata(
@@ -495,7 +501,10 @@ async fn fetch_artifact(
         || async {
             Ok(Peers::new(
                 &log,
-                discovery.discover_peers(&log).await?,
+                Box::new(HttpPeers::new(
+                    &log,
+                    discovery.discover_peers(&log).await?,
+                )),
                 Duration::from_secs(10),
             ))
         },
@@ -508,7 +517,7 @@ async fn fetch_artifact(
         log,
         "fetched {} bytes from {}",
         artifact.artifact.num_bytes(),
-        artifact.addr,
+        artifact.peer,
     );
 
     Ok(artifact)

installinator/src/mock_peers.rs

@@ -27,12 +27,12 @@ use uuid::Uuid;
 
 use crate::{
     errors::HttpError,
-    peers::{FetchReceiver, PeersImpl},
+    peers::{FetchReceiver, PeerAddress, PeersImpl},
 };
 
 struct MockPeersUniverse {
     artifact: Bytes,
-    peers: BTreeMap<SocketAddr, MockPeer>,
+    peers: BTreeMap<PeerAddress, MockPeer>,
     attempt_bitmaps: Vec<AttemptBitmap>,
 }
 
@@ -49,7 +49,7 @@ impl fmt::Debug for MockPeersUniverse {
 impl MockPeersUniverse {
     fn new(
         artifact: Bytes,
-        peers: BTreeMap<SocketAddr, MockPeer>,
+        peers: BTreeMap<PeerAddress, MockPeer>,
         attempt_bitmaps: Vec<AttemptBitmap>,
     ) -> Self {
         assert!(peers.len() <= 32, "this test only supports up to 32 peers");
@@ -70,7 +70,7 @@ impl MockPeersUniverse {
         // being unique identifiers. This means that this code can use a BTreeMap rather than a
         // fancier structure like an IndexMap.
         let peers_strategy = prop::collection::btree_map(
-            any::<SocketAddr>(),
+            any::<PeerAddress>(),
             any::<MockResponse_>(),
             0..max_peer_count,
         );
@@ -82,7 +82,7 @@ impl MockPeersUniverse {
         (artifact_strategy, peers_strategy, attempt_bitmaps_strategy).prop_map(
             |(artifact, peers, attempt_bitmaps): (
                 Vec<u8>,
-                BTreeMap<SocketAddr, MockResponse_>,
+                BTreeMap<PeerAddress, MockResponse_>,
                 Vec<AttemptBitmap>,
             )| {
                 let artifact = Bytes::from(artifact);
@@ -107,7 +107,7 @@ impl MockPeersUniverse {
     fn expected_result(
         &self,
         timeout: Duration,
-    ) -> Result<(usize, SocketAddr), usize> {
+    ) -> Result<(usize, PeerAddress), usize> {
         self.attempts()
             .enumerate()
             .filter_map(|(attempt, peers)| {
@@ -174,20 +174,20 @@ enum AttemptBitmap {
 struct MockPeers {
     artifact: Bytes,
     // Peers within the universe that have been selected
-    selected_peers: BTreeMap<SocketAddr, MockPeer>,
+    selected_peers: BTreeMap<PeerAddress, MockPeer>,
 }
 
 impl MockPeers {
-    fn get(&self, addr: SocketAddr) -> Option<&MockPeer> {
-        self.selected_peers.get(&addr)
+    fn get(&self, peer: PeerAddress) -> Option<&MockPeer> {
+        self.selected_peers.get(&peer)
     }
 
-    fn peers(&self) -> impl Iterator<Item = (&SocketAddr, &MockPeer)> + '_ {
+    fn peers(&self) -> impl Iterator<Item = (&PeerAddress, &MockPeer)> + '_ {
         self.selected_peers.iter()
     }
 
     /// Returns the peer that can return the entire dataset within the timeout.
-    fn successful_peer(&self, timeout: Duration) -> Option<SocketAddr> {
+    fn successful_peer(&self, timeout: Duration) -> Option<PeerAddress> {
         self.peers()
             .filter_map(|(addr, peer)| {
                 if peer.artifact != self.artifact {
@@ -233,7 +233,7 @@ impl MockPeers {
 
 #[async_trait]
 impl PeersImpl for MockPeers {
-    fn peers(&self) -> Box<dyn Iterator<Item = SocketAddr> + Send + '_> {
+    fn peers(&self) -> Box<dyn Iterator<Item = PeerAddress> + Send + '_> {
         Box::new(self.selected_peers.keys().copied())
     }
 
@@ -243,7 +243,7 @@ impl PeersImpl for MockPeers {
 
     async fn fetch_from_peer_impl(
         &self,
-        peer: SocketAddr,
+        peer: PeerAddress,
         // We don't (yet) use the artifact ID in MockPeers
         _artifact_hash_id: ArtifactHashId,
     ) -> Result<(u64, FetchReceiver), HttpError> {
@@ -261,7 +261,7 @@ impl PeersImpl for MockPeers {
 
     async fn report_progress_impl(
         &self,
-        _peer: SocketAddr,
+        _peer: PeerAddress,
         _update_id: Uuid,
         _report: EventReport,
     ) -> Result<(), ClientError> {
@@ -469,26 +469,31 @@ struct MockReportPeers {
 
 impl MockReportPeers {
     // SocketAddr::new is not a const fn in stable Rust as of this writing
-    fn valid_peer() -> SocketAddr {
-        SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 2000)
+    fn valid_peer() -> PeerAddress {
+        PeerAddress::new(SocketAddr::new(
+            IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
+            2000,
+        ))
     }
 
-    fn invalid_peer() -> SocketAddr {
-        SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)), 2000)
+    fn invalid_peer() -> PeerAddress {
+        PeerAddress::new(SocketAddr::new(
+            IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)),
+            2000,
+        ))
     }
 
-    fn unresponsive_peer() -> SocketAddr {
-        SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 3)), 2000)
-    }
-
-    fn new(update_id: Uuid, report_sender: mpsc::Sender<EventReport>) -> Self {
-        Self { update_id, report_sender }
+    fn unresponsive_peer() -> PeerAddress {
+        PeerAddress::new(SocketAddr::new(
+            IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 3)),
+            2000,
+        ))
     }
 }
 
 #[async_trait]
 impl PeersImpl for MockReportPeers {
-    fn peers(&self) -> Box<dyn Iterator<Item = SocketAddr> + Send + '_> {
+    fn peers(&self) -> Box<dyn Iterator<Item = PeerAddress> + Send + '_> {
         Box::new(
             [
                 Self::valid_peer(),
@@ -505,7 +510,7 @@ impl PeersImpl for MockReportPeers {
 
     async fn fetch_from_peer_impl(
         &self,
-        _peer: SocketAddr,
+        _peer: PeerAddress,
         _artifact_hash_id: ArtifactHashId,
     ) -> Result<(u64, FetchReceiver), HttpError> {
         unimplemented!(
@@ -516,7 +521,7 @@ impl PeersImpl for MockReportPeers {
 
     async fn report_progress_impl(
         &self,
-        peer: SocketAddr,
+        peer: PeerAddress,
         update_id: Uuid,
         report: EventReport,
     ) -> Result<(), ClientError> {
@@ -598,10 +603,10 @@ mod tests {
                     async move {
                         Ok(Peers::new(
                             &reporter_log,
-                            Box::new(MockReportPeers::new(
+                            Box::new(MockReportPeers {
                                 update_id,
                                 report_sender,
-                            )),
+                            }),
                             // The timeout is currently unused by broadcast_report.
                             Duration::from_secs(10),
                         ))
@@ -620,7 +625,7 @@ mod tests {
                         let artifact =
                             fetch_artifact(&cx, &log, attempts, timeout)
                                 .await?;
-                        let address = artifact.addr;
+                        let address = artifact.peer.address();
                         StepSuccess::new(artifact)
                             .with_metadata(
                                 InstallinatorCompletionMetadata::Download {
@@ -652,21 +657,21 @@ mod tests {
             match (expected_result, fetched_artifact) {
                 (
                     Ok((expected_attempt, expected_addr)),
-                    Ok(FetchedArtifact { attempt, addr, mut artifact }),
+                    Ok(FetchedArtifact { attempt, peer, mut artifact }),
                 ) => {
                     assert_eq!(
                         expected_attempt, attempt,
                         "expected successful attempt is the same as actual attempt"
                     );
                     assert_eq!(
-                        expected_addr, addr,
+                        expected_addr, peer,
                         "expected successful peer is the same as actual peer"
                     );
                     let artifact = artifact.copy_to_bytes(artifact.num_bytes());
                     assert_eq!(
                         expected_artifact, artifact,
                         "correct artifact fetched from peer {}",
-                        addr,
+                        peer,
                     );
                 }
                 (Err(_), Err(_)) => {}
@@ -716,7 +721,7 @@ mod tests {
 
     fn assert_reports(
         reports: &[EventReport],
-        expected_result: Result<(usize, SocketAddr), usize>,
+        expected_result: Result<(usize, PeerAddress), usize>,
     ) {
         let all_step_events: Vec<_> =
             reports.iter().flat_map(|report| &report.step_events).collect();
@@ -740,7 +745,7 @@ mod tests {
     fn assert_success_events(
         all_step_events: Vec<&StepEvent>,
         expected_attempt: usize,
-        expected_addr: SocketAddr,
+        expected_peer: PeerAddress,
     ) {
         let mut saw_success = false;
 
@@ -753,7 +758,8 @@ mod tests {
                                 peer,
                             } => {
                                 assert_ne!(
-                                    *peer, expected_addr,
+                                    *peer,
+                                    expected_peer.address(),
                                     "peer cannot match since this is the last attempt"
                                 );
                             }
@@ -795,7 +801,8 @@ mod tests {
                             ..
                         } => {
                             assert_eq!(
-                                *address, expected_addr,
+                                *address,
+                                expected_peer.address(),
                                 "address matches expected"
                             );
                         }