Perf improve (#59)

* improve perf

* fix forward
This commit is contained in:
Sijie.Sun
2024-04-26 23:02:07 +08:00
committed by GitHub
parent 096af6aa45
commit 69651ae3fd
16 changed files with 370 additions and 162 deletions

View File

@@ -9,11 +9,11 @@ use std::{
use futures::{stream::FuturesUnordered, Future, Sink, Stream};
use network_interface::NetworkInterfaceConfig as _;
use pin_project_lite::pin_project;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use bytes::{Buf, Bytes, BytesMut};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use tokio_stream::StreamExt;
use tokio_util::io::{poll_read_buf, poll_write_buf};
use tokio_util::io::poll_write_buf;
use zerocopy::FromBytes as _;
use crate::{
@@ -149,13 +149,18 @@ where
*self_mut.max_packet_size * 64,
);
match ready!(poll_read_buf(
self_mut.reader.as_mut(),
cx,
&mut self_mut.buf
)) {
Ok(size) => {
if size == 0 {
let cap = self_mut.buf.capacity() - self_mut.buf.len();
let buf = self_mut.buf.chunk_mut().as_mut_ptr();
let buf = unsafe { std::slice::from_raw_parts_mut(buf, cap) };
let mut buf = ReadBuf::new(buf);
let ret = ready!(self_mut.reader.as_mut().poll_read(cx, &mut buf));
let len = buf.filled().len();
unsafe { self_mut.buf.advance_mut(len) };
match ret {
Ok(_) => {
if len == 0 {
return Poll::Ready(None);
}
}
@@ -173,14 +178,16 @@ pub trait ZCPacketToBytes {
pub struct TcpZCPacketToBytes;
impl ZCPacketToBytes for TcpZCPacketToBytes {
fn into_bytes(&self, mut item: ZCPacket) -> Result<Bytes, TunnelError> {
fn into_bytes(&self, item: ZCPacket) -> Result<Bytes, TunnelError> {
let mut item = item.convert_type(ZCPacketType::TCP);
let tcp_len = PEER_MANAGER_HEADER_SIZE + item.payload_len();
let Some(header) = item.mut_tcp_tunnel_header() else {
return Err(TunnelError::InvalidPacket("packet too short".to_string()));
};
header.len.set(tcp_len.try_into().unwrap());
Ok(item.into_bytes(ZCPacketType::TCP))
Ok(item.into_bytes())
}
}

View File

@@ -96,23 +96,57 @@ const PAYLOAD_OFFSET_FOR_NIC_PACKET: usize = max(
WG_TUNNEL_HEADER_SIZE,
) + PEER_MANAGER_HEADER_SIZE;
const INVALID_OFFSET: usize = usize::MAX;
const fn get_converted_offset(old_hdr_size: usize, new_hdr_size: usize) -> usize {
if old_hdr_size < new_hdr_size {
INVALID_OFFSET
} else {
old_hdr_size - new_hdr_size
}
}
impl ZCPacketType {
pub fn get_packet_offsets(&self) -> ZCPacketOffsets {
match self {
ZCPacketType::TCP => ZCPacketOffsets {
payload_offset: TCP_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
peer_manager_header_offset: TCP_TUNNEL_HEADER_SIZE,
..Default::default()
tcp_tunnel_header_offset: 0,
udp_tunnel_header_offset: get_converted_offset(
TCP_TUNNEL_HEADER_SIZE,
UDP_TUNNEL_HEADER_SIZE,
),
wg_tunnel_header_offset: get_converted_offset(
TCP_TUNNEL_HEADER_SIZE,
WG_TUNNEL_HEADER_SIZE,
),
},
ZCPacketType::UDP => ZCPacketOffsets {
payload_offset: UDP_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
peer_manager_header_offset: UDP_TUNNEL_HEADER_SIZE,
..Default::default()
tcp_tunnel_header_offset: get_converted_offset(
UDP_TUNNEL_HEADER_SIZE,
TCP_TUNNEL_HEADER_SIZE,
),
udp_tunnel_header_offset: 0,
wg_tunnel_header_offset: get_converted_offset(
UDP_TUNNEL_HEADER_SIZE,
WG_TUNNEL_HEADER_SIZE,
),
},
ZCPacketType::WG => ZCPacketOffsets {
payload_offset: WG_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
peer_manager_header_offset: WG_TUNNEL_HEADER_SIZE,
..Default::default()
tcp_tunnel_header_offset: get_converted_offset(
WG_TUNNEL_HEADER_SIZE,
TCP_TUNNEL_HEADER_SIZE,
),
udp_tunnel_header_offset: get_converted_offset(
WG_TUNNEL_HEADER_SIZE,
UDP_TUNNEL_HEADER_SIZE,
),
wg_tunnel_header_offset: 0,
},
ZCPacketType::NIC => ZCPacketOffsets {
payload_offset: PAYLOAD_OFFSET_FOR_NIC_PACKET,
@@ -155,8 +189,10 @@ impl ZCPacket {
pub fn new_with_payload(payload: &[u8]) -> Self {
let mut ret = Self::new_nic_packet();
let total_len = ret.packet_type.get_packet_offsets().payload_offset + payload.len();
ret.inner.resize(total_len, 0);
let payload_off = ret.packet_type.get_packet_offsets().payload_offset;
let total_len = payload_off + payload.len();
ret.inner.reserve(total_len);
unsafe { ret.inner.set_len(total_len) };
ret.mut_payload()[..payload.len()].copy_from_slice(&payload);
ret
}
@@ -165,7 +201,7 @@ impl ZCPacket {
let mut ret = Self::new_nic_packet();
ret.inner.reserve(cap);
let total_len = ret.packet_type.get_packet_offsets().payload_offset - packet_info_len;
ret.inner.resize(total_len, 0);
unsafe { ret.inner.set_len(total_len) };
ret
}
@@ -275,45 +311,56 @@ impl ZCPacket {
hdr.len.set(payload_len as u32);
}
pub fn into_bytes(mut self, target_packet_type: ZCPacketType) -> Bytes {
fn tunnel_payload(&self) -> &[u8] {
&self.inner[self
.packet_type
.get_packet_offsets()
.peer_manager_header_offset..]
}
pub fn convert_type(mut self, target_packet_type: ZCPacketType) -> Self {
if target_packet_type == self.packet_type {
return self.inner.freeze();
} else {
assert_eq!(
self.packet_type,
ZCPacketType::NIC,
"only support NIC, got {:?}",
self
);
return self;
}
match target_packet_type {
ZCPacketType::TCP => self
.inner
.split_off(
self.packet_type
.get_packet_offsets()
.tcp_tunnel_header_offset,
)
.freeze(),
ZCPacketType::UDP => self
.inner
.split_off(
self.packet_type
.get_packet_offsets()
.udp_tunnel_header_offset,
)
.freeze(),
ZCPacketType::WG => self
.inner
.split_off(
self.packet_type
.get_packet_offsets()
.wg_tunnel_header_offset,
)
.freeze(),
let new_offset = match target_packet_type {
ZCPacketType::TCP => {
self.packet_type
.get_packet_offsets()
.tcp_tunnel_header_offset
}
ZCPacketType::UDP => {
self.packet_type
.get_packet_offsets()
.udp_tunnel_header_offset
}
ZCPacketType::WG => {
self.packet_type
.get_packet_offsets()
.wg_tunnel_header_offset
}
ZCPacketType::NIC => unreachable!(),
};
tracing::debug!(?self.packet_type, ?target_packet_type, ?new_offset, "convert zc packet type");
if new_offset == INVALID_OFFSET {
// copy peer manager header and payload to new buffer
let tunnel_payload = self.tunnel_payload();
let new_pm_offset = target_packet_type
.get_packet_offsets()
.peer_manager_header_offset;
let mut buf = BytesMut::with_capacity(new_pm_offset + tunnel_payload.len());
unsafe { buf.set_len(new_pm_offset) };
buf.extend_from_slice(tunnel_payload);
return Self::new_from_buf(buf, target_packet_type);
}
return Self::new_from_buf(self.inner.split_off(new_offset), target_packet_type);
}
pub fn into_bytes(self) -> Bytes {
self.inner.freeze()
}
pub fn inner(self) -> BytesMut {
@@ -349,7 +396,7 @@ mod tests {
assert_eq!(packet.payload_len(), 11);
println!("{:?}", packet.inner);
let tcp_packet = packet.into_bytes(ZCPacketType::TCP);
let tcp_packet = packet.convert_type(ZCPacketType::TCP).into_bytes();
assert_eq!(&tcp_packet[..1], b"\x0b");
println!("{:?}", tcp_packet);
}

View File

@@ -1,4 +1,4 @@
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering::Relaxed};
use std::sync::atomic::{AtomicU32, Ordering::Relaxed};
pub struct WindowLatency {
latency_us_window: Vec<AtomicU32>,
@@ -48,48 +48,49 @@ impl WindowLatency {
}
}
#[derive(Default)]
pub struct Throughput {
tx_bytes: AtomicU64,
rx_bytes: AtomicU64,
tx_bytes: u64,
rx_bytes: u64,
tx_packets: AtomicU64,
rx_packets: AtomicU64,
tx_packets: u64,
rx_packets: u64,
}
impl Throughput {
pub fn new() -> Self {
Self {
tx_bytes: AtomicU64::new(0),
rx_bytes: AtomicU64::new(0),
tx_packets: AtomicU64::new(0),
rx_packets: AtomicU64::new(0),
}
Self::default()
}
pub fn tx_bytes(&self) -> u64 {
self.tx_bytes.load(Relaxed)
self.tx_bytes
}
pub fn rx_bytes(&self) -> u64 {
self.rx_bytes.load(Relaxed)
self.rx_bytes
}
pub fn tx_packets(&self) -> u64 {
self.tx_packets.load(Relaxed)
self.tx_packets
}
pub fn rx_packets(&self) -> u64 {
self.rx_packets.load(Relaxed)
self.rx_packets
}
pub fn record_tx_bytes(&self, bytes: u64) {
self.tx_bytes.fetch_add(bytes, Relaxed);
self.tx_packets.fetch_add(1, Relaxed);
#[allow(invalid_reference_casting)]
unsafe {
*(&self.tx_bytes as *const u64 as *mut u64) += bytes;
*(&self.tx_packets as *const u64 as *mut u64) += 1;
}
}
pub fn record_rx_bytes(&self, bytes: u64) {
self.rx_bytes.fetch_add(bytes, Relaxed);
self.rx_packets.fetch_add(1, Relaxed);
#[allow(invalid_reference_casting)]
unsafe {
*(&self.rx_bytes as *const u64 as *mut u64) += bytes;
*(&self.rx_packets as *const u64 as *mut u64) += 1;
}
}
}

View File

@@ -125,20 +125,21 @@ async fn forward_from_ring_to_udp(
let Some(buf) = ring_recv.next().await else {
return None;
};
let mut packet = match buf {
let packet = match buf {
Ok(v) => v,
Err(e) => {
return Some(e);
}
};
let mut packet = packet.convert_type(ZCPacketType::UDP);
let udp_payload_len = packet.udp_payload().len();
let header = packet.mut_udp_tunnel_header().unwrap();
header.conn_id.set(conn_id);
header.len.set(udp_payload_len as u16);
header.msg_type = UdpPacketType::Data as u8;
let buf = packet.into_bytes(ZCPacketType::UDP);
let buf = packet.into_bytes();
tracing::trace!(?udp_payload_len, ?buf, "udp forward from ring to udp");
let ret = socket.send_to(&buf, &addr).await;
if ret.is_err() {
@@ -232,7 +233,7 @@ impl UdpTunnelListenerData {
tracing::info!(?conn_id, ?remote_addr, "udp connection accept handling",);
let socket = self.socket.as_ref().unwrap().clone();
let sack_buf = new_sack_packet(conn_id, magic).into_bytes(ZCPacketType::UDP);
let sack_buf = new_sack_packet(conn_id, magic).into_bytes();
if let Err(e) = socket.send_to(&sack_buf, remote_addr).await {
tracing::error!(?e, "udp send sack packet error");
return;
@@ -436,6 +437,7 @@ impl TunnelListener for UdpTunnelListener {
}
}
#[derive(Debug)]
pub struct UdpTunnelConnector {
addr: url::Url,
bind_addrs: Vec<SocketAddr>,
@@ -613,7 +615,7 @@ impl UdpTunnelConnector {
// send syn
let conn_id = rand::random();
let magic = rand::random();
let udp_packet = new_syn_packet(conn_id, magic).into_bytes(ZCPacketType::UDP);
let udp_packet = new_syn_packet(conn_id, magic).into_bytes();
let ret = socket.send_to(&udp_packet, &addr).await?;
tracing::warn!(?udp_packet, ?ret, "udp send syn");

View File

@@ -138,17 +138,15 @@ impl Debug for WgPeerData {
impl WgPeerData {
#[tracing::instrument]
async fn handle_one_packet_from_me(
&self,
mut zc_packet: ZCPacket,
) -> Result<(), anyhow::Error> {
async fn handle_one_packet_from_me(&self, zc_packet: ZCPacket) -> Result<(), anyhow::Error> {
let mut send_buf = vec![0u8; MAX_PACKET];
let packet = if matches!(self.wg_type, WgType::InternalUse) {
let mut zc_packet = zc_packet.convert_type(ZCPacketType::WG);
Self::fill_ip_header(&mut zc_packet);
zc_packet.into_bytes(ZCPacketType::WG)
zc_packet.into_bytes()
} else {
zc_packet.into_bytes(ZCPacketType::WG)
zc_packet.convert_type(ZCPacketType::WG).into_bytes()
};
tracing::trace!(?packet, "Sending packet to peer");
@@ -650,7 +648,7 @@ impl WgTunnelConnector {
let mut buf = vec![0u8; MAX_PACKET];
let (n, recv_addr) = data.udp.recv_from(&mut buf).await.unwrap();
if recv_addr != addr {
continue;
tracing::warn!(?recv_addr, "Received packet from changed address");
}
data.handle_one_packet_from_peer(&mut sink, &buf[..n]).await;
}