diff --git a/iroh-relay/src/relay_map.rs b/iroh-relay/src/relay_map.rs index 2c3b72d70c..355c3dc8f2 100644 --- a/iroh-relay/src/relay_map.rs +++ b/iroh-relay/src/relay_map.rs @@ -43,6 +43,13 @@ impl PartialEq for RelayMap { impl Eq for RelayMap {} impl RelayMap { + /// Creates an empty relay map. + pub fn empty() -> Self { + Self { + relays: Default::default(), + } + } + /// Returns the URLs of all servers in this relay map. /// /// This function is generic over the container to collect into. If you simply want a list @@ -59,13 +66,6 @@ impl RelayMap { .collect::() } - /// Creates an empty relay map. - pub fn empty() -> Self { - Self { - relays: Default::default(), - } - } - /// Returns a list with the [`RelayConfig`] for each relay in this relay map. /// /// This function is generic over the container to collect into. If you simply want a list @@ -111,14 +111,27 @@ impl RelayMap { pub fn remove(&self, url: &RelayUrl) -> Option> { self.relays.write().expect("poisoned").remove(url) } + + /// Extends this `RelayMap` with another one. + pub fn extend(&self, other: &RelayMap) { + let mut a = self.relays.write().expect("poisoned"); + let b = other.relays.read().expect("poisoned"); + a.extend(b.iter().map(|(a, b)| (a.clone(), b.clone()))); + } } impl FromIterator for RelayMap { fn from_iter>(iter: T) -> Self { + Self::from_iter(iter.into_iter().map(Arc::new)) + } +} + +impl FromIterator> for RelayMap { + fn from_iter>>(iter: T) -> Self { Self { relays: Arc::new(RwLock::new( iter.into_iter() - .map(|endpoint| (endpoint.url.clone(), Arc::new(endpoint))) + .map(|config| (config.url.clone(), config)) .collect(), )), } @@ -220,3 +233,43 @@ impl fmt::Display for RelayConfig { write!(f, "{}", self.url) } } + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use super::*; + + #[test] + fn relay_map_extend() { + let urls1 = vec![ + RelayUrl::from_str("https://hello-a-01.com").unwrap(), + RelayUrl::from_str("https://hello-b-01.com").unwrap(), + RelayUrl::from_str("https://hello-c-01-.com").unwrap(), + ]; + + let urls2 = vec![ + RelayUrl::from_str("https://hello-a-02.com").unwrap(), + RelayUrl::from_str("https://hello-b-02.com").unwrap(), + RelayUrl::from_str("https://hello-c-02-.com").unwrap(), + ]; + + let map1 = RelayMap::from_iter(urls1.clone().into_iter().map(RelayConfig::from)); + let map2 = RelayMap::from_iter(urls2.clone().into_iter().map(RelayConfig::from)); + + assert_ne!(map1, map2); + + // combine + + let map3 = RelayMap::from_iter( + map1.relays::>() + .into_iter() + .chain(map2.relays::>()), + ); + + assert_eq!(map3.len(), 6); + + map1.extend(&map2); + assert_eq!(map3, map1); + } +}