udpflow/
sockmap.rs

1use std::net::SocketAddr;
2use std::sync::{Arc, RwLock};
3use std::collections::HashMap;
4
5use tokio::sync::mpsc::Sender;
6
7pub(crate) type Packet = Vec<u8>;
8
9#[derive(Clone)]
10pub(crate) struct SockMap(Arc<RwLock<HashMap<SocketAddr, Sender<Packet>>>>);
11
12impl SockMap {
13    pub fn new() -> Self { Self(Arc::new(RwLock::new(HashMap::new()))) }
14
15    #[inline]
16    pub fn get(&self, addr: &SocketAddr) -> Option<Sender<Packet>> {
17        // fetch the lock
18
19        let sockmap = self.0.read().unwrap();
20
21        sockmap.get(addr).cloned()
22
23        // drop the lock
24    }
25
26    #[inline]
27    pub fn insert(&self, addr: SocketAddr, tx: Sender<Packet>) {
28        // fetch the lock
29        let mut sockmap = self.0.write().unwrap();
30
31        let _ = sockmap.insert(addr, tx);
32
33        // drop the lock
34    }
35
36    #[inline]
37    pub fn remove(&self, addr: &SocketAddr) {
38        // fetch the lock
39        let mut sockmap = self.0.write().unwrap();
40
41        let _ = sockmap.remove(addr);
42
43        // drop the lock
44    }
45}