use std::collections::HashMap;
use std::io::{BufReader, BufWriter};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use crossbeam::channel::Sender;
use log::{debug, info, warn};
use uuid::Uuid;
use super::codec::{read_msg, write_msg, WireMsg};
use crate::error::{Error, Result};
use crate::raft::{Envelope, NodeID, Request, Response};
#[derive(Clone)]
pub struct PeerOutbox {
addrs: Arc<HashMap<NodeID, SocketAddr>>,
conns: Arc<Mutex<HashMap<NodeID, TcpStream>>>,
}
impl PeerOutbox {
pub fn new(peers: Vec<(NodeID, SocketAddr)>) -> Self {
Self {
addrs: Arc::new(peers.into_iter().collect()),
conns: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn send_raft(&self, env: Envelope) -> Result<()> {
let to = env.to;
let msg = WireMsg::Raft(env);
if let Err(e) = self.send_to(to, &msg) {
self.invalidate(to);
if let Err(e2) = self.send_to(to, &msg) {
self.invalidate(to);
return Err(e2);
}
let _ = e;
}
Ok(())
}
fn send_to(&self, id: NodeID, msg: &WireMsg) -> Result<()> {
let addr = self
.addrs
.get(&id)
.copied()
.ok_or_else(|| Error::IO(format!("unknown peer {id}")))?;
let mut guard = self.conns.lock().expect("lock");
if !guard.contains_key(&id) {
let stream = Self::connect(addr)?;
guard.insert(id, stream);
}
let stream = guard.get_mut(&id).unwrap();
match write_msg(stream, msg) {
Ok(()) => Ok(()),
Err(e) => {
guard.remove(&id);
Err(e)
}
}
}
fn connect(addr: SocketAddr) -> Result<TcpStream> {
let stream = TcpStream::connect_timeout(&addr, Duration::from_secs(2))
.map_err(|e| Error::IO(format!("connect {addr}: {e}")))?;
stream.set_nodelay(true).ok();
stream.set_write_timeout(Some(Duration::from_secs(2))).ok();
stream.set_read_timeout(Some(Duration::from_secs(2))).ok();
Ok(stream)
}
pub fn invalidate(&self, id: NodeID) {
self.conns.lock().expect("lock").remove(&id);
}
}
pub enum Inbound {
Raft(Envelope),
Client {
id: Uuid,
request: Request,
reply: Sender<std::result::Result<Response, Error>>,
},
}
pub fn spawn_listener(addr: SocketAddr, inbound_tx: Sender<Inbound>) -> Result<()> {
let listener = TcpListener::bind(addr).map_err(|e| Error::IO(format!("bind {addr}: {e}")))?;
info!("Listening on {addr}");
thread::spawn(move || {
for conn in listener.incoming() {
match conn {
Ok(stream) => {
stream.set_nodelay(true).ok();
stream.set_read_timeout(Some(Duration::from_secs(30))).ok();
stream.set_write_timeout(Some(Duration::from_secs(2))).ok();
let tx = inbound_tx.clone();
thread::spawn(move || handle_conn(stream, tx));
}
Err(e) => warn!("accept error: {e}"),
}
}
});
Ok(())
}
fn handle_conn(stream: TcpStream, inbound_tx: Sender<Inbound>) {
let mut reader = BufReader::new(stream.try_clone().expect("clone"));
let writer = Arc::new(Mutex::new(BufWriter::new(stream)));
loop {
let msg = match read_msg(&mut reader) {
Ok(m) => m,
Err(e) => {
debug!("connection closed: {e}");
break;
}
};
match msg {
WireMsg::Raft(env) => {
if inbound_tx.send(Inbound::Raft(env)).is_err() {
break;
}
}
WireMsg::Client { id, request } => {
let (reply_tx, reply_rx) = crossbeam::channel::bounded(1);
if inbound_tx
.send(Inbound::Client { id, request, reply: reply_tx })
.is_err()
{
break;
}
match reply_rx.recv_timeout(Duration::from_secs(30)) {
Ok(response) => {
let mut w = writer.lock().expect("writer");
if write_msg(&mut *w, &WireMsg::ClientReply { id, response }).is_err() {
break;
}
}
Err(_) => {
let mut w = writer.lock().expect("writer");
let _ = write_msg(
&mut *w,
&WireMsg::ClientReply {
id,
response: Err(Error::IO("request timed out".into())),
},
);
}
}
}
WireMsg::ClientReply { .. } => {
}
}
}
}
pub fn run_client_request(
addr: SocketAddr,
request: Request,
timeout: Duration,
) -> Result<Response> {
let stream = TcpStream::connect_timeout(&addr, Duration::from_secs(2))
.map_err(|e| Error::IO(format!("connect {addr}: {e}")))?;
stream.set_nodelay(true).ok();
stream.set_read_timeout(Some(timeout)).ok();
stream.set_write_timeout(Some(timeout)).ok();
let mut stream = stream;
let id = Uuid::new_v4();
write_msg(&mut stream, &WireMsg::Client { id, request })?;
match read_msg(&mut stream)? {
WireMsg::ClientReply { id: rid, response } if rid == id => response,
other => Err(Error::InvalidData(format!("unexpected reply: {other:?}"))),
}
}