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 = TcpListener::bind(conf.bind.to_string())
110            .await
111            .map_err(|e| NetworkError::Bind(conf.bind.clone(), e))?;
112
113        let addr = listener.local_addr()?;
114        let node = conf.keypair.public_key();
115
116        // Command channel from application to network.
117        let (otx, orx) = mpsc::unbounded_channel();
118
119        // Channel of messages from peers to the application.
120        let (itx, irx) = mpsc::unbounded_channel();
121
122        let (etx, erx) = watch::channel(Slot::MIN);
123
124        let metr = conf.metrics.clone().unwrap_or_else(|| Arc::new(NoMetrics));
125        let conf = Arc::new(conf);
126
127        // The server ends when all `NetworkSender`s are dropped.
128        Server::spawn(
129            conf.clone(),
130            listener,
131            Role::Active,
132            itx,
133            orx,
134            erx,
135            metr.clone(),
136        );
137
138        let recv = NetworkReceiver { rx: irx };
139        let send = NetworkSender {
140            conf: conf.clone(),
141            node,
142            tx: otx,
143            next_slot: etx,
144            metrics: metr,
145        };
146
147        info!(name = %conf.name, %node, %addr, "listening");
148
149        Ok(Self { recv, send })
150    }
151
152    pub fn sender(&self) -> &NetworkSender {
153        &self.send
154    }
155
156    pub fn receiver(&self) -> &NetworkReceiver {
157        &self.recv
158    }
159
160    pub fn receiver_mut(&mut self) -> &mut NetworkReceiver {
161        &mut self.recv
162    }
163
164    pub fn split_into(self) -> (NetworkSender, NetworkReceiver) {
165        (self.send, self.recv)
166    }
167
168    pub async fn receive(&mut self) -> Option<(PublicKey, Bytes)> {
169        self.recv.receive().await
170    }
171}
172
173impl Deref for Network {
174    type Target = NetworkSender;
175
176    fn deref(&self) -> &Self::Target {
177        &self.send
178    }
179}
180
181impl NetworkReceiver {
182    /// Receive the next incoming message.
183    ///
184    /// The returned public key denotes the source where the message came from.
185    pub async fn receive(&mut self) -> Option<(PublicKey, Bytes)> {
186        let (k, b, _) = self.rx.recv().await?;
187        debug!(peer = %k, len = b.len(), "message received");
188        Some((k, b))
189    }
190}
191
192impl NetworkSender {
193    pub fn config(&self) -> &Config {
194        &self.conf
195    }
196
197    /// Send a message to a party, identified by the given public key.
198    pub fn unicast(&self, s: Slot, to: PublicKey, msg: Vec<u8>) -> Result<(), NetworkError> {
199        debug!(slot = %s, %to, len = msg.len(), "unicast");
200        self.length_check(&msg)?;
201        if self.lt_lower_bound(s) {
202            return Ok(());
203        }
204        let cmd = SendCommand::builder()
205            .slot(s)
206            .action(SendAction::Unicast(to, msg))
207            .build();
208        self.tx
209            .send(Command::Send(cmd))
210            .map_err(|_| NetworkError::ChannelClosed)
211    }
212
213    /// Send a message to all parties.
214    pub fn broadcast(&self, s: Slot, msg: Vec<u8>) -> Result<(), NetworkError> {
215        debug!(slot = %s, len = msg.len(), "broadcast");
216        self.length_check(&msg)?;
217        if self.lt_lower_bound(s) {
218            return Ok(());
219        }
220        let cmd = SendCommand::builder()
221            .slot(s)
222            .action(SendAction::Broadcast(msg))
223            .build();
224        self.tx
225            .send(Command::Send(cmd))
226            .map_err(|_| NetworkError::ChannelClosed)
227    }
228
229    /// Send a message to several parties, identified by their public keys.
230    pub fn multicast<P>(&self, s: Slot, to: P, msg: Vec<u8>) -> Result<(), NetworkError>
231    where
232        P: IntoIterator<Item = PublicKey>,
233    {
234        debug!(slot = %s, len = msg.len(), "multicast");
235        self.length_check(&msg)?;
236        if self.lt_lower_bound(s) {
237            return Ok(());
238        }
239        let cmd = SendCommand::builder()
240            .slot(s)
241            .action(SendAction::Multicast(to.into_iter().collect(), msg))
242            .build();
243        self.tx
244            .send(Command::Send(cmd))
245            .map_err(|_| NetworkError::ChannelClosed)
246    }
247
248    /// General send operation, supporting custom retry policies.
249    pub fn send(&self, cmd: SendCommand) -> Result<(), NetworkError> {
250        let bytes = msg_bytes(&cmd);
251        debug!(slot = %cmd.slot, len = %bytes.len(), "send");
252        self.length_check(bytes)?;
253        if self.lt_lower_bound(cmd.slot) {
254            return Ok(());
255        }
256        self.tx
257            .send(Command::Send(cmd))
258            .map_err(|_| NetworkError::ChannelClosed)
259    }
260
261    /// Add the given peers to the network.
262    pub fn add_peers<P>(&self, r: Role, peers: P) -> Result<(), NetworkError>
263    where
264        P: IntoIterator<Item = (PublicKey, NetAddr)>,
265    {
266        debug!(role = %r, "add_peers");
267        let peers = peers.into_iter().collect::<Vec<_>>();
268        self.tx
269            .send(Command::Peer(PeerCommand::Add(r, peers)))
270            .map_err(|_| NetworkError::ChannelClosed)
271    }
272
273    /// Remove the given peers from the network.
274    pub fn remove_peers<P>(&self, peers: P) -> Result<(), NetworkError>
275    where
276        P: IntoIterator<Item = PublicKey>,
277    {
278        debug!("remove_peers");
279        let peers = peers.into_iter().collect::<Vec<_>>();
280        for p in &peers {
281            self.metrics.del(p);
282        }
283        self.tx
284            .send(Command::Peer(PeerCommand::Remove(peers)))
285            .map_err(|_| NetworkError::ChannelClosed)
286    }
287
288    /// Assign the given role to the given peers.
289    pub fn assign_peers<P>(&self, r: Role, peers: P) -> Result<(), NetworkError>
290    where
291        P: IntoIterator<Item = PublicKey>,
292    {
293        debug!(role = %r, "assign_peers");
294        let peers = peers.into_iter().collect::<Vec<_>>();
295        self.tx
296            .send(Command::Peer(PeerCommand::Assign(r, peers)))
297            .map_err(|_| NetworkError::ChannelClosed)
298    }
299
300    /// Trigger garbage collection of messages below the given slot.
301    pub fn gc(&self, s: Slot) -> Result<(), NetworkError> {
302        debug!(slot = %s, "gc");
303        if self.next_slot.is_closed() {
304            return Err(NetworkError::ChannelClosed);
305        }
306        self.next_slot.send_if_modified(|lower_bound| {
307            if s > *lower_bound {
308                *lower_bound = s;
309                true
310            } else {
311                false
312            }
313        });
314        Ok(())
315    }
316
317    /// Trigger network shutdown.
318    ///
319    /// The returned future will resolve once the server task finished.
320    pub fn shutdown(&self) -> Result<impl Future<Output = ()> + use<>, NetworkError> {
321        warn!(name = %self.conf.name, node = %self.node, "shutdown");
322        let (tx, rx) = oneshot::channel();
323        self.tx
324            .send(Command::Shutdown(tx))
325            .map_err(|_| NetworkError::ChannelClosed)?;
326        Ok(async move {
327            let _ = rx.await;
328        })
329    }
330
331    /// Check the number of message bytes does not exceed the configured maximum.
332    fn length_check(&self, msg: &[u8]) -> Result<(), NetworkError> {
333        if msg.len() > self.conf.max_message_size.get() {
334            warn!(
335                name = %self.conf.name,
336                node = %self.node,
337                len  = %msg.len(),
338                max  = %self.conf.max_message_size,
339                "message too large to send"
340            );
341            return Err(NetworkError::MessageTooLarge);
342        }
343        Ok(())
344    }
345
346    /// Check if the given slot is less than our lower bound.
347    fn lt_lower_bound(&self, s: Slot) -> bool {
348        let lower_bound = *self.next_slot.borrow();
349        if s < lower_bound {
350            warn!(
351                name = %self.conf.name,
352                node = %self.node,
353                slot = %s,
354                %lower_bound,
355                "slot below lower bound"
356            );
357            return true;
358        }
359        false
360    }
361}
362
363fn msg_bytes(cmd: &SendCommand) -> &[u8] {
364    match &cmd.action {
365        SendAction::Unicast(_, b) => b,
366        SendAction::Multicast(_, b) => b,
367        SendAction::Broadcast(b) => b,
368    }
369}
cliquenet/net.rs

cliquenet/
net.rs
