connection changes and buffers changes

This commit is contained in:
2026-03-20 18:26:00 +07:00
parent b258db671b
commit 770ce7c63e
4 changed files with 97 additions and 42 deletions
+2
View File
@@ -2,3 +2,5 @@ pub mod dns;
pub mod ip_store; pub mod ip_store;
pub mod tcp_connection; pub mod tcp_connection;
pub mod udp_connection; pub mod udp_connection;
pub const CHANNEL_CAPACITY: usize = 2048;
+76 -36
View File
@@ -18,24 +18,33 @@ pub enum ConnectionState {
pub struct TcpConnection { pub struct TcpConnection {
handle: SocketHandle, handle: SocketHandle,
state: ConnectionState, state: ConnectionState,
tx: mpsc::UnboundedSender<Vec<u8>>, // UPLOAD: Ограниченный канал (Bounded) для Backpressure (128)
tx: mpsc::Sender<Vec<u8>>,
// DOWNLOAD: Безлимитный канал (Unbounded), чтобы Муксер НИКОГДА не зависал
rx: mpsc::UnboundedReceiver<Vec<u8>>, rx: mpsc::UnboundedReceiver<Vec<u8>>,
pending_data: Vec<u8>, pending_data: Vec<u8>,
handshake_rx: Option<oneshot::Receiver<()>>, handshake_rx: Option<oneshot::Receiver<()>>,
} }
const MAX_PENDING: usize = 2 * 1024 * 1024; // Подняли лимит буфера для загрузок до 32 Мегабайт.
// Этого хватит для Спидтеста на скорости ~200-300 Мбит/с без разрыва соединения.
const MAX_PENDING: usize = 32 * 1024 * 1024;
impl TcpConnection { impl TcpConnection {
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self { pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
let (tx_to_mux, mut rx_from_smol) = mpsc::unbounded_channel::<Vec<u8>>(); // UPLOAD: Ограничиваем очередь
let (tx_to_mux, mut rx_from_smol) = mpsc::channel::<Vec<u8>>(128);
// DOWNLOAD: Безлимитный канал до синхронного tick()
let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>(); let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>();
let (handshake_tx, handshake_rx) = oneshot::channel(); let (handshake_tx, handshake_rx) = oneshot::channel();
let stream_id = muxer.next_id(); let stream_id = muxer.next_id();
tokio::spawn(async move { tokio::spawn(async move {
let (v_tx, mut v_rx) = mpsc::channel(64); // ИСПРАВЛЕНИЕ: Даем Муксеру ограниченный канал, как он и просит (тип mpsc::Sender)
// Делаем его достаточно вместительным (1024)
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
muxer.register_stream(stream_id, v_tx).await; muxer.register_stream(stream_id, v_tx).await;
let connect_payload = target_addr.to_string(); let connect_payload = target_addr.to_string();
@@ -83,6 +92,9 @@ impl TcpConnection {
} }
}; };
// МАГИЯ ЗДЕСЬ: Мы мгновенно читаем из ограниченного v_rx и переливаем
// в безлимитный tx_to_smol. send() у безлимитного канала никогда не блокируется.
// Поэтому v_rx всегда пустой, и Муксер никогда не зависнет!
let from_proxy = async { let from_proxy = async {
while let Some(data) = v_rx.recv().await { while let Some(data) = v_rx.recv().await {
if data.is_empty() { if data.is_empty() {
@@ -113,7 +125,7 @@ impl TcpConnection {
handle, handle,
state: ConnectionState::Handshaking, state: ConnectionState::Handshaking,
tx: tx_to_mux, tx: tx_to_mux,
rx: rx_from_proxy, rx: rx_from_proxy, // Это UnboundedReceiver, здесь ничего менять не надо
pending_data: vec![], pending_data: vec![],
handshake_rx: Some(handshake_rx), handshake_rx: Some(handshake_rx),
} }
@@ -156,19 +168,17 @@ impl TcpConnection {
return false; return false;
} }
} }
ConnectionState::Closed => { ConnectionState::Closed => {
return false; return false;
} }
_ => {} _ => {}
} }
true true
} }
pub fn is_finished(&self, socket: &tcp::Socket) -> bool { pub fn is_finished(&self, socket: &tcp::Socket) -> bool {
use tcp::State; use tcp::State;
matches!(socket.state(), State::Closed | State::TimeWait) matches!(socket.state(), State::Closed | State::TimeWait)
} }
@@ -177,25 +187,69 @@ impl TcpConnection {
} }
fn poll_and_process(&mut self, socket: &mut tcp::Socket) { fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
// 1. UPLOAD: Чанкинг и Backpressure (Защита от краша сервера)
if socket.can_recv() { if socket.can_recv() {
let mut total_read = 0; while socket.can_recv() {
let _ = socket.recv(|data| { let mut channel_full = false;
if !data.is_empty() { let mut channel_closed = false;
if self.tx.send(data.to_vec()).is_ok() {
total_read = data.len(); let _ = socket.recv(|data| {
if data.is_empty() {
return (0, ());
}
let chunk_size = std::cmp::min(data.len(), 16000);
let chunk = &data[..chunk_size];
match self.tx.try_send(chunk.to_vec()) {
Ok(_) => (chunk_size, ()),
Err(mpsc::error::TrySendError::Full(_)) => {
channel_full = true;
(0, ()) // Оставляем данные в smoltcp
}
Err(_) => {
channel_closed = true;
(0, ())
}
}
});
if channel_full || channel_closed {
break;
}
}
}
// 2. DOWNLOAD: Сброс безлимитного канала в буфер
while let Ok(data) = self.rx.try_recv() {
if self.pending_data.is_empty() && socket.can_send() {
match socket.send_slice(&data) {
Ok(n) if n < data.len() => {
self.pending_data.extend_from_slice(&data[n..]);
}
Ok(_) => {}
Err(_) => {
self.pending_data.extend_from_slice(&data);
} }
} }
(total_read, ()) } else {
}); self.pending_data.extend_from_slice(&data);
} }
if self.pending_data.len() > MAX_PENDING { // Защита от OOM. 32MB более чем достаточно.
netrunner_logger::error!(%self.handle, "TCP Buffer overflow ({} bytes). Dropping connection.", self.pending_data.len()); if self.pending_data.len() > MAX_PENDING {
socket.abort(); netrunner_logger::error!(
self.state = ConnectionState::Closed; %self.handle,
return; "TCP Buffer overflow ({} bytes). Dropping connection.",
self.pending_data.len()
);
socket.abort();
self.state = ConnectionState::Closed;
return;
}
} }
// 3. DOWNLOAD: Отправка буфера в smoltcp
if !self.pending_data.is_empty() && socket.can_send() { if !self.pending_data.is_empty() && socket.can_send() {
match socket.send_slice(&self.pending_data) { match socket.send_slice(&self.pending_data) {
Ok(n) => { Ok(n) => {
@@ -204,19 +258,5 @@ impl TcpConnection {
Err(_) => {} Err(_) => {}
} }
} }
if self.pending_data.is_empty() && socket.can_send() {
if let Ok(data) = self.rx.try_recv() {
match socket.send_slice(&data) {
Ok(n) if n < data.len() => {
self.pending_data.extend_from_slice(&data[n..]);
}
Ok(_) => {}
Err(_) => {
self.pending_data = data;
}
}
}
}
} }
} }
+18 -5
View File
@@ -11,11 +11,13 @@ use netrunner_core::{
proxy::connection::muxer::{MuxMessage, Muxer}, proxy::connection::muxer::{MuxMessage, Muxer},
}; };
use crate::connections::CHANNEL_CAPACITY;
pub struct UdpConnection { pub struct UdpConnection {
pub handle: SocketHandle, pub handle: SocketHandle,
stream_id: u32, stream_id: u32,
tx_to_net: mpsc::Sender<MuxMessage>, tx_to_net: mpsc::Sender<MuxMessage>,
rx_from_net: mpsc::Receiver<Bytes>, rx_from_net: mpsc::UnboundedReceiver<Bytes>, // <-- ТУТ
client_endpoint: Option<IpEndpoint>, client_endpoint: Option<IpEndpoint>,
last_activity: Instant, last_activity: Instant,
token: CancellationToken, token: CancellationToken,
@@ -23,8 +25,6 @@ pub struct UdpConnection {
const UDP_TIMEOUT: Duration = Duration::from_secs(60); const UDP_TIMEOUT: Duration = Duration::from_secs(60);
const CHANNEL_CAPACITY: usize = 128;
impl UdpConnection { impl UdpConnection {
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self { pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
let stream_id = muxer.next_id(); let stream_id = muxer.next_id();
@@ -33,11 +33,14 @@ impl UdpConnection {
let (tx_to_net, mut rx_from_smol) = mpsc::channel::<MuxMessage>(CHANNEL_CAPACITY); let (tx_to_net, mut rx_from_smol) = mpsc::channel::<MuxMessage>(CHANNEL_CAPACITY);
let (v_tx, v_rx) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY); // ИСПРАВЛЕНИЕ: Внутренний безлимитный канал
let (internal_tx, internal_rx) = mpsc::unbounded_channel::<Bytes>();
let m_clone = muxer.clone(); let m_clone = muxer.clone();
tokio::spawn(async move { tokio::spawn(async move {
// Даем Муксеру ограниченный канал (тип Sender), чтобы компилятор пропустил
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
m_clone.register_stream(stream_id, v_tx).await; m_clone.register_stream(stream_id, v_tx).await;
let _ = m_clone let _ = m_clone
@@ -56,8 +59,18 @@ impl UdpConnection {
} }
}; };
// Мгновенно выкачиваем данные из Муксера в безлимитный внутренний канал
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if internal_tx.send(data).is_err() {
break;
}
}
};
tokio::select! { tokio::select! {
_ = to_proxy => {} _ = to_proxy => {}
_ = from_proxy => {} // Не забудь добавить этот блок в select!
_ = task_token.cancelled() => {} _ = task_token.cancelled() => {}
} }
@@ -75,7 +88,7 @@ impl UdpConnection {
handle, handle,
stream_id, stream_id,
tx_to_net, tx_to_net,
rx_from_net: v_rx, rx_from_net: internal_rx, // Сохраняем безлимитный ресивер
client_endpoint: None, client_endpoint: None,
last_activity: Instant::now(), last_activity: Instant::now(),
token, token,
+1 -1
View File
@@ -358,7 +358,7 @@ impl ConnectionManager {
} }
} }
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> { pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(256)); let mut sockets = SocketSet::new(Vec::with_capacity(1024));
for _ in 0..n_icmp { for _ in 0..n_icmp {
sockets.add(Self::create_icmp_socket()); sockets.add(Self::create_icmp_socket());