fix: rtmp ingest

fix: idle placeholder stream

crates/core/src/ingress/mod.rs

@@ -130,14 +130,13 @@ impl BufferedReader {
         }
     }
 
-    /// Read data from buffer, filling the entire output buffer before returning
+    /// Read data from buffer
     pub fn read_buffered(&mut self, buf: &mut [u8]) -> usize {
-        if self.buf.len() >= buf.len() {
-            let drain = self.buf.drain(..buf.len());
-            buf.copy_from_slice(drain.as_slice());
-            buf.len()
-        } else {
-            0
+        let to_drain = buf.len().min(self.buf.len());
+        if to_drain > 0 {
+            let drain = self.buf.drain(..to_drain);
+            buf[..to_drain].copy_from_slice(drain.as_slice());
         }
+        to_drain
     }
 }

crates/core/src/ingress/rtmp.rs

@@ -1,7 +1,8 @@
 use crate::ingress::{BufferedReader, ConnectionInfo};
 use crate::overseer::Overseer;
 use crate::pipeline::runner::PipelineRunner;
-use anyhow::{bail, Result};
+use anyhow::{anyhow, bail, Result};
+use bytes::{Bytes, BytesMut};
 use log::{error, info};
 use rml_rtmp::handshake::{Handshake, HandshakeProcessResult, PeerType};
 use rml_rtmp::sessions::{
@@ -16,6 +17,8 @@ use tokio::net::TcpListener;
 use tokio::runtime::Handle;
 use tokio::time::Instant;
 use uuid::Uuid;
+use xflv::errors::FlvMuxerError;
+use xflv::muxer::FlvMuxer;
 
 const MAX_MEDIA_BUFFER_SIZE: usize = 10 * 1024 * 1024; // 10MB limit
 
@@ -27,8 +30,8 @@ struct RtmpClient {
     buffer: BufferedReader,
     session: ServerSession,
     msg_queue: VecDeque<ServerSessionResult>,
-    reader_buf: [u8; 4096],
     pub published_stream: Option<RtmpPublishedStream>,
+    muxer: FlvMuxer,
 }
 
 impl RtmpClient {
@@ -41,8 +44,8 @@ impl RtmpClient {
             session: ses,
             buffer: BufferedReader::new(1024 * 1024, MAX_MEDIA_BUFFER_SIZE, "RTMP"),
             msg_queue: VecDeque::from(res),
-            reader_buf: [0; 4096],
             published_stream: None,
+            muxer: FlvMuxer::new(),
         })
     }
 
@@ -89,27 +92,28 @@ impl RtmpClient {
         Ok(())
     }
 
-    fn read_data(&mut self) -> Result<()> {
-        let r = match self.socket.read(&mut self.reader_buf) {
+    fn read_data(&mut self) -> Result<Option<usize>> {
+        let mut buf = [0; 4096];
+        let r = match self.socket.read(&mut buf) {
             Ok(r) => r,
             Err(e) => {
                 return match e.kind() {
-                    ErrorKind::WouldBlock => Ok(()),
-                    ErrorKind::Interrupted => Ok(()),
+                    ErrorKind::WouldBlock => Ok(None),
+                    ErrorKind::Interrupted => Ok(None),
                     _ => Err(anyhow::Error::new(e)),
                 };
             }
         };
         if r == 0 {
-            bail!("EOF");
+            return Ok(Some(0));
         }
 
-        let mx = self.session.handle_input(&self.reader_buf[..r])?;
+        let mx = self.session.handle_input(&buf[..r])?;
         if !mx.is_empty() {
             self.msg_queue.extend(mx);
             self.process_msg_queue()?;
         }
-        Ok(())
+        Ok(Some(r))
     }
 
     fn process_msg_queue(&mut self) -> Result<()> {
@@ -128,6 +132,44 @@ impl RtmpClient {
         Ok(())
     }
 
+    fn write_flv_header(&mut self, metadata: &rml_rtmp::sessions::StreamMetadata) -> Result<()> {
+        let has_video = metadata.video_codec_id.is_some();
+        let has_audio = metadata.audio_codec_id.is_some();
+
+        self.muxer
+            .write_flv_header(has_audio, has_video)
+            .map_err(|e| anyhow!("failed to write flv header {}", e))?;
+        self.muxer
+            .write_previous_tag_size(0)
+            .map_err(|e| anyhow!("failed to write flv header {}", e))?;
+
+        // Extract data from the muxer
+        let data = self.muxer.writer.extract_current_bytes();
+        self.buffer.add_data(&data);
+
+        info!(
+            "FLV header written with audio: {}, video: {}",
+            has_audio, has_video
+        );
+        Ok(())
+    }
+
+    fn write_flv_tag(
+        &mut self,
+        tag_type: u8,
+        timestamp: u32,
+        data: Bytes,
+    ) -> Result<(), FlvMuxerError> {
+        let body_len = data.len();
+        self.muxer
+            .write_flv_tag_header(tag_type, body_len as _, timestamp)?;
+        self.muxer.write_flv_tag_body(BytesMut::from(data))?;
+        self.muxer.write_previous_tag_size((11 + body_len) as _)?;
+        let flv_data = self.muxer.writer.extract_current_bytes();
+        self.buffer.add_data(&flv_data);
+        Ok(())
+    }
+
     fn handle_event(&mut self, event: ServerSessionEvent) -> Result<()> {
         match event {
             ServerSessionEvent::ClientChunkSizeChanged { new_chunk_size } => {
@@ -169,12 +211,19 @@ impl RtmpClient {
                     "Metadata configured: {}/{} {:?}",
                     app_name, stream_key, metadata
                 );
+                self.write_flv_header(&metadata)?;
             }
-            ServerSessionEvent::AudioDataReceived { data, .. } => {
-                self.buffer.add_data(&data);
+            ServerSessionEvent::AudioDataReceived {
+                data, timestamp, ..
+            } => {
+                self.write_flv_tag(8, timestamp.value, data)
+                    .map_err(|e| anyhow!("failed to write flv tag: {}", e))?;
             }
-            ServerSessionEvent::VideoDataReceived { data, .. } => {
-                self.buffer.add_data(&data);
+            ServerSessionEvent::VideoDataReceived {
+                data, timestamp, ..
+            } => {
+                self.write_flv_tag(9, timestamp.value, data)
+                    .map_err(|e| anyhow!("failed to write flv tag: {}", e))?;
             }
             ServerSessionEvent::UnhandleableAmf0Command { .. } => {}
             ServerSessionEvent::PlayStreamRequested { request_id, .. } => {
@@ -195,10 +244,20 @@ impl Read for RtmpClient {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
         // Block until we have enough data to fill the buffer
         while self.buffer.buf.len() < buf.len() {
-            if let Err(e) = self.read_data() {
-                error!("Error reading data: {}", e);
-                return Ok(0);
-            };
+            match self.read_data() {
+                Ok(Some(0)) => {
+                    let r = self.buffer.read_buffered(buf);
+                    if r == 0 {
+                        return Err(std::io::Error::other(anyhow!("EOF")));
+                    }
+                    return Ok(r);
+                }
+                Err(e) => {
+                    error!("Error reading data: {}", e);
+                    return Ok(0);
+                }
+                _ => continue,
+            }
         }
 
         Ok(self.buffer.read_buffered(buf))
@@ -239,13 +298,15 @@ pub async fn listen(out_dir: String, addr: String, overseer: Arc<dyn Overseer>)
                     overseer,
                     info,
                     Box::new(cc),
-                    Some("flv".to_string()),
+                    None,
                 ) {
                     Ok(pl) => pl,
                     Err(e) => {
                         bail!("Failed to create PipelineRunner {}", e)
                     }
                 };
+                //pl.set_demuxer_format("flv");
+                //pl.set_demuxer_buffer_size(1024 * 64);
                 pl.run();
                 Ok(())
             })?;

crates/core/src/ingress/test.rs

@@ -4,7 +4,7 @@ use crate::overseer::Overseer;
 use anyhow::Result;
 use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P;
 use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVSampleFormat::AV_SAMPLE_FMT_FLTP;
-use ffmpeg_rs_raw::ffmpeg_sys_the_third::{av_frame_free, av_packet_free, AV_PROFILE_H264_MAIN};
+use ffmpeg_rs_raw::ffmpeg_sys_the_third::{av_frame_free, av_packet_free, AV_PROFILE_H264_MAIN, AVRational};
 use ffmpeg_rs_raw::{Encoder, Muxer};
 use log::info;
 use ringbuf::traits::{Observer, Split};
@@ -115,6 +115,8 @@ impl TestPatternSrc {
                 SAMPLE_RATE,
                 frame_size,
                 1,
+                AVRational { num: 1, den: VIDEO_FPS as i32 },
+                AVRational { num: 1, den: SAMPLE_RATE as i32 },
             )?,
             video_encoder,
             audio_encoder,