diff --git a/.gitignore b/.gitignore index 6ea2ceb5..809acc62 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ /target /aur +.envrc +.direnv/ *.log .DS_Store *.json diff --git a/Cargo.lock b/Cargo.lock index 64e259a8..e2826c91 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15,7 +15,6 @@ dependencies = [ "clap_complete", "console", "dirs", - "env_logger", "fastrand", "ffprobe", "futures-util", @@ -26,6 +25,7 @@ dependencies = [ "pin-project-lite", "proptest", "rstest", + "rustls", "same-file", "serde", "serde_json", @@ -33,10 +33,13 @@ dependencies = [ "shell-escape", "sled", "test-case", - "thiserror", + "thiserror 2.0.18", "tokio", "tokio-process-tools", "tokio-stream", + "tokio-tungstenite", + "tracing", + "tracing-subscriber", ] [[package]] @@ -90,7 +93,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -101,7 +104,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -150,6 +153,29 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "aws-lc-rs" +version = "1.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4342d8937fc7e5dd9b1c60292261c0670c882a2cd1719cfc11b1af41731e32ad" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d9ceb1da931507a12f4fccea479dccd00da1943e1b4ae72d8e502d707361444" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + [[package]] name = "bit-set" version = "0.8.0" @@ -188,7 +214,16 @@ dependencies = [ "cc", "cfg-if", "constant_time_eq", - "cpufeatures", + "cpufeatures 0.3.0", +] + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", ] [[package]] @@ -216,6 +251,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -291,6 +328,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -306,7 +352,7 @@ dependencies = [ "encode_unicode", "libc", "unicode-width", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -315,6 +361,31 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "cpufeatures" version = "0.3.0" @@ -349,35 +420,29 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] -name = "defmt" -version = "1.1.0" +name = "crypto-common" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e524506490a1953d237cb87b1cfc1e46f88c18f10a22dfe0f507dc6bfc7f7f" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ - "bitflags 1.3.2", - "defmt-macros", + "generic-array", + "typenum", ] [[package]] -name = "defmt-macros" -version = "1.1.0" +name = "data-encoding" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0a27770e9c8f719a79d8b638281f4d828f77d8fd61e0bd94451b9b85e576a0b" -dependencies = [ - "defmt-parser", - "proc-macro-error2", - "proc-macro2", - "quote", - "syn", -] +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] -name = "defmt-parser" -version = "1.0.0" +name = "digest" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "thiserror", + "block-buffer", + "crypto-common", ] [[package]] @@ -398,37 +463,20 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] -name = "encode_unicode" -version = "1.0.0" +name = "dunce" +version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" [[package]] -name = "env_filter" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" -dependencies = [ - "log", - "regex", -] - -[[package]] -name = "env_logger" -version = "0.11.11" +name = "encode_unicode" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" -dependencies = [ - "anstream", - "anstyle", - "env_filter", - "jiff", - "log", -] +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" [[package]] name = "equivalent" @@ -443,7 +491,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -484,6 +532,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures-core" version = "0.3.32" @@ -512,6 +566,12 @@ dependencies = [ "syn", ] +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + [[package]] name = "futures-task" version = "0.3.32" @@ -532,6 +592,7 @@ checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-core", "futures-macro", + "futures-sink", "futures-task", "pin-project-lite", "slab", @@ -546,6 +607,16 @@ dependencies = [ "byteorder", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -598,6 +669,22 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + [[package]] name = "humantime" version = "2.3.0" @@ -655,28 +742,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] -name = "jiff" -version = "0.2.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34f877a98676d2fb664698d74cc6a51ce6c484ce8c770f05d0108ec9090aeb46" -dependencies = [ - "defmt", - "jiff-static", - "log", - "portable-atomic", - "portable-atomic-util", - "serde_core", -] - -[[package]] -name = "jiff-static" -version = "0.2.29" +name = "jobserver" +version = "0.1.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0666b5ab5ecaca213fc2a85b8c0083d9004e84ee2d5f9a7e0017aaf50986f25f" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" dependencies = [ - "proc-macro2", - "quote", - "syn", + "getrandom 0.4.3", + "libc", ] [[package]] @@ -690,6 +762,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -726,6 +804,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "memchr" version = "2.8.2" @@ -740,7 +827,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "wasi", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -755,6 +842,15 @@ dependencies = [ "libc", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -776,6 +872,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + [[package]] name = "option-ext" version = "0.2.0" @@ -837,19 +939,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] -name = "portable-atomic" -version = "1.13.1" +name = "pkg-config" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] -name = "portable-atomic-util" -version = "0.2.7" +name = "portable-atomic" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" -dependencies = [ - "portable-atomic", -] +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" [[package]] name = "ppv-lite86" @@ -869,28 +968,6 @@ dependencies = [ "toml_edit", ] -[[package]] -name = "proc-macro-error-attr2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" -dependencies = [ - "proc-macro2", - "quote", -] - -[[package]] -name = "proc-macro-error2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" -dependencies = [ - "proc-macro-error-attr2", - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "proc-macro2" version = "1.0.106" @@ -910,8 +987,8 @@ dependencies = [ "bit-vec", "bitflags 2.13.0", "num-traits", - "rand", - "rand_chacha", + "rand 0.9.4", + "rand_chacha 0.9.0", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -946,14 +1023,35 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -963,7 +1061,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -981,7 +1088,7 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -1010,7 +1117,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1048,6 +1155,20 @@ version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + [[package]] name = "rstest" version = "0.24.0" @@ -1097,7 +1218,56 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +dependencies = [ + "aws-lc-rs", + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "aws-lc-rs", + "ring", + "rustls-pki-types", + "untrusted", ] [[package]] @@ -1127,12 +1297,44 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags 2.13.0", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "semver" version = "1.0.28" @@ -1207,6 +1409,26 @@ dependencies = [ "syn", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shell-escape" version = "0.1.5" @@ -1257,12 +1479,28 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "strsim" version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.118" @@ -1284,7 +1522,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1294,7 +1532,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1330,13 +1568,33 @@ dependencies = [ "test-case-core", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -1350,6 +1608,15 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + [[package]] name = "tokio" version = "1.52.3" @@ -1361,8 +1628,9 @@ dependencies = [ "mio", "pin-project-lite", "signal-hook-registry", + "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1385,12 +1653,22 @@ dependencies = [ "bytes", "memchr", "nix", - "thiserror", + "thiserror 2.0.18", "tokio", "tracing", "typed-builder", "unwrap-infallible", - "windows-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", ] [[package]] @@ -1404,6 +1682,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -1441,9 +1735,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -1451,6 +1757,56 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.8.6", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror 1.0.69", + "utf-8", ] [[package]] @@ -1473,6 +1829,12 @@ dependencies = [ "syn", ] +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unarray" version = "0.1.4" @@ -1497,18 +1859,42 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "unwrap-infallible" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e497bb1f828cc9fb236722c2eaa100dcf201563f38f4da6252357a59037adf31" +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8parse" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "wait-timeout" version = "0.2.1" @@ -1610,7 +1996,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1625,6 +2011,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -1634,6 +2029,70 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "winnow" version = "1.0.3" @@ -1669,6 +2128,12 @@ dependencies = [ "syn", ] +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index 7f3e72bd..6fadc8b5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,10 +19,6 @@ clap-verbosity-flag = "3.0.2" clap_complete = "4.4.10" console = "0.16" dirs = "6" -env_logger = { version = "0.11.3", default-features = false, features = [ - "auto-color", - "humantime", -] } fastrand = "2" ffprobe = "0.4" futures-util = "0.3.19" @@ -37,15 +33,21 @@ serde_json = "1.0.105" shell-escape = "0.1.5" sled = "0.34.7" thiserror = "2" +rustls = { version = "0.23", features = ["ring"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } tokio = { version = "1.15", features = [ "rt", "macros", "process", "fs", + "net", "signal", + "time", ] } tokio-process-tools = "0.11.2" tokio-stream = "0.1" +tokio-tungstenite = { version = "0.24", features = ["rustls-tls-native-roots"] } [profile.release] lto = true diff --git a/src/command.rs b/src/command.rs index 874450d5..6a1dcfce 100644 --- a/src/command.rs +++ b/src/command.rs @@ -6,6 +6,9 @@ pub mod print_completions; pub(crate) mod rules; pub mod sample_encode; pub mod vmaf; +pub mod worker; +pub(crate) mod worker_protocol; +pub(crate) mod worker_transfer; pub mod xpsnr; pub use auto_encode::auto_encode; @@ -14,6 +17,7 @@ pub use encode::encode; pub use print_completions::print_completions; pub use sample_encode::sample_encode; pub use vmaf::vmaf; +pub use worker::worker; pub use xpsnr::xpsnr; const PROGRESS_CHARS: &str = "##-"; diff --git a/src/command/encode/mod.rs b/src/command/encode/mod.rs index 9faee027..a03d9e25 100644 --- a/src/command/encode/mod.rs +++ b/src/command/encode/mod.rs @@ -145,6 +145,51 @@ mod tests { assert!(Args::try_parse_from(["ab-av1", "--input", "input.mkv", "--crf", "NaN"]).is_err()); } + #[test] + fn parse_passthrough_errors_are_reported_by_clap() { + let svt_err = match Args::try_parse_from([ + "ab-av1", + "--input", + "input.mkv", + "--crf", + "30", + "--svt", + "crf=32", + ]) { + Ok(_) => panic!("reserved svt arg should fail"), + Err(err) => err, + }; + assert!(svt_err.to_string().contains("crf")); + + let enc_err = match Args::try_parse_from([ + "ab-av1", + "--input", + "input.mkv", + "--crf", + "30", + "--enc", + "-svtav1-params=crf=32", + ]) { + Ok(_) => panic!("reserved encoder arg should fail"), + Err(err) => err, + }; + assert!(enc_err.to_string().contains("svtav1-params")); + + let enc_input_err = match Args::try_parse_from([ + "ab-av1", + "--input", + "input.mkv", + "--crf", + "30", + "--enc-input", + "-svtav1-params=crf=32", + ]) { + Ok(_) => panic!("reserved encoder input arg should fail"), + Err(err) => err, + }; + assert!(enc_input_err.to_string().contains("svtav1-params")); + } + // ab-kgc.89: default output extension must preserve input container for webm/mov #[test_case("clip.mp4", false, "mp4"; "video mp4 keeps mp4")] #[test_case("clip.mkv", false, "mkv"; "video mkv keeps mkv")] diff --git a/src/command/worker.rs b/src/command/worker.rs new file mode 100644 index 00000000..77154f46 --- /dev/null +++ b/src/command/worker.rs @@ -0,0 +1,1283 @@ +use crate::command::worker_protocol::{ + AnnouncePayload, CRF_SEARCH_TOPIC, CancelPayload, Capabilities, ClientEvent, ClientFrame, + ErrorReplyPayload, JobResultPayload, ReplyBody, ServerPushFrame, ServerReply, +}; +use crate::command::{args, crf_search, sample_encode}; +use crate::ffprobe::Ffprobe; +use crate::temporary; +use anyhow::{Context, Result, anyhow, bail}; +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 tokio::net::TcpStream; +use tokio_tungstenite::{ + MaybeTlsStream, WebSocketStream, connect_async, + tungstenite::{Error as WsError, Message}, +}; +use tracing::debug; + +const PHOENIX_VSN: &str = "2.0.0"; +const SUPPORTED_PROTOCOL_VERSION: u64 = 1; + +/// Connect to a Reencodarr websocket worker endpoint and request one job. +#[derive(Parser, Debug, Clone)] +pub struct Args { + /// Reencodarr base URL, e.g. http://127.0.0.1:4000 + #[arg(long)] + connect: String, + + /// Worker authentication token. + #[arg(long, env = "REENCODARR_WORKER_TOKEN")] + token: String, + + /// Client worker id announced to Reencodarr. + #[arg(long)] + worker_id: String, + + /// Worker version announced to Reencodarr. + #[arg(long, default_value = env!("CARGO_PKG_VERSION"))] + version: String, + + /// Protocol version announced to Reencodarr. + #[arg(long, default_value_t = SUPPORTED_PROTOCOL_VERSION)] + protocol_version: u64, + + /// Exit after the first work poll instead of running as a long-lived worker. + #[arg(long)] + once: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerConfig { + connect: String, + token: String, + worker_id: String, + version: String, + protocol_version: u64, + once: bool, +} + +impl From for WorkerConfig { + fn from( + Args { + connect, + token, + worker_id, + version, + protocol_version, + once, + }: Args, + ) -> Self { + Self { + connect, + token, + worker_id, + version, + protocol_version, + once, + } + } +} + +#[derive(Debug, PartialEq)] +pub struct WorkerSession { + pub assigned_worker_id: String, + pub negotiated_protocol_version: u64, + pub work_status: String, + pub assigned_job: Option, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone)] +struct WorkerJob { + assignment: crate::command::worker_protocol::JobAssignedPayload, + input_dir: PathBuf, +} + +#[cfg_attr(not(test), allow(dead_code))] +impl WorkerJob { + fn new( + assignment: crate::command::worker_protocol::JobAssignedPayload, + input_dir: PathBuf, + ) -> Self { + Self { + assignment, + input_dir, + } + } + + fn input_path(&self) -> PathBuf { + self.input_dir.join(&self.assignment.source_name) + } + + fn crf_search_config(&self, encoder: args::Encoder) -> Result { + Ok(crf_search::CrfSearchConfig { + args: args::Encode { + encoder, + input: self.input_path(), + 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 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, + } + } +} + +#[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 +} + +#[cfg_attr(not(test), allow(dead_code))] +async fn run_worker_job_until( + job: WorkerJob, + probe: Arc, + shutdown: S, +) -> Result +where + S: std::future::Future, +{ + let config = job.crf_search_config("libsvtav1".parse().expect("default encoder"))?; + let mut run = std::pin::pin!(crf_search::run(config, probe)); + tokio::pin!(shutdown); + + loop { + tokio::select! { + _ = &mut shutdown => bail!("worker shutdown"), + update = run.next() => match update { + Some(Ok(crf_search::Update::Done(best))) => return Ok(best), + Some(Ok(crf_search::Update::Status { .. })) + | Some(Ok(crf_search::Update::SampleResult { .. })) + | Some(Ok(crf_search::Update::RunResult(_))) => {} + Some(Err(error)) => return Err(error.into()), + None => break, + }, + } + } + + unreachable!("crf-search stream should finish with Done") +} + +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(..) + )) +} + +#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] +struct JoinResponse { + worker_id: String, +} + +#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] +struct AnnounceResponse { + accepted: bool, + protocol_version: u64, +} + +#[derive(Debug, Clone, Copy)] +struct WorkerRuntime { + idle_delay: Duration, + reconnect_base_delay: Duration, + reconnect_max_delay: Duration, + max_pulls: Option, +} + +impl Default for WorkerRuntime { + fn default() -> Self { + Self { + idle_delay: Duration::from_secs(5), + reconnect_base_delay: Duration::ZERO, + reconnect_max_delay: Duration::ZERO, + max_pulls: None, + } + } +} + +#[derive(Debug, Clone, Copy)] +struct ReconnectBackoff { + current: Duration, + base: Duration, + max: Duration, +} + +impl ReconnectBackoff { + fn new(base: Duration, max: Duration) -> Self { + Self { + current: base, + base, + max, + } + } + + fn next_delay(&mut self) -> Duration { + let delay = self.current; + self.current = self.current.saturating_mul(2).min(self.max); + delay + } + + fn reset(&mut self) { + self.current = self.base; + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PendingJobOutcome { + Waiting, + Ready, + Canceled, +} + +type WorkerSocket = WebSocketStream>; + +struct ConnectedWorker { + assigned_worker_id: String, + negotiated_protocol_version: u64, + next_ref: u64, + socket: WorkerSocket, +} + +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))?; + + send_json(&mut socket, ClientFrame::new(1, ClientEvent::Join)).await?; + let join: JoinResponse = expect_reply(&mut socket, "1", "phx_join").await?; + + send_json( + &mut socket, + ClientFrame::new( + 2, + ClientEvent::Announce(AnnouncePayload { + worker_id: config.worker_id.clone(), + protocol_version: config.protocol_version, + version: config.version.clone(), + capabilities: Capabilities { crf_search: true }, + }), + ), + ) + .await?; + let announce: AnnounceResponse = expect_reply(&mut socket, "2", "announce").await?; + if !announce.accepted { + bail!("worker announcement was not accepted"); + } + + Ok(Self { + assigned_worker_id: join.worker_id, + negotiated_protocol_version: announce.protocol_version, + next_ref: 3, + socket, + }) + } + + async fn request_work(&mut self) -> Result { + let request_ref = self.next_ref; + self.next_ref += 1; + + send_json( + &mut self.socket, + ClientFrame::new(request_ref, ClientEvent::PullWork), + ) + .await?; + expect_reply(&mut self.socket, &request_ref.to_string(), "pull_work").await + } + + async fn wait_for_pending_job( + &mut self, + job: &WorkerJob, + idle_delay: Duration, + ) -> Result { + tokio::select! { + frame = self.socket.next() => { + match frame { + Some(Ok(Message::Ping(payload))) => { + self.socket + .send(Message::Pong(payload)) + .await + .context("send websocket pong")?; + Ok(PendingJobOutcome::Waiting) + } + 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 + { + eprintln!("worker job {} canceled: {}", cancel.job_id, cancel.reason); + return Ok(PendingJobOutcome::Canceled); + } + Ok(PendingJobOutcome::Waiting) + } + Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => { + Ok(PendingJobOutcome::Waiting) + } + Some(Ok(Message::Close(frame))) => { + bail!("websocket closed while waiting for worker input: {frame:?}") + } + Some(Err(error)) => Err(error).context("read websocket message"), + None => bail!("websocket ended while waiting for worker input"), + } + } + _ = tokio::time::sleep(idle_delay) => { + if job.input_path().exists() { + Ok(PendingJobOutcome::Ready) + } else { + Ok(PendingJobOutcome::Waiting) + } + } + } + } +} + +async fn run_worker_job_and_publish(job: &WorkerJob) -> Result<()> { + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "starting worker job" + ); + 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?; + + debug!(job_id = %job.assignment.job_id, "publishing worker result"); + println!( + "{}", + serde_json::to_string(&job.result_payload(&best)).context("serialize worker job result")? + ); + Ok(()) +} + +fn build_worker_job( + assignment: crate::command::worker_protocol::JobAssignedPayload, +) -> 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)) +} + +pub async fn worker(config: WorkerConfig) -> Result<()> { + if config.once { + let session = run_worker_session(&config).await?; + println!( + "connected worker {} via {} and received {}", + session.assigned_worker_id, session.negotiated_protocol_version, session.work_status + ); + return Ok(()); + } + + run_worker_until(&config, WorkerRuntime::default()).await?; + Ok(()) +} + +async fn run_worker_until(config: &WorkerConfig, runtime: WorkerRuntime) -> Result<()> { + let mut completed_pulls = 0usize; + let mut reconnect_backoff = + ReconnectBackoff::new(runtime.reconnect_base_delay, runtime.reconnect_max_delay); + + loop { + match run_connected_worker(config, runtime, &mut completed_pulls).await { + Ok(()) => { + reconnect_backoff.reset(); + return Ok(()); + } + Err(error) => { + eprintln!("worker connection lost: {error}"); + tokio::time::sleep(reconnect_backoff.next_delay()).await; + } + } + } +} + +async fn run_connected_worker( + config: &WorkerConfig, + runtime: WorkerRuntime, + completed_pulls: &mut usize, +) -> Result<()> { + debug!( + connect = %config.connect, + worker_id = %config.worker_id, + once = config.once, + "connecting worker" + ); + let mut worker = ConnectedWorker::connect(config).await?; + 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; + } + + 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"); + let work_status = worker.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 + ); + + if let ServerReply::JobAssigned(assignment) = work_status { + let job = build_worker_job(assignment)?; + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "job assigned" + ); + 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 { + debug!( + job_id = %job.assignment.job_id, + input = %job.input_path().display(), + "waiting for worker input file" + ); + pending_job = Some(job); + } + continue; + } + + if runtime.max_pulls == Some(*completed_pulls) { + return Ok(()); + } + + tokio::time::sleep(runtime.idle_delay).await; + } +} + +async fn run_worker_session(config: &WorkerConfig) -> Result { + let mut worker = ConnectedWorker::connect(config).await?; + let pull_work = worker.request_work().await?; + let work_status = work_status_label(&pull_work); + let assigned_job = match pull_work { + ServerReply::JobAssigned(payload) => Some(payload), + ServerReply::NoWork(_) => None, + }; + + Ok(WorkerSession { + assigned_worker_id: worker.assigned_worker_id, + negotiated_protocol_version: worker.negotiated_protocol_version, + work_status, + assigned_job, + }) +} + +fn work_status_label(reply: &ServerReply) -> String { + match reply { + ServerReply::NoWork(payload) => payload.status.as_str().into(), + ServerReply::JobAssigned(payload) => { + format!("{} (job_id={})", payload.status.as_str(), payload.job_id) + } + } +} + +fn decode_cancel_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" { + return Ok(None); + } + + let cancel = serde_json::from_value::(frame.4).context("decode cancel push")?; + Ok(Some(cancel)) +} + +fn websocket_connect_error(request_url: &str, error: WsError) -> anyhow::Error { + match error { + WsError::Http(response) => { + let status = response.status(); + let body = response + .body() + .as_ref() + .map(|bytes| String::from_utf8_lossy(bytes).into_owned()) + .unwrap_or_default(); + anyhow!( + "connect websocket {request_url}: HTTP {status} {}", + body.trim() + ) + } + other => anyhow!("connect websocket {request_url}: {other}"), + } +} + +fn worker_websocket_url(base_url: &str, token: &str) -> Result { + let base_url = base_url.trim_end_matches('/'); + let scheme = match () { + _ if base_url.starts_with("http://") => "ws://", + _ if base_url.starts_with("https://") => "wss://", + _ if base_url.starts_with("ws://") => "ws://", + _ if base_url.starts_with("wss://") => "wss://", + _ => bail!("unsupported websocket base URL: {base_url}"), + }; + let rest = base_url + .split_once("://") + .map(|(_, rest)| rest) + .ok_or_else(|| anyhow!("missing scheme in websocket base URL: {base_url}"))?; + Ok(format!( + "{scheme}{rest}/workers/socket/websocket?token={token}&vsn={PHOENIX_VSN}" + )) +} + +async fn send_json(writer: &mut W, value: T) -> Result<()> +where + W: SinkExt + Unpin, + T: serde::Serialize, +{ + writer + .send(Message::Text( + serde_json::to_string(&value).context("encode websocket message")?, + )) + .await + .context("send websocket message") +} + +async fn expect_reply(reader: &mut R, expected_ref: &str, expected_event: &str) -> Result +where + R: StreamExt> + + Unpin, + T: for<'de> Deserialize<'de>, +{ + while let Some(message) = reader.next().await { + match message.context("read websocket message")? { + Message::Text(text) => { + let ServerPushFrame(_, msg_ref, topic, event, body): ServerPushFrame = + serde_json::from_str(&text).context("decode phoenix frame")?; + if topic != CRF_SEARCH_TOPIC + || event != "phx_reply" + || msg_ref.as_deref() != Some(expected_ref) + { + continue; + } + + let ReplyBody { status, response } = serde_json::from_value::>(body) + .context("decode phoenix reply body")?; + return match status.as_str() { + "ok" => serde_json::from_value(response).context("decode phoenix ok reply"), + "error" => { + let error: ErrorReplyPayload = serde_json::from_value(response) + .context("decode phoenix error reply")?; + let supported_versions = match error.supported_protocol_versions.is_empty() + { + true => String::new(), + false => format!( + " (supported_protocol_versions={:?})", + error.supported_protocol_versions + ), + }; + Err(anyhow!( + "{expected_event} failed: {}{}", + error.reason, + supported_versions + )) + } + other => Err(anyhow!("unexpected phoenix status {other}")), + }; + } + Message::Close(frame) => { + bail!("websocket closed before {expected_event} reply: {frame:?}") + } + Message::Ping(_) | Message::Pong(_) | Message::Binary(_) | Message::Frame(_) => { + continue; + } + } + } + + bail!("websocket ended before {expected_event} reply") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::command::worker_protocol::{ + CancelPayload, ErrorReplyPayload, JobAssignedPayload, ReplyBody, ServerFrame, + ServerPushFrame, WorkStatus, + }; + use crate::{command::crf_search::test_hooks as crf_test_hooks, ffprobe::Ffprobe}; + use anyhow::Result; + use serde_json::{Value, json}; + use std::sync::Arc; + use tokio::net::TcpListener; + use tokio_tungstenite::{accept_async, tungstenite::Message}; + + #[derive(Clone, Copy)] + struct WorkerTestConfig { + once: bool, + protocol_version: u64, + } + + impl WorkerTestConfig { + fn continuous() -> Self { + Self { + once: false, + protocol_version: 1, + } + } + + fn with_protocol_version(protocol_version: u64) -> Self { + Self { + once: false, + protocol_version, + } + } + } + + struct FakeCoordinator { + address: std::net::SocketAddr, + server: tokio::task::JoinHandle<()>, + } + + impl FakeCoordinator { + async fn bind(address: &str) -> Result<(TcpListener, std::net::SocketAddr)> { + let listener = TcpListener::bind(address).await?; + let address = listener.local_addr()?; + Ok((listener, address)) + } + + async fn with_no_work_replies(no_work_replies: usize) -> Result { + let (listener, address) = Self::bind("127.0.0.1:0").await?; + + let server = tokio::spawn(async move { + serve_no_work_session(listener, no_work_replies).await; + }); + + Ok(Self { address, server }) + } + + async fn with_job_assignment() -> Result { + let (listener, address) = Self::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_assigned_reply(&mut writer).await; + }); + + Ok(Self { address, server }) + } + + fn worker_config(&self, config: WorkerTestConfig) -> WorkerConfig { + WorkerConfig { + connect: format!("http://{}", self.address), + token: "test-worker-token".into(), + worker_id: "abav1-dev".into(), + version: "0.11.4".into(), + protocol_version: config.protocol_version, + once: config.once, + } + } + + async fn finish(self) { + self.server.await.expect("server task"); + } + } + + #[test] + fn args_lowers_to_worker_config() { + let config = WorkerConfig::from(Args { + connect: "http://127.0.0.1:4000".into(), + token: "token".into(), + worker_id: "abav1-dev".into(), + version: "0.11.4".into(), + protocol_version: 1, + once: false, + }); + + assert_eq!(config.connect, "http://127.0.0.1:4000"); + assert_eq!(config.token, "token"); + assert_eq!(config.worker_id, "abav1-dev"); + assert_eq!(config.version, "0.11.4"); + assert_eq!(config.protocol_version, 1); + assert!(!config.once); + } + + #[test] + fn worker_websocket_url_rewrites_http_scheme() { + let url = worker_websocket_url("http://127.0.0.1:4000/", "secret").expect("url"); + + assert_eq!( + url, + "ws://127.0.0.1:4000/workers/socket/websocket?token=secret&vsn=2.0.0" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_session_joins_announces_and_requests_work() -> Result<()> { + let coordinator = FakeCoordinator::with_no_work_replies(1).await?; + + let session = + run_worker_session(&coordinator.worker_config(WorkerTestConfig::continuous())).await?; + + assert_eq!( + session, + WorkerSession { + assigned_worker_id: "worker-123".into(), + negotiated_protocol_version: 1, + work_status: "no_work".into(), + assigned_job: None, + } + ); + + coordinator.finish().await; + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_session_exposes_assigned_job_payload() -> Result<()> { + let coordinator = FakeCoordinator::with_job_assignment().await?; + + let session = + run_worker_session(&coordinator.worker_config(WorkerTestConfig::continuous())).await?; + + let job = session.assigned_job.expect("assigned job"); + assert_eq!(session.assigned_worker_id, "worker-123"); + assert_eq!(session.negotiated_protocol_version, 1); + assert_eq!(session.work_status, "job_assigned (job_id=job-123)"); + assert_eq!(job.job_id, "job-123"); + assert_eq!(job.video_id, 123); + assert_eq!(job.source_name, "movie.mkv"); + assert_eq!(job.size_bytes, 1024); + assert_eq!(job.chunk_size_bytes, 256); + assert_eq!(job.target_vmaf, 96.5); + + coordinator.finish().await; + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_stays_connected_and_pulls_work_again_after_no_work() -> Result<()> { + let coordinator = FakeCoordinator::with_no_work_replies(2).await?; + + run_worker_until( + &coordinator.worker_config(WorkerTestConfig::continuous()), + WorkerRuntime { + idle_delay: Duration::from_millis(1), + reconnect_base_delay: Duration::from_millis(1), + reconnect_max_delay: Duration::from_millis(1), + max_pulls: Some(2), + }, + ) + .await?; + + coordinator.finish().await; + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_reconnects_after_disconnect_and_continues_pulling_work() -> Result<()> { + let coordinator = FakeCoordinator::with_no_work_replies(1).await?; + let address = coordinator.address; + + let worker = async move { + run_worker_until( + &coordinator.worker_config(WorkerTestConfig::continuous()), + WorkerRuntime { + idle_delay: Duration::from_millis(1), + reconnect_base_delay: Duration::from_millis(1), + reconnect_max_delay: Duration::from_millis(2), + max_pulls: Some(2), + }, + ) + .await + }; + + let replacement = async move { + tokio::time::sleep(Duration::from_millis(5)).await; + let listener = TcpListener::bind(address) + .await + .expect("bind replacement coordinator"); + serve_no_work_session(listener, 1).await; + }; + + let (worker, _replacement) = tokio::join!(worker, replacement); + worker?; + 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?; + + 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, 99).await; + send_announce_error_reply( + &mut writer, + 2, + json!({ + "reason": "unsupported_protocol_version", + "supported_protocol_versions": [1] + }), + ) + .await; + }); + + let error = run_worker_session( + &FakeCoordinator { + address, + server: tokio::spawn(async {}), + } + .worker_config(WorkerTestConfig::with_protocol_version(99)), + ) + .await + .expect_err("protocol mismatch should fail"); + + assert!( + error.to_string().contains("unsupported_protocol_version"), + "unexpected error: {error}" + ); + assert!( + error.to_string().contains("[1]"), + "unexpected error: {error}" + ); + + server.await.expect("server task"); + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_ignores_cancel_push_while_waiting_for_pull_work_reply() -> Result<()> { + 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_cancel_push(&mut writer, "job-123", "shutdown").await; + send_no_work_reply(&mut writer, 3).await; + }); + + let session = run_worker_session( + &FakeCoordinator { + address, + server: tokio::spawn(async {}), + } + .worker_config(WorkerTestConfig::continuous()), + ) + .await?; + + assert_eq!(session.work_status, "no_work"); + + server.await.expect("server task"); + Ok(()) + } + + #[test] + fn worker_formats_assigned_job_status_with_job_id() { + let status = work_status_label(&ServerReply::JobAssigned(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, + })); + + assert_eq!(status, "job_assigned (job_id=job-123)"); + } + + #[test] + 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 job = WorkerJob::new( + 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, + }, + job_dir.clone(), + ); + + let config = job + .crf_search_config("libsvtav1".parse().expect("encoder")) + .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())); + assert_eq!(config.min_vmaf.expect("target vmaf").get(), 96.5); + assert!(config.cache); + } + + #[tokio::test(flavor = "current_thread")] + async fn worker_job_runs_crf_search_from_fake_probe() -> Result<()> { + crf_test_hooks::set(|_crf| sample_encode::Output { + vmaf_score: Some(97.0), + xpsnr_score: None, + predicted_encode_size: 100, + encode_percent: 50.0, + predicted_encode_time: Duration::from_secs(1), + from_cache: false, + }); + + let job_dir = + std::env::temp_dir().join(format!("ab-av1-worker-exec-{}", std::process::id())); + let job = WorkerJob::new( + 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, + }, + job_dir, + ); + + let probe = 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()), + }); + + let best = run_worker_job(job.clone(), probe).await?; + crf_test_hooks::clear(); + + 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, + } + ); + Ok(()) + } + + #[test] + fn reconnect_backoff_grows_and_caps() { + let mut backoff = + ReconnectBackoff::new(Duration::from_millis(100), Duration::from_millis(1_000)); + + assert_eq!(backoff.next_delay(), Duration::from_millis(100)); + assert_eq!(backoff.next_delay(), Duration::from_millis(200)); + assert_eq!(backoff.next_delay(), Duration::from_millis(400)); + assert_eq!(backoff.next_delay(), Duration::from_millis(800)); + assert_eq!(backoff.next_delay(), Duration::from_millis(1_000)); + assert_eq!(backoff.next_delay(), Duration::from_millis(1_000)); + } + + #[test] + fn reconnect_backoff_resets_after_success() { + let mut backoff = + ReconnectBackoff::new(Duration::from_millis(100), Duration::from_millis(1_000)); + + assert_eq!(backoff.next_delay(), Duration::from_millis(100)); + assert_eq!(backoff.next_delay(), Duration::from_millis(200)); + + backoff.reset(); + + assert_eq!(backoff.next_delay(), Duration::from_millis(100)); + } + + async fn expect_join(reader: &mut R) + where + R: StreamExt> + + Unpin, + { + assert_text_message( + reader + .next() + .await + .expect("join frame") + .expect("join message"), + serde_json::to_value(ClientFrame::new(1, ClientEvent::Join)).expect("join frame json"), + ); + } + + async fn expect_announce(reader: &mut R, protocol_version: u64) + where + R: StreamExt> + + Unpin, + { + assert_text_message( + reader + .next() + .await + .expect("announce frame") + .expect("announce message"), + serde_json::to_value(ClientFrame::new( + 2, + ClientEvent::Announce(AnnouncePayload { + worker_id: "abav1-dev".into(), + protocol_version, + version: "0.11.4".into(), + capabilities: Capabilities { crf_search: true }, + }), + )) + .expect("announce frame json"), + ); + } + + async fn expect_pull_work(reader: &mut R, request_ref: u64) + where + R: StreamExt> + + Unpin, + { + assert_text_message( + reader + .next() + .await + .expect("pull_work frame") + .expect("pull_work message"), + serde_json::to_value(ClientFrame::new(request_ref, ClientEvent::PullWork)) + .expect("pull_work frame json"), + ); + } + + async fn send_join_reply(writer: &mut W) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + json!([null, "1", CRF_SEARCH_TOPIC, "phx_reply", { + "status": "ok", + "response": {"worker_id": "worker-123"} + }]) + .to_string(), + )) + .await + .expect("send join reply"); + } + + async fn send_announce_reply(writer: &mut W) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + serde_json::to_string(&ServerFrame::::reply( + 2, + ReplyBody::ok(json!({"accepted": true, "protocol_version": 1})), + )) + .expect("announce reply json"), + )) + .await + .expect("send announce reply"); + } + + async fn send_no_work_reply(writer: &mut W, request_ref: u64) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + serde_json::to_string(&ServerFrame::reply( + request_ref, + ReplyBody::ok(ServerReply::NoWork( + crate::command::worker_protocol::NoWorkPayload { + status: WorkStatus::NoWork, + }, + )), + )) + .expect("no_work reply json"), + )) + .await + .expect("send pull_work reply"); + } + + async fn send_job_assigned_reply(writer: &mut W) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + serde_json::to_string(&ServerFrame::reply( + 3, + ReplyBody::ok(ServerReply::JobAssigned(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, + })), + )) + .expect("job assigned reply json"), + )) + .await + .expect("send job assigned reply"); + } + + async fn send_announce_error_reply(writer: &mut W, request_ref: u64, response: Value) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + serde_json::to_string(&crate::command::worker_protocol::ServerFrame::reply( + request_ref, + ReplyBody::error( + serde_json::from_value::(response) + .expect("error payload"), + ), + )) + .expect("announce error reply json"), + )) + .await + .expect("send announce error reply"); + } + + async fn send_cancel_push(writer: &mut W, job_id: &str, reason: &str) + where + W: SinkExt + Unpin, + { + writer + .send(Message::Text( + serde_json::to_string(&ServerPushFrame::new( + "cancel", + CancelPayload { + job_id: job_id.into(), + reason: reason.into(), + }, + )) + .expect("cancel push json"), + )) + .await + .expect("send cancel push"); + } + + async fn serve_no_work_session(listener: TcpListener, no_work_replies: usize) { + 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; + + for request_ref in 3..(3 + no_work_replies as u64) { + expect_pull_work(&mut reader, request_ref).await; + send_no_work_reply(&mut writer, request_ref).await; + } + } + + fn assert_text_message(message: Message, expected: Value) { + let Message::Text(text) = message else { + panic!("expected text frame, got {message:?}"); + }; + let actual: Value = serde_json::from_str(&text).expect("decode message"); + assert_eq!(actual, expected); + } +} diff --git a/src/command/worker_protocol.rs b/src/command/worker_protocol.rs new file mode 100644 index 00000000..f453b805 --- /dev/null +++ b/src/command/worker_protocol.rs @@ -0,0 +1,481 @@ +use serde::{Deserialize, Serialize}; + +pub(crate) const CRF_SEARCH_TOPIC: &str = "workers:crf_search"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub(crate) struct ClientFrame(String, String, String, String, ClientPayload); + +impl ClientFrame { + pub(crate) fn new(reference: u64, event: ClientEvent) -> Self { + let (event_name, payload) = event.into_parts(); + Self( + "1".into(), + reference.to_string(), + CRF_SEARCH_TOPIC.into(), + event_name.into(), + payload, + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ClientEvent { + Join, + Announce(AnnouncePayload), + PullWork, +} + +impl ClientEvent { + fn into_parts(self) -> (&'static str, ClientPayload) { + match self { + Self::Join => ("phx_join", ClientPayload::Empty(EmptyPayload {})), + Self::Announce(payload) => ("announce", ClientPayload::Announce(payload)), + Self::PullWork => ("pull_work", ClientPayload::Empty(EmptyPayload {})), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(untagged)] +enum ClientPayload { + Empty(EmptyPayload), + Announce(AnnouncePayload), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +struct EmptyPayload {} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct AnnouncePayload { + pub(crate) worker_id: String, + pub(crate) protocol_version: u64, + pub(crate) version: String, + pub(crate) capabilities: Capabilities, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct Capabilities { + pub(crate) crf_search: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct ServerFrame( + pub(crate) Option, + pub(crate) String, + pub(crate) String, + pub(crate) String, + pub(crate) ReplyBody, +); + +impl ServerFrame { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn reply(reference: u64, body: ReplyBody) -> Self { + Self( + None, + reference.to_string(), + CRF_SEARCH_TOPIC.into(), + "phx_reply".into(), + body, + ) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct ServerPushFrame( + pub(crate) Option, + pub(crate) Option, + pub(crate) String, + pub(crate) String, + pub(crate) T, +); + +impl ServerPushFrame { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn new(event: &str, payload: T) -> Self { + Self(None, None, CRF_SEARCH_TOPIC.into(), event.into(), payload) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct ReplyBody { + pub(crate) status: String, + pub(crate) response: T, +} + +impl ReplyBody { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn ok(response: T) -> Self { + Self { + status: "ok".into(), + response, + } + } + + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn error(response: T) -> Self { + Self { + status: "error".into(), + response, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum ServerReply { + JobAssigned(JobAssignedPayload), + NoWork(NoWorkPayload), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum WorkStatus { + NoWork, + JobAssigned, +} + +impl WorkStatus { + pub(crate) fn as_str(&self) -> &'static str { + match self { + Self::NoWork => "no_work", + Self::JobAssigned => "job_assigned", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct NoWorkPayload { + pub(crate) status: WorkStatus, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct JobAssignedPayload { + pub(crate) status: WorkStatus, + pub(crate) job_id: String, + pub(crate) video_id: u64, + pub(crate) source_name: String, + pub(crate) size_bytes: u64, + pub(crate) chunk_size_bytes: u64, + pub(crate) target_vmaf: f32, +} + +#[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 ErrorReplyPayload { + pub(crate) reason: String, + #[serde(default)] + pub(crate) supported_protocol_versions: Vec, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct CancelPayload { + pub(crate) job_id: String, + 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) size_bytes: u64, + pub(crate) checksum: String, +} + +#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct TransferProgressPayload { + pub(crate) job_id: String, + pub(crate) received_bytes: u64, + pub(crate) expected_bytes: Option, +} + +#[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, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct TransferFailurePayload { + pub(crate) job_id: String, + pub(crate) stage: TransferStage, + pub(crate) retriable: bool, + pub(crate) reason: String, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum TransferStage { + ReceiveChunk, + ValidateChunk, + FinalizeTransfer, + RunCrfSearch, +} + +impl ErrorReplyPayload { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + supported_protocol_versions: Vec::new(), + } + } + + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn with_supported_protocol_versions(mut self, versions: Vec) -> Self { + self.supported_protocol_versions = versions; + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn announce_request_serializes_to_current_reencodarr_contract() { + let frame = ClientFrame::new( + 2, + ClientEvent::Announce(AnnouncePayload { + worker_id: "abav1-dev".into(), + protocol_version: 1, + version: "0.11.4".into(), + capabilities: Capabilities { crf_search: true }, + }), + ); + + assert_eq!( + serde_json::to_value(frame).expect("serialize announce"), + json!([ + "1", + "2", + "workers:crf_search", + "announce", + { + "worker_id": "abav1-dev", + "protocol_version": 1, + "version": "0.11.4", + "capabilities": { "crf_search": true } + } + ]) + ); + } + + #[test] + fn server_reply_parses_current_no_work_payload() { + let reply: ServerFrame = serde_json::from_value(json!([ + null, + "3", + "workers:crf_search", + "phx_reply", + { + "status": "ok", + "response": { "status": "no_work" } + } + ])) + .expect("parse 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!([ + null, + "4", + "workers:crf_search", + "phx_reply", + { + "status": "ok", + "response": { + "status": "job_assigned", + "job_id": "job-123", + "video_id": 123, + "source_name": "movie.mkv", + "size_bytes": 1024, + "chunk_size_bytes": 256, + "target_vmaf": 96.5 + } + } + ])) + .expect("parse job_assigned reply"); + + assert_eq!( + reply, + ServerFrame::reply( + 4, + ReplyBody::ok(ServerReply::JobAssigned(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, + })), + ) + ); + } + + #[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!([ + null, + "2", + "workers:crf_search", + "phx_reply", + { + "status": "error", + "response": { + "reason": "unsupported_protocol_version", + "supported_protocol_versions": [1] + } + } + ])) + .expect("parse error reply"); + + assert_eq!( + reply, + ServerFrame::reply( + 2, + ReplyBody::error( + ErrorReplyPayload::new("unsupported_protocol_version") + .with_supported_protocol_versions(vec![1]), + ), + ) + ); + } + + #[test] + fn server_push_parses_cancel_payload() { + let push: ServerPushFrame = serde_json::from_value(json!([ + null, + null, + "workers:crf_search", + "cancel", + { + "job_id": "job-123", + "reason": "shutdown" + } + ])) + .expect("parse cancel push"); + + assert_eq!( + push, + ServerPushFrame::new( + "cancel", + CancelPayload { + job_id: "job-123".into(), + reason: "shutdown".into(), + }, + ) + ); + } + + #[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(), + }; + + 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", + }) + ); + } + + #[test] + fn transfer_failure_payload_serializes_stage_and_retry_hint() { + let payload = TransferFailurePayload { + job_id: "job-123".into(), + stage: TransferStage::FinalizeTransfer, + retriable: true, + reason: "disk full".into(), + }; + + assert_eq!( + serde_json::to_value(payload).expect("serialize transfer failure"), + json!({ + "job_id": "job-123", + "stage": "finalize_transfer", + "retriable": true, + "reason": "disk full", + }) + ); + } +} diff --git a/src/command/worker_transfer.rs b/src/command/worker_transfer.rs new file mode 100644 index 00000000..352a8ec0 --- /dev/null +++ b/src/command/worker_transfer.rs @@ -0,0 +1,384 @@ +use crate::temporary; +use anyhow::{Context, Result}; +use blake3::Hash; +use std::{ + fs::{self, File, OpenOptions}, + io::Write, + path::{Path, PathBuf}, +}; +use tracing::debug; + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Chunk { + pub index: u64, + pub offset: u64, + pub bytes: Vec, + pub checksum: Hash, +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug, thiserror::Error)] +pub(crate) enum ChunkReceiverError { + #[error("chunk {index} arrived out of order; expected index {expected_index}")] + OutOfOrder { index: u64, expected_index: u64 }, + #[error("chunk {index} starts at offset {offset}, expected {expected_offset}")] + UnexpectedOffset { + index: u64, + offset: u64, + expected_offset: u64, + }, + #[error("chunk {index} failed checksum validation")] + CorruptChunk { index: u64 }, + #[error("chunk {index} was already received")] + DuplicateChunk { index: u64 }, + #[error("final digest mismatch")] + DigestMismatch, + #[error("final size mismatch")] + SizeMismatch, + #[error("destination already exists")] + DestinationExists, + #[error("file would exceed max size {max_size} bytes")] + FileTooLarge { size: u64, max_size: u64 }, + #[error("receiver already finished")] + Finished, + #[error(transparent)] + Io(#[from] std::io::Error), +} + +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Debug)] +pub(crate) struct ChunkReceiver { + final_path: PathBuf, + temp_path: PathBuf, + file: Option, + max_size: Option, + next_index: u64, + next_offset: u64, + finished: bool, + hasher: blake3::Hasher, +} + +#[cfg_attr(not(test), allow(dead_code))] +impl ChunkReceiver { + pub(crate) fn new( + final_path: impl Into, + temp_dir: impl AsRef, + max_size: Option, + ) -> 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" + ); + + Ok(Self { + final_path, + temp_path, + file: Some(file), + max_size, + next_index: 0, + next_offset: 0, + finished: false, + hasher: blake3::Hasher::new(), + }) + } + + pub(crate) fn received_bytes(&self) -> u64 { + self.next_offset + } + + pub(crate) fn push(&mut self, chunk: Chunk) -> Result<(), ChunkReceiverError> { + if self.finished { + return Err(ChunkReceiverError::Finished); + } + debug!( + index = chunk.index, + offset = chunk.offset, + size = chunk.bytes.len(), + final_path = %self.final_path.display(), + "receiving chunk" + ); + if chunk.index < self.next_index { + return Err(ChunkReceiverError::DuplicateChunk { index: chunk.index }); + } + if chunk.index > self.next_index { + return Err(ChunkReceiverError::OutOfOrder { + index: chunk.index, + expected_index: self.next_index, + }); + } + if chunk.offset != self.next_offset { + return Err(ChunkReceiverError::UnexpectedOffset { + index: chunk.index, + offset: chunk.offset, + expected_offset: self.next_offset, + }); + } + if self.max_size.is_some_and(|max_size| { + self.next_offset.saturating_add(chunk.bytes.len() as u64) > max_size + }) { + return Err(ChunkReceiverError::FileTooLarge { + size: self.next_offset.saturating_add(chunk.bytes.len() as u64), + max_size: self.max_size.expect("checked max_size"), + }); + } + if blake3::hash(&chunk.bytes) != chunk.checksum { + return Err(ChunkReceiverError::CorruptChunk { index: chunk.index }); + } + + self.file + .as_mut() + .expect("open chunk file") + .write_all(&chunk.bytes)?; + self.hasher.update(&chunk.bytes); + self.next_index += 1; + self.next_offset += chunk.bytes.len() as u64; + Ok(()) + } + + pub(crate) fn finish( + mut self, + expected_size: Option, + expected_digest: Option, + ) -> std::result::Result { + if self.finished { + return Err(ChunkReceiverError::Finished); + } + if expected_size.is_some_and(|size| size != self.next_offset) { + return Err(ChunkReceiverError::SizeMismatch); + } + let digest = self.hasher.finalize(); + if expected_digest.is_some_and(|expected| expected != digest) { + return Err(ChunkReceiverError::DigestMismatch); + } + + self.file.as_mut().expect("open chunk file").sync_all()?; + let _ = self.file.take(); + if let Some(parent) = self.final_path.parent() { + fs::create_dir_all(parent)?; + } + if self.final_path.exists() { + return Err(ChunkReceiverError::DestinationExists); + } + debug!( + temp_path = %self.temp_path.display(), + final_path = %self.final_path.display(), + "finalizing chunk receiver" + ); + 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) + } +} + +impl Drop for ChunkReceiver { + fn drop(&mut self) { + if !self.finished { + let _ = self.file.take(); + let _ = fs::remove_file(&self.temp_path); + let _ = temporary::unadd(&self.temp_path); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serial_test::serial; + + fn temp_paths(label: &str) -> (PathBuf, PathBuf) { + let root = std::env::temp_dir().join(format!( + "ab-av1-worker-transfer-{}-{}-{}", + label, + std::process::id(), + fastrand::u64(..) + )); + let final_path = root.join("movie.mkv"); + (root, final_path) + } + + fn chunk(index: u64, offset: u64, bytes: &[u8]) -> Chunk { + Chunk { + index, + offset, + bytes: bytes.to_vec(), + checksum: blake3::hash(bytes), + } + } + + #[serial] + #[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"); + + receiver.push(chunk(0, 0, b"hello ")).expect("chunk 0"); + 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() { + let temp_dir = std::env::temp_dir().join(format!( + "ab-av1-worker-transfer-parent-{}-{}", + std::process::id(), + fastrand::u64(..) + )); + let final_path = temp_dir.join("nested").join("movie.mkv"); + let mut receiver = ChunkReceiver::new(&final_path, &temp_dir, None).expect("receiver"); + + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + receiver + .finish(Some(5), Some(blake3::hash(b"hello"))) + .expect("finish"); + + assert_eq!(fs::read(&final_path).expect("read final"), b"hello"); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[test] + fn finish_refuses_to_overwrite_existing_destination() { + let temp_dir = std::env::temp_dir().join(format!( + "ab-av1-worker-transfer-overwrite-{}-{}", + std::process::id(), + fastrand::u64(..) + )); + let final_path = temp_dir.join("nested").join("movie.mkv"); + 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"); + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + + assert!(matches!( + receiver.finish(Some(5), Some(blake3::hash(b"hello"))), + Err(ChunkReceiverError::DestinationExists) + )); + assert_eq!(fs::read(&final_path).expect("read final"), b"existing"); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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 bad = chunk(0, 0, b"hello"); + bad.checksum = blake3::hash(b"hell0"); + + assert!(matches!( + receiver.push(bad), + Err(ChunkReceiverError::CorruptChunk { index: 0 }) + )); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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"); + + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + assert!(matches!( + receiver.push(chunk(1, 10, b"world")), + Err(ChunkReceiverError::UnexpectedOffset { + index: 1, + offset: 10, + expected_offset: 5 + }) + )); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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"); + + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + assert!(matches!( + receiver.push(chunk(0, 0, b"hello")), + Err(ChunkReceiverError::DuplicateChunk { index: 0 }) + )); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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"); + + assert!(matches!( + receiver.push(chunk(1, 0, b"hello")), + Err(ChunkReceiverError::OutOfOrder { + index: 1, + expected_index: 0 + }) + )); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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"); + + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + assert!(matches!( + receiver.finish(Some(5), Some(blake3::hash(b"hell0"))), + Err(ChunkReceiverError::DigestMismatch) + )); + let _ = fs::remove_dir_all(temp_dir); + } + + #[serial] + #[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"); + + receiver.push(chunk(0, 0, b"hello")).expect("chunk 0"); + assert_eq!(receiver.received_bytes(), 5); + assert!(matches!( + receiver.push(chunk(1, 5, b"world!")), + Err(ChunkReceiverError::FileTooLarge { + size: 11, + max_size: 10 + }) + )); + let _ = fs::remove_dir_all(temp_dir); + } +} diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs index 417a71ae..06e8f6c5 100644 --- a/src/ffmpeg.rs +++ b/src/ffmpeg.rs @@ -378,10 +378,10 @@ mod tests { assert_ne!(hash_a.finish(), hash_b.finish()); } - // ab-kgc.24: duplicate ffmpeg flags should survive single removal + // ab-kgc.24: remove_arg mirrors main and only strips the first matching flag pair #[test] - fn remove_arg_strips_only_first_matching_pair() { - // setup — duplicate flags should require multiple removals or retain extras + fn remove_arg_strips_first_matching_pair() { + // setup — duplicate flags should leave later pairs alone let mut args = vec![ Arc::new("-preset".to_string()), Arc::new("8".to_string()), @@ -394,7 +394,7 @@ mod tests { // execute remove_arg(&mut args, "-preset"); - // assert — second -preset pair must survive one removal + // assert — only the first pair is removed assert_eq!( args.iter().map(|a| a.as_str()).collect::>(), vec!["-preset", "6", "-crf", "30"] diff --git a/src/main.rs b/src/main.rs index ed74c0f5..dc723307 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,12 +13,11 @@ mod test_support; mod vmaf; mod xpsnr; -use ::log::LevelFilter; use anyhow::anyhow; use clap::Parser; use futures_util::FutureExt; -use std::io::IsTerminal; use tokio::signal; +use tracing_subscriber::{fmt, EnvFilter}; #[derive(Parser)] #[command(version, about)] @@ -29,21 +28,24 @@ enum Command { Encode(command::encode::Args), CrfSearch(command::crf_search::Args), AutoEncode(command::auto_encode::Args), + Worker(command::worker::Args), PrintCompletions(command::print_completions::Args), } #[tokio::main(flavor = "current_thread")] async fn main() { - env_logger::builder() - .filter_module( - "ab_av1", - match std::io::stderr().is_terminal() { - true => LevelFilter::Off, - false => LevelFilter::Info, - }, - ) - .parse_default_env() - .init(); + rustls::crypto::ring::default_provider() + .install_default() + .expect("install rustls crypto provider"); + + let default_filter = if cfg!(debug_assertions) { + "ab_av1=debug" + } else { + "ab_av1=info" + }; + let filter = + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default_filter)); + fmt().with_env_filter(filter).init(); let action = Command::parse(); let keep = action.keep_temp_files(); @@ -65,6 +67,7 @@ async fn main() { Command::AutoEncode(args) => { command::auto_encode(command::auto_encode::AutoEncodeConfig::from(args)).boxed_local() } + Command::Worker(args) => command::worker(args.into()).boxed_local(), Command::PrintCompletions(args) => return command::print_completions(args.into()), }); @@ -131,6 +134,7 @@ mod tests { .expect("parse print-completions command"); match command { + Command::Worker(_) => {} Command::PrintCompletions(_) => {} _ => panic!("expected print-completions command"), } diff --git a/tests/cli_compat.rs b/tests/cli_compat.rs index 7ccb0914..989079b6 100644 --- a/tests/cli_compat.rs +++ b/tests/cli_compat.rs @@ -23,6 +23,7 @@ fn top_level_help_mentions_all_commands() { "encode", "crf-search", "auto-encode", + "worker", "print-completions", ] { assert!(help.contains(cmd), "missing command: {cmd}"); diff --git a/tests/worker_harness.rs b/tests/worker_harness.rs new file mode 100644 index 00000000..ed075864 --- /dev/null +++ b/tests/worker_harness.rs @@ -0,0 +1,127 @@ +#![allow(unused_crate_dependencies)] + +use anyhow::Result; +use futures_util::{Sink, SinkExt, StreamExt}; +use serde_json::{Value, json}; +use std::process::Stdio; +use tokio::{net::TcpListener, process::Command}; +use tokio_tungstenite::{accept_async, tungstenite::Message}; + +#[tokio::test(flavor = "current_thread")] +async fn worker_binary_handles_job_assignment() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + + 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_text_message( + reader + .next() + .await + .expect("join frame") + .expect("join message"), + json!(["1", "1", "workers:crf_search", "phx_join", {}]), + ); + send_text_message( + &mut writer, + json!([null, "1", "workers:crf_search", "phx_reply", { + "status": "ok", + "response": {"worker_id": "worker-123"} + }]), + ) + .await; + + expect_text_message( + reader + .next() + .await + .expect("announce frame") + .expect("announce message"), + json!(["1", "2", "workers:crf_search", "announce", { + "worker_id": "abav1-dev", + "protocol_version": 1, + "version": env!("CARGO_PKG_VERSION"), + "capabilities": {"crf_search": true} + }]), + ); + send_text_message( + &mut writer, + json!([null, "2", "workers:crf_search", "phx_reply", { + "status": "ok", + "response": {"accepted": true, "protocol_version": 1} + }]), + ) + .await; + + expect_text_message( + reader + .next() + .await + .expect("pull_work frame") + .expect("pull_work message"), + json!(["1", "3", "workers:crf_search", "pull_work", {}]), + ); + send_text_message( + &mut writer, + json!([null, "3", "workers:crf_search", "phx_reply", { + "status": "ok", + "response": { + "status": "job_assigned", + "job_id": "job-123", + "video_id": 123, + "source_name": "movie.mkv", + "size_bytes": 1024, + "chunk_size_bytes": 256, + "target_vmaf": 96.5 + } + }]), + ) + .await; + }); + + let output = Command::new(env!("CARGO_BIN_EXE_ab-av1")) + .args([ + "worker", + "--connect", + &format!("http://{address}"), + "--worker-id", + "abav1-dev", + "--once", + ]) + .env("REENCODARR_WORKER_TOKEN", "test-worker-token") + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await?; + + assert!(output.status.success(), "worker failed: {:?}", output); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + stdout.contains("received job_assigned (job_id=job-123)"), + "unexpected worker output: {stdout}" + ); + + server.await.expect("server task"); + Ok(()) +} + +async fn send_text_message(writer: &mut W, value: Value) +where + W: Sink + Unpin, +{ + writer + .send(Message::Text(value.to_string())) + .await + .expect("send websocket message"); +} + +fn expect_text_message(message: Message, expected: Value) { + let Message::Text(text) = message else { + panic!("expected text frame, got {message:?}"); + }; + let actual: Value = serde_json::from_str(&text).expect("decode websocket message"); + assert_eq!(actual, expected); +}