Hold C2ME byte permits through ME writer completion

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/proxy/middle_relay/session.rs

@@ -149,6 +149,7 @@ where
                             peer,
                             translated_local_addr,
                             payload,
+                            _permit,
                             flags,
                             effective_tag_array,
                         )

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<OwnedSemaphorePermit>,
+    },
     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)]

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)) => {

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<PooledBuffer> {
+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<OwnedSemaphorePermit> {
+    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<Self>,
         conn_id: u64,
@@ -76,6 +86,7 @@ impl MePool {
         data: &[u8],
         proto_flags: u32,
         tag_override: Option<&[u8]>,
+        mut payload_permit: Option<OwnedSemaphorePermit>,
     ) -> 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<Self>,
         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
     }

src/transport/middle_proxy/tests/send_adversarial_tests.rs

@@ -165,7 +165,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, 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
     });