From f3554b153de9685bde9b97c37530a3940b4267e0 Mon Sep 17 00:00:00 2001 From: open-trade Date: Thu, 14 May 2020 19:16:46 +0000 Subject: [PATCH] sled_async --- Cargo.lock | 43 ++++++++++++++++++++ Cargo.toml | 3 ++ src/lib.rs | 2 + src/rendezvous_server.rs | 53 ++++++++++++++++++------ src/sled_async.rs | 88 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 176 insertions(+), 13 deletions(-) create mode 100644 src/sled_async.rs diff --git a/Cargo.lock b/Cargo.lock index a93d919..1758e15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -319,6 +319,9 @@ name = "hbbs" version = "0.1.0" dependencies = [ "hbb_common 0.1.0", + "serde 1.0.110 (registry+https://github.com/rust-lang/crates.io-index)", + "serde_derive 1.0.110 (registry+https://github.com/rust-lang/crates.io-index)", + "serde_json 1.0.53 (registry+https://github.com/rust-lang/crates.io-index)", "sled 0.31.0 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -346,6 +349,11 @@ dependencies = [ "libc 0.2.69 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "itoa" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "jobserver" version = "0.1.21" @@ -774,6 +782,11 @@ dependencies = [ "syn 1.0.19 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "ryu" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "schannel" version = "0.1.19" @@ -818,6 +831,31 @@ dependencies = [ "libc 0.2.69 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "serde" +version = "1.0.110" +source = "registry+https://github.com/rust-lang/crates.io-index" + +[[package]] +name = "serde_derive" +version = "1.0.110" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "proc-macro2 1.0.12 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)", + "syn 1.0.19 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "serde_json" +version = "1.0.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "itoa 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)", + "ryu 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.110 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "signal-hook-registry" version = "1.2.0" @@ -1190,6 +1228,7 @@ dependencies = [ "checksum hermit-abi 0.1.12 (registry+https://github.com/rust-lang/crates.io-index)" = "61565ff7aaace3525556587bd2dc31d4a07071957be715e63ce7b1eccf51a8f4" "checksum humantime 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "df004cfca50ef23c36850aaaa59ad52cc70d0e90243c3c7737a4dd32dc7a3c4f" "checksum iovec 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)" = "b2b3ea6ff95e175473f8ffe6a7eb7c00d054240321b84c57051175fe3c1e075e" +"checksum itoa 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "b8b7a7c0c47db5545ed3fef7468ee7bb5b74691498139e4b3f6a20685dc6dd8e" "checksum jobserver 0.1.21 (registry+https://github.com/rust-lang/crates.io-index)" = "5c71313ebb9439f74b00d9d2dcec36440beaf57a6aa0623068441dd7cd81a7f2" "checksum js-sys 0.3.39 (registry+https://github.com/rust-lang/crates.io-index)" = "fa5a448de267e7358beaf4a5d849518fe9a0c13fce7afd44b06e68550e5562a7" "checksum kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "7507624b29483431c0ba2d82aece8ca6cdba9382bff4ddd0f7490560c056098d" @@ -1239,11 +1278,15 @@ dependencies = [ "checksum rustls 0.17.0 (registry+https://github.com/rust-lang/crates.io-index)" = "c0d4a31f5d68413404705d6982529b0e11a9aacd4839d1d6222ee3b8cb4015e1" "checksum rustls-native-certs 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a75ffeb84a6bd9d014713119542ce415db3a3e4748f0bfce1e1416cd224a23a5" "checksum rustversion 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b3bba175698996010c4f6dce5e7f173b6eb781fce25d2cfc45e27091ce0b79f6" +"checksum ryu 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "ed3d612bc64430efeb3f7ee6ef26d590dce0c43249217bddc62112540c7941e1" "checksum schannel 0.1.19 (registry+https://github.com/rust-lang/crates.io-index)" = "8f05ba609c234e60bee0d547fe94a4c7e9da733d1c962cf6e59efa4cd9c8bc75" "checksum scopeguard 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" "checksum sct 0.6.0 (registry+https://github.com/rust-lang/crates.io-index)" = "e3042af939fca8c3453b7af0f1c66e533a15a86169e39de2657310ade8f98d3c" "checksum security-framework 0.4.4 (registry+https://github.com/rust-lang/crates.io-index)" = "64808902d7d99f78eaddd2b4e2509713babc3dc3c85ad6f4c447680f3c01e535" "checksum security-framework-sys 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "17bf11d99252f512695eb468de5516e5cf75455521e69dfe343f3b74e4748405" +"checksum serde 1.0.110 (registry+https://github.com/rust-lang/crates.io-index)" = "99e7b308464d16b56eba9964e4972a3eee817760ab60d88c3f86e1fecb08204c" +"checksum serde_derive 1.0.110 (registry+https://github.com/rust-lang/crates.io-index)" = "818fbf6bfa9a42d3bfcaca148547aa00c7b915bec71d1757aa2d44ca68771984" +"checksum serde_json 1.0.53 (registry+https://github.com/rust-lang/crates.io-index)" = "993948e75b189211a9b31a7528f950c6adc21f9720b6438ff80a7fa2f864cea2" "checksum signal-hook-registry 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "94f478ede9f64724c5d173d7bb56099ec3e2d9fc2774aac65d34b8b890405f41" "checksum slab 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c111b5bd5695e56cffe5129854aa230b39c93a305372fdbb2668ca2394eea9f8" "checksum sled 0.31.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8fb6824dde66ad33bf20c6e8476f5b82b871bc8bc3c129a10ea2f7dae5060fa3" diff --git a/Cargo.toml b/Cargo.toml index 31a560c..cf87eff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,9 @@ edition = "2018" [dependencies] hbb_common = { path = "libs/hbb_common" } sled = "0.31" +serde_derive = "1.0" +serde = "1.0" +serde_json = "1.0" [workspace] members = ["libs/hbb_common"] 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>>, - db: sled::Db, + db: super::SledAsync, } impl PeerMap { fn new() -> ResultType { 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 { - self.map.read().unwrap().get(key).map(|x| x.clone()) + async fn get(&mut self, key: String) -> Option { + 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::(&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)> { // 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)), + Get((String, mpsc::Sender>)), + Close, +} + +#[derive(Clone)] +pub struct SledAsync { + db: sled::Db, + tx: Option>, +} + +impl SledAsync { + pub fn new(path: &str) -> ResultType { + Ok(Self { + db: sled::open(path)?, + tx: None, + }) + } + + pub fn run(&mut self) -> std::thread::JoinHandle<()> { + let (tx, rx) = mpsc::unbounded_channel::(); + 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) { + 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 { + if let Some(tx) = &self.tx { + let (tx_once, mut rx) = mpsc::channel::>(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) -> Option { + if let Some(v) = v { + if let Ok(v) = std::str::from_utf8(v) { + if let Ok(v) = serde_json::from_str::(&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)))); + } + } + } +}