From 4e6b7dfe6ef37b7e13d3d044b4fb97fe5cf91927 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 7 Jul 2026 18:08:46 -0600 Subject: [PATCH 1/7] Wire worker binary transfers --- .gitignore | 1 + Cargo.lock | 8 + Cargo.toml | 2 + src/command/worker.rs | 700 ++++++++++++++++++++++++++++++--- src/command/worker_protocol.rs | 85 +++- src/command/worker_transfer.rs | 12 +- src/main.rs | 2 +- 7 files changed, 721 insertions(+), 89 deletions(-) diff --git a/.gitignore b/.gitignore index 809acc62..814c2d71 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ /target /aur +/ab-av1-worker-*/ .envrc .direnv/ *.log diff --git a/Cargo.lock b/Cargo.lock index e2826c91..90090259 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9,11 +9,13 @@ dependencies = [ "allocation-counter", "anyhow", "async-stream", + "base64", "blake3", "clap", "clap-verbosity-flag", "clap_complete", "console", + "crc32fast", "dirs", "fastrand", "ffprobe", @@ -176,6 +178,12 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bit-set" version = "0.8.0" diff --git a/Cargo.toml b/Cargo.toml index 6fadc8b5..d11413e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ readme = "README.md" [dependencies] anyhow = "1.0.53" async-stream = "0.3.5" +base64 = "0.22" blake3 = "1.3.3" clap = { version = "4", features = ["derive", "env", "wrap_help"] } clap-verbosity-flag = "3.0.2" @@ -25,6 +26,7 @@ futures-util = "0.3.19" humantime = "2.1" indicatif = "0.18" infer = { version = "0.19", default-features = false } +crc32fast = "1" log = "0.4.21" pin-project-lite = "0.2.16" same-file = "1.0.6" diff --git a/src/command/worker.rs b/src/command/worker.rs index 77154f46..9ae7d947 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -1,25 +1,38 @@ use crate::command::worker_protocol::{ - AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ClientEvent, ClientFrame, - ErrorReplyPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, + AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, + ClientEvent, ClientFrame, ErrorReplyPayload, JobResultPayload, ReplyBody, ServerPushFrame, + ServerReply, TransferStartedPayload, }; +use crate::command::worker_transfer::{Chunk, ChunkReceiver}; use crate::command::{args, crf_search, sample_encode}; use crate::ffprobe::Ffprobe; use crate::temporary; use anyhow::{Context, Result, anyhow, bail}; +use base64::{Engine as _, engine::general_purpose::STANDARD}; use clap::Parser; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::{path::PathBuf, sync::Arc, time::Duration}; +use std::{ + path::{Path, PathBuf}, + sync::Arc, + time::Duration, +}; use tokio::net::TcpStream; use tokio_tungstenite::{ - MaybeTlsStream, WebSocketStream, connect_async, + MaybeTlsStream, WebSocketStream, tungstenite::{Error as WsError, Message}, + tungstenite::{client::IntoClientRequest, http::header::ORIGIN, protocol::WebSocketConfig}, }; -use tracing::debug; +use tracing::{debug, trace}; const PHOENIX_VSN: &str = "2.0.0"; const SUPPORTED_PROTOCOL_VERSION: u64 = 1; +const TRANSFER_CHUNK_MAGIC: &[u8; 4] = b"RAV1"; +const TRANSFER_CHUNK_VERSION: u8 = 1; +const TRANSFER_CHUNK_TYPE: u8 = 1; +const TRANSFER_CHUNK_HEADER_LEN: usize = 52; +const MAX_TRANSFER_FRAME_BYTES: usize = 640 * 1024 * 1024; /// Connect to a Reencodarr websocket worker endpoint and request one job. #[derive(Parser, Debug, Clone)] @@ -47,6 +60,10 @@ pub struct Args { /// Exit after the first work poll instead of running as a long-lived worker. #[arg(long)] once: bool, + + /// Use a local file instead of waiting for the server to transfer one over the socket. + #[arg(long)] + local_path: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -57,6 +74,7 @@ pub struct WorkerConfig { version: String, protocol_version: u64, once: bool, + local_path: Option, } impl From for WorkerConfig { @@ -68,6 +86,7 @@ impl From for WorkerConfig { version, protocol_version, once, + local_path, }: Args, ) -> Self { Self { @@ -77,6 +96,7 @@ impl From for WorkerConfig { version, protocol_version, once, + local_path, } } } @@ -94,6 +114,7 @@ pub struct WorkerSession { struct WorkerJob { assignment: crate::command::worker_protocol::JobAssignedPayload, input_dir: PathBuf, + input_path: PathBuf, } #[cfg_attr(not(test), allow(dead_code))] @@ -101,22 +122,24 @@ impl WorkerJob { fn new( assignment: crate::command::worker_protocol::JobAssignedPayload, input_dir: PathBuf, + input_path: PathBuf, ) -> Self { Self { assignment, input_dir, + input_path, } } - fn input_path(&self) -> PathBuf { - self.input_dir.join(&self.assignment.source_name) + fn input_path(&self) -> &Path { + &self.input_path } fn crf_search_config(&self, encoder: args::Encoder) -> Result { Ok(crf_search::CrfSearchConfig { args: args::Encode { encoder, - input: self.input_path(), + input: self.input_path().to_path_buf(), vfilter: None, pix_format: None, preset: None, @@ -173,6 +196,109 @@ impl WorkerJob { } } +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug)] +struct PendingJob { + job: WorkerJob, + receiver: Option, +} + +#[cfg_attr(not(test), allow(dead_code))] +impl PendingJob { + fn new(job: WorkerJob, receiver: ChunkReceiver) -> Self { + Self { + job, + receiver: Some(receiver), + } + } + + fn job(&self) -> &WorkerJob { + &self.job + } + + fn input_path(&self) -> &Path { + self.job.input_path() + } + + fn apply_chunk(&mut self, chunk: ChunkTransferPayload) -> Result<()> { + let bytes = STANDARD + .decode(chunk.data.as_bytes()) + .context("decode transfer chunk payload")?; + self.apply_raw_chunk(TransferChunk { + transfer_id: chunk.transfer_id, + video_id: chunk.video_id, + chunk_index: chunk.chunk_index, + total_chunks: chunk.total_chunks, + bytes_sent: chunk.bytes_sent, + total_bytes: chunk.total_bytes, + crc32: chunk.crc32, + bytes, + }) + } + + fn apply_raw_chunk(&mut self, chunk: TransferChunk) -> Result<()> { + if chunk.transfer_id != self.job.assignment.job_id { + bail!( + "chunk transfer job mismatch: expected {}, got {}", + self.job.assignment.job_id, + chunk.transfer_id + ); + } + if chunk.video_id != self.job.assignment.video_id { + bail!( + "chunk transfer video mismatch: expected {}, got {}", + self.job.assignment.video_id, + chunk.video_id + ); + } + if chunk.total_bytes != self.job.assignment.size_bytes { + bail!( + "chunk transfer size mismatch: expected {}, got {}", + self.job.assignment.size_bytes, + chunk.total_bytes + ); + } + let offset = self + .receiver + .as_ref() + .expect("pending receiver") + .received_bytes(); + if offset.saturating_add(chunk.bytes.len() as u64) != chunk.bytes_sent { + bail!( + "chunk {} size mismatch: expected cumulative {}, got {}", + chunk.chunk_index, + chunk.bytes_sent, + offset.saturating_add(chunk.bytes.len() as u64) + ); + } + + let receiver = self.receiver.as_mut().expect("pending receiver"); + receiver.push(Chunk { + index: chunk.chunk_index, + offset, + bytes: chunk.bytes, + checksum: chunk.crc32, + })?; + Ok(()) + } + + fn finish(&mut self) -> Result<()> { + let final_path = self.input_path().to_path_buf(); + let receiver = self.receiver.take().context("missing chunk receiver")?; + let written = receiver + .finish(Some(self.job.assignment.size_bytes), None) + .context("finalize worker input transfer")?; + if written != final_path { + bail!( + "transfer finished at unexpected path: expected {}, got {}", + final_path.display(), + written.display() + ); + } + Ok(()) + } +} + #[cfg_attr(not(test), allow(dead_code))] async fn run_worker_job(job: WorkerJob, probe: Arc) -> Result { run_worker_job_until(job, probe, std::future::pending::<()>()).await @@ -209,12 +335,14 @@ where } fn worker_job_input_dir(job_id: &str) -> PathBuf { - std::env::temp_dir().join(format!( - "ab-av1-worker-{}-{}-{}", - std::process::id(), - job_id, - fastrand::u64(..) - )) + std::env::current_dir() + .expect("current working directory") + .join(format!( + "ab-av1-worker-{}-{}-{}", + std::process::id(), + job_id, + fastrand::u64(..) + )) } #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] @@ -281,6 +409,25 @@ enum PendingJobOutcome { Canceled, } +#[derive(Debug)] +enum WorkerPush { + Cancel(CancelPayload), + Started(TransferStartedPayload), + Chunk(ChunkTransferPayload), +} + +#[derive(Debug)] +struct TransferChunk { + transfer_id: String, + video_id: u64, + chunk_index: u64, + total_chunks: u64, + bytes_sent: u64, + total_bytes: u64, + crc32: u64, + bytes: Vec, +} + type WorkerSocket = WebSocketStream>; struct ConnectedWorker { @@ -293,9 +440,25 @@ struct ConnectedWorker { impl ConnectedWorker { async fn connect(config: &WorkerConfig) -> Result { let request_url = worker_websocket_url(&config.connect, &config.token)?; - let (mut socket, _) = connect_async(&request_url) - .await - .map_err(|error| websocket_connect_error(&request_url, error))?; + let mut request = request_url + .clone() + .into_client_request() + .context("build websocket request")?; + request.headers_mut().insert( + ORIGIN, + config + .connect + .trim_end_matches('/') + .parse() + .context("build websocket origin")?, + ); + let (mut socket, _) = tokio_tungstenite::connect_async_with_config( + request, + Some(worker_websocket_config()), + false, + ) + .await + .map_err(|error| websocket_connect_error(&request_url, error))?; send_json(&mut socket, ClientFrame::new(1, ClientEvent::Join)).await?; let join: JoinResponse = expect_reply(&mut socket, "1", "phx_join").await?; @@ -340,7 +503,7 @@ impl ConnectedWorker { async fn wait_for_pending_job( &mut self, - job: &WorkerJob, + pending_job: &mut PendingJob, idle_delay: Duration, ) -> Result { tokio::select! { @@ -355,15 +518,109 @@ impl ConnectedWorker { } Some(Ok(Message::Pong(_))) => Ok(PendingJobOutcome::Waiting), Some(Ok(Message::Text(text))) => { - if let Some(cancel) = decode_cancel_push(&text)? - && cancel.job_id == job.assignment.job_id + match decode_worker_push(&text)? { + Some(WorkerPush::Cancel(cancel)) + if cancel.job_id == pending_job.job().assignment.job_id => + { + eprintln!( + "worker job {} canceled: {}", + cancel.job_id, cancel.reason + ); + return Ok(PendingJobOutcome::Canceled); + } + Some(WorkerPush::Started(started)) + if started.transfer_id == pending_job.job().assignment.job_id => + { + debug!( + job_id = %started.transfer_id, + source_name = %started.source_name, + chunk_size_bytes = started.chunk_size_bytes, + size_bytes = started.size_bytes, + total_bytes = started.total_bytes, + total_chunks = started.total_chunks, + "transfer started" + ); + Ok(PendingJobOutcome::Waiting) + } + Some(WorkerPush::Chunk(chunk)) + if chunk.transfer_id == pending_job.job().assignment.job_id => + { + if chunk.chunk_index == 0 || chunk.chunk_index % 256 == 0 { + debug!( + job_id = %chunk.transfer_id, + chunk_index = chunk.chunk_index, + bytes_sent = chunk.bytes_sent, + total_bytes = chunk.total_bytes, + total_chunks = chunk.total_chunks, + "received chunk" + ); + } else { + trace!( + job_id = %chunk.transfer_id, + chunk_index = chunk.chunk_index, + bytes_sent = chunk.bytes_sent, + total_bytes = chunk.total_bytes, + total_chunks = chunk.total_chunks, + "received chunk" + ); + } + pending_job.apply_chunk(chunk)?; + if pending_job.receiver.as_ref().is_some_and(|receiver| { + receiver.received_bytes() + == pending_job.job.assignment.size_bytes + }) { + debug!( + job_id = %pending_job.job().assignment.job_id, + "transfer complete" + ); + pending_job.finish()?; + return Ok(PendingJobOutcome::Ready); + } + Ok(PendingJobOutcome::Waiting) + } + Some(_) => Ok(PendingJobOutcome::Waiting), + None => Ok(PendingJobOutcome::Waiting), + } + } + Some(Ok(Message::Binary(bytes))) => { + let chunk = decode_binary_worker_push(&bytes)?; + if let Some(chunk) = chunk + && chunk.transfer_id == pending_job.job().assignment.job_id { - eprintln!("worker job {} canceled: {}", cancel.job_id, cancel.reason); - return Ok(PendingJobOutcome::Canceled); + if chunk.chunk_index == 0 || chunk.chunk_index % 16 == 0 { + debug!( + job_id = %chunk.transfer_id, + chunk_index = chunk.chunk_index, + bytes_sent = chunk.bytes_sent, + total_bytes = chunk.total_bytes, + total_chunks = chunk.total_chunks, + "received binary chunk" + ); + } else { + trace!( + job_id = %chunk.transfer_id, + chunk_index = chunk.chunk_index, + bytes_sent = chunk.bytes_sent, + total_bytes = chunk.total_bytes, + total_chunks = chunk.total_chunks, + "received binary chunk" + ); + } + pending_job.apply_raw_chunk(chunk)?; + if pending_job.receiver.as_ref().is_some_and(|receiver| { + receiver.received_bytes() == pending_job.job.assignment.size_bytes + }) { + debug!( + job_id = %pending_job.job().assignment.job_id, + "transfer complete" + ); + pending_job.finish()?; + return Ok(PendingJobOutcome::Ready); + } } Ok(PendingJobOutcome::Waiting) } - Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => { + Some(Ok(Message::Frame(_))) => { Ok(PendingJobOutcome::Waiting) } Some(Ok(Message::Close(frame))) => { @@ -374,11 +631,16 @@ impl ConnectedWorker { } } _ = tokio::time::sleep(idle_delay) => { - if job.input_path().exists() { - Ok(PendingJobOutcome::Ready) - } else { - Ok(PendingJobOutcome::Waiting) - } + debug!( + job_id = %pending_job.job().assignment.job_id, + received_bytes = pending_job + .receiver + .as_ref() + .map(|receiver| receiver.received_bytes()) + .unwrap_or_default(), + "still waiting on websocket transfer" + ); + Ok(PendingJobOutcome::Waiting) } } } @@ -390,12 +652,9 @@ async fn run_worker_job_and_publish(job: &WorkerJob) -> Result<()> { input = %job.input_path().display(), "starting worker job" ); - let probe = Arc::new(crate::ffprobe::probe(&job.input_path())); + let probe = Arc::new(crate::ffprobe::probe(job.input_path())); debug!(job_id = %job.assignment.job_id, "probe complete, running crf search"); - let best = run_worker_job(job.clone(), probe).await; - debug!(job_id = %job.assignment.job_id, "cleaning temp files"); - temporary::clean(true).await; - let best = best?; + let best = run_worker_job(job.clone(), probe).await?; debug!(job_id = %job.assignment.job_id, "publishing worker result"); println!( @@ -407,11 +666,15 @@ async fn run_worker_job_and_publish(job: &WorkerJob) -> Result<()> { fn build_worker_job( assignment: crate::command::worker_protocol::JobAssignedPayload, + local_path: Option<&Path>, ) -> Result { let input_dir = worker_job_input_dir(&assignment.job_id); std::fs::create_dir_all(&input_dir).context("create worker job dir")?; - temporary::add(&input_dir, temporary::TempKind::NotKeepable); - Ok(WorkerJob::new(assignment, input_dir)) + temporary::add(&input_dir, temporary::TempKind::Keepable); + let input_path = local_path + .map(Path::to_path_buf) + .unwrap_or_else(|| input_dir.join(&assignment.source_name)); + Ok(WorkerJob::new(assignment, input_dir, input_path)) } pub async fn worker(config: WorkerConfig) -> Result<()> { @@ -440,7 +703,7 @@ async fn run_worker_until(config: &WorkerConfig, runtime: WorkerRuntime) -> Resu return Ok(()); } Err(error) => { - eprintln!("worker connection lost: {error}"); + eprintln!("worker connection lost: {error:#}"); tokio::time::sleep(reconnect_backoff.next_delay()).await; } } @@ -456,34 +719,40 @@ async fn run_connected_worker( connect = %config.connect, worker_id = %config.worker_id, once = config.once, + local_path = ?config.local_path, "connecting worker" ); let mut worker = ConnectedWorker::connect(config).await?; - let mut pending_job: Option = None; + let mut pending_job: Option = None; loop { - if let Some(job) = pending_job.as_ref() { - debug!( - job_id = %job.assignment.job_id, - input = %job.input_path().display(), - "waiting for pending job input" - ); - let next = worker.wait_for_pending_job(job, runtime.idle_delay).await?; - if matches!(next, PendingJobOutcome::Waiting) { - debug!(job_id = %job.assignment.job_id, "pending job still waiting"); - continue; - } - if matches!(next, PendingJobOutcome::Canceled) { - debug!(job_id = %job.assignment.job_id, "pending job canceled"); - temporary::clean(true).await; - pending_job = None; - continue; + if pending_job.is_some() { + let next = { + let job = pending_job.as_mut().expect("pending job"); + trace!( + job_id = %job.job.assignment.job_id, + input = %job.input_path().display(), + "waiting for pending job input" + ); + worker.wait_for_pending_job(job, runtime.idle_delay).await? + }; + + match next { + PendingJobOutcome::Waiting => continue, + PendingJobOutcome::Canceled => { + if let Some(job) = pending_job.as_ref() { + debug!(job_id = %job.job.assignment.job_id, "pending job canceled"); + } + pending_job = None; + continue; + } + PendingJobOutcome::Ready => { + let job = pending_job.take().expect("pending job"); + debug!(job_id = %job.job.assignment.job_id, "pending job input arrived"); + run_worker_job_and_publish(&job.job).await?; + continue; + } } - - debug!(job_id = %job.assignment.job_id, "pending job input arrived"); - run_worker_job_and_publish(job).await?; - pending_job = None; - continue; } debug!("requesting work"); @@ -496,7 +765,7 @@ async fn run_connected_worker( ); if let ServerReply::JobAssigned(assignment) = work_status { - let job = build_worker_job(assignment)?; + let job = build_worker_job(assignment, config.local_path.as_deref())?; debug!( job_id = %job.assignment.job_id, input = %job.input_path().display(), @@ -505,13 +774,24 @@ async fn run_connected_worker( if job.input_path().exists() { debug!(job_id = %job.assignment.job_id, "input already present, starting job"); run_worker_job_and_publish(&job).await?; + } else if config.local_path.is_some() { + bail!( + "local input path does not exist: {}", + job.input_path().display() + ); } else { + let receiver = ChunkReceiver::new( + job.input_path(), + &job.input_dir, + Some(job.assignment.size_bytes), + ) + .context("prepare worker input transfer")?; debug!( job_id = %job.assignment.job_id, input = %job.input_path().display(), - "waiting for worker input file" + "waiting for worker input over websocket" ); - pending_job = Some(job); + pending_job = Some(PendingJob::new(job, receiver)); } continue; } @@ -550,17 +830,176 @@ fn work_status_label(reply: &ServerReply) -> String { } } -fn decode_cancel_push(text: &str) -> Result> { +fn decode_worker_push(text: &str) -> Result> { let frame: ServerPushFrame = match serde_json::from_str(text) { Ok(frame) => frame, Err(_) => return Ok(None), }; - if frame.2 != CRF_SEARCH_TOPIC || frame.3 != "cancel" { + if frame.2 != CRF_SEARCH_TOPIC { return Ok(None); } - let cancel = serde_json::from_value::(frame.4).context("decode cancel push")?; - Ok(Some(cancel)) + let payload = frame.4.clone(); + if matches!(frame.3.as_str(), "chunk_transfer" | "transfer_chunk") { + trace!( + topic = %frame.2, + event = %frame.3, + payload_bytes = text.len(), + "received worker push" + ); + } else { + debug!( + topic = %frame.2, + event = %frame.3, + payload_bytes = text.len(), + "received worker push" + ); + } + let push = match frame.3.as_str() { + "cancel" => WorkerPush::Cancel( + serde_json::from_value::(payload.clone()) + .context("decode cancel push")?, + ), + "transfer_started" => WorkerPush::Started( + serde_json::from_value::(payload.clone()) + .with_context(|| format!("decode transfer started push event={}", frame.3))?, + ), + "chunk_transfer" | "transfer_chunk" => WorkerPush::Chunk( + serde_json::from_value::(payload.clone()).with_context(|| { + format!( + "decode chunk transfer push event={} payload_bytes={}", + frame.3, + text.len() + ) + })?, + ), + _ => return Ok(None), + }; + + Ok(Some(push)) +} + +fn decode_binary_transfer_chunk(bytes: &[u8]) -> Result { + if bytes.len() < TRANSFER_CHUNK_HEADER_LEN { + bail!( + "binary transfer chunk too short: got {} bytes, need at least {}", + bytes.len(), + TRANSFER_CHUNK_HEADER_LEN + ); + } + if &bytes[0..4] != TRANSFER_CHUNK_MAGIC { + bail!( + "invalid binary transfer chunk magic: len={} prefix={}", + bytes.len(), + hex_prefix(bytes, 24) + ); + } + if bytes[4] != TRANSFER_CHUNK_VERSION { + bail!("unsupported binary transfer chunk version {}", bytes[4]); + } + if bytes[5] != TRANSFER_CHUNK_TYPE { + bail!("unsupported binary transfer chunk type {}", bytes[5]); + } + + let transfer_id_size = u16::from_be_bytes([bytes[6], bytes[7]]) as usize; + let transfer_id_start = TRANSFER_CHUNK_HEADER_LEN; + let data_start = transfer_id_start + .checked_add(transfer_id_size) + .context("binary transfer chunk transfer_id_size overflow")?; + if bytes.len() < data_start { + bail!( + "binary transfer chunk transfer_id truncated: got {} bytes, need {}", + bytes.len(), + data_start + ); + } + + let transfer_id = std::str::from_utf8(&bytes[transfer_id_start..data_start]) + .context("decode binary transfer chunk transfer_id")? + .to_owned(); + + Ok(TransferChunk { + transfer_id, + video_id: read_u64(&bytes, 8), + chunk_index: read_u64(&bytes, 16), + total_chunks: read_u64(&bytes, 24), + bytes_sent: read_u64(&bytes, 32), + total_bytes: read_u64(&bytes, 40), + crc32: read_u32(&bytes, 48) as u64, + bytes: bytes[data_start..].to_vec(), + }) +} + +fn decode_binary_worker_push(bytes: &[u8]) -> Result> { + let Some((topic, event, payload)) = decode_phoenix_binary_frame(bytes)? else { + return Ok(None); + }; + if topic != CRF_SEARCH_TOPIC || event != "transfer_chunk" { + return Ok(None); + } + decode_binary_transfer_chunk(payload).map(Some) +} + +fn decode_phoenix_binary_frame(bytes: &[u8]) -> Result> { + if bytes.len() < 4 { + return Ok(None); + } + + let join_ref_size = bytes[0] as usize; + let ref_size = bytes[1] as usize; + let topic_size = bytes[2] as usize; + let event_size = bytes[3] as usize; + let join_ref_start = 4usize; + let ref_start = join_ref_start + .checked_add(join_ref_size) + .context("phoenix binary join_ref_size overflow")?; + let topic_start = ref_start + .checked_add(ref_size) + .context("phoenix binary ref_size overflow")?; + let event_start = topic_start + .checked_add(topic_size) + .context("phoenix binary topic_size overflow")?; + let payload_start = event_start + .checked_add(event_size) + .context("phoenix binary event_size overflow")?; + if bytes.len() < payload_start { + bail!( + "phoenix binary frame truncated: len={} header={}", + bytes.len(), + hex_prefix(bytes, 8) + ); + } + + let topic = std::str::from_utf8(&bytes[topic_start..event_start]) + .context("decode phoenix binary topic")?; + let event = std::str::from_utf8(&bytes[event_start..payload_start]) + .context("decode phoenix binary event")?; + Ok(Some((topic, event, &bytes[payload_start..]))) +} + +fn read_u64(bytes: &[u8], offset: usize) -> u64 { + u64::from_be_bytes( + bytes[offset..offset + 8] + .try_into() + .expect("read_u64 offset validated by fixed header length"), + ) +} + +fn read_u32(bytes: &[u8], offset: usize) -> u32 { + u32::from_be_bytes( + bytes[offset..offset + 4] + .try_into() + .expect("read_u32 offset validated by fixed header length"), + ) +} + +fn hex_prefix(bytes: &[u8], max_len: usize) -> String { + bytes + .iter() + .take(max_len) + .map(|byte| format!("{byte:02x}")) + .collect::>() + .join(" ") } fn websocket_connect_error(request_url: &str, error: WsError) -> anyhow::Error { @@ -599,6 +1038,14 @@ fn worker_websocket_url(base_url: &str, token: &str) -> Result { )) } +fn worker_websocket_config() -> WebSocketConfig { + WebSocketConfig { + max_message_size: Some(MAX_TRANSFER_FRAME_BYTES), + max_frame_size: Some(MAX_TRANSFER_FRAME_BYTES), + ..WebSocketConfig::default() + } +} + async fn send_json(writer: &mut W, value: T) -> Result<()> where W: SinkExt + Unpin, @@ -751,6 +1198,7 @@ mod tests { version: "0.11.4".into(), protocol_version: config.protocol_version, once: config.once, + local_path: None, } } @@ -768,6 +1216,7 @@ mod tests { version: "0.11.4".into(), protocol_version: 1, once: false, + local_path: None, }); assert_eq!(config.connect, "http://127.0.0.1:4000"); @@ -979,6 +1428,7 @@ mod tests { fn worker_job_lowering_uses_an_isolated_temp_dir_and_target_vmaf() { let job_dir = std::env::temp_dir().join(format!("ab-av1-worker-job-{}", std::process::id())); + let input_path = job_dir.join("movie.mkv"); let job = WorkerJob::new( JobAssignedPayload { status: WorkStatus::JobAssigned, @@ -990,6 +1440,7 @@ mod tests { target_vmaf: 96.5, }, job_dir.clone(), + input_path, ); let config = job @@ -1002,6 +1453,26 @@ mod tests { assert!(config.cache); } + #[test] + fn build_worker_job_uses_local_path_only_when_requested() { + let assignment = JobAssignedPayload { + status: WorkStatus::JobAssigned, + job_id: "job-123".into(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 1024, + chunk_size_bytes: 256, + target_vmaf: 96.5, + }; + let local_path = std::env::temp_dir() + .join(format!("ab-av1-worker-local-{}", std::process::id())) + .join("movie.mkv"); + + let job = build_worker_job(assignment, Some(local_path.as_path())).expect("worker job"); + + assert_eq!(job.input_path(), local_path.as_path()); + } + #[tokio::test(flavor = "current_thread")] async fn worker_job_runs_crf_search_from_fake_probe() -> Result<()> { crf_test_hooks::set(|_crf| sample_encode::Output { @@ -1015,6 +1486,7 @@ mod tests { let job_dir = std::env::temp_dir().join(format!("ab-av1-worker-exec-{}", std::process::id())); + let input_path = job_dir.join("movie.mkv"); let job = WorkerJob::new( JobAssignedPayload { status: WorkStatus::JobAssigned, @@ -1026,6 +1498,7 @@ mod tests { target_vmaf: 96.5, }, job_dir, + input_path, ); let probe = Arc::new(Ffprobe { @@ -1088,6 +1561,103 @@ mod tests { assert_eq!(backoff.next_delay(), Duration::from_millis(100)); } + #[test] + fn binary_transfer_chunk_decodes_rav1_frame() { + let data = b"hello worker bytes".to_vec(); + let frame = binary_transfer_chunk_frame( + "job-123", + 123, + 7, + 10, + 8 * 1024 + data.len() as u64, + 64 * 1024, + &data, + ); + + let chunk = decode_binary_transfer_chunk(&frame).expect("decode binary transfer chunk"); + + assert_eq!(chunk.transfer_id, "job-123"); + assert_eq!(chunk.video_id, 123); + assert_eq!(chunk.chunk_index, 7); + assert_eq!(chunk.total_chunks, 10); + assert_eq!(chunk.bytes_sent, 8 * 1024 + data.len() as u64); + assert_eq!(chunk.total_bytes, 64 * 1024); + assert_eq!(chunk.crc32, crc32fast::hash(&data) as u64); + assert_eq!(chunk.bytes, data); + } + + #[test] + fn binary_transfer_chunk_rejects_bad_magic() { + let mut frame = binary_transfer_chunk_frame("job-123", 123, 0, 1, 5, 5, b"hello"); + frame[0..4].copy_from_slice(b"NOPE"); + + assert!(decode_binary_transfer_chunk(&frame).is_err()); + } + + #[test] + fn binary_worker_push_decodes_phoenix_enveloped_transfer_chunk() { + let data = b"hello worker bytes".to_vec(); + let chunk = binary_transfer_chunk_frame( + "job-123", + 123, + 7, + 10, + 8 * 1024 + data.len() as u64, + 64 * 1024, + &data, + ); + let frame = phoenix_binary_frame("1", CRF_SEARCH_TOPIC, "transfer_chunk", &chunk); + + let chunk = decode_binary_worker_push(&frame) + .expect("decode phoenix binary push") + .expect("transfer chunk"); + + assert_eq!(chunk.transfer_id, "job-123"); + assert_eq!(chunk.video_id, 123); + assert_eq!(chunk.chunk_index, 7); + assert_eq!(chunk.bytes, data); + } + + fn binary_transfer_chunk_frame( + transfer_id: &str, + video_id: u64, + chunk_index: u64, + total_chunks: u64, + bytes_sent: u64, + total_bytes: u64, + data: &[u8], + ) -> Vec { + let mut frame = + Vec::with_capacity(TRANSFER_CHUNK_HEADER_LEN + transfer_id.len() + data.len()); + frame.extend_from_slice(TRANSFER_CHUNK_MAGIC); + frame.push(TRANSFER_CHUNK_VERSION); + frame.push(TRANSFER_CHUNK_TYPE); + frame.extend_from_slice(&(transfer_id.len() as u16).to_be_bytes()); + frame.extend_from_slice(&video_id.to_be_bytes()); + frame.extend_from_slice(&chunk_index.to_be_bytes()); + frame.extend_from_slice(&total_chunks.to_be_bytes()); + frame.extend_from_slice(&bytes_sent.to_be_bytes()); + frame.extend_from_slice(&total_bytes.to_be_bytes()); + frame.extend_from_slice(&crc32fast::hash(data).to_be_bytes()); + frame.extend_from_slice(transfer_id.as_bytes()); + frame.extend_from_slice(data); + frame + } + + fn phoenix_binary_frame(reference: &str, topic: &str, event: &str, payload: &[u8]) -> Vec { + let mut frame = + Vec::with_capacity(4 + reference.len() + topic.len() + event.len() + payload.len()); + frame.push(0); + frame.push(reference.len() as u8); + frame.push(topic.len() as u8); + frame.push(event.len() as u8); + frame.extend_from_slice(reference.as_bytes()); + frame.extend_from_slice(topic.as_bytes()); + frame.extend_from_slice(event.as_bytes()); + frame.extend_from_slice(payload); + frame + } + async fn expect_join(reader: &mut R) where R: StreamExt> diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index f453b805..e7ad5f13 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -188,16 +188,31 @@ pub(crate) struct CancelPayload { pub(crate) reason: String, } -/// Chunk bytes travel in binary websocket frames. -/// These metadata messages stay on the text side channel. #[cfg_attr(not(test), allow(dead_code))] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub(crate) struct ChunkTransferPayload { - pub(crate) job_id: String, - pub(crate) index: u64, - pub(crate) offset: u64, +pub(crate) struct TransferStartedPayload { + pub(crate) chunk_size_bytes: u64, pub(crate) size_bytes: u64, - pub(crate) checksum: String, + pub(crate) source_name: String, + pub(crate) status: String, + pub(crate) total_bytes: u64, + pub(crate) total_chunks: u64, + pub(crate) transfer_id: String, + pub(crate) video_id: u64, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ChunkTransferPayload { + pub(crate) bytes_sent: u64, + pub(crate) chunk_index: u64, + pub(crate) crc32: u64, + pub(crate) data: String, + pub(crate) status: String, + pub(crate) total_bytes: u64, + pub(crate) total_chunks: u64, + pub(crate) transfer_id: String, + pub(crate) video_id: u64, } #[allow(dead_code)] @@ -437,24 +452,60 @@ mod tests { ); } + #[test] + fn transfer_started_payload_serializes_metadata_side_channel() { + let payload = TransferStartedPayload { + chunk_size_bytes: 1_048_576, + size_bytes: 9_560_739_312, + source_name: "movie.mkv".into(), + status: "transfer_started".into(), + total_bytes: 9_560_739_312, + total_chunks: 9_118, + transfer_id: "job-123".into(), + video_id: 123, + }; + + assert_eq!( + serde_json::to_value(payload).expect("serialize transfer started"), + json!({ + "chunk_size_bytes": 1_048_576, + "size_bytes": 9_560_739_312u64, + "source_name": "movie.mkv", + "status": "transfer_started", + "total_bytes": 9_560_739_312u64, + "total_chunks": 9_118, + "transfer_id": "job-123", + "video_id": 123, + }) + ); + } + #[test] fn chunk_transfer_payload_serializes_metadata_side_channel() { let payload = ChunkTransferPayload { - job_id: "job-123".into(), - index: 7, - offset: 8192, - size_bytes: 4096, - checksum: "deadbeef".into(), + bytes_sent: 4096, + chunk_index: 7, + crc32: 0xdead_beef, + data: "deadbeef".into(), + status: "transfer_chunk".into(), + total_bytes: 9_560_739_312, + total_chunks: 9_118, + transfer_id: "job-123".into(), + video_id: 123, }; assert_eq!( serde_json::to_value(payload).expect("serialize chunk transfer"), json!({ - "job_id": "job-123", - "index": 7, - "offset": 8192, - "size_bytes": 4096, - "checksum": "deadbeef", + "bytes_sent": 4096, + "chunk_index": 7, + "crc32": 3735928559u64, + "data": "deadbeef", + "status": "transfer_chunk", + "total_bytes": 9_560_739_312u64, + "total_chunks": 9_118, + "transfer_id": "job-123", + "video_id": 123, }) ); } diff --git a/src/command/worker_transfer.rs b/src/command/worker_transfer.rs index 352a8ec0..913cacea 100644 --- a/src/command/worker_transfer.rs +++ b/src/command/worker_transfer.rs @@ -6,7 +6,7 @@ use std::{ io::Write, path::{Path, PathBuf}, }; -use tracing::debug; +use tracing::{debug, trace}; #[cfg_attr(not(test), allow(dead_code))] #[derive(Debug, Clone, PartialEq, Eq)] @@ -14,7 +14,7 @@ pub(crate) struct Chunk { pub index: u64, pub offset: u64, pub bytes: Vec, - pub checksum: Hash, + pub checksum: u64, } #[cfg_attr(not(test), allow(dead_code))] @@ -107,7 +107,7 @@ impl ChunkReceiver { if self.finished { return Err(ChunkReceiverError::Finished); } - debug!( + trace!( index = chunk.index, offset = chunk.offset, size = chunk.bytes.len(), @@ -138,7 +138,7 @@ impl ChunkReceiver { max_size: self.max_size.expect("checked max_size"), }); } - if blake3::hash(&chunk.bytes) != chunk.checksum { + if crc32fast::hash(&chunk.bytes) as u64 != chunk.checksum { return Err(ChunkReceiverError::CorruptChunk { index: chunk.index }); } @@ -221,7 +221,7 @@ mod tests { index, offset, bytes: bytes.to_vec(), - checksum: blake3::hash(bytes), + checksum: crc32fast::hash(bytes) as u64, } } @@ -293,7 +293,7 @@ mod tests { let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); let mut bad = chunk(0, 0, b"hello"); - bad.checksum = blake3::hash(b"hell0"); + bad.checksum = crc32fast::hash(b"hell0") as u64; assert!(matches!( receiver.push(bad), diff --git a/src/main.rs b/src/main.rs index dc723307..6ad1fdba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,7 +17,7 @@ use anyhow::anyhow; use clap::Parser; use futures_util::FutureExt; use tokio::signal; -use tracing_subscriber::{fmt, EnvFilter}; +use tracing_subscriber::{EnvFilter, fmt}; #[derive(Parser)] #[command(version, about)] From 0c2d9a5df2f00f8f56be1f3a7a5e343be718fa29 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 7 Jul 2026 18:14:49 -0600 Subject: [PATCH 2/7] Report worker CRF stats --- Cargo.lock | 80 +++++++++++++++++ Cargo.toml | 1 + src/command/worker.rs | 124 +++++++++++++++++++++++++-- src/command/worker_protocol.rs | 151 ++++++++++++++++++++++++++++++++- 4 files changed, 345 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 90090259..1d2cd034 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,6 +34,7 @@ dependencies = [ "serial_test", "shell-escape", "sled", + "sysinfo", "test-case", "thiserror 2.0.18", "tokio", @@ -412,6 +413,16 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + [[package]] name = "crossbeam-epoch" version = "0.9.18" @@ -480,6 +491,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "encode_unicode" version = "1.0.0" @@ -850,6 +867,15 @@ dependencies = [ "libc", ] +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1099,6 +1125,26 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.2.16" @@ -1520,6 +1566,21 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "sysinfo" +version = "0.30.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a5b4ddaee55fb2bea2bf0e5000747e5f5c0de765e5a5ff87f4cd106439f4bb3" +dependencies = [ + "cfg-if", + "core-foundation-sys", + "libc", + "ntapi", + "once_cell", + "rayon", + "windows", +] + [[package]] name = "tempfile" version = "3.27.0" @@ -2013,6 +2074,25 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e48a53791691ab099e5e2ad123536d0fff50652600abaf43bbf952894110d0be" +dependencies = [ + "windows-core", + "windows-targets", +] + +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/Cargo.toml b/Cargo.toml index d11413e8..816bc441 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ serde = { version = "1.0.185", features = ["derive"] } serde_json = "1.0.105" shell-escape = "0.1.5" sled = "0.34.7" +sysinfo = "0.30" thiserror = "2" rustls = { version = "0.23", features = ["ring"] } tracing = "0.1" diff --git a/src/command/worker.rs b/src/command/worker.rs index 9ae7d947..d21208d4 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -1,7 +1,8 @@ use crate::command::worker_protocol::{ AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, - ClientEvent, ClientFrame, ErrorReplyPayload, JobResultPayload, ReplyBody, ServerPushFrame, - ServerReply, TransferStartedPayload, + ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, + HeartbeatPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, + TransferStartedPayload, }; use crate::command::worker_transfer::{Chunk, ChunkReceiver}; use crate::command::{args, crf_search, sample_encode}; @@ -12,12 +13,13 @@ use base64::{Engine as _, engine::general_purpose::STANDARD}; use clap::Parser; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Value, json}; use std::{ path::{Path, PathBuf}, sync::Arc, - time::Duration, + time::{Duration, Instant}, }; +use sysinfo::{Disks, Pid, System}; use tokio::net::TcpStream; use tokio_tungstenite::{ MaybeTlsStream, WebSocketStream, @@ -33,6 +35,7 @@ const TRANSFER_CHUNK_VERSION: u8 = 1; const TRANSFER_CHUNK_TYPE: u8 = 1; const TRANSFER_CHUNK_HEADER_LEN: usize = 52; const MAX_TRANSFER_FRAME_BYTES: usize = 640 * 1024 * 1024; +const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30); /// Connect to a Reencodarr websocket worker endpoint and request one job. #[derive(Parser, Debug, Clone)] @@ -194,6 +197,33 @@ impl WorkerJob { from_cache: best.enc.from_cache, } } + + fn progress_payload(&self, status: &sample_encode::Status) -> CrfSearchProgressPayload { + CrfSearchProgressPayload { + video_id: self.assignment.video_id, + percent: (status.progress.clamp(0.0, 1.0) * 100.0), + filename: self.assignment.source_name.clone(), + eta: None, + fps: status.fps, + } + } + + fn crf_result_payload( + &self, + sample: &crf_search::Sample, + chosen: bool, + ) -> CrfSearchResultPayload { + CrfSearchResultPayload { + crf: sample.crf, + score: sample.enc.single_score(), + percent: sample.enc.encode_percent, + size: sample.enc.predicted_encode_size, + time: sample.enc.predicted_encode_time.as_secs_f64(), + params: json!({ "encoder": "libsvtav1", "preset": 8 }), + target: self.assignment.target_vmaf, + chosen, + } + } } #[cfg_attr(not(test), allow(dead_code))] @@ -334,6 +364,78 @@ where unreachable!("crf-search stream should finish with Done") } +async fn run_worker_job_with_reporting( + job: WorkerJob, + probe: Arc, + worker: &mut ConnectedWorker, +) -> Result { + let config = job.crf_search_config("libsvtav1".parse().expect("default encoder"))?; + let mut run = std::pin::pin!(crf_search::run(config, probe)); + let mut last_heartbeat = Instant::now() - HEARTBEAT_INTERVAL; + + loop { + if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL { + worker + .send_event(ClientEvent::Heartbeat(heartbeat_payload(&job.input_dir))) + .await?; + last_heartbeat = Instant::now(); + } + + match run.next().await { + Some(Ok(crf_search::Update::Done(best))) => { + worker + .send_event(ClientEvent::CrfSearchResult( + job.crf_result_payload(&best, true), + )) + .await?; + return Ok(best); + } + Some(Ok(crf_search::Update::Status { sample, .. })) => { + worker + .send_event(ClientEvent::CrfSearchProgress( + job.progress_payload(&sample), + )) + .await?; + } + Some(Ok(crf_search::Update::SampleResult { .. })) => {} + Some(Ok(crf_search::Update::RunResult(sample))) => { + worker + .send_event(ClientEvent::CrfSearchResult( + job.crf_result_payload(&sample, false), + )) + .await?; + } + Some(Err(error)) => return Err(error.into()), + None => break, + } + } + + unreachable!("crf-search stream should finish with Done") +} + +fn heartbeat_payload(path: &Path) -> HeartbeatPayload { + let mut system = System::new(); + system.refresh_cpu(); + system.refresh_memory(); + + let pid = Pid::from_u32(std::process::id()); + let memory_rss_bytes = system.process(pid).map(|process| process.memory()); + + let disks = Disks::new_with_refreshed_list(); + let disk = disks + .iter() + .filter(|disk| path.starts_with(disk.mount_point())) + .max_by_key(|disk| disk.mount_point().as_os_str().len()); + + HeartbeatPayload { + cpu_percent: Some(system.global_cpu_info().cpu_usage()), + memory_rss_bytes, + memory_total_bytes: Some(system.total_memory()), + disk_free_bytes: disk.map(|disk| disk.available_space()), + disk_total_bytes: disk.map(|disk| disk.total_space()), + } +} + fn worker_job_input_dir(job_id: &str) -> PathBuf { std::env::current_dir() .expect("current working directory") @@ -501,6 +603,12 @@ impl ConnectedWorker { expect_reply(&mut self.socket, &request_ref.to_string(), "pull_work").await } + async fn send_event(&mut self, event: ClientEvent) -> Result<()> { + let request_ref = self.next_ref; + self.next_ref += 1; + send_json(&mut self.socket, ClientFrame::new(request_ref, event)).await + } + async fn wait_for_pending_job( &mut self, pending_job: &mut PendingJob, @@ -646,7 +754,7 @@ impl ConnectedWorker { } } -async fn run_worker_job_and_publish(job: &WorkerJob) -> Result<()> { +async fn run_worker_job_and_publish(worker: &mut ConnectedWorker, job: &WorkerJob) -> Result<()> { debug!( job_id = %job.assignment.job_id, input = %job.input_path().display(), @@ -654,7 +762,7 @@ async fn run_worker_job_and_publish(job: &WorkerJob) -> Result<()> { ); let probe = Arc::new(crate::ffprobe::probe(job.input_path())); debug!(job_id = %job.assignment.job_id, "probe complete, running crf search"); - let best = run_worker_job(job.clone(), probe).await?; + let best = run_worker_job_with_reporting(job.clone(), probe, worker).await?; debug!(job_id = %job.assignment.job_id, "publishing worker result"); println!( @@ -749,7 +857,7 @@ async fn run_connected_worker( PendingJobOutcome::Ready => { let job = pending_job.take().expect("pending job"); debug!(job_id = %job.job.assignment.job_id, "pending job input arrived"); - run_worker_job_and_publish(&job.job).await?; + run_worker_job_and_publish(&mut worker, &job.job).await?; continue; } } @@ -773,7 +881,7 @@ async fn run_connected_worker( ); if job.input_path().exists() { debug!(job_id = %job.assignment.job_id, "input already present, starting job"); - run_worker_job_and_publish(&job).await?; + run_worker_job_and_publish(&mut worker, &job).await?; } else if config.local_path.is_some() { bail!( "local input path does not exist: {}", diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index e7ad5f13..9a884acb 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize}; pub(crate) const CRF_SEARCH_TOPIC: &str = "workers:crf_search"; -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Serialize)] pub(crate) struct ClientFrame(String, String, String, String, ClientPayload); impl ClientFrame { @@ -18,11 +18,14 @@ impl ClientFrame { } } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub(crate) enum ClientEvent { Join, Announce(AnnouncePayload), PullWork, + Heartbeat(HeartbeatPayload), + CrfSearchProgress(CrfSearchProgressPayload), + CrfSearchResult(CrfSearchResultPayload), } impl ClientEvent { @@ -31,15 +34,23 @@ impl ClientEvent { Self::Join => ("phx_join", ClientPayload::Empty(EmptyPayload {})), Self::Announce(payload) => ("announce", ClientPayload::Announce(payload)), Self::PullWork => ("pull_work", ClientPayload::Empty(EmptyPayload {})), + Self::Heartbeat(payload) => ("heartbeat", ClientPayload::Heartbeat(payload)), + Self::CrfSearchProgress(payload) => { + ("crf_search_progress", ClientPayload::Progress(payload)) + } + Self::CrfSearchResult(payload) => ("crf_search_result", ClientPayload::Result(payload)), } } } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Serialize)] #[serde(untagged)] enum ClientPayload { Empty(EmptyPayload), Announce(AnnouncePayload), + Heartbeat(HeartbeatPayload), + Progress(CrfSearchProgressPayload), + Result(CrfSearchResultPayload), } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -58,6 +69,41 @@ pub(crate) struct Capabilities { pub(crate) crf_search: bool, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct HeartbeatPayload { + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) cpu_percent: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) memory_rss_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) memory_total_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) disk_free_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) disk_total_bytes: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct CrfSearchProgressPayload { + pub(crate) video_id: u64, + pub(crate) percent: f32, + pub(crate) filename: String, + pub(crate) eta: Option, + pub(crate) fps: f32, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct CrfSearchResultPayload { + pub(crate) crf: f32, + pub(crate) score: f32, + pub(crate) percent: f64, + pub(crate) size: u64, + pub(crate) time: f64, + pub(crate) params: serde_json::Value, + pub(crate) target: f32, + pub(crate) chosen: bool, +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct ServerFrame( pub(crate) Option, @@ -301,6 +347,105 @@ mod tests { ); } + #[test] + fn heartbeat_serializes_worker_telemetry_event() { + let frame = ClientFrame::new( + 4, + ClientEvent::Heartbeat(HeartbeatPayload { + cpu_percent: Some(12.5), + memory_rss_bytes: Some(1234), + memory_total_bytes: Some(8192), + disk_free_bytes: Some(4096), + disk_total_bytes: Some(16_384), + }), + ); + + assert_eq!( + serde_json::to_value(frame).expect("serialize heartbeat"), + json!([ + "1", + "4", + "workers:crf_search", + "heartbeat", + { + "cpu_percent": 12.5, + "memory_rss_bytes": 1234, + "memory_total_bytes": 8192, + "disk_free_bytes": 4096, + "disk_total_bytes": 16_384, + } + ]) + ); + } + + #[test] + fn crf_search_progress_serializes_reencodarr_progress_event() { + let frame = ClientFrame::new( + 5, + ClientEvent::CrfSearchProgress(CrfSearchProgressPayload { + video_id: 123, + percent: 42.5, + filename: "movie.mkv".into(), + eta: None, + fps: 27.25, + }), + ); + + assert_eq!( + serde_json::to_value(frame).expect("serialize crf progress"), + json!([ + "1", + "5", + "workers:crf_search", + "crf_search_progress", + { + "video_id": 123, + "percent": 42.5, + "filename": "movie.mkv", + "eta": null, + "fps": 27.25, + } + ]) + ); + } + + #[test] + fn crf_search_result_serializes_reencodarr_vmaf_model_event() { + let frame = ClientFrame::new( + 6, + ClientEvent::CrfSearchResult(CrfSearchResultPayload { + crf: 31.5, + score: 96.2, + percent: 42.5, + size: 123_456, + time: 87.5, + params: json!({ "encoder": "libsvtav1", "preset": 8 }), + target: 95.0, + chosen: true, + }), + ); + + assert_eq!( + serde_json::to_value(frame).expect("serialize crf result"), + json!([ + "1", + "6", + "workers:crf_search", + "crf_search_result", + { + "crf": 31.5, + "score": 96.19999694824219, + "percent": 42.5, + "size": 123456, + "time": 87.5, + "params": { "encoder": "libsvtav1", "preset": 8 }, + "target": 95.0, + "chosen": true, + } + ]) + ); + } + #[test] fn server_reply_parses_current_no_work_payload() { let reply: ServerFrame = serde_json::from_value(json!([ From fda92ab7e6709a76fef53fd27080aee552424c0c Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 7 Jul 2026 18:33:19 -0600 Subject: [PATCH 3/7] Keep worker socket alive during jobs --- src/command/worker.rs | 95 +++++++++++++++++++++++++++++-------------- 1 file changed, 64 insertions(+), 31 deletions(-) diff --git a/src/command/worker.rs b/src/command/worker.rs index d21208d4..919f8e06 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -17,7 +17,7 @@ use serde_json::{Value, json}; use std::{ path::{Path, PathBuf}, sync::Arc, - time::{Duration, Instant}, + time::Duration, }; use sysinfo::{Disks, Pid, System}; use tokio::net::TcpStream; @@ -371,48 +371,81 @@ async fn run_worker_job_with_reporting( ) -> Result { let config = job.crf_search_config("libsvtav1".parse().expect("default encoder"))?; let mut run = std::pin::pin!(crf_search::run(config, probe)); - let mut last_heartbeat = Instant::now() - HEARTBEAT_INTERVAL; + let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { - if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL { - worker - .send_event(ClientEvent::Heartbeat(heartbeat_payload(&job.input_dir))) - .await?; - last_heartbeat = Instant::now(); - } - - match run.next().await { - Some(Ok(crf_search::Update::Done(best))) => { - worker - .send_event(ClientEvent::CrfSearchResult( - job.crf_result_payload(&best, true), - )) - .await?; - return Ok(best); - } - Some(Ok(crf_search::Update::Status { sample, .. })) => { + tokio::select! { + _ = heartbeat.tick() => { worker - .send_event(ClientEvent::CrfSearchProgress( - job.progress_payload(&sample), - )) + .send_event(ClientEvent::Heartbeat(heartbeat_payload(&job.input_dir))) .await?; } - Some(Ok(crf_search::Update::SampleResult { .. })) => {} - Some(Ok(crf_search::Update::RunResult(sample))) => { - worker - .send_event(ClientEvent::CrfSearchResult( - job.crf_result_payload(&sample, false), - )) - .await?; + frame = worker.socket.next() => { + handle_job_websocket_frame(&mut worker.socket, frame, &job.assignment.job_id).await?; } - Some(Err(error)) => return Err(error.into()), - None => break, + update = run.next() => match update { + Some(Ok(crf_search::Update::Done(best))) => { + worker + .send_event(ClientEvent::CrfSearchResult( + job.crf_result_payload(&best, true), + )) + .await?; + return Ok(best); + } + Some(Ok(crf_search::Update::Status { sample, .. })) => { + worker + .send_event(ClientEvent::CrfSearchProgress( + job.progress_payload(&sample), + )) + .await?; + } + Some(Ok(crf_search::Update::SampleResult { .. })) => {} + Some(Ok(crf_search::Update::RunResult(sample))) => { + worker + .send_event(ClientEvent::CrfSearchResult( + job.crf_result_payload(&sample, false), + )) + .await?; + } + Some(Err(error)) => return Err(error.into()), + None => break, + }, } } unreachable!("crf-search stream should finish with Done") } +async fn handle_job_websocket_frame( + socket: &mut WorkerSocket, + frame: Option>, + job_id: &str, +) -> Result<()> { + match frame { + Some(Ok(Message::Ping(payload))) => { + socket + .send(Message::Pong(payload)) + .await + .context("send websocket pong during worker job")?; + } + Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Text(text))) => { + if let Some(WorkerPush::Cancel(cancel)) = decode_worker_push(&text)? + && cancel.job_id == job_id + { + bail!("worker job {} canceled: {}", cancel.job_id, cancel.reason); + } + } + Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => {} + Some(Ok(Message::Close(frame))) => bail!("websocket closed during worker job: {frame:?}"), + Some(Err(error)) => return Err(error).context("read websocket message during worker job"), + None => bail!("websocket ended during worker job"), + } + + Ok(()) +} + fn heartbeat_payload(path: &Path) -> HeartbeatPayload { let mut system = System::new(); system.refresh_cpu(); From 73f54afe9b5096e66ee61e650a9f2a6159d154e8 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 7 Jul 2026 18:41:30 -0600 Subject: [PATCH 4/7] Report worker transfer progress --- src/command/worker.rs | 58 ++++++++++++++++++++++++++++++-- src/command/worker_protocol.rs | 61 ++++++++++++++++++++++++++++++++-- 2 files changed, 115 insertions(+), 4 deletions(-) diff --git a/src/command/worker.rs b/src/command/worker.rs index 919f8e06..9abe932a 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -2,7 +2,7 @@ use crate::command::worker_protocol::{ AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, HeartbeatPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, - TransferStartedPayload, + TransferProgressPayload, TransferStartedPayload, }; use crate::command::worker_transfer::{Chunk, ChunkReceiver}; use crate::command::{args, crf_search, sample_encode}; @@ -17,7 +17,7 @@ use serde_json::{Value, json}; use std::{ path::{Path, PathBuf}, sync::Arc, - time::Duration, + time::{Duration, Instant}, }; use sysinfo::{Disks, Pid, System}; use tokio::net::TcpStream; @@ -231,6 +231,7 @@ impl WorkerJob { struct PendingJob { job: WorkerJob, receiver: Option, + transfer_started_at: Instant, } #[cfg_attr(not(test), allow(dead_code))] @@ -239,6 +240,7 @@ impl PendingJob { Self { job, receiver: Some(receiver), + transfer_started_at: Instant::now(), } } @@ -312,6 +314,46 @@ impl PendingJob { Ok(()) } + fn transfer_progress_payload( + &self, + chunk_index: u64, + total_chunks: u64, + ) -> TransferProgressPayload { + let received_bytes = self + .receiver + .as_ref() + .map(ChunkReceiver::received_bytes) + .unwrap_or(self.job.assignment.size_bytes); + let expected_bytes = Some(self.job.assignment.size_bytes); + let elapsed = self.transfer_started_at.elapsed().as_secs_f64().max(0.001); + let bytes_per_second = received_bytes as f64 / elapsed; + let remaining_bytes = self + .job + .assignment + .size_bytes + .saturating_sub(received_bytes); + let eta = (bytes_per_second > 0.0).then_some(remaining_bytes as f64 / bytes_per_second); + let percent = if self.job.assignment.size_bytes == 0 { + 100.0 + } else { + 100.0 * received_bytes as f64 / self.job.assignment.size_bytes as f64 + }; + + TransferProgressPayload { + job_id: self.job.assignment.job_id.clone(), + transfer_id: self.job.assignment.job_id.clone(), + video_id: self.job.assignment.video_id, + filename: self.job.assignment.source_name.clone(), + received_bytes, + expected_bytes, + percent, + bytes_per_second, + eta, + chunk_index, + total_chunks, + } + } + fn finish(&mut self) -> Result<()> { let final_path = self.input_path().to_path_buf(); let receiver = self.receiver.take().context("missing chunk receiver")?; @@ -705,7 +747,13 @@ impl ConnectedWorker { "received chunk" ); } + let chunk_index = chunk.chunk_index; + let total_chunks = chunk.total_chunks; pending_job.apply_chunk(chunk)?; + self.send_event(ClientEvent::TransferProgress( + pending_job.transfer_progress_payload(chunk_index, total_chunks), + )) + .await?; if pending_job.receiver.as_ref().is_some_and(|receiver| { receiver.received_bytes() == pending_job.job.assignment.size_bytes @@ -747,7 +795,13 @@ impl ConnectedWorker { "received binary chunk" ); } + let chunk_index = chunk.chunk_index; + let total_chunks = chunk.total_chunks; pending_job.apply_raw_chunk(chunk)?; + self.send_event(ClientEvent::TransferProgress( + pending_job.transfer_progress_payload(chunk_index, total_chunks), + )) + .await?; if pending_job.receiver.as_ref().is_some_and(|receiver| { receiver.received_bytes() == pending_job.job.assignment.size_bytes }) { diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index 9a884acb..f8587b57 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -24,6 +24,7 @@ pub(crate) enum ClientEvent { Announce(AnnouncePayload), PullWork, Heartbeat(HeartbeatPayload), + TransferProgress(TransferProgressPayload), CrfSearchProgress(CrfSearchProgressPayload), CrfSearchResult(CrfSearchResultPayload), } @@ -35,6 +36,10 @@ impl ClientEvent { Self::Announce(payload) => ("announce", ClientPayload::Announce(payload)), Self::PullWork => ("pull_work", ClientPayload::Empty(EmptyPayload {})), Self::Heartbeat(payload) => ("heartbeat", ClientPayload::Heartbeat(payload)), + Self::TransferProgress(payload) => ( + "transfer_progress", + ClientPayload::TransferProgress(payload), + ), Self::CrfSearchProgress(payload) => { ("crf_search_progress", ClientPayload::Progress(payload)) } @@ -49,6 +54,7 @@ enum ClientPayload { Empty(EmptyPayload), Announce(AnnouncePayload), Heartbeat(HeartbeatPayload), + TransferProgress(TransferProgressPayload), Progress(CrfSearchProgressPayload), Result(CrfSearchResultPayload), } @@ -261,12 +267,20 @@ pub(crate) struct ChunkTransferPayload { pub(crate) video_id: u64, } -#[allow(dead_code)] -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct TransferProgressPayload { pub(crate) job_id: String, + pub(crate) transfer_id: String, + pub(crate) video_id: u64, + pub(crate) filename: String, pub(crate) received_bytes: u64, pub(crate) expected_bytes: Option, + pub(crate) percent: f64, + pub(crate) bytes_per_second: f64, + pub(crate) eta: Option, + pub(crate) chunk_index: u64, + pub(crate) total_chunks: u64, } #[allow(dead_code)] @@ -378,6 +392,49 @@ mod tests { ); } + #[test] + fn transfer_progress_serializes_transfer_stats_event() { + let frame = ClientFrame::new( + 5, + ClientEvent::TransferProgress(TransferProgressPayload { + job_id: "job-123".into(), + transfer_id: "job-123".into(), + video_id: 123, + filename: "movie.mkv".into(), + received_bytes: 512, + expected_bytes: Some(1024), + percent: 50.0, + bytes_per_second: 256.0, + eta: Some(2.0), + chunk_index: 3, + total_chunks: 8, + }), + ); + + assert_eq!( + serde_json::to_value(frame).expect("serialize transfer progress"), + json!([ + "1", + "5", + "workers:crf_search", + "transfer_progress", + { + "job_id": "job-123", + "transfer_id": "job-123", + "video_id": 123, + "filename": "movie.mkv", + "received_bytes": 512, + "expected_bytes": 1024, + "percent": 50.0, + "bytes_per_second": 256.0, + "eta": 2.0, + "chunk_index": 3, + "total_chunks": 8, + } + ]) + ); + } + #[test] fn crf_search_progress_serializes_reencodarr_progress_event() { let frame = ClientFrame::new( From b2ec022fcadc73d2be2b14a5eff017cf0048de98 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 8 Jul 2026 21:47:38 -0600 Subject: [PATCH 5/7] Report worker transfer and CRF progress --- src/command/worker.rs | 819 ++++++++++++++++++++++++++------- src/command/worker_protocol.rs | 177 ++++++- src/command/worker_transfer.rs | 118 +++-- src/process/managed.rs | 15 +- 4 files changed, 932 insertions(+), 197 deletions(-) diff --git a/src/command/worker.rs b/src/command/worker.rs index 9abe932a..29d8cef6 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -2,12 +2,12 @@ use crate::command::worker_protocol::{ AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, HeartbeatPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, - TransferProgressPayload, TransferStartedPayload, + TransferFailurePayload, TransferProgressPayload, TransferStage, TransferStartedPayload, + WorkStatus, }; use crate::command::worker_transfer::{Chunk, ChunkReceiver}; -use crate::command::{args, crf_search, sample_encode}; +use crate::command::{crf_search, sample_encode}; use crate::ffprobe::Ffprobe; -use crate::temporary; use anyhow::{Context, Result, anyhow, bail}; use base64::{Engine as _, engine::general_purpose::STANDARD}; use clap::Parser; @@ -15,8 +15,9 @@ use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use std::{ + fs, path::{Path, PathBuf}, - sync::Arc, + sync::{Arc, Mutex, OnceLock}, time::{Duration, Instant}, }; use sysinfo::{Disks, Pid, System}; @@ -35,7 +36,8 @@ const TRANSFER_CHUNK_VERSION: u8 = 1; const TRANSFER_CHUNK_TYPE: u8 = 1; const TRANSFER_CHUNK_HEADER_LEN: usize = 52; const MAX_TRANSFER_FRAME_BYTES: usize = 640 * 1024 * 1024; -const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30); +const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); +static HEARTBEAT_SYSTEM: OnceLock> = OnceLock::new(); /// Connect to a Reencodarr websocket worker endpoint and request one job. #[derive(Parser, Debug, Clone)] @@ -138,49 +140,27 @@ impl WorkerJob { &self.input_path } - fn crf_search_config(&self, encoder: args::Encoder) -> Result { - Ok(crf_search::CrfSearchConfig { - args: args::Encode { - encoder, - input: self.input_path().to_path_buf(), - vfilter: None, - pix_format: None, - preset: None, - keyint: None, - scd: None, - svt_args: vec![], - enc_args: vec![], - enc_input_args: vec![], - }, - min_vmaf: Some(crf_search::MinScore::new(self.assignment.target_vmaf)?), - min_xpsnr: None, - max_encoded_percent: crf_search::MaxEncodedPercent::new(80.0)?, - min_crf: None, - max_crf: None, - thorough: false, - crf_increment: None, - high_crf_means_hq: None, - cache: true, - sample: args::Sample { - samples: None, - sample_every: args::SampleDuration::new(Duration::from_secs(12 * 60))?, - min_samples: None, - sample_duration: args::SampleDuration::new(Duration::from_secs(20))?, - keep: false, - temp_dir: Some(self.input_dir.clone()), - extension: None, - }, - scoring: sample_encode::ScoringConfig { - score: args::ScoreArgs { - reference_vfilter: None, - } - .into(), - vmaf: args::Vmaf::default().into(), - xpsnr: false, - xpsnr_opts: args::Xpsnr::default().into(), - }, - verbose: clap_verbosity_flag::Verbosity::new(0, 0), - }) + fn crf_search_config(&self) -> Result { + let mut argv = self.assignment.crf_search_args.clone(); + if argv.is_empty() { + bail!( + "job {} missing crf_search_args; server must provide CRF search arguments", + self.assignment.job_id + ); + } + + if argv.first().is_some_and(|arg| arg == "ab-av1") { + argv.remove(0); + } + if argv.first().is_some_and(|arg| arg == "crf-search") { + argv.remove(0); + } + argv.insert(0, "crf-search".into()); + + let mut config = crf_search::CrfSearchConfig::from(crf_search::Args::try_parse_from(argv)?); + config.args.input = self.input_path().to_path_buf(); + config.sample.temp_dir = Some(self.input_dir.clone()); + Ok(config) } fn result_payload(&self, best: &crf_search::Sample) -> JobResultPayload { @@ -214,7 +194,16 @@ impl WorkerJob { chosen: bool, ) -> CrfSearchResultPayload { CrfSearchResultPayload { + job_id: self.assignment.job_id.clone(), + video_id: self.assignment.video_id, + source_name: self.assignment.source_name.clone(), crf: sample.crf, + vmaf_score: sample.enc.vmaf_score, + xpsnr_score: sample.enc.xpsnr_score, + predicted_encode_size: sample.enc.predicted_encode_size, + encode_percent: sample.enc.encode_percent, + predicted_encode_time_secs: sample.enc.predicted_encode_time.as_secs_f64(), + from_cache: sample.enc.from_cache, score: sample.enc.single_score(), percent: sample.enc.encode_percent, size: sample.enc.predicted_encode_size, @@ -226,6 +215,18 @@ impl WorkerJob { } } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +struct WorkerJobReportState { + connected: bool, + #[serde(skip_serializing_if = "Option::is_none")] + heartbeat: Option, + #[serde(skip_serializing_if = "Option::is_none")] + transfer_progress: Option, + #[serde(skip_serializing_if = "Option::is_none")] + crf_progress: Option, + crf_results: Vec, +} + #[cfg_attr(not(test), allow(dead_code))] #[derive(Debug)] struct PendingJob { @@ -236,10 +237,10 @@ struct PendingJob { #[cfg_attr(not(test), allow(dead_code))] impl PendingJob { - fn new(job: WorkerJob, receiver: ChunkReceiver) -> Self { + fn waiting(job: WorkerJob) -> Self { Self { job, - receiver: Some(receiver), + receiver: None, transfer_started_at: Instant::now(), } } @@ -268,6 +269,22 @@ impl PendingJob { }) } + fn ensure_receiver(&mut self, chunk_size_bytes: u64) -> Result<()> { + if self.receiver.is_some() { + return Ok(()); + } + + let input_path = self.job.input_path().to_path_buf(); + let input_dir = self.job.input_dir.clone(); + let size_bytes = self.job.assignment.size_bytes; + self.receiver = Some( + ChunkReceiver::new(input_path, &input_dir, Some(size_bytes), chunk_size_bytes) + .context("prepare worker input transfer")?, + ); + self.transfer_started_at = Instant::now(); + Ok(()) + } + fn apply_raw_chunk(&mut self, chunk: TransferChunk) -> Result<()> { if chunk.transfer_id != self.job.assignment.job_id { bail!( @@ -290,6 +307,10 @@ impl PendingJob { chunk.total_bytes ); } + if self.receiver.is_none() { + self.ensure_receiver(chunk.bytes.len() as u64) + .context("prepare worker input transfer from first chunk")?; + } let offset = self .receiver .as_ref() @@ -385,7 +406,7 @@ async fn run_worker_job_until( where S: std::future::Future, { - let config = job.crf_search_config("libsvtav1".parse().expect("default encoder"))?; + let config = job.crf_search_config()?; let mut run = std::pin::pin!(crf_search::run(config, probe)); tokio::pin!(shutdown); @@ -407,93 +428,291 @@ where } async fn run_worker_job_with_reporting( + config: &WorkerConfig, job: WorkerJob, probe: Arc, - worker: &mut ConnectedWorker, + worker: &mut Option, ) -> Result { - let config = job.crf_search_config("libsvtav1".parse().expect("default encoder"))?; - let mut run = std::pin::pin!(crf_search::run(config, probe)); + let crf_config = job.crf_search_config()?; + let mut run = std::pin::pin!(crf_search::run(crf_config, probe)); let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL); heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut state = WorkerJobReportState { + connected: true, + ..WorkerJobReportState::default() + }; + let mut reconnect = tokio::time::interval(Duration::from_secs(5)); + reconnect.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { - tokio::select! { - _ = heartbeat.tick() => { - worker - .send_event(ClientEvent::Heartbeat(heartbeat_payload(&job.input_dir))) - .await?; - } - frame = worker.socket.next() => { - handle_job_websocket_frame(&mut worker.socket, frame, &job.assignment.job_id).await?; - } - update = run.next() => match update { - Some(Ok(crf_search::Update::Done(best))) => { - worker - .send_event(ClientEvent::CrfSearchResult( - job.crf_result_payload(&best, true), - )) - .await?; - return Ok(best); - } - Some(Ok(crf_search::Update::Status { sample, .. })) => { - worker - .send_event(ClientEvent::CrfSearchProgress( - job.progress_payload(&sample), - )) - .await?; + match worker.as_mut() { + Some(current_worker) => { + tokio::select! { + _ = heartbeat.tick() => { + let heartbeat = heartbeat_payload(&job.input_dir, Some(job.assignment.video_id)); + state.heartbeat = Some(heartbeat.clone()); + state.connected = true; + if let Some(current_worker) = worker.as_mut() { + debug!( + job_id = %job.assignment.job_id, + active_video_id = job.assignment.video_id, + "sending worker heartbeat" + ); + if let Err(error) = current_worker + .send_event(ClientEvent::Heartbeat(heartbeat)) + .await + { + debug!( + job_id = %job.assignment.job_id, + error = %error, + "worker heartbeat failed; reconnecting while job continues" + ); + state.connected = false; + *worker = None; + } + } + } + frame = current_worker.socket.next() => { + match frame { + Some(Ok(Message::Ping(payload))) => { + current_worker + .socket + .send(Message::Pong(payload)) + .await + .context("send websocket pong")?; + } + Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Text(text))) => { + match decode_worker_push(&text)? { + Some(WorkerPush::Cancel(cancel)) + if cancel.job_id == job.assignment.job_id => + { + eprintln!( + "worker job {} canceled: {}", + cancel.job_id, cancel.reason + ); + return Err(anyhow!( + "worker job {} canceled: {}", + cancel.job_id, cancel.reason + )); + } + _ => {} + } + } + Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => {} + Some(Ok(Message::Close(frame))) => { + debug!(job_id = %job.assignment.job_id, ?frame, "worker socket closed during job"); + state.connected = false; + *worker = None; + } + Some(Err(error)) => { + debug!(job_id = %job.assignment.job_id, error = %error, "worker socket lost during job"); + state.connected = false; + *worker = None; + } + None => { + debug!(job_id = %job.assignment.job_id, "worker websocket ended during job"); + state.connected = false; + *worker = None; + } + } + } + update = run.next() => { + let (best, disconnected) = handle_crf_update( + &job, + &mut state, + Some(current_worker), + update, + ).await?; + if disconnected { + state.connected = false; + *worker = None; + } + if let Some(best) = best { + return Ok(best); + } + } } - Some(Ok(crf_search::Update::SampleResult { .. })) => {} - Some(Ok(crf_search::Update::RunResult(sample))) => { - worker - .send_event(ClientEvent::CrfSearchResult( - job.crf_result_payload(&sample, false), - )) - .await?; + } + None => { + tokio::select! { + _ = heartbeat.tick() => { + state.heartbeat = Some(heartbeat_payload(&job.input_dir, Some(job.assignment.video_id))); + state.connected = false; + } + _ = reconnect.tick() => { + match ConnectedWorker::connect(config).await { + Ok(mut reconnected) => { + replay_worker_state(&mut reconnected, &state).await; + state.connected = true; + *worker = Some(reconnected); + } + Err(error) => { + trace!(job_id = %job.assignment.job_id, error = %error, "worker reconnect attempt failed"); + } + } + } + update = run.next() => { + let (best, _) = handle_crf_update( + &job, + &mut state, + None, + update, + ).await?; + if let Some(best) = best { + return Ok(best); + } + } } - Some(Err(error)) => return Err(error.into()), - None => break, - }, + } } } - - unreachable!("crf-search stream should finish with Done") } -async fn handle_job_websocket_frame( - socket: &mut WorkerSocket, - frame: Option>, - job_id: &str, -) -> Result<()> { - match frame { - Some(Ok(Message::Ping(payload))) => { - socket - .send(Message::Pong(payload)) - .await - .context("send websocket pong during worker job")?; +async fn handle_crf_update( + job: &WorkerJob, + state: &mut WorkerJobReportState, + worker: Option<&mut ConnectedWorker>, + update: Option>, +) -> Result<(Option, bool)> { + let Some(update) = update else { + return Ok((None, false)); + }; + + match update { + Ok(crf_search::Update::Done(best)) => { + let payload = job.crf_result_payload(&best, true); + state.crf_results.push(payload.clone()); + let mut disconnected = false; + if let Some(worker) = worker { + disconnected = !send_worker_event( + worker, + ClientEvent::CrfSearchResult(payload), + &job.assignment.job_id, + "crf_result", + ) + .await; + } + Ok((Some(best), disconnected)) } - Some(Ok(Message::Pong(_))) => {} - Some(Ok(Message::Text(text))) => { - if let Some(WorkerPush::Cancel(cancel)) = decode_worker_push(&text)? - && cancel.job_id == job_id - { - bail!("worker job {} canceled: {}", cancel.job_id, cancel.reason); + Ok(crf_search::Update::Status { sample, .. }) => { + let payload = job.progress_payload(&sample); + state.crf_progress = Some(payload.clone()); + let mut disconnected = false; + if let Some(worker) = worker { + disconnected = !send_worker_event( + worker, + ClientEvent::CrfSearchProgress(payload), + &job.assignment.job_id, + "crf_progress", + ) + .await; } + Ok((None, disconnected)) } - Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => {} - Some(Ok(Message::Close(frame))) => bail!("websocket closed during worker job: {frame:?}"), - Some(Err(error)) => return Err(error).context("read websocket message during worker job"), - None => bail!("websocket ended during worker job"), + Ok(crf_search::Update::SampleResult { + crf, + sample, + result, + }) => { + debug!( + job_id = %job.assignment.job_id, + crf, + sample, + vmaf = ?result.vmaf_score, + "recorded sample result" + ); + Ok((None, false)) + } + Ok(crf_search::Update::RunResult(sample)) => { + let payload = job.crf_result_payload(&sample, false); + state.crf_results.push(payload.clone()); + let mut disconnected = false; + if let Some(worker) = worker { + disconnected = !send_worker_event( + worker, + ClientEvent::CrfSearchResult(payload), + &job.assignment.job_id, + "crf_run_result", + ) + .await; + } + Ok((None, disconnected)) + } + Err(error) => Err(error.into()), } +} - Ok(()) +async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobReportState) { + if let Some(heartbeat) = &state.heartbeat { + send_worker_event( + worker, + ClientEvent::Heartbeat(heartbeat.clone()), + "state", + "heartbeat", + ) + .await; + } + if let Some(progress) = &state.transfer_progress { + send_worker_event( + worker, + ClientEvent::TransferProgress(progress.clone()), + "state", + "transfer_progress", + ) + .await; + } + if let Some(progress) = &state.crf_progress { + send_worker_event( + worker, + ClientEvent::CrfSearchProgress(progress.clone()), + "state", + "crf_progress", + ) + .await; + } + for result in &state.crf_results { + send_worker_event( + worker, + ClientEvent::CrfSearchResult(result.clone()), + "state", + "crf_result", + ) + .await; + } +} + +async fn send_worker_event( + worker: &mut ConnectedWorker, + event: ClientEvent, + job_id: &str, + event_name: &'static str, +) -> bool { + if let Err(error) = worker.send_event(event).await { + debug!( + job_id = %job_id, + event = event_name, + error = %error, + "worker event send failed; keeping job and waiting to reconnect" + ); + false + } else { + true + } } -fn heartbeat_payload(path: &Path) -> HeartbeatPayload { - let mut system = System::new(); +fn heartbeat_payload(path: &Path, active_video_id: Option) -> HeartbeatPayload { + let system = HEARTBEAT_SYSTEM.get_or_init(|| { + let mut system = System::new_all(); + system.refresh_cpu(); + Mutex::new(system) + }); + let mut system = system.lock().expect("heartbeat system lock"); system.refresh_cpu(); system.refresh_memory(); let pid = Pid::from_u32(std::process::id()); + system.refresh_process(pid); let memory_rss_bytes = system.process(pid).map(|process| process.memory()); let disks = Disks::new_with_refreshed_list(); @@ -508,18 +727,14 @@ fn heartbeat_payload(path: &Path) -> HeartbeatPayload { memory_total_bytes: Some(system.total_memory()), disk_free_bytes: disk.map(|disk| disk.available_space()), disk_total_bytes: disk.map(|disk| disk.total_space()), + active_video_id, } } fn worker_job_input_dir(job_id: &str) -> PathBuf { std::env::current_dir() .expect("current working directory") - .join(format!( - "ab-av1-worker-{}-{}-{}", - std::process::id(), - job_id, - fastrand::u64(..) - )) + .join(format!("ab-av1-worker-{}", job_id)) } #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] @@ -669,12 +884,14 @@ impl ConnectedWorker { async fn request_work(&mut self) -> Result { let request_ref = self.next_ref; self.next_ref += 1; + let frame = ClientFrame::new(request_ref, ClientEvent::PullWork); + debug!( + request_ref = request_ref, + frame = %serde_json::to_string(&frame).context("serialize pull_work frame")?, + "sending pull_work" + ); - send_json( - &mut self.socket, - ClientFrame::new(request_ref, ClientEvent::PullWork), - ) - .await?; + send_json(&mut self.socket, frame).await?; expect_reply(&mut self.socket, &request_ref.to_string(), "pull_work").await } @@ -684,6 +901,24 @@ impl ConnectedWorker { send_json(&mut self.socket, ClientFrame::new(request_ref, event)).await } + async fn send_transfer_progress(&mut self, payload: TransferProgressPayload) -> Result<()> { + let throughput = format_bytes_per_second(payload.bytes_per_second); + debug!( + job_id = %payload.job_id, + transfer_id = %payload.transfer_id, + video_id = payload.video_id, + received_bytes = payload.received_bytes, + expected_bytes = ?payload.expected_bytes, + percent = payload.percent, + bytes_per_second = %throughput, + chunk_index = payload.chunk_index, + total_chunks = payload.total_chunks, + "sending transfer progress" + ); + self.send_event(ClientEvent::TransferProgress(payload)) + .await + } + async fn wait_for_pending_job( &mut self, pending_job: &mut PendingJob, @@ -714,6 +949,7 @@ impl ConnectedWorker { Some(WorkerPush::Started(started)) if started.transfer_id == pending_job.job().assignment.job_id => { + pending_job.ensure_receiver(started.chunk_size_bytes)?; debug!( job_id = %started.transfer_id, source_name = %started.source_name, @@ -721,6 +957,11 @@ impl ConnectedWorker { size_bytes = started.size_bytes, total_bytes = started.total_bytes, total_chunks = started.total_chunks, + received_bytes = pending_job + .receiver + .as_ref() + .map(ChunkReceiver::received_bytes) + .unwrap_or_default(), "transfer started" ); Ok(PendingJobOutcome::Waiting) @@ -750,9 +991,9 @@ impl ConnectedWorker { let chunk_index = chunk.chunk_index; let total_chunks = chunk.total_chunks; pending_job.apply_chunk(chunk)?; - self.send_event(ClientEvent::TransferProgress( + self.send_transfer_progress( pending_job.transfer_progress_payload(chunk_index, total_chunks), - )) + ) .await?; if pending_job.receiver.as_ref().is_some_and(|receiver| { receiver.received_bytes() @@ -798,9 +1039,9 @@ impl ConnectedWorker { let chunk_index = chunk.chunk_index; let total_chunks = chunk.total_chunks; pending_job.apply_raw_chunk(chunk)?; - self.send_event(ClientEvent::TransferProgress( + self.send_transfer_progress( pending_job.transfer_progress_payload(chunk_index, total_chunks), - )) + ) .await?; if pending_job.receiver.as_ref().is_some_and(|receiver| { receiver.received_bytes() == pending_job.job.assignment.size_bytes @@ -826,6 +1067,12 @@ impl ConnectedWorker { } } _ = tokio::time::sleep(idle_delay) => { + self + .send_event(ClientEvent::Heartbeat(heartbeat_payload( + &pending_job.job.input_dir, + Some(pending_job.job.assignment.video_id), + ))) + .await?; debug!( job_id = %pending_job.job().assignment.job_id, received_bytes = pending_job @@ -833,7 +1080,7 @@ impl ConnectedWorker { .as_ref() .map(|receiver| receiver.received_bytes()) .unwrap_or_default(), - "still waiting on websocket transfer" + "pending job still waiting" ); Ok(PendingJobOutcome::Waiting) } @@ -841,7 +1088,16 @@ impl ConnectedWorker { } } -async fn run_worker_job_and_publish(worker: &mut ConnectedWorker, job: &WorkerJob) -> Result<()> { +fn format_bytes_per_second(bytes_per_second: f64) -> String { + const MIB: f64 = 1024.0 * 1024.0; + format!("{:.1} MiB/s", bytes_per_second / MIB) +} + +async fn run_worker_job_and_publish( + config: &WorkerConfig, + worker: &mut Option, + job: &WorkerJob, +) -> Result<()> { debug!( job_id = %job.assignment.job_id, input = %job.input_path().display(), @@ -849,7 +1105,7 @@ async fn run_worker_job_and_publish(worker: &mut ConnectedWorker, job: &WorkerJo ); let probe = Arc::new(crate::ffprobe::probe(job.input_path())); debug!(job_id = %job.assignment.job_id, "probe complete, running crf search"); - let best = run_worker_job_with_reporting(job.clone(), probe, worker).await?; + let best = run_worker_job_with_reporting(config, job.clone(), probe, worker).await?; debug!(job_id = %job.assignment.job_id, "publishing worker result"); println!( @@ -864,14 +1120,105 @@ fn build_worker_job( local_path: Option<&Path>, ) -> Result { let input_dir = worker_job_input_dir(&assignment.job_id); - std::fs::create_dir_all(&input_dir).context("create worker job dir")?; - temporary::add(&input_dir, temporary::TempKind::Keepable); + fs::create_dir_all(&input_dir).context("create worker job dir")?; let input_path = local_path .map(Path::to_path_buf) - .unwrap_or_else(|| input_dir.join(&assignment.source_name)); + .map(Ok) + .unwrap_or_else(|| worker_job_input_path(&input_dir, &assignment))?; Ok(WorkerJob::new(assignment, input_dir, input_path)) } +fn worker_job_input_path( + input_dir: &Path, + assignment: &crate::command::worker_protocol::JobAssignedPayload, +) -> Result { + let source_name = worker_source_file_name(assignment)?; + let expected = input_dir.join(&source_name); + if expected.exists() { + return Ok(expected); + } + + let mut candidates = Vec::new(); + for entry in fs::read_dir(input_dir).context("read worker job dir")? { + let entry = entry.context("read worker job dir entry")?; + let path = entry.path(); + let metadata = entry + .metadata() + .context("read worker job dir entry metadata")?; + debug!( + job_id = %assignment.job_id, + path = %path.display(), + is_file = metadata.is_file(), + len = metadata.len(), + expected_len = assignment.size_bytes, + "found worker input dir entry" + ); + if !metadata.is_file() || metadata.len() != assignment.size_bytes { + continue; + } + if path + .file_name() + .is_some_and(|name| name == ".ab-av1-worker.part") + { + continue; + } + candidates.push(path); + } + + match candidates.as_slice() { + [path] => { + debug!( + job_id = %assignment.job_id, + input = %path.display(), + expected = %expected.display(), + source_name = %assignment.source_name, + "using existing completed worker input" + ); + Ok(path.clone()) + } + [] => { + debug!( + job_id = %assignment.job_id, + expected = %expected.display(), + source_name = %assignment.source_name, + "no existing completed worker input found" + ); + Ok(expected) + } + _ => { + debug!( + job_id = %assignment.job_id, + expected = %expected.display(), + source_name = %assignment.source_name, + candidates = candidates.len(), + "multiple existing worker inputs matched expected size" + ); + Ok(expected) + } + } +} + +fn worker_source_file_name( + assignment: &crate::command::worker_protocol::JobAssignedPayload, +) -> Result { + Path::new(&assignment.source_name) + .file_name() + .map(PathBuf::from) + .or_else(|| { + assignment + .crf_search_args + .windows(2) + .find(|args| args[0] == "--input") + .and_then(|args| Path::new(&args[1]).file_name().map(PathBuf::from)) + }) + .with_context(|| { + format!( + "job {} has no source filename in source_name or crf_search_args --input", + assignment.job_id + ) + }) +} + pub async fn worker(config: WorkerConfig) -> Result<()> { if config.once { let session = run_worker_session(&config).await?; @@ -917,7 +1264,7 @@ async fn run_connected_worker( local_path = ?config.local_path, "connecting worker" ); - let mut worker = ConnectedWorker::connect(config).await?; + let mut worker = Some(ConnectedWorker::connect(config).await?); let mut pending_job: Option = None; loop { @@ -929,11 +1276,25 @@ async fn run_connected_worker( input = %job.input_path().display(), "waiting for pending job input" ); - worker.wait_for_pending_job(job, runtime.idle_delay).await? + worker + .as_mut() + .expect("connected worker") + .wait_for_pending_job(job, runtime.idle_delay) + .await? }; match next { - PendingJobOutcome::Waiting => continue, + PendingJobOutcome::Waiting => { + debug!( + job_id = %pending_job.as_ref().expect("pending job").job.assignment.job_id, + received_bytes = pending_job + .as_ref() + .and_then(|job| job.receiver.as_ref().map(ChunkReceiver::received_bytes)) + .unwrap_or_default(), + "pending job still waiting" + ); + continue; + } PendingJobOutcome::Canceled => { if let Some(job) = pending_job.as_ref() { debug!(job_id = %job.job.assignment.job_id, "pending job canceled"); @@ -943,50 +1304,100 @@ async fn run_connected_worker( } PendingJobOutcome::Ready => { let job = pending_job.take().expect("pending job"); - debug!(job_id = %job.job.assignment.job_id, "pending job input arrived"); - run_worker_job_and_publish(&mut worker, &job.job).await?; + debug!( + job_id = %job.job.assignment.job_id, + input = %job.input_path().display(), + "pending job input arrived" + ); + run_worker_job_and_publish(config, &mut worker, &job.job).await?; continue; } } } debug!("requesting work"); - let work_status = worker.request_work().await?; + let worker_ref = worker.as_mut().expect("connected worker"); + let work_status = worker_ref.request_work().await?; *completed_pulls += 1; let status = work_status_label(&work_status); println!( "connected worker {} via {} and received {}", - worker.assigned_worker_id, worker.negotiated_protocol_version, status + worker_ref.assigned_worker_id, worker_ref.negotiated_protocol_version, status ); if let ServerReply::JobAssigned(assignment) = work_status { let job = build_worker_job(assignment, config.local_path.as_deref())?; debug!( job_id = %job.assignment.job_id, + status = %job.assignment.status.as_str(), input = %job.input_path().display(), + already_present = job.input_path().exists(), + pending_transfer = job.assignment.status == WorkStatus::JobAssigned + && !job.input_path().exists() + && config.local_path.is_none(), "job assigned" ); + if let Some(current_worker) = worker.as_mut() { + current_worker + .send_event(ClientEvent::Heartbeat(heartbeat_payload( + &job.input_dir, + Some(job.assignment.video_id), + ))) + .await?; + } if job.input_path().exists() { - debug!(job_id = %job.assignment.job_id, "input already present, starting job"); - run_worker_job_and_publish(&mut worker, &job).await?; + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "input already present, starting job" + ); + run_worker_job_and_publish(config, &mut worker, &job).await?; } else if config.local_path.is_some() { bail!( "local input path does not exist: {}", job.input_path().display() ); + } else if job.assignment.status == WorkStatus::JobInProgress { + let reason = format!( + "worker input is missing at {}; worker cannot resume job_in_progress without local file", + job.input_path().display() + ); + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + reason = %reason, + "reporting retriable transfer failure" + ); + if let Some(current_worker) = worker.as_mut() { + current_worker + .send_event(ClientEvent::TransferFailure(TransferFailurePayload { + job_id: job.assignment.job_id.clone(), + stage: TransferStage::ReceiveChunk, + retriable: true, + reason, + })) + .await?; + } + pending_job = Some(PendingJob::waiting(job)); + debug!( + job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, + pending_job = true, + "stored pending job after requesting transfer resend" + ); } else { - let receiver = ChunkReceiver::new( - job.input_path(), - &job.input_dir, - Some(job.assignment.size_bytes), - ) - .context("prepare worker input transfer")?; debug!( job_id = %job.assignment.job_id, input = %job.input_path().display(), + temp_dir = %job.input_dir.display(), + receiver_ready = false, "waiting for worker input over websocket" ); - pending_job = Some(PendingJob::new(job, receiver)); + pending_job = Some(PendingJob::waiting(job)); + debug!( + job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, + pending_job = true, + "stored pending job" + ); } continue; } @@ -1033,6 +1444,9 @@ fn decode_worker_push(text: &str) -> Result> { if frame.2 != CRF_SEARCH_TOPIC { return Ok(None); } + if frame.3 == "phx_reply" { + return Ok(None); + } let payload = frame.4.clone(); if matches!(frame.3.as_str(), "chunk_transfer" | "transfer_chunk") { @@ -1271,11 +1685,26 @@ where { continue; } + debug!( + expected_event, + raw = %text, + "received phoenix reply" + ); - let ReplyBody { status, response } = serde_json::from_value::>(body) - .context("decode phoenix reply body")?; + let ReplyBody { status, response }: ReplyBody = + serde_json::from_value::>(body) + .context("decode phoenix reply body")?; + debug!( + expected_event, + status = %status, + response = %response, + "decoded phoenix reply" + ); return match status.as_str() { - "ok" => serde_json::from_value(response).context("decode phoenix ok reply"), + "ok" => serde_json::from_value(response.clone()).map_err(|error| { + let raw_response = response.to_string(); + anyhow!("decode phoenix ok reply: {error}; raw_response={raw_response}") + }), "error" => { let error: ErrorReplyPayload = serde_json::from_value(response) .context("decode phoenix error reply")?; @@ -1614,6 +2043,13 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], })); assert_eq!(status, "job_assigned (job_id=job-123)"); @@ -1633,14 +2069,19 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], }, job_dir.clone(), input_path, ); - let config = job - .crf_search_config("libsvtav1".parse().expect("encoder")) - .expect("job config"); + let config = job.crf_search_config().expect("job config"); assert_eq!(config.args.input, job_dir.join("movie.mkv")); assert_eq!(config.sample.temp_dir.as_deref(), Some(job_dir.as_path())); @@ -1658,6 +2099,13 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], }; let local_path = std::env::temp_dir() .join(format!("ab-av1-worker-local-{}", std::process::id())) @@ -1668,6 +2116,37 @@ mod tests { assert_eq!(job.input_path(), local_path.as_path()); } + #[test] + fn build_worker_job_uses_source_basename_for_worker_file_lookup() -> Result<()> { + let job_dir = worker_job_input_dir("job-path-source"); + fs::create_dir_all(&job_dir)?; + let input_path = job_dir.join("movie.mkv"); + fs::write(&input_path, [0_u8; 4])?; + + let job = build_worker_job( + JobAssignedPayload { + status: WorkStatus::JobInProgress, + job_id: "job-path-source".into(), + video_id: 123, + source_name: "/server/library/movie.mkv".into(), + size_bytes: 4, + chunk_size_bytes: 0, + target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/library/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], + }, + None, + )?; + + assert_eq!(job.input_path(), input_path.as_path()); + Ok(()) + } + #[tokio::test(flavor = "current_thread")] async fn worker_job_runs_crf_search_from_fake_probe() -> Result<()> { crf_test_hooks::set(|_crf| sample_encode::Output { @@ -1691,6 +2170,13 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], }, job_dir, input_path, @@ -1727,6 +2213,18 @@ mod tests { from_cache: false, } ); + let result = job.crf_result_payload(&best, true); + assert_eq!(result.job_id, "job-123"); + assert_eq!(result.video_id, 123); + assert_eq!(result.source_name, "movie.mkv"); + assert_eq!(result.crf, best.crf); + assert_eq!(result.vmaf_score, Some(97.0)); + assert_eq!(result.xpsnr_score, None); + assert_eq!(result.predicted_encode_size, 100); + assert_eq!(result.encode_percent, 50.0); + assert_eq!(result.predicted_encode_time_secs, 1.0); + assert!(!result.from_cache); + assert!(result.chosen); Ok(()) } @@ -1976,6 +2474,13 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], })), )) .expect("job assigned reply json"), diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index f8587b57..6fff98b0 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -25,6 +25,7 @@ pub(crate) enum ClientEvent { PullWork, Heartbeat(HeartbeatPayload), TransferProgress(TransferProgressPayload), + TransferFailure(TransferFailurePayload), CrfSearchProgress(CrfSearchProgressPayload), CrfSearchResult(CrfSearchResultPayload), } @@ -40,6 +41,9 @@ impl ClientEvent { "transfer_progress", ClientPayload::TransferProgress(payload), ), + Self::TransferFailure(payload) => { + ("transfer_failed", ClientPayload::TransferFailure(payload)) + } Self::CrfSearchProgress(payload) => { ("crf_search_progress", ClientPayload::Progress(payload)) } @@ -55,6 +59,7 @@ enum ClientPayload { Announce(AnnouncePayload), Heartbeat(HeartbeatPayload), TransferProgress(TransferProgressPayload), + TransferFailure(TransferFailurePayload), Progress(CrfSearchProgressPayload), Result(CrfSearchResultPayload), } @@ -87,6 +92,8 @@ pub(crate) struct HeartbeatPayload { pub(crate) disk_free_bytes: Option, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) disk_total_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) active_video_id: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -100,7 +107,16 @@ pub(crate) struct CrfSearchProgressPayload { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct CrfSearchResultPayload { + pub(crate) job_id: String, + pub(crate) video_id: u64, + pub(crate) source_name: String, pub(crate) crf: f32, + pub(crate) vmaf_score: Option, + pub(crate) xpsnr_score: Option, + pub(crate) predicted_encode_size: u64, + pub(crate) encode_percent: f64, + pub(crate) predicted_encode_time_secs: f64, + pub(crate) from_cache: bool, pub(crate) score: f32, pub(crate) percent: f64, pub(crate) size: u64, @@ -179,11 +195,12 @@ pub(crate) enum ServerReply { NoWork(NoWorkPayload), } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub(crate) enum WorkStatus { NoWork, JobAssigned, + JobInProgress, } impl WorkStatus { @@ -191,15 +208,55 @@ impl WorkStatus { match self { Self::NoWork => "no_work", Self::JobAssigned => "job_assigned", + Self::JobInProgress => "job_in_progress", } } } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +impl<'de> Deserialize<'de> for WorkStatus { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let status = String::deserialize(deserializer)?; + match status.as_str() { + "no_work" => Ok(Self::NoWork), + "job_assigned" => Ok(Self::JobAssigned), + "job_in_progress" => Ok(Self::JobInProgress), + other => Err(serde::de::Error::unknown_variant( + other, + &["no_work", "job_assigned", "job_in_progress"], + )), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub(crate) struct NoWorkPayload { pub(crate) status: WorkStatus, } +impl<'de> Deserialize<'de> for NoWorkPayload { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + #[serde(untagged)] + enum NoWorkRepr { + Null(()), + Explicit { status: WorkStatus }, + } + + match NoWorkRepr::deserialize(deserializer)? { + NoWorkRepr::Null(()) => Ok(NoWorkPayload { + status: WorkStatus::NoWork, + }), + NoWorkRepr::Explicit { status } => Ok(NoWorkPayload { status }), + } + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct JobAssignedPayload { pub(crate) status: WorkStatus, @@ -207,8 +264,11 @@ pub(crate) struct JobAssignedPayload { pub(crate) video_id: u64, pub(crate) source_name: String, pub(crate) size_bytes: u64, + #[serde(default)] pub(crate) chunk_size_bytes: u64, pub(crate) target_vmaf: f32, + #[serde(default)] + pub(crate) crf_search_args: Vec, } #[cfg_attr(not(test), allow(dead_code))] @@ -371,6 +431,7 @@ mod tests { memory_total_bytes: Some(8192), disk_free_bytes: Some(4096), disk_total_bytes: Some(16_384), + active_video_id: Some(123), }), ); @@ -387,6 +448,7 @@ mod tests { "memory_total_bytes": 8192, "disk_free_bytes": 4096, "disk_total_bytes": 16_384, + "active_video_id": 123, } ]) ); @@ -471,7 +533,16 @@ mod tests { let frame = ClientFrame::new( 6, ClientEvent::CrfSearchResult(CrfSearchResultPayload { + job_id: "job-123".into(), + video_id: 123, + source_name: "movie.mkv".into(), crf: 31.5, + vmaf_score: Some(96.2), + xpsnr_score: None, + predicted_encode_size: 123_456, + encode_percent: 42.5, + predicted_encode_time_secs: 87.5, + from_cache: true, score: 96.2, percent: 42.5, size: 123_456, @@ -490,7 +561,16 @@ mod tests { "workers:crf_search", "crf_search_result", { + "job_id": "job-123", + "video_id": 123, + "source_name": "movie.mkv", "crf": 31.5, + "vmaf_score": 96.19999694824219, + "xpsnr_score": null, + "predicted_encode_size": 123456, + "encode_percent": 42.5, + "predicted_encode_time_secs": 87.5, + "from_cache": true, "score": 96.19999694824219, "percent": 42.5, "size": 123456, @@ -528,6 +608,31 @@ mod tests { ); } + #[test] + fn server_reply_parses_null_no_work_payload() { + let reply: ServerFrame = serde_json::from_value(json!([ + null, + "3", + "workers:crf_search", + "phx_reply", + { + "status": "ok", + "response": null + } + ])) + .expect("parse null no_work reply"); + + assert_eq!( + reply, + ServerFrame::reply( + 3, + ReplyBody::ok(ServerReply::NoWork(NoWorkPayload { + status: WorkStatus::NoWork, + })), + ) + ); + } + #[test] fn server_reply_parses_future_job_assignment_payload() { let reply: ServerFrame = serde_json::from_value(json!([ @@ -544,7 +649,14 @@ mod tests { "source_name": "movie.mkv", "size_bytes": 1024, "chunk_size_bytes": 256, - "target_vmaf": 96.5 + "target_vmaf": 96.5, + "crf_search_args": [ + "crf-search", + "--input", + "/server/movie.mkv", + "--min-vmaf", + "96.5" + ] } } ])) @@ -562,6 +674,65 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "96.5".into(), + ], + })), + ) + ); + } + + #[test] + fn server_reply_parses_in_progress_assignment_without_chunk_size() { + let reply: ServerFrame = serde_json::from_value(json!([ + null, + "3", + "workers:crf_search", + "phx_reply", + { + "status": "ok", + "response": { + "status": "job_in_progress", + "source_name": "movie.mkv", + "video_id": 123, + "job_id": "job-123", + "size_bytes": 1024, + "target_vmaf": 95, + "crf_search_args": [ + "crf-search", + "--input", + "/server/movie.mkv", + "--min-vmaf", + "95" + ] + } + } + ])) + .expect("parse job_in_progress reply without chunk size"); + + assert_eq!( + reply, + ServerFrame::reply( + 3, + ReplyBody::ok(ServerReply::JobAssigned(JobAssignedPayload { + status: WorkStatus::JobInProgress, + job_id: "job-123".into(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 1024, + chunk_size_bytes: 0, + target_vmaf: 95.0, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "95".into(), + ], })), ) ); diff --git a/src/command/worker_transfer.rs b/src/command/worker_transfer.rs index 913cacea..9a1485ae 100644 --- a/src/command/worker_transfer.rs +++ b/src/command/worker_transfer.rs @@ -1,9 +1,8 @@ -use crate::temporary; use anyhow::{Context, Result}; use blake3::Hash; use std::{ fs::{self, File, OpenOptions}, - io::Write, + io::{Read, Write}, path::{Path, PathBuf}, }; use tracing::{debug, trace}; @@ -65,37 +64,68 @@ impl ChunkReceiver { final_path: impl Into, temp_dir: impl AsRef, max_size: Option, + chunk_size_bytes: u64, ) -> Result { let final_path = final_path.into(); let temp_dir = temp_dir.as_ref(); fs::create_dir_all(temp_dir).context("create chunk temp dir")?; - let temp_path = temp_dir.join(format!( - ".ab-av1-worker-{}-{}.part", - std::process::id(), - fastrand::u64(..) - )); - let file = OpenOptions::new() - .create_new(true) - .write(true) - .open(&temp_path) - .context("create chunk temp file")?; - temporary::add(&temp_path, temporary::TempKind::NotKeepable); - debug!( - final_path = %final_path.display(), - temp_path = %temp_path.display(), - max_size = ?max_size, - "created chunk receiver" - ); + let temp_path = temp_dir.join(".ab-av1-worker.part"); + let (file, next_index, next_offset, hasher) = if temp_path.exists() { + let file = OpenOptions::new() + .read(true) + .append(true) + .open(&temp_path) + .context("open chunk temp file for resume")?; + let metadata = file.metadata().context("inspect chunk temp file")?; + let next_offset = metadata.len(); + let next_index = if chunk_size_bytes == 0 { + 0 + } else { + next_offset / chunk_size_bytes + }; + let mut hasher = blake3::Hasher::new(); + let mut reader = File::open(&temp_path).context("read chunk temp file for resume")?; + let mut buf = [0u8; 8192]; + loop { + let read = reader.read(&mut buf).context("hash chunk temp file")?; + if read == 0 { + break; + } + hasher.update(&buf[..read]); + } + debug!( + final_path = %final_path.display(), + temp_path = %temp_path.display(), + received_bytes = next_offset, + chunk_size_bytes, + max_size = ?max_size, + "resumed chunk receiver" + ); + (Some(file), next_index, next_offset, hasher) + } else { + let file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temp_path) + .context("create chunk temp file")?; + debug!( + final_path = %final_path.display(), + temp_path = %temp_path.display(), + max_size = ?max_size, + "created chunk receiver" + ); + (Some(file), 0, 0, blake3::Hasher::new()) + }; Ok(Self { final_path, temp_path, - file: Some(file), + file, max_size, - next_index: 0, - next_offset: 0, + next_index, + next_offset, finished: false, - hasher: blake3::Hasher::new(), + hasher, }) } @@ -183,7 +213,6 @@ impl ChunkReceiver { ); fs::rename(&self.temp_path, &self.final_path)?; let final_path = self.final_path.clone(); - temporary::unadd(&self.temp_path); self.finished = true; debug!(final_path = %final_path.display(), "chunk receiver finished"); Ok(final_path) @@ -195,7 +224,6 @@ impl Drop for ChunkReceiver { if !self.finished { let _ = self.file.take(); let _ = fs::remove_file(&self.temp_path); - let _ = temporary::unadd(&self.temp_path); } } } @@ -229,7 +257,7 @@ mod tests { #[test] fn valid_transfer_writes_final_file() { let (temp_dir, final_path) = temp_paths("valid"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 6).expect("receiver"); receiver.push(chunk(0, 0, b"hello ")).expect("chunk 0"); receiver.push(chunk(1, 6, b"world")).expect("chunk 1"); @@ -243,6 +271,27 @@ mod tests { let _ = fs::remove_dir_all(temp_dir); } + #[serial] + #[test] + fn resumes_existing_partial_transfer() { + let (temp_dir, final_path) = temp_paths("resume"); + fs::create_dir_all(&temp_dir).expect("temp dir"); + let temp_path = temp_dir.join(".ab-av1-worker.part"); + fs::write(&temp_path, b"hello ").expect("seed temp file"); + + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 6).expect("receiver"); + assert_eq!(receiver.received_bytes(), 6); + + receiver.push(chunk(1, 6, b"world")).expect("chunk 1"); + let written = receiver + .finish(Some(11), Some(blake3::hash(b"hello world"))) + .expect("finish"); + + assert_eq!(written, final_path); + assert_eq!(fs::read(&final_path).expect("read final"), b"hello world"); + let _ = fs::remove_dir_all(temp_dir); + } + #[serial] #[test] fn finish_creates_parent_directory() { @@ -252,7 +301,7 @@ mod tests { fastrand::u64(..) )); let final_path = temp_dir.join("nested").join("movie.mkv"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); receiver @@ -275,7 +324,7 @@ mod tests { fs::create_dir_all(final_path.parent().expect("parent")).expect("create parent"); fs::write(&final_path, b"existing").expect("seed final"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); assert!(matches!( @@ -290,7 +339,7 @@ mod tests { #[test] fn corrupt_chunk_is_rejected() { let (temp_dir, final_path) = temp_paths("corrupt"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); let mut bad = chunk(0, 0, b"hello"); bad.checksum = crc32fast::hash(b"hell0") as u64; @@ -306,7 +355,7 @@ mod tests { #[test] fn missing_chunk_is_rejected_by_offset() { let (temp_dir, final_path) = temp_paths("missing"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); assert!(matches!( @@ -324,7 +373,7 @@ mod tests { #[test] fn duplicate_chunk_is_rejected() { let (temp_dir, final_path) = temp_paths("duplicate"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); assert!(matches!( @@ -338,7 +387,7 @@ mod tests { #[test] fn out_of_order_chunk_is_rejected() { let (temp_dir, final_path) = temp_paths("out-of-order"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); assert!(matches!( receiver.push(chunk(1, 0, b"hello")), @@ -354,7 +403,7 @@ mod tests { #[test] fn final_digest_mismatch_is_rejected() { let (temp_dir, final_path) = temp_paths("digest"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None, 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); assert!(matches!( @@ -368,7 +417,8 @@ mod tests { #[test] fn max_size_limit_is_enforced() { let (temp_dir, final_path) = temp_paths("max-size"); - let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, Some(10)).expect("receiver"); + let mut receiver = + ChunkReceiver::new(&final_path, &temp_dir, Some(10), 5).expect("receiver"); receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); assert_eq!(receiver.received_bytes(), 5); diff --git a/src/process/managed.rs b/src/process/managed.rs index c1c11144..3e809dfd 100644 --- a/src/process/managed.rs +++ b/src/process/managed.rs @@ -381,8 +381,11 @@ impl MustCompleteProcess { /// and only then yields `ProcessDone`. pub fn stderr_events(self) -> impl Stream> { async_stream::try_stream! { - let mut process = self.0; - let mut stderr = process.handle.stderr().try_subscribe()?; + let mut process = TerminateOnDropProcess(Some(self.0)); + let Some(inner) = process.0.as_mut() else { + return; + }; + let mut stderr = inner.handle.stderr().try_subscribe()?; while let Some(event) = stderr.next_event().await { match managed_event_from_stream_event(event)? { Some(ManagedEvent::RawStderr(chunk)) => yield ManagedEvent::RawStderr(chunk), @@ -392,7 +395,13 @@ impl MustCompleteProcess { } } - yield ManagedEvent::ProcessDone(wait_for_process_done(&mut process).await?); + drop(stderr); + let Some(inner) = process.0.as_mut() else { + return; + }; + let done = wait_for_process_done(inner).await?; + drop(process.0.take()); + yield ManagedEvent::ProcessDone(done); } } } From 39f1103f308a0f679313346fa5879681b94aac84 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 8 Jul 2026 22:39:40 -0600 Subject: [PATCH 6/7] Make worker input-resend state explicit --- src/command/worker.rs | 346 ++++++++++++++++++++++++++------- src/command/worker_protocol.rs | 35 +++- src/command/worker_transfer.rs | 8 +- 3 files changed, 313 insertions(+), 76 deletions(-) diff --git a/src/command/worker.rs b/src/command/worker.rs index 29d8cef6..82e454d3 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -1,7 +1,7 @@ use crate::command::worker_protocol::{ AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, - HeartbeatPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, + HeartbeatPayload, JobResultPayload, PullWorkPayload, ReplyBody, ServerPushFrame, ServerReply, TransferFailurePayload, TransferProgressPayload, TransferStage, TransferStartedPayload, WorkStatus, }; @@ -794,6 +794,43 @@ impl ReconnectBackoff { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum WorkerJobPhase { + ReceivingInput, + InputMissing, + InputReady, + CrfSearching, +} + +fn worker_job_phase(job: &WorkerJob, local_path: Option<&Path>) -> Result { + if job.input_path().exists() { + return match job.assignment.status { + WorkStatus::JobAssigned => Ok(WorkerJobPhase::InputReady), + WorkStatus::JobInProgress => Ok(WorkerJobPhase::CrfSearching), + WorkStatus::NoWork => bail!( + "assigned job {} has invalid no_work status", + job.assignment.job_id + ), + }; + } + + if local_path.is_some() { + bail!( + "local input path does not exist: {}", + job.input_path().display() + ); + } + + match job.assignment.status { + WorkStatus::JobAssigned => Ok(WorkerJobPhase::ReceivingInput), + WorkStatus::JobInProgress => Ok(WorkerJobPhase::InputMissing), + WorkStatus::NoWork => bail!( + "assigned job {} has invalid no_work status", + job.assignment.job_id + ), + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum PendingJobOutcome { Waiting, @@ -882,9 +919,13 @@ impl ConnectedWorker { } async fn request_work(&mut self) -> Result { + self.request_work_with(PullWorkPayload::default()).await + } + + async fn request_work_with(&mut self, payload: PullWorkPayload) -> Result { let request_ref = self.next_ref; self.next_ref += 1; - let frame = ClientFrame::new(request_ref, ClientEvent::PullWork); + let frame = ClientFrame::new(request_ref, ClientEvent::PullWork(payload)); debug!( request_ref = request_ref, frame = %serde_json::to_string(&frame).context("serialize pull_work frame")?, @@ -944,7 +985,7 @@ impl ConnectedWorker { "worker job {} canceled: {}", cancel.job_id, cancel.reason ); - return Ok(PendingJobOutcome::Canceled); + Ok(PendingJobOutcome::Canceled) } Some(WorkerPush::Started(started)) if started.transfer_id == pending_job.job().assignment.job_id => @@ -1252,6 +1293,63 @@ async fn run_worker_until(config: &WorkerConfig, runtime: WorkerRuntime) -> Resu } } +async fn request_input_resend( + worker: &mut ConnectedWorker, + job: &WorkerJob, + local_path: Option<&Path>, +) -> Result { + let reason = format!( + "worker input is missing at {}; worker cannot resume job_in_progress without local file", + job.input_path().display() + ); + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + reason = %reason, + "reporting retriable transfer failure" + ); + worker + .send_event(ClientEvent::TransferFailure(TransferFailurePayload { + job_id: job.assignment.job_id.clone(), + stage: TransferStage::ReceiveChunk, + retriable: true, + reason, + })) + .await?; + + debug!( + job_id = %job.assignment.job_id, + "requesting input resend for active job" + ); + let resend = worker + .request_work_with(PullWorkPayload::input_missing()) + .await?; + + let ServerReply::JobAssigned(assignment) = resend else { + bail!( + "server returned no_work after input_missing for job {}", + job.assignment.job_id + ); + }; + + if assignment.job_id != job.assignment.job_id { + bail!( + "server reassigned job {} after input_missing for job {}", + assignment.job_id, + job.assignment.job_id + ); + } + if assignment.status != WorkStatus::JobAssigned { + bail!( + "server kept job {} in {} after input_missing; refusing to wait without transfer", + assignment.job_id, + assignment.status.as_str() + ); + } + + build_worker_job(assignment, local_path) +} + async fn run_connected_worker( config: &WorkerConfig, runtime: WorkerRuntime, @@ -1345,59 +1443,55 @@ async fn run_connected_worker( ))) .await?; } - if job.input_path().exists() { - debug!( - job_id = %job.assignment.job_id, - input = %job.input_path().display(), - "input already present, starting job" - ); - run_worker_job_and_publish(config, &mut worker, &job).await?; - } else if config.local_path.is_some() { - bail!( - "local input path does not exist: {}", - job.input_path().display() - ); - } else if job.assignment.status == WorkStatus::JobInProgress { - let reason = format!( - "worker input is missing at {}; worker cannot resume job_in_progress without local file", - job.input_path().display() - ); - debug!( - job_id = %job.assignment.job_id, - input = %job.input_path().display(), - reason = %reason, - "reporting retriable transfer failure" - ); - if let Some(current_worker) = worker.as_mut() { - current_worker - .send_event(ClientEvent::TransferFailure(TransferFailurePayload { - job_id: job.assignment.job_id.clone(), - stage: TransferStage::ReceiveChunk, - retriable: true, - reason, - })) - .await?; + let phase = worker_job_phase(&job, config.local_path.as_deref())?; + match phase { + WorkerJobPhase::InputReady | WorkerJobPhase::CrfSearching => { + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + phase = ?phase, + "input already present, starting job" + ); + run_worker_job_and_publish(config, &mut worker, &job).await?; + } + WorkerJobPhase::InputMissing => { + let resend_job = request_input_resend( + worker.as_mut().expect("connected worker"), + &job, + config.local_path.as_deref(), + ) + .await?; + debug!( + job_id = %resend_job.assignment.job_id, + input = %resend_job.input_path().display(), + temp_dir = %resend_job.input_dir.display(), + phase = ?WorkerJobPhase::ReceivingInput, + receiver_ready = false, + "waiting for worker input over websocket" + ); + pending_job = Some(PendingJob::waiting(resend_job)); + debug!( + job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, + pending_job = true, + "stored pending job after input resend request" + ); + } + WorkerJobPhase::ReceivingInput => { + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + temp_dir = %job.input_dir.display(), + phase = ?phase, + receiver_ready = false, + "waiting for worker input over websocket" + ); + pending_job = Some(PendingJob::waiting(job)); + debug!( + job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, + pending_job = true, + "stored pending job" + ); } - pending_job = Some(PendingJob::waiting(job)); - debug!( - job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, - pending_job = true, - "stored pending job after requesting transfer resend" - ); - } else { - debug!( - job_id = %job.assignment.job_id, - input = %job.input_path().display(), - temp_dir = %job.input_dir.display(), - receiver_ready = false, - "waiting for worker input over websocket" - ); - pending_job = Some(PendingJob::waiting(job)); - debug!( - job_id = %pending_job.as_ref().unwrap().job.assignment.job_id, - pending_job = true, - "stored pending job" - ); } continue; } @@ -1529,12 +1623,12 @@ fn decode_binary_transfer_chunk(bytes: &[u8]) -> Result { Ok(TransferChunk { transfer_id, - video_id: read_u64(&bytes, 8), - chunk_index: read_u64(&bytes, 16), - total_chunks: read_u64(&bytes, 24), - bytes_sent: read_u64(&bytes, 32), - total_bytes: read_u64(&bytes, 40), - crc32: read_u32(&bytes, 48) as u64, + video_id: read_u64(bytes, 8), + chunk_index: read_u64(bytes, 16), + total_chunks: read_u64(bytes, 24), + bytes_sent: read_u64(bytes, 32), + total_bytes: read_u64(bytes, 40), + crc32: read_u32(bytes, 48) as u64, bytes: bytes[data_start..].to_vec(), }) } @@ -1954,6 +2048,64 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "current_thread")] + async fn worker_requests_input_resend_when_resumed_job_lacks_local_file() -> Result<()> { + let job_id = "missing-input-resend"; + let _ = fs::remove_dir_all(worker_job_input_dir(job_id)); + let (listener, address) = FakeCoordinator::bind("127.0.0.1:0").await?; + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept connection"); + let socket = accept_async(stream).await.expect("accept websocket"); + let (mut writer, mut reader) = socket.split(); + + expect_join(&mut reader).await; + send_join_reply(&mut writer).await; + expect_announce(&mut reader, 1).await; + send_announce_reply(&mut writer).await; + + expect_pull_work(&mut reader, 3).await; + send_job_reply_with_job_id(&mut writer, 3, WorkStatus::JobInProgress, job_id).await; + + let heartbeat = expect_client_event(&mut reader, 4, "heartbeat").await; + assert_eq!(heartbeat["active_video_id"], json!(123)); + + let failure = expect_client_event(&mut reader, 5, "transfer_failed").await; + assert_eq!(failure["job_id"], json!(job_id)); + assert_eq!(failure["stage"], json!("receive_chunk")); + assert_eq!(failure["retriable"], json!(true)); + + expect_pull_work_payload(&mut reader, 6, PullWorkPayload::input_missing()).await; + send_job_reply_with_job_id(&mut writer, 6, WorkStatus::JobAssigned, job_id).await; + }); + + let config = FakeCoordinator { + address, + server: tokio::spawn(async {}), + } + .worker_config(WorkerTestConfig::continuous()); + let mut completed_pulls = 0; + let error = run_connected_worker( + &config, + WorkerRuntime { + idle_delay: Duration::from_millis(50), + reconnect_base_delay: Duration::from_millis(1), + reconnect_max_delay: Duration::from_millis(1), + max_pulls: None, + }, + &mut completed_pulls, + ) + .await + .expect_err("server closes after resend assignment"); + + assert!( + error.to_string().contains("websocket"), + "unexpected error: {error}" + ); + server.await.expect("server task"); + Ok(()) + } + #[tokio::test(flavor = "current_thread")] async fn worker_reports_supported_versions_on_protocol_mismatch() -> Result<()> { let (listener, address) = FakeCoordinator::bind("127.0.0.1:0").await?; @@ -2391,6 +2543,14 @@ mod tests { } async fn expect_pull_work(reader: &mut R, request_ref: u64) + where + R: StreamExt> + + Unpin, + { + expect_pull_work_payload(reader, request_ref, PullWorkPayload::default()).await; + } + + async fn expect_pull_work_payload(reader: &mut R, request_ref: u64, payload: PullWorkPayload) where R: StreamExt> + Unpin, @@ -2401,11 +2561,36 @@ mod tests { .await .expect("pull_work frame") .expect("pull_work message"), - serde_json::to_value(ClientFrame::new(request_ref, ClientEvent::PullWork)) - .expect("pull_work frame json"), + serde_json::to_value(ClientFrame::new( + request_ref, + ClientEvent::PullWork(payload), + )) + .expect("pull_work frame json"), ); } + async fn expect_client_event(reader: &mut R, request_ref: u64, event: &str) -> Value + where + R: StreamExt> + + Unpin, + { + let Message::Text(text) = reader + .next() + .await + .expect("client event frame") + .expect("client event message") + else { + panic!("expected text client event"); + }; + let actual: Value = serde_json::from_str(&text).expect("decode client event"); + let frame = actual.as_array().expect("client event frame array"); + assert_eq!(frame[0], json!("1")); + assert_eq!(frame[1], json!(request_ref.to_string())); + assert_eq!(frame[2], json!(CRF_SEARCH_TOPIC)); + assert_eq!(frame[3], json!(event)); + frame[4].clone() + } + async fn send_join_reply(writer: &mut W) where W: SinkExt + Unpin, @@ -2462,17 +2647,40 @@ mod tests { where W: SinkExt + Unpin, { + send_job_reply(writer, 3, WorkStatus::JobAssigned).await; + } + + async fn send_job_reply(writer: &mut W, request_ref: u64, status: WorkStatus) + where + W: SinkExt + Unpin, + { + send_job_reply_with_job_id(writer, request_ref, status, "job-123").await; + } + + async fn send_job_reply_with_job_id( + writer: &mut W, + request_ref: u64, + status: WorkStatus, + job_id: &str, + ) where + W: SinkExt + Unpin, + { + let chunk_size_bytes = if status == WorkStatus::JobAssigned { + 256 + } else { + 0 + }; writer .send(Message::Text( serde_json::to_string(&ServerFrame::reply( - 3, + request_ref, ReplyBody::ok(ServerReply::JobAssigned(JobAssignedPayload { - status: WorkStatus::JobAssigned, - job_id: "job-123".into(), + status, + job_id: job_id.into(), video_id: 123, source_name: "movie.mkv".into(), size_bytes: 1024, - chunk_size_bytes: 256, + chunk_size_bytes, target_vmaf: 96.5, crf_search_args: vec![ "crf-search".into(), @@ -2483,10 +2691,10 @@ mod tests { ], })), )) - .expect("job assigned reply json"), + .expect("job reply json"), )) .await - .expect("send job assigned reply"); + .expect("send job reply"); } async fn send_announce_error_reply(writer: &mut W, request_ref: u64, response: Value) diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index 6fff98b0..eefe135f 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -22,7 +22,7 @@ impl ClientFrame { pub(crate) enum ClientEvent { Join, Announce(AnnouncePayload), - PullWork, + PullWork(PullWorkPayload), Heartbeat(HeartbeatPayload), TransferProgress(TransferProgressPayload), TransferFailure(TransferFailurePayload), @@ -35,7 +35,7 @@ impl ClientEvent { match self { Self::Join => ("phx_join", ClientPayload::Empty(EmptyPayload {})), Self::Announce(payload) => ("announce", ClientPayload::Announce(payload)), - Self::PullWork => ("pull_work", ClientPayload::Empty(EmptyPayload {})), + Self::PullWork(payload) => ("pull_work", ClientPayload::PullWork(payload)), Self::Heartbeat(payload) => ("heartbeat", ClientPayload::Heartbeat(payload)), Self::TransferProgress(payload) => ( "transfer_progress", @@ -57,6 +57,7 @@ impl ClientEvent { enum ClientPayload { Empty(EmptyPayload), Announce(AnnouncePayload), + PullWork(PullWorkPayload), Heartbeat(HeartbeatPayload), TransferProgress(TransferProgressPayload), TransferFailure(TransferFailurePayload), @@ -67,6 +68,24 @@ enum ClientPayload { #[derive(Debug, Clone, PartialEq, Eq, Serialize)] struct EmptyPayload {} +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct PullWorkPayload { + #[serde(default, skip_serializing_if = "is_false")] + pub(crate) input_missing: bool, +} + +impl PullWorkPayload { + pub(crate) fn input_missing() -> Self { + Self { + input_missing: true, + } + } +} + +fn is_false(value: &bool) -> bool { + !*value +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct AnnouncePayload { pub(crate) worker_id: String, @@ -883,6 +902,18 @@ mod tests { ); } + #[test] + fn pull_work_payload_omits_default_and_reports_missing_input() { + assert_eq!( + serde_json::to_value(PullWorkPayload::default()).expect("serialize pull_work"), + json!({}) + ); + assert_eq!( + serde_json::to_value(PullWorkPayload::input_missing()).expect("serialize pull_work"), + json!({"input_missing": true}) + ); + } + #[test] fn transfer_failure_payload_serializes_stage_and_retry_hint() { let payload = TransferFailurePayload { diff --git a/src/command/worker_transfer.rs b/src/command/worker_transfer.rs index 9a1485ae..09fff65a 100644 --- a/src/command/worker_transfer.rs +++ b/src/command/worker_transfer.rs @@ -78,11 +78,9 @@ impl ChunkReceiver { .context("open chunk temp file for resume")?; let metadata = file.metadata().context("inspect chunk temp file")?; let next_offset = metadata.len(); - let next_index = if chunk_size_bytes == 0 { - 0 - } else { - next_offset / chunk_size_bytes - }; + let next_index = next_offset + .checked_div(chunk_size_bytes) + .unwrap_or_default(); let mut hasher = blake3::Hasher::new(); let mut reader = File::open(&temp_path).context("read chunk temp file for resume")?; let mut buf = [0u8; 8192]; From 2a3b74f552e4e5618174f8949f566f4dbab1f94d Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Thu, 9 Jul 2026 09:04:25 -0600 Subject: [PATCH 7/7] Report worker CRF search failures --- src/command/worker.rs | 50 +++++++++++++++++++++++++++++++--- src/command/worker_protocol.rs | 15 ++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/src/command/worker.rs b/src/command/worker.rs index 82e454d3..2e83cbc0 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -1,9 +1,9 @@ use crate::command::worker_protocol::{ AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, - HeartbeatPayload, JobResultPayload, PullWorkPayload, ReplyBody, ServerPushFrame, ServerReply, - TransferFailurePayload, TransferProgressPayload, TransferStage, TransferStartedPayload, - WorkStatus, + FailureReportPayload, HeartbeatPayload, JobResultPayload, PullWorkPayload, ReplyBody, + ServerPushFrame, ServerReply, TransferFailurePayload, TransferProgressPayload, TransferStage, + TransferStartedPayload, WorkStatus, }; use crate::command::worker_transfer::{Chunk, ChunkReceiver}; use crate::command::{crf_search, sample_encode}; @@ -213,6 +213,23 @@ impl WorkerJob { chosen, } } + + fn failure_payload(&self, error: &anyhow::Error) -> FailureReportPayload { + FailureReportPayload { + video_id: self.assignment.video_id, + stage: "crf_search".into(), + category: "process_failure".into(), + message: error.to_string(), + code: "worker_crf_search_failed".into(), + context: json!({ + "job_id": self.assignment.job_id, + "source_name": self.assignment.source_name, + "error_chain": format!("{error:#}"), + }), + retriable: false, + stderr_excerpt: Some(format!("{error:#}")), + } + } } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -1146,7 +1163,13 @@ async fn run_worker_job_and_publish( ); let probe = Arc::new(crate::ffprobe::probe(job.input_path())); debug!(job_id = %job.assignment.job_id, "probe complete, running crf search"); - let best = run_worker_job_with_reporting(config, job.clone(), probe, worker).await?; + let best = match run_worker_job_with_reporting(config, job.clone(), probe, worker).await { + Ok(best) => best, + Err(error) => { + publish_worker_failure(worker, job, &error).await; + return Err(error); + } + }; debug!(job_id = %job.assignment.job_id, "publishing worker result"); println!( @@ -1156,6 +1179,25 @@ async fn run_worker_job_and_publish( Ok(()) } +async fn publish_worker_failure( + worker: &mut Option, + job: &WorkerJob, + error: &anyhow::Error, +) { + let Some(worker) = worker else { + return; + }; + + let payload = job.failure_payload(error); + let _ = send_worker_event( + worker, + ClientEvent::VideoFailed(payload), + &job.assignment.job_id, + "video_failed", + ) + .await; +} + fn build_worker_job( assignment: crate::command::worker_protocol::JobAssignedPayload, local_path: Option<&Path>, diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index eefe135f..d835a6e4 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -28,6 +28,7 @@ pub(crate) enum ClientEvent { TransferFailure(TransferFailurePayload), CrfSearchProgress(CrfSearchProgressPayload), CrfSearchResult(CrfSearchResultPayload), + VideoFailed(FailureReportPayload), } impl ClientEvent { @@ -48,6 +49,7 @@ impl ClientEvent { ("crf_search_progress", ClientPayload::Progress(payload)) } Self::CrfSearchResult(payload) => ("crf_search_result", ClientPayload::Result(payload)), + Self::VideoFailed(payload) => ("video_failed", ClientPayload::Failure(payload)), } } } @@ -63,6 +65,7 @@ enum ClientPayload { TransferFailure(TransferFailurePayload), Progress(CrfSearchProgressPayload), Result(CrfSearchResultPayload), + Failure(FailureReportPayload), } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -145,6 +148,18 @@ pub(crate) struct CrfSearchResultPayload { pub(crate) chosen: bool, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct FailureReportPayload { + pub(crate) video_id: u64, + pub(crate) stage: String, + pub(crate) category: String, + pub(crate) message: String, + pub(crate) code: String, + pub(crate) context: serde_json::Value, + pub(crate) retriable: bool, + pub(crate) stderr_excerpt: Option, +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct ServerFrame( pub(crate) Option,