Configurable encoder pipeline

src/demux/mod.rs

@@ -1,14 +1,19 @@
-use std::ffi::CStr;
-use std::pin::Pin;
 use std::ptr;
+use std::time::{Duration, SystemTime};
 
 use anyhow::Error;
-use async_trait::async_trait;
-use bytes::{BufMut, Bytes, BytesMut};
+use bytes::{Bytes, BytesMut};
+use ffmpeg_sys_next::AVMediaType::{AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_VIDEO};
 use ffmpeg_sys_next::*;
-use log::{debug, info, warn};
+use log::info;
+use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
+use tokio::time::Instant;
 
-use crate::pipeline::{PipelinePayload, PipelineStep};
+use crate::demux::info::{DemuxStreamInfo, StreamChannelType, StreamInfoChannel};
+use crate::pipeline::PipelinePayload;
+use crate::utils::get_ffmpeg_error_msg;
+
+pub mod info;
 
 ///
 /// Demuxer supports demuxing and decoding
@@ -20,8 +25,10 @@ use crate::pipeline::{PipelinePayload, PipelineStep};
 /// | Format | MPEG-TS                       |
 ///
 pub(crate) struct Demuxer {
-    buffer: BytesMut,
     ctx: *mut AVFormatContext,
+    chan_in: UnboundedReceiver<Bytes>,
+    chan_out: UnboundedSender<PipelinePayload>,
+    started: Instant,
 }
 
 unsafe impl Send for Demuxer {}
@@ -32,105 +39,141 @@ unsafe extern "C" fn read_data(
     buffer: *mut libc::c_uchar,
     size: libc::c_int,
 ) -> libc::c_int {
-    let muxer = opaque as *mut Demuxer;
-    let len = size.min((*muxer).buffer.len() as libc::c_int);
+    let chan = opaque as *mut UnboundedReceiver<Bytes>;
+    let mut data = (*chan).blocking_recv().expect("shit");
+    let buff_len = data.len();
+    let mut len = size.min(buff_len as libc::c_int);
+
     if len > 0 {
         memcpy(
             buffer as *mut libc::c_void,
-            (*muxer).buffer.as_ptr() as *const libc::c_void,
+            data.as_ptr() as *const libc::c_void,
             len as libc::c_ulonglong,
         );
-        _ = (*muxer).buffer.split_to(len as usize);
-        len
-    } else {
-        AVERROR_BUFFER_TOO_SMALL
     }
+    len
 }
 
 impl Demuxer {
-    const BUFFER_SIZE: usize = 1024 * 1024;
-    const INIT_BUFFER_THRESHOLD: usize = 2048;
-
-    pub fn new() -> Self {
+    pub fn new(
+        chan_in: UnboundedReceiver<Bytes>,
+        chan_out: UnboundedSender<PipelinePayload>,
+    ) -> Self {
         unsafe {
             let ps = avformat_alloc_context();
-            (*ps).probesize = Self::BUFFER_SIZE as i64;
             (*ps).flags |= AVFMT_FLAG_CUSTOM_IO;
+
             Self {
                 ctx: ps,
-                buffer: BytesMut::with_capacity(Self::BUFFER_SIZE),
+                chan_in,
+                chan_out,
+                started: Instant::now(),
             }
         }
     }
 
-    unsafe fn append_buffer(&mut self, bytes: &Bytes) {
-        self.buffer.extend_from_slice(bytes);
-    }
+    unsafe fn probe_input(&mut self) -> Result<DemuxStreamInfo, Error> {
+        let buf_ptr = ptr::from_mut(&mut self.chan_in) as *mut libc::c_void;
+        let pb = avio_alloc_context(
+            av_mallocz(4096) as *mut libc::c_uchar,
+            4096,
+            0,
+            buf_ptr,
+            Some(read_data),
+            None,
+            None,
+        );
 
-    unsafe fn probe_input(&mut self) -> Result<bool, Error> {
-        let size = self.buffer.len();
-        let score = (*self.ctx).probe_score;
-        if score == 0 && size >= Self::INIT_BUFFER_THRESHOLD {
-            let pb = avio_alloc_context(
-                av_mallocz(4096) as *mut libc::c_uchar,
-                4096,
-                0,
-                self as *const Self as *mut libc::c_void,
-                Some(read_data),
-                None,
-                None,
-            );
-
-            (*self.ctx).pb = pb;
-            let ret = avformat_open_input(
-                &mut self.ctx,
-                ptr::null_mut(),
-                ptr::null_mut(),
-                ptr::null_mut(),
-            );
-            if ret < 0 {
-                let msg = Self::get_ffmpeg_error_msg(ret);
-                return Err(Error::msg(msg));
-            }
-
-            if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 {
-                return Err(Error::msg("Could not find stream info"));
-            }
-
-            for x in 0..(*self.ctx).nb_streams {
-                av_dump_format(self.ctx, x as libc::c_int, ptr::null_mut(), 0);
-            }
-        }
-        Ok(score > 0)
-    }
-
-    unsafe fn decode_packet(&mut self) -> Result<Option<*mut AVPacket>, Error> {
-        let pkt: *mut AVPacket = av_packet_alloc();
-        av_init_packet(pkt);
-
-        let ret = av_read_frame(self.ctx, pkt);
-        if ret == AVERROR_BUFFER_TOO_SMALL {
-            return Ok(None);
-        }
+        (*self.ctx).pb = pb;
+        let ret = avformat_open_input(
+            &mut self.ctx,
+            ptr::null_mut(),
+            ptr::null_mut(),
+            ptr::null_mut(),
+        );
         if ret < 0 {
-            let msg = Self::get_ffmpeg_error_msg(ret);
+            let msg = get_ffmpeg_error_msg(ret);
             return Err(Error::msg(msg));
         }
-        Ok(Some(pkt))
+        if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 {
+            return Err(Error::msg("Could not find stream info"));
+        }
+        av_dump_format(self.ctx, 0, ptr::null_mut(), 0);
+
+        let mut channel_infos = vec![];
+        let video_stream_index =
+            av_find_best_stream(self.ctx, AVMEDIA_TYPE_VIDEO, -1, -1, ptr::null_mut(), 0) as usize;
+        if video_stream_index != AVERROR_STREAM_NOT_FOUND as usize {
+            let video_stream = *(*self.ctx).streams.add(video_stream_index);
+            let codec_copy = unsafe {
+                let ptr = avcodec_parameters_alloc();
+                avcodec_parameters_copy(ptr, (*video_stream).codecpar);
+                ptr
+            };
+            channel_infos.push(StreamInfoChannel {
+                index: video_stream_index,
+                channel_type: StreamChannelType::Video,
+                width: (*(*video_stream).codecpar).width as usize,
+                height: (*(*video_stream).codecpar).height as usize,
+                codec_params: codec_copy,
+            });
+        }
+
+        let audio_stream_index =
+            av_find_best_stream(self.ctx, AVMEDIA_TYPE_AUDIO, -1, -1, ptr::null_mut(), 0) as usize;
+        if audio_stream_index != AVERROR_STREAM_NOT_FOUND as usize {
+            let audio_stream = *(*self.ctx).streams.add(audio_stream_index);
+            let codec_copy = unsafe {
+                let ptr = avcodec_parameters_alloc();
+                avcodec_parameters_copy(ptr, (*audio_stream).codecpar);
+                ptr
+            };
+            channel_infos.push(StreamInfoChannel {
+                index: audio_stream_index,
+                channel_type: StreamChannelType::Audio,
+                width: (*(*audio_stream).codecpar).width as usize,
+                height: (*(*audio_stream).codecpar).height as usize,
+                codec_params: codec_copy,
+            });
+        }
+
+        let info = DemuxStreamInfo {
+            channels: channel_infos,
+        };
+        Ok(info)
     }
 
-    unsafe fn print_buffer_info(&mut self) {
-        let mut pb = (*self.ctx).pb;
-        let offset = (*pb).pos;
-        let remaining = (*pb).buffer_size as i64 - (*pb).pos;
-        info!("offset={}, remaining={}", offset, remaining);
+    unsafe fn get_packet(&mut self) -> Result<(), Error> {
+        let pkt: *mut AVPacket = av_packet_alloc();
+        let ret = av_read_frame(self.ctx, pkt);
+        if ret == AVERROR_EOF {
+            // reset EOF flag, stream never ends
+            (*(*self.ctx).pb).eof_reached = 0;
+            return Ok(());
+        }
+        if ret < 0 {
+            let msg = get_ffmpeg_error_msg(ret);
+            return Err(Error::msg(msg));
+        }
+        let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize);
+        (*pkt).time_base = (*stream).time_base;
+        (*pkt).opaque = stream as *mut libc::c_void;
+
+        self.chan_out.send(PipelinePayload::AvPacket(pkt))?;
+        Ok(())
     }
 
-    fn get_ffmpeg_error_msg(ret: libc::c_int) -> String {
+    pub fn process(&mut self) -> Result<Option<DemuxStreamInfo>, Error> {
         unsafe {
-            let mut buf: [libc::c_char; 255] = [0; 255];
-            av_make_error_string(buf.as_mut_ptr(), 255, ret);
-            String::from(CStr::from_ptr(buf.as_ptr()).to_str().unwrap())
+            let score = (*self.ctx).probe_score;
+            if score < 30 {
+                if (Instant::now() - self.started) > Duration::from_secs(1) {
+                    return Ok(Some(self.probe_input()?));
+                }
+                return Ok(None);
+            }
+            self.get_packet()?;
+            Ok(None)
         }
     }
 }
@@ -143,32 +186,3 @@ impl Drop for Demuxer {
         }
     }
 }
-
-#[async_trait]
-impl PipelineStep for Demuxer {
-    fn name(&self) -> String {
-        "Demuxer".to_owned()
-    }
-
-    async fn process(&mut self, pkg: PipelinePayload) -> Result<PipelinePayload, Error> {
-        match pkg {
-            PipelinePayload::Bytes(ref bb) => unsafe {
-                self.append_buffer(bb);
-                if !self.probe_input()? {
-                    return Ok(PipelinePayload::Empty);
-                }
-                match self.decode_packet() {
-                    Ok(pkt) => match pkt {
-                        Some(pkt) => Ok(PipelinePayload::AvPacket(pkt)),
-                        None => Ok(PipelinePayload::Empty),
-                    },
-                    Err(e) => {
-                        warn!("{}", e);
-                        Ok(PipelinePayload::Empty)
-                    }
-                }
-            },
-            _ => return Err(Error::msg("Wrong pkg format")),
-        }
-    }
-}