rutster/crates/rutster/src/tap_engine.rs

//! # TapEngine — per-session WSS task supervisor (spec §4.3, §5.1)
//!
//! Spawned by `session_map::drive_all_sessions` when the poll task observes
//! `channel.state == Connected && channel.tap.is_none()` (spec §5.1 step 3).
//! Owns the WSS connection lifecycle + the bounded-backoff reconnect policy.
//!
//! # Why this is NOT the step-4 forbidden "dedicated timing thread"
//!
//! ARCHITECTURE.md mandates "dedicated timing threads, not the shared
//! tokio pool" for the *timed media work* (the 20 ms loop). The TapEngine
//! is cold-path network I/O — connect_async, ws.send, ws.next — on tokio's
//! shared pool. Slice-1 §3.4's scoped deviation (media loop on tokio) is
//! unchanged; adding a cold-path I/O supervisor doesn't widen it.
//!
//! # Return shape
//!
//! `spawn_tap_engine` returns the tuple `(TapAudioPipe, TapConn)`. The
//! `TapAudioPipe` (the seam object) is wired into `RtcSession` via
//! `RtcSession::with_pipe(...)`; its mpsc ends are MOVED into the pipe.
//! `TapConn` holds only the engine-control handles (`close_tx`, `join`,
//! `metrics`) — the mpsc pairs are not split between the two (the earlier
//! brief sketch had `TapConn` carrying them; the structural review in
//! Task 7's plan corrected this so the pipe is the single owner of the
//! channel ends the loop_driver touches indirectly through `AudioSource`/
//! `AudioSink`).

use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use futures_util::FutureExt;
use rutster_call_model::ChannelId;
use rutster_tap::tap_client::{run_tap_client, TapClientError};
use rutster_tap::{TapAudioPipe, TapMetrics};
use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tracing::{info, warn};
use url::Url;

/// Capacity for the two mpsc channels between TapAudioPipe and TapClient.
/// Large enough that a slow brain tick doesn't drop on every cycle;
/// small enough that a runaway brain doesn't accumulate seconds of audio.
const TAP_MPSC_CAPACITY: usize = 32;

/// Connect timeout — bounded dial per spec §4.3 step 1 (2 s).
const CONNECT_TIMEOUT: Duration = Duration::from_secs(2);

/// The cold-path state the binary holds per session (looked up by ChannelId).
///
/// Holds ONLY the engine-control handles — the mpsc ends that drive
/// `TapAudioPipe` are owned by the `TapAudioPipe` itself (returned alongside
/// `TapConn` from `spawn_tap_engine`). This split keeps `TapConn` cheap
/// to store in a `DashMap<ChannelId, TapConn>` even though the actual
/// data flow happens through the `TapAudioPipe` wired into `RtcSession`.
pub struct TapConn {
    /// Close signal — fired on `Channel::Closing`. `oneshot::Sender` is the
    /// sending half; dropping it without `send(())` would also close the
    /// receiver (`Err(Closed)` on the await), but an explicit `send(())`
    /// is the documented teardown trigger (spec §5.1 step 5).
    pub close_tx: oneshot::Sender<()>,
    /// Task handle — abort on session teardown to guarantee the engine
    /// task stops even if the close signal is stuck in a `select!` arm.
    /// Per AGENTS.md: `abort()` is the safety net, `close_tx` is the gentle path.
    pub join: JoinHandle<()>,
    /// Shared metrics — also held by `TapAudioPipe` + `TapClient`. Storing
    /// here lets the binary read a `MetricsSnapshot` for log/CDR emission
    /// on teardown (spec §3.8 + the eventual CDR/ringbus emitter).
    pub metrics: Arc<TapMetrics>,
}

/// Spawn the TapEngine task for one session. Dials `tap_url`, runs the pump
/// loop, reconnects with bounded backoff on failure (spec §4.3, §5.2).
///
/// Returns BOTH the `TapAudioPipe` (the seam object to wire into
/// `RtcSession` via `RtcSession::with_pipe(pipe)`) AND the `TapConn`
/// (the engine-control handle to store in the binary's tap registry for
/// teardown).
///
/// # Why a tuple return (not a single struct holding both?)
///
/// The two values have different lifetimes + consumers: `TapAudioPipe`
/// is consumed by `RtcSession::with_pipe` immediately and lives behind
/// `RtcSession.pipe`; `TapConn` lives in the binary's tap registry and
/// outlives `TapAudioPipe` (it's the teardown handle). Coupling them in a
/// single struct would force the registry to also own the pipe, splitting
/// ownership of `RtcSession`'s internals across two modules — exactly the
/// pattern the slice-2 plan's structural review warned against.
pub fn spawn_tap_engine(session_id: ChannelId, tap_url: Url) -> (TapAudioPipe, TapConn) {
    // Two mpsc channels. The naming convention is "from the engine's POV":
    //  - `tx_pcm_in`/`rx_pcm_in`: peer PCM flowing INTO the engine (sink side
    //    of TapAudioPipe calls `tx_pcm_in.try_send(frame)`; engine task owns
    //    `rx_pcm_in` and feeds it to run_tap_client).
    //  - `tx_audio_out`/`rx_audio_out`: brain PCM flowing OUT of the engine
    //    (engine task owns `tx_audio_out` and feeds it to run_tap_client;
    //    TapAudioPipe holds `rx_audio_out` and drains it into the playout
    //    ring on each `next_pcm_frame` call).
    let (tx_pcm_in, rx_pcm_in) = mpsc::channel(TAP_MPSC_CAPACITY);
    let (tx_audio_out, rx_audio_out) = mpsc::channel(TAP_MPSC_CAPACITY);
    let (close_tx, close_rx) = oneshot::channel::<()>();
    let metrics = TapMetrics::new();

    // Clone metrics three ways: the engine task, the TapAudioPipe, and
    // the TapConn handle each hold their own `Arc<TapMetrics>` refcount.
    // All three views must observe the same counters — `Arc` is the
    // canonical Rust pattern for shared-ownership shared-mutation (the
    // counters live behind `AtomicU64` inside `TapMetrics`, so the
    // shared mutation is sound without a `Mutex`).
    let metrics_for_pipe = metrics.clone();
    let metrics_for_conn = metrics.clone();

    let join = tokio::spawn(async move {
        run_engine_loop(
            session_id,
            tap_url,
            rx_pcm_in,
            tx_audio_out,
            close_rx,
            metrics,
        )
        .await;
    });

    // The pipe owns `tx_pcm_in` (drains peer PCM, sends to engine via mpsc)
    // and `rx_audio_out` (drains engine's audio_out frames → playout ring).
    let pipe = TapAudioPipe::new(tx_pcm_in, rx_audio_out, metrics_for_pipe);

    let conn = TapConn {
        close_tx,
        join,
        metrics: metrics_for_conn,
    };

    (pipe, conn)
}

/// The engine's main loop: connect → handshake → pump, with bounded
/// exponential backoff (250 ms → 500 ms → 1 s → 2 s → cap at 5 s) on
/// any non-close failure. The `close` oneshot is shared across reconnect
/// attempts — one close signal per session, no fresh oneshot per dial.
///
/// # Deviation from the brief's verbatim sketch (Note A in plan)
///
/// The brief's `run_engine_loop` body had dead-code `let _ = pump_close;`
/// and convoluted `connect_async` future plumbing inside a `timeout(...)`
/// that wasn't actually awaited. The contract is what matters — connect,
/// run the pump, backoff, retry — so we use the cleaner multiply-select
/// pattern the plan's Note A spelled out: a single `tokio::select!` over
/// (a) the close receiver and (b) `connect_brain`, then a second select
/// over the close receiver and the backoff sleep.
async fn run_engine_loop(
    session_id: ChannelId,
    tap_url: Url,
    mut rx_pcm_in: mpsc::Receiver<rutster_media::PcmFrame>,
    tx_audio_out: mpsc::Sender<rutster_media::PcmFrame>,
    mut close: oneshot::Receiver<()>,
    metrics: Arc<TapMetrics>,
) {
    let mut backoff = Backoff::default();

    loop {
        // === Step 1: dial the brain, with the close signal racing the
        // connect future. `biased` makes close the priority arm — if the
        // binary tears the session down mid-dial we exit immediately
        // rather than waiting on connect_async to time out or succeed. ===
        let ws_result = tokio::select! {
            biased;
            _ = &mut close => {
                info!(%session_id, "close signal during connect; exiting");
                return;
            }
            // Boxed + pinned via `.map`/`boxed()` so the select arms share
            // a common Future output type (connect_async returns a different
            // concrete type than the close future; without boxing, select!
            // can't unify them on stable Rust).
            ws_result = connect_brain(&tap_url).boxed() => ws_result,
        };

        match ws_result {
            Ok(ws) => {
                // Successful connect → reset backoff + run the pump loop.
                backoff.reset();
                info!(%session_id, %tap_url, "brain WS connected; running pump");

                // === Step 2: run the pump loop until close/error. ===
                // `close` is a shared `&mut oneshot::Receiver` — same
                // signal across reconnect attempts (Task 4's design).
                let pump_result = run_tap_client(
                    ws,
                    session_id,
                    &mut rx_pcm_in,
                    tx_audio_out.clone(),
                    metrics.clone(),
                    &mut close,
                )
                .await;

                match pump_result {
                    Ok(()) => {
                        // Brain WS stream ended cleanly (e.g. brain sent bye).
                        // Per spec §4.3 step 4, we still backoff + retry:
                        // the brain may have restarted. Loop continues.
                        info!(%session_id, "tap client closed gracefully; backing off");
                    }
                    Err(TapClientError::Closed) => {
                        // Close signal from the binary (Channel::Closing).
                        // `run_tap_client` already sent `bye` and closed
                        // the WS; we're done — do NOT retry.
                        info!(%session_id, "tap client closed via core signal; exiting");
                        return;
                    }
                    Err(e) => {
                        // WS/protocol error → backoff + retry per spec §4.3.
                        warn!(error = ?e, %session_id, "tap client error; backing off");
                    }
                }
            }
            Err(e) => {
                warn!(error = ?e, %session_id, "tap connect failed; backing off");
            }
        }

        // === Step 3: backoff before retry (spec §4.3, §5.2). ===
        metrics.reconnect_attempts.fetch_add(1, Ordering::Relaxed);
        let delay = backoff.next_delay();
        tokio::select! {
            biased;
            _ = &mut close => {
                info!(%session_id, "close signal during backoff; exiting");
                return;
            }
            _ = tokio::time::sleep(delay) => {}
        }
    }
}

/// Dial the brain. Returns the live `WebSocketStream` (no TLS — slice-2 is
/// `ws://` loopback only per spec §4.4; `wss://` is rejected at
/// `POST /v1/sessions` time before we ever get here).
///
/// `connect_async` returns `WebSocketStream<MaybeTlsStream<TcpStream>>` for
/// `ws://` URLs (the `MaybeTlsStream` resolves to the no-TLS variant). The
/// concrete type is `tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>`.
/// We surface that explicit type rather than `impl Trait` because
/// `run_tap_client` is generic over `T: AsyncRead + AsyncWrite + Unpin`,
/// so a concrete type annotates + documents the wire shape for learners
/// without obscuring it behind `impl`.
///
/// # Why `IntoClientRequest` (not just `connect_async(&url)`)
///
/// `connect_async` accepts anything that implements `IntoClientRequest`.
/// `&Url` does, but going through `into_client_request()` first surfaces
/// any URL-shape error before the network round-trip — easier to attribute
/// (URL parse failure is a config error, not a network error).
async fn connect_brain(
    tap_url: &Url,
) -> Result<
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
    TapClientError,
> {
    let request = tap_url
        .as_str()
        .into_client_request()
        .map_err(TapClientError::Ws)?;
    // Bounded dial — `tokio::time::timeout` returns `Result<Ok(inner),
    // Elapsed>`; map the elapsed to a HandshakeTimeout-shaped Ws error so
    // the engine loop's backoff path sees it through the same `Err(e)`
    // arm as a true WS failure (one code path, two causes).
    let connect_fut = tokio_tungstenite::connect_async(request);
    let ws_result = tokio::time::timeout(CONNECT_TIMEOUT, connect_fut)
        .await
        .map_err(|_| TapClientError::Ws(tokio_tungstenite::tungstenite::Error::AlreadyClosed))?;
    let (ws, _response) = ws_result?;
    Ok(ws)
}

/// Bounded exponential backoff (spec §4.3, §5.2):
/// 250 ms → 500 ms → 1 s → 2 s → cap at 5 s. Infinite retries.
///
/// # Why not `std::time::Duration::saturating_mul`?
///
/// `Duration * 2` panics on overflow (Rust 1.x); `.min(cap)` keeps us
/// bounded. The explicit cap at 5 s means even after 1e6 retries the
/// delay stays at 5 s — important for the infinite-retry posture (spec
/// §5.2: "the brain WILL come back; the call SHOULD outlast the brain
/// outage").
struct Backoff {
    count: u32,
    current: Duration,
}

impl Default for Backoff {
    fn default() -> Self {
        Self {
            count: 0,
            current: Duration::from_millis(250),
        }
    }
}

impl Backoff {
    /// Return the current delay and advance the backoff state.
    ///
    /// Matches the spec §4.3 step 4 enumeration exactly:
    /// 250ms → 500ms → 1s → 2s → 5s cap (forever). The "cap at 5s" line
    /// means jump directly to 5s after the 2s step — NOT doubling to 4s
    /// first and then clamping (that would inject a 4s step the spec
    /// doesn't enumerate). The conditional below skips the intermediate
    /// 4s value that naïve `(current * 2).min(5s)` would produce.
    fn next_delay(&mut self) -> Duration {
        let d = self.current;
        // Below the 2s step: keep doubling (250 → 500 → 1s → 2s).
        // At or above 2s: jump straight to the 5s cap (no 4s step).
        self.current = if self.current < Duration::from_secs(2) {
            self.current * 2
        } else {
            Duration::from_secs(5)
        };
        self.count += 1;
        d
    }
    /// Reset after a successful connect + handshake.
    fn reset(&mut self) {
        self.count = 0;
        self.current = Duration::from_millis(250);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rutster_media::AudioSource;

    #[test]
    fn backoff_doubles_until_cap() {
        let mut b = Backoff::default();
        // Spec §4.3 step 4 + §5.2 enumeration:
        //   250ms → 500ms → 1s → 2s → 5s cap (no intermediate 4s step from
        //   naïve doubling). Infinite retries at the cap.
        assert_eq!(b.next_delay(), Duration::from_millis(250)); // step 1
        assert_eq!(b.next_delay(), Duration::from_millis(500)); // step 2
        assert_eq!(b.next_delay(), Duration::from_secs(1)); // step 3
        assert_eq!(b.next_delay(), Duration::from_secs(2)); // step 4
        assert_eq!(b.next_delay(), Duration::from_secs(5)); // step 5: cap
                                                            // Cap invariant: stays at 5s forever (infinite retries, per spec §5.2).
        assert_eq!(b.next_delay(), Duration::from_secs(5)); // step 6
        assert_eq!(b.next_delay(), Duration::from_secs(5)); // step 7
        assert_eq!(b.next_delay(), Duration::from_secs(5)); // step 8+
    }

    #[test]
    fn backoff_reset_returns_to_floor() {
        let mut b = Backoff::default();
        let _ = b.next_delay();
        let _ = b.next_delay();
        b.reset();
        assert_eq!(b.next_delay(), Duration::from_millis(250));
    }

    #[tokio::test]
    async fn spawn_returns_pipe_with_live_mpsc_ends() {
        // Smoke test: spawning the engine immediately returns a TapAudioPipe
        // the binary can wire into RtcSession. We don't actually wire a brain
        // here — the engine task will spin connect_brain (failing on a
        // non-existent server) and back off; the metrics counter should
        // increment. We abort the task on test drop to avoid leak.
        let id = ChannelId::new();
        let url = Url::parse("ws://127.0.0.1:1/echo").unwrap(); // port 1 = unreachable
        let (mut pipe, conn) = spawn_tap_engine(id, url);
        // TapAudioPipe is the seam object — should default to silent underflow.
        assert!(pipe.next_pcm_frame().is_none());
        // TapConn carries the close oneshot + JoinHandle + metrics.
        let _ = conn.close_tx.send(());
        conn.join.abort();
    }
}