0g-storage-node/node/network/tests/common/swarm.rs
2024-01-04 18:05:32 +08:00

100 lines
3.0 KiB
Rust

use std::collections::HashMap;
use std::pin::Pin;
use super::behaviour::{CallTraceBehaviour, MockBehaviour};
use futures::stream::Stream;
use futures::task::{Context, Poll};
use libp2p::swarm::handler::ConnectionHandler;
use libp2p::swarm::{IntoConnectionHandler, NetworkBehaviour, Swarm, SwarmBuilder, SwarmEvent};
use libp2p::{PeerId, Transport};
use futures::StreamExt;
pub fn new_test_swarm<B>(behaviour: B) -> Swarm<B>
where
B: NetworkBehaviour,
{
let id_keys = libp2p::identity::Keypair::generate_ed25519();
let local_public_key = id_keys.public();
let transport = libp2p::core::transport::MemoryTransport
.upgrade(libp2p::core::upgrade::Version::V1)
.authenticate(libp2p::plaintext::PlainText2Config {
local_public_key: local_public_key.clone(),
})
.multiplex(libp2p::yamux::YamuxConfig::default())
.boxed();
SwarmBuilder::new(transport, behaviour, local_public_key.into()).build()
}
pub fn random_multiaddr() -> libp2p::multiaddr::Multiaddr {
libp2p::multiaddr::Protocol::Memory(rand::random::<u64>()).into()
}
/// Bind a memory multiaddr to a compatible swarm.
pub async fn bind_listener<B: NetworkBehaviour>(
swarm: &mut Swarm<B>,
) -> libp2p::multiaddr::Multiaddr {
swarm.listen_on(random_multiaddr()).unwrap();
match swarm.select_next_some().await {
SwarmEvent::NewListenAddr {
listener_id: _,
address,
} => address,
_ => panic!("Testing swarm's first event should be a new listener"),
}
}
#[derive(Default)]
pub struct SwarmPool<B: NetworkBehaviour> {
swarms: HashMap<PeerId, Swarm<B>>,
}
impl<B: NetworkBehaviour> SwarmPool<B> {
pub fn with_capacity(capacity: usize) -> Self {
Self {
swarms: HashMap::with_capacity(capacity),
}
}
pub fn insert(&mut self, swarm: Swarm<B>) -> PeerId {
let peer_id = *swarm.local_peer_id();
self.swarms.insert(peer_id, swarm);
peer_id
}
pub fn remove(&mut self, peer_id: &PeerId) {
self.swarms.remove(peer_id);
}
pub fn get_mut(&mut self, peer_id: &PeerId) -> Option<&mut Swarm<B>> {
self.swarms.get_mut(peer_id)
}
pub fn swarms(&self) -> &HashMap<PeerId, Swarm<B>> {
&self.swarms
}
pub fn swarms_mut(&mut self) -> &mut HashMap<PeerId, Swarm<B>> {
&mut self.swarms
}
}
impl<B> Stream for SwarmPool<B>
where
B: NetworkBehaviour,
<B as NetworkBehaviour>::ConnectionHandler: ConnectionHandler,
{
type Item = (PeerId,
SwarmEvent<<B as NetworkBehaviour>::OutEvent, <<<B as NetworkBehaviour>::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::Error>);
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut polls = self
.get_mut()
.swarms
.iter_mut()
.map(|(&peer_id, swarm)| swarm.map(move |ev| (peer_id, ev)))
.collect::<futures::stream::SelectAll<_>>();
polls.poll_next_unpin(cx)
}
}