hotshot_libp2p_networking/network/node/
handle.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use std::{collections::HashSet, fmt::Debug, sync::Arc, time::Duration};
8
9use bimap::BiMap;
10use hotshot_types::traits::{
11    network::NetworkError, node_implementation::NodeType, signature_key::SignatureKey,
12};
13use libp2p::{Multiaddr, request_response::ResponseChannel};
14use libp2p_identity::PeerId;
15use parking_lot::Mutex;
16use tokio::{
17    sync::mpsc::{Receiver, UnboundedReceiver, UnboundedSender},
18    time::{sleep, timeout},
19};
20use tracing::{debug, info, instrument};
21
22use crate::network::{
23    ClientRequest, NetworkEvent, NetworkNode, NetworkNodeConfig, SwarmTaskHandle,
24    behaviours::dht::{
25        record::{Namespace, RecordKey, RecordValue},
26        store::persistent::DhtPersistentStorage,
27    },
28    gen_multiaddr, log_summary,
29};
30
31/// A handle containing:
32/// - A reference to the state
33/// - Controls for the swarm
34#[derive(Debug, Clone)]
35pub struct NetworkNodeHandle<T: NodeType> {
36    /// network configuration
37    network_config: NetworkNodeConfig,
38
39    /// send an action to the networkbehaviour
40    send_network: UnboundedSender<ClientRequest>,
41
42    /// The map from consensus keys to peer IDs
43    consensus_key_to_pid_map: Arc<Mutex<BiMap<T::SignatureKey, PeerId>>>,
44
45    /// the local address we're listening on
46    listen_addr: Multiaddr,
47
48    /// the peer id of the networkbehaviour
49    peer_id: PeerId,
50
51    /// human readable id
52    id: usize,
53
54    /// Handle to the spawned swarm event-loop task.
55    swarm_task: Arc<Mutex<Option<SwarmTaskHandle>>>,
56}
57
58/// internal network node receiver
59#[derive(Debug)]
60pub struct NetworkNodeReceiver {
61    /// the receiver
62    receiver: UnboundedReceiver<NetworkEvent>,
63
64    ///kill switch
65    recv_kill: Option<Receiver<()>>,
66}
67
68impl NetworkNodeReceiver {
69    /// recv a network event
70    /// # Errors
71    /// Errors if the receiver channel is closed
72    pub async fn recv(&mut self) -> Result<NetworkEvent, NetworkError> {
73        self.receiver
74            .recv()
75            .await
76            .ok_or(NetworkError::ChannelReceiveError(
77                "Receiver channel closed".to_string(),
78            ))
79    }
80    /// Add a kill switch to the receiver
81    pub fn set_kill_switch(&mut self, kill_switch: Receiver<()>) {
82        self.recv_kill = Some(kill_switch);
83    }
84
85    /// Take the kill switch to allow killing the receiver task
86    pub fn take_kill_switch(&mut self) -> Option<Receiver<()>> {
87        self.recv_kill.take()
88    }
89}
90
91/// Spawn a network node task task and return the handle and the receiver for it
92/// # Errors
93/// Errors if spawning the task fails
94pub async fn spawn_network_node<T: NodeType, D: DhtPersistentStorage>(
95    config: NetworkNodeConfig,
96    dht_persistent_storage: D,
97    consensus_key_to_pid_map: Arc<Mutex<BiMap<T::SignatureKey, PeerId>>>,
98    id: usize,
99) -> Result<(NetworkNodeReceiver, NetworkNodeHandle<T>), NetworkError> {
100    let mut network: NetworkNode<T, _> = NetworkNode::new(
101        config.clone(),
102        dht_persistent_storage,
103        Arc::clone(&consensus_key_to_pid_map),
104    )
105    .await
106    .map_err(|e| NetworkError::ConfigError(format!("failed to create network node: {e}")))?;
107    // randomly assigned port
108    let listen_addr = config
109        .bind_address
110        .clone()
111        .unwrap_or_else(|| gen_multiaddr(0));
112    let peer_id = network.peer_id();
113    let listen_addr = network.start_listen(listen_addr).await.map_err(|e| {
114        NetworkError::ListenError(format!("failed to start listening on Libp2p: {e}"))
115    })?;
116    // pin here to force the future onto the heap since it can be large
117    // in the case of flume
118    let (send_chan, recv_chan, swarm_task) = network.spawn_listeners().map_err(|err| {
119        NetworkError::ListenError(format!("failed to spawn listeners for Libp2p: {err}"))
120    })?;
121    log_summary::spawn_summary_task();
122    let receiver = NetworkNodeReceiver {
123        receiver: recv_chan,
124        recv_kill: None,
125    };
126
127    let handle = NetworkNodeHandle::<T> {
128        network_config: config,
129        send_network: send_chan,
130        consensus_key_to_pid_map,
131        listen_addr,
132        peer_id,
133        id,
134        swarm_task: Arc::new(Mutex::new(Some(swarm_task))),
135    };
136    Ok((receiver, handle))
137}
138
139impl<T: NodeType> NetworkNodeHandle<T> {
140    /// Cleanly shuts down a swarm node
141    /// This is done by sending a message to
142    /// the swarm itself to spin down
143    #[instrument]
144    pub async fn shutdown(&self) -> Result<(), NetworkError> {
145        self.send_request(ClientRequest::Shutdown)?;
146
147        // Wait for the swarm event-loop task to actually finish, so its
148        // listening socket is released before we return.
149        let task = self.swarm_task.lock().take();
150        if let Some(task) = task {
151            match timeout(Duration::from_secs(5), task).await {
152                Ok(Ok(_)) => {},
153                Ok(Err(err)) => debug!(%err, "swarm task ended with error during shutdown"),
154                Err(_) => {
155                    debug!("timed out waiting for swarm task to finish during shutdown");
156                },
157            }
158        }
159        Ok(())
160    }
161    /// Notify the network to begin the bootstrap process
162    /// # Errors
163    /// If unable to send via `send_network`. This should only happen
164    /// if the network is shut down.
165    pub fn begin_bootstrap(&self) -> Result<(), NetworkError> {
166        let req = ClientRequest::BeginBootstrap;
167        self.send_request(req)
168    }
169
170    /// Get a reference to the network node handle's listen addr.
171    #[must_use]
172    pub fn listen_addr(&self) -> Multiaddr {
173        self.listen_addr.clone()
174    }
175
176    /// Print out the routing table used by kademlia
177    /// NOTE: only for debugging purposes currently
178    /// # Errors
179    /// if the client has stopped listening for a response
180    pub async fn print_routing_table(&self) -> Result<(), NetworkError> {
181        let (s, r) = futures::channel::oneshot::channel();
182        let req = ClientRequest::GetRoutingTable(s);
183        self.send_request(req)?;
184        r.await
185            .map_err(|e| NetworkError::ChannelReceiveError(e.to_string()))
186    }
187    /// Wait until at least `num_peers` have connected
188    ///
189    /// # Errors
190    /// If the channel closes before the result can be sent back
191    pub async fn wait_to_connect(
192        &self,
193        num_required_peers: usize,
194        node_id: usize,
195    ) -> Result<(), NetworkError> {
196        // Wait for the required number of peers to connect
197        loop {
198            // Get the number of currently connected peers
199            let num_connected = self.num_connected().await?;
200            if num_connected >= num_required_peers {
201                break;
202            }
203
204            // Log the number of connected peers
205            info!(
206                "Node {} connected to {}/{} peers",
207                node_id, num_connected, num_required_peers
208            );
209
210            // Sleep for a second before checking again
211            sleep(Duration::from_secs(1)).await;
212        }
213
214        Ok(())
215    }
216
217    /// Look up a peer's addresses in kademlia
218    /// NOTE: this should always be called before any `request_response` is initiated
219    /// # Errors
220    /// if the client has stopped listening for a response
221    pub async fn lookup_pid(&self, peer_id: PeerId) -> Result<(), NetworkError> {
222        let (s, r) = futures::channel::oneshot::channel();
223        let req = ClientRequest::LookupPeer(peer_id, s);
224        self.send_request(req)?;
225        r.await
226            .map_err(|err| NetworkError::ChannelReceiveError(err.to_string()))
227    }
228
229    /// Looks up a node's `PeerId` by its consensus key.
230    ///
231    /// # Errors
232    /// If the DHT lookup fails
233    pub async fn lookup_node(
234        &self,
235        consensus_key: &T::SignatureKey,
236        dht_timeout: Duration,
237    ) -> Result<PeerId, NetworkError> {
238        // First check if we already have an open connection to the peer
239        if let Some(pid) = self
240            .consensus_key_to_pid_map
241            .lock()
242            .get_by_left(consensus_key)
243        {
244            return Ok(*pid);
245        }
246
247        // Create the record key
248        let key = RecordKey::new(Namespace::Lookup, consensus_key.to_bytes());
249
250        // Get the record from the DHT
251        let pid = self.get_record_timeout(key, dht_timeout).await?;
252
253        PeerId::from_bytes(&pid).map_err(|err| NetworkError::FailedToDeserialize(err.to_string()))
254    }
255
256    /// Insert a record into the kademlia DHT
257    /// # Errors
258    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize the key or value
259    pub async fn put_record(
260        &self,
261        key: RecordKey,
262        value: RecordValue<T::SignatureKey>,
263    ) -> Result<(), NetworkError> {
264        // Serialize the key
265        let key = key.to_bytes();
266
267        // Serialize the record
268        let value = bincode::serialize(&value)
269            .map_err(|e| NetworkError::FailedToSerialize(e.to_string()))?;
270
271        let (s, r) = futures::channel::oneshot::channel();
272        let req = ClientRequest::PutDHT {
273            key: key.clone(),
274            value,
275            notify: s,
276        };
277
278        self.send_request(req)?;
279
280        r.await.map_err(|_| NetworkError::RequestCancelled)
281    }
282
283    /// Receive a record from the kademlia DHT if it exists.
284    /// Must be replicated on at least 2 nodes
285    /// # Errors
286    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize the key
287    /// - Will return [`NetworkError::FailedToDeserialize`] when unable to deserialize the returned value
288    pub async fn get_record(
289        &self,
290        key: RecordKey,
291        retry_count: u8,
292    ) -> Result<Vec<u8>, NetworkError> {
293        // Serialize the key
294        let serialized_key = key.to_bytes();
295
296        let (s, r) = futures::channel::oneshot::channel();
297        let req = ClientRequest::GetDHT {
298            key: serialized_key.clone(),
299            notify: vec![s],
300            retry_count,
301        };
302        self.send_request(req)?;
303
304        // Map the error
305        let result = r.await.map_err(|_| NetworkError::RequestCancelled)?;
306
307        // Deserialize the record's value
308        let record: RecordValue<T::SignatureKey> = bincode::deserialize(&result)
309            .map_err(|e| NetworkError::FailedToDeserialize(e.to_string()))?;
310
311        Ok(record.value().to_vec())
312    }
313
314    /// Get a record from the kademlia DHT with a timeout
315    /// # Errors
316    /// - Will return [`NetworkError::Timeout`] when times out
317    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize the key or value
318    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
319    pub async fn get_record_timeout(
320        &self,
321        key: RecordKey,
322        timeout_duration: Duration,
323    ) -> Result<Vec<u8>, NetworkError> {
324        timeout(timeout_duration, self.get_record(key, 3))
325            .await
326            .map_err(|err| NetworkError::Timeout(err.to_string()))?
327    }
328
329    /// Insert a record into the kademlia DHT with a timeout
330    /// # Errors
331    /// - Will return [`NetworkError::Timeout`] when times out
332    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize the key or value
333    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
334    pub async fn put_record_timeout(
335        &self,
336        key: RecordKey,
337        value: RecordValue<T::SignatureKey>,
338        timeout_duration: Duration,
339    ) -> Result<(), NetworkError> {
340        timeout(timeout_duration, self.put_record(key, value))
341            .await
342            .map_err(|err| NetworkError::Timeout(err.to_string()))?
343    }
344
345    /// Subscribe to a topic
346    /// # Errors
347    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
348    pub async fn subscribe(&self, topic: String) -> Result<(), NetworkError> {
349        let (s, r) = futures::channel::oneshot::channel();
350        let req = ClientRequest::Subscribe(topic, Some(s));
351        self.send_request(req)?;
352        r.await
353            .map_err(|err| NetworkError::ChannelReceiveError(err.to_string()))
354    }
355
356    /// Unsubscribe from a topic
357    /// # Errors
358    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
359    pub async fn unsubscribe(&self, topic: String) -> Result<(), NetworkError> {
360        let (s, r) = futures::channel::oneshot::channel();
361        let req = ClientRequest::Unsubscribe(topic, Some(s));
362        self.send_request(req)?;
363        r.await
364            .map_err(|err| NetworkError::ChannelReceiveError(err.to_string()))
365    }
366
367    /// Ignore `peers` when pruning
368    /// e.g. maintain their connection
369    /// # Errors
370    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
371    pub fn ignore_peers(&self, peers: Vec<PeerId>) -> Result<(), NetworkError> {
372        let req = ClientRequest::IgnorePeers(peers);
373        self.send_request(req)
374    }
375
376    /// Make a direct request to `peer_id` containing `msg`
377    /// # Errors
378    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
379    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize `msg`
380    pub fn direct_request(&self, pid: PeerId, msg: &[u8]) -> Result<(), NetworkError> {
381        self.direct_request_no_serialize(pid, msg.to_vec())
382    }
383
384    /// Make a direct request to `peer_id` containing `msg` without serializing
385    /// # Errors
386    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
387    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize `msg`
388    pub fn direct_request_no_serialize(
389        &self,
390        pid: PeerId,
391        contents: Vec<u8>,
392    ) -> Result<(), NetworkError> {
393        let req = ClientRequest::DirectRequest {
394            pid,
395            contents,
396            retry_count: 1,
397        };
398        self.send_request(req)
399    }
400
401    /// Reply with `msg` to a request over `chan`
402    /// # Errors
403    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
404    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize `msg`
405    pub fn direct_response(
406        &self,
407        chan: ResponseChannel<Vec<u8>>,
408        msg: &[u8],
409    ) -> Result<(), NetworkError> {
410        let req = ClientRequest::DirectResponse(chan, msg.to_vec());
411        self.send_request(req)
412    }
413
414    /// Forcefully disconnect from a peer
415    /// # Errors
416    /// If the channel is closed somehow
417    /// Shouldnt' happen.
418    /// # Panics
419    /// If channel errors out
420    /// shouldn't happen.
421    pub fn prune_peer(&self, pid: PeerId) -> Result<(), NetworkError> {
422        let req = ClientRequest::Prune(pid);
423        self.send_request(req)
424    }
425
426    /// Gossip a message to peers
427    /// # Errors
428    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
429    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize `msg`
430    pub fn gossip(&self, topic: String, msg: &[u8]) -> Result<(), NetworkError> {
431        self.gossip_no_serialize(topic, msg.to_vec())
432    }
433
434    /// Gossip a message to peers without serializing
435    /// # Errors
436    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
437    /// - Will return [`NetworkError::FailedToSerialize`] when unable to serialize `msg`
438    pub fn gossip_no_serialize(&self, topic: String, msg: Vec<u8>) -> Result<(), NetworkError> {
439        let req = ClientRequest::GossipMsg(topic, msg);
440        self.send_request(req)
441    }
442
443    /// Tell libp2p about known network nodes
444    /// # Errors
445    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
446    pub fn add_known_peers(
447        &self,
448        known_peers: Vec<(PeerId, Multiaddr)>,
449    ) -> Result<(), NetworkError> {
450        debug!("Adding {} known peers", known_peers.len());
451        let req = ClientRequest::AddKnownPeers(known_peers);
452        self.send_request(req)
453    }
454
455    /// Send a client request to the network
456    ///
457    /// # Errors
458    /// - Will return [`NetworkError::ChannelSendError`] when underlying `NetworkNode` has been killed
459    fn send_request(&self, req: ClientRequest) -> Result<(), NetworkError> {
460        self.send_network
461            .send(req)
462            .map_err(|err| NetworkError::ChannelSendError(err.to_string()))
463    }
464
465    /// Returns number of peers this node is connected to
466    /// # Errors
467    /// If the channel is closed somehow
468    /// Shouldnt' happen.
469    /// # Panics
470    /// If channel errors out
471    /// shouldn't happen.
472    pub async fn num_connected(&self) -> Result<usize, NetworkError> {
473        let (s, r) = futures::channel::oneshot::channel();
474        let req = ClientRequest::GetConnectedPeerNum(s);
475        self.send_request(req)?;
476        Ok(r.await.unwrap())
477    }
478
479    /// return hashset of PIDs this node is connected to
480    /// # Errors
481    /// If the channel is closed somehow
482    /// Shouldnt' happen.
483    /// # Panics
484    /// If channel errors out
485    /// shouldn't happen.
486    pub async fn connected_pids(&self) -> Result<HashSet<PeerId>, NetworkError> {
487        let (s, r) = futures::channel::oneshot::channel();
488        let req = ClientRequest::GetConnectedPeers(s);
489        self.send_request(req)?;
490        Ok(r.await.unwrap())
491    }
492
493    /// Return the set of peer IDs in the Kademlia routing table.
494    /// # Errors
495    /// If the channel is closed somehow
496    pub async fn kad_routing_peers(&self) -> Result<HashSet<PeerId>, NetworkError> {
497        let (s, r) = futures::channel::oneshot::channel();
498        let req = ClientRequest::GetKadRoutingPeers(s);
499        self.send_request(req)?;
500        Ok(r.await.unwrap())
501    }
502
503    /// Get a reference to the network node handle's id.
504    #[must_use]
505    pub fn id(&self) -> usize {
506        self.id
507    }
508
509    /// Get a reference to the network node handle's peer id.
510    #[must_use]
511    pub fn peer_id(&self) -> PeerId {
512        self.peer_id
513    }
514
515    /// Return a reference to the network config
516    #[must_use]
517    pub fn config(&self) -> &NetworkNodeConfig {
518        &self.network_config
519    }
520}
hotshot_libp2p_networking/network/node/handle.rs

hotshot_libp2p_networking/network/node/
handle.rs
