perf(fake-tcp) reduce the number of clone() calls in hot path

This commit is contained in:
Datong Sun 2021-11-18 20:31:57 -08:00
parent 91988520e5
commit 583cdbe300
3 changed files with 23 additions and 39 deletions

View File

@ -15,7 +15,6 @@ use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, RwLock};
use tokio::sync::broadcast;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tokio::sync::watch;
use tokio::sync::Mutex as AsyncMutex;
use tokio::time;
use tokio_tun::Tun;
@ -58,6 +57,7 @@ pub enum State {
SynSent,
SynReceived,
Established,
Closed,
}
pub struct Socket {
@ -69,8 +69,6 @@ pub struct Socket {
seq: AtomicU32,
ack: AtomicU32,
state: State,
closing_tx: watch::Sender<()>,
closing_rx: watch::Receiver<()>,
}
impl Socket {
@ -83,7 +81,6 @@ impl Socket {
state: State,
) -> (Socket, Sender<Bytes>) {
let (incoming_tx, incoming_rx) = mpsc::channel(MPSC_BUFFER_LEN);
let (closing_tx, closing_rx) = watch::channel(());
(
Socket {
@ -95,8 +92,6 @@ impl Socket {
seq: AtomicU32::new(0),
ack: AtomicU32::new(ack.unwrap_or(0)),
state,
closing_tx,
closing_rx,
},
incoming_tx,
)
@ -114,8 +109,6 @@ impl Socket {
}
pub async fn send(&self, payload: &[u8]) -> Option<()> {
let mut closing = self.closing_rx.clone();
match self.state {
State::Established => {
let buf = self.build_tcp_packet(tcp::TcpFlags::ACK, Some(payload));
@ -126,28 +119,22 @@ impl Socket {
res.unwrap();
Some(())
},
_ = closing.changed() => {
None
}
}
}
State::Closed => None,
_ => unreachable!(),
}
}
pub async fn recv(&self, buf: &mut [u8]) -> Option<usize> {
let mut closing = self.closing_rx.clone();
match self.state {
State::Established => {
let mut incoming = self.incoming.lock().await;
tokio::select! {
Some(raw_buf) = incoming.recv() => {
incoming.recv().await.and_then(|raw_buf| {
let (_v4_packet, tcp_packet) = parse_ipv4_packet(&raw_buf);
if (tcp_packet.get_flags() & tcp::TcpFlags::RST) != 0 {
info!("Connection {} reset by peer", self);
self.close();
return None;
}
@ -159,20 +146,13 @@ impl Socket {
buf[..payload.len()].copy_from_slice(payload);
Some(payload.len())
},
_ = closing.changed() => {
None
}
}
})
}
State::Closed => None,
_ => unreachable!(),
}
}
pub fn close(&self) {
self.closing_tx.send(()).unwrap();
}
async fn accept(mut self) {
for _ in 0..RETRIES {
match self.state {
@ -272,6 +252,8 @@ impl Socket {
impl Drop for Socket {
fn drop(&mut self) {
self.state = State::Closed;
let tuple = AddrTuple::new(self.local_addr, self.remote_addr);
// dissociates ourself from the dispatch map
assert!(self.shared.tuples.write().unwrap().remove(&tuple).is_some());
@ -282,7 +264,7 @@ impl Drop for Socket {
if let Err(e) = self.tun.try_send(&buf) {
warn!("Unable to send RST to remote end: {}", e);
}
self.close();
info!("Fake TCP connection to {} closed", self);
}
}

View File

@ -108,6 +108,7 @@ async fn main() {
} else {
panic!("only IPv4 remote address is supported");
};
info!("Remote address is: {}", remote_addr);
let tun_local: Ipv4Addr = matches
.value_of("tun_local")

View File

@ -77,6 +77,7 @@ async fn main() {
.expect("bad remote address or host")
.next()
.expect("unable to resolve remote host name");
info!("Remote address is: {}", remote_addr);
let tun_local: Ipv4Addr = matches
.value_of("tun_local")