datacake_node/rpc/
network.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::sync::Arc;
4
5use datacake_rpc::Channel;
6use parking_lot::RwLock;
7use tracing::trace;
8
9#[derive(Clone, Default)]
10/// A collection of RPC client connections which can be reused and multiplexed.
11pub struct RpcNetwork {
12    clients: Arc<RwLock<HashMap<SocketAddr, Channel>>>,
13}
14
15impl RpcNetwork {
16    /// Attempts to get an already existing connection or creates a new connection.
17    pub fn get_or_connect(&self, addr: SocketAddr) -> Channel {
18        {
19            let guard = self.clients.read();
20            if let Some(channel) = guard.get(&addr) {
21                return channel.clone();
22            }
23        }
24
25        trace!(addr = %addr, "Connect client to network.");
26        self.connect(addr)
27    }
28
29    /// Connects to a given address and adds it to the clients.
30    pub fn connect(&self, addr: SocketAddr) -> Channel {
31        let channel = Channel::connect(addr);
32
33        {
34            let mut guard = self.clients.write();
35            guard.insert(addr, channel.clone());
36        }
37
38        channel
39    }
40
41    /// Removes a client from the network.
42    pub fn disconnect(&self, addr: SocketAddr) {
43        let mut guard = self.clients.write();
44        guard.remove(&addr);
45    }
46}
datacake_node/rpc/network.rs

datacake_node/rpc/
network.rs
