diff --git a/Cargo.lock b/Cargo.lock index bf5df632a..b7b89f2a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -929,8 +929,7 @@ dependencies = [ [[package]] name = "cadeau" version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7abc663bd17ad2e21ad335c0d2c3fd51b341113c4165b8b5cc078dd15b97f3f9" +source = "git+https://github.com/Devolutions/cadeau?rev=d990b0a96e60226f5107e7838dde2487139c2207#d990b0a96e60226f5107e7838dde2487139c2207" dependencies = [ "xmf-sys", ] @@ -8572,6 +8571,7 @@ version = "0.0.0" dependencies = [ "anyhow", "axum 0.8.9", + "bytes 1.12.1", "cadeau", "criterion", "ebml-iterable", @@ -9522,8 +9522,7 @@ dependencies = [ [[package]] name = "xmf-sys" version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe719e326e963b00276c45c7d8a10567330106f48ea56ad25976b4be0e6caebc" +source = "git+https://github.com/Devolutions/cadeau?rev=d990b0a96e60226f5107e7838dde2487139c2207#d990b0a96e60226f5107e7838dde2487139c2207" dependencies = [ "dlib", ] diff --git a/Cargo.toml b/Cargo.toml index 2cc707e17..c61e538e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ codegen-units = 1 lto = true [patch.crates-io] +cadeau = { git = "https://github.com/Devolutions/cadeau", rev = "d990b0a96e60226f5107e7838dde2487139c2207" } ebml-iterable = { git = "https://github.com/irvingoujAtDevolution/ebml-iterable", tag = "v0.6.3-devo1" } tracing-appender = { git = "https://github.com/CBenoit/tracing.git", rev = "42097daf92e683cf18da7639ddccb056721a796c" } diff --git a/crates/video-streamer/Cargo.toml b/crates/video-streamer/Cargo.toml index dd93154e4..17d010d40 100644 --- a/crates/video-streamer/Cargo.toml +++ b/crates/video-streamer/Cargo.toml @@ -13,6 +13,7 @@ bench = ["perf-diagnostics"] [dependencies] anyhow = "1.0" +bytes = "1" futures-util = { version = "0.3", features = ["sink"] } tokio = { version = "1.52", features = [ "io-util", diff --git a/crates/video-streamer/README.md b/crates/video-streamer/README.md index 1151304f2..023f0fa86 100644 --- a/crates/video-streamer/README.md +++ b/crates/video-streamer/README.md @@ -1,58 +1,78 @@ # video-streamer -This crate takes an unseekable WebM recording (typically from Chrome CaptureStream) and rewrites it into a “fresh” WebM stream that can start playing immediately. -It does this by parsing the incoming WebM, finding the correct cut point, and re-encoding frames. -The output stream begins with a keyframe and valid headers. +`video-streamer` converts one logical recording session into a pull-driven stream of independent VP8 WebM segments. -## Prerequisites - -This crate relies on `cadeau` and its XMF backend for VP8/VP9 decode+encode. -To override which XMF implementation is used at runtime, set `DGATEWAY_LIB_XMF_PATH` to an `xmf.dll` path before running tests or benches. +The input may contain several append-only clips. +Each clip may contain VP8 or VP9 and may change resolution. +The output keeps one transport connection, always uses VP8, and starts a new fixed-size WebM segment at every clip or resolution boundary. -Example: +## Interface -`$env:DGATEWAY_LIB_XMF_PATH = 'D:\library\cadeau\xmf.dll'` +Call `stream_session` with a recording event stream and a message transport. +The transport must implement `Stream> + Sink`. -## Tests +```rust +stream_session(recording_events, transport, SessionConfig::default()).await?; +``` -Run all tests: +The input must follow this grammar: -`cargo test -p video-streamer` +```text +(ClipStarted Bytes* CaughtUp Bytes* ClipEnded)* SessionEnded +``` -Run the WebM streaming correctness suite: +Use `StartAt::LiveEdge` for the clip that was already growing when a consumer joined. +The streamer retains only that clip's latest group of pictures until `CaughtUp` arrives. +Use `StartAt::Beginning` for clips that start after the consumer joins. -`cargo test -p video-streamer --test webm_stream_correctness -- --nocapture` +The incremental decoder owns incomplete EBML bytes between `Bytes` events. +The caller never seeks, rolls back, or retries an incomplete element. +The decoder limits one buffered EBML element and one retained group of pictures to 64 MiB each. -Some tests are marked `#[ignore]` because they require large local assets or are intended for local investigation. -Run ignored tests with: +## Wire protocol -`cargo test -p video-streamer -- --ignored --nocapture` +Each transport item is one complete protocol message. +For WebSocket use, one item maps to one binary WebSocket message. +The first byte is its type code. -Test assets live under `testing-assets\`. +Client messages: -## Logging and diagnostics +| Code | Message | Payload | +| ---: | --- | --- | +| `0` | Start | Empty | +| `1` | Pull | Empty | -Most detailed diagnostics are compiled out by default to keep production logs clean. -To include extra diagnostics, build with `perf-diagnostics`: +Server messages: -`cargo test -p video-streamer --features perf-diagnostics -- --nocapture` +| Code | Message | Payload | +| ---: | --- | --- | +| `0` | Chunk | WebM bytes | +| `1` | Segment started | `{"codec":"vp8","sequence":N,"width":W,"height":H}` | +| `2` | Error | `{"error":"UnexpectedError"}` | +| `3` | Stream ended | Empty | -Then set `RUST_LOG` as needed. -Example: +`Start` requests the first `Segment started` message. +Each `Pull` requests exactly one later server message. +The next `Segment started` message ends the previous segment implicitly. +`Stream ended` ends the final segment and the session. -`$env:RUST_LOG = 'video_streamer=trace'` +Every segment has its own EBML and Tracks headers. +Every segment begins with a keyframe and keeps one resolution. -## Benchmarks - -The main benchmark is `benches\vpx_reencode.rs`. -Run it with: +## Prerequisites -`cargo bench -p video-streamer --bench vpx_reencode --features bench -- --nocapture` +This crate uses `cadeau` and its XMF backend for VP8 and VP9 decoding and VP8 encoding. +The streamer reads the dimensions of every decoded image so source resolution changes do not depend on codec header parsing. +Set `DGATEWAY_LIB_XMF_PATH` when the default XMF library is unavailable. -Benchmark output is intentionally quiet by default. -To print detailed per-run results, set `VIDEO_STREAMER_BENCH_VERBOSE`: +```powershell +$env:DGATEWAY_LIB_XMF_PATH = 'D:\library\cadeau\xmf.dll' +``` -`$env:VIDEO_STREAMER_BENCH_VERBOSE = '1'` +## Checks -To correlate benchmark results with internal timing, also enable `perf-diagnostics` (the `bench` feature enables it). -This is intentionally a build-time gate so production logs stay clean. +```powershell +cargo +nightly fmt --all +cargo check -p video-streamer --tests +cargo clippy -p video-streamer --tests -- -D warnings +``` diff --git a/crates/video-streamer/src/decoder.rs b/crates/video-streamer/src/decoder.rs new file mode 100644 index 000000000..b8fce099e --- /dev/null +++ b/crates/video-streamer/src/decoder.rs @@ -0,0 +1,55 @@ +use anyhow::Context as _; +use cadeau::xmf::vpx::{VpxCodec, VpxDecoder, VpxImage}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct Dimensions { + pub width: u32, + pub height: u32, +} + +pub(crate) struct DecodedFrame<'decoder> { + pub image: VpxImage<'decoder>, + pub dimensions: Dimensions, +} + +pub(crate) struct InputDecoder { + codec: VpxCodec, + threads: u32, + decoder: Option, +} + +impl InputDecoder { + pub(crate) fn new(codec: VpxCodec, threads: u32) -> Self { + Self { + codec, + threads, + decoder: None, + } + } + + pub(crate) fn decode<'decoder>(&'decoder mut self, data: &[u8]) -> anyhow::Result> { + if self.decoder.is_none() { + self.decoder = Some( + VpxDecoder::builder() + .threads(self.threads) + .width(0) + .height(0) + .codec(self.codec) + .build()?, + ); + } + + let decoder = self.decoder.as_mut().context("input decoder is missing")?; + decoder.decode(data)?; + let image = decoder.next_frame()?; + let dimensions = Dimensions { + width: image.width(), + height: image.height(), + }; + anyhow::ensure!( + dimensions.width > 0 && dimensions.height > 0, + "decoder returned invalid frame dimensions" + ); + Ok(DecodedFrame { image, dimensions }) + } +} diff --git a/crates/video-streamer/src/lib.rs b/crates/video-streamer/src/lib.rs index e568689a0..04e480466 100644 --- a/crates/video-streamer/src/lib.rs +++ b/crates/video-streamer/src/lib.rs @@ -25,7 +25,11 @@ macro_rules! perf_debug { pub mod config; pub mod debug; +mod decoder; +mod normalizer; +mod protocol; pub mod reopenable; +mod session; pub(crate) mod streamer; #[macro_use] @@ -39,6 +43,8 @@ pub use streamer::reopenable_file::ReOpenableFile; pub use streamer::signal_writer::SignalWriter; #[rustfmt::skip] pub use streamer::webm_stream; +#[rustfmt::skip] +pub use session::{RecordingEvent, SessionConfig, StartAt, stream_session}; #[cfg(feature = "bench")] pub mod bench_support; diff --git a/crates/video-streamer/src/normalizer.rs b/crates/video-streamer/src/normalizer.rs new file mode 100644 index 000000000..1630bde5c --- /dev/null +++ b/crates/video-streamer/src/normalizer.rs @@ -0,0 +1,639 @@ +use std::io::{self, Write}; +use std::pin::Pin; +use std::task::{Context as TaskContext, Poll}; + +use anyhow::Context; +use bytes::{Bytes, BytesMut}; +use cadeau::xmf::vpx::{VpxCodec, VpxEncoder, VpxEncoderPreset, VpxImage}; +use ebml_iterable::TagDecoder; +use futures_util::{Stream, StreamExt}; +use tokio::sync::mpsc; +use webm_iterable::matroska_spec::{Master, MatroskaSpec, SimpleBlock}; +use webm_iterable::{WebmWriter, WriteOptions}; + +use crate::decoder::{Dimensions, InputDecoder}; +use crate::session::{RecordingEvent, SessionConfig, StartAt}; +use crate::streamer::block_tag::{VideoBlock, is_vpx_key_frame}; + +const OUTPUT_CHANNEL_CAPACITY: usize = 1; +const INPUT_CHANNEL_CAPACITY: usize = 1; +const INPUT_CHUNK_SIZE: usize = 64 * 1024; +const MAX_BUFFERED_TAG_BYTES: usize = 64 * 1024 * 1024; +const MAX_PENDING_GOP_BYTES: usize = 64 * 1024 * 1024; +const OUTPUT_BITRATE: u32 = 256 * 1024; +const VPX_EFLAG_FORCE_KF: u32 = 0x0000_0001; +const WEBM_TIMESTAMP_SCALE_NS: u64 = 1_000_000; +const MAX_WEBM_BLOCK_TIMESTAMP: u64 = 32_767; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct SegmentInfo { + pub sequence: u64, + pub width: u32, + pub height: u32, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum SegmentEvent { + Begin(SegmentInfo), + Data(Bytes), + End, +} + +pub(crate) struct NormalizedSession { + receiver: mpsc::Receiver>, + supervisor: Option>, +} + +impl Stream for NormalizedSession { + type Item = anyhow::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + self.receiver.poll_recv(cx) + } +} + +impl NormalizedSession { + pub(crate) async fn shutdown(mut self) -> anyhow::Result<()> { + self.receiver.close(); + let supervisor = self.supervisor.take().context("normalizer supervisor is missing")?; + supervisor.await.context("normalizer supervisor failed") + } +} + +impl Drop for NormalizedSession { + fn drop(&mut self) { + if let Some(supervisor) = self.supervisor.take() { + supervisor.abort(); + } + } +} + +pub(crate) fn normalize(source: S, config: SessionConfig) -> NormalizedSession +where + S: Stream> + Send + 'static, +{ + let (output_sender, output_receiver) = mpsc::channel(OUTPUT_CHANNEL_CAPACITY); + let (input_sender, input_receiver) = mpsc::channel(INPUT_CHANNEL_CAPACITY); + + let supervisor = tokio::spawn(async move { + let worker_sender = output_sender.clone(); + let mut worker = tokio::task::spawn_blocking(move || normalize_events(input_receiver, worker_sender, config)); + let mut forward = Box::pin(async move { + tokio::pin!(source); + while let Some(event) = source.next().await { + if input_sender.send(event).await.is_err() { + break; + } + } + }); + + tokio::select! { + result = &mut worker => publish_worker_result(result, &output_sender).await, + () = output_sender.closed() => { + drop(forward); + let _ = worker.await; + } + () = &mut forward => { + drop(forward); + publish_worker_result(worker.await, &output_sender).await; + } + }; + }); + + NormalizedSession { + receiver: output_receiver, + supervisor: Some(supervisor), + } +} + +async fn publish_worker_result( + result: Result, tokio::task::JoinError>, + sender: &mpsc::Sender>, +) { + let error = match result { + Ok(Ok(())) => return, + Ok(Err(error)) => error.context("session normalization failed"), + Err(error) => anyhow::Error::new(error).context("normalizer worker failed"), + }; + let _ = sender.send(Err(error)).await; +} + +fn normalize_events( + mut receiver: mpsc::Receiver>, + sender: mpsc::Sender>, + config: SessionConfig, +) -> anyhow::Result<()> { + let mut phase = SessionPhase::AwaitClip; + let mut next_segment_sequence = 0; + + while let Some(event) = receiver.blocking_recv() { + match event.context("recording source failed")? { + RecordingEvent::ClipStarted { sequence, start_at } => { + anyhow::ensure!( + matches!(phase, SessionPhase::AwaitClip), + "clip {sequence} started before the previous clip ended" + ); + phase = SessionPhase::InClip(Box::new(ClipNormalizer::new( + sequence, + start_at, + sender.clone(), + config, + next_segment_sequence, + ))); + } + RecordingEvent::Bytes(bytes) => { + let SessionPhase::InClip(clip) = &mut phase else { + anyhow::bail!("recording bytes arrived outside a clip"); + }; + clip.push(&bytes)?; + } + RecordingEvent::CaughtUp => { + let SessionPhase::InClip(clip) = &mut phase else { + anyhow::bail!("caught-up arrived outside a clip"); + }; + clip.caught_up()?; + } + RecordingEvent::ClipEnded => { + let SessionPhase::InClip(current) = std::mem::replace(&mut phase, SessionPhase::AwaitClip) else { + anyhow::bail!("clip end arrived outside a clip"); + }; + next_segment_sequence = (*current).finish()?; + } + RecordingEvent::SessionEnded => { + anyhow::ensure!( + matches!(phase, SessionPhase::AwaitClip), + "session ended before the active clip ended" + ); + phase = SessionPhase::Ended; + break; + } + } + } + + anyhow::ensure!( + matches!(phase, SessionPhase::Ended), + "recording source ended before the session end event" + ); + Ok(()) +} + +enum SessionPhase { + AwaitClip, + InClip(Box), + Ended, +} + +#[derive(Clone, Copy)] +struct SourceVideo { + track: u64, + codec: VpxCodec, +} + +struct PendingFrame { + data: Vec, + timestamp: u64, + codec: VpxCodec, + key_frame: bool, +} + +enum ClipPhase { + History(HistoryPolicy), + Live, +} + +enum HistoryPolicy { + EmitAll, + KeepLatestGop(PendingGop), +} + +#[derive(Default)] +struct PendingGop { + frames: Vec, + bytes: usize, +} + +impl PendingGop { + fn push(&mut self, frame: PendingFrame) -> anyhow::Result<()> { + if frame.key_frame { + self.frames.clear(); + self.bytes = 0; + } else if self.frames.is_empty() { + return Ok(()); + } + + let bytes = self + .bytes + .checked_add(frame.data.len()) + .context("pending GOP size overflow")?; + anyhow::ensure!(bytes <= MAX_PENDING_GOP_BYTES, "pending GOP exceeds the resource limit"); + self.frames.push(frame); + self.bytes = bytes; + Ok(()) + } +} + +struct ClipNormalizer { + clip_sequence: u64, + decoder: TagDecoder, + input: BytesMut, + source_video: Option, + cluster_timestamp: Option, + timestamp_scale_ns: u64, + phase: ClipPhase, + input_decoder: Option, + output_segment: Option, + next_segment_sequence: u64, + sender: mpsc::Sender>, + config: SessionConfig, +} + +impl ClipNormalizer { + fn new( + clip_sequence: u64, + start_at: StartAt, + sender: mpsc::Sender>, + config: SessionConfig, + next_segment_sequence: u64, + ) -> Self { + let targets = [ + MatroskaSpec::TrackEntry(Master::Start), + MatroskaSpec::BlockGroup(Master::Start), + ]; + let mut decoder = TagDecoder::new(&targets); + decoder.set_max_allowable_tag_size(Some(MAX_BUFFERED_TAG_BYTES)); + let phase = match start_at { + StartAt::Beginning => ClipPhase::History(HistoryPolicy::EmitAll), + StartAt::LiveEdge => ClipPhase::History(HistoryPolicy::KeepLatestGop(PendingGop::default())), + }; + Self { + clip_sequence, + decoder, + input: BytesMut::new(), + source_video: None, + cluster_timestamp: None, + timestamp_scale_ns: WEBM_TIMESTAMP_SCALE_NS, + phase, + input_decoder: None, + output_segment: None, + next_segment_sequence, + sender, + config, + } + } + + fn push(&mut self, bytes: &[u8]) -> anyhow::Result<()> { + for chunk in bytes.chunks(INPUT_CHUNK_SIZE) { + self.input.extend_from_slice(chunk); + while let Some(positioned) = self.decoder.decode(&mut self.input)? { + self.handle_tag(positioned.tag)?; + } + } + Ok(()) + } + + fn caught_up(&mut self) -> anyhow::Result<()> { + let history = match std::mem::replace(&mut self.phase, ClipPhase::Live) { + ClipPhase::History(history) => history, + ClipPhase::Live => anyhow::bail!("clip {} sent caught-up twice", self.clip_sequence), + }; + if let HistoryPolicy::KeepLatestGop(pending) = history { + for frame in pending.frames { + self.process_frame(frame)?; + } + } + Ok(()) + } + + fn finish(mut self) -> anyhow::Result { + anyhow::ensure!( + matches!(self.phase, ClipPhase::Live), + "clip {} ended before caught-up", + self.clip_sequence + ); + loop { + match self.decoder.decode_eof(&mut self.input)? { + Some(positioned) => self.handle_tag(positioned.tag)?, + None if self.decoder.is_finished() => break, + None => continue, + } + } + + if let Some(segment) = self.output_segment.take() { + segment.finish()?; + } + Ok(self.next_segment_sequence) + } + + fn handle_tag(&mut self, tag: MatroskaSpec) -> anyhow::Result<()> { + match tag { + MatroskaSpec::TrackEntry(Master::Full(children)) => { + if let Some(video) = parse_video_track(&children)? { + anyhow::ensure!(self.source_video.is_none(), "multiple video tracks are not supported"); + self.source_video = Some(video); + } + } + MatroskaSpec::TimestampScale(value) => self.timestamp_scale_ns = value, + MatroskaSpec::Cluster(Master::Start) => self.cluster_timestamp = None, + MatroskaSpec::Timestamp(value) => self.cluster_timestamp = Some(value), + tag @ (MatroskaSpec::SimpleBlock(_) | MatroskaSpec::BlockGroup(Master::Full(_))) => { + self.handle_block(tag)?; + } + _ => {} + } + Ok(()) + } + + fn handle_block(&mut self, tag: MatroskaSpec) -> anyhow::Result<()> { + let video = self + .source_video + .context("video track header not found before video data")?; + let block = VideoBlock::new(tag, self.cluster_timestamp, video.codec)?; + if block.track != video.track { + return Ok(()); + } + + let data = block.get_frame()?; + let key_frame = is_vpx_key_frame(&data, video.codec); + let timestamp = scale_timestamp(block.absolute_timestamp()?, self.timestamp_scale_ns)?; + let frame = PendingFrame { + data, + timestamp, + codec: video.codec, + key_frame, + }; + + match &mut self.phase { + ClipPhase::History(HistoryPolicy::KeepLatestGop(pending)) => pending.push(frame), + ClipPhase::History(HistoryPolicy::EmitAll) | ClipPhase::Live => self.process_frame(frame), + } + } + + fn process_frame(&mut self, frame: PendingFrame) -> anyhow::Result<()> { + let input_decoder = self + .input_decoder + .get_or_insert_with(|| InputDecoder::new(frame.codec, self.config.encoder_threads)); + let decoded = input_decoder.decode(&frame.data)?; + let dimensions = decoded.dimensions; + let size_changed = self + .output_segment + .as_ref() + .is_some_and(|current| current.dimensions != dimensions); + if size_changed { + self.output_segment + .take() + .context("missing active output segment")? + .finish()?; + } + + let new_segment = if self.output_segment.is_none() { + Some(SegmentInfo { + sequence: self.next_segment_sequence, + width: dimensions.width, + height: dimensions.height, + }) + } else { + None + }; + + if let Some(info) = new_segment { + self.output_segment = Some(OutputSegment::new(self.sender.clone(), info, self.config)?); + self.next_segment_sequence = self + .next_segment_sequence + .checked_add(1) + .context("segment sequence overflow")?; + } + self.output_segment + .as_mut() + .context("output segment is missing")? + .encode(&decoded.image, frame.timestamp)?; + Ok(()) + } +} + +fn parse_video_track(children: &[MatroskaSpec]) -> anyhow::Result> { + let is_video = children + .iter() + .find_map(|tag| match tag { + MatroskaSpec::TrackType(value) => Some(*value == 1), + _ => None, + }) + .unwrap_or(false); + + if !is_video { + return Ok(None); + } + + let track = children + .iter() + .find_map(|tag| match tag { + MatroskaSpec::TrackNumber(value) => Some(*value), + _ => None, + }) + .context("video track number is missing")?; + let codec_id = children + .iter() + .find_map(|tag| match tag { + MatroskaSpec::CodecID(value) => Some(value.as_str()), + _ => None, + }) + .context("video codec ID is missing")?; + let codec = match codec_id { + "V_VP8" | "vp8" => VpxCodec::VP8, + "V_VP9" | "vp9" => VpxCodec::VP9, + _ => anyhow::bail!("unsupported video codec: {codec_id}"), + }; + + Ok(Some(SourceVideo { track, codec })) +} + +fn scale_timestamp(value: u64, timestamp_scale_ns: u64) -> anyhow::Result { + let nanoseconds = u128::from(value) + .checked_mul(u128::from(timestamp_scale_ns)) + .context("video timestamp overflow")?; + u64::try_from(nanoseconds / u128::from(WEBM_TIMESTAMP_SCALE_NS)).context("video timestamp is too large") +} + +struct OutputSegment { + info: SegmentInfo, + dimensions: Dimensions, + origin_timestamp: Option, + previous_timestamp: Option, + cluster_timestamp: Option, + encoder: VpxEncoder, + writer: WebmWriter, +} + +impl OutputSegment { + fn new( + sender: mpsc::Sender>, + info: SegmentInfo, + config: SessionConfig, + ) -> anyhow::Result { + send_event(&sender, SegmentEvent::Begin(info))?; + + let encoder = VpxEncoder::builder() + .timebase_num(1) + .timebase_den(1000) + .codec(VpxCodec::VP8) + .width(info.width) + .height(info.height) + .threads(config.encoder_threads) + .bitrate(OUTPUT_BITRATE) + .preset(VpxEncoderPreset::BestPerformance) + .build()?; + let mut writer = WebmWriter::new(EventWriter { sender }); + write_header(&mut writer, info.width, info.height)?; + + Ok(Self { + info, + dimensions: Dimensions { + width: info.width, + height: info.height, + }, + origin_timestamp: None, + previous_timestamp: None, + cluster_timestamp: None, + encoder, + writer, + }) + } + + fn encode(&mut self, image: &VpxImage<'_>, timestamp: u64) -> anyhow::Result<()> { + let origin = *self.origin_timestamp.get_or_insert(timestamp); + let relative_timestamp = timestamp.saturating_sub(origin); + let duration = self + .previous_timestamp + .map_or(30, |previous| timestamp.saturating_sub(previous).max(1)); + self.previous_timestamp = Some(timestamp); + + let cluster_timestamp_expired = self.cluster_timestamp.is_some_and(|cluster_timestamp| { + relative_timestamp.saturating_sub(cluster_timestamp) > MAX_WEBM_BLOCK_TIMESTAMP + }); + let flags = if relative_timestamp == 0 || cluster_timestamp_expired { + VPX_EFLAG_FORCE_KF + } else { + 0 + }; + self.encoder.encode_frame( + image, + i64::try_from(relative_timestamp).context("relative timestamp is too large")?, + usize::try_from(duration).unwrap_or(usize::MAX), + flags, + )?; + self.write_encoded_frames() + } + + fn write_encoded_frames(&mut self) -> anyhow::Result<()> { + let frames = self + .encoder + .packet_iterator() + .filter_map(|packet| packet.frame()) + .map(|frame| { + let timestamp = u64::try_from(frame.pts()).context("encoder returned a negative timestamp")?; + let data = frame.buffer().context("encoder returned a frame without data")?; + Ok((timestamp, data)) + }) + .collect::>>()?; + + for (timestamp, data) in frames { + let is_key_frame = is_vpx_key_frame(&data, VpxCodec::VP8); + anyhow::ensure!( + self.cluster_timestamp.is_some() || is_key_frame, + "output segment does not begin with a key frame" + ); + if self.cluster_timestamp.is_none() || is_key_frame { + if self.cluster_timestamp.is_some() { + self.writer.write(&MatroskaSpec::Cluster(Master::End))?; + } + self.writer.write_advanced( + &MatroskaSpec::Cluster(Master::Start), + WriteOptions::is_unknown_sized_element(), + )?; + self.writer.write(&MatroskaSpec::Timestamp(timestamp))?; + self.cluster_timestamp = Some(timestamp); + } + + let cluster_timestamp = self.cluster_timestamp.context("output cluster timestamp is missing")?; + let block_timestamp = timestamp + .checked_sub(cluster_timestamp) + .context("output frame timestamp precedes its cluster")?; + let block_timestamp = + i16::try_from(block_timestamp).context("output cluster exceeds block timestamp range")?; + let block = SimpleBlock::new_uncheked(&data, 1, block_timestamp, false, None, false, is_key_frame); + self.writer.write(&MatroskaSpec::from(block))?; + } + + Ok(()) + } + + fn finish(mut self) -> anyhow::Result<()> { + self.encoder.flush()?; + self.write_encoded_frames()?; + if self.cluster_timestamp.is_some() { + self.writer.write(&MatroskaSpec::Cluster(Master::End))?; + } + let event_writer = self.writer.into_inner()?; + send_event(&event_writer.sender, SegmentEvent::End) + .with_context(|| format!("failed to finish segment {}", self.info.sequence)) + } +} + +fn write_header(writer: &mut WebmWriter, width: u32, height: u32) -> anyhow::Result<()> { + writer.write(&MatroskaSpec::Ebml(Master::Full(vec![ + MatroskaSpec::EbmlVersion(1), + MatroskaSpec::EbmlReadVersion(1), + MatroskaSpec::EbmlMaxIdLength(4), + MatroskaSpec::EbmlMaxSizeLength(8), + MatroskaSpec::DocType("webm".to_owned()), + MatroskaSpec::DocTypeVersion(4), + MatroskaSpec::DocTypeReadVersion(2), + ])))?; + writer.write_advanced( + &MatroskaSpec::Segment(Master::Start), + WriteOptions::is_unknown_sized_element(), + )?; + writer.write(&MatroskaSpec::Info(Master::Full(vec![ + MatroskaSpec::TimestampScale(WEBM_TIMESTAMP_SCALE_NS), + MatroskaSpec::MuxingApp("Devolutions Gateway".to_owned()), + MatroskaSpec::WritingApp("Devolutions Gateway".to_owned()), + ])))?; + writer.write(&MatroskaSpec::Tracks(Master::Full(vec![MatroskaSpec::TrackEntry( + Master::Full(vec![ + MatroskaSpec::TrackNumber(1), + MatroskaSpec::TrackUID(1), + MatroskaSpec::TrackType(1), + MatroskaSpec::FlagEnabled(1), + MatroskaSpec::FlagDefault(1), + MatroskaSpec::FlagLacing(0), + MatroskaSpec::CodecID("V_VP8".to_owned()), + MatroskaSpec::Video(Master::Full(vec![ + MatroskaSpec::PixelWidth(u64::from(width)), + MatroskaSpec::PixelHeight(u64::from(height)), + ])), + ]), + )])))?; + Ok(()) +} + +fn send_event(sender: &mpsc::Sender>, event: SegmentEvent) -> anyhow::Result<()> { + sender + .blocking_send(Ok(event)) + .map_err(|_| anyhow::anyhow!("segment event receiver closed")) +} + +struct EventWriter { + sender: mpsc::Sender>, +} + +impl Write for EventWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.sender + .blocking_send(Ok(SegmentEvent::Data(Bytes::copy_from_slice(buffer)))) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "segment event receiver closed"))?; + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/crates/video-streamer/src/protocol.rs b/crates/video-streamer/src/protocol.rs new file mode 100644 index 000000000..5eecbc0d3 --- /dev/null +++ b/crates/video-streamer/src/protocol.rs @@ -0,0 +1,290 @@ +use std::error::Error; +use std::pin::Pin; + +use anyhow::Context as _; +use bytes::{BufMut as _, Bytes, BytesMut}; +use futures_util::{Sink, SinkExt as _, Stream, StreamExt as _}; + +use crate::normalizer::{SegmentEvent, SegmentInfo}; + +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum ServerMessage { + Chunk(Bytes), + SegmentStarted(SegmentInfo), + Error(UserFriendlyError), + StreamEnded, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ClientMessage { + Start, + Pull, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum UserFriendlyError { + UnexpectedError, +} + +impl UserFriendlyError { + fn as_str(&self) -> &'static str { + match self { + Self::UnexpectedError => "UnexpectedError", + } + } +} + +pub(crate) async fn stream_segments(mut transport: T, segments: S) -> anyhow::Result<()> +where + T: Stream> + Sink + Unpin, + S: Stream>, + E: Error + Send + Sync + 'static, +{ + tokio::pin!(segments); + let mut expected = ClientMessage::Start; + let mut segment_state = SegmentState::AwaitingBegin; + + loop { + let Some(message) = transport.next().await else { + return Ok(()); + }; + let message = message + .map_err(anyhow::Error::new) + .context("read client stream message")?; + let message = match decode_client_message(&message) { + Ok(message) if message == expected => message, + Ok(_) | Err(_) => { + let _ = + send_server_message(&mut transport, ServerMessage::Error(UserFriendlyError::UnexpectedError)).await; + anyhow::bail!("invalid client stream state"); + } + }; + + let response = match wait_for_response(&mut transport, segments.as_mut(), &mut segment_state).await { + Ok(Some(response)) => response, + Ok(None) => return Ok(()), + Err(error) => { + let _ = + send_server_message(&mut transport, ServerMessage::Error(UserFriendlyError::UnexpectedError)).await; + return Err(error); + } + }; + + let ended = response == ServerMessage::StreamEnded; + send_server_message(&mut transport, response).await?; + if ended { + return Ok(()); + } + + expected = match message { + ClientMessage::Start | ClientMessage::Pull => ClientMessage::Pull, + }; + } +} + +async fn send_server_message(transport: &mut T, message: ServerMessage) -> anyhow::Result<()> +where + T: Sink + Unpin, + E: Error + Send + Sync + 'static, +{ + transport + .send(encode_server_message(message)) + .await + .map_err(anyhow::Error::new) + .context("write server stream message") +} + +fn decode_client_message(message: &[u8]) -> anyhow::Result { + match message { + [0] => Ok(ClientMessage::Start), + [1] => Ok(ClientMessage::Pull), + _ => anyhow::bail!("invalid client message"), + } +} + +fn encode_server_message(message: ServerMessage) -> Bytes { + let mut encoded = BytesMut::new(); + match message { + ServerMessage::Chunk(chunk) => { + encoded.reserve(1 + chunk.len()); + encoded.put_u8(0); + encoded.put(chunk); + } + ServerMessage::SegmentStarted(info) => { + encoded.put_u8(1); + let json = format!( + "{{\"codec\":\"vp8\",\"sequence\":{},\"width\":{},\"height\":{}}}", + info.sequence, info.width, info.height + ); + encoded.put(json.as_bytes()); + } + ServerMessage::Error(error) => { + encoded.put_u8(2); + let json = format!("{{\"error\":\"{}\"}}", error.as_str()); + encoded.put(json.as_bytes()); + } + ServerMessage::StreamEnded => encoded.put_u8(3), + } + encoded.freeze() +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SegmentState { + AwaitingBegin, + Streaming, +} + +async fn wait_for_response( + transport: &mut T, + segments: Pin<&mut S>, + state: &mut SegmentState, +) -> anyhow::Result> +where + T: Stream> + Unpin, + S: Stream>, + E: Error + Send + Sync + 'static, +{ + tokio::select! { + response = next_segment_message(segments, state) => response, + message = transport.next() => match message { + None => Ok(None), + Some(Ok(_)) => anyhow::bail!("client sent another request before receiving a response"), + Some(Err(error)) => Err(anyhow::Error::new(error).context("read client stream message")), + }, + } +} + +async fn next_segment_message( + mut segments: Pin<&mut S>, + state: &mut SegmentState, +) -> anyhow::Result> +where + S: Stream>, +{ + loop { + let Some(event) = segments.as_mut().next().await else { + anyhow::ensure!( + *state == SegmentState::AwaitingBegin, + "segment stream ended inside a segment" + ); + return Ok(Some(ServerMessage::StreamEnded)); + }; + + match event? { + SegmentEvent::Begin(info) => { + anyhow::ensure!( + *state == SegmentState::AwaitingBegin, + "segment began before the previous segment ended" + ); + *state = SegmentState::Streaming; + return Ok(Some(ServerMessage::SegmentStarted(info))); + } + SegmentEvent::Data(data) => { + anyhow::ensure!( + *state == SegmentState::Streaming, + "segment data arrived outside a segment" + ); + return Ok(Some(ServerMessage::Chunk(data))); + } + SegmentEvent::End => { + anyhow::ensure!(*state == SegmentState::Streaming, "segment ended outside a segment"); + *state = SegmentState::AwaitingBegin; + } + } + } +} + +#[cfg(test)] +mod tests { + use futures_util::stream; + + use super::*; + + #[test] + fn protocol_codes_are_stable() { + assert_eq!( + encode_server_message(ServerMessage::SegmentStarted(SegmentInfo { + sequence: 7, + width: 1920, + height: 1080, + })), + Bytes::from_static(b"\x01{\"codec\":\"vp8\",\"sequence\":7,\"width\":1920,\"height\":1080}") + ); + assert_eq!( + encode_server_message(ServerMessage::Chunk(Bytes::from_static(b"webm"))), + Bytes::from_static(b"\x00webm") + ); + assert_eq!( + encode_server_message(ServerMessage::StreamEnded), + Bytes::from_static(b"\x03") + ); + } + + #[test] + fn client_messages_require_one_complete_transport_message() { + assert_eq!( + decode_client_message(b"\x00").expect("decode start"), + ClientMessage::Start + ); + assert_eq!( + decode_client_message(b"\x01").expect("decode pull"), + ClientMessage::Pull + ); + assert!(decode_client_message(b"\x00\x01").is_err()); + assert!(decode_client_message(b"").is_err()); + } + + #[tokio::test] + async fn segment_end_is_implicit_on_the_wire() { + let events = [ + Ok(SegmentEvent::Begin(SegmentInfo { + sequence: 0, + width: 640, + height: 480, + })), + Ok(SegmentEvent::Data(Bytes::from_static(b"first"))), + Ok(SegmentEvent::End), + Ok(SegmentEvent::Begin(SegmentInfo { + sequence: 1, + width: 800, + height: 600, + })), + Ok(SegmentEvent::Data(Bytes::from_static(b"second"))), + Ok(SegmentEvent::End), + ]; + let segments = stream::iter(events); + tokio::pin!(segments); + let mut state = SegmentState::AwaitingBegin; + + assert!(matches!( + next_segment_message(segments.as_mut(), &mut state) + .await + .expect("first begin"), + Some(ServerMessage::SegmentStarted(SegmentInfo { sequence: 0, .. })) + )); + assert_eq!( + next_segment_message(segments.as_mut(), &mut state) + .await + .expect("first data"), + Some(ServerMessage::Chunk(Bytes::from_static(b"first"))) + ); + assert!(matches!( + next_segment_message(segments.as_mut(), &mut state) + .await + .expect("second begin"), + Some(ServerMessage::SegmentStarted(SegmentInfo { sequence: 1, .. })) + )); + assert_eq!( + next_segment_message(segments.as_mut(), &mut state) + .await + .expect("second data"), + Some(ServerMessage::Chunk(Bytes::from_static(b"second"))) + ); + assert_eq!( + next_segment_message(segments.as_mut(), &mut state) + .await + .expect("stream end"), + Some(ServerMessage::StreamEnded) + ); + } +} diff --git a/crates/video-streamer/src/session.rs b/crates/video-streamer/src/session.rs new file mode 100644 index 000000000..a8711b36f --- /dev/null +++ b/crates/video-streamer/src/session.rs @@ -0,0 +1,55 @@ +use std::error::Error; + +use bytes::Bytes; +use futures_util::{Sink, Stream}; + +/// A structural event from one append-only recording session. +/// +/// A clip starts, receives zero or more byte events, catches up exactly once, receives more bytes, +/// and ends before another clip starts. +/// The session ends only when no clip is active. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum RecordingEvent { + ClipStarted { sequence: u64, start_at: StartAt }, + Bytes(Bytes), + CaughtUp, + ClipEnded, + SessionEnded, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StartAt { + Beginning, + LiveEdge, +} + +#[derive(Clone, Copy, Debug)] +pub struct SessionConfig { + pub encoder_threads: u32, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + encoder_threads: u32::try_from(num_cpus::get()).unwrap_or(1).max(1), + } + } +} + +/// Converts a recording session into fixed-size VP8 WebM segments over one pull-driven stream. +pub async fn stream_session(source: S, transport: T, config: SessionConfig) -> anyhow::Result<()> +where + S: Stream> + Send + 'static, + T: Stream> + Sink + Unpin, + E: Error + Send + Sync + 'static, +{ + let mut segments = crate::normalizer::normalize(source, config); + let stream_result = crate::protocol::stream_segments(transport, &mut segments).await; + let shutdown_result = segments.shutdown().await; + + match (stream_result, shutdown_result) { + (Err(error), _) => Err(error), + (Ok(()), Err(error)) => Err(error), + (Ok(()), Ok(())) => Ok(()), + } +} diff --git a/crates/video-streamer/src/streamer/block_tag.rs b/crates/video-streamer/src/streamer/block_tag.rs index 483dfc108..769657ed6 100644 --- a/crates/video-streamer/src/streamer/block_tag.rs +++ b/crates/video-streamer/src/streamer/block_tag.rs @@ -12,6 +12,7 @@ pub(crate) enum BlockTag { #[derive(Clone)] pub(crate) struct VideoBlock { + pub(crate) track: u64, pub(crate) cluster_timestamp: Option, pub(crate) timestamp: i16, pub(crate) is_key_frame: bool, @@ -22,6 +23,7 @@ impl fmt::Debug for VideoBlock { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("VideoBlock") .field("cluster_timestamp", &self.cluster_timestamp) + .field("track", &self.track) .field("timestamp", &self.timestamp) .field("is_key_frame", &self.is_key_frame) .field( @@ -58,6 +60,7 @@ impl VideoBlock { .any(|frame| is_vpx_key_frame(frame.data, codec)); Self { + track: block.track, cluster_timestamp, block_tag: BlockTag::BlockGroup(children), timestamp, @@ -67,6 +70,7 @@ impl VideoBlock { MatroskaSpec::SimpleBlock(data) => { let simple_block = SimpleBlock::try_from(&data)?; Self { + track: simple_block.track, cluster_timestamp, timestamp: simple_block.timestamp, is_key_frame: simple_block.keyframe, @@ -80,11 +84,13 @@ impl VideoBlock { } pub(crate) fn absolute_timestamp(&self) -> anyhow::Result { - let timestamp = u64::try_from(self.timestamp)?; - Ok(self + let cluster_timestamp = self .cluster_timestamp - .with_context(|| format!("Cluster timestamp not found for timestamp: {}", self.timestamp))? - + timestamp) + .with_context(|| format!("Cluster timestamp not found for timestamp: {}", self.timestamp))?; + let timestamp = i64::try_from(cluster_timestamp)? + .checked_add(i64::from(self.timestamp)) + .context("block timestamp overflow")?; + u64::try_from(timestamp).context("negative absolute block timestamp") } // We only handle non-lacing frames for now @@ -120,7 +126,7 @@ impl VideoBlock { } }; - assert!(frame.len() == 1); + anyhow::ensure!(frame.len() == 1, "laced video blocks are not supported"); Ok(frame[0].clone()) } }