ack_udp/
lib.rs

1use std::{
2  sync::Arc, 
3  net::SocketAddr,
4  io, 
5  collections::{HashMap, VecDeque, HashSet}
6};
7use tokio::{net::UdpSocket, sync::mpsc::{Sender, self}};
8
9
10use chrono::Utc;
11use parking_lot::Mutex;
12use rand::Rng;
13use sock_send::SockSend;
14use types::{
15  AckUdpDatagram, 
16  AckUdpDatagramOutStatus, 
17  AckUdpDatagramOutStatusEnum, 
18  AckUdpPacket
19};
20
21mod types;
22mod sock_send;
23mod methods;
24
25pub struct AckUdp {
26  pub sock: Arc<UdpSocket>,
27
28  pub ready_to_read_datagrams: Arc<Mutex<VecDeque<(SocketAddr, Vec<u8>)>>>,
29  pub pending_in_datagrams: Arc<Mutex<HashMap<[u8; 5], AckUdpDatagram>>>,
30  pub pending_out_datagrams: Arc<Mutex<HashMap<[u8; 5], AckUdpDatagram>>>,
31
32  pub out_datagrams_status_links: Arc<Mutex<HashMap<[u8; 5], Arc<Mutex<AckUdpDatagramOutStatus>>>>>,
33
34  pub kill_listener_channel_sender: Sender<()>,
35  pub kill_income_checker_channel_sender: Sender<()>,
36  pub kill_outcome_checker_listener_channel_sender: Sender<()>,
37  pub kill_incoming_queue_channel_sender: Sender<()>,
38
39  pub incoming_queue: Arc<Mutex<VecDeque<(SocketAddr, Vec<u8>)>>>
40}
41
42impl AckUdp {
43  pub async fn new(address: SocketAddr) -> io::Result<AckUdp> {
44    let sock = Arc::new(UdpSocket::bind(address).await?);
45
46    let (listener_sender, listener_receiver) = mpsc::channel(1);
47    let (income_checker_sender, income_checker_receiver) = mpsc::channel(1);
48    let (outcome_checker_sender, outcome_checker_receiver) = mpsc::channel(1);
49    let (incoming_queue_sender, incoming_queue_receiver) = mpsc::channel(1);
50
51    let instance = AckUdp {
52      sock: sock.clone(),
53      ready_to_read_datagrams: Arc::new(Mutex::new(VecDeque::new())),
54      pending_in_datagrams: Arc::new(Mutex::new(HashMap::new())),
55      pending_out_datagrams: Arc::new(Mutex::new(HashMap::new())),
56      out_datagrams_status_links: Arc::new(Mutex::new(HashMap::new())),
57      kill_listener_channel_sender: listener_sender,
58      kill_income_checker_channel_sender: income_checker_sender,
59      kill_outcome_checker_listener_channel_sender: outcome_checker_sender,
60      kill_incoming_queue_channel_sender: incoming_queue_sender,
61      incoming_queue: Arc::new(Mutex::new(VecDeque::new()))
62    };
63
64    // Check for dropped INcome datagrams
65    tokio::spawn(AckUdp::check_dropped_income(
66      income_checker_receiver, 
67      instance.pending_in_datagrams.clone())
68    );
69
70    // Check for dropped OUTcome datagrams and resend segments for datagrams that are not already dropped
71    tokio::spawn(AckUdp::check_dropped_outcome(
72      outcome_checker_receiver, 
73      instance.pending_out_datagrams.clone(), 
74      instance.sock.clone(), 
75      instance.out_datagrams_status_links.clone()
76    ));
77
78    // Listen for incoming packets on socket
79    tokio::spawn(AckUdp::listen_packets(listener_receiver, instance.sock.clone(), instance.incoming_queue.clone()));
80
81    // Proccess packet 
82    tokio::spawn(AckUdp::process_packets(
83      incoming_queue_receiver, 
84      instance.incoming_queue.clone(), 
85      instance.sock.clone(), 
86      instance.ready_to_read_datagrams.clone(), 
87      instance.pending_in_datagrams.clone(), 
88      instance.pending_out_datagrams.clone(), 
89      instance.out_datagrams_status_links.clone()
90    ));
91
92    Ok(instance)
93  }
94
95  pub fn recv(&mut self) -> Option<(SocketAddr, Vec<u8>)> {
96    self.ready_to_read_datagrams.lock().pop_front()
97  }
98
99  pub fn send(&mut self, buf: &[u8], address: SocketAddr ) -> io::Result<Arc<Mutex<AckUdpDatagramOutStatus>>> {
100    let datagram_id = rand::thread_rng().gen::<[u8; 5]>();
101    let status = Arc::new(Mutex::new(AckUdpDatagramOutStatus(AckUdpDatagramOutStatusEnum::Pending)));
102    if buf.len() > 400 {
103      let segments_count = {
104        let a = buf.len() as f64 / 400 as f64;
105        a.ceil() as u32
106      };
107
108      let segments = Arc::new(Mutex::new(HashMap::new()));
109
110      let datagram = AckUdpDatagram {
111        id: datagram_id,
112        address,
113        segments_count: segments_count,
114        segments: segments.clone(),
115        segments_got: Arc::new(Mutex::new(vec![])),
116        segments_acks: Arc::new(Mutex::new(HashSet::new())),
117        last_active: Utc::now(),
118        checks_failure_count: 0
119      };
120
121      self.pending_out_datagrams.lock().insert(datagram_id, datagram.clone());
122      self.out_datagrams_status_links.lock().insert(datagram_id, status.clone());
123
124      let buffer_length = buf.len();
125      for index in 0..segments_count {
126        let start = 400 * index as usize;
127        let end = {
128          let v = start + 400;
129          if v > buffer_length {
130            buffer_length
131          }
132          else {
133            v
134          }
135        };
136        let payload = &buf[start..end];
137        let packet = AckUdpPacket { 
138          datagram_id, 
139          seg_index: index, 
140          total_segments: segments_count, 
141          ack: 0, 
142          payload_size: payload.len() as u16, 
143          payload: payload.to_vec()
144        };
145        segments.lock().insert(index, packet.clone());
146        
147        let packet_bytes: Vec<u8> = packet.into();
148        self.sock.sock_send(packet_bytes, address);
149      }
150
151      let mut sent_datagram = self.pending_out_datagrams.lock().get(&datagram_id).unwrap().clone();
152      sent_datagram.last_active = Utc::now();
153      self.pending_out_datagrams.lock().insert(datagram_id, sent_datagram);
154
155      return Ok(status.clone());
156    }
157    else {
158      let packet = AckUdpPacket { 
159        datagram_id,
160        seg_index: 0,
161        total_segments: 1,
162        ack: 0,
163        payload_size: buf.len() as u16, 
164        payload: buf.to_vec()
165      };
166
167      let segments= Arc::new(Mutex::new(HashMap::new()));
168      segments.lock().insert(0, packet.clone());
169
170      let datagram = AckUdpDatagram {
171        id: datagram_id,
172        address,
173        segments_count: 1,
174        segments,
175        segments_got: Arc::new(Mutex::new(vec![])),
176        segments_acks: Arc::new(Mutex::new(HashSet::new())),
177        last_active: Utc::now(),
178        checks_failure_count: 0
179      };
180
181      self.pending_out_datagrams.lock().insert(datagram_id, datagram);
182      self.out_datagrams_status_links.lock().insert(datagram_id, status.clone());
183      
184      let packet_bytes: Vec<u8> = packet.into();
185      self.sock.sock_send(packet_bytes, address);
186
187      Ok(status.clone())
188    }
189  }
190}
ack_udp/lib.rs

ack_udp/
lib.rs
