diff --git a/Cargo.lock b/Cargo.lock index 5959e87..080b707 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -74,7 +74,7 @@ version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cexpr", "clang-sys", "itertools", @@ -86,6 +86,12 @@ dependencies = [ "syn", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.11.1" @@ -170,7 +176,7 @@ dependencies = [ "semver", "serde", "serde_json", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -194,6 +200,16 @@ dependencies = [ "nom", ] +[[package]] +name = "cfg-expr" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02" +dependencies = [ + "smallvec", + "target-lexicon", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -251,7 +267,7 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fa961b519f0b462e3a3b4a34b64d119eeaca1d59af726fe450bbba07a9fc0a1" dependencies = [ - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -467,6 +483,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dircpy" +version = "0.3.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebcbec2b9a580ddee352ac38523d2ecd4dcaad53532957034394556909e27f4b" +dependencies = [ + "jwalk", + "log", + "walkdir", +] + [[package]] name = "directories" version = "6.0.0" @@ -534,7 +561,16 @@ dependencies = [ name = "engine" version = "0.3.0" dependencies = [ + "async-channel", + "chrono", + "dashmap", "enginelib", + "futures", + "serde", + "tmq", + "tokio", + "toml 1.1.2+spec-1.1.0", + "tracing", ] [[package]] @@ -555,7 +591,7 @@ dependencies = [ "rust-rocksdb", "serde", "tokio", - "toml", + "toml 1.1.2+spec-1.1.0", "tracing", "tracing-subscriber", "tracing-test", @@ -650,12 +686,94 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", + "futures-sink", +] + [[package]] name = "futures-core" version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -739,7 +857,7 @@ dependencies = [ "parking_lot", "signal-hook 0.3.18", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -752,7 +870,7 @@ dependencies = [ "gix-date", "gix-utils", "itoa", - "thiserror", + "thiserror 2.0.18", "winnow 0.7.15", ] @@ -769,7 +887,7 @@ dependencies = [ "gix-trace", "kstring", "smallvec", - "thiserror", + "thiserror 2.0.18", "unicode-bom", ] @@ -779,7 +897,7 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d982fc7ef0608e669851d0d2a6141dae74c60d5a27e8daa451f2a4857bbf41e2" dependencies = [ - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -788,7 +906,7 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c356b3825677cb6ff579551bb8311a81821e184453cbd105e2fc5311b288eeb" dependencies = [ - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -814,7 +932,7 @@ dependencies = [ "gix-chunk", "gix-hash", "memmap2", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -832,7 +950,7 @@ dependencies = [ "gix-sec", "memchr", "smallvec", - "thiserror", + "thiserror 2.0.18", "unicode-bom", "winnow 0.7.15", ] @@ -843,11 +961,11 @@ version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2409cffa4fe8b303847d5b6ba8df9da9ba65d302fc5ee474ea0cac5afde79840" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bstr", "gix-path", "libc", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -860,7 +978,7 @@ dependencies = [ "itoa", "jiff", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -884,7 +1002,7 @@ dependencies = [ "gix-traverse", "gix-worktree", "imara-diff", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -904,7 +1022,7 @@ dependencies = [ "gix-trace", "gix-utils", "gix-worktree", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -920,7 +1038,7 @@ dependencies = [ "gix-path", "gix-ref", "gix-sec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -936,7 +1054,7 @@ dependencies = [ "libc", "once_cell", "prodash", - "thiserror", + "thiserror 2.0.18", "walkdir", "zlib-rs", ] @@ -959,7 +1077,7 @@ dependencies = [ "gix-trace", "gix-utils", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -973,7 +1091,7 @@ dependencies = [ "gix-features", "gix-path", "gix-utils", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -982,7 +1100,7 @@ version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8546300aee4c65c5862c22a3e321124a69b654a61a8b60de546a9284812b7e2" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bstr", "gix-features", "gix-path", @@ -997,7 +1115,7 @@ dependencies = [ "faster-hex", "gix-features", "sha1-checked", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1030,7 +1148,7 @@ version = "0.45.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea6d3e9e11647ba49f441dea0782494cc6d2875ff43fa4ad9094e6957f42051" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bstr", "filetime", "fnv", @@ -1049,7 +1167,7 @@ dependencies = [ "memmap2", "rustix", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1060,7 +1178,7 @@ checksum = "115268ae5e3b3b7bc7fc77260eecee05acca458e45318ca45d35467fa81a3ac5" dependencies = [ "gix-tempfile", "gix-utils", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1080,7 +1198,7 @@ dependencies = [ "gix-validate", "itoa", "smallvec", - "thiserror", + "thiserror 2.0.18", "winnow 0.7.15", ] @@ -1102,7 +1220,7 @@ dependencies = [ "gix-quote", "parking_lot", "tempfile", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1120,7 +1238,7 @@ dependencies = [ "gix-path", "memmap2", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1132,7 +1250,7 @@ dependencies = [ "bstr", "faster-hex", "gix-trace", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1144,7 +1262,7 @@ dependencies = [ "bstr", "gix-trace", "gix-validate", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1153,13 +1271,13 @@ version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed9e0c881933c37a7ef45288d6c5779c4a7b3ad240b4c37657e1d9829eb90085" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bstr", "gix-attributes", "gix-config-value", "gix-glob", "gix-path", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1177,7 +1295,7 @@ dependencies = [ "gix-transport", "gix-utils", "maybe-async", - "thiserror", + "thiserror 2.0.18", "winnow 0.7.15", ] @@ -1189,7 +1307,7 @@ checksum = "96fc2ff2ec8cc0c92807f02eab1f00eb02619fc2810d13dc42679492fcc36757" dependencies = [ "bstr", "gix-utils", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1209,7 +1327,7 @@ dependencies = [ "gix-utils", "gix-validate", "memmap2", - "thiserror", + "thiserror 2.0.18", "winnow 0.7.15", ] @@ -1225,7 +1343,7 @@ dependencies = [ "gix-revision", "gix-validate", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1234,7 +1352,7 @@ version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "91898c83b18c635696f7355d171cfa74a52f38022ff89581f567768935ebc4c8" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bstr", "gix-commitgraph", "gix-date", @@ -1243,7 +1361,7 @@ dependencies = [ "gix-object", "gix-revwalk", "gix-trace", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1258,7 +1376,7 @@ dependencies = [ "gix-hashtable", "gix-object", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1267,7 +1385,7 @@ version = "0.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea9962ed6d9114f7f100efe038752f41283c225bb507a2888903ac593dffa6be" dependencies = [ - "bitflags", + "bitflags 2.11.1", "gix-path", "libc", "windows-sys", @@ -1282,7 +1400,7 @@ dependencies = [ "bstr", "gix-hash", "gix-lock", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1305,7 +1423,7 @@ dependencies = [ "gix-pathspec", "gix-worktree", "portable-atomic", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1320,7 +1438,7 @@ dependencies = [ "gix-pathspec", "gix-refspec", "gix-url", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1357,7 +1475,7 @@ dependencies = [ "gix-quote", "gix-sec", "gix-url", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1366,7 +1484,7 @@ version = "0.51.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d052b83d1d1744be95ac6448ac02f95f370a8f6720e466be9ce57146e39f5280" dependencies = [ - "bitflags", + "bitflags 2.11.1", "gix-commitgraph", "gix-date", "gix-hash", @@ -1374,7 +1492,7 @@ dependencies = [ "gix-object", "gix-revwalk", "smallvec", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1387,7 +1505,7 @@ dependencies = [ "gix-features", "gix-path", "percent-encoding", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1408,7 +1526,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b1e63a5b516e970a594f870ed4571a8fdcb8a344e7bd407a20db8bd61dbfde4" dependencies = [ "bstr", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1658,6 +1776,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jwalk" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2735847566356cd2179a2a38264839308f7079fa96e6bd5a42d740460e003c56" +dependencies = [ + "crossbeam", + "rayon", +] + [[package]] name = "kstring" version = "2.0.2" @@ -1711,7 +1839,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" dependencies = [ - "bitflags", + "bitflags 2.11.1", "libc", "plain", "redox_syscall 0.7.4", @@ -2032,13 +2160,33 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -2047,7 +2195,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f450ad9c3b1da563fb6948a8e0fb0fb9269711c9c73d9ea1de5058c79c8d643a" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -2058,7 +2206,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -2147,7 +2295,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.11.1", "errno", "libc", "linux-raw-sys", @@ -2228,6 +2376,15 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_spanned" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" +dependencies = [ + "serde", +] + [[package]] name = "serde_spanned" version = "1.1.1" @@ -2309,6 +2466,12 @@ dependencies = [ "libc", ] +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "smallvec" version = "1.15.1" @@ -2363,6 +2526,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "system-deps" +version = "6.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3e535eb8dded36d55ec13eddacd30dec501792ff23a0b1682c38601b8cf2349" +dependencies = [ + "cfg-expr", + "heck", + "pkg-config", + "toml 0.8.23", + "version-compare", +] + [[package]] name = "tar" version = "0.4.45" @@ -2374,6 +2550,12 @@ dependencies = [ "xattr", ] +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + [[package]] name = "tempfile" version = "3.27.0" @@ -2387,13 +2569,33 @@ dependencies = [ "windows-sys", ] +[[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]] @@ -2474,6 +2676,19 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tmq" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3f41ac3a42f65436eed7e1afe80dbe8a982dcac2ea4581bf61bc2d3dcfb19a1" +dependencies = [ + "futures", + "log", + "thiserror 1.0.69", + "tokio", + "zmq", +] + [[package]] name = "tokio" version = "1.52.0" @@ -2502,6 +2717,18 @@ dependencies = [ "syn", ] +[[package]] +name = "toml" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" +dependencies = [ + "serde", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", + "toml_edit", +] + [[package]] name = "toml" version = "1.1.2+spec-1.1.0" @@ -2510,13 +2737,22 @@ checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" dependencies = [ "indexmap", "serde_core", - "serde_spanned", - "toml_datetime", + "serde_spanned 1.1.1", + "toml_datetime 1.1.1+spec-1.1.0", "toml_parser", "toml_writer", "winnow 1.0.1", ] +[[package]] +name = "toml_datetime" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" +dependencies = [ + "serde", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -2526,6 +2762,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "toml_edit" +version = "0.22.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" +dependencies = [ + "indexmap", + "serde", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", + "winnow 0.7.15", +] + [[package]] name = "toml_parser" version = "1.1.2+spec-1.1.0" @@ -2721,6 +2970,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "version-compare" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03c2856837ef78f57382f06b2b8563a2f512f7185d732608fd9176cb3b8edf0e" + [[package]] name = "version_check" version = "0.9.5" @@ -2834,7 +3089,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags", + "bitflags 2.11.1", "hashbrown 0.15.5", "indexmap", "semver", @@ -2990,7 +3245,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags", + "bitflags 2.11.1", "indexmap", "log", "serde", @@ -3030,6 +3285,16 @@ dependencies = [ "rustix", ] +[[package]] +name = "zeromq-src" +version = "0.2.6+4.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc120b771270365d5ed0dfb4baf1005f2243ae1ae83703265cb3504070f4160b" +dependencies = [ + "cc", + "dircpy", +] + [[package]] name = "zlib-rs" version = "0.5.5" @@ -3042,6 +3307,28 @@ version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zmq" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd3091dd571fb84a9b3e5e5c6a807d186c411c812c8618786c3c30e5349234e7" +dependencies = [ + "bitflags 1.3.2", + "libc", + "zmq-sys", +] + +[[package]] +name = "zmq-sys" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e8351dc72494b4d7f5652a681c33634063bbad58046c1689e75270908fdc864" +dependencies = [ + "libc", + "system-deps", + "zeromq-src", +] + [[package]] name = "zstd-sys" version = "2.0.16+zstd.1.5.7" diff --git a/engine/Cargo.toml b/engine/Cargo.toml index 40b1d9e..c97b436 100644 --- a/engine/Cargo.toml +++ b/engine/Cargo.toml @@ -8,6 +8,34 @@ description = "A Blazingly fast distributed task system" [features] default = [] dev = [] + +[[bin]] +name = "server" +path = "src/bin/server.rs" + +[[bin]] +name = "client" +path = "src/bin/client.rs" + +[[bin]] +name = "bench" +path = "src/bin/bench.rs" + +[[bin]] +name = "bench_server" +path = "src/bin/bench_server.rs" + [dependencies] enginelib = { path = "../enginelib" } +tokio = { version = "1.50.0", features = ["rt-multi-thread", "macros", "sync", "signal", "time"] } +tmq = "0.5.0" +futures = "0.3" +serde = { workspace = true } +toml = { workspace = true } +tracing = "0.1" + +[dev-dependencies] +chrono = { version = "0.4.44", features = ["serde"] } +dashmap = "6.2.1" +async-channel = "2.5.0" diff --git a/engine/src/bench_task.rs b/engine/src/bench_task.rs new file mode 100644 index 0000000..23f6b5f --- /dev/null +++ b/engine/src/bench_task.rs @@ -0,0 +1,57 @@ +//! Small statically linked task used only by the process benchmark. + +use enginelib::Identifier; +use enginelib::task::{Task, Verifiable}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct FibTask { + pub iter: u64, + pub result: u64, +} + +impl Verifiable for FibTask { + fn verify(&self, bytes: &[u8]) -> bool { + enginelib::api::from_bytes::(bytes).is_ok() + } +} + +impl Task for FibTask { + fn to_toml(&self) -> String { + toml::to_string(self).unwrap_or_default() + } + + fn from_toml(&self, data: String) -> Box { + Box::new(toml::from_str::(&data).unwrap_or_default()) + } + + fn get_id(&self) -> Identifier { + ("engine_mod".to_string(), "fib".to_string()) + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + + fn run_cpu(&mut self) { + // Keep work negligible: this benchmark measures engine and transport + // throughput rather than Fibonacci performance. + let iterations = self.iter.min(16); + let mut a = 0u64; + let mut b = 1u64; + for _ in 0..iterations { + let previous = a; + a = b; + b = b.wrapping_add(previous); + } + self.result = a; + } + + fn from_bytes(&self, bytes: &[u8]) -> Box { + Box::new(enginelib::api::from_bytes::(bytes).unwrap_or_default()) + } + + fn to_bytes(&self) -> Vec { + enginelib::api::to_allocvec(self).unwrap_or_default() + } +} diff --git a/engine/src/bin/bench.rs b/engine/src/bin/bench.rs new file mode 100644 index 0000000..f74b11e --- /dev/null +++ b/engine/src/bin/bench.rs @@ -0,0 +1,164 @@ +//! Deployment-faithful throughput benchmark (load generator). +//! +//! Spawns the `bench_server` binary as a **separate OS process** and drives it +//! over real tcp: one submitter connection plus N worker connections, each a real +//! libzmq DEALER. No in-process server, no unbounded channel — the server runs the +//! real bounded pipeline, fed by persistent per-task-type loaders started by the +//! benchmark server (submit writes to the DB; loaders stream records into the +//! lease channel). Stall detection is a safety net: if progress stops for +//! `BENCH_STALL_SECS`, the run reports `stalled` rather than hanging or lying. +//! +//! Usage: bench [submit_batch] [lease_batch] [workers] [port] +//! Emits one JSON line on stdout. + +use std::process::{Command, Stdio}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use engine::bench_task::FibTask; +use engine::client::Client; +use enginelib::Identifier; +use enginelib::task::Task; + +#[tokio::main] +async fn main() { + let mut args = std::env::args().skip(1); + let n: u64 = args.next().and_then(|s| s.parse().ok()).unwrap_or(100_000); + let submit_batch: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(1000); + let lease_batch: u32 = args.next().and_then(|s| s.parse().ok()).unwrap_or(256); + let workers: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(16); + let port: u16 = args.next().and_then(|s| s.parse().ok()).unwrap_or(55610); + + let endpoint = format!("tcp://127.0.0.1:{port}"); + let task_type: Identifier = ("engine_mod".to_string(), "fib".to_string()); + let db_path = format!("target/bench_db_{port}_{}", std::process::id()); + + // Launch the real server as its own process. + let server_bin = std::env::current_exe() + .ok() + .and_then(|p| p.parent().map(|d| d.join("bench_server"))) + .expect("locate bench_server next to bench"); + let mut server = Command::new(&server_bin) + .arg(port.to_string()) + .arg(&db_path) + .stdout(Stdio::null()) + .spawn() + .expect("spawn bench_server"); + // Give it time to bind the ROUTER. + tokio::time::sleep(Duration::from_millis(800)).await; + + let completed = Arc::new(AtomicU64::new(0)); + + // Workers (separate DEALER connections): lease → decode + run → complete. + let mut worker_handles = Vec::new(); + for _ in 0..workers { + let endpoint = endpoint.clone(); + let task_type = task_type.clone(); + let completed = completed.clone(); + worker_handles.push(tokio::spawn(async move { + let mut client = match Client::connect(&endpoint, String::new()) { + Ok(c) => c, + Err(_) => return, + }; + loop { + let leased = match client + .lease(task_type.clone(), "w".into(), lease_batch) + .await + { + Ok(v) => v, + Err(_) => break, + }; + if leased.is_empty() { + continue; + } + let mut results = Vec::with_capacity(leased.len()); + for st in &leased { + let mut task: FibTask = + enginelib::api::from_bytes(&st.bytes).unwrap_or_default(); + task.run_cpu(); + results.push((st.task_id.clone(), task.to_bytes())); + } + match client.complete(task_type.clone(), results).await { + Ok(ok) => { + completed.fetch_add(ok as u64, Ordering::Relaxed); + } + Err(_) => break, + } + } + })); + } + + let payload = FibTask { + iter: 20, + result: 0, + } + .to_bytes(); + + // Submitter: submit all N over its own connection, timed. + let e2e_start = Instant::now(); + let submit_start = Instant::now(); + let submit_task = { + let endpoint = endpoint.clone(); + let task_type = task_type.clone(); + let payload = payload.clone(); + tokio::spawn(async move { + let mut client = Client::connect(&endpoint, String::new()).expect("submitter connect"); + let mut sent = 0u64; + while sent < n { + let this = std::cmp::min(submit_batch as u64, n - sent) as usize; + if client + .submit(task_type.clone(), vec![payload.clone(); this]) + .await + .is_err() + { + break; + } + sent += this as u64; + } + sent + }) + }; + let submitted = submit_task.await.unwrap_or(0); + let submit_secs = submit_start.elapsed().as_secs_f64(); + + // Wait for completion. Stall detection is a safety net (e.g. a wedged loader, + // a crashed worker, or a genuine bug) — it reports `stalled` instead of + // hanging forever, rather than describing any expected behavior. + let stall_secs: u64 = std::env::var("BENCH_STALL_SECS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(20); + let mut last = 0u64; + let mut last_change = Instant::now(); + let mut stalled = false; + loop { + let c = completed.load(Ordering::Relaxed); + if c >= n { + break; + } + if c != last { + last = c; + last_change = Instant::now(); + } else if last_change.elapsed() > Duration::from_secs(stall_secs) { + stalled = true; + break; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + let done = completed.load(Ordering::Relaxed); + let e2e_secs = e2e_start.elapsed().as_secs_f64(); + + for h in worker_handles { + h.abort(); + } + let _ = server.kill(); + let _ = server.wait(); + let _ = std::fs::remove_dir_all(&db_path); + + let submit_tps = submitted as f64 / submit_secs; + let e2e_tps = done as f64 / e2e_secs; + println!( + "{{\"n\":{n},\"submitted\":{submitted},\"completed\":{done},\"stalled\":{stalled},\"workers\":{workers},\"submit_batch\":{submit_batch},\"lease_batch\":{lease_batch},\"submit_secs\":{submit_secs:.4},\"submit_tps\":{submit_tps:.1},\"e2e_secs\":{e2e_secs:.4},\"e2e_tps\":{e2e_tps:.1}}}" + ); +} diff --git a/engine/src/bin/bench_server.rs b/engine/src/bin/bench_server.rs new file mode 100644 index 0000000..34f8279 --- /dev/null +++ b/engine/src/bin/bench_server.rs @@ -0,0 +1,48 @@ +//! Standalone server process for the deployment-faithful benchmark. +//! +//! Mirrors `bin/server.rs` (the deployed pipeline): `populate()` builds the real +//! `bounded(8192)` channels, a startup `load()` recovers any persisted backlog, +//! then `serve()` runs the ROUTER over tcp. The only deviation from a real deploy +//! is that `FibTask` is statically linked here instead of dynamically loaded from +//! a `.rf` mod — that changes registration, not per-task runtime cost. +//! +//! Usage: bench_server + +use std::sync::Arc; + +use engine::bench_task::FibTask; +use enginelib::api::ServerAPI; +use enginelib::task::Task; + +#[tokio::main] +async fn main() { + let mut args = std::env::args().skip(1); + let port: u16 = args.next().and_then(|s| s.parse().ok()).unwrap_or(55610); + let db_path = args.next().unwrap_or_else(|| "target/bench_db".to_string()); + + let task_type = ("engine_mod".to_string(), "fib".to_string()); + + let api = ServerAPI::with_path(&db_path); + let fib: Arc = Arc::new(FibTask::default()); + api.task_registry.tasks.insert(task_type.clone(), fib); + let api = Arc::new(api); + // Real bounded queue + dedup set, as the deployed server builds. Depth is + // configurable (ENGINE_QUEUE_SIZE, default 8192) for the sweep. + let queue_size: usize = std::env::var("ENGINE_QUEUE_SIZE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(8192); + ServerAPI::populate_with(&api, queue_size); + // Match init()'s deployment behavior: run the lease reaper. (At the 3600s TTL + // it never fires within a bench, but it's free and keeps this faithful.) + ServerAPI::spawn_reaper(&api); + // Submit only writes to the DB; one loader per task type feeds the bounded + // lease queue. Transport shards must not create duplicate loaders. + ServerAPI::spawn_loaders(&api); + + let endpoint = format!("tcp://127.0.0.1:{port}"); + eprintln!("bench_server listening on {endpoint} (db={db_path})"); + if let Err(e) = engine::server::serve(api, &endpoint).await { + eprintln!("bench_server exited: {e}"); + } +} diff --git a/engine/src/bin/client.rs b/engine/src/bin/client.rs new file mode 100644 index 0000000..e2626e4 --- /dev/null +++ b/engine/src/bin/client.rs @@ -0,0 +1,85 @@ +//! Thin CLI over [`engine::client::Client`] for driving the engine by hand. +//! +//! Usage: +//! client submit submit one task +//! client lease [max] lease up to max (default 1) +//! client complete complete one leased task +//! client cancel cancel (requeue) a lease +//! +//! Endpoint and auth token come from config.toml (same as the server). + +use engine::client::Client; +use enginelib::config::Config; + +#[tokio::main] +async fn main() { + let args: Vec = std::env::args().skip(1).collect(); + if args.is_empty() { + eprintln!("usage: client ..."); + std::process::exit(2); + } + + let cfg = Config::new(); + let endpoint = format!("tcp://{}", cfg.config_toml.host); + let auth = cfg.config_toml.auth_token.clone().unwrap_or_default(); + + let mut client = match Client::connect(&endpoint, auth) { + Ok(c) => c, + Err(e) => { + eprintln!("connect failed: {e}"); + std::process::exit(1); + } + }; + + let result = match args[0].as_str() { + "submit" if args.len() >= 4 => { + let task_type = (args[1].clone(), args[2].clone()); + let payload = args[3].clone().into_bytes(); + client + .submit(task_type, vec![payload]) + .await + .map(|ids| format!("submitted: {}", ids.join(", "))) + } + "lease" if args.len() >= 3 => { + let task_type = (args[1].clone(), args[2].clone()); + let max: u32 = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(1); + client + .lease(task_type, "cli".to_string(), max) + .await + .map(|tasks| { + let lines: Vec = tasks + .iter() + .map(|t| format!(" {} -> {} bytes", t.task_id, t.bytes.len())) + .collect(); + format!("leased {}:\n{}", tasks.len(), lines.join("\n")) + }) + } + "complete" if args.len() >= 5 => { + let task_type = (args[1].clone(), args[2].clone()); + let results = vec![(args[3].clone(), args[4].clone().into_bytes())]; + client + .complete(task_type, results) + .await + .map(|ok| format!("completed: {ok}")) + } + "cancel" if args.len() >= 4 => { + let task_type = (args[1].clone(), args[2].clone()); + client + .cancel(task_type, args[3].clone()) + .await + .map(|()| "cancelled".to_string()) + } + other => { + eprintln!("unknown or malformed command: {other}"); + std::process::exit(2); + } + }; + + match result { + Ok(msg) => println!("{msg}"), + Err(e) => { + eprintln!("error: {e}"); + std::process::exit(1); + } + } +} diff --git a/engine/src/bin/server.rs b/engine/src/bin/server.rs index 4a78f83..d031367 100644 --- a/engine/src/bin/server.rs +++ b/engine/src/bin/server.rs @@ -1,7 +1,67 @@ +use std::sync::Arc; + +use engine::transport::TransportStats; use enginelib::api::ServerAPI; +use enginelib::config::Config; +use enginelib::events::Events; +use tokio::sync::watch; +use tracing::{error, info}; + +#[tokio::main] +async fn main() { + let config = match Config::load_checked() { + Ok(config) => config, + Err(err) => { + eprintln!("server configuration error: {err}"); + std::process::exit(2); + } + }; + let transport = config.config_toml.transport.clone(); + let endpoint = format!("tcp://{}", config.config_toml.host); + + let api = ServerAPI::init_with_config(config); + Events::ServerStart(&api); + ServerAPI::spawn_loaders(&api); + + let (shutdown_tx, shutdown_rx) = watch::channel(false); + tokio::spawn(async move { + wait_for_shutdown().await; + let _ = shutdown_tx.send(true); + }); + + let stats = Arc::new(TransportStats::default()); + info!("engine server listening on {endpoint}"); + if let Err(err) = engine::transport::serve_with_shutdown( + api, + &endpoint, + transport, + shutdown_rx, + stats.clone(), + ) + .await + { + error!("server exited: {err}"); + } + info!( + admitted = stats.admitted(), + overloaded = stats.overloaded(), + dropped_replies = stats.dropped_replies(), + "server transport stopped" + ); +} -fn main() { - let mut api = ServerAPI::init(); - // let server = enginelib::server::RPC::new(&api); - // server.run(); +async fn wait_for_shutdown() { + #[cfg(unix)] + { + use tokio::signal::unix::{SignalKind, signal}; + let mut terminate = signal(SignalKind::terminate()).expect("install SIGTERM handler"); + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = terminate.recv() => {} + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + } } diff --git a/engine/src/client.rs b/engine/src/client.rs new file mode 100644 index 0000000..746a62f --- /dev/null +++ b/engine/src/client.rs @@ -0,0 +1,146 @@ +//! Reusable async DEALER client for the engine wire protocol. + +use enginelib::Identifier; +use enginelib::api::StoredTask; +use enginelib::error::Error; +use enginelib::nucleus::query::{Query, QueryResult}; +use enginelib::protocol::{Envelope, Request, Response, decode, encode}; +use futures::{SinkExt, StreamExt}; +use tmq::dealer::Dealer; +use tmq::{Context, Multipart, dealer}; + +/// A connected DEALER client. Each call sends one request and awaits one reply; +/// DEALER allows many in flight, but these methods are sequential per `&mut self`. +pub struct Client { + sock: Dealer, + auth: String, +} + +impl Client { + /// Connect a DEALER socket to `endpoint` (e.g. `tcp://[::1]:50051`). + pub fn connect(endpoint: &str, auth: String) -> Result { + let ctx = Context::new(); + let sock = dealer(&ctx) + .connect(endpoint) + .map_err(|e| Error::io_error(format!("Failed to connect DEALER to {endpoint}: {e}")))?; + Ok(Client { sock, auth }) + } + + /// Send one `Request`, await and decode the `Response`. + async fn call(&mut self, req: Request) -> Result { + let env = Envelope { + auth: self.auth.clone(), + req, + }; + let mp: Multipart = vec![encode(&env)?].into(); + self.sock + .send(mp) + .await + .map_err(|e| Error::io_error(format!("send failed: {e}")))?; + + let reply = self + .sock + .next() + .await + .ok_or_else(|| Error::new("connection closed".into()))? + .map_err(|e| Error::io_error(format!("recv failed: {e}")))?; + let frame = reply + .into_iter() + .last() + .ok_or_else(|| Error::new("empty reply".into()))?; + decode(&frame) + } + + /// Submit a batch of task payloads; returns the assigned ids. + pub async fn submit( + &mut self, + task_type: Identifier, + tasks: Vec>, + ) -> Result, Error> { + match self.call(Request::SubmitBatch { task_type, tasks }).await? { + Response::Submitted { task_ids } => Ok(task_ids), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } + + /// Lease up to `max` tasks of one type. The server long-polls for the first + /// task but gives up after its poll window; an empty reply just means + /// "nothing arrived in time" — lease again. + pub async fn lease( + &mut self, + task_type: Identifier, + user_id: String, + max: u32, + ) -> Result, Error> { + match self + .call(Request::LeaseBatch { + task_type, + user_id, + max, + }) + .await? + { + Response::Leased(tasks) => Ok(tasks), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } + + /// Complete a batch of leased tasks; returns how many succeeded. + pub async fn complete( + &mut self, + task_type: Identifier, + results: Vec<(String, Vec)>, + ) -> Result { + match self + .call(Request::CompleteBatch { task_type, results }) + .await? + { + Response::Completed { ok } => Ok(ok), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } + + /// Renew a single lease's TTL. + pub async fn renew(&mut self, task_type: Identifier, task_id: String) -> Result<(), Error> { + match self.call(Request::Renew { task_type, task_id }).await? { + Response::Renewed => Ok(()), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } + + /// Cancel a single lease (requeues the task). + pub async fn cancel(&mut self, task_type: Identifier, task_id: String) -> Result<(), Error> { + match self.call(Request::Cancel { task_type, task_id }).await? { + Response::Cancelled => Ok(()), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } + + /// Run an administrative query. + pub async fn query(&mut self, q: Query) -> Result { + match self.call(Request::Query(q)).await? { + Response::Query(r) => Ok(r), + Response::Overloaded { retry_after_ms } => Err(overloaded(retry_after_ms)), + Response::Err(e) => Err(e), + other => Err(unexpected(other)), + } + } +} + +fn overloaded(retry_after_ms: u64) -> Error { + Error::overloaded(format!("server overloaded; retry after {retry_after_ms}ms")) +} + +fn unexpected(resp: Response) -> Error { + Error::new(format!("unexpected response variant: {resp:?}")) +} diff --git a/engine/src/lib.rs b/engine/src/lib.rs new file mode 100644 index 0000000..24a5770 --- /dev/null +++ b/engine/src/lib.rs @@ -0,0 +1,10 @@ +//! Transport layer for the engine (ZeroMQ over tcp). +//! +//! `tmq`/libzmq lives here, deliberately kept out of `enginelib` so that mods +//! (which link `enginelib`) don't inherit a libzmq build dependency. The wire +//! types are shared via [`enginelib::protocol`]. + +pub mod bench_task; +pub mod client; +pub mod server; +pub mod transport; diff --git a/engine/src/server.rs b/engine/src/server.rs new file mode 100644 index 0000000..1fb1d7c --- /dev/null +++ b/engine/src/server.rs @@ -0,0 +1,139 @@ +//! Request authentication and nucleus dispatch for the ZeroMQ transport. + +use std::sync::Arc; +use std::time::Duration; + +use enginelib::api::ServerAPI; +use enginelib::config::TransportConfig; +use enginelib::error::{Error, ErrorKind}; +use enginelib::events::Events; +use enginelib::nucleus::{cancel, complete, lease, query, renew, submit}; +use enginelib::protocol::{Envelope, Request, Response}; + +/// Serve with the transport limits loaded into `ServerAPI`. +/// +/// Background task loaders are deliberately started by the process/bootstrap +/// code, not here, so adding transport shards cannot accidentally duplicate +/// them. +pub async fn serve(api: Arc, endpoint: &str) -> Result<(), Error> { + crate::transport::serve(api, endpoint).await +} + +fn unauthorized() -> Response { + Response::Err(Error::with_kind(ErrorKind::NotSupported, "auth rejected")) +} + +fn handler_failed() -> Response { + Response::Err(Error::new("request handler failed".into())) +} + +/// Authenticate and dispatch one already decoded, admitted request. +pub(crate) async fn dispatch( + api: Arc, + env: Envelope, + limits: TransportConfig, +) -> Response { + let Envelope { auth, req } = env; + + match req { + // Lease is the only operation that waits asynchronously. Authentication + // is synchronous and mod-provided, so run it on Tokio's blocking pool. + Request::LeaseBatch { + task_type, + user_id, + max, + } => { + let auth_api = api.clone(); + let auth_db = api.db.clone(); + let auth_user = user_id.clone(); + let authorized = tokio::task::spawn_blocking(move || { + Events::CheckAuth(&auth_api, auth_user, auth, auth_db) + }) + .await + .unwrap_or(false); + if !authorized { + return unauthorized(); + } + + // Leave encoding overhead beyond StoredTask's fields. The transport + // performs a final exact encoded-size check as a second line of defense. + let lease_bytes = limits.max_reply_bytes.saturating_sub(1024); + match lease::lease_batch_limited( + api, + task_type, + user_id, + max, + lease_bytes, + Duration::from_millis(limits.lease_long_poll_ms), + ) + .await + { + Ok(tasks) => { + Response::Leased(tasks.into_iter().map(|task| (*task).clone()).collect()) + } + Err(err) => Response::Err(err), + } + } + request => { + // RocksDB and event handlers are synchronous. Keep them off Tokio + // runtime workers; admission limits bound outstanding calls. + tokio::task::spawn_blocking(move || dispatch_blocking(api, auth, request)) + .await + .unwrap_or_else(|_| handler_failed()) + } + } +} + +fn dispatch_blocking(api: Arc, auth: String, req: Request) -> Response { + match req { + Request::SubmitBatch { task_type, tasks } => { + if !Events::CheckAuth(&api, String::new(), auth, api.db.clone()) { + return unauthorized(); + } + match submit::submit_batch(api, task_type, tasks) { + Ok(task_ids) => Response::Submitted { task_ids }, + Err(err) => Response::Err(err), + } + } + Request::CompleteBatch { task_type, results } => { + if !Events::CheckAuth(&api, String::new(), auth, api.db.clone()) { + return unauthorized(); + } + match complete::complete_batch(api, task_type, results) { + Ok(ok) => Response::Completed { ok }, + Err(err) => Response::Err(err), + } + } + Request::Renew { task_type, task_id } => { + if !Events::CheckAuth(&api, String::new(), auth, api.db.clone()) { + return unauthorized(); + } + match renew::renew(api, task_type, &task_id) { + Ok(()) => Response::Renewed, + Err(err) => Response::Err(err), + } + } + Request::Cancel { task_type, task_id } => { + if !Events::CheckAuth(&api, String::new(), auth, api.db.clone()) { + return unauthorized(); + } + match cancel::cancel(api, task_type, task_id) { + Ok(()) => Response::Cancelled, + Err(err) => Response::Err(err), + } + } + Request::Query(query_request) => { + let target = ("core".to_string(), "query".to_string()); + if !Events::CheckAdminAuth(&api, auth, target, api.db.clone()) { + return unauthorized(); + } + match query::query(api, query_request) { + Ok(result) => Response::Query(result), + Err(err) => Response::Err(err), + } + } + Request::LeaseBatch { .. } => { + Response::Err(Error::new("lease dispatched on blocking path".into())) + } + } +} diff --git a/engine/src/transport.rs b/engine/src/transport.rs new file mode 100644 index 0000000..35251ef --- /dev/null +++ b/engine/src/transport.rs @@ -0,0 +1,566 @@ +//! Bounded, single-owner ZeroMQ transport pipeline. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use enginelib::api::ServerAPI; +use enginelib::config::TransportConfig; +use enginelib::error::Error; +use enginelib::nucleus::query::Query; +use enginelib::protocol::{Envelope, Request, Response, decode, encode}; +use futures::{SinkExt, StreamExt}; +use tmq::{Context, Multipart, router}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, watch}; +use tokio::task::JoinSet; +use tokio::time::Instant; +use tracing::{error, warn}; + +#[derive(Debug, Clone, Copy)] +enum OperationClass { + LeasePoll, + Mutation, + Query, +} + +impl OperationClass { + fn of(request: &Request) -> Self { + match request { + Request::LeaseBatch { .. } => Self::LeasePoll, + Request::Query(_) => Self::Query, + _ => Self::Mutation, + } + } +} + +/// In-process counters are always available to tests and operators even when no +/// external metrics recorder is installed. Labels are intentionally bounded. +#[derive(Debug, Default)] +pub struct TransportStats { + admitted: AtomicUsize, + overloaded: AtomicUsize, + dropped_replies: AtomicUsize, + send_failures: AtomicUsize, + active: AtomicUsize, + max_active: AtomicUsize, + active_long_polls: AtomicUsize, + max_active_long_polls: AtomicUsize, + reply_count: AtomicUsize, + reply_bytes: AtomicUsize, + max_reply_count: AtomicUsize, + max_reply_bytes: AtomicUsize, + shutdown_drops: AtomicUsize, +} + +impl TransportStats { + pub fn admitted(&self) -> usize { + self.admitted.load(Ordering::Relaxed) + } + pub fn overloaded(&self) -> usize { + self.overloaded.load(Ordering::Relaxed) + } + pub fn dropped_replies(&self) -> usize { + self.dropped_replies.load(Ordering::Relaxed) + } + pub fn active(&self) -> usize { + self.active.load(Ordering::Relaxed) + } + pub fn max_active(&self) -> usize { + self.max_active.load(Ordering::Relaxed) + } + pub fn max_active_long_polls(&self) -> usize { + self.max_active_long_polls.load(Ordering::Relaxed) + } + pub fn reply_count(&self) -> usize { + self.reply_count.load(Ordering::Relaxed) + } + pub fn reply_bytes(&self) -> usize { + self.reply_bytes.load(Ordering::Relaxed) + } + pub fn max_reply_count(&self) -> usize { + self.max_reply_count.load(Ordering::Relaxed) + } + pub fn max_reply_bytes(&self) -> usize { + self.max_reply_bytes.load(Ordering::Relaxed) + } +} + +fn observe_max(target: &AtomicUsize, value: usize) { + target.fetch_max(value, Ordering::Relaxed); +} + +#[derive(Clone)] +struct AdmissionController { + global: Arc, + long_polls: Arc, + mutations: Arc, + queries: Arc, + stats: Arc, +} + +impl AdmissionController { + fn new(config: &TransportConfig, stats: Arc) -> Self { + Self { + global: Arc::new(Semaphore::new(config.max_active_requests)), + long_polls: Arc::new(Semaphore::new(config.max_active_long_polls)), + mutations: Arc::new(Semaphore::new(config.max_active_mutations)), + queries: Arc::new(Semaphore::new(config.max_active_queries)), + stats, + } + } + + fn try_admit(&self, class: OperationClass) -> Option { + let global = self.global.clone().try_acquire_owned().ok()?; + let class_permit = match class { + OperationClass::LeasePoll => self.long_polls.clone().try_acquire_owned(), + OperationClass::Mutation => self.mutations.clone().try_acquire_owned(), + OperationClass::Query => self.queries.clone().try_acquire_owned(), + } + .ok()?; + + self.stats.admitted.fetch_add(1, Ordering::Relaxed); + let active = self.stats.active.fetch_add(1, Ordering::Relaxed) + 1; + observe_max(&self.stats.max_active, active); + if matches!(class, OperationClass::LeasePoll) { + let active = self.stats.active_long_polls.fetch_add(1, Ordering::Relaxed) + 1; + observe_max(&self.stats.max_active_long_polls, active); + } + + Some(AdmissionGuard { + _global: global, + _class: class_permit, + class, + stats: self.stats.clone(), + }) + } +} + +struct AdmissionGuard { + _global: OwnedSemaphorePermit, + _class: OwnedSemaphorePermit, + class: OperationClass, + stats: Arc, +} + +impl Drop for AdmissionGuard { + fn drop(&mut self) { + self.stats.active.fetch_sub(1, Ordering::Relaxed); + if matches!(self.class, OperationClass::LeasePoll) { + self.stats.active_long_polls.fetch_sub(1, Ordering::Relaxed); + } + } +} + +#[derive(Clone)] +struct ReplyBudget { + slots: Arc, + bytes: Arc, +} + +struct ReplyBudgetGuard { + _slot: OwnedSemaphorePermit, + _bytes: OwnedSemaphorePermit, + bytes: usize, + stats: Arc, +} + +impl Drop for ReplyBudgetGuard { + fn drop(&mut self) { + self.stats.reply_count.fetch_sub(1, Ordering::Relaxed); + self.stats + .reply_bytes + .fetch_sub(self.bytes, Ordering::Relaxed); + } +} + +struct ReplyEnvelope { + identity: Vec, + payload: Vec, + _admission: Option, + _budget: ReplyBudgetGuard, +} + +/// Serve indefinitely using the configuration stored on `ServerAPI`. +pub async fn serve(api: Arc, endpoint: &str) -> Result<(), Error> { + let config = api.cfg.config_toml.transport.clone(); + let stats = Arc::new(TransportStats::default()); + let (_shutdown_tx, shutdown_rx) = watch::channel(false); + serve_with_shutdown(api, endpoint, config, shutdown_rx, stats).await +} + +/// Serve one transport shard until `shutdown` becomes true. +pub async fn serve_with_shutdown( + api: Arc, + endpoint: &str, + config: TransportConfig, + mut shutdown: watch::Receiver, + stats: Arc, +) -> Result<(), Error> { + config.validate().map_err(Error::invalid_argument)?; + + let context = Context::new(); + let mut socket = router(&context) + .set_maxmsgsize(config.max_wire_bytes as i64) + .set_sndhwm(config.zmq_sndhwm) + .set_rcvhwm(config.zmq_rcvhwm) + .set_linger(config.zmq_linger_ms) + .bind(endpoint) + .map_err(|err| Error::io_error(format!("Failed to bind ROUTER on {endpoint}: {err}")))?; + + let admission = AdmissionController::new(&config, stats.clone()); + let reply_budget = ReplyBudget { + slots: Arc::new(Semaphore::new(config.reply_queue_count)), + bytes: Arc::new(Semaphore::new(config.reply_queue_bytes)), + }; + let (reply_tx, mut reply_rx) = mpsc::channel::(config.reply_queue_count); + let mut handlers = JoinSet::new(); + + loop { + tokio::select! { + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + break; + } + } + Some(joined) = handlers.join_next(), if !handlers.is_empty() => { + if let Err(err) = joined { + error!("transport handler failed: {err}"); + } + } + Some(reply) = reply_rx.recv() => { + send_reply(&mut socket, reply, &stats).await; + } + incoming = socket.next() => { + let Some(incoming) = incoming else { break; }; + let multipart = match incoming { + Ok(multipart) => multipart, + Err(err) => { + error!("ROUTER receive failed: {err}"); + continue; + } + }; + admit_message( + multipart, + api.clone(), + &config, + &admission, + &reply_tx, + &reply_budget, + &stats, + &mut handlers, + ); + } + } + } + + // No new work is admitted. Existing handlers may still enqueue through + // their sender clones; drain them until the configured deadline. + drop(reply_tx); + let deadline = Instant::now() + Duration::from_millis(config.shutdown_timeout_ms); + while !handlers.is_empty() || !reply_rx.is_empty() { + tokio::select! { + Some(joined) = handlers.join_next(), if !handlers.is_empty() => { + if let Err(err) = joined { + error!("transport handler failed while draining: {err}"); + } + } + Some(reply) = reply_rx.recv() => { + let ReplyEnvelope { + identity, + payload, + _admission, + _budget, + } = reply; + let multipart: Multipart = vec![identity, payload].into(); + let send = socket.send(multipart); + tokio::pin!(send); + tokio::select! { + result = &mut send => { + if let Err(err) = result { + stats.send_failures.fetch_add(1, Ordering::Relaxed); + error!("ROUTER send failed while draining: {err}"); + } + } + _ = tokio::time::sleep_until(deadline) => { + // Cancellation is safe only because the socket is dropped + // immediately and will never flush tmq's retained multipart. + stats.shutdown_drops.fetch_add(1, Ordering::Relaxed); + handlers.abort_all(); + return Ok(()); + } + } + } + _ = tokio::time::sleep_until(deadline) => { + break; + } + } + } + + if !handlers.is_empty() { + handlers.abort_all(); + } + while let Ok(reply) = reply_rx.try_recv() { + stats.shutdown_drops.fetch_add(1, Ordering::Relaxed); + drop(reply); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn admit_message( + mut multipart: Multipart, + api: Arc, + config: &TransportConfig, + admission: &AdmissionController, + reply_tx: &mpsc::Sender, + reply_budget: &ReplyBudget, + stats: &Arc, + handlers: &mut JoinSet<()>, +) { + if multipart.len() != 2 { + warn!("dropping malformed message with {} frames", multipart.len()); + return; + } + let identity = multipart.pop_front().expect("two frames").to_vec(); + let payload = multipart.pop_front().expect("two frames").to_vec(); + if payload.len() > config.max_wire_bytes { + try_enqueue_immediate( + identity, + Response::Err(Error::invalid_argument("request exceeds wire size limit")), + reply_tx, + reply_budget, + stats, + config, + ); + return; + } + + let envelope: Envelope = match decode(&payload) { + Ok(envelope) => envelope, + Err(err) => { + try_enqueue_immediate( + identity, + Response::Err(err), + reply_tx, + reply_budget, + stats, + config, + ); + return; + } + }; + if let Err(err) = validate_request(&envelope.req, config) { + try_enqueue_immediate( + identity, + Response::Err(err), + reply_tx, + reply_budget, + stats, + config, + ); + return; + } + + let class = OperationClass::of(&envelope.req); + let Some(guard) = admission.try_admit(class) else { + stats.overloaded.fetch_add(1, Ordering::Relaxed); + try_enqueue_immediate( + identity, + Response::Overloaded { retry_after_ms: 50 }, + reply_tx, + reply_budget, + stats, + config, + ); + return; + }; + + let reply_tx = reply_tx.clone(); + let reply_budget = reply_budget.clone(); + let stats = stats.clone(); + let limits = config.clone(); + handlers.spawn(async move { + let response = crate::server::dispatch(api, envelope, limits.clone()).await; + let payload = encode_limited(response, limits.max_reply_bytes); + let Ok(payload) = payload else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + enqueue_reply( + identity, + payload, + Some(guard), + &reply_tx, + reply_budget, + stats, + ) + .await; + }); +} + +fn validate_request(request: &Request, config: &TransportConfig) -> Result<(), Error> { + match request { + Request::SubmitBatch { tasks, .. } => { + validate_batch(tasks.iter().map(Vec::len), tasks.len(), config) + } + Request::CompleteBatch { results, .. } => validate_batch( + results + .iter() + .map(|(task_id, bytes)| task_id.len().saturating_add(bytes.len())), + results.len(), + config, + ), + Request::LeaseBatch { max, .. } if *max as usize > config.max_batch_items => Err( + Error::invalid_argument("lease batch exceeds configured item limit"), + ), + Request::Query(Query::Scan { limit, .. }) if *limit > config.max_batch_items => Err( + Error::invalid_argument("query page exceeds configured item limit"), + ), + _ => Ok(()), + } +} + +fn validate_batch( + sizes: impl Iterator, + count: usize, + config: &TransportConfig, +) -> Result<(), Error> { + if count > config.max_batch_items { + return Err(Error::invalid_argument( + "batch exceeds configured item limit", + )); + } + let mut bytes = 0usize; + for size in sizes { + bytes = bytes + .checked_add(size) + .ok_or_else(|| Error::invalid_argument("batch byte size overflow"))?; + if bytes > config.max_batch_bytes { + return Err(Error::invalid_argument( + "batch exceeds configured byte limit", + )); + } + } + Ok(()) +} + +fn encode_limited(response: Response, max_reply_bytes: usize) -> Result, Error> { + let payload = encode(&response)?; + if payload.len() <= max_reply_bytes { + return Ok(payload); + } + let fallback = Response::Err(Error::invalid_argument( + "response exceeds configured byte limit", + )); + let payload = encode(&fallback)?; + if payload.len() <= max_reply_bytes { + Ok(payload) + } else { + Err(Error::invalid_argument("reply limit is too small")) + } +} + +fn try_enqueue_immediate( + identity: Vec, + response: Response, + reply_tx: &mpsc::Sender, + reply_budget: &ReplyBudget, + stats: &Arc, + config: &TransportConfig, +) { + let Ok(payload) = encode_limited(response, config.max_reply_bytes) else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + let Ok(slot) = reply_budget.slots.clone().try_acquire_owned() else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + let Ok(bytes) = reply_budget + .bytes + .clone() + .try_acquire_many_owned(payload.len() as u32) + else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + let budget = track_reply_budget(slot, bytes, payload.len(), stats.clone()); + let reply = ReplyEnvelope { + identity, + payload, + _admission: None, + _budget: budget, + }; + if reply_tx.try_send(reply).is_err() { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + } +} + +async fn enqueue_reply( + identity: Vec, + payload: Vec, + admission: Option, + reply_tx: &mpsc::Sender, + reply_budget: ReplyBudget, + stats: Arc, +) { + let byte_count = payload.len(); + let Ok(slot) = reply_budget.slots.acquire_owned().await else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + let Ok(bytes) = reply_budget + .bytes + .acquire_many_owned(byte_count as u32) + .await + else { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + return; + }; + let budget = track_reply_budget(slot, bytes, byte_count, stats.clone()); + let reply = ReplyEnvelope { + identity, + payload, + _admission: admission, + _budget: budget, + }; + if reply_tx.send(reply).await.is_err() { + stats.dropped_replies.fetch_add(1, Ordering::Relaxed); + } +} + +fn track_reply_budget( + slot: OwnedSemaphorePermit, + byte_permit: OwnedSemaphorePermit, + bytes: usize, + stats: Arc, +) -> ReplyBudgetGuard { + let count = stats.reply_count.fetch_add(1, Ordering::Relaxed) + 1; + let total_bytes = stats.reply_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes; + observe_max(&stats.max_reply_count, count); + observe_max(&stats.max_reply_bytes, total_bytes); + ReplyBudgetGuard { + _slot: slot, + _bytes: byte_permit, + bytes, + stats, + } +} + +async fn send_reply( + socket: &mut tmq::router::Router, + reply: ReplyEnvelope, + stats: &Arc, +) { + let ReplyEnvelope { + identity, + payload, + _admission, + _budget, + } = reply; + let multipart: Multipart = vec![identity, payload].into(); + if let Err(err) = socket.send(multipart).await { + stats.send_failures.fetch_add(1, Ordering::Relaxed); + error!("ROUTER send failed: {err}"); + } +} diff --git a/engine/tests/e2e.rs b/engine/tests/e2e.rs new file mode 100644 index 0000000..cc366af --- /dev/null +++ b/engine/tests/e2e.rs @@ -0,0 +1,122 @@ +//! End-to-end test of the tcp ZeroMQ channel: submit → lease → complete over the +//! wire, plus a query to confirm the persisted record transitions. + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use engine::client::Client; +use enginelib::Identifier; +use enginelib::api::{ServerAPI, finished_task_key, task_key}; +use enginelib::nucleus::query::{Query, QueryResult}; +use enginelib::task::{Task, Verifiable}; + +/// Minimal registered task type so `submit`/`complete` pass verification. +#[derive(Debug, Clone)] +struct EchoTask; + +impl Verifiable for EchoTask { + fn verify(&self, _bytes: &[u8]) -> bool { + true + } +} + +impl Task for EchoTask { + fn get_id(&self) -> Identifier { + ("test".to_string(), "echo".to_string()) + } + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + fn to_bytes(&self) -> Vec { + Vec::new() + } + fn from_bytes(&self, _bytes: &[u8]) -> Box { + Box::new(self.clone()) + } + fn from_toml(&self, _d: String) -> Box { + Box::new(self.clone()) + } + fn to_toml(&self) -> String { + String::new() + } +} + +fn unique_tmp() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + std::env::temp_dir() + .join(format!("engine_e2e_{nanos}")) + .to_string_lossy() + .into_owned() +} + +#[tokio::test] +async fn submit_lease_complete_over_tcp() { + let task_type: Identifier = ("test".to_string(), "echo".to_string()); + let endpoint = "tcp://127.0.0.1:55571"; + + // Build an isolated engine with the echo task registered. + let api = ServerAPI::with_path(&unique_tmp()); + api.task_registry + .tasks + .insert(task_type.clone(), Arc::new(EchoTask)); + let api = Arc::new(api); + ServerAPI::populate(&api); + ServerAPI::spawn_loaders(&api); + + // Serve in the background. + let serve_api = api.clone(); + tokio::spawn(async move { + let _ = engine::server::serve(serve_api, endpoint).await; + }); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + let mut client = Client::connect(endpoint, String::new()).expect("connect"); + + // Submit. + let payload = b"hello".to_vec(); + let ids = client + .submit(task_type.clone(), vec![payload.clone()]) + .await + .expect("submit"); + assert_eq!(ids.len(), 1); + let task_id = ids[0].clone(); + + // Lease it back. + let leased = client + .lease(task_type.clone(), "worker".to_string(), 4) + .await + .expect("lease"); + assert_eq!(leased.len(), 1); + assert_eq!(leased[0].task_id, task_id); + assert_eq!(leased[0].bytes, payload); + + // Complete it. + let ok = client + .complete(task_type.clone(), vec![(task_id.clone(), b"done".to_vec())]) + .await + .expect("complete"); + assert_eq!(ok, 1); + + // Pending record gone, finished record present. + let pending = client + .query(Query::Get { + key: task_key(&task_type, &task_id), + }) + .await + .expect("query pending"); + assert!(matches!(pending, QueryResult::Record(None))); + + let finished = client + .query(Query::Get { + key: finished_task_key(&task_type, &task_id), + }) + .await + .expect("query finished"); + match finished { + QueryResult::Record(Some(rec)) => assert_eq!(rec.value, b"done".to_vec()), + other => panic!("expected finished record, got {other:?}"), + } +} diff --git a/engine/tests/recovery.rs b/engine/tests/recovery.rs new file mode 100644 index 0000000..065c022 --- /dev/null +++ b/engine/tests/recovery.rs @@ -0,0 +1,259 @@ +//! In-process tests of the lease-recovery paths: reaping requeues expired +//! leases, the loader's rescan flag un-strands tasks dropped back to +//! "persisted only", complete_batch counts duplicate ids once, and the +//! registry refuses identifiers that would collide key prefixes. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use chrono::Utc; +use dashmap::DashSet; +use enginelib::Identifier; +use enginelib::Registry; +use enginelib::api::{LeasedTaskQueue, ServerAPI}; +use enginelib::nucleus::{complete, lease, submit}; +use enginelib::task::{Task, Verifiable}; + +#[derive(Debug, Clone)] +struct EchoTask; + +impl Verifiable for EchoTask { + fn verify(&self, _bytes: &[u8]) -> bool { + true + } +} + +impl Task for EchoTask { + fn get_id(&self) -> Identifier { + ("test".to_string(), "echo".to_string()) + } + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + fn to_bytes(&self) -> Vec { + Vec::new() + } + fn from_bytes(&self, _bytes: &[u8]) -> Box { + Box::new(self.clone()) + } + fn from_toml(&self, _d: String) -> Box { + Box::new(self.clone()) + } + fn to_toml(&self) -> String { + String::new() + } +} + +fn unique_tmp(tag: &str) -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + std::env::temp_dir() + .join(format!("engine_recovery_{tag}_{nanos}")) + .to_string_lossy() + .into_owned() +} + +fn test_api(tag: &str, task_type: &Identifier) -> Arc { + let api = ServerAPI::with_path(&unique_tmp(tag)); + api.task_registry + .tasks + .insert(task_type.clone(), Arc::new(EchoTask)); + let api = Arc::new(api); + ServerAPI::populate(&api); + api +} + +/// Backdate every lease for `task_type` far enough that the reaper sees it as +/// expired. +fn expire_leases(api: &Arc, task_type: &Identifier) { + let mut leases = api.leased_tasks.tasks.get_mut(task_type).expect("leases"); + for lease in leases.iter_mut() { + lease.given_at = Utc::now() - chrono::Duration::seconds(7200); + } +} + +/// An expired lease is reaped and its task handed straight back to the queue, +/// so the next lease call gets it again. +#[tokio::test] +async fn reaped_lease_is_requeued() { + let task_type: Identifier = ("test".to_string(), "echo".to_string()); + let api = test_api("reap", &task_type); + + submit::submit(api.clone(), b"payload", task_type.clone()).expect("submit"); + ServerAPI::load_all(&api, task_type.clone()) + .await + .expect("load"); + + let leased = lease::lease_batch(api.clone(), task_type.clone(), "w1".into(), 1) + .await + .expect("lease"); + assert_eq!(leased.len(), 1); + let task_id = leased[0].task_id.clone(); + + expire_leases(&api, &task_type); + LeasedTaskQueue::reap_expired(&api); + + // The lease is gone and the task is back in the channel for the next worker. + assert!( + api.leased_tasks + .tasks + .get(&task_type) + .is_none_or(|l| l.is_empty()) + ); + let released = lease::lease_batch(api.clone(), task_type.clone(), "w2".into(), 1) + .await + .expect("re-lease"); + assert_eq!(released.len(), 1); + assert_eq!(released[0].task_id, task_id); +} + +/// When the reaper can't requeue (channel full), the task drops back to +/// "persisted only" and the rescan flag makes the running loader rewind its +/// cursor and re-enqueue it — the stranded-task case. +#[tokio::test] +async fn full_channel_reap_recovers_via_loader_rescan() { + let task_type: Identifier = ("test".to_string(), "echo".to_string()); + let api = ServerAPI::with_path(&unique_tmp("rescan")); + api.task_registry + .tasks + .insert(task_type.clone(), Arc::new(EchoTask)); + let api = Arc::new(api); + // Capacity-1 channel (inserted before populate, whose or_insert keeps it) + // so a single queued task makes the reaper's try_send fail. + let (tx, rx) = async_channel::bounded(1); + api.task_queue.tasks.insert( + task_type.clone(), + (tx, rx, DashSet::default(), Arc::new(AtomicBool::new(false))), + ); + ServerAPI::populate(&api); + tokio::spawn(ServerAPI::run_loader(api.clone(), task_type.clone())); + + let ids = submit::submit_batch( + api.clone(), + task_type.clone(), + vec![b"a".to_vec(), b"b".to_vec()], + ) + .expect("submit"); + + // Lease the first task; the loader then fills the channel with the second. + let first = tokio::time::timeout( + Duration::from_secs(5), + lease::lease_batch(api.clone(), task_type.clone(), "w1".into(), 1), + ) + .await + .expect("lease timed out") + .expect("lease"); + assert_eq!(first.len(), 1); + + // Wait until the channel is full so the reaper's try_send must fail. + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while api + .task_queue + .tasks + .get(&task_type) + .expect("queue") + .0 + .is_empty() + { + assert!( + tokio::time::Instant::now() < deadline, + "loader never enqueued" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + + expire_leases(&api, &task_type); + LeasedTaskQueue::reap_expired(&api); + + // Both tasks must come back out: the queued one, and the reaped one that + // only the rescan can recover. + let mut seen = Vec::new(); + while seen.len() < 2 { + let batch = tokio::time::timeout( + Duration::from_secs(5), + lease::lease_batch(api.clone(), task_type.clone(), "w2".into(), 2), + ) + .await + .expect("recovery lease timed out") + .expect("lease"); + seen.extend(batch.iter().map(|t| t.task_id.clone())); + } + seen.sort(); + let mut expected = ids.clone(); + expected.sort(); + assert_eq!(seen, expected); +} + +/// A closed loader channel must not leave records marked in-flight when none +/// were enqueued. The rescan flag allows recovery after the queue is recreated. +#[tokio::test] +async fn load_after_closed_channel_clears_dedup_and_requests_rescan() { + let task_type: Identifier = ("test".to_string(), "echo".to_string()); + let api = test_api("closed_loader", &task_type); + let ids = submit::submit_batch( + api.clone(), + task_type.clone(), + vec![b"a".to_vec(), b"b".to_vec()], + ) + .expect("submit"); + + let queue = api.task_queue.tasks.get(&task_type).expect("queue"); + queue.1.close(); + drop(queue); + + let err = ServerAPI::load_after(&api, &task_type, None) + .await + .expect_err("closed queue must reject loader sends"); + assert!(err.as_ref().contains("Failed to send stored task")); + + let queue = api.task_queue.tasks.get(&task_type).expect("queue"); + for id in ids { + assert!(!queue.2.contains(&id), "task remained in dedup set"); + } + assert!(queue.3.load(Ordering::Acquire), "rescan was not requested"); +} + +/// A duplicate task id inside one complete_batch request is completed and +/// counted once, not twice. +#[tokio::test] +async fn complete_batch_counts_duplicate_ids_once() { + let task_type: Identifier = ("test".to_string(), "echo".to_string()); + let api = test_api("dup", &task_type); + + submit::submit(api.clone(), b"payload", task_type.clone()).expect("submit"); + ServerAPI::load_all(&api, task_type.clone()) + .await + .expect("load"); + let leased = lease::lease_batch(api.clone(), task_type.clone(), "w".into(), 1) + .await + .expect("lease"); + let task_id = leased[0].task_id.clone(); + + let ok = complete::complete_batch( + api.clone(), + task_type.clone(), + vec![ + (task_id.clone(), b"done".to_vec()), + (task_id.clone(), b"done again".to_vec()), + ], + ) + .expect("complete"); + assert_eq!(ok, 1); +} + +/// Identifiers containing ':' would collide RocksDB key prefixes across task +/// types; the registry must refuse them. +#[test] +fn registry_refuses_colliding_identifiers() { + let mut registry = enginelib::api::EngineTaskRegistry::default(); + let bad: Identifier = ("ns:x".to_string(), "name".to_string()); + registry.register(Arc::new(EchoTask), bad.clone()); + assert!(registry.get(&bad).is_none()); + + let good: Identifier = ("ns".to_string(), "name".to_string()); + registry.register(Arc::new(EchoTask), good.clone()); + assert!(registry.get(&good).is_some()); +} diff --git a/engine/tests/transport_bounds.rs b/engine/tests/transport_bounds.rs new file mode 100644 index 0000000..a504e29 --- /dev/null +++ b/engine/tests/transport_bounds.rs @@ -0,0 +1,212 @@ +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use engine::client::Client; +use engine::transport::{TransportStats, serve_with_shutdown}; +use enginelib::Identifier; +use enginelib::api::ServerAPI; +use enginelib::config::TransportConfig; +use enginelib::error::ErrorKind; +use enginelib::task::{Task, Verifiable}; +use tokio::sync::watch; + +#[derive(Debug, Clone)] +struct EchoTask; + +impl Verifiable for EchoTask { + fn verify(&self, _bytes: &[u8]) -> bool { + true + } +} + +impl Task for EchoTask { + fn get_id(&self) -> Identifier { + ("test".into(), "echo".into()) + } + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + fn to_bytes(&self) -> Vec { + Vec::new() + } + fn from_bytes(&self, _bytes: &[u8]) -> Box { + Box::new(self.clone()) + } + fn from_toml(&self, _data: String) -> Box { + Box::new(self.clone()) + } + fn to_toml(&self) -> String { + String::new() + } +} + +fn unique_tmp() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + std::env::temp_dir() + .join(format!("engine_transport_bounds_{nanos}")) + .to_string_lossy() + .into_owned() +} + +fn test_config() -> TransportConfig { + TransportConfig { + max_active_requests: 2, + max_active_long_polls: 1, + max_active_mutations: 2, + max_active_queries: 1, + reply_queue_count: 4, + reply_queue_bytes: 64 * 1024, + max_wire_bytes: 32 * 1024, + max_reply_bytes: 16 * 1024, + max_batch_items: 2, + max_batch_bytes: 8 * 1024, + lease_long_poll_ms: 500, + zmq_sndhwm: 8, + zmq_rcvhwm: 8, + zmq_linger_ms: 0, + shutdown_timeout_ms: 1000, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn long_poll_bulkhead_and_request_limits_are_enforced() { + let endpoint = "tcp://127.0.0.1:55572"; + let task_type: Identifier = ("test".into(), "echo".into()); + let api = ServerAPI::with_path(&unique_tmp()); + api.task_registry + .tasks + .insert(task_type.clone(), Arc::new(EchoTask)); + let api = Arc::new(api); + ServerAPI::populate(&api); + ServerAPI::spawn_loaders(&api); + + let config = test_config(); + let stats = Arc::new(TransportStats::default()); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let server = tokio::spawn(serve_with_shutdown( + api, + endpoint, + config.clone(), + shutdown_rx, + stats.clone(), + )); + tokio::time::sleep(Duration::from_millis(100)).await; + + let mut waiting_client = Client::connect(endpoint, String::new()).expect("connect waiter"); + let waiting_type = task_type.clone(); + let waiting = tokio::spawn(async move { + waiting_client + .lease(waiting_type, "worker-1".into(), 1) + .await + }); + tokio::time::sleep(Duration::from_millis(50)).await; + + // The dedicated lease-poll pool is full, but mutation capacity remains. + let mut rejected_client = Client::connect(endpoint, String::new()).expect("connect rejected"); + let err = tokio::time::timeout( + Duration::from_secs(1), + rejected_client.lease(task_type.clone(), "worker-2".into(), 1), + ) + .await + .expect("overload response timed out") + .expect_err("second long poll should be rejected"); + assert_eq!(err.kind(), ErrorKind::Overloaded); + + let mut submitter = Client::connect(endpoint, String::new()).expect("connect submitter"); + let ids = submitter + .submit(task_type.clone(), vec![b"payload".to_vec()]) + .await + .expect("mutation should not be starved by long poll"); + assert_eq!(ids.len(), 1); + + let leased = tokio::time::timeout(Duration::from_secs(2), waiting) + .await + .expect("waiting lease timed out") + .expect("waiting task panicked") + .expect("waiting lease failed"); + assert_eq!(leased.len(), 1); + assert_eq!(leased[0].task_id, ids[0]); + + // Batch validation happens before authentication or nucleus mutation. + let err = submitter + .submit(task_type, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]) + .await + .expect_err("oversized batch should be rejected"); + assert_eq!(err.kind(), ErrorKind::InvalidArgument); + + shutdown_tx.send(true).expect("request shutdown"); + tokio::time::timeout(Duration::from_secs(2), server) + .await + .expect("server did not shut down") + .expect("server task panicked") + .expect("server returned error"); + + assert!(stats.overloaded() >= 1); + assert!(stats.max_active() <= config.max_active_requests); + assert!(stats.max_active_long_polls() <= config.max_active_long_polls); + assert!(stats.max_reply_count() <= config.reply_queue_count); + assert!(stats.max_reply_bytes() <= config.reply_queue_bytes); + assert_eq!(stats.reply_count(), 0); + assert_eq!(stats.reply_bytes(), 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn shutdown_aborts_long_polls_at_the_drain_deadline_without_leaking_permits() { + let endpoint = "tcp://127.0.0.1:55573"; + let task_type: Identifier = ("test".into(), "echo".into()); + let api = ServerAPI::with_path(&unique_tmp()); + api.task_registry + .tasks + .insert(task_type.clone(), Arc::new(EchoTask)); + let api = Arc::new(api); + ServerAPI::populate(&api); + ServerAPI::spawn_loaders(&api); + + let mut config = test_config(); + config.lease_long_poll_ms = 5000; + config.shutdown_timeout_ms = 100; + let stats = Arc::new(TransportStats::default()); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let server = tokio::spawn(serve_with_shutdown( + api, + endpoint, + config, + shutdown_rx, + stats.clone(), + )); + tokio::time::sleep(Duration::from_millis(100)).await; + + let mut client = Client::connect(endpoint, String::new()).expect("connect waiter"); + let waiter = tokio::spawn(async move { client.lease(task_type, "worker".into(), 1).await }); + + let wait_deadline = tokio::time::Instant::now() + Duration::from_secs(1); + while stats.max_active_long_polls() == 0 { + assert!( + tokio::time::Instant::now() < wait_deadline, + "long poll was not admitted" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + + shutdown_tx.send(true).expect("request shutdown"); + tokio::time::timeout(Duration::from_secs(1), server) + .await + .expect("server exceeded shutdown deadline") + .expect("server task panicked") + .expect("server returned error"); + waiter.abort(); + + let release_deadline = tokio::time::Instant::now() + Duration::from_secs(1); + while stats.active() != 0 { + assert!( + tokio::time::Instant::now() < release_deadline, + "admission permit leaked after shutdown" + ); + tokio::task::yield_now().await; + } + assert_eq!(stats.reply_count(), 0); + assert_eq!(stats.reply_bytes(), 0); +} diff --git a/enginelib/macros/src/lib.rs b/enginelib/macros/src/lib.rs index a9f1515..c3c07a6 100644 --- a/enginelib/macros/src/lib.rs +++ b/enginelib/macros/src/lib.rs @@ -64,8 +64,8 @@ pub fn derive_verifiable(input: TokenStream) -> TokenStream { let expanded = quote! { impl ::enginelib::task::Verifiable for #name { - fn verify(&self, b: Vec) -> bool { - let k: Result<#name, _> = enginelib::api::postcard::from_bytes(b.as_slice()); + fn verify(&self, b: &[u8]) -> bool { + let k: Result<#name, _> = enginelib::api::postcard::from_bytes(b); k.is_ok() } } diff --git a/enginelib/src/api.rs b/enginelib/src/api.rs index 0220a5c..8be25bb 100644 --- a/enginelib/src/api.rs +++ b/enginelib/src/api.rs @@ -6,9 +6,11 @@ use dashmap::{DashMap, DashSet}; pub use postcard; pub use postcard::from_bytes; pub use postcard::to_allocvec; +use rust_rocksdb::{Direction, IteratorMode}; use serde::{Deserialize, Serialize}; use std::sync::Arc; -use tracing::{Level, debug, instrument}; +use std::sync::atomic::{AtomicBool, Ordering}; +use tracing::{Level, debug, error, instrument}; // t:namespace:task_name: -> Serialized Task Record pub fn task_key_prefix(task_type: &Identifier) -> String { @@ -42,39 +44,46 @@ pub struct ServerAPI { pub task_registry: EngineTaskRegistry, // RW } -impl Default for ServerAPI { - fn default() -> Self { - let path = "./engine_db"; - let mut opts = rust_rocksdb::Options::default(); - opts.create_if_missing(true); - - let db = Arc::new( - rust_rocksdb::DB::open(&opts, path) - .expect("Failed to open RocksDB — check that the path is writable and no other process holds the lock"), - ); - +impl ServerAPI { + /// Initialize the global tracing subscriber once (idempotent). Referenced by + /// the `#[module]` macro's generated `run` prologue as well as `with_path`. + pub fn setup_logger() { use std::sync::OnceLock; - static INIT: OnceLock<()> = OnceLock::new(); INIT.get_or_init(|| { #[cfg(debug_assertions)] let _ = tracing_subscriber::FmtSubscriber::builder() - // all spans/events with a level higher than TRACE (e.g, debug, info, warn, etc.) - // will be written to stdout. .with_max_level(Level::DEBUG) - // builds the subscriber. .try_init(); #[cfg(not(debug_assertions))] let _ = tracing_subscriber::FmtSubscriber::builder() - // all spans/events with a level higher than TRACE (e.g, debug, info, warn, etc.) - // will be written to stdout. .with_max_level(Level::INFO) - // builds the subscriber. .try_init(); }); + } + + /// Open (or create) a `ServerAPI` backed by the RocksDB store at `path`. + /// [`Default`] uses the standard `./engine_db`; tests pass a temp path so + /// they don't collide on the shared store. + pub fn with_path(path: &str) -> Self { + Self::with_path_and_config(path, Config::default()) + } + + /// Open a database with an explicitly loaded configuration. Production + /// startup uses this path so transport and auth settings are not discarded. + pub fn with_path_and_config(path: &str, config: Config) -> Self { + let mut opts = rust_rocksdb::Options::default(); + opts.create_if_missing(true); + + let db = Arc::new( + rust_rocksdb::DB::open(&opts, path) + .expect("Failed to open RocksDB — check that the path is writable and no other process holds the lock"), + ); + + Self::setup_logger(); let mut k = Self { - cfg: Config::default(), + cfg: config, event_bus: EventBus::default(), lib_manager: LibraryManager::default(), db, @@ -84,11 +93,21 @@ impl Default for ServerAPI { }; LibraryManager::load_modules(&mut k); crate::event::register_inventory_handlers(&mut k); - return k; + k + } +} + +impl Default for ServerAPI { + fn default() -> Self { + Self::with_path("./engine_db") } } impl ServerAPI { - pub async fn load(api: &Arc, task_type: Identifier) -> Result<(), Error> { + /// Enqueue up to one batch of persisted-but-unqueued records for `task_type`, + /// returning how many were enqueued. Used for startup recovery: call it in a + /// loop until it returns 0 to drain the whole backlog (see + /// [`ServerAPI::load_all`]). + pub async fn load(api: &Arc, task_type: Identifier) -> Result { const LOAD_BATCH_SIZE: usize = 4096; let prefix = task_key_prefix(&task_type); @@ -133,46 +152,224 @@ impl ServerAPI { (k.0.clone(), batch) }; + let count = batch.len(); let mut batch = batch.into_iter(); while let Some(task) = batch.next() { if let Err(err) = sender.send(task).await { - // Un-mark everything that never made it into the queue so a - // later load() can retry those tasks. + // Un-mark everything that never made it into the queue, and set + // the rescan flag so a later load() or the running loader can + // retry those tasks. if let Some(k) = api.task_queue.tasks.get(&task_type) { k.2.remove(&err.0.task_id); for task in batch { k.2.remove(&task.task_id); } + k.3.store(true, Ordering::Release); } return Err(Error::new(format!("Failed to send stored task: {err}"))); } } - Ok(()) + Ok(count) + } + + /// Drain the entire persisted backlog for `task_type` into the queue by + /// looping [`load`] until it enqueues nothing. Enqueue is backpressured, so + /// this streams the backlog in bounded memory. Returns the total enqueued. + pub async fn load_all(api: &Arc, task_type: Identifier) -> Result { + let mut total = 0; + loop { + let n = Self::load(api, task_type.clone()).await?; + total += n; + if n == 0 { + break; + } + } + Ok(total) + } + + /// Enqueue one batch of records for `task_type` whose key sorts strictly after + /// `cursor` (or from the start of the prefix when `cursor` is `None`), + /// returning `(count, new_cursor)`. + /// + /// This exploits the fact that `druid` ids are timestamp-ordered (big-endian + /// nanosecond prefix), so the `t:` keys are lexicographically ordered by + /// submission time. A newer submit always sorts *after* an older one, so the + /// loader resumes from a cursor with a direct RocksDB seek — no re-scan of the + /// prefix, no walking delete-tombstones (completed records sit behind the + /// cursor), and each record is read exactly once (no duplicate enqueue). + pub async fn load_after( + api: &Arc, + task_type: &Identifier, + cursor: Option, + ) -> Result<(usize, Option), Error> { + const LOAD_BATCH_SIZE: usize = 4096; + let prefix = task_key_prefix(task_type); + let seek = cursor.clone().unwrap_or_else(|| prefix.clone()); + + let (sender, batch, last) = { + let k = api + .task_queue + .tasks + .get(task_type) + .ok_or(Error::not_found("TaskTypeNotFound"))?; + + let mut batch = Vec::new(); + let mut last = cursor.clone(); + let iter = api + .db + .iterator(IteratorMode::From(seek.as_bytes(), Direction::Forward)); + for result in iter { + if batch.len() >= LOAD_BATCH_SIZE { + break; + } + let (key, value) = match result { + Ok(kv) => kv, + Err(err) => { + eprintln!("RocksDB read error: {err}"); + continue; + } + }; + if !key.starts_with(prefix.as_bytes()) { + break; + } + let Ok(key) = String::from_utf8(key.to_vec()) else { + continue; + }; + // `From` is inclusive; never re-emit the cursor key itself. + if cursor.as_deref() == Some(key.as_str()) { + continue; + } + if let Some(task_id) = task_id_from_key(&key) + && k.2.insert(task_id.to_string()) + { + batch.push(StoredTask { + bytes: value.into(), + task_id: task_id.to_string(), + task_type: task_type.clone(), + }); + } + // Advance the cursor past every key seen, so skipped keys aren't + // revisited next call. + last = Some(key); + } + (k.0.clone(), batch, last) + }; + + let count = batch.len(); + let mut batch = batch.into_iter(); + while let Some(task) = batch.next() { + if let Err(err) = sender.send(task).await { + // None of these tasks reached the channel. Clear their in-flight + // markers and rewind the loader so the persisted records can be + // recovered if the queue is recreated. + if let Some(k) = api.task_queue.tasks.get(task_type) { + k.2.remove(&err.0.task_id); + for task in batch { + k.2.remove(&task.task_id); + } + k.3.store(true, Ordering::Release); + } + return Err(Error::new(format!("Failed to send stored task: {err}"))); + } + } + Ok((count, last)) + } + + /// Run the loader for `task_type` forever: continuously stream + /// persisted-but-unqueued records into the bounded channel via [`load_after`]. + /// + /// This is the queue's engine. `submit` only writes to the DB; this task moves + /// those records into the lease channel — with backpressure, since + /// `send().await` blocks while the channel is full. It carries a forward cursor + /// (works because ids are time-ordered), so in steady state it never re-scans; + /// startup/restart recovery is the same path with `cursor = None`. + /// + /// The one exception: recovery paths (lease reaping, cancel with a full + /// channel, complete's DB-failure paths) drop a task back to "persisted only" + /// by clearing its dedup entry — and that record now sits *behind* the cursor. + /// Those paths set the per-type rescan flag after clearing the entry; when the + /// flag is set the loader rewinds its cursor and re-walks the prefix. The + /// dedup set makes the re-scan duplicate-free, so the only cost is the walk + /// itself, paid only when a recovery actually happened. When the DB is + /// momentarily caught up it sleeps briefly to avoid busy-spinning. + pub async fn run_loader(api: Arc, task_type: Identifier) { + use std::time::Duration; + let mut cursor: Option = None; + loop { + if let Some(queue) = api.task_queue.tasks.get(&task_type) { + // Recovery signal: a cleared dedup entry is behind the cursor. + if queue.3.swap(false, Ordering::AcqRel) { + cursor = None; + } + } + match Self::load_after(&api, &task_type, cursor.clone()).await { + // Keep the cursor even when nothing was enqueued: it advanced + // past dedup-skipped keys, and dropping it would make the next + // iteration re-walk them (a busy re-scan after a rewind). + Ok((0, next)) => { + cursor = next; + tokio::time::sleep(Duration::from_millis(1)).await + } + Ok((_, next)) => cursor = next, // advanced; loop immediately + Err(err) => { + eprintln!("loader error for {task_type:?}: {err}"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } + } + + /// Spawn one [`run_loader`] task per registered task type. Call once, inside a + /// tokio runtime, after [`populate`]. + pub fn spawn_loaders(api: &Arc) { + for entry in api.task_registry.tasks.iter() { + tokio::spawn(Self::run_loader(api.clone(), entry.key().clone())); + } } pub fn populate(api: &Arc) { + Self::populate_with(api, 8192); + } + + /// Like [`populate`] but with a configurable per-type lease-channel capacity. + /// Used to sweep queue depth in benchmarks. + pub fn populate_with(api: &Arc, queue_size: usize) { api.task_registry.tasks.iter().for_each(|f| { let key = f.key(); - let (tx, rx) = async_channel::bounded(8192); // Add to config or make unbound ? - api.task_queue - .tasks - .entry(key.clone()) - .or_insert((tx, rx, DashSet::default())); + let (tx, rx) = async_channel::bounded(queue_size.max(1)); + api.task_queue.tasks.entry(key.clone()).or_insert(( + tx, + rx, + DashSet::default(), + Arc::new(AtomicBool::new(false)), + )); api.leased_tasks.tasks.entry(key.clone()).or_default(); // task reg should be populated by mods }); } pub fn init() -> Arc { - let api = Arc::new(Self::default()); + Self::init_with_config(Config::default()) + } + + pub fn init_with_config(config: Config) -> Arc { + let api = Arc::new(Self::with_path_and_config("./engine_db", config)); Self::populate(&api); - let dapi = api.clone(); + Self::spawn_reaper(&api); + api + } + + /// Spawn the background reaper thread: every [`REAP_INTERVAL_SECS`] it reclaims + /// leases whose TTL has elapsed and requeues their tasks. Called by [`init`]; + /// call it directly when building a `ServerAPI` via [`with_path`] and you want + /// the same behavior. + pub fn spawn_reaper(api: &Arc) { + let api = api.clone(); std::thread::spawn(move || { loop { - std::thread::sleep(std::time::Duration::from_secs(3600)); - LeasedTaskQueue::reap_expired(&dapi); + std::thread::sleep(std::time::Duration::from_secs(REAP_INTERVAL_SECS)); + LeasedTaskQueue::reap_expired(&api); } }); - api } } @@ -183,6 +380,15 @@ pub struct EngineTaskRegistry { impl Registry for EngineTaskRegistry { #[instrument(skip_all, fields(namespace = %identifier.0, name = %identifier.1))] fn register(&mut self, task: Arc, identifier: Identifier) { + // ':' is the DB key separator (`t:ns:name:id`), so an identifier + // containing it would collide with other types' key prefixes. Refuse it. + if identifier.0.contains(':') || identifier.1.contains(':') { + error!( + "TaskRegistry: refusing to register {}.{}: identifiers must not contain ':'", + identifier.0, identifier.1 + ); + return; + } // Insert the task into the hashmap with (mod_id, identifier) as the key debug!( "TaskRegistry: Registering task {}.{}", @@ -195,16 +401,21 @@ impl Registry for EngineTaskRegistry { self.tasks.get(identifier).map(|obj| obj.clone_box()) } } +/// Per task type: the bounded lease channel, the in-flight dedup set (ids that +/// are queued or leased — blocks the loader from re-enqueuing them), and the +/// rescan flag. Recovery paths set the flag *after* removing an id from the +/// dedup set to tell [`ServerAPI::run_loader`] its cursor has passed a record +/// that must be re-enqueued, so it rewinds and re-walks the prefix. +pub type TaskQueueEntry = ( + async_channel::Sender, + async_channel::Receiver, + dashmap::DashSet, + Arc, +); + #[derive(Debug, Default, Clone)] pub struct TaskQueue { - pub tasks: DashMap< - Identifier, - ( - async_channel::Sender, - async_channel::Receiver, - dashmap::DashSet, - ), - >, + pub tasks: DashMap, } #[derive(Debug, Default, Clone)] pub struct LeasedTaskQueue { @@ -224,6 +435,9 @@ pub struct LeasedTask { } const LEASE_TTL_SECS: i64 = 3600; +/// How often the reap thread wakes. Much smaller than the TTL so an expired +/// lease is noticed within ~a minute of expiring, not up to a whole extra TTL. +const REAP_INTERVAL_SECS: u64 = 60; impl LeasedTask { fn expired(&self) -> bool { @@ -232,12 +446,20 @@ impl LeasedTask { } impl LeasedTaskQueue { - fn reap_expired(api: &Arc) { - let mut reaped: Vec<(Identifier, String)> = Vec::new(); + /// Drop every expired lease and hand its task back to the queue. + /// + /// The reaped lease still holds the `StoredTask`, so the task is requeued + /// directly with a non-blocking `try_send` (dedup entry kept: the task goes + /// straight from "leased" back to "queued"). Only if the channel is full or + /// closed does it fall back to clearing the dedup entry and setting the + /// rescan flag so [`ServerAPI::run_loader`] re-enqueues the still-persisted + /// record. Public so tests can drive reaping without waiting for the timer. + pub fn reap_expired(api: &Arc) { + let mut reaped: Vec<(Identifier, Arc)> = Vec::new(); api.leased_tasks.tasks.retain(|task_type, tasks| { tasks.retain(|task| { if task.expired() { - reaped.push((task_type.clone(), task.stored_task.task_id.clone())); + reaped.push((task_type.clone(), task.stored_task.clone())); false } else { true @@ -245,11 +467,15 @@ impl LeasedTaskQueue { }); !tasks.is_empty() }); - // Clear reaped ids from the dedup set outside the leased_tasks locks so - // load() can re-enqueue the tasks from their still-persisted records. - for (task_type, task_id) in reaped { - if let Some(queue) = api.task_queue.tasks.get(&task_type) { - queue.2.remove(&task_id); + // Requeue outside the leased_tasks locks. + for (task_type, stored_task) in reaped { + if let Some(queue) = api.task_queue.tasks.get(&task_type) + && queue.0.try_send((*stored_task).clone()).is_err() + { + // Channel full/closed: drop back to "persisted only" and + // signal the loader to rescan for the record. + queue.2.remove(&stored_task.task_id); + queue.3.store(true, Ordering::Release); } } } diff --git a/enginelib/src/config.rs b/enginelib/src/config.rs index b779bfd..a09610a 100644 --- a/enginelib/src/config.rs +++ b/enginelib/src/config.rs @@ -1,4 +1,4 @@ -use std::{fs, io::Error}; +use std::{fs, io::ErrorKind}; use serde::{Deserialize, Serialize}; use tracing::{error, instrument}; @@ -7,42 +7,147 @@ fn default_host() -> String { "[::1]:50051".into() } +/// Resource limits for one transport shard. +/// +/// Defaults are deliberately finite. They are starting points that should be +/// tuned with the process benchmark for the deployment's payload distribution. #[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(default)] +pub struct TransportConfig { + pub max_active_requests: usize, + pub max_active_long_polls: usize, + pub max_active_mutations: usize, + pub max_active_queries: usize, + pub reply_queue_count: usize, + pub reply_queue_bytes: usize, + pub max_wire_bytes: usize, + pub max_reply_bytes: usize, + pub max_batch_items: usize, + pub max_batch_bytes: usize, + pub lease_long_poll_ms: u64, + pub zmq_sndhwm: i32, + pub zmq_rcvhwm: i32, + pub zmq_linger_ms: i32, + pub shutdown_timeout_ms: u64, +} + +impl Default for TransportConfig { + fn default() -> Self { + Self { + max_active_requests: 1024, + max_active_long_polls: 512, + max_active_mutations: 384, + max_active_queries: 128, + reply_queue_count: 1024, + reply_queue_bytes: 64 * 1024 * 1024, + max_wire_bytes: 8 * 1024 * 1024, + max_reply_bytes: 8 * 1024 * 1024, + max_batch_items: 4096, + max_batch_bytes: 4 * 1024 * 1024, + lease_long_poll_ms: 30_000, + zmq_sndhwm: 1024, + zmq_rcvhwm: 1024, + zmq_linger_ms: 1000, + shutdown_timeout_ms: 30_000, + } + } +} + +impl TransportConfig { + pub fn validate(&self) -> Result<(), String> { + let nonzero = [ + ("max_active_requests", self.max_active_requests), + ("max_active_long_polls", self.max_active_long_polls), + ("max_active_mutations", self.max_active_mutations), + ("max_active_queries", self.max_active_queries), + ("reply_queue_count", self.reply_queue_count), + ("reply_queue_bytes", self.reply_queue_bytes), + ("max_wire_bytes", self.max_wire_bytes), + ("max_reply_bytes", self.max_reply_bytes), + ("max_batch_items", self.max_batch_items), + ("max_batch_bytes", self.max_batch_bytes), + ]; + for (name, value) in nonzero { + if value == 0 { + return Err(format!("transport.{name} must be greater than zero")); + } + } + if self.max_active_long_polls >= self.max_active_requests { + return Err( + "transport.max_active_long_polls must be lower than max_active_requests".into(), + ); + } + if self.max_reply_bytes > self.reply_queue_bytes { + return Err("transport.max_reply_bytes must not exceed reply_queue_bytes".into()); + } + if self.reply_queue_bytes > u32::MAX as usize { + return Err("transport.reply_queue_bytes must fit in a u32 semaphore budget".into()); + } + if self.max_wire_bytes > i64::MAX as usize { + return Err("transport.max_wire_bytes is too large for ZeroMQ".into()); + } + if self.max_batch_items > u32::MAX as usize { + return Err("transport.max_batch_items must fit in u32".into()); + } + if self.lease_long_poll_ms == 0 || self.shutdown_timeout_ms == 0 { + return Err("transport timeouts must be greater than zero".into()); + } + if self.zmq_sndhwm <= 0 || self.zmq_rcvhwm <= 0 || self.zmq_linger_ms < 0 { + return Err("transport ZMQ HWM values must be positive and linger non-negative".into()); + } + Ok(()) + } +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(default)] pub struct ConfigTomlServer { - #[serde(default = "default_host")] pub host: String, // Renamed from cgrpc_token; keep the alias so existing configs still authenticate. #[serde(alias = "cgrpc_token")] pub auth_token: Option, + + pub transport: TransportConfig, } + impl Default for ConfigTomlServer { fn default() -> Self { Self { host: default_host(), - auth_token: Option::None, + auth_token: None, + transport: TransportConfig::default(), } } } + #[derive(Debug, Clone, Default)] pub struct Config { pub config_toml: ConfigTomlServer, } impl Config { + /// Load `config.toml`, using defaults only when the file is absent. + /// Malformed or unsafe production configuration is returned as an error. + pub fn load_checked() -> Result { + let config_toml = match fs::read_to_string("config.toml") { + Ok(content) => toml::from_str::(&content) + .map_err(|err| format!("Failed to parse config.toml: {err}"))?, + Err(err) if err.kind() == ErrorKind::NotFound => ConfigTomlServer::default(), + Err(err) => return Err(format!("Failed to read config.toml: {err}")), + }; + config_toml.transport.validate()?; + Ok(Self { config_toml }) + } + + /// Compatibility loader for library/CLI callers. The production server uses + /// [`Config::load_checked`] and fails closed on malformed configuration. #[allow(clippy::new_without_default)] #[instrument] pub fn new() -> Self { - let mut content: String = "".to_owned(); - let result: Result = fs::read_to_string("config.toml"); - if let Ok(file_content) = result { - content = file_content; - } - let config_toml: ConfigTomlServer = toml::from_str(&content).unwrap_or_else(|err| { - error!("Failed to parse config file."); - error!("{:#?}", err); - ConfigTomlServer::default() - }); - Self { config_toml } + Self::load_checked().unwrap_or_else(|err| { + error!("{err}"); + Self::default() + }) } } diff --git a/enginelib/src/error.rs b/enginelib/src/error.rs index 7fae89c..c662f7a 100644 --- a/enginelib/src/error.rs +++ b/enginelib/src/error.rs @@ -1,15 +1,18 @@ use std::fmt; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum ErrorKind { NotFound, NotSupported, InvalidArgument, IOError, + Overloaded, Unknown, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Error { kind: ErrorKind, message: String, @@ -42,6 +45,10 @@ impl Error { Self::with_kind(ErrorKind::IOError, message) } + pub fn overloaded(message: impl Into) -> Error { + Self::with_kind(ErrorKind::Overloaded, message) + } + pub fn kind(&self) -> ErrorKind { self.kind } diff --git a/enginelib/src/lib.rs b/enginelib/src/lib.rs index 0b375e7..7eaa5ba 100644 --- a/enginelib/src/lib.rs +++ b/enginelib/src/lib.rs @@ -4,6 +4,7 @@ pub mod config; pub mod event; pub mod events; pub mod nucleus; +pub mod protocol; extern crate self as enginelib; pub use inventory; pub mod plugin; diff --git a/enginelib/src/nucleus/cancel.rs b/enginelib/src/nucleus/cancel.rs index e4b7e8b..daafde5 100644 --- a/enginelib/src/nucleus/cancel.rs +++ b/enginelib/src/nucleus/cancel.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::atomic::Ordering; use crate::{Identifier, api::ServerAPI, error::Error}; @@ -6,12 +7,13 @@ use crate::{Identifier, api::ServerAPI, error::Error}; /// list and resends it into the queue so another worker can pick it up. /// /// The persisted record is never touched, and the task_id stays in the dedup set -/// the whole time, so no concurrent `load()` can enqueue a duplicate. The requeue +/// the whole time, so no concurrent load can enqueue a duplicate. The requeue /// is a non-blocking `try_send` so there is no await window in which the lease is /// gone but the task isn't yet back in the channel: on success the task moves /// straight from "leased" to "queued"; if the channel is full or closed we clear -/// the dedup entry instead, dropping the task back to "persisted only" so a later -/// `load()` recovers it. Either way the task is never stranded. +/// the dedup entry and set the rescan flag instead, dropping the task back to +/// "persisted only" so the loader rewinds and re-enqueues it. Either way the +/// task is never stranded. pub fn cancel(api: Arc, task_type: Identifier, task_id: String) -> Result<(), Error> { // Remove the lease, capturing the stored task so we can requeue it. let stored_task = { @@ -38,10 +40,11 @@ pub fn cancel(api: Arc, task_type: Identifier, task_id: String) -> Re match queue.0.try_send((*stored_task).clone()) { // Queued again; task_id stays in the dedup set (still in-flight). Ok(()) => Ok(()), - // Channel full/closed: clear the dedup entry so the still-persisted - // record is re-enqueued by a later load(). Never left stranded. + // Channel full/closed: clear the dedup entry, then signal the loader to + // rescan so the still-persisted record is re-enqueued. Never stranded. Err(_) => { queue.2.remove(&task_id); + queue.3.store(true, Ordering::Release); Ok(()) } } diff --git a/enginelib/src/nucleus/complete.rs b/enginelib/src/nucleus/complete.rs index 191fc57..3ff8cc0 100644 --- a/enginelib/src/nucleus/complete.rs +++ b/enginelib/src/nucleus/complete.rs @@ -1,4 +1,8 @@ +use std::collections::HashSet; use std::sync::Arc; +use std::sync::atomic::Ordering; + +use rust_rocksdb::{WriteBatch, WriteOptions}; use crate::api::{finished_task_key, task_key}; use crate::{Identifier, Registry, api::ServerAPI, error::Error}; @@ -30,13 +34,17 @@ pub fn complete( } // The lease is already gone, so every remaining exit path must clear the - // dedup entry. On success it prevents load() from re-enqueuing a task whose - // pending record is deleted; on DB failure it lets a later load() re-enqueue - // the *still-persisted* pending record, instead of stranding it (no lease, no - // channel slot, but blocked from re-enqueue by the dedup entry) until restart. - let clear_dedup = || { + // dedup entry. On success it prevents the loader from re-enqueuing a task + // whose pending record is deleted. On DB failure the pending record still + // exists but sits *behind* the loader's cursor, so clearing the entry alone + // would strand it (no lease, no channel slot) until restart — the rescan + // flag tells the loader to rewind and re-enqueue it. + let clear_dedup = |rescan: bool| { if let Some(queue) = api.task_queue.tasks.get(&task_type) { queue.2.remove(&task_id); + if rescan { + queue.3.store(true, Ordering::Release); + } } }; @@ -46,17 +54,98 @@ pub fn complete( .db .put(finished_task_key(&task_type, &task_id), task_bytes) { - clear_dedup(); + clear_dedup(true); return Err(Error::io_error(format!( "Failed to persist finished task: {err}" ))); } if let Err(err) = api.db.delete(task_key(&task_type, &task_id)) { - clear_dedup(); - return Err(Error::io_error(format!("Failed to delete task record: {err}"))); + clear_dedup(true); + return Err(Error::io_error(format!( + "Failed to delete task record: {err}" + ))); } - clear_dedup(); + clear_dedup(false); Ok(()) } + +/// Complete a batch of leased tasks in one fused write, returning how many +/// succeeded. +/// +/// All `f:` puts and `t:` deletes go into a single RocksDB `WriteBatch` (one WAL +/// append instead of two per task). Only tasks actually leased are written; +/// unleased/duplicate ids are skipped. Deleting the `t:` record is what stops the +/// loader from ever re-enqueuing a completed task; clearing the dedup entry keeps +/// the loader's in-flight set accurate. On write failure the dedup entries are +/// cleared so the loader re-enqueues the still-persisted records rather than +/// stranding them (the rescan flag makes the loader rewind to find them). +pub fn complete_batch( + api: Arc, + task_type: Identifier, + results: Vec<(String, Vec)>, +) -> Result { + let task = api + .task_registry + .get(&task_type) + .ok_or(Error::not_found("TaskTypeNotFound"))?; + for (_, bytes) in &results { + if !task.verify(bytes) { + return Err(Error::invalid_argument("Failed to verify task bytes")); + } + } + + // Drop the lease for every requested id in one pass; record which were leased. + let requested: HashSet<&str> = results.iter().map(|(id, _)| id.as_str()).collect(); + let mut leased: HashSet = HashSet::new(); + if let Some(mut leases) = api.leased_tasks.tasks.get_mut(&task_type) { + leases.retain(|lease| { + let id = &lease.stored_task.task_id; + if requested.contains(id.as_str()) { + leased.insert(id.clone()); + false + } else { + true + } + }); + } + if leased.is_empty() { + return Ok(0); + } + + let mut batch = WriteBatch::default(); + let mut written: Vec = Vec::with_capacity(leased.len()); + for (task_id, bytes) in &results { + // `remove` (not `contains`) so a duplicate id in one request is written + // and counted exactly once. + if leased.remove(task_id) { + batch.put(finished_task_key(&task_type, task_id), bytes); + batch.delete(task_key(&task_type, task_id)); + written.push(task_id.clone()); + } + } + let mut wopts = WriteOptions::default(); + wopts.set_sync(false); + + let clear_dedup = |rescan: bool| { + if let Some(queue) = api.task_queue.tasks.get(&task_type) { + for task_id in &written { + queue.2.remove(task_id); + } + if rescan { + queue.3.store(true, Ordering::Release); + } + } + }; + + if let Err(err) = api.db.write_opt(&batch, &wopts) { + clear_dedup(true); + return Err(Error::io_error(format!( + "Failed to persist completed batch: {err}" + ))); + } + clear_dedup(false); + + Ok(written.len() as u32) +} diff --git a/enginelib/src/nucleus/lease.rs b/enginelib/src/nucleus/lease.rs index 9ba0a93..b34a29c 100644 --- a/enginelib/src/nucleus/lease.rs +++ b/enginelib/src/nucleus/lease.rs @@ -1,4 +1,6 @@ use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; use chrono::Utc; @@ -38,3 +40,127 @@ pub async fn lease( Ok(stored_task) } + +/// How long `lease_batch` waits for the first task before returning empty. +/// Bounds the damage of a dead client: a handler parked on an empty queue gives +/// up after this window instead of sitting on `recv()` forever and having a +/// future task handed to a peer that disconnected long ago (which would sit +/// leased to nobody for the whole TTL). Clients treat an empty reply as "poll +/// again". +const LEASE_LONG_POLL_SECS: u64 = 30; + +/// Lease up to `max` tasks of one type in a single call. +/// +/// Long-poll: awaits the first task (up to [`LEASE_LONG_POLL_SECS`]) so idle +/// workers don't busy-poll, then drains the rest of what's already queued +/// without waiting. Returns 0..=`max` tasks; empty means the poll window +/// elapsed (or `max == 0`) and the caller should lease again. Reads only the +/// in-memory channel — the loader is what keeps that channel fed from the DB. +pub async fn lease_batch( + api: Arc, + task_type: Identifier, + user_id: String, + max: u32, +) -> Result>, Error> { + lease_batch_limited( + api, + task_type, + user_id, + max, + usize::MAX, + Duration::from_secs(LEASE_LONG_POLL_SECS), + ) + .await +} + +/// Transport-facing lease with explicit response-byte and long-poll bounds. +/// Tasks that would exceed the byte budget are put back without changing their +/// dedup state, so the caller never receives an undeliverable lease. +pub async fn lease_batch_limited( + api: Arc, + task_type: Identifier, + user_id: String, + max: u32, + max_bytes: usize, + poll_timeout: Duration, +) -> Result>, Error> { + if max == 0 { + return Ok(Vec::new()); + } + + // Clone both channel halves out of the map guard before awaiting. + let (sender, receiver) = { + let queue = api + .task_queue + .tasks + .get(&task_type) + .ok_or(Error::not_found("TaskTypeNotFound"))?; + (queue.0.clone(), queue.1.clone()) + }; + + let first = match tokio::time::timeout(poll_timeout, receiver.recv()).await { + Err(_elapsed) => return Ok(Vec::new()), + Ok(recv) => recv.map_err(|err| Error::new(format!("Task queue closed: {err}")))?, + }; + + let mut used_bytes = estimated_wire_bytes(&first); + if used_bytes > max_bytes { + requeue_unleased(&api, &task_type, &sender, first); + return Err(Error::invalid_argument( + "Stored task exceeds the configured lease response limit", + )); + } + + let mut out: Vec> = vec![Arc::new(first)]; + while (out.len() as u32) < max { + match receiver.try_recv() { + Ok(task) => { + let task_bytes = estimated_wire_bytes(&task); + if used_bytes.saturating_add(task_bytes) > max_bytes { + requeue_unleased(&api, &task_type, &sender, task); + break; + } + used_bytes = used_bytes.saturating_add(task_bytes); + out.push(Arc::new(task)); + } + Err(_) => break, + } + } + + { + let now = Utc::now(); + let mut leases = api.leased_tasks.tasks.entry(task_type).or_default(); + for stored_task in &out { + leases.push(LeasedTask { + stored_task: stored_task.clone(), + user_id: user_id.clone(), + given_at: now, + }); + } + } + + Ok(out) +} + +fn estimated_wire_bytes(task: &StoredTask) -> usize { + task.bytes + .len() + .saturating_add(task.task_id.len()) + .saturating_add(task.task_type.0.len()) + .saturating_add(task.task_type.1.len()) + .saturating_add(64) +} + +fn requeue_unleased( + api: &Arc, + task_type: &Identifier, + sender: &async_channel::Sender, + task: StoredTask, +) { + if let Err(err) = sender.try_send(task) + && let Some(queue) = api.task_queue.tasks.get(task_type) + { + queue.2.remove(&err.into_inner().task_id); + queue.3.store(true, Ordering::Release); + } +} diff --git a/enginelib/src/nucleus/query.rs b/enginelib/src/nucleus/query.rs index b7915d6..c3e21e3 100644 --- a/enginelib/src/nucleus/query.rs +++ b/enginelib/src/nucleus/query.rs @@ -16,7 +16,7 @@ use crate::{api::ServerAPI, error::Error}; /// [`crate::api`] (`task_key_prefix`, `finished_task_key_prefix`, `task_key`, /// `finished_task_key`) to scope to a task type, its pending records, its /// finished records, or a single record. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum Query { /// Fetch a single record by its exact key. O(1) point read (`db.get`, not a /// scan) — the primitive for addressing one known task, e.g. the id returned @@ -73,7 +73,7 @@ pub struct Page { } /// The outcome of a [`Query`]. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum QueryResult { /// A single record, or `None` when the key is absent. Record(Option), @@ -96,7 +96,9 @@ pub fn query(api: Arc, query: Query) -> Result { .db .get(key.as_bytes()) .map_err(|err| Error::io_error(format!("Failed to read record: {err}")))?; - Ok(QueryResult::Record(value.map(|value| Record { key, value }))) + Ok(QueryResult::Record( + value.map(|value| Record { key, value }), + )) } Query::EstimateKeys => { let keys = api diff --git a/enginelib/src/nucleus/submit.rs b/enginelib/src/nucleus/submit.rs index 30ada41..84c7a05 100644 --- a/enginelib/src/nucleus/submit.rs +++ b/enginelib/src/nucleus/submit.rs @@ -1,24 +1,64 @@ use std::sync::Arc; +use rust_rocksdb::{WriteBatch, WriteOptions}; + use crate::api::task_key; use crate::{Identifier, Registry, api::ServerAPI, error::Error}; // t:namespace:task_name: -> Serialized Task Record // f:namespace:task_name: -> Finished Task Record + +/// Submit a single task. Thin wrapper over [`submit_batch`]; returns its id. pub fn submit( api: Arc, task_bytes: &[u8], task_type: Identifier, ) -> Result { + let mut ids = submit_batch(api, task_type, vec![task_bytes.to_vec()])?; + Ok(ids + .pop() + .expect("submit_batch returns one id per submitted task")) +} + +/// Persist a batch of task payloads of one `task_type`. +/// +/// **Write-only.** Submit's whole job is to make the tasks durable — one RocksDB +/// `WriteBatch` (one WAL append). It does *not* touch the in-memory queue. +/// Moving persisted records into the lease channel is the loader's job +/// ([`ServerAPI::run_loader`]), which streams them in with backpressure. That +/// keeps the DB the single source of truth: nothing is enqueued that isn't +/// durable, nothing is lost if a producer outruns consumers, and restart +/// recovery is the same load path — never a stranded task. +pub fn submit_batch( + api: Arc, + task_type: Identifier, + payloads: Vec>, +) -> Result, Error> { let task = api .task_registry .get(&task_type) .ok_or(Error::not_found("TaskTypeNotFound"))?; - if !task.verify(task_bytes) { - return Err(Error::invalid_argument("Failed to verify task bytes")); + + // Verify every payload before persisting anything — a batch is all-or-nothing. + for bytes in &payloads { + if !task.verify(bytes) { + return Err(Error::invalid_argument("Failed to verify task bytes")); + } } - let task_id = druid::Druid::default().to_hex(); + + let ids: Vec = payloads + .iter() + .map(|_| druid::Druid::default().to_hex()) + .collect(); + let mut batch = WriteBatch::default(); + for (id, bytes) in ids.iter().zip(&payloads) { + batch.put(task_key(&task_type, id), bytes); + } + let mut wopts = WriteOptions::default(); + // Do not acknowledge submissions until the WAL is durable on stable storage. + wopts.set_sync(true); api.db - .put(task_key(&task_type, &task_id), task_bytes) - .map_err(|err| Error::io_error(err.to_string()))?; - Ok(task_id) + .write_opt(&batch, &wopts) + .map_err(|err| Error::io_error(format!("Failed to persist task batch: {err}")))?; + + Ok(ids) } diff --git a/enginelib/src/prelude.rs b/enginelib/src/prelude.rs index 8e84693..f499c4e 100644 --- a/enginelib/src/prelude.rs +++ b/enginelib/src/prelude.rs @@ -29,6 +29,7 @@ pub use crate::events::start_event::StartEvent; pub use crate::events::task_block_acquired_event::TaskBlockAcquiredEvent; pub use crate::events::{Events, ID}; pub use crate::plugin::{LibraryDependency, LibraryMetadata}; +pub use crate::protocol::{Envelope, Request, Response}; pub use crate::task::{Runner, Task, Verifiable}; pub use crate::{Identifier, RawIdentifier, Registry}; pub use tracing; diff --git a/enginelib/src/protocol.rs b/enginelib/src/protocol.rs new file mode 100644 index 0000000..85603b6 --- /dev/null +++ b/enginelib/src/protocol.rs @@ -0,0 +1,101 @@ +//! Wire protocol shared by the ZeroMQ server and client. +//! +//! Everything here is pure `serde` — no transport dependency — so both the +//! `engine` server/client and any other consumer share one source of truth for +//! the wire format. Messages are postcard-serialized via the project serializer +//! ([`crate::api::to_allocvec`] / [`crate::api::from_bytes`]). +//! +//! The protocol is **batch-first**: a singular submit/lease/complete is just a +//! batch of size one, so a single ZeroMQ message can carry many tasks. That is +//! the key lever for sustaining 100k+ tasks/s — it collapses the per-message +//! ZeroMQ + syscall cost across a whole batch. + +use serde::{Deserialize, Serialize}; + +use crate::Identifier; +use crate::api::StoredTask; +use crate::error::Error; +use crate::nucleus::query::{Query, QueryResult}; + +/// A request plus the caller's auth token, as sent on the wire. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Envelope { + /// Caller auth token, checked server-side before dispatch. + pub auth: String, + /// The requested operation. + pub req: Request, +} + +/// A client → server operation. Batch variants carry N items in one message. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum Request { + /// Submit N task payloads of one type. Singular submit = a one-element vec. + SubmitBatch { + task_type: Identifier, + tasks: Vec>, + }, + /// Lease up to `max` tasks of one type for `user_id` (long-poll: blocks for + /// the first, then drains the rest without waiting). + LeaseBatch { + task_type: Identifier, + user_id: String, + max: u32, + }, + /// Complete N leased tasks: `(task_id, result_bytes)` pairs. + CompleteBatch { + task_type: Identifier, + results: Vec<(String, Vec)>, + }, + /// Renew a single lease's TTL. + Renew { + task_type: Identifier, + task_id: String, + }, + /// Cancel a single lease (requeues the task). + Cancel { + task_type: Identifier, + task_id: String, + }, + /// Administrative / CRUD query over the task store. + Query(Query), +} + +/// A server → client reply. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum Response { + /// Ids assigned to a `SubmitBatch`, in request order. + Submitted { task_ids: Vec }, + /// Tasks handed out by a `LeaseBatch` (may be fewer than `max`). + Leased(Vec), + /// A `CompleteBatch` finished; `ok` is how many were completed. + Completed { ok: u32 }, + /// A `Renew` succeeded. + Renewed, + /// A `Cancel` succeeded. + Cancelled, + /// A `Query` result. + Query(QueryResult), + /// The transport is at capacity. Mutations were not executed and may be + /// retried after the suggested delay. + Overloaded { retry_after_ms: u64 }, + /// The operation failed. + Err(Error), +} + +/// postcard-encode any wire value, mapping serializer failure to [`Error`]. +pub fn encode(value: &T) -> Result, Error> { + crate::api::to_allocvec(value).map_err(|err| Error::new(format!("encode failed: {err}"))) +} + +/// postcard-decode a wire value, mapping malformed input to an invalid-argument +/// [`Error`]. +pub fn decode<'a, T: Deserialize<'a>>(bytes: &'a [u8]) -> Result { + crate::api::from_bytes(bytes) + .map_err(|err| Error::invalid_argument(format!("decode failed: {err}"))) +} + +impl From for Response { + fn from(err: Error) -> Self { + Response::Err(err) + } +}