[add] ping jitter, receiver backpressure, stream frame limits

transport/Cargo.toml

@@ -16,6 +16,7 @@ tokio = { version = "1", features = ["full"] }
 rustls-native-certs = "0.8.4"
 log = "0.4"
 rcgen = "0.14"
+tracing = "0.1"
 
 [dev-dependencies]
 

transport/src/connection.rs

@@ -2,9 +2,10 @@ use crate::ConnectionHandle;
 use mtp_codec::CommunicationValue;
 use mtp_common::CommunicationError;
 use std::sync::Arc;
-use tokio::sync::{mpsc, Mutex, RwLock, Semaphore};
+use tokio::sync::{mpsc, Mutex, Notify, RwLock, Semaphore};
 use tokio::time::{Duration, sleep, timeout};
 use wtransport::Connection;
+use tracing::{debug, info, instrument, trace};
 
 const APPLICATION_CLOSE_REASON: &str = "mtp-close";
 
@@ -33,6 +34,7 @@ pub struct Policy {
     pub persistent_stream_retry_backoff: Duration,
     pub receiver_queue_capacity: usize,
     pub max_concurrent_stream_tasks: usize,
+    pub max_frames_per_stream: Option<usize>,
 }
 
 impl Default for Policy {
@@ -55,6 +57,7 @@ impl Default for Policy {
             persistent_stream_retry_backoff: Duration::from_millis(20),
             receiver_queue_capacity: 1000,
             max_concurrent_stream_tasks: 128,
+            max_frames_per_stream: None,
         }
     }
 }
@@ -114,6 +117,11 @@ impl Policy {
         self.max_concurrent_stream_tasks = max_concurrent_stream_tasks;
         self
     }
+
+    pub fn with_max_frames_per_stream(mut self, max_frames_per_stream: Option<usize>) -> Self {
+        self.max_frames_per_stream = max_frames_per_stream;
+        self
+    }
 }
 
 enum ReceivedFrame {
@@ -142,6 +150,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(stream, data, policy), level = "trace")]
     async fn write_frame(
         stream: &mut wtransport::SendStream,
         data: &CommunicationValue,
@@ -189,6 +198,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(conn, policy), level = "trace")]
     async fn open_uni_stream(
         conn: &Connection,
         policy: &Policy,
@@ -221,6 +231,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(conn, stream_opt, data, policy), level = "trace")]
     async fn send_on_persistent_stream(
         conn: &Connection,
         stream_opt: &mut Option<wtransport::SendStream>,
@@ -254,6 +265,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(conn, data, policy), level = "trace")]
     async fn send_on_single_stream(
         conn: &Connection,
         data: &CommunicationValue,
@@ -279,6 +291,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(conn, policy), level = "trace")]
     async fn send_close_frame(
         conn: &Connection,
         policy: &Policy,
@@ -319,6 +332,7 @@ impl Sender {
         Ok(())
     }
 
+    #[instrument(skip(self, data), level = "trace")]
     pub async fn send(&self, data: &CommunicationValue) -> Result<(), CommunicationError> {
         if self.handle.is_closed() {
             return Err(self
@@ -374,6 +388,7 @@ impl Sender {
         }
     }
 
+    #[instrument(skip(self), level = "trace")]
     pub async fn finish_stream(&self) -> Result<(), CommunicationError> {
         let _send_lock = self.send_guard.lock().await;
         let mut stream_opt = self.stream_guard.lock().await;
@@ -403,7 +418,9 @@ impl Sender {
         &self.handle
     }
 
+    #[instrument(skip(self), level = "trace")]
     pub fn close(&self) {
+        info!(target = "mtp.transport", "fire-and-forget close requested");
         let connection = self.connection.clone();
         let handle = self.handle.clone();
         let policy = self.policy.clone();
@@ -431,6 +448,7 @@ impl Sender {
             let _ = Self::send_close_frame(&connection, &policy).await;
 
             handle.close(Some(CommunicationError::StreamClosed));
+            info!(target = "mtp.transport", "connection closed");
 
             sleep(policy.force_close_delay).await;
             if connection.quic_connection().close_reason().is_none() {
@@ -442,6 +460,51 @@ impl Sender {
         });
     }
 
+    #[instrument(skip(self), level = "trace")]
+    /// Initiate a best-effort graceful close and wait for the configured force-close delay.
+    pub async fn close_and_wait(&self) {
+        info!(target = "mtp.transport", "graceful close initiated");
+        let connection = self.connection.clone();
+        let handle = self.handle.clone();
+        let policy = self.policy.clone();
+        let mut stream_opt = self.stream_guard.lock().await;
+
+        if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
+            handle.close(Some(CommunicationError::StreamClosed));
+            return;
+        }
+
+        if let Some(mut stream) = stream_opt.take() {
+            let close_bytes = policy.close_frame_len.to_be_bytes();
+            let close_write = async {
+                stream.write_all(&close_bytes).await?;
+                stream.finish().await
+            };
+
+            match timeout(policy.write_timeout, close_write).await {
+                Ok(Ok(())) => {}
+                Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => log::warn!(
+                    "[Sender] close_and_wait failed: peer sent STOP_SENDING (error code {code})"
+                ),
+                Ok(Err(e)) => log::warn!("[Sender] close_and_wait failed: {e}"),
+                Err(_) => log::warn!("[Sender] close_and_wait timed out"),
+            }
+        } else {
+            let _ = Self::send_close_frame(&connection, &policy).await;
+        }
+
+        handle.close(Some(CommunicationError::StreamClosed));
+        info!(target = "mtp.transport", "connection closed");
+
+        sleep(policy.force_close_delay).await;
+        if connection.quic_connection().close_reason().is_none() {
+            connection.quic_connection().close(
+                policy.application_close_code.into(),
+                APPLICATION_CLOSE_REASON.as_bytes(),
+            );
+        }
+    }
+
     pub fn is_open(&self) -> bool {
         self.handle.is_open()
     }
@@ -455,11 +518,17 @@ impl Sender {
     }
 }
 
+/// Single-consumer framed message receiver.
+///
+/// `receive()` is intended to be driven by one task at a time. Internally the
+/// underlying `mpsc::Receiver` is protected by a mutex so the type remains
+/// `Sync`, but it is not a multi-consumer queue.
 pub struct Receiver {
     rx: Mutex<mpsc::Receiver<Result<CommunicationValue, CommunicationError>>>,
     _accept_task: tokio::task::JoinHandle<()>,
     handle: Arc<ConnectionHandle>,
     ping_control: Arc<RwLock<PingControl>>,
+    queue_notify: Arc<Notify>,
 }
 
 #[derive(Clone, Default)]
@@ -491,13 +560,40 @@ impl Receiver {
         let accept_policy = policy.clone();
         let ping_control = Arc::new(RwLock::new(PingControl::default()));
         let accept_ping_control = ping_control.clone();
+        let queue_notify = Arc::new(Notify::new());
+        let accept_queue_notify = queue_notify.clone();
         let stream_limit = Arc::new(Semaphore::new(policy.max_concurrent_stream_tasks.max(1)));
         let accept_stream_limit = stream_limit.clone();
+        debug!(
+            target = "mtp.transport",
+            max_concurrent_stream_tasks = policy.max_concurrent_stream_tasks,
+            receiver_queue_capacity = policy.receiver_queue_capacity,
+            "receiver accept loop started"
+        );
+        info!(
+            target = "mtp.transport",
+            max_concurrent_stream_tasks = policy.max_concurrent_stream_tasks,
+            receiver_queue_capacity = policy.receiver_queue_capacity,
+            "connection accepted"
+        );
 
         let accept_task = tokio::spawn(async move {
             let mut close_rx = conn_handle.subscribe_close();
 
             loop {
+                if tx.capacity() == 0 {
+                    trace!(target = "mtp.transport", "accept loop paused: receiver queue full");
+                    tokio::select! {
+                        _ = close_rx.changed() => {
+                            if close_rx.borrow().is_some() {
+                                break;
+                            }
+                        }
+                        _ = accept_queue_notify.notified() => {}
+                    }
+                    continue;
+                }
+
                 tokio::select! {
                     _ = close_rx.changed() => {
                         if close_rx.borrow().is_some() {
@@ -523,9 +619,20 @@ impl Receiver {
                                 tokio::spawn(async move {
                                     let _permit = permit;
                                     let mut s = stream;
+                                    let mut frame_count = 0usize;
                                     loop {
+                                        if let Some(max_frames) = stream_policy.max_frames_per_stream
+                                            && frame_count >= max_frames
+                                        {
+                                            let close_error = CommunicationError::StreamError;
+                                            let _ = tx_stream.send(Err(close_error.clone())).await;
+                                            stream_handle.close(Some(close_error));
+                                            break;
+                                        }
+
                                         match Self::read_one_frame(&mut s, &stream_policy).await {
                                             Ok(ReceivedFrame::Message(msg)) => {
+                                                frame_count += 1;
                                                 let ping_type = mtp_codec::CommunicationType::Ping
                                                     .to_id(&mtp_codec::TypeMap::latest());
                                                 let pong_type = mtp_codec::CommunicationType::Pong
@@ -630,6 +737,7 @@ impl Receiver {
             _accept_task: accept_task,
             handle,
             ping_control,
+            queue_notify,
         }
     }
 
@@ -651,6 +759,7 @@ impl Receiver {
         }
     }
 
+    #[instrument(skip(stream, policy), level = "trace")]
     async fn read_one_frame(
         stream: &mut wtransport::RecvStream,
         policy: &Policy,
@@ -726,6 +835,7 @@ impl Receiver {
         Ok(ReceivedFrame::Message(message))
     }
 
+    #[instrument(skip(self), level = "trace")]
     pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
         if self.handle.is_closed() {
             return Err(self
@@ -736,7 +846,10 @@ impl Receiver {
 
         let mut rx = self.rx.lock().await;
         match rx.recv().await {
-            Some(result) => result,
+            Some(result) => {
+                self.queue_notify.notify_one();
+                result
+            }
             _ => Err(self
                 .handle
                 .close_reason()
@@ -796,6 +909,7 @@ mod tests {
         assert_eq!(p.persistent_stream_retry_backoff, Duration::from_millis(20));
         assert_eq!(p.receiver_queue_capacity, 1000);
         assert_eq!(p.max_concurrent_stream_tasks, 128);
+        assert_eq!(p.max_frames_per_stream, None);
     }
 
     #[test]

transport/tests/integration.rs

@@ -199,3 +199,149 @@ async fn test_drop_receiver_keeps_sender_alive() {
     client_tx.close();
     host_tx.close();
 }
+
+#[tokio::test]
+async fn test_persistent_stream_reopens_after_local_finish() {
+    let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await;
+
+    let tm = TypeMap::latest();
+    let msg1 = numbered_message(CommunicationType::Ping, 11, &tm);
+    client_tx.send(&msg1).await.unwrap();
+    let received1 = host_rx.receive().await.unwrap();
+    assert_numbered_message(&received1, CommunicationType::Ping, 11, &tm);
+
+    client_tx.finish_stream().await.unwrap();
+
+    let msg2 = numbered_message(CommunicationType::Pong, 22, &tm);
+    client_tx.send(&msg2).await.unwrap();
+    let received2 = host_rx.receive().await.unwrap();
+    assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm);
+
+    client_tx.close();
+    host_tx.close();
+    drop(client_rx);
+}
+
+#[tokio::test]
+async fn test_receiver_backpressure_with_small_queue() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let mut h = start_test_host(cert_pem.clone(), key_pem).await;
+    let url = format!("https://127.0.0.1:{}", h.local_addr().port());
+    let policy = Policy::default().with_receiver_queue_capacity(1);
+    let (client_tx, client_rx) = connect(&url, Some(cert_pem), policy.clone())
+        .await
+        .unwrap();
+    let (_host_tx, host_rx) = h.next().await.unwrap();
+
+    let tm = TypeMap::latest();
+    for i in 0..8u128 {
+        client_tx
+            .send(&numbered_message(CommunicationType::Ping, i, &tm))
+            .await
+            .unwrap();
+    }
+
+    for i in 0..8u128 {
+        let received = tokio::time::timeout(
+            std::time::Duration::from_secs(5),
+            host_rx.receive(),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
+    }
+
+    client_tx.close();
+    drop(client_rx);
+    h.shutdown();
+}
+
+#[tokio::test]
+async fn test_max_frames_per_stream_enforced() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let policy = Policy::default().with_max_frames_per_stream(Some(1));
+    let mut h = host(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        cert_pem.clone(),
+        key_pem,
+        policy,
+    )
+    .await
+    .unwrap();
+    let url = format!("https://127.0.0.1:{}", h.local_addr().port());
+    let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
+    let (_host_tx, host_rx) = h.next().await.unwrap();
+
+    let tm = TypeMap::latest();
+    client_tx
+        .send(&numbered_message(CommunicationType::Ping, 1, &tm))
+        .await
+        .unwrap();
+    let first = host_rx.receive().await.unwrap();
+    assert_numbered_message(&first, CommunicationType::Ping, 1, &tm);
+
+    client_tx
+        .send(&numbered_message(CommunicationType::Ping, 2, &tm))
+        .await
+        .unwrap();
+    let second = host_rx.receive().await;
+    assert!(second.is_err(), "stream should be closed after frame limit");
+
+    client_tx.close();
+    h.shutdown();
+}
+
+#[tokio::test]
+async fn test_semaphore_saturation_with_concurrent_streams() {
+    let (cert_pem, key_pem) = generate_self_signed_cert();
+    let policy = Policy::default()
+        .with_send_mode(mtp_transport::SendMode::SingleStreamPerMessage)
+        .with_receiver_queue_capacity(1)
+        .with_max_concurrent_stream_tasks(1);
+    let mut h = host(
+        IpAddr::V4(Ipv4Addr::LOCALHOST),
+        0,
+        cert_pem.clone(),
+        key_pem,
+        policy,
+    )
+    .await
+    .unwrap();
+    let url = format!("https://127.0.0.1:{}", h.local_addr().port());
+    let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
+        .await
+        .unwrap();
+    let (_host_tx, host_rx) = h.next().await.unwrap();
+
+    let tm = TypeMap::latest();
+    let mut joins = Vec::new();
+    for i in 0..6u128 {
+        let tx = client_tx.clone();
+        let msg = numbered_message(CommunicationType::Ping, i, &tm);
+        joins.push(tokio::spawn(async move { tx.send(&msg).await }));
+    }
+
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+    for join in joins {
+        join.await.unwrap().unwrap();
+    }
+
+    for i in 0..6u128 {
+        let received = tokio::time::timeout(
+            std::time::Duration::from_secs(5),
+            host_rx.receive(),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
+    }
+
+    client_tx.close();
+    h.shutdown();
+}