phase 2
This commit is contained in:
@@ -9,3 +9,4 @@ tokio = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
ffmpeg = { package = "ffmpeg-next", version = "8" }
|
||||
|
||||
128
media/client/src/capture.rs
Normal file
128
media/client/src/capture.rs
Normal file
@@ -0,0 +1,128 @@
|
||||
//! Screen capture via KMS/DRM using ffmpeg's kmsgrab device.
|
||||
//!
|
||||
//! Opens `/dev/dri/card0` via ffmpeg's kmsgrab input format,
|
||||
//! reads DRM/DMA-BUF frames at the compositor's refresh rate,
|
||||
//! and decimates to the target FPS.
|
||||
|
||||
use std::ffi::CString;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use tracing::info;
|
||||
|
||||
/// Configuration for KMS screen capture.
|
||||
pub struct CaptureConfig {
|
||||
/// DRM device path (e.g. "/dev/dri/card0")
|
||||
pub device: String,
|
||||
/// Target framerate (source frames are decimated to this)
|
||||
pub fps: u32,
|
||||
/// Output width (0 = native)
|
||||
pub width: u32,
|
||||
/// Output height (0 = native)
|
||||
pub height: u32,
|
||||
}
|
||||
|
||||
impl Default for CaptureConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
device: "/dev/dri/card0".into(),
|
||||
fps: 30,
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// KMS screen capture source.
|
||||
///
|
||||
/// Uses ffmpeg's kmsgrab input device. Requires DRM master or root.
|
||||
/// Outputs DRM_PRIME frames that can be mapped to VAAPI for zero-copy encode.
|
||||
pub struct KmsCapture {
|
||||
input_ctx: ffmpeg::format::context::Input,
|
||||
video_stream_idx: usize,
|
||||
decoder: ffmpeg::decoder::Video,
|
||||
}
|
||||
|
||||
impl KmsCapture {
|
||||
/// Open KMS capture with the given config.
|
||||
///
|
||||
/// This is a blocking call (ffmpeg device open). Run on a dedicated thread.
|
||||
pub fn open(config: &CaptureConfig) -> Result<Self> {
|
||||
ffmpeg::init().context("ffmpeg init")?;
|
||||
|
||||
// Set up the kmsgrab input device
|
||||
let mut opts = ffmpeg::Dictionary::new();
|
||||
opts.set("device", &config.device);
|
||||
opts.set("framerate", &config.fps.to_string());
|
||||
|
||||
let fmt_name = CString::new("kmsgrab").unwrap();
|
||||
let fmt_ptr = unsafe { ffmpeg::ffi::av_find_input_format(fmt_name.as_ptr()) };
|
||||
anyhow::ensure!(
|
||||
!fmt_ptr.is_null(),
|
||||
"kmsgrab input format not found — is ffmpeg built with --enable-libdrm?"
|
||||
);
|
||||
let fmt = unsafe { ffmpeg::format::Input::wrap(fmt_ptr as *mut _) };
|
||||
|
||||
let input_ctx = ffmpeg::format::open_with("-", &ffmpeg::Format::Input(fmt), opts)
|
||||
.context("failed to open kmsgrab device")?
|
||||
.input();
|
||||
|
||||
let video_stream = input_ctx
|
||||
.streams()
|
||||
.best(ffmpeg::media::Type::Video)
|
||||
.context("no video stream from kmsgrab")?;
|
||||
|
||||
let video_stream_idx = video_stream.index();
|
||||
|
||||
// Set up decoder with DRM hardware context
|
||||
let decoder_ctx = ffmpeg::codec::context::Context::from_parameters(video_stream.parameters())
|
||||
.context("decoder context from parameters")?;
|
||||
|
||||
let decoder = decoder_ctx.decoder().video().context("open video decoder")?;
|
||||
|
||||
info!(
|
||||
"KMS capture opened: {}x{} @ {}fps, stream idx={}",
|
||||
decoder.width(),
|
||||
decoder.height(),
|
||||
config.fps,
|
||||
video_stream_idx,
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
input_ctx,
|
||||
video_stream_idx,
|
||||
decoder,
|
||||
})
|
||||
}
|
||||
|
||||
/// Read the next raw (DRM_PRIME) frame from the capture device.
|
||||
///
|
||||
/// Blocking — call from a dedicated thread.
|
||||
/// Returns None at EOF (shouldn't happen for live capture).
|
||||
pub fn next_frame(&mut self) -> Result<Option<ffmpeg::frame::Video>> {
|
||||
loop {
|
||||
match self.input_ctx.packets().next() {
|
||||
Some((stream, packet)) => {
|
||||
if stream.index() != self.video_stream_idx {
|
||||
continue;
|
||||
}
|
||||
self.decoder.send_packet(&packet)?;
|
||||
let mut frame = ffmpeg::frame::Video::empty();
|
||||
match self.decoder.receive_frame(&mut frame) {
|
||||
Ok(()) => return Ok(Some(frame)),
|
||||
Err(ffmpeg::Error::Other { errno: ffmpeg::error::EAGAIN }) => continue,
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
}
|
||||
None => return Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn width(&self) -> u32 {
|
||||
self.decoder.width()
|
||||
}
|
||||
|
||||
pub fn height(&self) -> u32 {
|
||||
self.decoder.height()
|
||||
}
|
||||
}
|
||||
285
media/client/src/encoder.rs
Normal file
285
media/client/src/encoder.rs
Normal file
@@ -0,0 +1,285 @@
|
||||
//! VAAPI H.264 hardware encoding.
|
||||
//!
|
||||
//! Two-phase initialization:
|
||||
//! - Phase 1 (`VaapiEncoder::new`): store config, no ffmpeg calls.
|
||||
//! - Phase 2 (first call to `encode`): build filter graph using `hw_frames_ctx`
|
||||
//! from the first DRM_PRIME frame to wire up hwmap, then open the codec.
|
||||
//!
|
||||
//! The split exists because the filter graph needs the DRM hardware device
|
||||
//! context that kmsgrab attaches to each frame — that context isn't available
|
||||
//! until the first frame arrives, so the graph can't be validated at
|
||||
//! construction time.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use tracing::info;
|
||||
|
||||
/// Encoding configuration.
|
||||
pub struct EncoderConfig {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub fps: u32,
|
||||
/// Quantization parameter (lower = higher quality). Default: 20
|
||||
pub qp: u32,
|
||||
/// Keyframe interval in frames. Default: 30 (1 keyframe/sec at 30fps)
|
||||
pub gop_size: u32,
|
||||
}
|
||||
|
||||
impl Default for EncoderConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
fps: 30,
|
||||
qp: 20,
|
||||
gop_size: 30,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// VAAPI H.264 encoder with lazy filter graph initialization.
|
||||
pub struct VaapiEncoder {
|
||||
config: EncoderConfig,
|
||||
inner: Option<EncoderInner>,
|
||||
}
|
||||
|
||||
struct EncoderInner {
|
||||
encoder: ffmpeg::encoder::Video,
|
||||
filter_graph: ffmpeg::filter::Graph,
|
||||
frame_count: u64,
|
||||
time_base: ffmpeg::Rational,
|
||||
}
|
||||
|
||||
impl VaapiEncoder {
|
||||
/// Create an encoder from config. No ffmpeg resources are allocated yet.
|
||||
pub fn new(config: EncoderConfig) -> Self {
|
||||
Self { config, inner: None }
|
||||
}
|
||||
|
||||
/// Encode a frame. On the first call, initializes the filter graph and
|
||||
/// codec using the frame's `hw_frames_ctx`. Returns 0 or more packets.
|
||||
pub fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
|
||||
if self.inner.is_none() {
|
||||
self.inner = Some(EncoderInner::open(&self.config, frame).context("encoder init")?);
|
||||
}
|
||||
self.inner.as_mut().unwrap().encode(frame)
|
||||
}
|
||||
|
||||
/// Flush remaining packets out of the encoder.
|
||||
pub fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
|
||||
match self.inner.as_mut() {
|
||||
Some(inner) => inner.flush(),
|
||||
None => Ok(vec![]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl EncoderInner {
|
||||
fn open(config: &EncoderConfig, first_frame: &ffmpeg::frame::Video) -> Result<Self> {
|
||||
let time_base = ffmpeg::Rational::new(1, config.fps as i32);
|
||||
let filter_graph = Self::build_filter_graph(config, first_frame)?;
|
||||
|
||||
let codec = ffmpeg::encoder::find_by_name("h264_vaapi")
|
||||
.context("h264_vaapi encoder not found")?;
|
||||
|
||||
let mut encoder_ctx = ffmpeg::codec::context::Context::new_with_codec(codec)
|
||||
.encoder()
|
||||
.video()?;
|
||||
|
||||
encoder_ctx.set_width(config.width);
|
||||
encoder_ctx.set_height(config.height);
|
||||
encoder_ctx.set_time_base(time_base);
|
||||
encoder_ctx.set_frame_rate(Some(ffmpeg::Rational::new(config.fps as i32, 1)));
|
||||
encoder_ctx.set_gop(config.gop_size);
|
||||
encoder_ctx.set_max_b_frames(0);
|
||||
encoder_ctx.set_format(ffmpeg::format::Pixel::VAAPI);
|
||||
|
||||
let mut opts = ffmpeg::Dictionary::new();
|
||||
opts.set("qp", &config.qp.to_string());
|
||||
|
||||
let encoder = encoder_ctx
|
||||
.open_with(opts)
|
||||
.context("open h264_vaapi encoder")?;
|
||||
|
||||
info!(
|
||||
"VAAPI encoder opened: {}x{} @ {}fps, qp={}, gop={}",
|
||||
config.width, config.height, config.fps, config.qp, config.gop_size,
|
||||
);
|
||||
|
||||
Ok(Self { encoder, filter_graph, frame_count: 0, time_base })
|
||||
}
|
||||
|
||||
/// Build the filter graph: DRM_PRIME → hwmap(vaapi) → scale_vaapi → buffersink.
|
||||
///
|
||||
/// Mirrors the ffmpeg CLI approach:
|
||||
/// -init_hw_device drm=drm:/dev/dri/card0
|
||||
/// -init_hw_device vaapi=va@drm
|
||||
/// -filter_hw_device va
|
||||
///
|
||||
/// We create both devices explicitly and set the VAAPI device on the graph,
|
||||
/// then attach the frame's hw_frames_ctx to the buffersrc so hwmap can map
|
||||
/// DRM_PRIME frames to VAAPI surfaces.
|
||||
fn build_filter_graph(
|
||||
config: &EncoderConfig,
|
||||
first_frame: &ffmpeg::frame::Video,
|
||||
) -> Result<ffmpeg::filter::Graph> {
|
||||
let input_width = first_frame.width();
|
||||
let input_height = first_frame.height();
|
||||
|
||||
// Create DRM device, then derive VAAPI from it (like -init_hw_device vaapi=va@drm)
|
||||
let (drm_device, vaapi_device) = unsafe {
|
||||
let mut drm_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
|
||||
let dev_path = std::ffi::CString::new("/dev/dri/card0").unwrap();
|
||||
let ret = ffmpeg::ffi::av_hwdevice_ctx_create(
|
||||
&mut drm_ref,
|
||||
ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_DRM,
|
||||
dev_path.as_ptr(),
|
||||
std::ptr::null_mut(),
|
||||
0,
|
||||
);
|
||||
anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create(drm): {ret}");
|
||||
info!("DRM device created");
|
||||
|
||||
let mut vaapi_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
|
||||
let ret = ffmpeg::ffi::av_hwdevice_ctx_create_derived(
|
||||
&mut vaapi_ref,
|
||||
ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
|
||||
drm_ref,
|
||||
0,
|
||||
);
|
||||
anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create_derived(vaapi): {ret}");
|
||||
info!("VAAPI device derived from DRM");
|
||||
|
||||
(drm_ref, vaapi_ref)
|
||||
};
|
||||
|
||||
let mut graph = ffmpeg::filter::Graph::new();
|
||||
|
||||
let args = format!(
|
||||
"video_size={}x{}:pix_fmt={}:time_base=1/{}:pixel_aspect=1/1",
|
||||
input_width,
|
||||
input_height,
|
||||
ffmpeg::format::Pixel::DRM_PRIME as i32,
|
||||
config.fps,
|
||||
);
|
||||
|
||||
graph
|
||||
.add(&ffmpeg::filter::find("buffer").unwrap(), "in", &args)
|
||||
.context("add buffersrc")?;
|
||||
graph
|
||||
.add(&ffmpeg::filter::find("buffersink").unwrap(), "out", "")
|
||||
.context("add buffersink")?;
|
||||
|
||||
let filter_spec = format!(
|
||||
"hwmap=derive_device=vaapi,scale_vaapi=w={}:h={}:format=nv12",
|
||||
config.width, config.height,
|
||||
);
|
||||
|
||||
graph.output("in", 0)?.input("out", 0)?.parse(&filter_spec)?;
|
||||
|
||||
// Set the VAAPI device on filter contexts that need it
|
||||
// (equivalent to -filter_hw_device va in the CLI)
|
||||
unsafe {
|
||||
let graph_ptr = graph.as_mut_ptr();
|
||||
for i in 0..(*graph_ptr).nb_filters {
|
||||
let fctx = *(*graph_ptr).filters.add(i as usize);
|
||||
(*fctx).hw_device_ctx = ffmpeg::ffi::av_buffer_ref(vaapi_device);
|
||||
}
|
||||
}
|
||||
|
||||
// Attach hw_frames_ctx from the first DRM_PRIME frame to the buffersrc
|
||||
unsafe {
|
||||
let hw_frames_ctx = (*first_frame.as_ptr()).hw_frames_ctx;
|
||||
anyhow::ensure!(
|
||||
!hw_frames_ctx.is_null(),
|
||||
"first KMS frame has no hw_frames_ctx"
|
||||
);
|
||||
|
||||
let mut buffersrc = graph.get("in").unwrap();
|
||||
let par = ffmpeg::ffi::av_buffersrc_parameters_alloc();
|
||||
anyhow::ensure!(!par.is_null(), "av_buffersrc_parameters_alloc OOM");
|
||||
(*par).hw_frames_ctx = ffmpeg::ffi::av_buffer_ref(hw_frames_ctx);
|
||||
let ret = ffmpeg::ffi::av_buffersrc_parameters_set(buffersrc.as_mut_ptr(), par);
|
||||
ffmpeg::ffi::av_free(par as *mut _);
|
||||
anyhow::ensure!(ret >= 0, "av_buffersrc_parameters_set: {ret}");
|
||||
}
|
||||
|
||||
graph.validate().context("validate filter graph")?;
|
||||
info!("Filter graph ready: {filter_spec}");
|
||||
|
||||
// Keep device refs alive — they're ref-counted by the graph now,
|
||||
// but we drop our refs here.
|
||||
unsafe {
|
||||
ffmpeg::ffi::av_buffer_unref(&mut (drm_device as *mut _));
|
||||
ffmpeg::ffi::av_buffer_unref(&mut (vaapi_device as *mut _));
|
||||
}
|
||||
|
||||
Ok(graph)
|
||||
}
|
||||
|
||||
fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
|
||||
self.filter_graph
|
||||
.get("in")
|
||||
.unwrap()
|
||||
.source()
|
||||
.add(frame)
|
||||
.context("buffersrc add")?;
|
||||
|
||||
let mut packets = Vec::new();
|
||||
let mut filtered = ffmpeg::frame::Video::empty();
|
||||
|
||||
while self
|
||||
.filter_graph
|
||||
.get("out")
|
||||
.unwrap()
|
||||
.sink()
|
||||
.frame(&mut filtered)
|
||||
.is_ok()
|
||||
{
|
||||
filtered.set_pts(Some(self.frame_count as i64));
|
||||
self.frame_count += 1;
|
||||
|
||||
self.encoder.send_frame(&filtered)?;
|
||||
|
||||
let mut encoded = ffmpeg::Packet::empty();
|
||||
while self.encoder.receive_packet(&mut encoded).is_ok() {
|
||||
packets.push(EncodedPacket {
|
||||
data: encoded.data().unwrap_or(&[]).to_vec(),
|
||||
pts: encoded.pts().unwrap_or(0),
|
||||
dts: encoded.dts().unwrap_or(0),
|
||||
keyframe: encoded.is_key(),
|
||||
time_base_num: self.time_base.numerator() as u32,
|
||||
time_base_den: self.time_base.denominator() as u32,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(packets)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
|
||||
self.encoder.send_eof()?;
|
||||
let mut packets = Vec::new();
|
||||
let mut encoded = ffmpeg::Packet::empty();
|
||||
while self.encoder.receive_packet(&mut encoded).is_ok() {
|
||||
packets.push(EncodedPacket {
|
||||
data: encoded.data().unwrap_or(&[]).to_vec(),
|
||||
pts: encoded.pts().unwrap_or(0),
|
||||
dts: encoded.dts().unwrap_or(0),
|
||||
keyframe: encoded.is_key(),
|
||||
time_base_num: self.time_base.numerator() as u32,
|
||||
time_base_den: self.time_base.denominator() as u32,
|
||||
});
|
||||
}
|
||||
Ok(packets)
|
||||
}
|
||||
}
|
||||
|
||||
/// An encoded video packet ready for transport.
|
||||
pub struct EncodedPacket {
|
||||
pub data: Vec<u8>,
|
||||
pub pts: i64,
|
||||
pub dts: i64,
|
||||
pub keyframe: bool,
|
||||
pub time_base_num: u32,
|
||||
pub time_base_den: u32,
|
||||
}
|
||||
3
media/client/src/lib.rs
Normal file
3
media/client/src/lib.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
pub mod capture;
|
||||
pub mod encoder;
|
||||
pub mod pipeline;
|
||||
@@ -7,6 +7,10 @@ use tokio::io::{AsyncWriteExt, BufWriter};
|
||||
use tokio::net::TcpStream;
|
||||
use tracing::info;
|
||||
|
||||
use cht_client::capture::CaptureConfig;
|
||||
use cht_client::encoder::EncoderConfig;
|
||||
use cht_client::pipeline::Pipeline;
|
||||
|
||||
const DEFAULT_SERVER: &str = "mcrndeb:4444";
|
||||
|
||||
#[tokio::main]
|
||||
@@ -23,14 +27,17 @@ async fn main() -> Result<()> {
|
||||
|
||||
let mut writer = BufWriter::new(stream);
|
||||
|
||||
let capture_config = CaptureConfig::default();
|
||||
let encoder_config = EncoderConfig::default();
|
||||
|
||||
// Send session_start
|
||||
let session_start = ControlMessage::SessionStart {
|
||||
id: chrono_session_id(),
|
||||
id: session_id(),
|
||||
video: VideoParams {
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
width: encoder_config.width,
|
||||
height: encoder_config.height,
|
||||
codec: "h264".into(),
|
||||
fps: 30,
|
||||
fps: encoder_config.fps,
|
||||
},
|
||||
audio: AudioParams {
|
||||
sample_rate: 48000,
|
||||
@@ -39,60 +46,80 @@ async fn main() -> Result<()> {
|
||||
},
|
||||
};
|
||||
protocol::write_packet(&mut writer, &session_start.to_wire_packet()?).await?;
|
||||
writer.flush().await?;
|
||||
info!("Sent session_start");
|
||||
|
||||
// Send test packets (placeholder — will be replaced by real capture)
|
||||
let frame_interval_ns = 33_333_333u64; // ~30fps
|
||||
for i in 0u64..300 {
|
||||
let ts = i * frame_interval_ns;
|
||||
let keyframe = i % 30 == 0;
|
||||
// Start capture pipeline on a dedicated thread
|
||||
let (packet_tx, mut packet_rx) = tokio::sync::mpsc::channel(64);
|
||||
let mut pipeline = Pipeline::start(capture_config, encoder_config, packet_tx);
|
||||
|
||||
// Fake video packet
|
||||
let video = WirePacket {
|
||||
header: PacketHeader {
|
||||
packet_type: PacketType::Video,
|
||||
flags: if keyframe { FLAG_KEYFRAME } else { 0 },
|
||||
length: 1024,
|
||||
timestamp_ns: ts,
|
||||
},
|
||||
payload: vec![0u8; 1024],
|
||||
};
|
||||
protocol::write_packet(&mut writer, &video).await?;
|
||||
// Forward encoded packets to the server
|
||||
let mut video_count = 0u64;
|
||||
let mut keepalive_interval = tokio::time::interval(std::time::Duration::from_secs(5));
|
||||
|
||||
// Fake audio packet every 3 video frames
|
||||
if i % 3 == 0 {
|
||||
let audio = WirePacket {
|
||||
header: PacketHeader {
|
||||
packet_type: PacketType::Audio,
|
||||
flags: 0,
|
||||
length: 512,
|
||||
timestamp_ns: ts,
|
||||
},
|
||||
payload: vec![0u8; 512],
|
||||
};
|
||||
protocol::write_packet(&mut writer, &audio).await?;
|
||||
loop {
|
||||
tokio::select! {
|
||||
pkt = packet_rx.recv() => {
|
||||
match pkt {
|
||||
Some(encoded) => {
|
||||
let wire = WirePacket {
|
||||
header: PacketHeader {
|
||||
packet_type: PacketType::Video,
|
||||
flags: if encoded.keyframe { FLAG_KEYFRAME } else { 0 },
|
||||
length: encoded.data.len() as u32,
|
||||
timestamp_ns: pts_to_ns(
|
||||
encoded.pts,
|
||||
encoded.time_base_num,
|
||||
encoded.time_base_den,
|
||||
),
|
||||
},
|
||||
payload: encoded.data,
|
||||
};
|
||||
protocol::write_packet(&mut writer, &wire).await?;
|
||||
video_count += 1;
|
||||
|
||||
if video_count % 300 == 1 {
|
||||
info!("Sent {video_count} video packets");
|
||||
writer.flush().await?;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
info!("Pipeline channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = keepalive_interval.tick() => {
|
||||
let keepalive = ControlMessage::Keepalive;
|
||||
protocol::write_packet(&mut writer, &keepalive.to_wire_packet()?).await?;
|
||||
writer.flush().await?;
|
||||
}
|
||||
_ = tokio::signal::ctrl_c() => {
|
||||
info!("Ctrl+C received, stopping...");
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Keepalive every 150 frames (~5s)
|
||||
if i % 150 == 0 && i > 0 {
|
||||
let keepalive = ControlMessage::Keepalive;
|
||||
protocol::write_packet(&mut writer, &keepalive.to_wire_packet()?).await?;
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_nanos(frame_interval_ns)).await;
|
||||
}
|
||||
|
||||
// Send session_stop and flush
|
||||
pipeline.stop();
|
||||
|
||||
let stop = ControlMessage::SessionStop;
|
||||
protocol::write_packet(&mut writer, &stop.to_wire_packet()?).await?;
|
||||
writer.flush().await?;
|
||||
writer.shutdown().await?;
|
||||
info!("Sent session_stop, done");
|
||||
info!("Sent session_stop, {video_count} video packets total");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn chrono_session_id() -> String {
|
||||
fn pts_to_ns(pts: i64, tb_num: u32, tb_den: u32) -> u64 {
|
||||
if tb_den == 0 {
|
||||
return 0;
|
||||
}
|
||||
((pts as u128 * tb_num as u128 * 1_000_000_000) / tb_den as u128) as u64
|
||||
}
|
||||
|
||||
fn session_id() -> String {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
|
||||
106
media/client/src/pipeline.rs
Normal file
106
media/client/src/pipeline.rs
Normal file
@@ -0,0 +1,106 @@
|
||||
//! Capture pipeline: ties capture → encode → transport on a dedicated thread.
|
||||
//!
|
||||
//! This is the main loop that runs on a blocking thread, reading frames
|
||||
//! from KMS, encoding with VAAPI, and sending encoded packets through a channel.
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::capture::{CaptureConfig, KmsCapture};
|
||||
use crate::encoder::{EncodedPacket, EncoderConfig, VaapiEncoder};
|
||||
|
||||
/// A running capture pipeline that produces encoded packets.
|
||||
pub struct Pipeline {
|
||||
thread: Option<std::thread::JoinHandle<()>>,
|
||||
stop: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Pipeline {
|
||||
/// Start the capture → encode pipeline on a dedicated thread.
|
||||
///
|
||||
/// Encoded packets are sent through `packet_tx`.
|
||||
/// The pipeline runs until `stop()` is called or an error occurs.
|
||||
pub fn start(
|
||||
capture_config: CaptureConfig,
|
||||
encoder_config: EncoderConfig,
|
||||
packet_tx: tokio::sync::mpsc::Sender<EncodedPacket>,
|
||||
) -> Self {
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let stop_clone = stop.clone();
|
||||
|
||||
let thread = std::thread::Builder::new()
|
||||
.name("capture-pipeline".into())
|
||||
.spawn(move || {
|
||||
if let Err(e) = run_pipeline(capture_config, encoder_config, packet_tx, stop_clone) {
|
||||
error!("Pipeline error: {e:#}");
|
||||
}
|
||||
})
|
||||
.expect("spawn capture thread");
|
||||
|
||||
Self {
|
||||
thread: Some(thread),
|
||||
stop,
|
||||
}
|
||||
}
|
||||
|
||||
/// Signal the pipeline to stop and wait for the thread to finish.
|
||||
pub fn stop(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
if let Some(handle) = self.thread.take() {
|
||||
let _ = handle.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Pipeline {
|
||||
fn drop(&mut self) {
|
||||
self.stop();
|
||||
}
|
||||
}
|
||||
|
||||
fn run_pipeline(
|
||||
capture_config: CaptureConfig,
|
||||
encoder_config: EncoderConfig,
|
||||
packet_tx: tokio::sync::mpsc::Sender<EncodedPacket>,
|
||||
stop: Arc<AtomicBool>,
|
||||
) -> Result<()> {
|
||||
info!("Opening KMS capture...");
|
||||
let mut capture = KmsCapture::open(&capture_config).context("open capture")?;
|
||||
|
||||
info!("Encoder ready (will initialize on first frame)");
|
||||
let mut encoder = VaapiEncoder::new(encoder_config);
|
||||
|
||||
info!("Pipeline running");
|
||||
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
let frame = match capture.next_frame()? {
|
||||
Some(f) => f,
|
||||
None => {
|
||||
warn!("Capture returned EOF");
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let packets = encoder.encode(&frame).context("encode frame")?;
|
||||
|
||||
for pkt in packets {
|
||||
if packet_tx.blocking_send(pkt).is_err() {
|
||||
info!("Packet channel closed, stopping pipeline");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush encoder
|
||||
info!("Flushing encoder...");
|
||||
let flush_pkts = encoder.flush()?;
|
||||
for pkt in flush_pkts {
|
||||
let _ = packet_tx.blocking_send(pkt);
|
||||
}
|
||||
|
||||
info!("Pipeline stopped");
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user