refactor(tunnel): rename aether-proxy to aether-tunnel

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

@@ -0,0 +1,308 @@
+//! Dedicated WebSocket writer task.
+//!
+//! All frame writes go through an mpsc channel to a single writer task,
+//! avoiding contention on the WebSocket sink.  The writer also sends
+//! periodic WebSocket Ping frames to keep the connection alive through
+//! intermediary proxies (Nginx, Cloudflare, etc.).
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use aether_contracts::tunnel::{MsgType, HEADER_SIZE};
+#[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;
+use tracing::{debug, error, trace};
+
+use crate::state::TunnelMetrics;
+
+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.
+#[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.
+///
+/// `ping_interval` controls WebSocket-level Ping frequency (typically 15s).
+/// This keeps the connection alive through intermediary proxies/load-balancers.
+#[cfg(test)]
+pub fn spawn_writer<S>(sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
+where
+    S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
+{
+    spawn_writer_with_metrics(sink, ping_interval, None)
+}
+
+/// Spawn the writer task with optional tunnel metrics instrumentation.
+pub fn spawn_writer_with_metrics<S>(
+    mut sink: S,
+    ping_interval: Duration,
+    tunnel_metrics: Option<Arc<TunnelMetrics>>,
+) -> (FrameSender, JoinHandle<()>)
+where
+    S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
+{
+    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, tunnel_metrics.as_deref()).await {
+                    break;
+                }
+                continue;
+            }
+            if !high_open && !normal_open {
+                break;
+            }
+
+            tokio::select! {
+                biased;
+                frame = high_rx.recv(), if high_open => {
+                    match frame {
+                        Some(frame) => {
+                            if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
+                                break;
+                            }
+                        }
+                        None => high_open = false,
+                    }
+                }
+                _ = 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");
+                        if let Some(metrics) = tunnel_metrics.as_deref() {
+                            metrics.record_error("ws_ping_error", &e.to_string());
+                        }
+                        break;
+                    }
+                    trace!("sent WebSocket ping");
+                }
+                frame = normal_rx.recv(), if normal_open => {
+                    match frame {
+                        Some(frame) => {
+                            if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
+                                break;
+                            }
+                        }
+                        None => normal_open = false,
+                    }
+                }
+            }
+        }
+        debug!("writer task exiting");
+        let _ = sink.close().await;
+    });
+
+    (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, tunnel_metrics: Option<&TunnelMetrics>) -> bool
+where
+    S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
+{
+    let stream_id = frame.stream_id;
+    let msg_type = frame.msg_type;
+    let flags = frame.flags;
+    let data = frame.encode();
+    let wire_len = data.len().max(HEADER_SIZE);
+    if let Err(e) = sink.send(Message::Binary(data.into())).await {
+        error!(
+            stream_id = stream_id,
+            msg_type = ?msg_type,
+            flags = flags,
+            wire_len = wire_len,
+            error = %e,
+            "failed to write frame to WebSocket"
+        );
+        if let Some(metrics) = tunnel_metrics {
+            metrics.record_error("ws_write_error", &e.to_string());
+        }
+        return false;
+    }
+    if let Some(metrics) = tunnel_metrics {
+        metrics.record_ws_outgoing_frame(wire_len);
+    }
+    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);
+    }
+}