mariano/mitus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

almost back to working state with rust transport

Browse Source
This commit is contained in:
buenosairesam
2026-04-09 22:15:16 -03:00
parent ff96dcb4f7
commit 512d8ecef8
13 changed files with 1504 additions and 488 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
media/server/src/main.rs
View File

@@ -1,24 +1,40 @@
mod session;
use std::path::PathBuf;
use anyhow::Result;
use cht_common::protocol::{self, ControlMessage, PacketType};
use session::Session;
use tokio::io::BufReader;
use tokio::net::TcpListener;
use tracing::{error, info};
use tracing::{error, info, warn};
const LISTEN_ADDR: &str = "0.0.0.0:4444";
const LISTEN_ADDR: &str = "0.0.0.0:4447";
const DEFAULT_SESSIONS_DIR: &str = "/home/mariano/wdir/cht/data/sessions";
fn sessions_dir() -> PathBuf {
std::env::var("CHT_SESSIONS_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from(DEFAULT_SESSIONS_DIR))
}
#[tokio::main]
async fn main() -> Result<()> {
cht_common::logging::init("server");
let sessions_dir = sessions_dir();
info!("Sessions dir: {}", sessions_dir.display());
let listener = TcpListener::bind(LISTEN_ADDR).await?;
info!("Server listening on {LISTEN_ADDR}");
loop {
let (stream, addr) = listener.accept().await?;
info!("Client connected from {addr}");
let sdir = sessions_dir.clone();
tokio::spawn(async move {
if let Err(e) = handle_client(stream).await {
if let Err(e) = handle_client(stream, sdir).await {
error!("Client {addr} error: {e:#}");
}
info!("Client {addr} disconnected");
@@ -26,17 +42,19 @@ async fn main() -> Result<()> {
}
}
async fn handle_client(stream: tokio::net::TcpStream) -> Result<()> {
async fn handle_client(
stream: tokio::net::TcpStream,
sessions_dir: PathBuf,
) -> Result<()> {
let mut reader = BufReader::new(stream);
let mut video_packets = 0u64;
let mut audio_packets = 0u64;
let mut session: Option<Session> = None;
let mut video_count = 0u64;
let mut audio_count = 0u64;
loop {
let packet = match protocol::read_packet(&mut reader).await {
Ok(p) => p,
Err(e) => {
// Any read error at the header boundary is a clean disconnect
// (includes EOF from flush + shutdown)
let msg = format!("{e:#}");
if msg.contains("eof") || msg.contains("Eof")
|| msg.contains("connection reset")
@@ -50,25 +68,60 @@ async fn handle_client(stream: tokio::net::TcpStream) -> Result<()> {
match packet.header.packet_type {
PacketType::Video => {
video_packets += 1;
if video_packets % 300 == 1 {
info!(
"video: {video_packets} packets, ts={}ms, keyframe={}",
packet.header.timestamp_ns / 1_000_000,
packet.header.is_keyframe(),
);
if let Some(s) = &mut session {
// Blocking write — offload to blocking thread to avoid stalling tokio.
let data = packet.payload;
let keyframe = packet.header.is_keyframe();
tokio::task::block_in_place(|| s.write_video(&data, keyframe))?;
video_count += 1;
if video_count % 300 == 1 {
info!("video: {video_count} packets, ts={}ms, keyframe={}",
packet.header.timestamp_ns / 1_000_000,
packet.header.is_keyframe());
}
} else {
warn!("Video packet before SessionStart — dropped");
}
}
PacketType::Audio => {
audio_packets += 1;
if let Some(s) = &mut session {
let data = packet.payload;
tokio::task::block_in_place(|| s.write_audio(&data))?;
audio_count += 1;
if audio_count % 500 == 1 {
info!("audio: {audio_count} packets");
}
}
}
PacketType::Control => {
let ctrl = ControlMessage::from_payload(&packet.payload)?;
info!("control: {ctrl:?}");
match ctrl {
ControlMessage::SessionStart { id, video, .. } => {
let s = tokio::task::block_in_place(|| {
Session::start(&id, &sessions_dir, video.fps)
})?;
session = Some(s);
}
ControlMessage::SessionStop => {
if let Some(s) = session.take() {
tokio::task::block_in_place(|| s.close());
}
break;
}
ControlMessage::Keepalive
| ControlMessage::Reconnect { .. }
| ControlMessage::ParamChange { .. } => {}
}
}
}
}
info!("Session totals: {video_packets} video, {audio_packets} audio packets");
if let Some(s) = session.take() {
tokio::task::block_in_place(|| s.close());
}
info!("Session totals: {video_count} video, {audio_count} audio packets");
Ok(())
}

306
media/server/src/session.rs Normal file
View File

@@ -0,0 +1,306 @@
//! Session: manages the ffmpeg recording subprocess for one client connection.
//!
//! Receives raw H.264 NAL units and AAC audio from the transport:
//! - Video: piped into ffmpeg → fragmented MP4 + UDP relay for live display
//! - Audio: written to raw AAC file for Python post-processing
//!
//! Also provides a Unix domain socket at `stream/scene.sock` carrying a copy
//! of the raw H.264 stream for Python's GPU scene detection. The socket is
//! fire-and-forget: if nobody connects, data is silently dropped; if the
//! reader is slow, old frames are dropped rather than stalling recording.
//!
//! Creates the session directory and writes its path to `data/active-session`
//! so the Python app can pick it up for SessionProcessor (audio extraction, etc).
use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::thread;
use anyhow::{Context, Result};
use tokio::io::AsyncWriteExt;
use tracing::{debug, info, warn};
// Written next to the sessions/ directory so everything stays under data/.
// Python reads this to discover the session dir created by cht-server.
const ACTIVE_SESSION_FILENAME: &str = "active-session";
const RELAY_URL: &str = "udp://127.0.0.1:4445";
const SCENE_SOCKET_NAME: &str = "scene.sock";
struct ScenePacket {
data: Vec<u8>,
keyframe: bool,
}
pub struct Session {
#[allow(dead_code)]
session_dir: PathBuf,
active_session_file: PathBuf,
ffmpeg: Child,
video_stdin: Option<ChildStdin>,
audio_file: Option<File>,
scene_tx: Option<tokio::sync::mpsc::Sender<ScenePacket>>,
#[allow(dead_code)]
fps: u32,
}
impl Session {
pub fn start(session_id: &str, sessions_dir: &Path, fps: u32) -> Result<Self> {
let active_session_file = sessions_dir
.parent()
.unwrap_or(sessions_dir)
.join(ACTIVE_SESSION_FILENAME);
let session_dir = sessions_dir.join(session_id);
let stream_dir = session_dir.join("stream");
fs::create_dir_all(&stream_dir)
.with_context(|| format!("create session dir: {}", stream_dir.display()))?;
let recording_path = stream_dir.join("recording_000.mp4");
let audio_path = stream_dir.join("audio.aac");
info!("Session {session_id}: recording → {}", recording_path.display());
let mut child = Command::new("ffmpeg")
.args([
"-f", "h264",
"-framerate", &fps.to_string(),
"-i", "pipe:0",
// fMP4 — same flags as Python StreamRecorder
"-c:v", "copy",
"-f", "mp4",
"-movflags", "frag_keyframe+empty_moov+default_base_moof",
"-flush_packets", "1",
recording_path.to_str().unwrap(),
// UDP relay for live display
"-c:v", "copy",
"-f", "mpegts",
RELAY_URL,
"-hide_banner", "-loglevel", "warning",
])
.stdin(Stdio::piped())
.stdout(Stdio::null())
.stderr(Stdio::piped())
.spawn()
.context("spawn ffmpeg recorder")?;
let video_stdin = child.stdin.take().expect("stdin piped");
// Drain stderr so ffmpeg never blocks on a full pipe.
let stderr = child.stderr.take().expect("stderr piped");
let sid = session_id.to_string();
thread::Builder::new()
.name("ffmpeg-recorder-stderr".into())
.spawn(move || {
use std::io::{BufRead, BufReader};
for line in BufReader::new(stderr).lines().map_while(Result::ok) {
if !line.is_empty() {
debug!("[recorder/{sid}] {line}");
}
}
})
.expect("spawn stderr thread");
// Open audio file for raw AAC frames from client
let audio_file = File::create(&audio_path)
.map(Some)
.unwrap_or_else(|e| {
warn!("Could not create audio file: {e}");
None
});
// Scene relay: Unix socket for Python scene detection.
let socket_path = stream_dir.join(SCENE_SOCKET_NAME);
let (scene_tx, scene_rx) = tokio::sync::mpsc::channel(32);
tokio::spawn(scene_relay_task(socket_path, scene_rx));
// Tell Python which session dir to watch.
if let Err(e) = fs::write(&active_session_file, session_dir.to_str().unwrap_or("")) {
warn!("Could not write {}: {e}", active_session_file.display());
}
info!("Session {session_id}: ffmpeg pid={}, audio → {}",
child.id(), audio_path.display());
Ok(Self {
session_dir,
active_session_file,
ffmpeg: child,
video_stdin: Some(video_stdin),
audio_file,
scene_tx: Some(scene_tx),
fps,
})
}
pub fn write_video(&mut self, data: &[u8], keyframe: bool) -> Result<()> {
if let Some(stdin) = &mut self.video_stdin {
stdin.write_all(data).context("write H.264 to ffmpeg")?;
}
// Best-effort relay to scene detector — drop if channel full.
if let Some(tx) = &self.scene_tx {
let _ = tx.try_send(ScenePacket { data: data.to_vec(), keyframe });
}
Ok(())
}
pub fn write_audio(&mut self, data: &[u8]) -> Result<()> {
if let Some(f) = &mut self.audio_file {
// Wrap raw AAC frame with ADTS header so the file is playable/parseable.
// Assumes AAC-LC, 48kHz, stereo (matches client's encoder config).
write_adts_frame(f, data)?;
}
Ok(())
}
#[allow(dead_code)]
pub fn session_dir(&self) -> &Path {
&self.session_dir
}
pub fn close(mut self) {
// Drop stdin → ffmpeg gets EOF → flushes and exits cleanly.
drop(self.video_stdin.take());
drop(self.audio_file.take());
// Drop scene_tx → relay task sees channel closed → exits.
drop(self.scene_tx.take());
match self.ffmpeg.wait() {
Ok(s) => info!("ffmpeg recorder exited: {s}"),
Err(e) => warn!("ffmpeg recorder wait error: {e}"),
}
// Clear the active session marker.
let _ = fs::remove_file(&self.active_session_file);
}
}
impl Drop for Session {
fn drop(&mut self) {
if self.video_stdin.is_some() {
drop(self.video_stdin.take());
drop(self.audio_file.take());
drop(self.scene_tx.take());
let _ = self.ffmpeg.kill();
}
}
}
// ---------------------------------------------------------------------------
// Scene relay: serves raw H.264 over a Unix domain socket
// ---------------------------------------------------------------------------
async fn scene_relay_task(
socket_path: PathBuf,
mut rx: tokio::sync::mpsc::Receiver<ScenePacket>,
) {
// Remove stale socket from a previous session.
let _ = fs::remove_file(&socket_path);
let listener = match tokio::net::UnixListener::bind(&socket_path) {
Ok(l) => l,
Err(e) => {
warn!("Scene relay: bind failed on {}: {e}", socket_path.display());
return;
}
};
info!("Scene relay: listening on {}", socket_path.display());
let mut client: Option<tokio::net::UnixStream> = None;
// Buffer the latest keyframe so new clients start with a valid decoder state.
let mut last_keyframe: Option<Vec<u8>> = None;
loop {
if client.is_some() {
// We have a connected reader — forward data.
match rx.recv().await {
Some(pkt) => {
if pkt.keyframe {
last_keyframe = Some(pkt.data.clone());
}
let stream = client.as_mut().unwrap();
if stream.write_all(&pkt.data).await.is_err() {
info!("Scene relay: client disconnected");
client = None;
}
}
None => break, // Channel closed, session ending.
}
} else {
// No reader — accept connections while draining the channel.
tokio::select! {
biased;
result = listener.accept() => {
match result {
Ok((mut stream, _)) => {
info!("Scene relay: client connected");
// Send the last keyframe so the decoder can initialize.
if let Some(ref kf) = last_keyframe {
if stream.write_all(kf).await.is_err() {
warn!("Scene relay: failed to send keyframe");
continue;
}
info!("Scene relay: sent keyframe ({} bytes)", kf.len());
}
client = Some(stream);
}
Err(e) => warn!("Scene relay: accept error: {e}"),
}
}
pkt = rx.recv() => {
match pkt {
Some(pkt) => {
if pkt.keyframe {
last_keyframe = Some(pkt.data);
}
// Discard — no reader connected.
}
None => break, // Channel closed.
}
}
}
}
}
drop(client);
let _ = fs::remove_file(&socket_path);
info!("Scene relay: stopped");
}
// ---------------------------------------------------------------------------
// ADTS header for raw AAC framing
// ---------------------------------------------------------------------------
/// Write a raw AAC frame wrapped in a 7-byte ADTS header.
///
/// Fixed params: AAC-LC profile, 48 kHz sample rate, 2 channels (stereo).
/// These match the client's `-c:a aac -b:a 128k` default config.
fn write_adts_frame(w: &mut impl Write, aac_data: &[u8]) -> Result<()> {
// ADTS fixed header fields:
// profile: AAC-LC = 1 (stored as profile-1 = 0 in MPEG-4 ID mode)
// sample_rate: 48000 → index 3
// channels: 2 → channel_configuration 2
const PROFILE_MINUS1: u8 = 1; // AAC-LC
const SR_IDX: u8 = 3; // 48 kHz
const CH_CFG: u8 = 2; // stereo
let frame_len = (aac_data.len() + 7) as u16; // total ADTS frame = header + payload
let header: [u8; 7] = [
// byte 0-1: syncword(12) | ID(1)=0(MPEG4) | layer(2)=0 | protection(1)=1(no CRC)
0xFF,
0xF1,
// byte 2: profile(2) | sr_idx(4) | private(1)=0 | ch_cfg[2](1)
(PROFILE_MINUS1 << 6) | (SR_IDX << 2) | ((CH_CFG >> 2) & 1),
// byte 3: ch_cfg[1:0](2) | orig(1)=0 | home(1)=0 | copyright_id(1)=0 | copyright_start(1)=0 | frame_len[12:11](2)
((CH_CFG & 3) << 6) | ((frame_len >> 11) as u8 & 0x03),
// byte 4: frame_len[10:3](8)
((frame_len >> 3) & 0xFF) as u8,
// byte 5: frame_len[2:0](3) | buffer_fullness[10:6](5)
((frame_len & 0x07) << 5) as u8 | 0x1F,
// byte 6: buffer_fullness[5:0](6) | num_aac_frames_minus1(2)=0
0xFC,
];
w.write_all(&header).context("ADTS header")?;
w.write_all(aac_data).context("AAC frame")?;
Ok(())
}