mirror of
https://github.com/rustdesk/rustdesk-server.git
synced 2025-07-03 16:25:27 +00:00
sled_async
This commit is contained in:
open-trade
parent
1f4f1cc8e2
commit
f3554b153d
43
Cargo.lock
generated
43
Cargo.lock
generated
@ -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"
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -1,2 +1,4 @@
|
||||
mod rendezvous_server;
|
||||
mod sled_async;
|
||||
use sled_async::*;
|
||||
pub use rendezvous_server::*;
|
||||
|
||||
@ -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((
|
||||
|
||||
88
src/sled_async.rs
Normal file
88
src/sled_async.rs
Normal file
@ -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))));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user