authority-discovery: Support multiple authority ids per peer id

client/authority-discovery/src/lib.rs

@@ -18,6 +18,7 @@
 
 #![warn(missing_docs)]
 #![recursion_limit = "1024"]
+
 //! Substrate authority discovery.
 //!
 //! This crate enables Substrate authorities to discover and directly connect to
@@ -31,7 +32,7 @@ pub use crate::{
 	worker::{NetworkProvider, Role, Worker},
 };
 
-use std::{sync::Arc, time::Duration};
+use std::{collections::HashSet, sync::Arc, time::Duration};
 
 use futures::{
 	channel::{mpsc, oneshot},
@@ -58,11 +59,13 @@ pub struct WorkerConfig {
 	///
 	/// By default this is set to 1 hour.
 	pub max_publish_interval: Duration,
+
 	/// Interval at which the keystore is queried. If the keys have changed, unconditionally
 	/// re-publish its addresses on the DHT.
 	///
 	/// By default this is set to 1 minute.
 	pub keystore_refresh_interval: Duration,
+
 	/// The maximum interval in which the node will query the DHT for new entries.
 	///
 	/// By default this is set to 10 minutes.
@@ -156,7 +159,7 @@ where
 /// Message send from the [`Service`] to the [`Worker`].
 pub(crate) enum ServicetoWorkerMsg {
 	/// See [`Service::get_addresses_by_authority_id`].
-	GetAddressesByAuthorityId(AuthorityId, oneshot::Sender<Option<Vec<Multiaddr>>>),
-	/// See [`Service::get_authority_id_by_peer_id`].
-	GetAuthorityIdByPeerId(PeerId, oneshot::Sender<Option<AuthorityId>>),
+	GetAddressesByAuthorityId(AuthorityId, oneshot::Sender<Option<HashSet<Multiaddr>>>),
+	/// See [`Service::get_authority_ids_by_peer_id`].
+	GetAuthorityIdsByPeerId(PeerId, oneshot::Sender<Option<HashSet<AuthorityId>>>),
 }

client/authority-discovery/src/service.rs

@@ -16,7 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-use std::fmt::Debug;
+use std::{collections::HashSet, fmt::Debug};
 
 use crate::ServicetoWorkerMsg;
 
@@ -62,7 +62,7 @@ impl Service {
 	pub async fn get_addresses_by_authority_id(
 		&mut self,
 		authority: AuthorityId,
-	) -> Option<Vec<Multiaddr>> {
+	) -> Option<HashSet<Multiaddr>> {
 		let (tx, rx) = oneshot::channel();
 
 		self.to_worker
@@ -78,11 +78,14 @@ impl Service {
 	///
 	/// Returns `None` if no entry was present or connection to the
 	/// [`crate::Worker`] failed.
-	pub async fn get_authority_id_by_peer_id(&mut self, peer_id: PeerId) -> Option<AuthorityId> {
+	pub async fn get_authority_ids_by_peer_id(
+		&mut self,
+		peer_id: PeerId,
+	) -> Option<HashSet<AuthorityId>> {
 		let (tx, rx) = oneshot::channel();
 
 		self.to_worker
-			.send(ServicetoWorkerMsg::GetAuthorityIdByPeerId(peer_id, tx))
+			.send(ServicetoWorkerMsg::GetAuthorityIdsByPeerId(peer_id, tx))
 			.await
 			.ok()?;
 

client/authority-discovery/src/tests.rs

@@ -29,7 +29,7 @@ use libp2p::core::{
 	multiaddr::{Multiaddr, Protocol},
 	PeerId,
 };
-use std::sync::Arc;
+use std::{collections::HashSet, sync::Arc};
 
 use sp_authority_discovery::AuthorityId;
 use sp_core::crypto::key_types;
@@ -73,12 +73,12 @@ fn get_addresses_and_authority_id() {
 
 	pool.run_until(async {
 		assert_eq!(
-			Some(vec![remote_addr]),
+			Some(HashSet::from([remote_addr])),
 			service.get_addresses_by_authority_id(remote_authority_id.clone()).await,
 		);
 		assert_eq!(
-			Some(remote_authority_id),
-			service.get_authority_id_by_peer_id(remote_peer_id).await,
+			Some(HashSet::from([remote_authority_id])),
+			service.get_authority_ids_by_peer_id(remote_peer_id).await,
 		);
 	});
 }

client/authority-discovery/src/worker.rs

@@ -259,9 +259,9 @@ where
 					self.addr_cache.get_addresses_by_authority_id(&authority).map(Clone::clone),
 				);
 			},
-			ServicetoWorkerMsg::GetAuthorityIdByPeerId(peer_id, sender) => {
+			ServicetoWorkerMsg::GetAuthorityIdsByPeerId(peer_id, sender) => {
 				let _ = sender
-					.send(self.addr_cache.get_authority_id_by_peer_id(&peer_id).map(Clone::clone));
+					.send(self.addr_cache.get_authority_ids_by_peer_id(&peer_id).map(Clone::clone));
 			},
 		}
 	}
@@ -374,7 +374,7 @@ where
 			.map_err(|e| Error::CallingRuntime(e.into()))?
 			.into_iter()
 			.filter(|id| !local_keys.contains(id.as_ref()))
-			.collect();
+			.collect::<Vec<_>>();
 
 		self.addr_cache.retain_ids(&authorities);
 
@@ -548,7 +548,7 @@ where
 			if let Some(metrics) = &self.metrics {
 				metrics
 					.known_authorities_count
-					.set(self.addr_cache.num_ids().try_into().unwrap_or(std::u64::MAX));
+					.set(self.addr_cache.num_authority_ids().try_into().unwrap_or(std::u64::MAX));
 			}
 		}
 		Ok(())