more something

src/egress/hls.rs

@@ -1,27 +1,22 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::{HashMap, HashSet, VecDeque};
 use std::fmt::{Display, Formatter};
 use std::mem::transmute;
 use std::ptr;
 
 use anyhow::Error;
+use ffmpeg_sys_next::{AV_CH_LAYOUT_STEREO, av_channel_layout_copy, av_dump_format, av_get_sample_fmt, av_interleaved_write_frame, av_opt_set, av_packet_clone, av_packet_copy_props, AVChannelLayout, AVChannelLayout__bindgen_ty_1, avcodec_find_encoder, avcodec_parameters_from_context, avcodec_parameters_to_context, AVCodecContext, avformat_alloc_output_context2, avformat_free_context, avformat_new_stream, avformat_write_header, AVFormatContext, AVPacket, AVRational};
 use ffmpeg_sys_next::AVChannelOrder::AV_CHANNEL_ORDER_NATIVE;
 use ffmpeg_sys_next::AVColorSpace::AVCOL_SPC_BT709;
 use ffmpeg_sys_next::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO};
 use ffmpeg_sys_next::AVPixelFormat::AV_PIX_FMT_YUV420P;
-use ffmpeg_sys_next::{
-    av_dump_format, av_get_sample_fmt, av_interleaved_write_frame, av_opt_set,
-    avcodec_find_encoder, avcodec_parameters_from_context, avformat_alloc_output_context2,
-    avformat_free_context, avformat_new_stream, avformat_write_header, AVChannelLayout,
-    AVChannelLayout__bindgen_ty_1, AVCodecContext, AVFormatContext, AVPacket, AVRational,
-    AV_CH_LAYOUT_STEREO,
-};
+use futures_util::SinkExt;
 use itertools::Itertools;
 use log::info;
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc::UnboundedReceiver;
 use uuid::Uuid;
 
-use crate::egress::{map_variants_to_streams, EgressConfig, update_pkt_for_muxer, get_pkt_variant};
+use crate::egress::{EgressConfig, get_pkt_variant, map_variants_to_streams, update_pkt_for_muxer};
 use crate::encode::dump_pkt_info;
 use crate::pipeline::{PipelinePayload, PipelineProcessor};
 use crate::utils::{get_ffmpeg_error_msg, id_ref_to_uuid};
@@ -32,6 +27,9 @@ pub struct HlsEgress {
     config: EgressConfig,
     ctx: *mut AVFormatContext,
     chan_in: UnboundedReceiver<PipelinePayload>,
+    stream_init: HashSet<usize>,
+    init: bool,
+    packet_buffer: VecDeque<PipelinePayload>,
 }
 
 unsafe impl Send for HlsEgress {}
@@ -58,6 +56,9 @@ impl HlsEgress {
             config,
             ctx: ptr::null_mut(),
             chan_in,
+            init: false,
+            stream_init: HashSet::new(),
+            packet_buffer: VecDeque::new(),
         }
     }
 
@@ -75,7 +76,6 @@ impl HlsEgress {
         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,
@@ -146,27 +146,62 @@ impl HlsEgress {
 
         map_variants_to_streams(ctx, &mut self.config.variants)?;
 
-        let ret = avformat_write_header(ctx, ptr::null_mut());
-        if ret < 0 {
-            return Err(Error::msg(get_ffmpeg_error_msg(ret)));
-        }
-
         self.ctx = ctx;
         Ok(())
     }
 
-    unsafe fn process_pkt(&mut self, pkt: *mut AVPacket) -> Result<(), Error> {
+    unsafe fn process_pkt_internal(&mut self, pkt: *mut AVPacket) -> Result<(), Error> {
         let variant = get_pkt_variant(&self.config.variants, pkt)?;
         update_pkt_for_muxer(self.ctx, pkt, &variant);
-
         //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)));
         }
-
         Ok(())
     }
+
+    unsafe fn process_pkt(&mut self, pkt: *mut AVPacket) -> Result<(), Error> {
+        let variant = get_pkt_variant(&self.config.variants, pkt)?;
+        if !self.stream_init.contains(&variant.dst_index()) {
+            let encoder_ctx = (*pkt).opaque as *mut AVCodecContext;
+            let out_stream = *(*self.ctx).streams.add(variant.dst_index());
+            avcodec_parameters_from_context((*out_stream).codecpar, encoder_ctx);
+            self.stream_init.insert(variant.dst_index());
+        }
+        if !self.init {
+            let pkt_clone = av_packet_clone(pkt);
+            av_packet_copy_props(pkt_clone, pkt);
+            self.packet_buffer.push_back(PipelinePayload::AvPacket(
+                "Buffered Muxer Packet".to_string(),
+                pkt_clone,
+            ));
+        }
+
+        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)));
+            }
+
+            av_dump_format(self.ctx, 0, ptr::null(), 1);
+            self.init = true;
+            // push in pkts from buffer
+            while let Some(pkt) = self.packet_buffer.pop_front() {
+                match pkt {
+                    PipelinePayload::AvPacket(_, pkt) => {
+                        self.process_pkt_internal(pkt)?;
+                    }
+                    _ => return Err(Error::msg("")),
+                }
+            }
+            return Ok(());
+        } else if !self.init {
+            return Ok(());
+        }
+
+        self.process_pkt_internal(pkt)
+    }
 }
 
 impl PipelineProcessor for HlsEgress {

src/egress/mod.rs

@@ -2,8 +2,7 @@ use std::fmt::{Display, Formatter};
 use std::ptr;
 
 use anyhow::Error;
-use ffmpeg_sys_next::{av_dump_format, avformat_new_stream, AVFormatContext, AVPacket};
-use log::info;
+use ffmpeg_sys_next::{av_packet_rescale_ts, avformat_new_stream, AVFormatContext, AVPacket};
 use serde::{Deserialize, Serialize};
 
 use crate::utils::id_ref_to_uuid;
@@ -12,6 +11,7 @@ use crate::variant::{VariantStream, VariantStreamType};
 pub mod hls;
 pub mod http;
 pub mod mpegts;
+pub mod recorder;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct EgressConfig {
@@ -61,11 +61,10 @@ pub unsafe fn map_variants_to_streams(
             }
         }
     }
-
-    av_dump_format(ctx, 0, ptr::null(), 1);
     Ok(())
 }
 
+/// Get variant of this packet
 pub unsafe fn get_pkt_variant(
     vars: &Vec<VariantStream>,
     pkt: *mut AVPacket,
@@ -81,6 +80,7 @@ pub unsafe fn get_pkt_variant(
     Ok(variant.unwrap())
 }
 
+/// Update packet stream index to match muxer stream
 pub unsafe fn update_pkt_for_muxer(
     ctx: *mut AVFormatContext,
     pkt: *mut AVPacket,
@@ -91,4 +91,7 @@ pub unsafe fn update_pkt_for_muxer(
     if idx != (*pkt).stream_index {
         (*pkt).stream_index = idx;
     }
+    // match stream timebase in muxer
+    av_packet_rescale_ts(pkt, var.time_base(), (*stream).time_base);
+    (*pkt).time_base = (*stream).time_base;
 }

src/egress/recorder.rs

@@ -0,0 +1,117 @@
+use std::{fs, ptr};
+use std::collections::HashSet;
+use std::fmt::Display;
+
+use anyhow::Error;
+use ffmpeg_sys_next::{av_dump_format, av_guess_format, av_interleaved_write_frame, av_strdup, avformat_alloc_context, avformat_alloc_output_context2, avformat_free_context, avformat_write_header, AVFormatContext, AVIO_FLAG_READ_WRITE, avio_flush, avio_open2, AVPacket};
+use tokio::sync::mpsc::UnboundedReceiver;
+use uuid::Uuid;
+
+use crate::egress::{EgressConfig, get_pkt_variant, map_variants_to_streams, update_pkt_for_muxer};
+use crate::pipeline::{PipelinePayload, PipelineProcessor};
+use crate::utils::get_ffmpeg_error_msg;
+
+pub struct RecorderEgress {
+    id: Uuid,
+    config: EgressConfig,
+    ctx: *mut AVFormatContext,
+    chan_in: UnboundedReceiver<PipelinePayload>,
+    stream_init: HashSet<i32>,
+}
+
+unsafe impl Send for RecorderEgress {}
+
+unsafe impl Sync for RecorderEgress {}
+
+impl Drop for RecorderEgress {
+    fn drop(&mut self) {
+        unsafe {
+            avformat_free_context(self.ctx);
+            self.ctx = ptr::null_mut();
+        }
+    }
+}
+
+impl RecorderEgress {
+    pub fn new(
+        chan_in: UnboundedReceiver<PipelinePayload>,
+        id: Uuid,
+        config: EgressConfig,
+    ) -> Self {
+        Self {
+            id,
+            config,
+            ctx: ptr::null_mut(),
+            chan_in,
+            stream_init: HashSet::new(),
+        }
+    }
+
+    unsafe fn setup_muxer(&mut self) -> Result<(), Error> {
+        let mut ctx = avformat_alloc_context();
+        if ctx.is_null() {
+            return Err(Error::msg("Failed to create muxer context"));
+        }
+        let base = format!("{}/{}", self.config.out_dir, self.id);
+
+        let out_file = format!("{}/recording.mkv\0", base).as_ptr() as *const libc::c_char;
+        fs::create_dir_all(base.clone())?;
+        let ret = avio_open2(
+            &mut (*ctx).pb,
+            out_file,
+            AVIO_FLAG_READ_WRITE,
+            ptr::null(),
+            ptr::null_mut(),
+        );
+        if ret < 0 {
+            return Err(Error::msg(get_ffmpeg_error_msg(ret)));
+        }
+        (*ctx).oformat = av_guess_format(
+            "matroska\0".as_ptr() as *const libc::c_char,
+            out_file,
+            ptr::null(),
+        );
+        if (*ctx).oformat.is_null() {
+            return Err(Error::msg("Output format not found"));
+        }
+        (*ctx).url = av_strdup(out_file);
+        map_variants_to_streams(ctx, &mut self.config.variants)?;
+
+        let ret = avformat_write_header(ctx, ptr::null_mut());
+        if ret < 0 {
+            return Err(Error::msg(get_ffmpeg_error_msg(ret)));
+        }
+        av_dump_format(ctx, 0, ptr::null(), 1);
+        self.ctx = ctx;
+        Ok(())
+    }
+
+    unsafe fn process_pkt(&mut self, pkt: *mut AVPacket) -> Result<(), Error> {
+        let variant = get_pkt_variant(&self.config.variants, pkt)?;
+        update_pkt_for_muxer(self.ctx, pkt, &variant);
+
+        let ret = av_interleaved_write_frame(self.ctx, pkt);
+        if ret < 0 {
+            return Err(Error::msg(get_ffmpeg_error_msg(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) => unsafe {
+                    if self.ctx.is_null() {
+                        self.setup_muxer()?;
+                    }
+                    self.process_pkt(pkt)?;
+                },
+                _ => return Err(Error::msg("Payload not supported")),
+            }
+        }
+        Ok(())
+    }
+}