use std::fs::File;
use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::{mpsc, Arc};
use std::{cmp, thread, time};
use crate::core::ser;
use crate::core::ser::FixedLength;
use crate::msg::{read_body, read_header, read_item, write_to_buf, MsgHeader, Type};
use crate::types::Error;
use crate::util::read_write::{read_exact, write_all};
use crate::util::{RateCounter, RwLock};
pub trait MessageHandler: Send + 'static {
fn consume<'a>(
&self,
msg: Message<'a>,
writer: &'a mut dyn Write,
received_bytes: Arc<RwLock<RateCounter>>,
) -> Result<Option<Response<'a>>, Error>;
}
macro_rules! try_break {
($chan:ident, $inner:expr) => {
match $inner {
Ok(v) => Some(v),
Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::WouldBlock => None,
Err(e) => {
let _ = $chan.send(e);
break;
}
}
};
}
pub struct Message<'a> {
pub header: MsgHeader,
stream: &'a mut dyn Read,
}
impl<'a> Message<'a> {
fn from_header(header: MsgHeader, stream: &'a mut dyn Read) -> Message<'a> {
Message { header, stream }
}
pub fn body<T: ser::Readable>(&mut self) -> Result<T, Error> {
read_body(&self.header, self.stream)
}
pub fn streaming_read<T: ser::Readable>(&mut self) -> Result<(T, u64), Error> {
read_item(self.stream)
}
pub fn copy_attachment(&mut self, len: usize, writer: &mut dyn Write) -> Result<usize, Error> {
let mut written = 0;
while written < len {
let read_len = cmp::min(8000, len - written);
let mut buf = vec![0u8; read_len];
read_exact(
&mut self.stream,
&mut buf[..],
time::Duration::from_secs(10),
true,
)?;
writer.write_all(&mut buf)?;
written += read_len;
}
Ok(written)
}
}
pub struct Response<'a> {
resp_type: Type,
body: Vec<u8>,
stream: &'a mut dyn Write,
attachment: Option<File>,
}
impl<'a> Response<'a> {
pub fn new<T: ser::Writeable>(
resp_type: Type,
body: T,
stream: &'a mut dyn Write,
) -> Response<'a> {
let body = ser::ser_vec(&body).unwrap();
Response {
resp_type,
body,
stream,
attachment: None,
}
}
fn write(mut self, sent_bytes: Arc<RwLock<RateCounter>>) -> Result<(), Error> {
let mut msg =
ser::ser_vec(&MsgHeader::new(self.resp_type, self.body.len() as u64)).unwrap();
msg.append(&mut self.body);
write_all(&mut self.stream, &msg[..], time::Duration::from_secs(10))?;
{
let mut sent_bytes = sent_bytes.write();
sent_bytes.inc(msg.len() as u64);
}
if let Some(mut file) = self.attachment {
let mut buf = [0u8; 8000];
loop {
match file.read(&mut buf[..]) {
Ok(0) => break,
Ok(n) => {
write_all(&mut self.stream, &buf[..n], time::Duration::from_secs(10))?;
let mut sent_bytes = sent_bytes.write();
sent_bytes.inc_quiet(n as u64);
}
Err(e) => return Err(From::from(e)),
}
}
}
Ok(())
}
pub fn add_attachment(&mut self, file: File) {
self.attachment = Some(file);
}
}
pub const SEND_CHANNEL_CAP: usize = 10;
pub struct Tracker {
pub sent_bytes: Arc<RwLock<RateCounter>>,
pub received_bytes: Arc<RwLock<RateCounter>>,
pub send_channel: mpsc::SyncSender<Vec<u8>>,
pub close_channel: mpsc::Sender<()>,
pub error_channel: mpsc::Receiver<Error>,
}
impl Tracker {
pub fn send<T>(&self, body: T, msg_type: Type) -> Result<(), Error>
where
T: ser::Writeable,
{
let buf = write_to_buf(body, msg_type);
let buf_len = buf.len();
self.send_channel.try_send(buf)?;
let mut sent_bytes = self.sent_bytes.write();
sent_bytes.inc(buf_len as u64);
Ok(())
}
}
pub fn listen<H>(stream: TcpStream, handler: H) -> Tracker
where
H: MessageHandler,
{
let (send_tx, send_rx) = mpsc::sync_channel(SEND_CHANNEL_CAP);
let (close_tx, close_rx) = mpsc::channel();
let (error_tx, error_rx) = mpsc::channel();
let received_bytes = Arc::new(RwLock::new(RateCounter::new()));
let sent_bytes = Arc::new(RwLock::new(RateCounter::new()));
stream
.set_nonblocking(true)
.expect("Non-blocking IO not available.");
poll(
stream,
handler,
send_rx,
error_tx,
close_rx,
received_bytes.clone(),
sent_bytes.clone(),
);
Tracker {
sent_bytes: sent_bytes.clone(),
received_bytes: received_bytes.clone(),
send_channel: send_tx,
close_channel: close_tx,
error_channel: error_rx,
}
}
fn poll<H>(
conn: TcpStream,
handler: H,
send_rx: mpsc::Receiver<Vec<u8>>,
error_tx: mpsc::Sender<Error>,
close_rx: mpsc::Receiver<()>,
received_bytes: Arc<RwLock<RateCounter>>,
sent_bytes: Arc<RwLock<RateCounter>>,
) where
H: MessageHandler,
{
let mut reader = conn.try_clone().expect("clone conn for reader failed");
let mut writer = conn.try_clone().expect("clone conn for writer failed");
let _ = thread::Builder::new()
.name("peer".to_string())
.spawn(move || {
let sleep_time = time::Duration::from_millis(1);
let mut retry_send = Err(());
loop {
if let Some(h) = try_break!(error_tx, read_header(&mut reader, None)) {
let msg = Message::from_header(h, &mut reader);
trace!(
"Received message header, type {:?}, len {}.",
msg.header.msg_type,
msg.header.msg_len
);
let received = received_bytes.clone();
{
let mut received_bytes = received_bytes.write();
received_bytes.inc(MsgHeader::LEN as u64 + msg.header.msg_len);
}
if let Some(Some(resp)) =
try_break!(error_tx, handler.consume(msg, &mut writer, received))
{
try_break!(error_tx, resp.write(sent_bytes.clone()));
}
}
let maybe_data = retry_send.or_else(|_| send_rx.try_recv());
retry_send = Err(());
if let Ok(data) = maybe_data {
let written =
try_break!(error_tx, writer.write_all(&data[..]).map_err(&From::from));
if written.is_none() {
retry_send = Ok(data);
}
}
if let Ok(_) = close_rx.try_recv() {
debug!(
"Connection close with {} initiated by us",
conn.peer_addr()
.map(|a| a.to_string())
.unwrap_or("?".to_owned())
);
break;
}
thread::sleep(sleep_time);
}
let _ = conn.shutdown(Shutdown::Both);
});
}