This commit is contained in:
2026-04-09 18:19:03 -03:00
parent 5921cd6562
commit 5b467ffba8
18 changed files with 1793 additions and 43 deletions

View File

@@ -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
View 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
View 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
View File

@@ -0,0 +1,3 @@
pub mod capture;
pub mod encoder;
pub mod pipeline;

View File

@@ -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)

View 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(())
}