Refactor pipeline

This commit is contained in:
2024-09-03 14:09:30 +01:00
parent 65d8964632
commit 2c7d2dc9d1
27 changed files with 1465 additions and 1349 deletions

View File

@ -1,27 +1,27 @@
use std::collections::{HashMap, HashSet, VecDeque};
use std::collections::{HashSet, VecDeque};
use std::fmt::Display;
use std::ptr;
use anyhow::Error;
use ffmpeg_sys_next::{
av_dump_format, av_interleaved_write_frame, av_opt_set, av_packet_clone, av_packet_copy_props,
avcodec_parameters_from_context, avformat_alloc_output_context2, avformat_free_context,
avformat_write_header, AVFormatContext, AVPacket,
av_dump_format, av_interleaved_write_frame, av_opt_set, avcodec_parameters_copy,
avcodec_parameters_from_context, avformat_alloc_output_context2, avformat_free_context, avformat_write_header, AVFormatContext, AVPacket, AVStream,
};
use itertools::Itertools;
use log::info;
use tokio::sync::mpsc::UnboundedReceiver;
use uuid::Uuid;
use crate::egress::{EgressConfig, map_variants_to_streams};
use crate::egress::{map_variants_to_streams, EgressConfig};
use crate::pipeline::{AVPacketSource, PipelinePayload, PipelineProcessor};
use crate::return_ffmpeg_error;
use crate::utils::get_ffmpeg_error_msg;
use crate::variant::{VariantStream, VariantStreamType};
use crate::variant::{find_stream, StreamMapping, VariantStream};
pub struct HlsEgress {
id: Uuid,
config: EgressConfig,
variants: Vec<VariantStream>,
ctx: *mut AVFormatContext,
chan_in: UnboundedReceiver<PipelinePayload>,
stream_init: HashSet<Uuid>,
init: bool,
packet_buffer: VecDeque<PipelinePayload>,
@ -40,212 +40,279 @@ impl Drop for HlsEgress {
}
}
enum HlsMapEntry {
Video(usize),
Audio(usize),
Subtitle(usize),
}
impl Display for HlsMapEntry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
HlsMapEntry::Video(i) => write!(f, "v:{}", i),
HlsMapEntry::Audio(i) => write!(f, "a:{}", i),
HlsMapEntry::Subtitle(i) => write!(f, "s:{}", i),
}
}
}
struct HlsStream {
name: String,
entries: Vec<HlsMapEntry>,
}
impl Display for HlsStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{},name:{}", self.entries.iter().join(","), self.name)
}
}
impl HlsEgress {
pub fn new(
chan_in: UnboundedReceiver<PipelinePayload>,
id: Uuid,
config: EgressConfig,
) -> Self {
pub fn new(id: Uuid, config: EgressConfig, variants: Vec<VariantStream>) -> Self {
let filtered_vars: Vec<VariantStream> = config
.variants
.iter()
.filter_map(|x| variants.iter().find(|y| y.id() == *x))
.cloned()
.collect();
Self {
id,
config,
variants: filtered_vars,
ctx: ptr::null_mut(),
chan_in,
init: false,
stream_init: HashSet::new(),
packet_buffer: VecDeque::new(),
}
}
unsafe fn setup_muxer(&mut self) -> Result<(), Error> {
let mut ctx = ptr::null_mut();
pub(crate) fn setup_muxer(&mut self) -> Result<(), Error> {
unsafe {
let mut ctx = ptr::null_mut();
let base = format!("{}/{}", self.config.out_dir, self.id);
let base = format!("{}/{}", self.config.out_dir, self.id);
let ret = avformat_alloc_output_context2(
&mut ctx,
ptr::null(),
"hls\0".as_ptr() as *const libc::c_char,
format!("{}/%v/live.m3u8\0", base).as_ptr() as *const libc::c_char,
);
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
av_opt_set(
(*ctx).priv_data,
"hls_segment_filename\0".as_ptr() as *const libc::c_char,
format!("{}/%v/%05d.ts\0", base).as_ptr() as *const libc::c_char,
0,
);
av_opt_set(
(*ctx).priv_data,
"master_pl_name\0".as_ptr() as *const libc::c_char,
"live.m3u8\0".as_ptr() as *const libc::c_char,
0,
);
av_opt_set(
(*ctx).priv_data,
"master_pl_publish_rate\0".as_ptr() as *const libc::c_char,
"10\0".as_ptr() as *const libc::c_char,
0,
);
if let Some(first_video_track) = self.config.variants.iter().find_map(|v| {
if let VariantStream::Video(vv) = v {
Some(vv)
} else {
None
}
}) {
let ret = avformat_alloc_output_context2(
&mut ctx,
ptr::null(),
"hls\0".as_ptr() as *const libc::c_char,
format!("{}/%v/live.m3u8\0", base).as_ptr() as *const libc::c_char,
);
return_ffmpeg_error!(ret);
av_opt_set(
(*ctx).priv_data,
"hls_time\0".as_ptr() as *const libc::c_char,
format!("{}\0", first_video_track.keyframe_interval).as_ptr()
as *const libc::c_char,
"hls_segment_filename\0".as_ptr() as *const libc::c_char,
format!("{}/%v/%05d.ts\0", base).as_ptr() as *const libc::c_char,
0,
);
}
av_opt_set(
(*ctx).priv_data,
"hls_flags\0".as_ptr() as *const libc::c_char,
"delete_segments\0".as_ptr() as *const libc::c_char,
0,
);
av_opt_set(
(*ctx).priv_data,
"master_pl_name\0".as_ptr() as *const libc::c_char,
"live.m3u8\0".as_ptr() as *const libc::c_char,
0,
);
av_opt_set(
(*ctx).priv_data,
"master_pl_publish_rate\0".as_ptr() as *const libc::c_char,
"10\0".as_ptr() as *const libc::c_char,
0,
);
if let Some(first_video_track) = self.variants.iter().find_map(|v| {
if let VariantStream::Video(vv) = v {
Some(vv)
} else {
None
}
}) {
av_opt_set(
(*ctx).priv_data,
"hls_time\0".as_ptr() as *const libc::c_char,
format!(
"{}\0",
first_video_track.keyframe_interval / first_video_track.fps
)
.as_ptr() as *const libc::c_char,
0,
);
}
av_opt_set(
(*ctx).priv_data,
"hls_flags\0".as_ptr() as *const libc::c_char,
"delete_segments\0".as_ptr() as *const libc::c_char,
0,
);
map_variants_to_streams(ctx, &self.variants)?;
self.ctx = ctx;
Ok(())
}
}
unsafe fn setup_hls_mapping(&mut self) -> Result<(), Error> {
if self.ctx.is_null() {
return Err(Error::msg("Context not setup"));
}
// configure mapping
let mut stream_map: HashMap<usize, Vec<String>> = HashMap::new();
for var in &self.config.variants {
let cfg = match var {
VariantStream::Video(vx) => format!("v:{}", vx.dst_index),
VariantStream::Audio(ax) => format!("a:{}", ax.dst_index),
let mut stream_map = Vec::new();
for (g, vars) in &self
.variants
.iter()
.sorted_by(|a, b| a.group_id().cmp(&b.group_id()))
.group_by(|x| x.group_id())
{
let mut group = HlsStream {
name: format!("stream_{}", g),
entries: Vec::new(),
};
if let Some(out_stream) = stream_map.get_mut(&var.dst_index()) {
out_stream.push(cfg);
} else {
stream_map.insert(var.dst_index(), vec![cfg]);
for var in vars {
let n = Self::get_as_nth_stream_type(self.ctx, var);
match var {
VariantStream::Video(_) => group.entries.push(HlsMapEntry::Video(n)),
VariantStream::Audio(_) => group.entries.push(HlsMapEntry::Audio(n)),
VariantStream::CopyVideo(_) => group.entries.push(HlsMapEntry::Video(n)),
VariantStream::CopyAudio(_) => group.entries.push(HlsMapEntry::Audio(n)),
};
}
stream_map.push(group);
}
let stream_map = stream_map.values().map(|v| v.join(",")).join(" ");
let stream_map = stream_map.iter().join(" ");
info!("map_str={}", stream_map);
av_opt_set(
(*ctx).priv_data,
(*self.ctx).priv_data,
"var_stream_map\0".as_ptr() as *const libc::c_char,
format!("{}\0", stream_map).as_ptr() as *const libc::c_char,
0,
);
map_variants_to_streams(ctx, &mut self.config.variants)?;
self.ctx = ctx;
av_dump_format(self.ctx, 0, ptr::null(), 1);
Ok(())
}
unsafe fn process_pkt_internal(
unsafe fn process_av_packet_internal(
&mut self,
pkt: *mut AVPacket,
src: &AVPacketSource,
) -> Result<(), Error> {
let variant = match src {
AVPacketSource::Encoder(v) => self
.config
AVPacketSource::Encoder(v) => find_stream(&self.variants, v)?,
AVPacketSource::Demuxer(v) => self
.variants
.iter()
.find(|x| x.id() == *v)
.ok_or(Error::msg("Variant does not exist"))?,
AVPacketSource::Muxer(v) => self
.config
.variants
.iter()
.find(|x| x.id() == *v)
.ok_or(Error::msg("Variant does not exist"))?,
_ => return Err(Error::msg(format!("Cannot mux packet from {:?}", src))),
.find(|x| x.src_index() == (*(*v)).index as usize)
.ok_or(Error::msg("Demuxer packet didn't match any variant"))?,
};
(*pkt).stream_index = variant.dst_index() as libc::c_int;
//dump_pkt_info(pkt);
let ret = av_interleaved_write_frame(self.ctx, pkt);
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
return_ffmpeg_error!(ret);
Ok(())
}
fn process_payload_internal(&mut self, pkg: PipelinePayload) -> Result<(), Error> {
if let PipelinePayload::AvPacket(p, ref s) = pkg {
unsafe {
self.process_av_packet_internal(p, s)?;
}
}
Ok(())
}
unsafe fn process_pkt(
&mut self,
pkt: *mut AVPacket,
src: &AVPacketSource,
) -> Result<(), Error> {
let variant = match &src {
AVPacketSource::Encoder(v) => v,
AVPacketSource::Muxer(v) => v,
_ => return Err(Error::msg(format!("Cannot mux packet from {:?}", src))),
};
if !self.init {
let pkt_clone = av_packet_clone(pkt);
av_packet_copy_props(pkt_clone, pkt);
self.packet_buffer.push_back(PipelinePayload::AvPacket(
pkt_clone,
AVPacketSource::Muxer(*variant),
));
}
unsafe fn process_payload(&mut self, pkg: PipelinePayload) -> Result<(), Error> {
if !self.init && self.stream_init.len() == self.config.variants.len() {
let ret = avformat_write_header(self.ctx, ptr::null_mut());
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
self.setup_hls_mapping()?;
let ret = avformat_write_header(self.ctx, ptr::null_mut());
return_ffmpeg_error!(ret);
av_dump_format(self.ctx, 0, ptr::null(), 1);
self.init = true;
// push in pkts from buffer
// dequeue buffer
while let Some(pkt) = self.packet_buffer.pop_front() {
match pkt {
PipelinePayload::AvPacket(pkt, ref src) => {
self.process_pkt_internal(pkt, src)?;
}
_ => return Err(Error::msg("")),
}
self.process_payload_internal(pkt)?;
}
return Ok(());
} else if !self.init {
self.packet_buffer.push_back(pkg);
return Ok(());
}
self.process_pkt_internal(pkt, src)
self.process_payload_internal(pkg)
}
unsafe fn get_dst_stream(ctx: *const AVFormatContext, idx: usize) -> *mut AVStream {
for x in 0..(*ctx).nb_streams {
let stream = *(*ctx).streams.add(x as usize);
if (*stream).index as usize == idx {
return stream;
}
}
panic!("Stream index not found in output")
}
unsafe fn get_as_nth_stream_type(ctx: *const AVFormatContext, var: &VariantStream) -> usize {
let stream = Self::get_dst_stream(ctx, var.dst_index());
let mut ctr = 0;
for x in 0..(*ctx).nb_streams {
let stream_x = *(*ctx).streams.add(x as usize);
if (*(*stream).codecpar).codec_type == (*(*stream_x).codecpar).codec_type {
if (*stream_x).index == (*stream).index {
break;
}
ctr += 1;
}
}
ctr
}
}
impl PipelineProcessor for HlsEgress {
fn process(&mut self) -> Result<(), Error> {
while let Ok(pkg) = self.chan_in.try_recv() {
match pkg {
PipelinePayload::AvPacket(pkt, ref src) => unsafe {
self.process_pkt(pkt, src)?;
},
PipelinePayload::EncoderInfo(ref var, ctx) => unsafe {
if self.ctx.is_null() {
self.setup_muxer()?;
fn process(&mut self, pkg: PipelinePayload) -> Result<Vec<PipelinePayload>, Error> {
match pkg {
PipelinePayload::AvPacket(_, _) => unsafe {
self.process_payload(pkg)?;
},
PipelinePayload::SourceInfo(ref d) => unsafe {
for var in &self.variants {
match var {
VariantStream::CopyVideo(cv) => {
let src = *(*d.ctx).streams.add(cv.src_index);
let dst = Self::get_dst_stream(self.ctx, cv.dst_index);
let ret = avcodec_parameters_copy((*dst).codecpar, (*src).codecpar);
return_ffmpeg_error!(ret);
self.stream_init.insert(var.id());
}
VariantStream::CopyAudio(ca) => {
let src = *(*d.ctx).streams.add(ca.src_index);
let dst = Self::get_dst_stream(self.ctx, ca.dst_index);
let ret = avcodec_parameters_copy((*dst).codecpar, (*src).codecpar);
return_ffmpeg_error!(ret);
self.stream_init.insert(var.id());
}
_ => {}
}
}
},
PipelinePayload::EncoderInfo(ref var, ctx) => unsafe {
if let Some(my_var) = self.variants.iter().find(|x| x.id() == *var) {
if !self.stream_init.contains(var) {
let variant = self
.config
.variants
.iter()
.find(|x| x.id() == *var)
.ok_or(Error::msg("Variant does not exist"))?;
let out_stream = *(*self.ctx).streams.add(variant.dst_index());
let out_stream = Self::get_dst_stream(self.ctx, my_var.dst_index());
avcodec_parameters_from_context((*out_stream).codecpar, ctx);
self.stream_init.insert(*var);
}
},
_ => return Err(Error::msg(format!("Payload not supported: {:?}", pkg))),
}
}
},
_ => return Err(Error::msg(format!("Payload not supported: {:?}", pkg))),
}
Ok(())
// Muxer never returns anything
Ok(vec![])
}
}

View File

@ -1,11 +1,11 @@
use std::fmt::{Display, Formatter};
use std::ptr;
use crate::variant::{StreamMapping, VariantStream};
use anyhow::Error;
use ffmpeg_sys_next::{avformat_new_stream, AVFormatContext};
use serde::{Deserialize, Serialize};
use crate::variant::{VariantStream, VariantStreamType};
use uuid::Uuid;
pub mod hls;
pub mod http;
@ -15,7 +15,8 @@ pub mod recorder;
pub struct EgressConfig {
pub name: String,
pub out_dir: String,
pub variants: Vec<VariantStream>,
/// Which variants will be used in this muxer
pub variants: Vec<Uuid>,
}
impl Display for EgressConfig {
@ -33,31 +34,18 @@ impl Display for EgressConfig {
pub unsafe fn map_variants_to_streams(
ctx: *mut AVFormatContext,
variants: &mut Vec<VariantStream>,
variants: &Vec<VariantStream>,
) -> Result<(), Error> {
for var in variants {
match var {
VariantStream::Video(vs) => {
let stream = avformat_new_stream(ctx, ptr::null());
if stream.is_null() {
return Err(Error::msg("Failed to add stream to output"));
}
// overwrite dst_index to match output stream
vs.dst_index = (*stream).index as usize;
vs.to_stream(stream);
}
VariantStream::Audio(va) => {
let stream = avformat_new_stream(ctx, ptr::null());
if stream.is_null() {
return Err(Error::msg("Failed to add stream to output"));
}
// overwrite dst_index to match output stream
va.dst_index = (*stream).index as usize;
va.to_stream(stream);
}
let stream = avformat_new_stream(ctx, ptr::null());
if stream.is_null() {
return Err(Error::msg("Failed to add stream to output"));
}
// replace stream index value with variant dst_index
(*stream).index = var.dst_index() as libc::c_int;
var.to_stream(stream);
}
Ok(())
}
}

View File

@ -3,27 +3,27 @@ use std::{fs, ptr};
use anyhow::Error;
use ffmpeg_sys_next::{
av_dump_format, av_interleaved_write_frame, av_opt_set, avformat_alloc_output_context2, avformat_free_context,
avio_open2, AVFormatContext, AVPacket, AVIO_FLAG_WRITE,
av_dump_format, av_interleaved_write_frame, av_opt_set, avformat_alloc_output_context2,
avformat_free_context, avio_open2, AVFormatContext, AVPacket, AVIO_FLAG_WRITE,
};
use ffmpeg_sys_next::{
avcodec_parameters_from_context, avformat_write_header, AVFMT_GLOBALHEADER,
AV_CODEC_FLAG_GLOBAL_HEADER,
};
use log::info;
use tokio::sync::mpsc::UnboundedReceiver;
use uuid::Uuid;
use crate::egress::{map_variants_to_streams, EgressConfig};
use crate::pipeline::{PipelinePayload, PipelineProcessor};
use crate::return_ffmpeg_error;
use crate::utils::get_ffmpeg_error_msg;
use crate::variant::VariantStreamType;
use crate::variant::{find_stream, StreamMapping, VariantStream};
pub struct RecorderEgress {
id: Uuid,
config: EgressConfig,
variants: Vec<VariantStream>,
ctx: *mut AVFormatContext,
chan_in: UnboundedReceiver<PipelinePayload>,
stream_init: HashSet<Uuid>,
init: bool,
packet_buffer: VecDeque<PipelinePayload>,
@ -43,16 +43,18 @@ impl Drop for RecorderEgress {
}
impl RecorderEgress {
pub fn new(
chan_in: UnboundedReceiver<PipelinePayload>,
id: Uuid,
config: EgressConfig,
) -> Self {
pub fn new(id: Uuid, config: EgressConfig, variants: Vec<VariantStream>) -> Self {
let filtered_vars: Vec<VariantStream> = config
.variants
.iter()
.filter_map(|x| variants.iter().find(|y| y.id() == *x))
.cloned()
.collect();
Self {
id,
config,
variants: filtered_vars,
ctx: ptr::null_mut(),
chan_in,
stream_init: HashSet::new(),
init: false,
packet_buffer: VecDeque::new(),
@ -72,10 +74,8 @@ impl RecorderEgress {
ptr::null_mut(),
out_file.as_ptr() as *const libc::c_char,
);
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
map_variants_to_streams(ctx, &mut self.config.variants)?;
return_ffmpeg_error!(ret);
map_variants_to_streams(ctx, &self.variants)?;
if (*(*ctx).oformat).flags & AVFMT_GLOBALHEADER != 0 {
(*ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int;
@ -99,15 +99,12 @@ impl RecorderEgress {
ptr::null_mut(),
ptr::null_mut(),
);
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
return_ffmpeg_error!(ret);
av_dump_format(self.ctx, 0, ptr::null(), 1);
let ret = avformat_write_header(self.ctx, ptr::null_mut());
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
return_ffmpeg_error!(ret);
self.init = true;
Ok(true)
} else {
@ -118,54 +115,46 @@ impl RecorderEgress {
unsafe fn process_pkt(&mut self, pkt: *mut AVPacket) -> Result<(), Error> {
//dump_pkt_info(pkt);
let ret = av_interleaved_write_frame(self.ctx, pkt);
if ret < 0 {
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
}
return_ffmpeg_error!(ret);
Ok(())
}
}
impl PipelineProcessor for RecorderEgress {
fn process(&mut self) -> Result<(), Error> {
while let Ok(pkg) = self.chan_in.try_recv() {
match pkg {
PipelinePayload::AvPacket(pkt, ref src) => unsafe {
if self.open_muxer()? {
while let Some(pkt) = self.packet_buffer.pop_front() {
match pkt {
PipelinePayload::AvPacket(pkt, ref src) => {
self.process_pkt(pkt)?;
}
_ => return Err(Error::msg("")),
fn process(&mut self, pkg: PipelinePayload) -> Result<Vec<PipelinePayload>, Error> {
match pkg {
PipelinePayload::AvPacket(pkt, ref src) => unsafe {
if self.open_muxer()? {
while let Some(pkt) = self.packet_buffer.pop_front() {
match pkt {
PipelinePayload::AvPacket(pkt, ref src) => {
self.process_pkt(pkt)?;
}
_ => return Err(Error::msg("")),
}
self.process_pkt(pkt)?;
} else {
self.packet_buffer.push_back(pkg);
}
},
PipelinePayload::EncoderInfo(ref var, ctx) => unsafe {
if self.ctx.is_null() {
self.setup_muxer()?;
}
if !self.stream_init.contains(var) {
let my_var = self
.config
.variants
.iter()
.find(|x| x.id() == *var)
.ok_or(Error::msg("Variant does not exist"))?;
let out_stream = *(*self.ctx).streams.add(my_var.dst_index());
avcodec_parameters_from_context((*out_stream).codecpar, ctx);
(*(*out_stream).codecpar).codec_tag = 0;
self.process_pkt(pkt)?;
} else {
self.packet_buffer.push_back(pkg);
}
},
PipelinePayload::EncoderInfo(ref var, ctx) => unsafe {
if self.ctx.is_null() {
self.setup_muxer()?;
}
if !self.stream_init.contains(var) {
let my_var = find_stream(&self.variants, var)?;
let out_stream = *(*self.ctx).streams.add(my_var.dst_index());
avcodec_parameters_from_context((*out_stream).codecpar, ctx);
(*(*out_stream).codecpar).codec_tag = 0;
self.stream_init.insert(*var);
info!("Setup encoder info: {}", my_var);
}
},
_ => return Err(Error::msg("Payload not supported")),
}
self.stream_init.insert(*var);
info!("Setup encoder info: {}", my_var);
}
},
_ => return Err(Error::msg("Payload not supported")),
}
Ok(())
// Muxer never returns anything
Ok(vec![])
}
}