netapp/src/client.rs

181 lines
4.5 KiB
Rust
Raw Normal View History

2020-12-02 12:30:47 +00:00
use std::collections::HashMap;
use std::net::SocketAddr;
2021-10-12 15:59:46 +00:00
use std::sync::atomic::{self, AtomicBool, AtomicU32};
use std::sync::{Arc, Mutex};
2020-12-02 12:30:47 +00:00
use log::{debug, error, trace};
2020-12-02 12:30:47 +00:00
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot, watch};
use tokio_util::compat::*;
use futures::io::AsyncReadExt;
2020-12-02 12:30:47 +00:00
use async_trait::async_trait;
use kuska_handshake::async_std::{handshake_client, BoxStream};
2020-12-02 12:30:47 +00:00
2021-10-12 15:59:46 +00:00
use crate::endpoint::*;
2020-12-02 12:30:47 +00:00
use crate::error::*;
use crate::netapp::*;
use crate::proto::*;
use crate::util::*;
2021-10-12 15:59:46 +00:00
2020-12-02 19:12:24 +00:00
pub(crate) struct ClientConn {
pub(crate) remote_addr: SocketAddr,
pub(crate) peer_id: NodeID,
2020-12-02 19:12:24 +00:00
query_send: mpsc::UnboundedSender<Option<(RequestID, RequestPriority, Vec<u8>)>>,
2021-10-12 15:59:46 +00:00
next_query_number: AtomicU32,
inflight: Mutex<HashMap<RequestID, oneshot::Sender<Vec<u8>>>>,
must_exit: AtomicBool,
stop_recv_loop: watch::Sender<bool>,
2020-12-02 12:30:47 +00:00
}
impl ClientConn {
pub(crate) async fn init(
netapp: Arc<NetApp>,
socket: TcpStream,
peer_id: NodeID,
2020-12-02 12:30:47 +00:00
) -> Result<(), Error> {
let remote_addr = socket.peer_addr()?;
let mut socket = socket.compat();
2020-12-02 12:30:47 +00:00
let handshake = handshake_client(
&mut socket,
2020-12-02 12:30:47 +00:00
netapp.netid.clone(),
2021-10-12 11:18:24 +00:00
netapp.id,
2020-12-02 12:30:47 +00:00
netapp.privkey.clone(),
2021-10-12 11:18:24 +00:00
peer_id,
2020-12-02 12:30:47 +00:00
)
.await?;
debug!(
"Handshake complete (client) with {}@{}",
hex::encode(&peer_id),
2020-12-02 12:30:47 +00:00
remote_addr
);
let (read, write) = socket.split();
2020-12-02 12:30:47 +00:00
2020-12-07 17:07:55 +00:00
let (read, write) =
2020-12-02 12:30:47 +00:00
BoxStream::from_handshake(read, write, handshake, 0x8000).split_read_write();
let (query_send, query_recv) = mpsc::unbounded_channel();
let (stop_recv_loop, stop_recv_loop_recv) = watch::channel(false);
2020-12-02 12:30:47 +00:00
let conn = Arc::new(ClientConn {
remote_addr,
2021-10-12 11:18:24 +00:00
peer_id,
2021-10-12 15:59:46 +00:00
next_query_number: AtomicU32::from(RequestID::default()),
2020-12-02 12:30:47 +00:00
query_send,
inflight: Mutex::new(HashMap::new()),
must_exit: AtomicBool::new(false),
stop_recv_loop,
2020-12-02 12:30:47 +00:00
});
2021-10-12 11:18:24 +00:00
netapp.connected_as_client(peer_id, conn.clone());
2020-12-02 12:30:47 +00:00
tokio::spawn(async move {
let conn2 = conn.clone();
let conn3 = conn.clone();
2020-12-07 17:07:55 +00:00
tokio::try_join!(conn2.send_loop(query_recv, write), async move {
tokio::select!(
r = conn3.recv_loop(read) => r,
_ = await_exit(stop_recv_loop_recv) => Ok(()),
)
})
2020-12-02 12:30:47 +00:00
.map(|_| ())
.log_err("ClientConn send_loop/recv_loop/dispatch_loop");
netapp.disconnected_as_client(&peer_id, conn);
2020-12-02 12:30:47 +00:00
});
Ok(())
}
pub fn close(&self) {
self.must_exit.store(true, atomic::Ordering::SeqCst);
self.query_send
.send(None)
.log_err("could not write None in query_send");
if self.inflight.lock().unwrap().is_empty() {
self.stop_recv_loop
2021-10-12 11:07:34 +00:00
.send(true)
.log_err("could not write true to stop_recv_loop");
2020-12-02 12:30:47 +00:00
}
}
2021-10-12 15:59:46 +00:00
pub(crate) async fn call<T>(
2020-12-02 12:30:47 +00:00
self: Arc<Self>,
rq: T,
2021-10-12 15:59:46 +00:00
path: &str,
2020-12-02 12:30:47 +00:00
prio: RequestPriority,
) -> Result<<T as Message>::Response, Error>
where
T: Message,
{
let id = self
.next_query_number
2021-10-12 15:59:46 +00:00
.fetch_add(1, atomic::Ordering::Relaxed);
2021-10-12 15:59:46 +00:00
let mut bytes = vec![prio, path.as_bytes().len() as u8];
bytes.extend_from_slice(path.as_bytes());
2020-12-02 12:30:47 +00:00
bytes.extend_from_slice(&rmp_to_vec_all_named(&rq)?[..]);
let (resp_send, resp_recv) = oneshot::channel();
let old = self.inflight.lock().unwrap().insert(id, resp_send);
if let Some(old_ch) = old {
error!(
"Too many inflight requests! RequestID collision. Interrupting previous request."
);
if old_ch.send(vec![]).is_err() {
debug!("Could not send empty response to collisionned request, probably because request was interrupted. Dropping response.");
}
}
2020-12-02 12:30:47 +00:00
trace!("request: query_send {}, {} bytes", id, bytes.len());
self.query_send.send(Some((id, prio, bytes)))?;
2020-12-02 12:30:47 +00:00
let resp = resp_recv.await?;
if resp.len() == 0 {
return Err(Error::Message("Response is 0 bytes, either a collision or a protocol error".into()));
}
2020-12-02 12:30:47 +00:00
2021-10-12 15:59:46 +00:00
let code = resp[0];
if code == 0 {
rmp_serde::decode::from_read_ref::<_, Result<<T as Message>::Response, String>>(
&resp[1..],
)?
2020-12-02 12:30:47 +00:00
.map_err(Error::Remote)
2021-10-12 15:59:46 +00:00
} else {
Err(Error::Remote(format!("Remote error code {}", code)))
}
2020-12-02 12:30:47 +00:00
}
}
impl SendLoop for ClientConn {}
#[async_trait]
impl RecvLoop for ClientConn {
async fn recv_handler(self: Arc<Self>, id: RequestID, msg: Vec<u8>) {
2020-12-11 12:34:04 +00:00
trace!("ClientConn recv_handler {} ({} bytes)", id, msg.len());
2020-12-07 15:00:12 +00:00
let mut inflight = self.inflight.lock().unwrap();
if let Some(ch) = inflight.remove(&id) {
if ch.send(msg).is_err() {
debug!("Could not send request response, probably because request was interrupted. Dropping response.");
}
}
if inflight.is_empty() && self.must_exit.load(atomic::Ordering::SeqCst) {
self.stop_recv_loop
2021-10-12 11:07:34 +00:00
.send(true)
.log_err("could not write true to stop_recv_loop");
}
2020-12-02 12:30:47 +00:00
}
}