From 67178a636fc3a05c8398c26a6915e361c7e834ab Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Fri, 18 Sep 2026 17:40:18 +0530 Subject: [PATCH] feat(stream): implement RTMP livestream network broadcaster and worker pipeline --- src/bin/scythe-daemon.rs | 193 ++++++++++++++++++++++++++++++++++++--- src/muxer.rs | 49 ++++++++-- 2 files changed, 224 insertions(+), 18 deletions(-) diff --git a/src/bin/scythe-daemon.rs b/src/bin/scythe-daemon.rs index c3f9379..7d8f6ef 100644 --- a/src/bin/scythe-daemon.rs +++ b/src/bin/scythe-daemon.rs @@ -14,6 +14,69 @@ use std::io::{Read, Write}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; +fn mask_stream_url(url: &str) -> String { + if let Some((base, key)) = url.rsplit_once('/') + && key.len() > 4 { + let masked = format!("{}****", &key[..2]); + return format!("{}/{}", base, masked); + } + url.to_string() +} + +fn spawn_stream_worker( + rtmp_url: String, + codec_ctx_ptr: usize, + audio_codec_ctx_ptr: Option, + stream_state: Arc, + stream_start: Arc, +) -> crossbeam_channel::Sender { + let (tx, rx) = crossbeam_channel::bounded::(120); + thread::spawn(move || { + let codec_ctx = codec_ctx_ptr as *mut ffmpeg_next::ffi::AVCodecContext; + let audio_codec_ctx = audio_codec_ctx_ptr.map(|p| p as *mut ffmpeg_next::ffi::AVCodecContext); + + println!("Connecting to RTMP destination: {}...", mask_stream_url(&rtmp_url)); + let muxer_res = unsafe { Muxer::new(&rtmp_url, codec_ctx, audio_codec_ctx) }; + let mut muxer = match muxer_res { + Ok(m) => { + println!("Connected to RTMP stream successfully."); + m + } + Err(e) => { + eprintln!("Failed to connect to RTMP destination: {}", e); + stream_state.store(false, Ordering::SeqCst); + stream_start.store(0, Ordering::SeqCst); + return; + } + }; + + let mut write_errors = 0; + while let Ok(pkt) = rx.recv() { + if !stream_state.load(Ordering::Relaxed) { + break; + } + if let Err(e) = muxer.write_packet(&pkt) { + write_errors += 1; + if write_errors % 30 == 1 { + eprintln!("Error writing packet to RTMP stream (errors: {}): {}", write_errors, e); + } + if write_errors >= 120 { + eprintln!("Too many consecutive write errors on RTMP stream. Aborting broadcast."); + break; + } + } else { + write_errors = 0; + } + } + + let _ = muxer.finalize(); + stream_state.store(false, Ordering::SeqCst); + stream_start.store(0, Ordering::SeqCst); + println!("RTMP stream broadcast finalized."); + }); + tx +} + fn ensure_wayland_env() { scythe::overlay::ensure_wayland_env(); } @@ -401,6 +464,11 @@ async fn main() -> Result<(), Box> { let mut rec_base_video_pts: i64 = 0; let mut rec_base_audio_pts: i64 = -1; + let mut stream_sender: Option> = None; + let mut stream_waiting_keyframe = false; + let mut stream_base_video_pts: i64 = 0; + let mut stream_base_audio_pts: i64 = -1; + thread::spawn(move || { while let Ok(drain) = mux_rx.recv() { let filename = scythe::config::ScytheConfig::format_video_filename("Replay", "mp4"); @@ -527,6 +595,7 @@ async fn main() -> Result<(), Box> { if let Some(mut m) = normal_muxer.take() { let _ = m.finalize(); } + let _ = stream_sender.take(); if let Ok(exe) = std::env::current_exe() { let mut cmd = std::process::Command::new(exe); cmd.stdin(std::process::Stdio::null()) @@ -645,27 +714,77 @@ async fn main() -> Result<(), Box> { println!("Audio mute toggled: {}", !cur); }, Command::StartStreaming => { - let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); - stream_state_clone.store(true, Ordering::SeqCst); - stream_start_clone.store(now, Ordering::SeqCst); - println!("Livestream started at timestamp {}", now); + let cfg = scythe::config::ScytheConfig::load(); + let key = cfg.stream_key.trim(); + if key.is_empty() { + eprintln!("Cannot start livestream: stream key is empty in config."); + stream_state_clone.store(false, Ordering::SeqCst); + stream_start_clone.store(0, Ordering::SeqCst); + } else if stream_sender.is_some() { + println!("Livestream is already active."); + } else { + let rtmp_url = format!("{}/{}", cfg.stream_url.trim_end_matches('/'), key); + let tx = spawn_stream_worker( + rtmp_url, + codec_ctx_ptr, + audio_codec_ctx_ptr, + Arc::clone(&stream_state_clone), + Arc::clone(&stream_start_clone), + ); + stream_sender = Some(tx); + stream_waiting_keyframe = true; + stream_base_video_pts = 0; + stream_base_audio_pts = -1; + let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); + stream_state_clone.store(true, Ordering::SeqCst); + stream_start_clone.store(now, Ordering::SeqCst); + println!("Livestream started at timestamp {}, waiting for video keyframe...", now); + } }, Command::StopStreaming => { + if stream_sender.take().is_some() { + println!("Livestream stopped."); + } + stream_waiting_keyframe = false; + stream_base_video_pts = 0; + stream_base_audio_pts = -1; stream_state_clone.store(false, Ordering::SeqCst); stream_start_clone.store(0, Ordering::SeqCst); - println!("Livestream stopped."); }, Command::ToggleStreaming => { - let cur = stream_state_clone.load(Ordering::SeqCst); - if cur { + if stream_sender.is_some() { + stream_sender = None; + stream_waiting_keyframe = false; + stream_base_video_pts = 0; + stream_base_audio_pts = -1; stream_state_clone.store(false, Ordering::SeqCst); stream_start_clone.store(0, Ordering::SeqCst); println!("ToggleStreaming: Livestream stopped."); } else { - let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); - stream_state_clone.store(true, Ordering::SeqCst); - stream_start_clone.store(now, Ordering::SeqCst); - println!("ToggleStreaming: Livestream started at timestamp {}", now); + let cfg = scythe::config::ScytheConfig::load(); + let key = cfg.stream_key.trim(); + if key.is_empty() { + eprintln!("Cannot start livestream: stream key is empty in config."); + stream_state_clone.store(false, Ordering::SeqCst); + stream_start_clone.store(0, Ordering::SeqCst); + } else { + let rtmp_url = format!("{}/{}", cfg.stream_url.trim_end_matches('/'), key); + let tx = spawn_stream_worker( + rtmp_url, + codec_ctx_ptr, + audio_codec_ctx_ptr, + Arc::clone(&stream_state_clone), + Arc::clone(&stream_start_clone), + ); + stream_sender = Some(tx); + stream_waiting_keyframe = true; + stream_base_video_pts = 0; + stream_base_audio_pts = -1; + let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); + stream_state_clone.store(true, Ordering::SeqCst); + stream_start_clone.store(now, Ordering::SeqCst); + println!("ToggleStreaming: Livestream started at timestamp {}, waiting for video keyframe...", now); + } } }, _ => {} @@ -711,6 +830,34 @@ async fn main() -> Result<(), Box> { let _ = muxer.write_packet(&rebased); } } + if let Some(ref tx) = stream_sender + && !stream_waiting_keyframe { + let stream_base_video_time_us = unsafe { + ffmpeg_next::ffi::av_rescale_q( + stream_base_video_pts, + video_time_base, + ffmpeg_next::ffi::AVRational { num: 1, den: 1_000_000 }, + ) + }; + let a_time_us = if let Some(a_tb) = audio_time_base { + unsafe { + ffmpeg_next::ffi::av_rescale_q( + pkt.pts(), + a_tb, + ffmpeg_next::ffi::AVRational { num: 1, den: 1_000_000 }, + ) + } + } else { + 0 + }; + if a_time_us >= stream_base_video_time_us { + if stream_base_audio_pts < 0 { + stream_base_audio_pts = pkt.pts(); + } + let rebased = pkt.rebased(stream_base_audio_pts); + let _ = tx.try_send(rebased); + } + } if config.replay_enabled { ring.push_overwrite(pkt); } @@ -731,7 +878,15 @@ async fn main() -> Result<(), Box> { recv(ticker) -> _ => { while ticker.try_recv().is_ok() {} - if !normal_recording && !config.replay_enabled { + if stream_sender.is_some() && !stream_state_clone.load(Ordering::Relaxed) { + stream_sender = None; + stream_waiting_keyframe = false; + stream_base_video_pts = 0; + stream_base_audio_pts = -1; + } + + let is_streaming = stream_sender.is_some(); + if !normal_recording && !config.replay_enabled && !is_streaming { if last_malloc_trim.elapsed().as_secs() >= 5 { last_malloc_trim = std::time::Instant::now(); #[cfg(target_os = "linux")] @@ -805,6 +960,20 @@ async fn main() -> Result<(), Box> { let _ = muxer.write_packet(&rebased); } } + + if let Some(ref tx) = stream_sender { + if stream_waiting_keyframe && pkt.is_keyframe() { + stream_waiting_keyframe = false; + stream_base_video_pts = pkt.pts(); + stream_base_audio_pts = -1; + println!("Livestream keyframe received, broadcasting video..."); + } + if !stream_waiting_keyframe { + let rebased = pkt.rebased(stream_base_video_pts); + let _ = tx.try_send(rebased); + } + } + if config.replay_enabled { ring.push_overwrite(pkt); } diff --git a/src/muxer.rs b/src/muxer.rs index 6171555..aedf42e 100644 --- a/src/muxer.rs +++ b/src/muxer.rs @@ -21,14 +21,21 @@ impl Muxer { let path_cstr = CString::new(path).map_err(|e| e.to_string())?; let mut fmt_ctx: *mut AVFormatContext = std::ptr::null_mut(); + let is_rtmp = path.starts_with("rtmp://") || path.starts_with("rtmps://"); + let format_name = if is_rtmp { + c"flv".as_ptr() + } else { + std::ptr::null() + }; + let ret = avformat_alloc_output_context2( &mut fmt_ctx, std::ptr::null(), - std::ptr::null(), + format_name, path_cstr.as_ptr(), ); if ret < 0 || fmt_ctx.is_null() { - return Err("Failed to alloc output context".into()); + return Err(format!("Failed to alloc output context (code {})", ret)); } let stream = avformat_new_stream(fmt_ctx, std::ptr::null()); @@ -61,15 +68,31 @@ impl Muxer { } if ((*(*fmt_ctx).oformat).flags & AVFMT_NOFILE) == 0 { - let ret = avio_open(&mut (*fmt_ctx).pb, path_cstr.as_ptr(), AVIO_FLAG_WRITE); + let mut io_opts: *mut AVDictionary = std::ptr::null_mut(); + if is_rtmp { + av_dict_set(&mut io_opts, c"rw_timeout".as_ptr(), c"5000000".as_ptr(), 0); + av_dict_set(&mut io_opts, c"tcp_nodelay".as_ptr(), c"1".as_ptr(), 0); + } + let ret = avio_open2( + &mut (*fmt_ctx).pb, + path_cstr.as_ptr(), + AVIO_FLAG_WRITE, + std::ptr::null(), + &mut io_opts, + ); + av_dict_free(&mut io_opts); if ret < 0 { avformat_free_context(fmt_ctx); - return Err("Failed to open output file".into()); + return Err(format!("Failed to open output URL/file (code {})", ret)); } } let mut opts: *mut AVDictionary = std::ptr::null_mut(); - av_dict_set(&mut opts, c"movflags".as_ptr(), c"+faststart".as_ptr(), 0); + if is_rtmp { + av_dict_set(&mut opts, c"flvflags".as_ptr(), c"no_duration_filesize".as_ptr(), 0); + } else { + av_dict_set(&mut opts, c"movflags".as_ptr(), c"+faststart".as_ptr(), 0); + } (*fmt_ctx).avoid_negative_ts = 2; // AVFMT_AVOID_NEG_TS_MAKE_ZERO: Shift timestamps so that they start at 0 (*fmt_ctx).max_interleave_delta = 1_000_000; let ret = avformat_write_header(fmt_ctx, &mut opts); @@ -79,7 +102,7 @@ impl Muxer { avio_closep(&mut (*fmt_ctx).pb); } avformat_free_context(fmt_ctx); - return Err("Failed to write header".into()); + return Err(format!("Failed to write header (code {})", ret)); } Ok(Self { @@ -155,3 +178,17 @@ impl Drop for Muxer { } } } + +#[cfg(test)] +mod tests { + #[test] + fn test_rtmp_url_classification() { + let rtmp_twitch = "rtmp://live.twitch.tv/app/live_test_key"; + let rtmps_youtube = "rtmps://a.rtmp.youtube.com/live2/xxxx-yyyy-zzzz"; + let local_mp4 = "/home/user/Videos/recording.mp4"; + + assert!(rtmp_twitch.starts_with("rtmp://") || rtmp_twitch.starts_with("rtmps://")); + assert!(rtmps_youtube.starts_with("rtmp://") || rtmps_youtube.starts_with("rtmps://")); + assert!(!local_mp4.starts_with("rtmp://") && !local_mp4.starts_with("rtmps://")); + } +}