refactor(state_portal): extract bitrate helpers + thread loops to submodules

Step 3: split state_portal.rs (1241 -> 829 LOC) into three modules.

- src/state_portal.rs (829 LOC): keeps StatePortal struct + impl (with
  poll_and_encode / handle_pw_frame / shutdown / etc.) + Drop + PortalStage
  enum + DRM helpers + DRM tests. Per Oracle/Explore audit, all 21
  StatePortal fields are private and poll_and_encode interleaves three
  channel reads with state-machine transitions; moving it would force
  pub(crate) on every field, so it stays in mod.rs.
- src/state_portal/bitrate.rs (144 LOC): RESOLUTION_TIERS + 4 pure fns
  (resolution_bitrate_bps / webrtc_startup_bitrate_bps / select_resolution /
  next_upscale_tier) + 10 tests that exercise them. Pure fns with no
  StatePortal field access — the cleanest possible extract.
- src/state_portal/threads.rs (287 LOC): the 5 thread-related types
  (EncodeThreadTiming / EncodeThread / WebrtcThread / WebRtcThreadConfig /
  WebRtcThreadChannels) + the two free fns encode_thread_loop /
  webrtc_thread_loop + the 3 channel-semantics regression tests
  (try_send_* / shutdown_rx_drop_*) that document crossbeam invariants
  the shutdown logic relies on. Struct fields widened to pub(super) so
  StatePortal in mod.rs can construct and join them.

Test preservation:
- state_portal test count: 17 (mod.rs=4 drm tests + bitrate.rs=10 +
  threads.rs=3 channel tests) — matches baseline.

Verification (all green):
- cargo build / cargo build --release
- cargo test (79 lib + 3 integration = 82 pass, 1 ignored — unchanged)
- cargo clippy --all-targets -- -D warnings
- cargo fmt --check
This commit is contained in:
2026-07-13 16:35:44 +08:00
parent 60d6e7f046
commit bcfbd93f5a
3 changed files with 441 additions and 414 deletions
+287
View File
@@ -0,0 +1,287 @@
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(_))
));
}
}