sled_async

3 files changed, 130 insertions, 13 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 6506a01..ae049bb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,2 +1,4 @@
 mod rendezvous_server;
+mod sled_async;
+use sled_async::*;
 pub use rendezvous_server::*;
diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index 32105aa..494a7ec 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -1,7 +1,6 @@
 use hbb_common::{
     allow_err,
-    bytes::Bytes,
-    bytes::BytesMut,
+    bytes::{Bytes, BytesMut},
     bytes_codec::BytesCodec,
     futures_util::{
         sink::SinkExt,
@@ -16,6 +15,7 @@ use hbb_common::{
     udp::FramedSocket,
     AddrMangle, ResultType,
 };
+use serde_derive::{Deserialize, Serialize};
 use std::{
     collections::HashMap,
     net::SocketAddr,
@@ -29,28 +29,47 @@ struct Peer {
     last_reg_time: Instant,
 }
 
+#[derive(Debug, Serialize, Deserialize, Default)]
+struct PeerSerde {
+    #[serde(default)]
+    ip: String,
+}
+
 #[derive(Clone)]
 struct PeerMap {
     map: Arc<RwLock<HashMap<String, Peer>>>,
-    db: sled::Db,
+    db: super::SledAsync,
 }
 
 impl PeerMap {
     fn new() -> ResultType<Self> {
         Ok(Self {
             map: Default::default(),
-            db: sled::open("./sled.db")?,
+            db: super::SledAsync::new("./sled.db")?,
         })
     }
 
     #[inline]
     fn insert(&mut self, key: String, peer: Peer) {
-        self.map.write().unwrap().insert(key, peer);
+        if self.map.write().unwrap().insert(key, peer).is_none() {}
     }
 
     #[inline]
-    fn get(&self, key: &str) -> Option<Peer> {
-        self.map.read().unwrap().get(key).map(|x| x.clone())
+    async fn get(&mut self, key: String) -> Option<Peer> {
+        let p = self.map.read().unwrap().get(&key).map(|x| x.clone());
+        if p.is_some() {
+            return p;
+        } else {
+            if let Some(_) = self.db.get(key).await {
+                // to-do
+            }
+        }
+        None
+    }
+
+    #[inline]
+    fn is_in_memory(&self, key: &str) -> bool {
+        self.map.read().unwrap().contains_key(key)
     }
 }
 
@@ -93,7 +112,7 @@ impl RendezvousServer {
                             if let Ok(msg_in) = parse_from_bytes::<RendezvousMessage>(&bytes) {
                                 match msg_in.union {
                                     Some(rendezvous_message::Union::punch_hole_request(ph)) => {
-                                        allow_err!(rs.handle_tcp_punch_hole_request(addr, &ph.id).await);
+                                        allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await);
                                     }
                                     Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
                                         allow_err!(rs.handle_hole_sent(&phs, addr, None).await);
@@ -141,7 +160,15 @@ impl RendezvousServer {
                     }
                 }
                 Some(rendezvous_message::Union::punch_hole_request(ph)) => {
-                    self.handle_udp_punch_hole_request(addr, &ph.id).await?;
+                    let id = ph.id;
+                    if self.pm.is_in_memory(&id) {
+                        self.handle_udp_punch_hole_request(addr, id).await?;
+                    } else {
+                        let mut me = self.clone();
+                        tokio::spawn(async move {
+                            allow_err!(me.handle_udp_punch_hole_request(addr, id).await);
+                        });
+                    }
                 }
                 Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
                     self.handle_hole_sent(&phs, addr, Some(socket)).await?;
@@ -215,14 +242,14 @@ impl RendezvousServer {
     async fn handle_punch_hole_request(
         &mut self,
         addr: SocketAddr,
-        id: &str,
+        id: String,
     ) -> ResultType<(RendezvousMessage, Option<SocketAddr>)> {
         // punch hole request from A, forward to B,
         // check if in same intranet first,
         // fetch local addrs if in same intranet.
         // because punch hole won't work if in the same intranet,
         // all routers will drop such self-connections.
-        if let Some(peer) = self.pm.get(id) {
+        if let Some(peer) = self.pm.get(id.clone()).await {
             if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT {
                 let mut msg_out = RendezvousMessage::new();
                 msg_out.set_punch_hole_response(PunchHoleResponse {
@@ -308,7 +335,7 @@ impl RendezvousServer {
     async fn handle_tcp_punch_hole_request(
         &mut self,
         addr: SocketAddr,
-        id: &str,
+        id: String,
     ) -> ResultType<()> {
         let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
         if let Some(addr) = to_addr {
@@ -323,7 +350,7 @@ impl RendezvousServer {
     async fn handle_udp_punch_hole_request(
         &mut self,
         addr: SocketAddr,
-        id: &str,
+        id: String,
     ) -> ResultType<()> {
         let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
         self.tx.send((
diff --git a/src/sled_async.rs b/src/sled_async.rs
new file mode 100644
index 0000000..1653ea1
--- /dev/null
+++ b/src/sled_async.rs
@@ -0,0 +1,88 @@
+use hbb_common::{
+    allow_err, log,
+    tokio::{self, sync::mpsc},
+    ResultType,
+};
+
+#[derive(Debug)]
+enum Action {
+    Insert((String, Vec<u8>)),
+    Get((String, mpsc::Sender<Option<sled::IVec>>)),
+    Close,
+}
+
+#[derive(Clone)]
+pub struct SledAsync {
+    db: sled::Db,
+    tx: Option<mpsc::UnboundedSender<Action>>,
+}
+
+impl SledAsync {
+    pub fn new(path: &str) -> ResultType<Self> {
+        Ok(Self {
+            db: sled::open(path)?,
+            tx: None,
+        })
+    }
+
+    pub fn run(&mut self) -> std::thread::JoinHandle<()> {
+        let (tx, rx) = mpsc::unbounded_channel::<Action>();
+        self.tx = Some(tx);
+        let db = self.db.clone();
+        std::thread::spawn(move || {
+            Self::io_loop(db, rx);
+        })
+    }
+
+    #[tokio::main(basic_scheduler)]
+    async fn io_loop(db: sled::Db, rx: mpsc::UnboundedReceiver<Action>) {
+        let mut rx = rx;
+        while let Some(x) = rx.recv().await {
+            match x {
+                Action::Insert((key, value)) => {
+                    allow_err!(db.insert(key, value));
+                }
+                Action::Get((key, sender)) => {
+                    let mut sender = sender;
+                    allow_err!(
+                        sender
+                            .send(if let Ok(v) = db.get(key) { v } else { None })
+                            .await
+                    );
+                }
+                Action::Close => break,
+            }
+        }
+    }
+
+    pub async fn get(&mut self, key: String) -> Option<sled::IVec> {
+        if let Some(tx) = &self.tx {
+            let (tx_once, mut rx) = mpsc::channel::<Option<sled::IVec>>(1);
+            allow_err!(tx.send(Action::Get((key, tx_once))));
+            if let Some(v) = rx.recv().await {
+                return v;
+            }
+        }
+        None
+    }
+
+    #[inline]
+    pub fn deserialize<'a, T: serde::Deserialize<'a>>(v: &'a Option<sled::IVec>) -> Option<T> {
+        if let Some(v) = v {
+            if let Ok(v) = std::str::from_utf8(v) {
+                if let Ok(v) = serde_json::from_str::<T>(&v) {
+                    return Some(v);
+                }
+            }
+        }
+        None
+    }
+
+    pub fn insert<'a, T: serde::Serialize>(&mut self, key: String, v: &T) {
+        if let Some(tx) = &self.tx {
+            if let Ok(v) = serde_json::to_vec(v) {
+                allow_err!(tx.send(Action::Insert((key, v))));
+            }
+        }
+    }
+}