cliquenet/
net.rs

1pub mod peer;
2pub mod server;
3
4use std::{fmt, ops::Deref, sync::Arc};
5
6use bon::Builder;
7use bytes::Bytes;
8use tokio::{
9    net::TcpListener,
10    sync::{
11        OwnedSemaphorePermit,
12        mpsc::{self, UnboundedReceiver, UnboundedSender},
13        oneshot, watch,
14    },
15};
16use tracing::{debug, info, warn};
17
18use crate::{
19    Config, Metrics, NetAddr, Role, error::NetworkError, metrics::NoMetrics, msg::Slot,
20    net::server::Server, x25519::PublicKey,
21};
22
23type PeerMessage = (PublicKey, Bytes, Option<OwnedSemaphorePermit>);
24
25#[derive(Debug)]
26pub struct Network {
27    recv: NetworkReceiver,
28    send: NetworkSender,
29}
30
31#[derive(Debug)]
32pub struct NetworkReceiver {
33    rx: UnboundedReceiver<PeerMessage>,
34}
35
36#[derive(Clone)]
37pub struct NetworkSender {
38    conf: Arc<Config>,
39    node: PublicKey,
40    tx: UnboundedSender<Command>,
41    next_slot: watch::Sender<Slot>,
42    metrics: Arc<dyn Metrics>,
43}
44
45impl fmt::Debug for NetworkSender {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        f.debug_struct("NetworkSender")
48            .field("node", &self.node)
49            .field("lower_bound", &*self.next_slot.borrow())
50            .field("conf", &self.conf)
51            .finish()
52    }
53}
54
55/// Server task instructions.
56#[derive(Debug)]
57enum Command {
58    Peer(PeerCommand),
59    Send(SendCommand),
60    Shutdown(oneshot::Sender<()>),
61}
62
63/// Update network peers.
64#[derive(Debug)]
65enum PeerCommand {
66    /// Add the given peers.
67    Add(Role, Vec<(PublicKey, NetAddr)>),
68    /// Remove the given peers.
69    Remove(Vec<PublicKey>),
70    /// Assign a `Role` to the given peers.
71    Assign(Role, Vec<PublicKey>),
72}
73
74/// Send to peer(s).
75#[derive(Clone, Debug, Builder)]
76pub struct SendCommand {
77    slot: Slot,
78    action: SendAction,
79    #[builder(default)]
80    retry: RetryPolicy,
81}
82
83/// Specify if a message should be retried if no ACK is received.
84#[derive(Clone, Copy, Debug, Default)]
85pub enum RetryPolicy {
86    #[default]
87    Default,
88    NoRetry,
89}
90
91impl RetryPolicy {
92    pub fn is_retry(self) -> bool {
93        matches!(self, Self::Default)
94    }
95}
96
97#[derive(Clone, Debug)]
98pub enum SendAction {
99    /// Send a message to one peer.
100    Unicast(PublicKey, Vec<u8>),
101    /// Send a message to some peers.
102    Multicast(Vec<PublicKey>, Vec<u8>),
103    /// Send a message to all peers with `Role::Active`.
104    Broadcast(Vec<u8>),
105}
106
107impl Network {
108    pub async fn create(conf: Config) -> Result<Self, NetworkError> {
109        let listener = {
110            let result = match &conf.bind {
111                NetAddr::Inet(ip, port) => TcpListener::bind((*ip, *port)).await,
112                NetAddr::Name(hs, port) => TcpListener::bind((&**hs, *port)).await,
113            };
114            result.map_err(|e| NetworkError::Bind(conf.bind.clone(), e))?
115        };
116
117        let addr = listener.local_addr()?;
118        let node = conf.keypair.public_key();
119
120        // Command channel from application to network.
121        let (otx, orx) = mpsc::unbounded_channel();
122
123        // Channel of messages from peers to the application.
124        let (itx, irx) = mpsc::unbounded_channel();
125
126        let (etx, erx) = watch::channel(Slot::MIN);
127
128        let metr = conf.metrics.clone().unwrap_or_else(|| Arc::new(NoMetrics));
129        let conf = Arc::new(conf);
130
131        // The server ends when all `NetworkSender`s are dropped.
132        Server::spawn(
133            conf.clone(),
134            listener,
135            Role::Active,
136            itx,
137            orx,
138            erx,
139            metr.clone(),
140        );
141
142        let recv = NetworkReceiver { rx: irx };
143        let send = NetworkSender {
144            conf: conf.clone(),
145            node,
146            tx: otx,
147            next_slot: etx,
148            metrics: metr,
149        };
150
151        info!(name = %conf.name, %node, %addr, max_message_size = %conf.max_message_size, "listening");
152
153        Ok(Self { recv, send })
154    }
155
156    pub fn sender(&self) -> &NetworkSender {
157        &self.send
158    }
159
160    pub fn receiver(&self) -> &NetworkReceiver {
161        &self.recv
162    }
163
164    pub fn receiver_mut(&mut self) -> &mut NetworkReceiver {
165        &mut self.recv
166    }
167
168    pub fn split_into(self) -> (NetworkSender, NetworkReceiver) {
169        (self.send, self.recv)
170    }
171
172    pub async fn receive(&mut self) -> Option<(PublicKey, Bytes)> {
173        self.recv.receive().await
174    }
175}
176
177impl Deref for Network {
178    type Target = NetworkSender;
179
180    fn deref(&self) -> &Self::Target {
181        &self.send
182    }
183}
184
185impl NetworkReceiver {
186    /// Receive the next incoming message.
187    ///
188    /// The returned public key denotes the source where the message came from.
189    pub async fn receive(&mut self) -> Option<(PublicKey, Bytes)> {
190        let (k, b, _) = self.rx.recv().await?;
191        debug!(peer = %k, len = b.len(), "message received");
192        Some((k, b))
193    }
194}
195
196impl NetworkSender {
197    pub fn config(&self) -> &Config {
198        &self.conf
199    }
200
201    /// Send a message to a party, identified by the given public key.
202    pub fn unicast(&self, s: Slot, to: PublicKey, msg: Vec<u8>) -> Result<(), NetworkError> {
203        debug!(slot = %s, %to, len = msg.len(), "unicast");
204        self.length_check(&msg)?;
205        if self.lt_lower_bound(s) {
206            return Ok(());
207        }
208        let cmd = SendCommand::builder()
209            .slot(s)
210            .action(SendAction::Unicast(to, msg))
211            .build();
212        self.tx
213            .send(Command::Send(cmd))
214            .map_err(|_| NetworkError::ChannelClosed)
215    }
216
217    /// Send a message to all parties.
218    pub fn broadcast(&self, s: Slot, msg: Vec<u8>) -> Result<(), NetworkError> {
219        debug!(slot = %s, len = msg.len(), "broadcast");
220        self.length_check(&msg)?;
221        if self.lt_lower_bound(s) {
222            return Ok(());
223        }
224        let cmd = SendCommand::builder()
225            .slot(s)
226            .action(SendAction::Broadcast(msg))
227            .build();
228        self.tx
229            .send(Command::Send(cmd))
230            .map_err(|_| NetworkError::ChannelClosed)
231    }
232
233    /// Send a message to several parties, identified by their public keys.
234    pub fn multicast<P>(&self, s: Slot, to: P, msg: Vec<u8>) -> Result<(), NetworkError>
235    where
236        P: IntoIterator<Item = PublicKey>,
237    {
238        debug!(slot = %s, len = msg.len(), "multicast");
239        self.length_check(&msg)?;
240        if self.lt_lower_bound(s) {
241            return Ok(());
242        }
243        let cmd = SendCommand::builder()
244            .slot(s)
245            .action(SendAction::Multicast(to.into_iter().collect(), msg))
246            .build();
247        self.tx
248            .send(Command::Send(cmd))
249            .map_err(|_| NetworkError::ChannelClosed)
250    }
251
252    /// General send operation, supporting custom retry policies.
253    pub fn send(&self, cmd: SendCommand) -> Result<(), NetworkError> {
254        let bytes = msg_bytes(&cmd);
255        debug!(slot = %cmd.slot, len = %bytes.len(), "send");
256        self.length_check(bytes)?;
257        if self.lt_lower_bound(cmd.slot) {
258            return Ok(());
259        }
260        self.tx
261            .send(Command::Send(cmd))
262            .map_err(|_| NetworkError::ChannelClosed)
263    }
264
265    /// Add the given peers to the network.
266    pub fn add_peers<P>(&self, r: Role, peers: P) -> Result<(), NetworkError>
267    where
268        P: IntoIterator<Item = (PublicKey, NetAddr)>,
269    {
270        debug!(role = %r, "add_peers");
271        let peers = peers.into_iter().collect::<Vec<_>>();
272        self.tx
273            .send(Command::Peer(PeerCommand::Add(r, peers)))
274            .map_err(|_| NetworkError::ChannelClosed)
275    }
276
277    /// Remove the given peers from the network.
278    pub fn remove_peers<P>(&self, peers: P) -> Result<(), NetworkError>
279    where
280        P: IntoIterator<Item = PublicKey>,
281    {
282        debug!("remove_peers");
283        let peers = peers.into_iter().collect::<Vec<_>>();
284        for p in &peers {
285            self.metrics.del(p);
286        }
287        self.tx
288            .send(Command::Peer(PeerCommand::Remove(peers)))
289            .map_err(|_| NetworkError::ChannelClosed)
290    }
291
292    /// Assign the given role to the given peers.
293    pub fn assign_peers<P>(&self, r: Role, peers: P) -> Result<(), NetworkError>
294    where
295        P: IntoIterator<Item = PublicKey>,
296    {
297        debug!(role = %r, "assign_peers");
298        let peers = peers.into_iter().collect::<Vec<_>>();
299        self.tx
300            .send(Command::Peer(PeerCommand::Assign(r, peers)))
301            .map_err(|_| NetworkError::ChannelClosed)
302    }
303
304    /// Trigger garbage collection of messages below the given slot.
305    pub fn gc(&self, s: Slot) -> Result<(), NetworkError> {
306        debug!(slot = %s, "gc");
307        if self.next_slot.is_closed() {
308            return Err(NetworkError::ChannelClosed);
309        }
310        self.next_slot.send_if_modified(|lower_bound| {
311            if s > *lower_bound {
312                *lower_bound = s;
313                true
314            } else {
315                false
316            }
317        });
318        Ok(())
319    }
320
321    /// Trigger network shutdown.
322    ///
323    /// The returned future will resolve once the server task finished.
324    pub fn shutdown(&self) -> Result<impl Future<Output = ()> + use<>, NetworkError> {
325        warn!(name = %self.conf.name, node = %self.node, "shutdown");
326        let (tx, rx) = oneshot::channel();
327        self.tx
328            .send(Command::Shutdown(tx))
329            .map_err(|_| NetworkError::ChannelClosed)?;
330        Ok(async move {
331            let _ = rx.await;
332        })
333    }
334
335    /// Check the number of message bytes does not exceed the configured maximum.
336    fn length_check(&self, msg: &[u8]) -> Result<(), NetworkError> {
337        if msg.len() > self.conf.max_message_size.get() {
338            warn!(
339                name = %self.conf.name,
340                node = %self.node,
341                len  = %msg.len(),
342                max  = %self.conf.max_message_size,
343                "message too large to send"
344            );
345            return Err(NetworkError::MessageTooLarge);
346        }
347        Ok(())
348    }
349
350    /// Check if the given slot is less than our lower bound.
351    fn lt_lower_bound(&self, s: Slot) -> bool {
352        let lower_bound = *self.next_slot.borrow();
353        if s < lower_bound {
354            warn!(
355                name = %self.conf.name,
356                node = %self.node,
357                slot = %s,
358                %lower_bound,
359                "slot below lower bound"
360            );
361            return true;
362        }
363        false
364    }
365}
366
367fn msg_bytes(cmd: &SendCommand) -> &[u8] {
368    match &cmd.action {
369        SendAction::Unicast(_, b) => b,
370        SendAction::Multicast(_, b) => b,
371        SendAction::Broadcast(b) => b,
372    }
373}
cliquenet/net.rs

cliquenet/
net.rs
