feat(tunnel/usage): proxy writer 双优先级队列、hub 连接压力感知选择、usage 请求记录级别控制及 trace 页面 proxy timing 增强

apps/aether-proxy/src/tunnel/dispatcher.rs

@@ -404,10 +404,12 @@ mod tests {
 
     #[tokio::test]
     async fn try_send_stream_error_emits_stream_error_frame() {
-        let (frame_tx, mut frame_rx) = bounded_queue::<Frame>(4);
+        let (high_tx, mut high_rx) = bounded_queue::<Frame>(4);
+        let (normal_tx, _normal_rx) = bounded_queue::<Frame>(4);
+        let frame_tx = FrameSender::from_test_queues(high_tx, normal_tx);
         try_send_stream_error(&frame_tx, 9, "proxy request body dispatch stalled");
 
-        let frame = frame_rx
+        let frame = high_rx
             .recv()
             .await
             .expect("stream error frame should enqueue");

apps/aether-proxy/src/tunnel/stream_handler.rs

@@ -1164,15 +1164,20 @@ fn build_prefixed_request_body(
 mod tests {
     use std::collections::HashMap;
     use std::net::SocketAddr;
+    use std::pin::Pin;
     use std::sync::atomic::AtomicU64;
-    use std::sync::Once;
+    use std::sync::{Mutex, Once};
+    use std::task::{Context, Poll};
 
-    use aether_runtime::{bounded_queue, ConcurrencyGate, DistributedConcurrencyGate};
+    use aether_runtime::{ConcurrencyGate, DistributedConcurrencyGate};
     use arc_swap::ArcSwap;
     use axum::body::Body;
     use axum::http::{header, Response, StatusCode};
     use axum::routing::{get, post};
     use axum::Router;
+    use futures_util::Sink;
+    use tokio::task::JoinHandle;
+    use tokio_tungstenite::tungstenite::{Error as WebSocketError, Message};
 
     use super::*;
     use crate::config::Config;
@@ -1370,14 +1375,22 @@ mod tests {
         let state = sample_state_for_port(addr.port());
         cache_test_host(&state, host, addr).await;
         let server_ctx = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (_body_tx, body_rx) = mpsc::channel(1);
 
         let mut meta = sample_request_meta();
         meta.url = format!("http://{host}:{}/start", addr.port());
 
-        handle_stream(Arc::clone(&state), server_ctx, 5, meta, body_rx, frame_tx).await;
-        let result = collect_stream_result(&mut frame_rx).await;
+        handle_stream(
+            Arc::clone(&state),
+            server_ctx,
+            5,
+            meta,
+            body_rx,
+            frame_tx.clone(),
+        )
+        .await;
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
         server.abort();
 
         assert!(
@@ -1423,14 +1436,22 @@ mod tests {
         let state = sample_state_for_port(addr.port());
         cache_test_host(&state, host, addr).await;
         let server_ctx = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (_body_tx, body_rx) = mpsc::channel(1);
 
         let mut meta = sample_request_meta();
         meta.url = format!("http://{host}:{}/ok", addr.port());
 
-        handle_stream(Arc::clone(&state), server_ctx, 3, meta, body_rx, frame_tx).await;
-        let result = collect_stream_result(&mut frame_rx).await;
+        handle_stream(
+            Arc::clone(&state),
+            server_ctx,
+            3,
+            meta,
+            body_rx,
+            frame_tx.clone(),
+        )
+        .await;
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
         server.abort();
 
         assert!(
@@ -1486,7 +1507,7 @@ mod tests {
         let state = sample_state_for_port(addr.port());
         cache_test_host(&state, host, addr).await;
         let server_ctx = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (body_tx, body_rx) = mpsc::channel(4);
         body_tx
             .send(TunnelFrame::new(
@@ -1504,8 +1525,16 @@ mod tests {
         meta.url = format!("http://{host}:{}/start", addr.port());
         meta.follow_redirects = Some(true);
 
-        handle_stream(Arc::clone(&state), server_ctx, 1, meta, body_rx, frame_tx).await;
-        let result = collect_stream_result(&mut frame_rx).await;
+        handle_stream(
+            Arc::clone(&state),
+            server_ctx,
+            1,
+            meta,
+            body_rx,
+            frame_tx.clone(),
+        )
+        .await;
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
         server.abort();
 
         assert!(
@@ -1553,15 +1582,23 @@ mod tests {
         let state = sample_state_for_port(addr.port());
         cache_test_host(&state, host, addr).await;
         let server_ctx = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (_body_tx, body_rx) = mpsc::channel(1);
 
         let mut meta = sample_request_meta();
         meta.url = format!("http://{host}:{}/start", addr.port());
         meta.follow_redirects = Some(false);
 
-        handle_stream(Arc::clone(&state), server_ctx, 7, meta, body_rx, frame_tx).await;
-        let result = collect_stream_result(&mut frame_rx).await;
+        handle_stream(
+            Arc::clone(&state),
+            server_ctx,
+            7,
+            meta,
+            body_rx,
+            frame_tx.clone(),
+        )
+        .await;
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
         server.abort();
 
         assert!(
@@ -1618,7 +1655,7 @@ mod tests {
         let state = sample_state_for_budget(addr.port(), 0);
         cache_test_host(&state, host, addr).await;
         let server_ctx = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (body_tx, body_rx) = mpsc::channel(4);
         body_tx
             .send(TunnelFrame::new(
@@ -1636,8 +1673,16 @@ mod tests {
         meta.url = format!("http://{host}:{}/start", addr.port());
         meta.follow_redirects = Some(true);
 
-        handle_stream(Arc::clone(&state), server_ctx, 11, meta, body_rx, frame_tx).await;
-        let result = collect_stream_result(&mut frame_rx).await;
+        handle_stream(
+            Arc::clone(&state),
+            server_ctx,
+            11,
+            meta,
+            body_rx,
+            frame_tx.clone(),
+        )
+        .await;
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
         server.abort();
 
         assert!(
@@ -1663,7 +1708,7 @@ mod tests {
         let _permit = gate.try_acquire().expect("first permit");
         let state = sample_state(Some(gate), None);
         let server = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (_body_tx, body_rx) = mpsc::channel(1);
 
         handle_stream(
@@ -1672,11 +1717,15 @@ mod tests {
             7,
             sample_request_meta(),
             body_rx,
-            frame_tx,
+            frame_tx.clone(),
         )
         .await;
 
-        let frame = frame_rx.recv().await.expect("overload frame");
+        let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
+            .await
+            .into_iter()
+            .find(|frame| frame.msg_type == MsgType::StreamError)
+            .expect("overload frame");
         assert_eq!(frame.stream_id, 7);
         assert_eq!(frame.msg_type, MsgType::StreamError);
         assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
@@ -1700,7 +1749,7 @@ mod tests {
         let _permit = gate.try_acquire().await.expect("first permit");
         let state = sample_state(None, Some(gate));
         let server = sample_server(&state);
-        let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
         let (_body_tx, body_rx) = mpsc::channel(1);
 
         handle_stream(
@@ -1709,11 +1758,15 @@ mod tests {
             9,
             sample_request_meta(),
             body_rx,
-            frame_tx,
+            frame_tx.clone(),
         )
         .await;
 
-        let frame = frame_rx.recv().await.expect("overload frame");
+        let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
+            .await
+            .into_iter()
+            .find(|frame| frame.msg_type == MsgType::StreamError)
+            .expect("overload frame");
         assert_eq!(frame.stream_id, 9);
         assert_eq!(frame.msg_type, MsgType::StreamError);
         assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
@@ -1881,20 +1934,85 @@ mod tests {
             .await;
     }
 
+    #[derive(Clone, Default)]
+    struct VecSink {
+        sent: Arc<Mutex<Vec<Message>>>,
+    }
+
+    impl Sink<Message> for VecSink {
+        type Error = WebSocketError;
+
+        fn poll_ready(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
+            self.sent.lock().expect("sink lock").push(item);
+            Ok(())
+        }
+
+        fn poll_flush(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    fn spawn_test_writer() -> (FrameSender, Arc<Mutex<Vec<Message>>>, JoinHandle<()>) {
+        let sink = VecSink::default();
+        let sent = Arc::clone(&sink.sent);
+        let (frame_tx, handle) = crate::tunnel::writer::spawn_writer(sink, Duration::from_secs(60));
+        (frame_tx, sent, handle)
+    }
+
     struct StreamResult {
         response: Option<ResponseMeta>,
         body: Bytes,
         error: Option<String>,
     }
 
+    async fn collect_emitted_frames(
+        frame_tx: FrameSender,
+        sent: Arc<Mutex<Vec<Message>>>,
+        writer_handle: JoinHandle<()>,
+    ) -> Vec<TunnelFrame> {
+        drop(frame_tx);
+        writer_handle.await.expect("writer should exit cleanly");
+
+        sent.lock()
+            .expect("sink lock")
+            .iter()
+            .filter_map(|message| match message {
+                Message::Binary(data) => {
+                    Some(TunnelFrame::decode(data.clone().into()).expect("frame should decode"))
+                }
+                Message::Ping(_) | Message::Pong(_) | Message::Close(_) => None,
+                other => panic!("unexpected writer message: {other:?}"),
+            })
+            .collect()
+    }
+
     async fn collect_stream_result(
-        frame_rx: &mut aether_runtime::BoundedQueueReceiver<TunnelFrame>,
+        frame_tx: FrameSender,
+        sent: Arc<Mutex<Vec<Message>>>,
+        writer_handle: JoinHandle<()>,
     ) -> StreamResult {
         let mut response = None;
         let mut body = BytesMut::new();
         let mut error = None;
 
-        while let Some(frame) = frame_rx.recv().await {
+        for frame in collect_emitted_frames(frame_tx, sent, writer_handle).await {
             match frame.msg_type {
                 MsgType::ResponseHeaders => {
                     let payload = decompress_if_gzip(&frame).expect("headers payload");

apps/aether-proxy/src/tunnel/writer.rs

@@ -7,7 +7,10 @@
 
 use std::time::Duration;
 
-use aether_runtime::{bounded_queue, BoundedQueueSender};
+use aether_contracts::tunnel::MsgType;
+#[cfg(test)]
+use aether_runtime::QueueSnapshot;
+use aether_runtime::{bounded_queue, BoundedQueueSender, QueueSendError};
 use futures_util::SinkExt;
 use tokio::task::JoinHandle;
 use tokio_tungstenite::tungstenite::Message;
@@ -15,8 +18,60 @@ use tracing::{debug, error, trace};
 
 use super::protocol::Frame;
 
+const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
+const NORMAL_PRIORITY_QUEUE_CAPACITY: usize = 256;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum FramePriority {
+    High,
+    Normal,
+}
+
+#[cfg(test)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct FrameQueueSnapshots {
+    pub high: QueueSnapshot,
+    pub normal: QueueSnapshot,
+}
+
 /// Sender half — cloned by stream handlers and heartbeat.
-pub type FrameSender = BoundedQueueSender<Frame>;
+#[derive(Debug, Clone)]
+pub struct FrameSender {
+    high_tx: BoundedQueueSender<Frame>,
+    normal_tx: BoundedQueueSender<Frame>,
+}
+
+impl FrameSender {
+    pub async fn send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
+        match classify_frame_priority(&frame) {
+            FramePriority::High => self.high_tx.send(frame).await,
+            FramePriority::Normal => self.normal_tx.send(frame).await,
+        }
+    }
+
+    pub fn try_send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
+        match classify_frame_priority(&frame) {
+            FramePriority::High => self.high_tx.try_send(frame),
+            FramePriority::Normal => self.normal_tx.try_send(frame),
+        }
+    }
+
+    #[cfg(test)]
+    pub fn snapshots(&self) -> FrameQueueSnapshots {
+        FrameQueueSnapshots {
+            high: self.high_tx.snapshot(),
+            normal: self.normal_tx.snapshot(),
+        }
+    }
+
+    #[cfg(test)]
+    pub(crate) fn from_test_queues(
+        high_tx: BoundedQueueSender<Frame>,
+        normal_tx: BoundedQueueSender<Frame>,
+    ) -> Self {
+        Self { high_tx, normal_tx }
+    }
+}
 
 /// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
 ///
@@ -26,33 +81,56 @@ pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, Jo
 where
     S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
 {
-    let (tx, mut rx) = bounded_queue::<Frame>(256);
+    let (high_tx, mut high_rx) = bounded_queue::<Frame>(HIGH_PRIORITY_QUEUE_CAPACITY);
+    let (normal_tx, mut normal_rx) = bounded_queue::<Frame>(NORMAL_PRIORITY_QUEUE_CAPACITY);
+    let tx = FrameSender { high_tx, normal_tx };
 
     let handle = tokio::spawn(async move {
         let mut ping_ticker = tokio::time::interval(ping_interval);
+        let mut high_open = true;
+        let mut normal_open = true;
         ping_ticker.tick().await; // skip first immediate tick
 
         loop {
+            if let Ok(frame) = high_rx.try_recv() {
+                if !write_frame(&mut sink, frame).await {
+                    break;
+                }
+                continue;
+            }
+            if !high_open && !normal_open {
+                break;
+            }
+
             tokio::select! {
-                frame = rx.recv() => {
+                biased;
+                frame = high_rx.recv(), if high_open => {
                     match frame {
                         Some(frame) => {
-                            let data = frame.encode();
-                            if let Err(e) = sink.send(Message::Binary(data.into())).await {
-                                error!(error = %e, "failed to write frame to WebSocket");
+                            if !write_frame(&mut sink, frame).await {
                                 break;
                             }
                         }
-                        None => break, // all senders dropped
+                        None => high_open = false,
                     }
                 }
-                _ = ping_ticker.tick() => {
+                _ = ping_ticker.tick(), if high_open || normal_open => {
                     if let Err(e) = sink.send(Message::Ping(vec![])).await {
                         error!(error = %e, "failed to send WebSocket ping");
                         break;
                     }
                     trace!("sent WebSocket ping");
                 }
+                frame = normal_rx.recv(), if normal_open => {
+                    match frame {
+                        Some(frame) => {
+                            if !write_frame(&mut sink, frame).await {
+                                break;
+                            }
+                        }
+                        None => normal_open = false,
+                    }
+                }
             }
         }
         debug!("writer task exiting");
@@ -61,3 +139,134 @@ where
 
     (tx, handle)
 }
+
+fn classify_frame_priority(frame: &Frame) -> FramePriority {
+    match frame.msg_type {
+        MsgType::ResponseHeaders
+        | MsgType::StreamError
+        | MsgType::Ping
+        | MsgType::Pong
+        | MsgType::GoAway
+        | MsgType::HeartbeatData
+        | MsgType::HeartbeatAck => FramePriority::High,
+        MsgType::RequestHeaders
+        | MsgType::RequestBody
+        | MsgType::ResponseBody
+        | MsgType::StreamEnd => FramePriority::Normal,
+    }
+}
+
+async fn write_frame<S>(sink: &mut S, frame: Frame) -> bool
+where
+    S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
+{
+    let data = frame.encode();
+    if let Err(e) = sink.send(Message::Binary(data.into())).await {
+        error!(error = %e, "failed to write frame to WebSocket");
+        return false;
+    }
+    true
+}
+
+#[cfg(test)]
+mod tests {
+    use std::pin::Pin;
+    use std::sync::{Arc, Mutex};
+    use std::task::{Context, Poll};
+    use std::time::Duration;
+
+    use futures_util::Sink;
+    use tokio_tungstenite::tungstenite::{Error, Message};
+
+    use super::spawn_writer;
+    use crate::tunnel::protocol::Frame;
+    use aether_contracts::tunnel::MsgType;
+
+    #[derive(Clone, Default)]
+    struct VecSink {
+        sent: Arc<Mutex<Vec<Message>>>,
+    }
+
+    impl Sink<Message> for VecSink {
+        type Error = Error;
+
+        fn poll_ready(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
+            self.sent.lock().expect("sink lock").push(item);
+            Ok(())
+        }
+
+        fn poll_flush(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    #[tokio::test]
+    async fn prioritizes_control_frames_ahead_of_buffered_body_frames() {
+        let sink = VecSink::default();
+        let sent = Arc::clone(&sink.sent);
+        let (sender, handle) = spawn_writer(sink, Duration::from_secs(60));
+
+        for idx in 0..8u8 {
+            sender
+                .try_send(Frame::new(
+                    7,
+                    MsgType::ResponseBody,
+                    0,
+                    bytes::Bytes::from(vec![idx; 32]),
+                ))
+                .expect("frame send should succeed");
+        }
+        sender
+            .try_send(Frame::new(
+                7,
+                MsgType::StreamError,
+                0,
+                bytes::Bytes::from_static(b"boom"),
+            ))
+            .expect("frame send should succeed");
+        let snapshots = sender.snapshots();
+        assert!(snapshots.high.enqueued_total >= 1);
+        assert!(snapshots.normal.enqueued_total >= 8);
+
+        tokio::time::sleep(Duration::from_millis(30)).await;
+        drop(sender);
+        handle.await.expect("writer should exit cleanly");
+
+        let sent = sent.lock().expect("sink lock");
+        assert!(
+            sent.len() >= 2,
+            "writer should flush both body and control frames"
+        );
+        let first = match &sent[0] {
+            Message::Binary(data) => {
+                Frame::decode(data.clone().into()).expect("frame should decode")
+            }
+            other => panic!("unexpected first message: {other:?}"),
+        };
+        let second = match &sent[1] {
+            Message::Binary(data) => {
+                Frame::decode(data.clone().into()).expect("frame should decode")
+            }
+            other => panic!("unexpected second message: {other:?}"),
+        };
+        assert_eq!(first.msg_type, MsgType::StreamError);
+        assert_eq!(second.msg_type, MsgType::ResponseBody);
+    }
+}