diff --git a/Cargo.lock b/Cargo.lock index 1d2cd034..b99dd902 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9,7 +9,6 @@ dependencies = [ "allocation-counter", "anyhow", "async-stream", - "base64", "blake3", "clap", "clap-verbosity-flag", @@ -43,6 +42,7 @@ dependencies = [ "tokio-tungstenite", "tracing", "tracing-subscriber", + "ureq", ] [[package]] @@ -485,6 +485,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dunce" version = "1.0.5" @@ -547,6 +558,15 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + [[package]] name = "fs2" version = "0.4.3" @@ -716,6 +736,109 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -814,6 +937,12 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + [[package]] name = "lock_api" version = "0.4.14" @@ -966,6 +1095,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -984,6 +1119,15 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1543,6 +1687,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "strsim" version = "0.11.1" @@ -1566,6 +1716,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "sysinfo" version = "0.30.13" @@ -1686,6 +1847,16 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tokio" version = "1.52.3" @@ -1940,12 +2111,42 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e497bb1f828cc9fb236722c2eaa100dcf201563f38f4da6252357a59037adf31" +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64", + "log", + "once_cell", + "url", +] + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + [[package]] name = "utf-8" version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "utf8parse" version = "0.2.2" @@ -2196,6 +2397,35 @@ version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.8.52" @@ -2216,12 +2446,66 @@ dependencies = [ "syn", ] +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + [[package]] name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index 816bc441..c2effee2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,6 @@ 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" @@ -36,6 +35,7 @@ shell-escape = "0.1.5" sled = "0.34.7" sysinfo = "0.30" thiserror = "2" +ureq = { version = "2.12", default-features = false } rustls = { version = "0.23", features = ["ring"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/src/command/worker.rs b/src/command/worker.rs index 2e83cbc0..61f8e706 100644 --- a/src/command/worker.rs +++ b/src/command/worker.rs @@ -1,23 +1,26 @@ use crate::command::worker_protocol::{ - AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ChunkTransferPayload, - ClientEvent, ClientFrame, CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, - FailureReportPayload, HeartbeatPayload, JobResultPayload, PullWorkPayload, ReplyBody, - ServerPushFrame, ServerReply, TransferFailurePayload, TransferProgressPayload, TransferStage, - TransferStartedPayload, WorkStatus, + AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ClientEvent, ClientFrame, + CrfSearchProgressPayload, CrfSearchResultPayload, ErrorReplyPayload, FailureReportPayload, + HeartbeatPayload, PullWorkPayload, ReplyBody, ServerPushFrame, ServerReply, + TransferFailurePayload, TransferProgressPayload, TransferStage, TransferStartedPayload, + WorkStatus, }; use crate::command::worker_transfer::{Chunk, ChunkReceiver}; use crate::command::{crf_search, sample_encode}; use crate::ffprobe::Ffprobe; 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, json}; use std::{ - fs, + fs, io, + io::Read, path::{Path, PathBuf}, - sync::{Arc, Mutex, OnceLock}, + sync::{ + Arc, Mutex, OnceLock, + atomic::{AtomicU64, Ordering}, + }, time::{Duration, Instant}, }; use sysinfo::{Disks, Pid, System}; @@ -37,6 +40,7 @@ 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(10); +const HTTP_TRANSFER_PROGRESS_INTERVAL: Duration = Duration::from_millis(500); static HEARTBEAT_SYSTEM: OnceLock> = OnceLock::new(); /// Connect to a Reencodarr websocket worker endpoint and request one job. @@ -163,28 +167,20 @@ impl WorkerJob { Ok(config) } - fn result_payload(&self, best: &crf_search::Sample) -> JobResultPayload { - JobResultPayload { - job_id: self.assignment.job_id.clone(), - video_id: self.assignment.video_id, - source_name: self.assignment.source_name.clone(), - crf: best.crf, - vmaf_score: best.enc.vmaf_score, - xpsnr_score: best.enc.xpsnr_score, - predicted_encode_size: best.enc.predicted_encode_size, - encode_percent: best.enc.encode_percent, - predicted_encode_time_secs: best.enc.predicted_encode_time.as_secs_f64(), - from_cache: best.enc.from_cache, - } - } - - fn progress_payload(&self, status: &sample_encode::Status) -> CrfSearchProgressPayload { + fn progress_payload( + &self, + crf: f32, + 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, + crf, + sample_num: status.sample, + total_samples: status.samples, } } @@ -234,7 +230,6 @@ 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")] @@ -270,22 +265,6 @@ impl PendingJob { 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 ensure_receiver(&mut self, chunk_size_bytes: u64) -> Result<()> { if self.receiver.is_some() { return Ok(()); @@ -363,14 +342,14 @@ impl PendingJob { .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 elapsed = self.transfer_started_at.elapsed().as_secs().max(1); + let bytes_per_second = received_bytes / 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 eta = (bytes_per_second > 0).then_some(remaining_bytes / bytes_per_second); let percent = if self.job.assignment.size_bytes == 0 { 100.0 } else { @@ -454,12 +433,10 @@ async fn run_worker_job_with_reporting( 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 state = WorkerJobReportState::default(); let mut reconnect = tokio::time::interval(Duration::from_secs(5)); reconnect.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut completed_best = None; loop { match worker.as_mut() { @@ -468,7 +445,6 @@ async fn run_worker_job_with_reporting( _ = 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, @@ -484,7 +460,6 @@ async fn run_worker_job_with_reporting( error = %error, "worker heartbeat failed; reconnecting while job continues" ); - state.connected = false; *worker = None; } } @@ -519,17 +494,14 @@ async fn run_worker_job_with_reporting( 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; } } @@ -542,11 +514,14 @@ async fn run_worker_job_with_reporting( update, ).await?; if disconnected { - state.connected = false; *worker = None; } if let Some(best) = best { - return Ok(best); + if disconnected || worker.is_none() { + completed_best = Some(best); + } else { + return Ok(best); + } } } } @@ -555,14 +530,16 @@ async fn run_worker_job_with_reporting( 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); + if replay_worker_state(&mut reconnected, &state).await { + *worker = Some(reconnected); + if let Some(best) = completed_best.take() { + return Ok(best); + } + } } Err(error) => { trace!(job_id = %job.assignment.job_id, error = %error, "worker reconnect attempt failed"); @@ -577,7 +554,7 @@ async fn run_worker_job_with_reporting( update, ).await?; if let Some(best) = best { - return Ok(best); + completed_best = Some(best); } } } @@ -612,8 +589,8 @@ async fn handle_crf_update( } Ok((Some(best), disconnected)) } - Ok(crf_search::Update::Status { sample, .. }) => { - let payload = job.progress_payload(&sample); + Ok(crf_search::Update::Status { crf, sample, .. }) => { + let payload = job.progress_payload(crf, &sample); state.crf_progress = Some(payload.clone()); let mut disconnected = false; if let Some(worker) = worker { @@ -660,9 +637,11 @@ async fn handle_crf_update( } } -async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobReportState) { +async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobReportState) -> bool { + let mut delivered = true; + if let Some(heartbeat) = &state.heartbeat { - send_worker_event( + delivered &= send_worker_event( worker, ClientEvent::Heartbeat(heartbeat.clone()), "state", @@ -671,7 +650,7 @@ async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobRepo .await; } if let Some(progress) = &state.transfer_progress { - send_worker_event( + delivered &= send_worker_event( worker, ClientEvent::TransferProgress(progress.clone()), "state", @@ -680,7 +659,7 @@ async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobRepo .await; } if let Some(progress) = &state.crf_progress { - send_worker_event( + delivered &= send_worker_event( worker, ClientEvent::CrfSearchProgress(progress.clone()), "state", @@ -689,7 +668,7 @@ async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobRepo .await; } for result in &state.crf_results { - send_worker_event( + delivered &= send_worker_event( worker, ClientEvent::CrfSearchResult(result.clone()), "state", @@ -697,6 +676,8 @@ async fn replay_worker_state(worker: &mut ConnectedWorker, state: &WorkerJobRepo ) .await; } + + delivered } async fn send_worker_event( @@ -813,12 +794,18 @@ impl ReconnectBackoff { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum WorkerJobPhase { - ReceivingInput, - InputMissing, + AwaitingInput(InputDelivery), InputReady, CrfSearching, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum InputDelivery { + Http, + Resend, + Websocket, +} + fn worker_job_phase(job: &WorkerJob, local_path: Option<&Path>) -> Result { if job.input_path().exists() { return match job.assignment.status { @@ -838,9 +825,13 @@ fn worker_job_phase(job: &WorkerJob, local_path: Option<&Path>) -> Result Ok(WorkerJobPhase::ReceivingInput), - WorkStatus::JobInProgress => Ok(WorkerJobPhase::InputMissing), + WorkStatus::JobAssigned => Ok(WorkerJobPhase::AwaitingInput(InputDelivery::Websocket)), + WorkStatus::JobInProgress => Ok(WorkerJobPhase::AwaitingInput(InputDelivery::Resend)), WorkStatus::NoWork => bail!( "assigned job {} has invalid no_work status", job.assignment.job_id @@ -859,7 +850,6 @@ enum PendingJobOutcome { enum WorkerPush { Cancel(CancelPayload), Started(TransferStartedPayload), - Chunk(ChunkTransferPayload), } #[derive(Debug)] @@ -1024,48 +1014,6 @@ impl ConnectedWorker { ); 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" - ); - } - let chunk_index = chunk.chunk_index; - let total_chunks = chunk.total_chunks; - pending_job.apply_chunk(chunk)?; - 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 - }) { - 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), } @@ -1146,9 +1094,9 @@ impl ConnectedWorker { } } -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) +fn format_bytes_per_second(bytes_per_second: u64) -> String { + const MIB: u64 = 1024 * 1024; + format!("{} MiB/s", bytes_per_second / MIB) } async fn run_worker_job_and_publish( @@ -1163,19 +1111,33 @@ 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 = match run_worker_job_with_reporting(config, job.clone(), probe, worker).await { - Ok(best) => best, + match run_worker_job_with_reporting(config, job.clone(), probe, worker).await { + Ok(_) => {} Err(error) => { publish_worker_failure(worker, job, &error).await; return Err(error); } - }; + } + remove_completed_worker_input(config, job)?; + Ok(()) +} - debug!(job_id = %job.assignment.job_id, "publishing worker result"); - println!( - "{}", - serde_json::to_string(&job.result_payload(&best)).context("serialize worker job result")? +fn remove_completed_worker_input(config: &WorkerConfig, job: &WorkerJob) -> Result<()> { + if config.local_path.is_some() { + return Ok(()); + } + + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "removing completed worker input" ); + fs::remove_file(job.input_path()).with_context(|| { + format!( + "remove completed worker input {}", + job.input_path().display() + ) + })?; Ok(()) } @@ -1302,6 +1264,140 @@ fn worker_source_file_name( }) } +async fn download_worker_input(worker: &mut ConnectedWorker, job: &WorkerJob) -> Result { + let Some(transfer) = job.assignment.transfer.clone() else { + return Ok(false); + }; + + if job.input_path().exists() { + return Ok(true); + } + + let parent = job + .input_path() + .parent() + .with_context(|| format!("worker input has no parent: {}", job.input_path().display()))?; + fs::create_dir_all(parent).context("create worker input dir")?; + + let part_path = parent.join(".ab-av1-http.part"); + let input_path = job.input_path().to_path_buf(); + let job_id = job.assignment.job_id.clone(); + let expected_size = job.assignment.size_bytes; + let received = Arc::new(AtomicU64::new(0)); + let copy_received = Arc::clone(&received); + + let mut copy = tokio::task::spawn_blocking(move || -> Result { + let response = ureq::get(&transfer.url) + .set(&transfer.auth.header, &transfer.auth.value) + .call() + .map_err(|error| anyhow!("HTTP input download failed for job {job_id}: {error}"))?; + + let reader = CountingReader { + inner: response.into_reader(), + received: copy_received, + }; + let mut output = + fs::File::create(&part_path).context("create HTTP worker input part file")?; + let bytes = io::copy(&mut reader.take(expected_size), &mut output) + .context("write HTTP worker input")?; + + if expected_size > 0 && bytes != expected_size { + bail!( + "HTTP input download for job {job_id} wrote {bytes} bytes, expected {expected_size}" + ); + } + + fs::rename(&part_path, &input_path).context("move HTTP worker input into place")?; + Ok(bytes) + }); + + let started_at = Instant::now(); + let mut progress = tokio::time::interval(HTTP_TRANSFER_PROGRESS_INTERVAL); + progress.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + _ = progress.tick() => { + let bytes = received.load(Ordering::Relaxed); + if bytes > 0 { + worker + .send_transfer_progress(http_transfer_progress_payload( + &job.assignment.job_id, + job.assignment.video_id, + &job.assignment.source_name, + job.assignment.size_bytes, + bytes, + started_at, + )) + .await?; + } + } + result = &mut copy => { + let bytes = result.context("join HTTP worker input download task")??; + worker + .send_transfer_progress(http_transfer_progress_payload( + &job.assignment.job_id, + job.assignment.video_id, + &job.assignment.source_name, + job.assignment.size_bytes, + bytes, + started_at, + )) + .await?; + break; + } + } + } + + Ok(true) +} + +struct CountingReader { + inner: R, + received: Arc, +} + +impl Read for CountingReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let read = self.inner.read(buf)?; + self.received.fetch_add(read as u64, Ordering::Relaxed); + Ok(read) + } +} + +fn http_transfer_progress_payload( + job_id: &str, + video_id: u64, + filename: &str, + expected_size: u64, + received_bytes: u64, + started_at: Instant, +) -> TransferProgressPayload { + let elapsed = started_at.elapsed().as_secs().max(1); + let bytes_per_second = received_bytes / elapsed; + let remaining_bytes = expected_size.saturating_sub(received_bytes); + let eta = (bytes_per_second > 0).then_some(remaining_bytes / bytes_per_second); + let percent = if expected_size == 0 { + 100.0 + } else { + 100.0 * received_bytes as f64 / expected_size as f64 + }; + + TransferProgressPayload { + job_id: job_id.to_owned(), + transfer_id: job_id.to_owned(), + video_id, + filename: filename.to_owned(), + received_bytes, + expected_bytes: Some(expected_size), + percent, + bytes_per_second, + eta, + chunk_index: 0, + total_chunks: 0, + } +} + pub async fn worker(config: WorkerConfig) -> Result<()> { if config.once { let session = run_worker_session(&config).await?; @@ -1392,6 +1488,33 @@ async fn request_input_resend( build_worker_job(assignment, local_path) } +async fn request_pending_input( + worker: &mut ConnectedWorker, + job: &WorkerJob, + local_path: Option<&Path>, +) -> Result { + let resend_job = request_input_resend(worker, job, local_path).await?; + Ok(PendingJob::waiting(resend_job)) +} + +async fn download_or_wait_for_input( + worker: &mut ConnectedWorker, + job: &WorkerJob, + local_path: Option<&Path>, +) -> Result> { + match download_worker_input(worker, job).await { + Ok(true) => return Ok(None), + Ok(false) => {} + Err(error) => debug!( + job_id = %job.assignment.job_id, + error = %error, + "HTTP worker input download failed, falling back to websocket transfer" + ), + } + + Ok(Some(request_pending_input(worker, job, local_path).await?)) +} + async fn run_connected_worker( config: &WorkerConfig, runtime: WorkerRuntime, @@ -1496,43 +1619,51 @@ async fn run_connected_worker( ); 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" - ); + WorkerJobPhase::AwaitingInput(delivery) => { + let pending = match delivery { + InputDelivery::Http => { + match download_or_wait_for_input( + worker.as_mut().expect("connected worker"), + &job, + config.local_path.as_deref(), + ) + .await? + { + None => { + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "downloaded worker input over HTTP, starting job" + ); + run_worker_job_and_publish(config, &mut worker, &job).await?; + None + } + Some(pending) => Some(pending), + } + } + InputDelivery::Resend => Some( + request_pending_input( + worker.as_mut().expect("connected worker"), + &job, + config.local_path.as_deref(), + ) + .await?, + ), + InputDelivery::Websocket => Some(PendingJob::waiting(job)), + }; + + if let Some(pending) = pending { + debug!( + job_id = %pending.job.assignment.job_id, + input = %pending.input_path().display(), + temp_dir = %pending.job.input_dir.display(), + phase = ?phase, + receiver_ready = false, + pending_job = true, + "waiting for worker input" + ); + pending_job = Some(pending); + } } } continue; @@ -1585,21 +1716,12 @@ fn decode_worker_push(text: &str) -> Result> { } 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" - ); - } + 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()) @@ -1609,15 +1731,6 @@ fn decode_worker_push(text: &str) -> Result> { 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), }; @@ -1878,7 +1991,7 @@ mod tests { use super::*; use crate::command::worker_protocol::{ CancelPayload, ErrorReplyPayload, JobAssignedPayload, ReplyBody, ServerFrame, - ServerPushFrame, WorkStatus, + ServerPushFrame, TransferAuth, TransferSpec, WorkStatus, }; use crate::{command::crf_search::test_hooks as crf_test_hooks, ffprobe::Ffprobe}; use anyhow::Result; @@ -2237,6 +2350,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -2263,6 +2377,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -2293,6 +2408,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -2326,6 +2442,7 @@ mod tests { size_bytes: 4, chunk_size_bytes: 0, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -2364,6 +2481,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -2392,21 +2510,6 @@ mod tests { assert!(best.crf.is_finite()); assert_eq!(best.enc.vmaf_score, Some(97.0)); assert_eq!(best.enc.encode_percent, 50.0); - assert_eq!( - job.result_payload(&best), - JobResultPayload { - job_id: "job-123".into(), - video_id: 123, - source_name: "movie.mkv".into(), - crf: best.crf, - vmaf_score: Some(97.0), - xpsnr_score: None, - predicted_encode_size: 100, - encode_percent: 50.0, - predicted_encode_time_secs: 1.0, - from_cache: false, - } - ); let result = job.crf_result_payload(&best, true); assert_eq!(result.job_id, "job-123"); assert_eq!(result.video_id, 123); @@ -2422,6 +2525,264 @@ mod tests { Ok(()) } + #[test] + fn pending_job_finalizes_chunk_to_worker_input_and_reports_progress() -> Result<()> { + let job_id = format!("worker-flow-chunk-{}", std::process::id()); + let input_dir = worker_job_input_dir(&job_id); + let input_path = input_dir.join("movie.mkv"); + let _ = fs::remove_dir_all(&input_dir); + let job = WorkerJob::new( + JobAssignedPayload { + status: WorkStatus::JobAssigned, + job_id: job_id.clone(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 4, + chunk_size_bytes: 4, + target_vmaf: 95.0, + transfer: None, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "95".into(), + ], + }, + input_dir.clone(), + input_path.clone(), + ); + let mut pending = PendingJob::waiting(job); + pending.apply_raw_chunk(TransferChunk { + transfer_id: job_id, + video_id: 123, + chunk_index: 0, + total_chunks: 1, + bytes_sent: 4, + total_bytes: 4, + crc32: crc32fast::hash(b"data") as u64, + bytes: b"data".to_vec(), + })?; + + let progress = pending.transfer_progress_payload(0, 1); + assert_eq!(progress.received_bytes, 4); + assert_eq!(progress.expected_bytes, Some(4)); + assert_eq!(progress.percent, 100.0); + assert_eq!(progress.chunk_index, 0); + assert_eq!(progress.total_chunks, 1); + + pending.finish()?; + assert_eq!(fs::read(&input_path)?, b"data"); + fs::remove_dir_all(input_dir)?; + Ok(()) + } + + #[test] + fn completed_worker_input_cleanup_respects_local_path() -> Result<()> { + let root = + std::env::temp_dir().join(format!("ab-av1-worker-cleanup-{}", std::process::id())); + let input_path = root.join("movie.mkv"); + fs::create_dir_all(&root)?; + fs::write(&input_path, b"data")?; + let job = WorkerJob::new( + JobAssignedPayload { + status: WorkStatus::JobAssigned, + job_id: "cleanup-job".into(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 4, + chunk_size_bytes: 4, + target_vmaf: 95.0, + transfer: None, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "95".into(), + ], + }, + root.clone(), + input_path.clone(), + ); + let config = WorkerConfig { + connect: String::new(), + token: String::new(), + worker_id: String::new(), + version: String::new(), + protocol_version: 1, + once: false, + local_path: Some(input_path.clone()), + }; + + remove_completed_worker_input(&config, &job)?; + assert!(input_path.exists()); + let config = WorkerConfig { + local_path: None, + ..config + }; + remove_completed_worker_input(&config, &job)?; + assert!(!input_path.exists()); + fs::remove_dir_all(root)?; + Ok(()) + } + + #[test] + fn worker_job_phase_selects_input_delivery_and_resume_behavior() -> Result<()> { + let root = std::env::temp_dir().join(format!("ab-av1-worker-phase-{}", std::process::id())); + let input_path = root.join("movie.mkv"); + let make_job = |status, transfer| { + WorkerJob::new( + JobAssignedPayload { + status, + job_id: "phase-job".into(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 4, + chunk_size_bytes: 4, + target_vmaf: 95.0, + transfer, + crf_search_args: vec![], + }, + root.clone(), + input_path.clone(), + ) + }; + + let _ = fs::remove_dir_all(&root); + assert_eq!( + worker_job_phase(&make_job(WorkStatus::JobAssigned, None), None)?, + WorkerJobPhase::AwaitingInput(InputDelivery::Websocket) + ); + assert_eq!( + worker_job_phase( + &make_job( + WorkStatus::JobAssigned, + Some(TransferSpec { + url: "http://server/input".into(), + auth: TransferAuth { + scheme: "Bearer".into(), + header: "authorization".into(), + value: "token".into(), + }, + }) + ), + None + )?, + WorkerJobPhase::AwaitingInput(InputDelivery::Http) + ); + assert_eq!( + worker_job_phase(&make_job(WorkStatus::JobInProgress, None), None)?, + WorkerJobPhase::AwaitingInput(InputDelivery::Resend) + ); + + fs::create_dir_all(&root)?; + fs::write(&input_path, b"data")?; + assert_eq!( + worker_job_phase(&make_job(WorkStatus::JobAssigned, None), None)?, + WorkerJobPhase::InputReady + ); + assert_eq!( + worker_job_phase(&make_job(WorkStatus::JobInProgress, None), None)?, + WorkerJobPhase::CrfSearching + ); + + fs::remove_file(&input_path)?; + assert!( + worker_job_phase(&make_job(WorkStatus::JobAssigned, None), Some(&input_path)).is_err() + ); + fs::remove_dir_all(root)?; + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn crf_updates_are_retained_for_reconnect_reporting() -> Result<()> { + let job = WorkerJob::new( + JobAssignedPayload { + status: WorkStatus::JobAssigned, + job_id: "reporting-job".into(), + video_id: 123, + source_name: "movie.mkv".into(), + size_bytes: 4, + chunk_size_bytes: 4, + target_vmaf: 95.0, + transfer: None, + crf_search_args: vec![ + "crf-search".into(), + "--input".into(), + "/server/movie.mkv".into(), + "--min-vmaf".into(), + "95".into(), + ], + }, + std::env::temp_dir(), + std::env::temp_dir().join("movie.mkv"), + ); + let mut state = WorkerJobReportState::default(); + let status = sample_encode::Status { + work: sample_encode::Work::Encode, + fps: 24.0, + progress: 0.5, + sample: 2, + samples: 4, + full_pass: false, + }; + let (_, disconnected) = handle_crf_update( + &job, + &mut state, + None, + Some(Ok(crf_search::Update::Status { + crf_run: 1, + crf: 31.0, + sample: status, + })), + ) + .await?; + assert!(!disconnected); + let progress = state.crf_progress.as_ref().expect("stored progress"); + assert_eq!(progress.video_id, 123); + assert_eq!(progress.percent, 50.0); + assert_eq!(progress.fps, 24.0); + assert_eq!(progress.crf, 31.0); + assert_eq!(progress.sample_num, 2); + assert_eq!(progress.total_samples, 4); + + crf_test_hooks::set(|_crf| sample_encode::Output { + vmaf_score: Some(96.0), + xpsnr_score: None, + predicted_encode_size: 100, + encode_percent: 50.0, + predicted_encode_time: Duration::from_secs(1), + from_cache: false, + }); + let sample = run_worker_job( + job.clone(), + Arc::new(Ffprobe { + duration: Ok(Duration::from_secs(600)), + has_audio: false, + max_audio_channels: None, + fps: Ok(24.0), + resolution: Some((1280, 720)), + is_image: false, + pix_fmt: Some("yuv420p10le".into()), + }), + ) + .await?; + crf_test_hooks::clear(); + let (best, disconnected) = handle_crf_update( + &job, + &mut state, + None, + Some(Ok(crf_search::Update::Done(sample))), + ) + .await?; + assert!(!disconnected); + assert!(best.is_some()); + assert_eq!(state.crf_results.len(), 1); + assert!(state.crf_results[0].chosen); + Ok(()) + } + #[test] fn reconnect_backoff_grows_and_caps() { let mut backoff = @@ -2724,6 +3085,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes, target_vmaf: 96.5, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs index d835a6e4..5207a640 100644 --- a/src/command/worker_protocol.rs +++ b/src/command/worker_protocol.rs @@ -1,4 +1,4 @@ -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, Serializer}; pub(crate) const CRF_SEARCH_TOPIC: &str = "workers:crf_search"; @@ -105,6 +105,7 @@ pub(crate) struct Capabilities { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct HeartbeatPayload { #[serde(skip_serializing_if = "Option::is_none")] + #[serde(serialize_with = "serialize_rounded_option_f32")] pub(crate) cpu_percent: Option, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) memory_rss_bytes: Option, @@ -121,10 +122,16 @@ pub(crate) struct HeartbeatPayload { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct CrfSearchProgressPayload { pub(crate) video_id: u64, + #[serde(serialize_with = "serialize_rounded_f32")] pub(crate) percent: f32, pub(crate) filename: String, - pub(crate) eta: Option, + pub(crate) eta: Option, + #[serde(serialize_with = "serialize_rounded_f32")] pub(crate) fps: f32, + #[serde(serialize_with = "serialize_rounded_f32")] + pub(crate) crf: f32, + pub(crate) sample_num: u64, + pub(crate) total_samples: u64, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -132,18 +139,27 @@ pub(crate) struct CrfSearchResultPayload { pub(crate) job_id: String, pub(crate) video_id: u64, pub(crate) source_name: String, + #[serde(serialize_with = "serialize_rounded_f32")] pub(crate) crf: f32, + #[serde(serialize_with = "serialize_rounded_option_f32")] pub(crate) vmaf_score: Option, + #[serde(serialize_with = "serialize_rounded_option_f32")] pub(crate) xpsnr_score: Option, pub(crate) predicted_encode_size: u64, + #[serde(serialize_with = "serialize_rounded_f64")] pub(crate) encode_percent: f64, + #[serde(serialize_with = "serialize_rounded_f64")] pub(crate) predicted_encode_time_secs: f64, pub(crate) from_cache: bool, + #[serde(serialize_with = "serialize_rounded_f32")] pub(crate) score: f32, + #[serde(serialize_with = "serialize_rounded_f64")] pub(crate) percent: f64, pub(crate) size: u64, + #[serde(serialize_with = "serialize_rounded_f64")] pub(crate) time: f64, pub(crate) params: serde_json::Value, + #[serde(serialize_with = "serialize_rounded_f32")] pub(crate) target: f32, pub(crate) chosen: bool, } @@ -303,21 +319,21 @@ pub(crate) struct JobAssignedPayload { pub(crate) target_vmaf: f32, #[serde(default)] pub(crate) crf_search_args: Vec, + #[serde(default)] + pub(crate) transfer: Option, } -#[cfg_attr(not(test), allow(dead_code))] -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub(crate) struct JobResultPayload { - 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, +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct TransferSpec { + pub(crate) url: String, + pub(crate) auth: TransferAuth, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct TransferAuth { + pub(crate) scheme: String, + pub(crate) header: String, + pub(crate) value: String, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -347,20 +363,6 @@ pub(crate) struct TransferStartedPayload { 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, -} - #[cfg_attr(not(test), allow(dead_code))] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct TransferProgressPayload { @@ -370,20 +372,39 @@ pub(crate) struct TransferProgressPayload { pub(crate) filename: String, pub(crate) received_bytes: u64, pub(crate) expected_bytes: Option, + #[serde(serialize_with = "serialize_rounded_f64")] pub(crate) percent: f64, - pub(crate) bytes_per_second: f64, - pub(crate) eta: Option, + pub(crate) bytes_per_second: u64, + pub(crate) eta: Option, pub(crate) chunk_index: u64, pub(crate) total_chunks: u64, } -#[allow(dead_code)] -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub(crate) struct TransferCompletePayload { - pub(crate) job_id: String, - pub(crate) final_path: String, - pub(crate) final_size_bytes: u64, - pub(crate) final_digest: String, +fn serialize_rounded_f32(value: &f32, serializer: S) -> Result +where + S: Serializer, +{ + serializer.serialize_f64(round_two(*value as f64)) +} + +fn serialize_rounded_option_f32(value: &Option, serializer: S) -> Result +where + S: Serializer, +{ + value + .map(|value| round_two(value as f64)) + .serialize(serializer) +} + +fn serialize_rounded_f64(value: &f64, serializer: S) -> Result +where + S: Serializer, +{ + serializer.serialize_f64(round_two(*value)) +} + +fn round_two(value: f64) -> f64 { + (value * 100.0).round() / 100.0 } #[cfg_attr(not(test), allow(dead_code))] @@ -488,6 +509,118 @@ mod tests { ); } + #[test] + fn outbound_float_payloads_round_to_two_places() { + assert_eq!( + serde_json::to_value(HeartbeatPayload { + cpu_percent: Some(12.346), + memory_rss_bytes: None, + memory_total_bytes: None, + disk_free_bytes: None, + disk_total_bytes: None, + active_video_id: None, + }) + .expect("serialize heartbeat floats"), + json!({ "cpu_percent": 12.35 }) + ); + + assert_eq!( + serde_json::to_value(CrfSearchProgressPayload { + video_id: 123, + percent: 42.346, + filename: "movie.mkv".into(), + eta: None, + fps: 27.346, + crf: 31.346, + sample_num: 2, + total_samples: 4, + }) + .expect("serialize crf progress floats"), + json!({ + "video_id": 123, + "percent": 42.35, + "filename": "movie.mkv", + "eta": null, + "fps": 27.35, + "crf": 31.35, + "sample_num": 2, + "total_samples": 4, + }) + ); + + assert_eq!( + serde_json::to_value(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: 66.666, + bytes_per_second: 256, + eta: Some(2), + chunk_index: 3, + total_chunks: 8, + }) + .expect("serialize transfer progress floats"), + json!({ + "job_id": "job-123", + "transfer_id": "job-123", + "video_id": 123, + "filename": "movie.mkv", + "received_bytes": 512, + "expected_bytes": 1024, + "percent": 66.67, + "bytes_per_second": 256, + "eta": 2, + "chunk_index": 3, + "total_chunks": 8, + }) + ); + + assert_eq!( + serde_json::to_value(CrfSearchResultPayload { + job_id: "job-123".into(), + video_id: 123, + source_name: "movie.mkv".into(), + crf: 31.346, + vmaf_score: Some(96.346), + xpsnr_score: Some(95.346), + predicted_encode_size: 123, + encode_percent: 42.346, + predicted_encode_time_secs: 87.346, + from_cache: false, + score: 96.346, + percent: 42.346, + size: 123, + time: 87.346, + params: json!({}), + target: 95.346, + chosen: true, + }) + .expect("serialize crf result floats"), + json!({ + "job_id": "job-123", + "video_id": 123, + "source_name": "movie.mkv", + "crf": 31.35, + "vmaf_score": 96.35, + "xpsnr_score": 95.35, + "predicted_encode_size": 123, + "encode_percent": 42.35, + "predicted_encode_time_secs": 87.35, + "from_cache": false, + "score": 96.35, + "percent": 42.35, + "size": 123, + "time": 87.35, + "params": {}, + "target": 95.35, + "chosen": true, + }) + ); + } + #[test] fn transfer_progress_serializes_transfer_stats_event() { let frame = ClientFrame::new( @@ -500,8 +633,8 @@ mod tests { received_bytes: 512, expected_bytes: Some(1024), percent: 50.0, - bytes_per_second: 256.0, - eta: Some(2.0), + bytes_per_second: 256, + eta: Some(2), chunk_index: 3, total_chunks: 8, }), @@ -522,8 +655,8 @@ mod tests { "received_bytes": 512, "expected_bytes": 1024, "percent": 50.0, - "bytes_per_second": 256.0, - "eta": 2.0, + "bytes_per_second": 256, + "eta": 2, "chunk_index": 3, "total_chunks": 8, } @@ -541,6 +674,9 @@ mod tests { filename: "movie.mkv".into(), eta: None, fps: 27.25, + crf: 31.0, + sample_num: 2, + total_samples: 4, }), ); @@ -557,6 +693,9 @@ mod tests { "filename": "movie.mkv", "eta": null, "fps": 27.25, + "crf": 31.0, + "sample_num": 2, + "total_samples": 4, } ]) ); @@ -599,13 +738,13 @@ mod tests { "video_id": 123, "source_name": "movie.mkv", "crf": 31.5, - "vmaf_score": 96.19999694824219, + "vmaf_score": 96.2, "xpsnr_score": null, "predicted_encode_size": 123456, "encode_percent": 42.5, "predicted_encode_time_secs": 87.5, "from_cache": true, - "score": 96.19999694824219, + "score": 96.2, "percent": 42.5, "size": 123456, "time": 87.5, @@ -684,6 +823,14 @@ mod tests { "size_bytes": 1024, "chunk_size_bytes": 256, "target_vmaf": 96.5, + "transfer": { + "url": "http://10.0.0.10:4000/workers/files/123", + "auth": { + "scheme": "bearer", + "header": "authorization", + "value": "Bearer transfer-token" + } + }, "crf_search_args": [ "crf-search", "--input", @@ -708,6 +855,14 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 256, target_vmaf: 96.5, + transfer: Some(super::TransferSpec { + url: "http://10.0.0.10:4000/workers/files/123".into(), + auth: super::TransferAuth { + scheme: "bearer".into(), + header: "authorization".into(), + value: "Bearer transfer-token".into(), + }, + }), crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -760,6 +915,7 @@ mod tests { size_bytes: 1024, chunk_size_bytes: 0, target_vmaf: 95.0, + transfer: None, crf_search_args: vec![ "crf-search".into(), "--input".into(), @@ -772,38 +928,6 @@ mod tests { ); } - #[test] - fn job_result_payload_serializes_structured_result_summary() { - let payload = JobResultPayload { - 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: false, - }; - - assert_eq!( - serde_json::to_value(payload).expect("serialize job result"), - json!({ - "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": false, - }) - ); - } - #[test] fn server_error_reply_parses_protocol_mismatch_payload() { let reply: ServerFrame = serde_json::from_value(json!([ @@ -887,36 +1011,6 @@ mod tests { ); } - #[test] - fn chunk_transfer_payload_serializes_metadata_side_channel() { - let payload = ChunkTransferPayload { - 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!({ - "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, - }) - ); - } - #[test] fn pull_work_payload_omits_default_and_reports_missing_input() { assert_eq!(