Skip to main content

arcly_stream/bus/
handle.rs

1//! The live stream handle — a lock-free, zero-copy broadcast fan-out bus.
2
3use crate::observe::{NoopObserver, Observer};
4use crate::{frame::FrameFlags, AppName, MediaFrame, Result, StreamId, StreamKey};
5use arc_swap::{ArcSwap, ArcSwapOption};
6use std::net::SocketAddr;
7use std::ops::ControlFlow;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::sync::{Arc, Mutex as StdMutex, OnceLock};
10use std::time::{Instant, SystemTime, UNIX_EPOCH};
11use tokio::sync::broadcast::error::RecvError;
12use tokio::sync::{broadcast, Notify, RwLock};
13use tokio_util::sync::CancellationToken;
14
15/// A per-frame egress sink driven by [`StreamHandle::drive_to`].
16///
17/// Implementors send one frame to a single downstream peer (an RTSP player, an
18/// upstream RTMP server, …). Returning [`ControlFlow::Break`] stops the drive
19/// cleanly — typically because the peer disconnected — while `Err` aborts it.
20#[async_trait::async_trait]
21pub trait FrameSink: Send {
22    /// Forward one frame; `Break` ends the session gracefully.
23    async fn send(&mut self, frame: Arc<MediaFrame>) -> Result<ControlFlow<()>>;
24}
25
26/// Current wall-clock time in Unix milliseconds (saturating to 0 pre-epoch).
27///
28/// Use this only for *displayed* timestamps (e.g. `started_at_ms`); it is
29/// subject to NTP steps and operator clock changes, so it must never drive
30/// elapsed-time decisions. For those, use [`mono_ms`].
31pub(crate) fn now_ms() -> u64 {
32    SystemTime::now()
33        .duration_since(UNIX_EPOCH)
34        .map(|d| d.as_millis() as u64)
35        .unwrap_or(0)
36}
37
38/// Milliseconds elapsed on a process-local **monotonic** clock since the first
39/// call. Unlike [`now_ms`], this never jumps backward (or forward) when the wall
40/// clock is adjusted, so it is the correct basis for QoS windows and the idle
41/// reaper: a leap-second or NTP correction can't spuriously reap a live stream
42/// or distort a measured bitrate.
43pub(crate) fn mono_ms() -> u64 {
44    static EPOCH: OnceLock<Instant> = OnceLock::new();
45    EPOCH.get_or_init(Instant::now).elapsed().as_millis() as u64
46}
47
48/// Current lifecycle state of a stream.
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum StreamState {
51    /// No publisher yet.
52    Idle,
53    /// A publisher is connected and sending data.
54    Publishing,
55    /// The stream is being transcoded into one or more renditions.
56    Transcoding,
57    /// The stream is being recorded.
58    Recording,
59    /// The publisher has disconnected; the stream has ended.
60    Ended,
61}
62
63impl std::fmt::Display for StreamState {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.write_str(match self {
66            StreamState::Idle => "idle",
67            StreamState::Publishing => "publishing",
68            StreamState::Transcoding => "transcoding",
69            StreamState::Recording => "recording",
70            StreamState::Ended => "ended",
71        })
72    }
73}
74
75/// Runtime metadata about a stream, updated continuously while publishing.
76///
77/// Resolution and the ingest protocol are set by the protocol handler (e.g. via
78/// [`StreamHandle::update_metadata`] after parsing the codec config), while the
79/// `fps` and `*_bitrate_bps` fields are overlaid live from measured throughput
80/// by [`StreamHandle::metadata_snapshot`].
81#[derive(Debug, Clone)]
82pub struct StreamMetadata {
83    /// The `(app, stream_id)` this metadata describes.
84    pub key: StreamKey,
85    /// Publisher remote address.
86    pub publisher_addr: Option<SocketAddr>,
87    /// Video width in pixels (0 = unknown).
88    pub width: u32,
89    /// Video height in pixels (0 = unknown).
90    pub height: u32,
91    /// Video frames per second (0 = unknown). Overlaid from measured throughput.
92    pub fps: f64,
93    /// Measured video ingest bitrate in bits-per-second.
94    pub video_bitrate_bps: u64,
95    /// Measured audio ingest bitrate in bits-per-second.
96    pub audio_bitrate_bps: u64,
97    /// Timestamp of the first frame received (Unix ms).
98    pub started_at_ms: u64,
99    /// Protocol used for ingest (e.g. `"rtmp"`).
100    pub ingest_protocol: String,
101}
102
103impl StreamMetadata {
104    /// Create zeroed metadata for `(app, stream_id)`.
105    pub fn new(app: AppName, stream_id: StreamId) -> Self {
106        Self {
107            key: StreamKey::new(app, stream_id),
108            publisher_addr: None,
109            width: 0,
110            height: 0,
111            fps: 0.0,
112            video_bitrate_bps: 0,
113            audio_bitrate_bps: 0,
114            started_at_ms: 0,
115            ingest_protocol: String::new(),
116        }
117    }
118}
119
120/// A point-in-time snapshot of a stream's measured quality of service.
121#[derive(Debug, Clone, Copy, Default)]
122pub struct Qos {
123    /// Video bitrate over the last ~1s window (bits/sec).
124    pub video_bitrate_bps: u64,
125    /// Audio bitrate over the last ~1s window (bits/sec).
126    pub audio_bitrate_bps: u64,
127    /// Video frames per second over the last ~1s window.
128    pub fps: f64,
129    /// Cumulative frames published on this stream.
130    pub total_frames: u64,
131    /// Cumulative payload bytes published on this stream.
132    pub total_bytes: u64,
133}
134
135/// Lock-free throughput counters folded into [`Qos`] / [`StreamMetadata`].
136///
137/// A stream has a single publisher (enforced by `start_publish`), so the
138/// read-modify-write here is effectively single-writer and `Relaxed` is sound.
139#[derive(Default)]
140struct QosCounters {
141    total_frames: AtomicU64,
142    total_bytes: AtomicU64,
143    window_start_ms: AtomicU64,
144    window_video_bytes: AtomicU64,
145    window_audio_bytes: AtomicU64,
146    window_video_frames: AtomicU64,
147    cur_video_bitrate: AtomicU64,
148    cur_audio_bitrate: AtomicU64,
149    cur_fps_milli: AtomicU64, // fps × 1000, integer-encoded
150    last_frame_ms: AtomicU64,
151}
152
153/// Keyframe-anchored replay buffer for instant playback start.
154struct GopBuffer {
155    /// Frames since (and including) the most recent keyframe.
156    frames: Vec<Arc<MediaFrame>>,
157    /// Hard cap on buffered frames (frame-count bound between keyframes).
158    capacity: usize,
159    /// Hard cap on buffered payload **bytes** between keyframes (0 = unbounded).
160    /// Frame count alone does not bound memory — one 4K keyframe can be megabytes
161    /// — so a high-bitrate stream can be capped by size independently of count.
162    byte_capacity: usize,
163    /// Payload bytes currently buffered in `frames` (kept in sync on push/clear).
164    bytes: usize,
165    /// Whether the current GOP already overflowed `capacity`/`byte_capacity` (so
166    /// the truncation signal fires once per GOP, not once per dropped frame).
167    overflowed: bool,
168}
169
170/// A live handle to a single active stream.
171///
172/// Multiple subscribers (HLS packager, DASH packager, WebRTC SFU, recorders …)
173/// call [`StreamHandle::subscribe_resilient`] to receive every [`MediaFrame`]
174/// cheaply via a `broadcast` channel (zero-copy `Bytes` cloning).
175///
176/// Each broadcast slot holds one `Arc<MediaFrame>` pointer (8 bytes), so e.g.
177/// 4096 slots ≈ 32 KB per stream.
178///
179/// # Backpressure model
180///
181/// Fan-out uses a **single, fixed-capacity ring buffer per stream** (a
182/// `tokio::broadcast` channel sized by `AppSpec::broadcast_capacity`). This is a
183/// deliberate design choice, with consequences worth understanding:
184///
185/// - **The publisher never blocks on a slow subscriber.** Publishing is a
186///   non-awaiting pointer write; one subscriber falling behind can never apply
187///   backpressure to the publisher or to its peers. This is what keeps the hot
188///   path lock-free and the fast publisher isolated from the slow viewer.
189/// - **Backpressure is resolved by dropping, not stalling.** A subscriber that
190///   can't keep up overruns the ring and observes lag.
191///   [`Subscription::recv`] resynchronizes to the oldest still-buffered frame
192///   and reports the gap via [`Observer::on_subscriber_lagged`]; with
193///   [`Subscription::max_lag`] a chronically slow consumer is evicted
194///   ([`Observer::on_subscriber_evicted`]) rather than churning forever.
195/// - **Capacity is the tuning knob**, traded per stream: larger capacity
196///   tolerates burstier consumers at higher per-stream memory, smaller capacity
197///   sheds laggards sooner. There is intentionally no per-subscriber queue —
198///   that would reintroduce unbounded memory growth and per-consumer locking,
199///   the very things this design avoids.
200///
201/// In short: a slow subscriber degrades only *its own* view (lag, then
202/// eviction), never the publisher's or another subscriber's. Wire an
203/// [`Observer`] to see lag and eviction as they happen.
204#[derive(Clone)]
205pub struct StreamHandle {
206    metadata: Arc<RwLock<StreamMetadata>>,
207    state: Arc<RwLock<StreamState>>,
208    key: StreamKey,
209    /// The frame-bus sender, held *indirectly* so its lifetime is owned by the
210    /// stream's lifecycle — not by how many `StreamHandle` clones happen to be
211    /// alive. Cloning a handle shares this cell; it does not mint a new sender.
212    ///
213    /// This is the structural fix for a sharp edge: previously the handle stored
214    /// the `broadcast::Sender` by value, so any consumer that merely retained a
215    /// handle (to subscribe, read metadata, request a keyframe) silently pinned
216    /// the channel open and defeated the `Closed` shutdown signal. Now
217    /// [`close`](Self::close) — called by the registry when a publish ends — empties
218    /// this cell, dropping the sole sender and closing the channel regardless of
219    /// any lingering handle clones.
220    tx: Arc<ArcSwapOption<broadcast::Sender<Arc<MediaFrame>>>>,
221    /// Latest video CONFIG (AVCDecoderConfigurationRecord) frame, if seen.
222    /// Uses `ArcSwap` for lock-free reads from multiple subscriber tasks.
223    video_config: Arc<ArcSwap<Option<Arc<MediaFrame>>>>,
224    /// Latest audio CONFIG (AudioSpecificConfig) frame, if seen.
225    audio_config: Arc<ArcSwap<Option<Arc<MediaFrame>>>>,
226    /// Rolling GOP buffer for instant-start (empty when `gop_capacity == 0`).
227    gop: Arc<StdMutex<GopBuffer>>,
228    gop_capacity: usize,
229    /// Default cumulative-lag budget applied to [`subscribe_resilient`]
230    /// subscriptions: a consumer that drops more than this many frames is evicted.
231    /// `None` = no default bound (a caller may still set one per subscription).
232    default_max_lag: Option<u64>,
233    /// Live throughput counters.
234    qos: Arc<QosCounters>,
235    /// Injected telemetry hook (no-op by default).
236    observer: Arc<dyn Observer>,
237    /// Upstream keyframe-request signal, shared across this handle's clones. A
238    /// playback consumer (e.g. a WHEP viewer whose decoder needs an IDR) calls
239    /// [`request_keyframe`](Self::request_keyframe); the *publisher's* ingest loop
240    /// awaits [`keyframe_requested`](Self::keyframe_requested) and asks its source
241    /// for a fresh keyframe (e.g. an RTCP PLI to a WHIP publisher).
242    keyframe_request: Arc<Notify>,
243}
244
245impl StreamHandle {
246    /// Create a handle with the no-op observer and no GOP cache.
247    pub fn new(app: AppName, stream_id: StreamId, capacity: usize) -> Self {
248        Self::with_observer(app, stream_id, capacity, 0, Arc::new(NoopObserver))
249    }
250
251    /// Create a handle wired to a host-supplied observer.
252    ///
253    /// `gop_capacity` bounds the keyframe-anchored replay buffer (0 disables it).
254    pub fn with_observer(
255        app: AppName,
256        stream_id: StreamId,
257        capacity: usize,
258        gop_capacity: usize,
259        observer: Arc<dyn Observer>,
260    ) -> Self {
261        Self::with_config(app, stream_id, capacity, gop_capacity, 0, None, observer)
262    }
263
264    /// Create a handle with the full per-stream policy: an optional GOP byte cap
265    /// (`gop_byte_capacity`, 0 = unbounded) and an optional default subscriber
266    /// lag budget (`default_max_lag`). The engine threads these from
267    /// [`AppSpec`](crate::AppSpec); [`with_observer`](Self::with_observer) is the
268    /// shorthand with both disabled.
269    #[allow(clippy::too_many_arguments)]
270    pub fn with_config(
271        app: AppName,
272        stream_id: StreamId,
273        capacity: usize,
274        gop_capacity: usize,
275        gop_byte_capacity: usize,
276        default_max_lag: Option<u64>,
277        observer: Arc<dyn Observer>,
278    ) -> Self {
279        let (tx, _) = broadcast::channel(capacity);
280        let qos = QosCounters::default();
281        // Treat creation as the last activity so a just-claimed stream is not
282        // instantly considered idle before its first frame arrives. Monotonic, to
283        // match the idle reaper (see `mono_ms`).
284        qos.last_frame_ms.store(mono_ms(), Ordering::Relaxed);
285        Self {
286            metadata: Arc::new(RwLock::new(StreamMetadata::new(
287                app.clone(),
288                stream_id.clone(),
289            ))),
290            state: Arc::new(RwLock::new(StreamState::Idle)),
291            key: StreamKey::new(app, stream_id),
292            tx: Arc::new(ArcSwapOption::new(Some(Arc::new(tx)))),
293            video_config: Arc::new(ArcSwap::new(Arc::new(None))),
294            audio_config: Arc::new(ArcSwap::new(Arc::new(None))),
295            gop: Arc::new(StdMutex::new(GopBuffer {
296                frames: Vec::new(),
297                capacity: gop_capacity,
298                byte_capacity: gop_byte_capacity,
299                bytes: 0,
300                overflowed: false,
301            })),
302            gop_capacity,
303            default_max_lag,
304            qos: Arc::new(qos),
305            observer,
306            keyframe_request: Arc::new(Notify::new()),
307        }
308    }
309
310    /// Request a fresh keyframe from the stream's publisher.
311    ///
312    /// Wakes the publisher's ingest loop (see [`keyframe_requested`]); a no-op if
313    /// nothing is currently awaiting. Used by playback consumers — e.g. a WHEP
314    /// viewer relaying a browser PLI — to recover when the GOP cache holds no
315    /// keyframe a new subscriber can start from.
316    ///
317    /// [`keyframe_requested`]: Self::keyframe_requested
318    pub fn request_keyframe(&self) {
319        self.keyframe_request.notify_one();
320    }
321
322    /// Await the next keyframe request raised by [`request_keyframe`](Self::request_keyframe).
323    ///
324    /// The publisher's ingest loop selects on this and, when it fires, asks its
325    /// source for a fresh keyframe (e.g. an RTCP PLI to a WHIP publisher).
326    pub async fn keyframe_requested(&self) {
327        self.keyframe_request.notified().await;
328    }
329
330    /// The `(app, stream_id)` this handle belongs to.
331    pub fn key(&self) -> &StreamKey {
332        &self.key
333    }
334
335    /// Publish a frame to all current subscribers.  Returns the number of
336    /// active receivers; returns `Ok(0)` when there are no subscribers.
337    pub fn publish_frame(&self, frame: MediaFrame) -> crate::Result<usize> {
338        self.observer.on_frame(&self.key, &frame);
339        let len = frame.data.len() as u64;
340        let is_audio = frame.is_audio();
341        let is_key = frame.is_keyframe();
342        let is_config = frame.flags.contains(FrameFlags::CONFIG);
343        let arc = Arc::new(frame);
344
345        // Cache the latest CONFIG frame for late-joining subscribers.
346        if is_config {
347            if is_audio {
348                self.audio_config.store(Arc::new(Some(Arc::clone(&arc))));
349            } else {
350                self.video_config.store(Arc::new(Some(Arc::clone(&arc))));
351            }
352        }
353
354        // Maintain the keyframe-anchored GOP replay buffer.
355        let mut gop_truncated = false;
356        if self.gop_capacity > 0 {
357            if let Ok(mut g) = self.gop.lock() {
358                if is_key {
359                    g.frames.clear();
360                    g.bytes = len as usize;
361                    g.overflowed = false;
362                    g.frames.push(Arc::clone(&arc));
363                } else if !is_config && !g.frames.is_empty() {
364                    // Only buffer once a keyframe anchors the GOP; CONFIG frames
365                    // are replayed separately via `cached_configs`.
366                    //
367                    // Admit the frame only if it stays within *both* the frame
368                    // count and the byte budget (0 = byte cap disabled). Frame
369                    // count alone does not bound memory, so the byte cap keeps a
370                    // high-bitrate GOP from growing without limit.
371                    let fits_count = g.frames.len() < g.capacity;
372                    let fits_bytes =
373                        g.byte_capacity == 0 || g.bytes + len as usize <= g.byte_capacity;
374                    if fits_count && fits_bytes {
375                        g.bytes += len as usize;
376                        g.frames.push(Arc::clone(&arc));
377                    } else if !g.overflowed {
378                        // First overflow of this GOP: the replay will be truncated.
379                        // Latch it so the signal fires once, not per dropped frame.
380                        g.overflowed = true;
381                        gop_truncated = true;
382                    }
383                }
384            }
385        }
386        if gop_truncated {
387            self.observer.on_gop_truncated(&self.key, self.gop_capacity);
388        }
389
390        self.record_qos(len, is_audio, is_key);
391
392        // The sender is gone once the stream is closed; publishing then is a
393        // no-op (Ok(0)) rather than an error, so a publisher racing teardown
394        // winds down cleanly.
395        let count = match self.tx.load_full() {
396            Some(tx) => tx.send(arc).unwrap_or(0),
397            None => 0,
398        };
399        Ok(count)
400    }
401
402    /// Fold one frame into the rolling throughput window.
403    fn record_qos(&self, len: u64, is_audio: bool, _is_key: bool) {
404        let q = &self.qos;
405        let now = mono_ms();
406        q.total_frames.fetch_add(1, Ordering::Relaxed);
407        q.total_bytes.fetch_add(len, Ordering::Relaxed);
408        q.last_frame_ms.store(now, Ordering::Relaxed);
409        if is_audio {
410            q.window_audio_bytes.fetch_add(len, Ordering::Relaxed);
411        } else {
412            q.window_video_bytes.fetch_add(len, Ordering::Relaxed);
413            q.window_video_frames.fetch_add(1, Ordering::Relaxed);
414        }
415
416        let ws = q.window_start_ms.load(Ordering::Relaxed);
417        if ws == 0 {
418            q.window_start_ms.store(now, Ordering::Relaxed);
419        } else if now.saturating_sub(ws) >= 1000 {
420            let elapsed = (now - ws) as f64 / 1000.0;
421            let vbytes = q.window_video_bytes.swap(0, Ordering::Relaxed);
422            let abytes = q.window_audio_bytes.swap(0, Ordering::Relaxed);
423            let vframes = q.window_video_frames.swap(0, Ordering::Relaxed);
424            q.cur_video_bitrate
425                .store((vbytes as f64 * 8.0 / elapsed) as u64, Ordering::Relaxed);
426            q.cur_audio_bitrate
427                .store((abytes as f64 * 8.0 / elapsed) as u64, Ordering::Relaxed);
428            q.cur_fps_milli.store(
429                (vframes as f64 / elapsed * 1000.0) as u64,
430                Ordering::Relaxed,
431            );
432            q.window_start_ms.store(now, Ordering::Relaxed);
433        }
434    }
435
436    /// A snapshot of measured throughput (bitrate, fps, totals).
437    pub fn qos(&self) -> Qos {
438        let q = &self.qos;
439        Qos {
440            video_bitrate_bps: q.cur_video_bitrate.load(Ordering::Relaxed),
441            audio_bitrate_bps: q.cur_audio_bitrate.load(Ordering::Relaxed),
442            fps: q.cur_fps_milli.load(Ordering::Relaxed) as f64 / 1000.0,
443            total_frames: q.total_frames.load(Ordering::Relaxed),
444            total_bytes: q.total_bytes.load(Ordering::Relaxed),
445        }
446    }
447
448    /// Monotonic-clock timestamp (process-local milliseconds) of the most
449    /// recently published frame (or stream creation if none yet). Used by the
450    /// engine's idle reaper; this is elapsed monotonic time, not wall-clock time,
451    /// so compare it only against other readings of the same monotonic clock.
452    pub fn last_frame_ms(&self) -> u64 {
453        self.qos.last_frame_ms.load(Ordering::Relaxed)
454    }
455
456    /// Returns the most recently seen video and audio CONFIG frames,
457    /// for replaying to late-joining subscribers.
458    pub fn cached_configs(&self) -> (Option<Arc<MediaFrame>>, Option<Arc<MediaFrame>>) {
459        let video = (**self.video_config.load()).clone();
460        let audio = (**self.audio_config.load()).clone();
461        (video, audio)
462    }
463
464    /// The frames a late joiner should be handed before going live: cached
465    /// decoder configs followed by the current GOP (keyframe + trailing deltas).
466    ///
467    /// Replaying these lets a new subscriber start decoding immediately rather
468    /// than waiting for the next keyframe — sub-second join times at scale.
469    /// Requires the app to have enabled a GOP cache; otherwise only the cached
470    /// configs are returned.
471    pub fn replay_buffer(&self) -> Vec<Arc<MediaFrame>> {
472        let (vcfg, acfg) = self.cached_configs();
473        let mut out = Vec::new();
474        out.extend(vcfg.clone());
475        out.extend(acfg.clone());
476        if self.gop_capacity > 0 {
477            if let Ok(g) = self.gop.lock() {
478                out.reserve(g.frames.len());
479                for f in &g.frames {
480                    // Avoid duplicating a config frame already pushed above. Only
481                    // the (≤2) cached configs can collide with a GOP frame — GOP
482                    // frames are themselves distinct — so compare against just
483                    // those, keeping this O(n) rather than O(n²) over the GOP.
484                    let dup = vcfg.as_ref().is_some_and(|c| Arc::ptr_eq(c, f))
485                        || acfg.as_ref().is_some_and(|c| Arc::ptr_eq(c, f));
486                    if !dup {
487                        out.push(Arc::clone(f));
488                    }
489                }
490            }
491        }
492        out
493    }
494
495    /// Subscribe to this stream's frame bus.
496    ///
497    /// The returned raw [`broadcast::Receiver`] surfaces [`RecvError::Lagged`]
498    /// when a slow consumer falls behind the channel capacity — callers that
499    /// `while let Ok(_) = rx.recv().await` will silently terminate on the first
500    /// lag. Prefer [`subscribe_resilient`](Self::subscribe_resilient) unless you
501    /// are deliberately handling lag yourself.
502    pub fn subscribe(&self) -> broadcast::Receiver<Arc<MediaFrame>> {
503        match self.tx.load_full() {
504            Some(tx) => tx.subscribe(),
505            // The stream has already closed: hand back a receiver on a spent
506            // channel so the caller's `recv` loop terminates immediately with
507            // `Closed` rather than blocking forever.
508            None => {
509                let (_, rx) = broadcast::channel(1);
510                rx
511            }
512        }
513    }
514
515    /// Subscribe with a [`Subscription`] that resynchronizes after lag instead
516    /// of terminating, reporting each gap to the installed [`Observer`] via
517    /// [`Observer::on_subscriber_lagged`].
518    ///
519    /// If the app configured [`AppSpec::subscriber_max_lag`](crate::AppSpec::subscriber_max_lag),
520    /// that bound is applied by default, so a chronically slow consumer is evicted
521    /// automatically; a caller may still tighten it with [`Subscription::max_lag`].
522    pub fn subscribe_resilient(&self) -> Subscription {
523        Subscription {
524            rx: self.subscribe(),
525            key: self.key.clone(),
526            observer: Arc::clone(&self.observer),
527            max_lag: self.default_max_lag,
528            skipped: 0,
529        }
530    }
531
532    /// Drive this stream to a [`FrameSink`]: replay the instant-start buffer
533    /// (cached configs + cached GOP) so the consumer can decode immediately,
534    /// then forward live frames until the stream ends or `shutdown` fires.
535    ///
536    /// This is the playback loop every egress that pulls from a single stream
537    /// shares (RTSP server, RTMP relay, …); each supplies only its per-frame
538    /// [`FrameSink::send`].
539    pub async fn drive_to(
540        &self,
541        shutdown: &CancellationToken,
542        sink: &mut dyn FrameSink,
543    ) -> Result<()> {
544        for frame in self.replay_buffer() {
545            if sink.send(frame).await?.is_break() {
546                return Ok(());
547            }
548        }
549        let mut sub = self.subscribe_resilient();
550        loop {
551            tokio::select! {
552                _ = shutdown.cancelled() => break,
553                next = sub.recv() => match next {
554                    Some(frame) => {
555                        if sink.send(frame).await?.is_break() {
556                            break;
557                        }
558                    }
559                    None => break, // stream ended
560                }
561            }
562        }
563        Ok(())
564    }
565
566    /// Drive this stream into a [`Packager`](crate::packager::Packager) (HLS
567    /// segmenter, DASH packager, …): replay the instant-start buffer — **which
568    /// carries the cached CONFIG access unit (SPS/PPS)** — then forward live
569    /// frames until the stream ends or `shutdown` fires, and `finish`.
570    ///
571    /// Priming with [`replay_buffer`](Self::replay_buffer) is essential, not just
572    /// an optimization: the CONFIG frame is published once at stream start, so a
573    /// packager that only `subscribe_resilient`s (subscribing *after* the event
574    /// that spawned it) never sees it — and then a fragmented-MP4 muxer never
575    /// builds its `init.m4s` and the master playlist never learns its codec
576    /// string. Replaying configs first makes both deterministic.
577    #[cfg(feature = "hls")]
578    pub async fn package_to(
579        &self,
580        shutdown: &CancellationToken,
581        packager: &mut dyn crate::packager::Packager,
582    ) -> Result<()> {
583        for frame in self.replay_buffer() {
584            packager.push(&frame).await?;
585        }
586        let mut sub = self.subscribe_resilient();
587        loop {
588            tokio::select! {
589                _ = shutdown.cancelled() => break,
590                next = sub.recv() => match next {
591                    Some(frame) => packager.push(&frame).await?,
592                    None => break, // stream ended
593                }
594            }
595        }
596        packager.finish().await
597    }
598
599    /// Number of active subscribers (0 once the stream is closed).
600    pub fn subscriber_count(&self) -> usize {
601        self.tx
602            .load_full()
603            .map(|tx| tx.receiver_count())
604            .unwrap_or(0)
605    }
606
607    /// Close the frame bus: drop the sole sender so every subscriber's `recv`
608    /// observes `Closed` and terminates, *regardless* of how many `StreamHandle`
609    /// clones are still alive.
610    ///
611    /// Called by the registry when a publish ends (see
612    /// `Application::end_publish`). Idempotent. This is what makes the channel's
613    /// lifetime track the stream's lifecycle rather than handle reachability.
614    pub fn close(&self) {
615        self.tx.store(None);
616    }
617
618    /// Transition to a new state.
619    pub async fn set_state(&self, state: StreamState) {
620        let mut guard = self.state.write().await;
621        *guard = state;
622    }
623
624    /// The current lifecycle state.
625    pub async fn current_state(&self) -> StreamState {
626        self.state.read().await.clone()
627    }
628
629    /// A consistent point-in-time copy of this stream's [`StreamMetadata`], with
630    /// the live measured `fps`/bitrate overlaid from [`qos`](Self::qos).
631    ///
632    /// Cloning the snapshot releases the lock immediately, so callers never hold
633    /// the metadata `RwLock` across an `.await`.
634    pub async fn metadata_snapshot(&self) -> StreamMetadata {
635        let mut m = self.metadata.read().await.clone();
636        let q = self.qos();
637        m.video_bitrate_bps = q.video_bitrate_bps;
638        m.audio_bitrate_bps = q.audio_bitrate_bps;
639        if q.fps > 0.0 {
640            m.fps = q.fps;
641        }
642        m
643    }
644
645    /// Mutate this stream's [`StreamMetadata`] under the write lock.
646    ///
647    /// Ingest handlers call this as they parse the stream — e.g. on the first
648    /// keyframe to record resolution from the codec config, or to set the
649    /// publisher address — so the metadata exposed to operators and the control
650    /// plane stays live rather than frozen at its zeroed defaults.
651    ///
652    /// ```no_run
653    /// # use arcly_stream::StreamHandle;
654    /// # async fn demo(handle: &StreamHandle, addr: std::net::SocketAddr) {
655    /// handle
656    ///     .update_metadata(|m| {
657    ///         m.publisher_addr = Some(addr);
658    ///         m.width = 1920;
659    ///         m.height = 1080;
660    ///         m.ingest_protocol = "rtmp".to_string();
661    ///     })
662    ///     .await;
663    /// # }
664    /// ```
665    pub async fn update_metadata(&self, f: impl FnOnce(&mut StreamMetadata)) {
666        let mut guard = self.metadata.write().await;
667        f(&mut guard);
668    }
669}
670
671impl std::fmt::Debug for StreamHandle {
672    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
673        f.debug_struct("StreamHandle")
674            .field("key", &self.key)
675            .field("subscribers", &self.subscriber_count())
676            .finish()
677    }
678}
679
680/// A lag-tolerant subscription to a stream's frame bus.
681///
682/// Returned by [`StreamHandle::subscribe_resilient`]. Unlike a raw
683/// [`broadcast::Receiver`], [`recv`](Self::recv) does not terminate when the
684/// consumer falls behind: the dropped span is reported to the [`Observer`] as
685/// [`on_subscriber_lagged`](Observer::on_subscriber_lagged) and reception
686/// continues from the oldest still-buffered frame. This is the recommended
687/// consumer loop for packagers, recorders, and SFUs.
688///
689/// An optional [`max_lag`](Self::max_lag) bound turns chronic lag into
690/// eviction: once cumulative dropped frames exceed the bound, `recv` returns
691/// `None` (after an [`on_subscriber_evicted`](Observer::on_subscriber_evicted)
692/// notification) so a hopelessly slow consumer is shed rather than wasting
693/// buffer churn forever.
694pub struct Subscription {
695    rx: broadcast::Receiver<Arc<MediaFrame>>,
696    key: StreamKey,
697    observer: Arc<dyn Observer>,
698    max_lag: Option<u64>,
699    skipped: u64,
700}
701
702impl Subscription {
703    /// Evict this subscriber once cumulative dropped frames exceed `max`.
704    ///
705    /// ```no_run
706    /// # use arcly_stream::StreamHandle;
707    /// # fn demo(handle: &StreamHandle) {
708    /// let sub = handle.subscribe_resilient().max_lag(10_000);
709    /// # let _ = sub;
710    /// # }
711    /// ```
712    pub fn max_lag(mut self, max: u64) -> Self {
713        self.max_lag = Some(max);
714        self
715    }
716
717    /// Total frames dropped from this subscriber's view so far.
718    pub fn dropped(&self) -> u64 {
719        self.skipped
720    }
721
722    /// Receive the next frame, resynchronizing past any lag.
723    ///
724    /// Returns `None` when the stream's sender is dropped (the publisher ended)
725    /// or when the `max_lag` eviction threshold is crossed:
726    ///
727    /// ```no_run
728    /// # async fn run(sub: &mut arcly_stream::bus::Subscription) {
729    /// while let Some(frame) = sub.recv().await {
730    ///     // packetize `frame` …
731    /// }
732    /// # }
733    /// ```
734    pub async fn recv(&mut self) -> Option<Arc<MediaFrame>> {
735        loop {
736            match self.rx.recv().await {
737                Ok(frame) => return Some(frame),
738                Err(RecvError::Lagged(skipped)) => {
739                    self.skipped = self.skipped.saturating_add(skipped);
740                    self.observer.on_subscriber_lagged(&self.key, skipped);
741                    if let Some(max) = self.max_lag {
742                        if self.skipped > max {
743                            self.observer.on_subscriber_evicted(&self.key);
744                            return None;
745                        }
746                    }
747                    continue;
748                }
749                Err(RecvError::Closed) => return None,
750            }
751        }
752    }
753
754    /// Borrow the underlying raw receiver, for callers that need
755    /// [`broadcast::Receiver`] APIs directly.
756    pub fn raw(&mut self) -> &mut broadcast::Receiver<Arc<MediaFrame>> {
757        &mut self.rx
758    }
759}
760
761#[cfg(test)]
762mod tests {
763    use super::*;
764    use crate::CodecId;
765
766    fn video(pts: i64, key: bool) -> MediaFrame {
767        MediaFrame::new_video(
768            pts,
769            pts,
770            bytes::Bytes::from_static(&[0, 0, 0, 1, 0x65]),
771            CodecId::H264,
772            key,
773        )
774    }
775
776    /// Regression guard for the sender-lifetime fix: `close()` must terminate a
777    /// subscriber's `recv` even while a `StreamHandle` clone is still held — the
778    /// exact shape that previously hung WHEP egress forever.
779    #[tokio::test]
780    async fn close_terminates_recv_while_a_handle_clone_is_held() {
781        let handle = StreamHandle::new("live".into(), "show".into(), 16);
782        let mut sub = handle.subscribe_resilient();
783
784        // A retained clone (e.g. an egress pump) must NOT pin the channel open.
785        let retained = handle.clone();
786        handle.publish_frame(video(0, true)).unwrap();
787        assert!(sub.recv().await.is_some(), "frame delivered before close");
788
789        retained.close();
790
791        // Without the registry-owned sender this would block forever; bound it so
792        // a regression fails loudly instead of hanging the suite.
793        let got = tokio::time::timeout(std::time::Duration::from_secs(5), sub.recv())
794            .await
795            .expect("recv resolved after close (no hang)");
796        assert!(got.is_none(), "recv returns None once the stream is closed");
797
798        // Publishing post-close is a clean no-op, not a panic.
799        assert_eq!(retained.publish_frame(video(1, false)).unwrap(), 0);
800        assert_eq!(retained.subscriber_count(), 0);
801    }
802
803    /// `package_to` must replay the cached CONFIG access unit even though it was
804    /// published *before* the packager subscribed — the regression behind DASH's
805    /// missing `init.m4s` and the absent master playlist.
806    #[cfg(feature = "hls")]
807    #[tokio::test]
808    async fn package_to_replays_config_published_before_subscribe() {
809        use crate::packager::Packager;
810        use tokio_util::sync::CancellationToken;
811
812        // A trivial packager that records the flags of every frame it is pushed.
813        struct RecordingPackager(std::sync::Arc<std::sync::Mutex<Vec<FrameFlags>>>);
814        #[async_trait::async_trait]
815        impl Packager for RecordingPackager {
816            async fn push(&mut self, frame: &MediaFrame) -> crate::Result<()> {
817                self.0.lock().unwrap().push(frame.flags);
818                Ok(())
819            }
820            async fn finish(&mut self) -> crate::Result<()> {
821                Ok(())
822            }
823        }
824
825        let handle = StreamHandle::new("live".into(), "cam".into(), 16);
826
827        // The one-shot CONFIG frame arrives BEFORE any packager subscribes.
828        let mut cfg = video(0, true);
829        cfg.flags |= FrameFlags::CONFIG;
830        handle.publish_frame(cfg).unwrap();
831
832        let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
833        let mut pkg = RecordingPackager(seen.clone());
834        let h = handle.clone();
835        let shutdown = CancellationToken::new();
836        let task = tokio::spawn(async move { h.package_to(&shutdown, &mut pkg).await });
837
838        // Give the replay + subscribe a moment, then end the stream.
839        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
840        handle.publish_frame(video(1, false)).unwrap();
841        handle.close();
842        task.await.unwrap().unwrap();
843
844        assert!(
845            seen.lock()
846                .unwrap()
847                .iter()
848                .any(|f| f.contains(FrameFlags::CONFIG)),
849            "package_to must replay the cached CONFIG frame to the packager"
850        );
851    }
852
853    /// The GOP-overflow signal fires once per GOP (not per dropped frame) and
854    /// resets on the next keyframe so a fresh GOP can report again.
855    #[test]
856    fn gop_truncation_signals_once_per_gop() {
857        use crate::StreamKey;
858        use std::sync::atomic::{AtomicUsize, Ordering};
859
860        #[derive(Default)]
861        struct CountingObserver {
862            truncations: AtomicUsize,
863        }
864        impl Observer for CountingObserver {
865            fn on_gop_truncated(&self, _key: &StreamKey, _capacity: usize) {
866                self.truncations.fetch_add(1, Ordering::Relaxed);
867            }
868        }
869
870        let obs = Arc::new(CountingObserver::default());
871        // Capacity 2: keyframe + 1 delta fit; the 2nd delta onward overflows.
872        let handle = StreamHandle::with_observer(
873            "live".into(),
874            "g".into(),
875            64,
876            2,
877            Arc::clone(&obs) as Arc<dyn Observer>,
878        );
879
880        handle.publish_frame(video(0, true)).unwrap(); // keyframe anchors GOP
881        handle.publish_frame(video(1, false)).unwrap(); // fills to capacity
882        handle.publish_frame(video(2, false)).unwrap(); // first overflow → signal
883        handle.publish_frame(video(3, false)).unwrap(); // still overflowing → silent
884        assert_eq!(
885            obs.truncations.load(Ordering::Relaxed),
886            1,
887            "one signal per GOP"
888        );
889
890        // A new keyframe resets the latch; overflowing again signals once more.
891        handle.publish_frame(video(4, true)).unwrap();
892        handle.publish_frame(video(5, false)).unwrap();
893        handle.publish_frame(video(6, false)).unwrap();
894        handle.publish_frame(video(7, false)).unwrap();
895        assert_eq!(
896            obs.truncations.load(Ordering::Relaxed),
897            2,
898            "next GOP signals again"
899        );
900    }
901
902    /// The GOP byte cap truncates a high-bitrate GOP independently of the frame
903    /// count: a buffer with a generous frame count but a tight byte budget stops
904    /// admitting frames once the budget is reached.
905    #[test]
906    fn gop_byte_cap_bounds_replay_buffer_memory() {
907        use crate::StreamKey;
908        use std::sync::atomic::{AtomicUsize, Ordering};
909
910        #[derive(Default)]
911        struct CountingObserver {
912            truncations: AtomicUsize,
913        }
914        impl Observer for CountingObserver {
915            fn on_gop_truncated(&self, _key: &StreamKey, _capacity: usize) {
916                self.truncations.fetch_add(1, Ordering::Relaxed);
917            }
918        }
919
920        fn frame(pts: i64, key: bool, len: usize) -> MediaFrame {
921            MediaFrame::new_video(
922                pts,
923                pts,
924                bytes::Bytes::from(vec![0u8; len]),
925                CodecId::H264,
926                key,
927            )
928        }
929
930        let obs = Arc::new(CountingObserver::default());
931        // Frame count is generous (100) but the byte budget admits only ~250 B.
932        let handle = StreamHandle::with_config(
933            "live".into(),
934            "hb".into(),
935            64,
936            100,
937            250,
938            None,
939            Arc::clone(&obs) as Arc<dyn Observer>,
940        );
941
942        handle.publish_frame(frame(0, true, 100)).unwrap(); // keyframe: 100 B
943        handle.publish_frame(frame(1, false, 100)).unwrap(); // +100 → 200 B (fits)
944        handle.publish_frame(frame(2, false, 100)).unwrap(); // +100 → 300 B > 250 → drop
945        handle.publish_frame(frame(3, false, 100)).unwrap(); // still over → silent
946
947        // Only keyframe + one delta were buffered, despite the 100-frame count cap.
948        assert_eq!(handle.replay_buffer().len(), 2);
949        assert_eq!(
950            obs.truncations.load(Ordering::Relaxed),
951            1,
952            "byte-cap overflow signals once per GOP"
953        );
954
955        // A new keyframe resets the byte accounting so the next GOP buffers again.
956        handle.publish_frame(frame(4, true, 100)).unwrap();
957        handle.publish_frame(frame(5, false, 100)).unwrap();
958        assert_eq!(handle.replay_buffer().len(), 2, "fresh GOP buffers anew");
959    }
960
961    /// A default `max_lag` configured on the handle is applied to every resilient
962    /// subscription, so a chronically slow consumer is evicted without the caller
963    /// setting a per-subscription bound.
964    #[tokio::test]
965    async fn default_max_lag_evicts_slow_subscriber() {
966        // Capacity 4, default lag budget 4 frames.
967        let handle = StreamHandle::with_config(
968            "live".into(),
969            "slow".into(),
970            4,
971            0,
972            0,
973            Some(4),
974            Arc::new(NoopObserver),
975        );
976        let mut sub = handle.subscribe_resilient();
977
978        // Overrun the ring far past the budget without ever receiving.
979        for i in 0..50 {
980            handle.publish_frame(video(i, i == 0)).unwrap();
981        }
982        // The first recv observes lag beyond the default budget → eviction (None).
983        assert!(
984            sub.recv().await.is_none(),
985            "subscriber evicted by the app default max_lag"
986        );
987    }
988
989    /// `mono_ms` is monotonic and drives `last_frame_ms`, so the idle reaper's
990    /// elapsed-time math is independent of the wall clock.
991    #[test]
992    fn mono_clock_is_monotonic_and_drives_last_frame() {
993        let a = mono_ms();
994        let b = mono_ms();
995        assert!(b >= a, "monotonic clock never goes backward");
996
997        let handle = StreamHandle::new("live".into(), "m".into(), 4);
998        let before = mono_ms();
999        handle.publish_frame(video(0, true)).unwrap();
1000        assert!(
1001            handle.last_frame_ms() >= before,
1002            "last_frame_ms advances on the monotonic clock"
1003        );
1004    }
1005}