wl-webrtc/src/state_portal/threads.rs

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::avhw::{BitrateCommand, CpuNv12Frame, EncodeOutcome, EncodedH264Frame, SwEncEncode};
use crate::webrtc::WebRtcState;

use super::bitrate::{next_upscale_tier, resolution_bitrate_bps, select_resolution};

pub(super) struct EncodeThreadTiming {
    pub(super) sws_us: u64,
    pub(super) encode_us: u64,
    pub(super) output_bytes: usize,
}

pub(super) struct EncodeThread {
    pub(super) handle: Option<std::thread::JoinHandle<()>>,
    pub(super) input_tx: crossbeam_channel::Sender<CpuNv12Frame>,
    pub(super) timing_rx: crossbeam_channel::Receiver<EncodeThreadTiming>,
    pub(super) duplicate_count: std::sync::Arc<std::sync::atomic::AtomicU64>,
}

pub(super) struct WebrtcThread {
    pub(super) handle: Option<std::thread::JoinHandle<()>>,
    pub(super) sent_gap_rx: crossbeam_channel::Receiver<(f64, Option<f64>)>,
}

/// Static configuration handed to the WebRTC sender thread. Immutable for the
/// thread's lifetime; a resolution tier change rebuilds the whole pipeline
/// (and spawns a new thread) rather than mutating this.
pub(super) struct WebRtcThreadConfig {
    pub(super) fps: u32,
    pub(super) enc_width: u32,
    pub(super) enc_height: u32,
    pub(super) max_bitrate: u64,
}

/// Channel endpoints owned exclusively by the WebRTC sender thread after spawn.
/// The reverse endpoints stay with StatePortal (or the encode thread) for
/// inbound/outbound traffic.
pub(super) struct WebRtcThreadChannels {
    pub(super) webrtc_rx: crossbeam_channel::Receiver<EncodedH264Frame>,
    pub(super) sent_gap_tx: crossbeam_channel::Sender<(f64, Option<f64>)>,
    pub(super) bitrate_tx: crossbeam_channel::Sender<BitrateCommand>,
    pub(super) resolution_tx: crossbeam_channel::Sender<BitrateCommand>,
}

pub(super) fn encode_thread_loop(
    mut encode: SwEncEncode,
    input_rx: crossbeam_channel::Receiver<CpuNv12Frame>,
    timing_tx: crossbeam_channel::Sender<EncodeThreadTiming>,
    duplicate_count: std::sync::Arc<std::sync::atomic::AtomicU64>,
) {
    loop {
        match input_rx.recv() {
            Ok(frame) => {
                match encode.encode_cpu_frame(&frame) {
                    Ok(EncodeOutcome::Encoded) => {
                        let t = encode.take_timing();
                        let _ = timing_tx.try_send(EncodeThreadTiming {
                            sws_us: t.sws_us,
                            encode_us: t.encode_us,
                            output_bytes: t.output_bytes,
                        });
                    }
                    Ok(EncodeOutcome::SkippedDuplicate) => {
                        duplicate_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    }
                    Ok(_) => {
                        // SkippedPaused / SkippedDisconnected — no counter needed
                    }
                    Err(e) => {
                        tracing::error!("Encode thread error: {e}");
                        break;
                    }
                }
            }
            Err(_) => {
                tracing::info!("Encode thread input closed, flushing encoder");
                if let Err(e) = encode.flush() {
                    tracing::error!("Encode thread flush error: {e}");
                }
                break;
            }
        }
    }
    tracing::info!("Encode thread exiting");
}

pub(super) fn webrtc_thread_loop(
    mut wrtc: WebRtcState,
    config: WebRtcThreadConfig,
    channels: WebRtcThreadChannels,
    paused: Arc<AtomicBool>,
) {
    let WebRtcThreadConfig {
        fps,
        enc_width,
        enc_height,
        max_bitrate,
    } = config;
    let WebRtcThreadChannels {
        webrtc_rx,
        sent_gap_tx,
        bitrate_tx,
        resolution_tx,
    } = channels;
    let mut frames_sent: u64 = 0;
    let mut last_send: Option<std::time::Instant> = None;
    let mut last_sent_bitrate: Option<u64> = None;
    let initial_tier = (enc_width, enc_height);
    let mut current_tier = initial_tier;
    let mut upscale_counter = 0u32;
    let mut last_resolution_eval = Instant::now();
    let timeout = Duration::from_millis(1);

    loop {
        if let Err(e) = wrtc.handle_signaling() {
            tracing::error!("WebRTC signaling error: {e}");
            break;
        }
        if let Err(e) = wrtc.poll_and_feed() {
            tracing::error!("WebRTC poll error: {e}");
            break;
        }

        if wrtc.take_force_keyframe() {
            let _ = bitrate_tx.try_send(BitrateCommand::ForceKeyframe);
        }

        let connected = wrtc.is_connected();
        let was_paused = paused.load(Ordering::Relaxed);
        let now_paused = !connected;
        if was_paused && !now_paused {
            tracing::info!("WebRTC client connected, resuming encoding");
        } else if !was_paused && now_paused {
            tracing::warn!("WebRTC client disconnected, pausing encoding");
        }
        paused.store(now_paused, Ordering::Relaxed);

        if let Some(bwe) = wrtc.get_bwe_estimate() {
            // #23: Cap BWE to prevent runaway bitrate escalation. Without this, BWE
            // estimates can rise to 10+ Mbps, causing IDR bursts and PLI storms.
            let effective_bwe = bwe.min(max_bitrate);
            if effective_bwe != bwe {
                tracing::debug!(
                    bwe,
                    effective_bwe,
                    max_bitrate,
                    "BWE exceeds --max-bitrate cap, clamping"
                );
            }
            let bwe = effective_bwe;

            let should_send = match last_sent_bitrate {
                None => true,
                Some(last) => {
                    let diff = bwe.abs_diff(last);
                    diff * 10 > last
                }
            };
            if should_send {
                let _ = bitrate_tx.try_send(BitrateCommand::UpdateBitrate { target_bps: bwe });
                last_sent_bitrate = Some(bwe);
            }

            if last_resolution_eval.elapsed() >= Duration::from_secs(1) {
                last_resolution_eval = Instant::now();
                let selected = select_resolution(current_tier.0, current_tier.1, bwe, fps);
                if selected != current_tier {
                    current_tier = selected;
                    upscale_counter = 0;
                    let _ = resolution_tx.try_send(BitrateCommand::UpdateResolution {
                        width: current_tier.0,
                        height: current_tier.1,
                    });
                    wrtc.set_need_keyframe();
                } else if let Some(next_tier) = next_upscale_tier(current_tier, initial_tier) {
                    let needed = resolution_bitrate_bps(next_tier.0, next_tier.1, fps);
                    if bwe > needed.saturating_mul(120) / 100 {
                        upscale_counter = upscale_counter.saturating_add(1);
                        if upscale_counter >= 10 {
                            current_tier = next_tier;
                            upscale_counter = 0;
                            let _ = resolution_tx.try_send(BitrateCommand::UpdateResolution {
                                width: current_tier.0,
                                height: current_tier.1,
                            });
                            wrtc.set_need_keyframe();
                        }
                    } else {
                        upscale_counter = 0;
                    }
                } else {
                    upscale_counter = 0;
                }
            }
        }

        if connected {
            while let Ok(enc_frame) = webrtc_rx.try_recv() {
                if let Err(e) = wrtc.write_h264_frame(&enc_frame.data, enc_frame.pts_ticks) {
                    tracing::debug!("WebRTC write frame error: {e}");
                }
                frames_sent = frames_sent.saturating_add(1);
                let gap_ms = last_send
                    .map(|l| l.elapsed().as_secs_f64() * 1000.0)
                    .unwrap_or(0.0);
                // Compute capture-to-send age on the sending thread so the
                // frame_age stat stays accurate when batch-drained later.
                let age_ms = Some(enc_frame.capture_time.elapsed().as_secs_f64() * 1000.0);
                last_send = Some(std::time::Instant::now());
                let _ = sent_gap_tx.try_send((gap_ms, age_ms));
            }
        } else {
            while webrtc_rx.try_recv().is_ok() {}
        }

        match webrtc_rx.recv_timeout(timeout) {
            Ok(enc_frame) => {
                if wrtc.is_connected() {
                    if let Err(e) = wrtc.write_h264_frame(&enc_frame.data, enc_frame.pts_ticks) {
                        tracing::debug!("WebRTC write frame error: {e}");
                    }
                    frames_sent = frames_sent.saturating_add(1);
                    let gap_ms = last_send
                        .map(|l| l.elapsed().as_secs_f64() * 1000.0)
                        .unwrap_or(0.0);
                    let age_ms = Some(enc_frame.capture_time.elapsed().as_secs_f64() * 1000.0);
                    last_send = Some(std::time::Instant::now());
                    let _ = sent_gap_tx.try_send((gap_ms, age_ms));
                }
            }
            Err(crossbeam_channel::RecvTimeoutError::Timeout) => {}
            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
                tracing::info!("WebRTC channel disconnected, exiting thread");
                return;
            }
        }
    }

    tracing::info!("WebRTC thread exiting");
}

#[cfg(test)]
mod tests {

    #[test]
    fn try_send_full_channel_returns_full_not_block() {
        let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
        tx.send(vec![1]).unwrap();
        tx.send(vec![2]).unwrap();

        assert!(matches!(
            tx.try_send(vec![3]),
            Err(crossbeam_channel::TrySendError::Full(_))
        ));
        assert_eq!(rx.len(), 2);
    }

    #[test]
    fn try_send_after_rx_dropped_returns_disconnected() {
        let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
        drop(rx);

        assert!(matches!(
            tx.try_send(vec![1]),
            Err(crossbeam_channel::TrySendError::Disconnected(_))
        ));
    }

    // given: full bounded channel
    // when:  rx is dropped, then try_send
    // expect: Disconnected, not blocking
    #[test]
    fn shutdown_rx_drop_prevents_deadlock_on_full_channel() {
        let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
        tx.send(vec![1]).unwrap();
        tx.send(vec![2]).unwrap();
        drop(rx);

        assert!(matches!(
            tx.try_send(vec![3]),
            Err(crossbeam_channel::TrySendError::Disconnected(_))
        ));
    }
}