LCOV - code coverage report
Current view: top level - libs/desim/src - network.rs (source / functions) Coverage Total Hit
Test: b9d67f908f91f00e353a27440ba89f642a869959.info Lines: 92.9 % 269 250
Test Date: 2024-11-19 21:44:13 Functions: 95.8 % 24 23

            Line data    Source code
       1              : use std::{
       2              :     cmp::Ordering,
       3              :     collections::{BinaryHeap, VecDeque},
       4              :     fmt::{self, Debug},
       5              :     ops::DerefMut,
       6              :     sync::{mpsc, Arc},
       7              : };
       8              : 
       9              : use parking_lot::{
      10              :     lock_api::{MappedMutexGuard, MutexGuard},
      11              :     Mutex, RawMutex,
      12              : };
      13              : use rand::rngs::StdRng;
      14              : use tracing::debug;
      15              : 
      16              : use crate::{
      17              :     executor::{self, ThreadContext},
      18              :     options::NetworkOptions,
      19              :     proto::NetEvent,
      20              :     proto::NodeEvent,
      21              : };
      22              : 
      23              : use super::{chan::Chan, proto::AnyMessage};
      24              : 
      25              : pub struct NetworkTask {
      26              :     options: Arc<NetworkOptions>,
      27              :     connections: Mutex<Vec<VirtualConnection>>,
      28              :     /// min-heap of connections having something to deliver.
      29              :     events: Mutex<BinaryHeap<Event>>,
      30              :     task_context: Arc<ThreadContext>,
      31              : }
      32              : 
      33              : impl NetworkTask {
      34          528 :     pub fn start_new(options: Arc<NetworkOptions>, tx: mpsc::Sender<Arc<NetworkTask>>) {
      35          528 :         let ctx = executor::get_thread_ctx();
      36          528 :         let task = Arc::new(Self {
      37          528 :             options,
      38          528 :             connections: Mutex::new(Vec::new()),
      39          528 :             events: Mutex::new(BinaryHeap::new()),
      40          528 :             task_context: ctx,
      41          528 :         });
      42          528 : 
      43          528 :         // send the task upstream
      44          528 :         tx.send(task.clone()).unwrap();
      45          528 : 
      46          528 :         // start the task
      47          528 :         task.start();
      48          528 :     }
      49              : 
      50        35330 :     pub fn start_new_connection(self: &Arc<Self>, rng: StdRng, dst_accept: Chan<NodeEvent>) -> TCP {
      51        35330 :         let now = executor::now();
      52        35330 :         let connection_id = self.connections.lock().len();
      53        35330 : 
      54        35330 :         let vc = VirtualConnection {
      55        35330 :             connection_id,
      56        35330 :             dst_accept,
      57        35330 :             dst_sockets: [Chan::new(), Chan::new()],
      58        35330 :             state: Mutex::new(ConnectionState {
      59        35330 :                 buffers: [NetworkBuffer::new(None), NetworkBuffer::new(Some(now))],
      60        35330 :                 rng,
      61        35330 :             }),
      62        35330 :         };
      63        35330 :         vc.schedule_timeout(self);
      64        35330 :         vc.send_connect(self);
      65        35330 : 
      66        35330 :         let recv_chan = vc.dst_sockets[0].clone();
      67        35330 :         self.connections.lock().push(vc);
      68        35330 : 
      69        35330 :         TCP {
      70        35330 :             net: self.clone(),
      71        35330 :             conn_id: connection_id,
      72        35330 :             dir: 0,
      73        35330 :             recv_chan,
      74        35330 :         }
      75        35330 :     }
      76              : }
      77              : 
      78              : // private functions
      79              : impl NetworkTask {
      80              :     /// Schedule to wakeup network task (self) `after_ms` later to deliver
      81              :     /// messages of connection `id`.
      82       187049 :     fn schedule(&self, id: usize, after_ms: u64) {
      83       187049 :         self.events.lock().push(Event {
      84       187049 :             time: executor::now() + after_ms,
      85       187049 :             conn_id: id,
      86       187049 :         });
      87       187049 :         self.task_context.schedule_wakeup(after_ms);
      88       187049 :     }
      89              : 
      90              :     /// Get locked connection `id`.
      91       250624 :     fn get(&self, id: usize) -> MappedMutexGuard<'_, RawMutex, VirtualConnection> {
      92       250624 :         MutexGuard::map(self.connections.lock(), |connections| {
      93       250624 :             connections.get_mut(id).unwrap()
      94       250624 :         })
      95       250624 :     }
      96              : 
      97       175988 :     fn collect_pending_events(&self, now: u64, vec: &mut Vec<Event>) {
      98       175988 :         vec.clear();
      99       175988 :         let mut events = self.events.lock();
     100       351449 :         while let Some(event) = events.peek() {
     101       348519 :             if event.time > now {
     102       173058 :                 break;
     103       175461 :             }
     104       175461 :             let event = events.pop().unwrap();
     105       175461 :             vec.push(event);
     106              :         }
     107       175988 :     }
     108              : 
     109          528 :     fn start(self: &Arc<Self>) {
     110          528 :         debug!("started network task");
     111              : 
     112          528 :         let mut events = Vec::new();
     113              :         loop {
     114       176516 :             let now = executor::now();
     115       176516 :             self.collect_pending_events(now, &mut events);
     116              : 
     117       176516 :             for event in events.drain(..) {
     118       175461 :                 let conn = self.get(event.conn_id);
     119       175461 :                 conn.process(self);
     120       175461 :             }
     121              : 
     122              :             // block until wakeup
     123       175988 :             executor::yield_me(-1);
     124              :         }
     125              :     }
     126              : }
     127              : 
     128              : // 0 - from node(0) to node(1)
     129              : // 1 - from node(1) to node(0)
     130              : type MessageDirection = u8;
     131              : 
     132         1285 : fn sender_str(dir: MessageDirection) -> &'static str {
     133         1285 :     match dir {
     134          211 :         0 => "client",
     135         1074 :         1 => "server",
     136            0 :         _ => unreachable!(),
     137              :     }
     138         1285 : }
     139              : 
     140          347 : fn receiver_str(dir: MessageDirection) -> &'static str {
     141          347 :     match dir {
     142          158 :         0 => "server",
     143          189 :         1 => "client",
     144            0 :         _ => unreachable!(),
     145              :     }
     146          347 : }
     147              : 
     148              : /// Virtual connection between two nodes.
     149              : /// Node 0 is the creator of the connection (client),
     150              : /// and node 1 is the acceptor (server).
     151              : struct VirtualConnection {
     152              :     connection_id: usize,
     153              :     /// one-off chan, used to deliver Accept message to dst
     154              :     dst_accept: Chan<NodeEvent>,
     155              :     /// message sinks
     156              :     dst_sockets: [Chan<NetEvent>; 2],
     157              :     state: Mutex<ConnectionState>,
     158              : }
     159              : 
     160              : struct ConnectionState {
     161              :     buffers: [NetworkBuffer; 2],
     162              :     rng: StdRng,
     163              : }
     164              : 
     165              : impl VirtualConnection {
     166              :     /// Notify the future about the possible timeout.
     167       107428 :     fn schedule_timeout(&self, net: &NetworkTask) {
     168       107428 :         if let Some(timeout) = net.options.keepalive_timeout {
     169       107428 :             net.schedule(self.connection_id, timeout);
     170       107428 :         }
     171       107428 :     }
     172              : 
     173              :     /// Send the handshake (Accept) to the server.
     174        35330 :     fn send_connect(&self, net: &NetworkTask) {
     175        35330 :         let now = executor::now();
     176        35330 :         let mut state = self.state.lock();
     177        35330 :         let delay = net.options.connect_delay.delay(&mut state.rng);
     178        35330 :         let buffer = &mut state.buffers[0];
     179        35330 :         assert!(buffer.buf.is_empty());
     180        35330 :         assert!(!buffer.recv_closed);
     181        35330 :         assert!(!buffer.send_closed);
     182        35330 :         assert!(buffer.last_recv.is_none());
     183              : 
     184        35330 :         let delay = if let Some(ms) = delay {
     185        27116 :             ms
     186              :         } else {
     187         8214 :             debug!("NET: TCP #{} dropped connect", self.connection_id);
     188         8214 :             buffer.send_closed = true;
     189         8214 :             return;
     190              :         };
     191              : 
     192              :         // Send a message into the future.
     193        27116 :         buffer
     194        27116 :             .buf
     195        27116 :             .push_back((now + delay, AnyMessage::InternalConnect));
     196        27116 :         net.schedule(self.connection_id, delay);
     197        35330 :     }
     198              : 
     199              :     /// Transmit some of the messages from the buffer to the nodes.
     200       175461 :     fn process(&self, net: &Arc<NetworkTask>) {
     201       175461 :         let now = executor::now();
     202       175461 : 
     203       175461 :         let mut state = self.state.lock();
     204              : 
     205       526383 :         for direction in 0..2 {
     206       350922 :             self.process_direction(
     207       350922 :                 net,
     208       350922 :                 state.deref_mut(),
     209       350922 :                 now,
     210       350922 :                 direction as MessageDirection,
     211       350922 :                 &self.dst_sockets[direction ^ 1],
     212       350922 :             );
     213       350922 :         }
     214              : 
     215              :         // Close the one side of the connection by timeout if the node
     216              :         // has not received any messages for a long time.
     217       175461 :         if let Some(timeout) = net.options.keepalive_timeout {
     218       175461 :             let mut to_close = [false, false];
     219       526383 :             for direction in 0..2 {
     220       350922 :                 let buffer = &mut state.buffers[direction];
     221       350922 :                 if buffer.recv_closed {
     222        68700 :                     continue;
     223       282222 :                 }
     224       282222 :                 if let Some(last_recv) = buffer.last_recv {
     225       256948 :                     if now - last_recv >= timeout {
     226        52812 :                         debug!(
     227            0 :                             "NET: connection {} timed out at {}",
     228            0 :                             self.connection_id,
     229            0 :                             receiver_str(direction as MessageDirection)
     230              :                         );
     231        52812 :                         let node_idx = direction ^ 1;
     232        52812 :                         to_close[node_idx] = true;
     233       204136 :                     }
     234        25274 :                 }
     235              :             }
     236       175461 :             drop(state);
     237              : 
     238       350922 :             for (node_idx, should_close) in to_close.iter().enumerate() {
     239       350922 :                 if *should_close {
     240        52812 :                     self.close(node_idx);
     241       298110 :                 }
     242              :             }
     243            0 :         }
     244       175461 :     }
     245              : 
     246              :     /// Process messages in the buffer in the given direction.
     247       350922 :     fn process_direction(
     248       350922 :         &self,
     249       350922 :         net: &Arc<NetworkTask>,
     250       350922 :         state: &mut ConnectionState,
     251       350922 :         now: u64,
     252       350922 :         direction: MessageDirection,
     253       350922 :         to_socket: &Chan<NetEvent>,
     254       350922 :     ) {
     255       350922 :         let buffer = &mut state.buffers[direction as usize];
     256       350922 :         if buffer.recv_closed {
     257        68700 :             assert!(buffer.buf.is_empty());
     258       282222 :         }
     259              : 
     260       423020 :         while !buffer.buf.is_empty() && buffer.buf.front().unwrap().0 <= now {
     261        72098 :             let msg = buffer.buf.pop_front().unwrap().1;
     262        72098 : 
     263        72098 :             buffer.last_recv = Some(now);
     264        72098 :             self.schedule_timeout(net);
     265        72098 : 
     266        72098 :             if let AnyMessage::InternalConnect = msg {
     267        25900 :                 // TODO: assert to_socket is the server
     268        25900 :                 let server_to_client = TCP {
     269        25900 :                     net: net.clone(),
     270        25900 :                     conn_id: self.connection_id,
     271        25900 :                     dir: direction ^ 1,
     272        25900 :                     recv_chan: to_socket.clone(),
     273        25900 :                 };
     274        25900 :                 // special case, we need to deliver new connection to a separate channel
     275        25900 :                 self.dst_accept.send(NodeEvent::Accept(server_to_client));
     276        46198 :             } else {
     277        46198 :                 to_socket.send(NetEvent::Message(msg));
     278        46198 :             }
     279              :         }
     280       350922 :     }
     281              : 
     282              :     /// Try to send a message to the buffer, optionally dropping it and
     283              :     /// determining delivery timestamp.
     284        70966 :     fn send(&self, net: &NetworkTask, direction: MessageDirection, msg: AnyMessage) {
     285        70966 :         let now = executor::now();
     286        70966 :         let mut state = self.state.lock();
     287              : 
     288        70966 :         let (delay, close) = if let Some(ms) = net.options.send_delay.delay(&mut state.rng) {
     289        64846 :             (ms, false)
     290              :         } else {
     291         6120 :             (0, true)
     292              :         };
     293              : 
     294        70966 :         let buffer = &mut state.buffers[direction as usize];
     295        70966 :         if buffer.send_closed {
     296         7818 :             debug!(
     297            0 :                 "NET: TCP #{} dropped message {:?} (broken pipe)",
     298              :                 self.connection_id, msg
     299              :             );
     300         7818 :             return;
     301        63148 :         }
     302        63148 : 
     303        63148 :         if close {
     304         4568 :             debug!(
     305            0 :                 "NET: TCP #{} dropped message {:?} (pipe just broke)",
     306              :                 self.connection_id, msg
     307              :             );
     308         4568 :             buffer.send_closed = true;
     309         4568 :             return;
     310        58580 :         }
     311        58580 : 
     312        58580 :         if buffer.recv_closed {
     313         6075 :             debug!(
     314            0 :                 "NET: TCP #{} dropped message {:?} (recv closed)",
     315              :                 self.connection_id, msg
     316              :             );
     317         6075 :             return;
     318        52505 :         }
     319        52505 : 
     320        52505 :         // Send a message into the future.
     321        52505 :         buffer.buf.push_back((now + delay, msg));
     322        52505 :         net.schedule(self.connection_id, delay);
     323        70966 :     }
     324              : 
     325              :     /// Close the connection. Only one side of the connection will be closed,
     326              :     /// and no further messages will be delivered. The other side will not be notified.
     327        57009 :     fn close(&self, node_idx: usize) {
     328        57009 :         let mut state = self.state.lock();
     329        57009 :         let recv_buffer = &mut state.buffers[1 ^ node_idx];
     330        57009 :         if recv_buffer.recv_closed {
     331          347 :             debug!(
     332            0 :                 "NET: TCP #{} closed twice at {}",
     333            0 :                 self.connection_id,
     334            0 :                 sender_str(node_idx as MessageDirection),
     335              :             );
     336          347 :             return;
     337        56662 :         }
     338        56662 : 
     339        56662 :         debug!(
     340            0 :             "NET: TCP #{} closed at {}",
     341            0 :             self.connection_id,
     342            0 :             sender_str(node_idx as MessageDirection),
     343              :         );
     344        56662 :         recv_buffer.recv_closed = true;
     345        56662 :         for msg in recv_buffer.buf.drain(..) {
     346         4841 :             debug!(
     347            0 :                 "NET: TCP #{} dropped message {:?} (closed)",
     348              :                 self.connection_id, msg
     349              :             );
     350              :         }
     351              : 
     352        56662 :         let send_buffer = &mut state.buffers[node_idx];
     353        56662 :         send_buffer.send_closed = true;
     354        56662 :         drop(state);
     355        56662 : 
     356        56662 :         // TODO: notify the other side?
     357        56662 : 
     358        56662 :         self.dst_sockets[node_idx].send(NetEvent::Closed);
     359        57009 :     }
     360              : }
     361              : 
     362              : struct NetworkBuffer {
     363              :     /// Messages paired with time of delivery
     364              :     buf: VecDeque<(u64, AnyMessage)>,
     365              :     /// True if the connection is closed on the receiving side,
     366              :     /// i.e. no more messages from the buffer will be delivered.
     367              :     recv_closed: bool,
     368              :     /// True if the connection is closed on the sending side,
     369              :     /// i.e. no more messages will be added to the buffer.
     370              :     send_closed: bool,
     371              :     /// Last time a message was delivered from the buffer.
     372              :     /// If None, it means that the server is the receiver and
     373              :     /// it has not yet aware of this connection (i.e. has not
     374              :     /// received the Accept).
     375              :     last_recv: Option<u64>,
     376              : }
     377              : 
     378              : impl NetworkBuffer {
     379        70660 :     fn new(last_recv: Option<u64>) -> Self {
     380        70660 :         Self {
     381        70660 :             buf: VecDeque::new(),
     382        70660 :             recv_closed: false,
     383        70660 :             send_closed: false,
     384        70660 :             last_recv,
     385        70660 :         }
     386        70660 :     }
     387              : }
     388              : 
     389              : /// Single end of a bidirectional network stream without reordering (TCP-like).
     390              : /// Reads are implemented using channels, writes go to the buffer inside VirtualConnection.
     391              : pub struct TCP {
     392              :     net: Arc<NetworkTask>,
     393              :     conn_id: usize,
     394              :     dir: MessageDirection,
     395              :     recv_chan: Chan<NetEvent>,
     396              : }
     397              : 
     398              : impl Debug for TCP {
     399          901 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     400          901 :         write!(f, "TCP #{} ({})", self.conn_id, sender_str(self.dir),)
     401          901 :     }
     402              : }
     403              : 
     404              : impl TCP {
     405              :     /// Send a message to the other side. It's guaranteed that it will not arrive
     406              :     /// before the arrival of all messages sent earlier.
     407        70966 :     pub fn send(&self, msg: AnyMessage) {
     408        70966 :         let conn = self.net.get(self.conn_id);
     409        70966 :         conn.send(&self.net, self.dir, msg);
     410        70966 :     }
     411              : 
     412              :     /// Get a channel to receive incoming messages.
     413       456048 :     pub fn recv_chan(&self) -> Chan<NetEvent> {
     414       456048 :         self.recv_chan.clone()
     415       456048 :     }
     416              : 
     417       310112 :     pub fn connection_id(&self) -> usize {
     418       310112 :         self.conn_id
     419       310112 :     }
     420              : 
     421         4197 :     pub fn close(&self) {
     422         4197 :         let conn = self.net.get(self.conn_id);
     423         4197 :         conn.close(self.dir as usize);
     424         4197 :     }
     425              : }
     426              : struct Event {
     427              :     time: u64,
     428              :     conn_id: usize,
     429              : }
     430              : 
     431              : // BinaryHeap is a max-heap, and we want a min-heap. Reverse the ordering here
     432              : // to get that.
     433              : impl PartialOrd for Event {
     434      1065275 :     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
     435      1065275 :         Some(self.cmp(other))
     436      1065275 :     }
     437              : }
     438              : 
     439              : impl Ord for Event {
     440      1065275 :     fn cmp(&self, other: &Self) -> Ordering {
     441      1065275 :         (other.time, other.conn_id).cmp(&(self.time, self.conn_id))
     442      1065275 :     }
     443              : }
     444              : 
     445              : impl PartialEq for Event {
     446            0 :     fn eq(&self, other: &Self) -> bool {
     447            0 :         (other.time, other.conn_id) == (self.time, self.conn_id)
     448            0 :     }
     449              : }
     450              : 
     451              : impl Eq for Event {}
        

Generated by: LCOV version 2.1-beta