From 893ce0cf36a7d7429a288d1b73f9b596145ea7bc Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Fri, 10 Jul 2026 16:35:39 +0300 Subject: [PATCH] Hold C2ME byte permits through ME writer completion Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/proxy/middle_relay/session.rs | 1 + src/transport/middle_proxy/codec.rs | 7 ++- src/transport/middle_proxy/pool_writer.rs | 2 +- src/transport/middle_proxy/send.rs | 51 +++++++++++++++---- .../tests/send_adversarial_tests.rs | 8 ++- 5 files changed, 55 insertions(+), 14 deletions(-) diff --git a/src/proxy/middle_relay/session.rs b/src/proxy/middle_relay/session.rs index acded2b..ab90b97 100644 --- a/src/proxy/middle_relay/session.rs +++ b/src/proxy/middle_relay/session.rs @@ -149,6 +149,7 @@ where peer, translated_local_addr, payload, + _permit, flags, effective_tag_array, ) diff --git a/src/transport/middle_proxy/codec.rs b/src/transport/middle_proxy/codec.rs index f603293..235eb3a 100644 --- a/src/transport/middle_proxy/codec.rs +++ b/src/transport/middle_proxy/codec.rs @@ -1,5 +1,6 @@ use bytes::Bytes; use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::sync::OwnedSemaphorePermit; use crate::crypto::{AesCbc, crc32, crc32c}; use crate::error::{ProxyError, Result}; @@ -13,7 +14,10 @@ const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024; /// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes. pub(crate) enum WriterCommand { - Data(Bytes), + Data { + payload: Bytes, + _permit: Option, + }, DataAndFlush(Bytes), ProxyReq(ProxyReqCommand), ControlAndFlush([u8; 12]), @@ -28,6 +32,7 @@ pub(crate) struct ProxyReqCommand { pub(crate) proto_flags: u32, pub(crate) proxy_tag: Option<[u8; 16]>, pub(crate) payload: PooledBuffer, + pub(crate) _permit: OwnedSemaphorePermit, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index 4c10a3b..34e52c2 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -63,7 +63,7 @@ async fn writer_command_loop( _ = cancel.cancelled() => return Ok(()), cmd = rx.recv() => { match cmd { - Some(WriterCommand::Data(payload)) => { + Some(WriterCommand::Data { payload, _permit }) => { rpc_writer.send(&payload).await?; } Some(WriterCommand::DataAndFlush(payload)) => { diff --git a/src/transport/middle_proxy/send.rs b/src/transport/middle_proxy/send.rs index b830309..27120fe 100644 --- a/src/transport/middle_proxy/send.rs +++ b/src/transport/middle_proxy/send.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; -use tokio::sync::mpsc; +use tokio::sync::{OwnedSemaphorePermit, mpsc}; use tokio::sync::mpsc::error::TrySendError; use tracing::{debug, warn}; @@ -43,9 +43,18 @@ fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> { tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok()) } -fn proxy_req_payload_from_command(cmd: WriterCommand) -> Option { +fn proxy_req_payload_from_command( + cmd: WriterCommand, +) -> Option<(PooledBuffer, OwnedSemaphorePermit)> { match cmd { - WriterCommand::ProxyReq(command) => Some(command.payload), + WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)), + _ => None, + } +} + +fn payload_permit_from_data_command(cmd: WriterCommand) -> Option { + match cmd { + WriterCommand::Data { _permit, .. } => _permit, _ => None, } } @@ -67,6 +76,7 @@ async fn reserve_writer_command_slot( impl MePool { /// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default. + /// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command. pub async fn send_proxy_req( self: &Arc, conn_id: u64, @@ -76,6 +86,7 @@ impl MePool { data: &[u8], proto_flags: u32, tag_override: Option<&[u8]>, + mut payload_permit: Option, ) -> Result<()> { let tag = tag_override.or(self.proxy_tag.as_deref()); let build_routed_payload = |effective_our_addr: SocketAddr| { @@ -119,7 +130,11 @@ impl MePool { if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { let (current_payload, _) = build_routed_payload(current_meta.our_addr); - match current.tx.try_send(WriterCommand::Data(current_payload)) { + let command = WriterCommand::Data { + payload: current_payload, + _permit: payload_permit.take(), + }; + match current.tx.try_send(command) { Ok(()) => { self.note_hybrid_route_success(); return Ok(()); @@ -143,14 +158,17 @@ impl MePool { "ME writer channel full within blocking send timeout".into(), )); } - Err(WriterCommandReserveError::Closed) => {} + Err(WriterCommandReserveError::Closed) => { + payload_permit = payload_permit_from_data_command(cmd); + } } warn!(writer_id = current.writer_id, "ME writer channel closed"); self.remove_writer_and_close_clients(current.writer_id) .await; continue; } - Err(TrySendError::Closed(_)) => { + Err(TrySendError::Closed(cmd)) => { + payload_permit = payload_permit_from_data_command(cmd); warn!(writer_id = current.writer_id, "ME writer channel closed"); self.remove_writer_and_close_clients(current.writer_id) .await; @@ -479,7 +497,10 @@ impl MePool { self.remove_writer_and_close_clients(w.id).await; continue; } - permit.send(WriterCommand::Data(payload)); + permit.send(WriterCommand::Data { + payload, + _permit: payload_permit.take(), + }); self.stats .increment_me_writer_pick_success_try_total(pick_mode); if w.generation < self.current_generation() { @@ -548,7 +569,10 @@ impl MePool { self.remove_writer_and_close_clients(w.id).await; continue; } - permit.send(WriterCommand::Data(payload)); + permit.send(WriterCommand::Data { + payload, + _permit: payload_permit.take(), + }); self.stats .increment_me_writer_pick_success_fallback_total(pick_mode); if w.generation < self.current_generation() { @@ -567,6 +591,7 @@ impl MePool { } /// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light. + /// The client byte permit follows the payload until writer completion or command drop. pub async fn send_proxy_req_pooled( self: &Arc, conn_id: u64, @@ -574,6 +599,7 @@ impl MePool { client_addr: SocketAddr, our_addr: SocketAddr, payload: PooledBuffer, + _permit: OwnedSemaphorePermit, proto_flags: u32, tag_override: Option<[u8; 16]>, ) -> Result<()> { @@ -587,6 +613,7 @@ impl MePool { proto_flags, proxy_tag: tag, payload, + _permit, }); match current.tx.try_send(command) { Ok(()) => { @@ -613,7 +640,8 @@ impl MePool { )); } Err(WriterCommandReserveError::Closed) => { - let Some(payload) = proxy_req_payload_from_command(cmd) else { + let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) + else { return Err(ProxyError::Proxy( "ME writer rejected unexpected command type".into(), )); @@ -630,13 +658,14 @@ impl MePool { payload.as_ref(), proto_flags, tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), ) .await; } } } Err(TrySendError::Closed(cmd)) => { - let Some(payload) = proxy_req_payload_from_command(cmd) else { + let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else { return Err(ProxyError::Proxy( "ME writer rejected unexpected command type".into(), )); @@ -653,6 +682,7 @@ impl MePool { payload.as_ref(), proto_flags, tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), ) .await; } @@ -667,6 +697,7 @@ impl MePool { payload.as_ref(), proto_flags, tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), ) .await } diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 6c64555..61f7aa7 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -165,7 +165,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver, budget: Duratio while Instant::now().duration_since(start) < budget { let remaining = budget.saturating_sub(Instant::now().duration_since(start)); match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await { - Ok(Some(WriterCommand::Data(_))) => data_count += 1, + Ok(Some(WriterCommand::Data { .. })) => data_count += 1, Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1, Ok(Some(WriterCommand::ProxyReq(_))) => data_count += 1, Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1, @@ -185,7 +185,7 @@ async fn recv_first_data_payload( while Instant::now().duration_since(start) < budget { let remaining = budget.saturating_sub(Instant::now().duration_since(start)); match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await { - Ok(Some(WriterCommand::Data(payload))) => return Some(payload.to_vec()), + Ok(Some(WriterCommand::Data { payload, .. })) => return Some(payload.to_vec()), Ok(Some(WriterCommand::DataAndFlush(payload))) => return Some(payload.to_vec()), Ok(Some(_)) => {} Ok(None) => break, @@ -240,6 +240,7 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() { b"hello", 0, None, + None, ) .await; @@ -299,6 +300,7 @@ async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay b"storm", 0, None, + None, ) .await; @@ -356,6 +358,7 @@ async fn send_proxy_req_uses_writer_source_ip_when_advertised_our_addr_differs() b"route", 0, None, + None, ) .await; @@ -411,6 +414,7 @@ async fn send_proxy_req_blocking_fallback_uses_writer_source_ip() { b"blocking", 0, None, + None, ) .await });