fix(state_portal): make shutdown() idempotent to prevent duplicate log lines (#22)

shutdown() fired twice on exit (explicit call in main.rs + Drop impl), causing "Total: N frames" and "StatePortal shutdown complete" to log twice 23us apart. Add `shutdown_started: bool` guard at function entry.

Plain bool (not AtomicBool) because &mut self already grants exclusive access. Guard is set BEFORE cleanup so Drop re-entry during panic unwinding is suppressed.

Includes scripts/test_shutdown_idempotency.sh as a live regression test (requires Wayland session). Verified PASS on KWin: both lines print exactly once.

scripts/test_shutdown_idempotency.sh

@@ -0,0 +1,150 @@
+#!/usr/bin/env bash
+# Regression test for issue #22: "Total frames" log printed twice on shutdown.
+#
+# Runs the wl-webrtc binary briefly, sends SIGINT, then asserts that
+#   - "Total: N frames in ..." appears exactly once
+#   - "StatePortal shutdown complete" appears exactly once
+#
+# Pre-fix:    both lines printed twice (explicit shutdown + Drop re-entry).
+# Post-fix:   both lines printed once (shutdown_started guard).
+#
+# Usage:
+#   ./scripts/test_shutdown_idempotency.sh              # WebRTC mode, build first
+#   ./scripts/test_shutdown_idempotency.sh --mode file  # --output mode instead
+#   ./scripts/test_shutdown_idempotency.sh --skip-build # skip cargo build --release
+#   ./scripts/test_shutdown_idempotency.sh --signal TERM # use SIGTERM instead of SIGINT
+#
+# Requires: a Wayland session (WAYLAND_DISPLAY). The script will warn but
+# proceed if unset; capture will simply fail and the test will report FAIL.
+
+set -euo pipefail
+
+MODE="webrtc"
+SKIP_BUILD=0
+SIGNAL="INT"
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --mode)        MODE="$2"; shift 2 ;;
+        --mode=*)      MODE="${1#*=}"; shift ;;
+        --skip-build)  SKIP_BUILD=1; shift ;;
+        --signal)      SIGNAL="$2"; shift 2 ;;
+        --signal=*)    SIGNAL="${1#*=}"; shift ;;
+        -h|--help)
+            sed -n '2,18p' "$0"; exit 0 ;;
+        *) echo "Unknown arg: $1" >&2; exit 2 ;;
+    esac
+done
+
+if [[ -t 1 ]]; then
+    GREEN=$'\e[32m'; RED=$'\e[31m'; YELLOW=$'\e[33m'; BOLD=$'\e[1m'; RESET=$'\e[0m'
+else
+    GREEN=""; RED=""; YELLOW=""; BOLD=""; RESET=""
+fi
+
+REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+cd "$REPO_ROOT"
+
+if [[ -z "${WAYLAND_DISPLAY:-}" ]]; then
+    echo "${YELLOW}WARNING${RESET}: WAYLAND_DISPLAY not set; live capture likely to fail." >&2
+fi
+
+if [[ $SKIP_BUILD -eq 0 ]]; then
+    echo "${BOLD}Building release binary...${RESET}"
+    cargo build --release
+fi
+
+BIN="$REPO_ROOT/target/release/wl-webrtc"
+if [[ ! -x "$BIN" ]]; then
+    echo "${RED}FAIL${RESET}: $BIN not found. Run without --skip-build first." >&2
+    exit 1
+fi
+
+case "$MODE" in
+    webrtc)
+        PORT=56666
+        RUN_ARGS=(--port "$PORT" -v)
+        EXTRA_CLEANUP=()
+        ;;
+    file)
+        OUTPUT_FILE="$(mktemp --tmpdir "wl22-test-XXXXXX.mp4")"
+        RUN_ARGS=(--output "$OUTPUT_FILE" -v)
+        EXTRA_CLEANUP=("rm -f "$OUTPUT_FILE"")
+        ;;
+    *)
+        echo "Invalid --mode: $MODE (use 'webrtc' or 'file')" >&2
+        exit 2
+        ;;
+esac
+
+LOG="$(mktemp --tmpdir "wl22-test-XXXXXX.log")"
+SERVER_PID=""
+
+cleanup() {
+    if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
+        kill "$SERVER_PID" 2>/dev/null || true
+        wait "$SERVER_PID" 2>/dev/null || true
+    fi
+    rm -f "$LOG"
+    for cmd in "${EXTRA_CLEANUP[@]}"; do eval "$cmd" 2>/dev/null || true; done
+}
+trap cleanup EXIT
+
+echo "${BOLD}Running${RESET}: $BIN ${RUN_ARGS[*]}"
+"$BIN" "${RUN_ARGS[@]}" >"$LOG" 2>&1 &
+SERVER_PID=$!
+
+# Give the server time to initialize, capture at least one frame, and stabilize.
+sleep 3
+
+if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+    echo "${RED}FAIL${RESET}: server exited before SIGINT could be sent." >&2
+    echo "----- Log -----" >&2
+    cat "$LOG" >&2
+    exit 1
+fi
+
+echo "${BOLD}Sending SIG${SIGNAL}...${RESET}"
+kill -"$SIGNAL" "$SERVER_PID"
+
+# Wait up to 3s for graceful exit.
+for _ in $(seq 1 30); do
+    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+        break
+    fi
+    sleep 0.1
+done
+
+if kill -0 "$SERVER_PID" 2>/dev/null; then
+    echo "${YELLOW}WARN${RESET}: process still alive 3s after SIG${SIGNAL}; force-killing"
+    kill -TERM "$SERVER_PID" 2>/dev/null || true
+    wait "$SERVER_PID" 2>/dev/null || true
+fi
+
+SERVER_PID=""
+
+# grep -c exits 1 on zero matches, which would abort under `set -e`; swallow it.
+TOTAL_COUNT=$(grep -c 'Total:.*frames in' "$LOG" || true)
+COMPLETE_COUNT=$(grep -c 'StatePortal shutdown complete' "$LOG" || true)
+
+echo
+echo "${BOLD}----- Last 8 log lines -----${RESET}"
+tail -n 8 "$LOG"
+echo "${BOLD}------------------------------${RESET}"
+echo
+echo "\"Total: N frames in ...\":        $TOTAL_COUNT occurrence(s) (expected 1)"
+echo "\"StatePortal shutdown complete\": $COMPLETE_COUNT occurrence(s) (expected 1)"
+
+if [[ "$TOTAL_COUNT" -eq 1 && "$COMPLETE_COUNT" -eq 1 ]]; then
+    echo
+    echo "${GREEN}${BOLD}PASS${RESET}: shutdown is idempotent (issue #22 fixed)"
+    exit 0
+fi
+
+echo
+echo "${RED}${BOLD}FAIL${RESET}: shutdown logged multiple times (issue #22 not fixed)"
+[[ "$TOTAL_COUNT" -ge 2 ]]    && echo "  - \"Total:\" printed $TOTAL_COUNT times"
+[[ "$COMPLETE_COUNT" -ge 2 ]] && echo "  - \"shutdown complete\" printed $COMPLETE_COUNT times"
+echo
+echo "Debug hint: to compare before/after the fix, run:"
+echo "  git stash && cargo build --release && $0 --skip-build && git stash pop"
+exit 1

src/avhw.rs

@@ -7,6 +7,7 @@ use std::ptr;
 use std::slice;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
+use std::time::Instant;
 
 use anyhow::{bail, Result};
 use ffmpeg_next as ff;
@@ -25,6 +26,9 @@ use crate::transform::{transpose_if_transform_transposed, Transform};
 pub enum BitrateCommand {
     UpdateBitrate { target_bps: u64 },
     UpdateResolution { width: u32, height: u32 },
+    /// Force the next encoded frame to be an IDR. Sent by the WebRTC thread
+    /// in response to str0m `Event::KeyframeRequest` or a resolution change.
+    ForceKeyframe,
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -33,6 +37,17 @@ pub struct ResolutionChange {
     pub height: u32,
 }
 
+/// Per-frame timing snapshot for the software encoder, consumed by the stats
+/// thread. `sws_us` measures NV12→YUV420P conversion, `encode_us` measures
+/// `avcodec_send_frame` + drain, and `output_bytes` counts encoded bytes
+/// produced by libavcodec (even if downstream delivery later drops them).
+#[derive(Default, Clone, Copy, Debug)]
+pub struct SwEncodeTiming {
+    pub sws_us: u64,
+    pub encode_us: u64,
+    pub output_bytes: usize,
+}
+
 // ---------------------------------------------------------------------------
 // AvHwDevCtx
 // ---------------------------------------------------------------------------
@@ -856,6 +871,7 @@ impl SwEncImport {
                     requested = Some((width & !1, height & !1));
                 }
                 BitrateCommand::UpdateBitrate { .. } => {}
+                BitrateCommand::ForceKeyframe => {}
             }
         }
 
@@ -966,6 +982,14 @@ pub struct SwEncEncode {
     fps: u32,
     bitrate: u64,
     gop_size: u32,
+    /// Set true when WebRTC requests a keyframe. Forces the next frame to
+    /// `AV_PICTURE_TYPE_I` and bypasses the dedup hash check. Cleared only
+    /// after `avcodec_send_frame` accepts the forced frame.
+    force_keyframe_pending: bool,
+    /// Last per-frame timing snapshot. Reset to `Default` at the start of
+    /// every `encode_cpu_frame` call (even on early returns) so stale values
+    /// from a previous frame can never leak out.
+    last_timing: SwEncodeTiming,
 }
 
 const FNV1A_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
@@ -1028,6 +1052,8 @@ impl SwEncEncode {
             fps,
             bitrate,
             gop_size,
+            force_keyframe_pending: false,
+            last_timing: SwEncodeTiming::default(),
         })
     }
 
@@ -1066,6 +1092,8 @@ impl SwEncEncode {
             fps,
             bitrate,
             gop_size,
+            force_keyframe_pending: false,
+            last_timing: SwEncodeTiming::default(),
         })
     }
 
@@ -1079,12 +1107,18 @@ impl SwEncEncode {
             }
         }
         let start_ts = self.starting_timestamp.unwrap_or(0);
-        self.drain_encoder(start_ts)?;
+        let _ = self.drain_encoder(start_ts)?;
 
         Ok(())
     }
 
+    pub fn take_timing(&mut self) -> SwEncodeTiming {
+        mem::take(&mut self.last_timing)
+    }
+
     pub fn encode_cpu_frame(&mut self, frame: &CpuNv12Frame) -> Result<()> {
+        self.last_timing = SwEncodeTiming::default();
+
         if self.webrtc_disconnected {
             return Ok(());
         }
@@ -1103,9 +1137,15 @@ impl SwEncEncode {
                     }
                 }
                 BitrateCommand::UpdateResolution { .. } => {}
+                BitrateCommand::ForceKeyframe => {
+                    self.force_keyframe_pending = true;
+                    tracing::debug!("encode thread: ForceKeyframe requested");
+                }
             }
         }
 
+        let force_this_frame = self.force_keyframe_pending;
+
         while let Ok(change) = self.resolution_rx.try_recv() {
             self.recreate_encoder(change.width, change.height)?;
         }
@@ -1130,13 +1170,14 @@ impl SwEncEncode {
         self.frame_count = self.frame_count.saturating_add(1);
         let current_hash = hash_sampled_y_plane(&frame.y_data, width, height, frame.y_stride);
         let force_gop_frame = self.gop_size > 0 && frame_index % u64::from(self.gop_size) == 0;
-        if frame_index > 0 && !force_gop_frame && current_hash == self.last_frame_hash {
+        if frame_index > 0 && !force_gop_frame && !force_this_frame && current_hash == self.last_frame_hash {
             tracing::debug!(frame_index, "skipping duplicate frame");
             self.last_frame_hash = current_hash;
             return Ok(());
         }
         self.last_frame_hash = current_hash;
 
+        let sws_start = Instant::now();
         // SAFETY: yuv_frame is an owned reusable YUV420P frame at the same dimensions as sw_nv12;
         // sws_ctx was created for NV12 -> YUV420P with no resize, so sws_scale only converts format.
         unsafe {
@@ -1164,6 +1205,7 @@ impl SwEncEncode {
                 bail!("sws_scale failed for software encoder: {scaled}");
             }
         }
+        let sws_us = sws_start.elapsed().as_micros() as u64;
 
         let pts = frame.pts;
         if self.starting_timestamp.is_none() {
@@ -1171,9 +1213,18 @@ impl SwEncEncode {
         }
         let start_ts = self.starting_timestamp.unwrap_or(0);
 
+        let enc_start = Instant::now();
         // SAFETY: yuv_frame is initialized, writable, and matches the opened encoder format.
+        // pict_type is reset every frame: the AVFrame is reused, so without resetting to NONE
+        // a previously-forced I-type would leak into subsequent P-frames. With forced-idr=1
+        // set on the encoder, AV_PICTURE_TYPE_I produces a true IDR NALU.
         unsafe {
             (*self.yuv_frame).pts = pts;
+            (*self.yuv_frame).pict_type = if force_this_frame {
+                ffi::AVPictureType::AV_PICTURE_TYPE_I
+            } else {
+                ffi::AVPictureType::AV_PICTURE_TYPE_NONE
+            };
             let ret = ffi::avcodec_send_frame(self.enc_video.as_mut_ptr(), self.yuv_frame);
             if ret < 0 {
                 bail!(
@@ -1183,7 +1234,20 @@ impl SwEncEncode {
             }
         }
 
-        self.drain_encoder(start_ts)
+        if force_this_frame {
+            self.force_keyframe_pending = false;
+        }
+
+        let output_bytes = self.drain_encoder(start_ts)?;
+        let encode_us = enc_start.elapsed().as_micros() as u64;
+
+        self.last_timing = SwEncodeTiming {
+            sws_us,
+            encode_us,
+            output_bytes,
+        };
+
+        Ok(())
     }
 
     fn recreate_encoder(&mut self, width: u32, height: u32) -> Result<()> {
@@ -1215,6 +1279,7 @@ impl SwEncEncode {
         self.enc_height = height;
         self.last_frame_hash = 0;
         self.frame_count = 0;
+        self.force_keyframe_pending = true;
 
         Ok(())
     }
@@ -1229,7 +1294,8 @@ impl SwEncEncode {
         Ok(())
     }
 
-    fn drain_encoder(&mut self, start_ts: i64) -> Result<()> {
+    fn drain_encoder(&mut self, start_ts: i64) -> Result<usize> {
+        let mut total_bytes = 0usize;
         loop {
             let mut pkt = ff::Packet::empty();
             // SAFETY: enc_video is an open encoder; pkt is writable packet storage.
@@ -1243,6 +1309,15 @@ impl SwEncEncode {
                 bail!("avcodec_receive_packet failed: {}", ff_err(ret));
             }
 
+            // Count encoded bytes produced before the Muxer/Channel match to
+            // avoid branch duplication and handle multi-packet drain correctly.
+            // SAFETY: pkt was just filled by a successful avcodec_receive_packet;
+            // the size field is valid and initialized.
+            let pkt_size = unsafe { (*pkt.as_mut_ptr()).size };
+            if pkt_size > 0 {
+                total_bytes += pkt_size as usize;
+            }
+
             match self.output {
                 Some(FrameOutput::Muxer(ref mut octx)) => {
                     let enc_tb = self.enc_video.time_base();
@@ -1305,7 +1380,7 @@ impl SwEncEncode {
                 None => {}
             }
         }
-        Ok(())
+        Ok(total_bytes)
     }
 }
 
@@ -1741,6 +1816,14 @@ fn create_software_h264_encoder(
             // avcodec_alloc_context3. Setting level is a simple i32 field
             // assignment on a properly aligned struct.
             (*enc.as_mut_ptr()).level = 42; // H.264 Level 4.2 (up to 1440p@30)
+            // SAFETY: priv_data belongs to the unopened libx264 encoder context.
+            // `forced-idr` is an FFmpeg-level private option (not x264-native),
+            // so it must be set via av_opt_set, NOT via the x264opts string.
+            // With forced-idr=1, setting AV_PICTURE_TYPE_I on an input frame
+            // produces a true IDR NALU with inline SPS/PPS (repeat_headers=1).
+            let key = CString::new("forced-idr").unwrap();
+            let val = CString::new("1").unwrap();
+            ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0);
             let key = CString::new("x264opts").unwrap();
             let vbv_maxrate = bitrate;
             let vbv_bufsize = bitrate / 4;

src/state_portal.rs

@@ -66,6 +66,7 @@ pub struct StatePortal {
     last_fillable_frame: Option<CpuNv12Frame>, // cached last frame for filler duplication
     next_filler_at: Option<Instant>,       // when to send next filler frame
     filler_frames_sent: u64,
+    shutdown_started: bool, // idempotency guard; plain bool because &mut self is exclusive (not AtomicBool)
 }
 
 impl StatePortal {
@@ -112,6 +113,7 @@ impl StatePortal {
             last_fillable_frame: None,
             next_filler_at: None,
             filler_frames_sent: 0,
+            shutdown_started: false,
         })
     }
 
@@ -586,10 +588,15 @@ impl StatePortal {
         Ok(())
     }
 
-    /// 关闭状态：刷新编码器并清理资源
+    /// 关闭状态：刷新编码器并清理资源（幂等）。
     ///
-    /// 使用 `enc.take()` 确保编码器只被 flush 一次，即使多次调用也安全（幂等）。
+    /// `shutdown_started` 守卫在清理之前置位——防止 panic 时 `Drop` 重入 unwinding。
     pub fn shutdown(&mut self) {
+        if self.shutdown_started {
+            return;
+        }
+        self.shutdown_started = true;
+
         self.last_fillable_frame = None;
         // 1. Stop encode thread (drops webrtc_tx → signals WebRTC thread to exit)
         if let Some(mut enc_thread) = self.enc_thread.take() {
@@ -644,14 +651,13 @@ fn encode_thread_loop(
     loop {
         match input_rx.recv() {
             Ok(frame) => {
-                let t_start = Instant::now();
                 match encode.encode_cpu_frame(&frame) {
                     Ok(()) => {
-                        let elapsed = t_start.elapsed().as_micros() as u64;
+                        let t = encode.take_timing();
                         let _ = timing_tx.try_send(EncodeThreadTiming {
-                            sws_us: 0,
-                            encode_us: elapsed,
-                            output_bytes: 0,
+                            sws_us: t.sws_us,
+                            encode_us: t.encode_us,
+                            output_bytes: t.output_bytes,
                         });
                     }
                     Err(e) => {
@@ -702,6 +708,10 @@ fn webrtc_thread_loop(
             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;

src/webrtc.rs

@@ -200,6 +200,7 @@ struct WebRtcInner {
     video_pt: Option<Pt>,
     connected: bool,
     need_keyframe: bool,
+    force_keyframe_to_encode: bool,
     current_bwe_estimate: Option<Bitrate>,
     rtp_clock: u32,
     buf: Vec<u8>,
@@ -353,6 +354,17 @@ impl WebRtcState {
     pub fn set_need_keyframe(&mut self) {
         if let Some(inner) = self.inner.as_mut() {
             inner.need_keyframe = true;
+            inner.force_keyframe_to_encode = true;
+        }
+    }
+
+    pub fn take_force_keyframe(&mut self) -> bool {
+        if let Some(inner) = self.inner.as_mut() {
+            let v = inner.force_keyframe_to_encode;
+            inner.force_keyframe_to_encode = false;
+            v
+        } else {
+            false
         }
     }
 }
@@ -429,6 +441,7 @@ impl WebRtcInner {
             video_pt: None,
             connected: false,
             need_keyframe: false,
+            force_keyframe_to_encode: false,
             current_bwe_estimate: None,
             rtp_clock: 0,
             buf: vec![0u8; 65535],
@@ -446,6 +459,7 @@ impl WebRtcInner {
             .map_err(|e| anyhow::anyhow!("accept_offer: {e}"))?;
 
         self.need_keyframe = true;
+        self.force_keyframe_to_encode = true;
         tracing::info!("SDP exchange complete, waiting for ICE/DTLS...");
 
         self.discover_video_params();
@@ -503,6 +517,7 @@ impl WebRtcInner {
                             tracing::info!("WebRTC connected!");
                             self.connected = true;
                             self.need_keyframe = true;
+                            self.force_keyframe_to_encode = true;
                             self.discover_video_params();
                         }
                         Event::IceConnectionStateChange(IceConnectionState::Disconnected) => {
@@ -525,6 +540,7 @@ impl WebRtcInner {
                         Event::KeyframeRequest(_) => {
                             tracing::info!("received keyframe request from viewer");
                             self.need_keyframe = true;
+                            self.force_keyframe_to_encode = true;
                         }
                         Event::EgressBitrateEstimate(est) => {
                             let bitrate = match est {