From 169aa58e0e3d39da1201378adc286595ae1157b3 Mon Sep 17 00:00:00 2001 From: Vladimir Petrzhikovskii Date: Tue, 8 Apr 2025 16:36:27 +0200 Subject: [PATCH 1/2] feat(util): add rate-limited `AsyncWrite` adapter --- Cargo.lock | 284 ++++++++++++++++++---- Cargo.toml | 1 + util/Cargo.toml | 10 +- util/src/{io.rs => io/mod.rs} | 2 + util/src/io/ratelimit.rs | 433 ++++++++++++++++++++++++++++++++++ 5 files changed, 677 insertions(+), 53 deletions(-) rename util/src/{io.rs => io/mod.rs} (94%) create mode 100644 util/src/io/ratelimit.rs diff --git a/Cargo.lock b/Cargo.lock index 856bafd34b..e10e0c5503 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,10 +34,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", - "getrandom", + "getrandom 0.2.15", "once_cell", "version_check", - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -49,6 +49,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anes" version = "0.1.6" @@ -780,6 +786,20 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "data-encoding" version = "2.6.0" @@ -947,7 +967,7 @@ dependencies = [ "curve25519-dalek", "generic-array", "hex", - "rand", + "rand 0.8.5", "serde", "sha2", "tl-proto", @@ -972,7 +992,7 @@ dependencies = [ "hex", "num-bigint", "num-traits", - "rand", + "rand 0.8.5", "rayon", "scc", "serde", @@ -1057,6 +1077,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "form_urlencoded" version = "1.2.1" @@ -1143,6 +1169,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.31" @@ -1191,7 +1223,21 @@ checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", +] + +[[package]] +name = "getrandom" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73fea8450eea4bac3940448fb7ae50d91f034f941199fcd9d909a5a07aa455f0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi", + "wasi 0.14.2+wasi-0.2.4", + "wasm-bindgen", ] [[package]] @@ -1206,6 +1252,29 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +[[package]] +name = "governor" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbe789d04bf14543f03c4b60cd494148aa79438c8440ae7d81a7778147745c3" +dependencies = [ + "cfg-if", + "dashmap 6.1.0", + "futures-sink", + "futures-timer", + "futures-util", + "getrandom 0.3.2", + "hashbrown 0.15.2", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand 0.9.0", + "smallvec", + "spinning_top", + "web-time", +] + [[package]] name = "h2" version = "0.4.6" @@ -1252,9 +1321,14 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.0" +version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" +checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] [[package]] name = "heck" @@ -1293,7 +1367,7 @@ dependencies = [ "hickory-proto", "once_cell", "radix_trie", - "rand", + "rand 0.8.5", "thiserror 1.0.66", "tokio", "tracing", @@ -1315,7 +1389,7 @@ dependencies = [ "idna 0.4.0", "ipnet", "once_cell", - "rand", + "rand 0.8.5", "thiserror 1.0.66", "tinyvec", "tokio", @@ -1487,7 +1561,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" dependencies = [ "equivalent", - "hashbrown 0.15.0", + "hashbrown 0.15.2", ] [[package]] @@ -1557,10 +1631,11 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.70" +version = "0.3.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1868808506b929d7b0cfa8f75951347aa71bb21144b7791bae35d9bccfcfe37a" +checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" dependencies = [ + "once_cell", "wasm-bindgen", ] @@ -1757,7 +1832,7 @@ checksum = "80e04d1dcff3aae0704555fe5fee3bcfaf3d1fdf8a7e521d5b9d2b42acb52cec" dependencies = [ "hermit-abi 0.3.9", "libc", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.52.0", ] @@ -1812,6 +1887,12 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "ntapi" version = "0.4.1" @@ -1941,7 +2022,7 @@ dependencies = [ "once_cell", "opentelemetry_api", "percent-encoding", - "rand", + "rand 0.8.5", "thiserror 1.0.66", ] @@ -2155,7 +2236,7 @@ version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" dependencies = [ - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -2240,7 +2321,7 @@ dependencies = [ "libc", "once_cell", "raw-cpuid", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", "web-sys", "winapi", ] @@ -2282,7 +2363,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fadfaed2cd7f389d0161bb73eeb07b7b78f8691047a6f3e73caaeae55310a4a6" dependencies = [ "bytes", - "rand", + "rand 0.8.5", "ring", "rustc-hash 2.1.0", "rustls", @@ -2314,6 +2395,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" + [[package]] name = "radix_trie" version = "0.2.1" @@ -2331,8 +2418,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.3", + "zerocopy 0.8.24", ] [[package]] @@ -2342,7 +2440,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.3", ] [[package]] @@ -2351,7 +2459,16 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.15", +] + +[[package]] +name = "rand_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +dependencies = [ + "getrandom 0.3.2", ] [[package]] @@ -2360,7 +2477,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59cad018caf63deb318e5a4586d99a24424a364f40f1e5778c29aca23f4fc73e" dependencies = [ - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -2407,7 +2524,7 @@ version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" dependencies = [ - "getrandom", + "getrandom 0.2.15", "libredox", "thiserror 1.0.66", ] @@ -2507,7 +2624,7 @@ checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" dependencies = [ "cc", "cfg-if", - "getrandom", + "getrandom 0.2.15", "libc", "spin", "untrusted", @@ -2805,7 +2922,7 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -2845,6 +2962,15 @@ version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "spki" version = "0.7.3" @@ -2958,7 +3084,7 @@ dependencies = [ "humantime", "opentelemetry", "pin-project", - "rand", + "rand 0.8.5", "serde", "static_assertions", "tarpc-plugins", @@ -3512,7 +3638,7 @@ dependencies = [ "metrics", "num-bigint", "parking_lot", - "rand", + "rand 0.8.5", "rayon", "reqwest", "rustc_version", @@ -3558,7 +3684,7 @@ dependencies = [ "indexmap 2.6.0", "metrics", "parking_lot", - "rand", + "rand 0.8.5", "rayon", "scc", "scopeguard", @@ -3593,7 +3719,7 @@ dependencies = [ "bumpalo", "bytes", "clap", - "dashmap", + "dashmap 5.5.3", "everscale-crypto", "everscale-types", "futures-util", @@ -3602,7 +3728,7 @@ dependencies = [ "itertools 0.12.1", "metrics", "parking_lot", - "rand", + "rand 0.8.5", "rand_pcg", "rayon", "scopeguard", @@ -3665,7 +3791,7 @@ dependencies = [ "metrics", "parking_lot", "pin-project-lite", - "rand", + "rand 0.8.5", "scopeguard", "serde", "tempfile", @@ -3708,7 +3834,7 @@ dependencies = [ "clap", "everscale-crypto", "everscale-types", - "rand", + "rand 0.8.5", "serde", "serde_json", "tokio", @@ -3733,7 +3859,7 @@ dependencies = [ "bytesize", "castaway", "clap", - "dashmap", + "dashmap 5.5.3", "ed25519", "everscale-crypto", "exponential-backoff", @@ -3746,7 +3872,7 @@ dependencies = [ "pin-project-lite", "pkcs8", "quinn", - "rand", + "rand 0.8.5", "ring", "rustls", "rustls-webpki", @@ -3780,7 +3906,7 @@ dependencies = [ "num-bigint", "parking_lot", "prost", - "rand", + "rand 0.8.5", "scopeguard", "serde", "serde_json", @@ -3823,7 +3949,7 @@ dependencies = [ "bytes", "bytesize", "crc32c", - "dashmap", + "dashmap 5.5.3", "everscale-types", "fdlimit", "futures-util", @@ -3834,7 +3960,7 @@ dependencies = [ "parking_lot", "parking_lot_core", "quick_cache", - "rand", + "rand 0.8.5", "rlimit", "scopeguard", "serde", @@ -3860,16 +3986,19 @@ dependencies = [ "anyhow", "base64", "bytes", + "bytesize", "castaway", "criterion", - "dashmap", + "dashmap 5.5.3", "futures-util", "getip", + "governor", "humantime", "libc", "metrics", "metrics-exporter-prometheus", - "rand", + "pin-project-lite", + "rand 0.8.5", "rayon", "serde", "serde_json", @@ -3993,7 +4122,7 @@ version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" dependencies = [ - "getrandom", + "getrandom 0.2.15", ] [[package]] @@ -4039,26 +4168,35 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wasi" +version = "0.14.2+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" +dependencies = [ + "wit-bindgen-rt", +] + [[package]] name = "wasm-bindgen" -version = "0.2.93" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a82edfc16a6c469f5f44dc7b571814045d60404b55a0ee849f9bcfa2e63dd9b5" +checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" dependencies = [ "cfg-if", "once_cell", + "rustversion", "wasm-bindgen-macro", ] [[package]] name = "wasm-bindgen-backend" -version = "0.2.93" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9de396da306523044d3302746f1208fa71d7532227f15e347e2d93e4145dd77b" +checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" dependencies = [ "bumpalo", "log", - "once_cell", "proc-macro2", "quote", "syn 2.0.96", @@ -4079,9 +4217,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.93" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "585c4c91a46b072c92e908d99cb1dcdf95c5218eeb6f3bf1efa991ee7a68cccf" +checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4089,9 +4227,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.93" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "afc340c74d9005395cf9dd098506f7f44e38f2b4a21c6aaacf9a105ea5e1e836" +checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" dependencies = [ "proc-macro2", "quote", @@ -4102,9 +4240,12 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.93" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c62a0a307cb4a311d3a07867860911ca130c3494e8c2719593806c08bc5d0484" +checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +dependencies = [ + "unicode-ident", +] [[package]] name = "web-sys" @@ -4116,6 +4257,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-roots" version = "0.26.6" @@ -4378,6 +4529,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wit-bindgen-rt" +version = "0.39.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" +dependencies = [ + "bitflags", +] + [[package]] name = "zerocopy" version = "0.7.35" @@ -4385,7 +4545,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" dependencies = [ "byteorder", - "zerocopy-derive", + "zerocopy-derive 0.7.35", +] + +[[package]] +name = "zerocopy" +version = "0.8.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879" +dependencies = [ + "zerocopy-derive 0.8.24", ] [[package]] @@ -4399,6 +4568,17 @@ dependencies = [ "syn 2.0.96", ] +[[package]] +name = "zerocopy-derive" +version = "0.8.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.96", +] + [[package]] name = "zeroize" version = "1.8.1" diff --git a/Cargo.toml b/Cargo.toml index 1c099e639b..735870038c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,6 +53,7 @@ exponential-backoff = "1" fdlimit = "0.3.0" futures-util = "0.3" getip = "0.1" +governor = { version = "0.10.0" } hex = "0.4" humantime = "2" indexmap = "2.2" diff --git a/util/Cargo.toml b/util/Cargo.toml index 611d648f72..feea8b4fc6 100644 --- a/util/Cargo.toml +++ b/util/Cargo.toml @@ -15,10 +15,12 @@ ahash = { workspace = true } anyhow = { workspace = true } base64 = { workspace = true } bytes = { workspace = true } +bytesize = { workspace = true, optional = true } castaway = { workspace = true } dashmap = { workspace = true } futures-util = { workspace = true } getip = { workspace = true, optional = true } +governor = { workspace = true, optional = true, features = ["std", "jitter"] } humantime = { workspace = true } libc = { workspace = true, optional = true } metrics = { workspace = true } @@ -39,11 +41,12 @@ zstd-sys = { workspace = true } metrics-exporter-prometheus = { workspace = true, optional = true } tikv-jemalloc-ctl = { workspace = true, optional = true } +pin-project-lite = { workspace = true, optional = true } [dev-dependencies] criterion = "0.5.1" tempfile = { workspace = true } -tokio = { workspace = true, features = ["time", "sync", "rt-multi-thread", "macros"] } +tokio = { workspace = true, features = ["time", "sync", "rt-multi-thread", "macros", "io-util"] } tracing-subscriber = { workspace = true, features = ["env-filter"] } [features] @@ -58,6 +61,11 @@ cli = [ "metrics-exporter-prometheus", "tikv-jemalloc-ctl", ] +governor = [ + "dep:governor", + "dep:pin-project-lite", + "dep:bytesize", +] [[bench]] name = "p2" diff --git a/util/src/io.rs b/util/src/io/mod.rs similarity index 94% rename from util/src/io.rs rename to util/src/io/mod.rs index 865bacbbe2..a657c7b9ff 100644 --- a/util/src/io.rs +++ b/util/src/io/mod.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "governor")] +pub mod ratelimit; // TODO: Extend with required methods. pub trait ByteOrderRead { fn read_be_uint(&mut self, bytes: usize) -> std::io::Result; diff --git a/util/src/io/ratelimit.rs b/util/src/io/ratelimit.rs new file mode 100644 index 0000000000..a0610eeb84 --- /dev/null +++ b/util/src/io/ratelimit.rs @@ -0,0 +1,433 @@ +use std::future::Future; +use std::io::{self}; +use std::num::NonZeroU32; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +use governor::clock::{self, QuantaClock, ReasonablyRealtime}; +use governor::middleware::RateLimitingMiddleware; +use governor::state::{DirectStateStore, InMemoryState, NotKeyed}; +use governor::{NotUntil, RateLimiter}; +use pin_project_lite::pin_project; +use tokio::io::AsyncWrite; +use tokio::time::Sleep; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum State { + CheckPermit, + Waiting, + Writing, +} + +pin_project! { + #[derive(Debug)] + pub struct RatelimitedWriter< + 'a, + W, + D: DirectStateStore, + C: ReasonablyRealtime, + MW: RateLimitingMiddleware> + > { + #[pin] + inner: W, + state: State, + limiter: &'a RateLimiter, + #[pin] + delay: Sleep, + } +} + +impl<'a, W, D, C, MW> RatelimitedWriter<'a, W, D, C, MW> +where + W: AsyncWrite, + D: DirectStateStore, + C: clock::Clock + ReasonablyRealtime, + MW: RateLimitingMiddleware>, +{ + pub fn new(inner: W, limiter: &'a RateLimiter) -> Self { + let delay = tokio::time::sleep(Duration::ZERO); + RatelimitedWriter { + inner, + limiter, + state: State::CheckPermit, + delay, + } + } + + pub fn get_ref(&self) -> &W { + &self.inner + } + + pub fn get_mut(self: Pin<&mut Self>) -> Pin<&mut W> { + self.project().inner + } + + pub fn into_inner(self) -> W { + self.inner + } +} + +impl AsyncWrite for RatelimitedWriter<'_, W, D, C, MW> +where + W: AsyncWrite, + D: DirectStateStore, + C: clock::Clock + ReasonablyRealtime, + MW: RateLimitingMiddleware>, +{ + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let mut this = self.project(); + + loop { + match this.state { + State::CheckPermit => match this.limiter.check() { + Ok(_) => { + *this.state = State::Writing; + } + Err(negative) => { + *this.state = State::Waiting; + + let delay = negative.wait_time_from(this.limiter.clock().reference_point()); + + if delay > Duration::ZERO { + let deadline = tokio::time::Instant::now() + delay; + this.delay.as_mut().reset(deadline); + } else { + *this.state = State::CheckPermit; + } + } + }, + + State::Waiting => match this.delay.as_mut().poll(cx) { + Poll::Pending => { + return Poll::Pending; + } + Poll::Ready(_) => { + *this.state = State::CheckPermit; + } + }, + + State::Writing => { + let offer_len = std::cmp::min(buf.len(), 1024); + let offer_buf = &buf[..offer_len]; + + let result = this.inner.as_mut().poll_write(cx, offer_buf); + + return match result { + Poll::Ready(Ok(bytes_written)) => { + *this.state = State::CheckPermit; + Poll::Ready(Ok(bytes_written)) + } + Poll::Ready(Err(e)) => { + *this.state = State::CheckPermit; + Poll::Ready(Err(e)) + } + Poll::Pending => Poll::Pending, + }; + } + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_shutdown(cx) + } + + fn is_write_vectored(&self) -> bool { + false + } +} + +pub fn rate_limiter( + speed: bytesize::ByteSize, +) -> Result, RlConstructError> { + let num_cells = speed.as_u64() / 1024; + + let speed: u32 = match num_cells.try_into() { + Ok(num_cells) => num_cells, + Err(_) => return Err(RlConstructError::InvalidRateLimit(speed)), + }; + let speed = NonZeroU32::new(speed).ok_or(RlConstructError::ZeroRateLimit)?; + let quota = governor::Quota::per_second(speed); + + Ok(governor::RateLimiter::direct(quota)) +} + +#[derive(thiserror::Error, Debug)] +pub enum RlConstructError { + #[error("Invalid rate limit: {0} / sec. Must be less than 35.18 terabits / second")] + InvalidRateLimit(bytesize::ByteSize), + #[error("Rate limit is zero")] + ZeroRateLimit, +} + +pub trait RatelimitAsyncWriteExt: AsyncWrite + Sized { + /// Creates a new `RatelimitedWriter` that wraps the current writer and applies the given rate limiter. + /// ONE RATE LIMITER CELL IS EQUIVALENT TO 1024 bytes written. + /// So if you want to ratelimit to 1MB/s, you need to set the rate limiter to 1024 cells per second. + fn ratelimit_write( + self, + limiter: &RateLimiter, + ) -> RatelimitedWriter<'_, Self, D, C, MW> + where + D: DirectStateStore, + C: clock::Clock + ReasonablyRealtime, + MW: RateLimitingMiddleware>, + { + RatelimitedWriter::new(self, limiter) + } +} + +impl RatelimitAsyncWriteExt for W where W: AsyncWrite + Sized {} + +#[cfg(test)] +mod tests { + use std::io; + use std::io::{Cursor, Read}; + use std::num::NonZeroU32; + use std::pin::Pin; + use std::task::{Context, Poll}; + + use governor::{Quota, RateLimiter}; + use tokio::io::{AsyncWrite, AsyncWriteExt}; + use tokio::time::{Duration, Instant}; + + use super::*; + + // --- Mock Writers --- + + #[derive(Debug)] + struct PendingOnceWriter { + first_poll: bool, + total_written: usize, + inner: Vec, + } + + impl PendingOnceWriter { + fn new() -> Self { + PendingOnceWriter { + first_poll: true, + total_written: 0, + inner: Vec::new(), + } + } + } + + impl AsyncWrite for PendingOnceWriter { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + if self.first_poll { + self.first_poll = false; + cx.waker().wake_by_ref(); + Poll::Pending + } else { + let n = buf.len(); + self.inner.extend_from_slice(buf); + self.total_written += n; + Poll::Ready(Ok(n)) + } + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + #[derive(Debug)] + struct ErrorWriter { + error_kind: io::ErrorKind, + } + + impl ErrorWriter { + fn new(kind: io::ErrorKind) -> Self { + ErrorWriter { error_kind: kind } + } + } + + impl AsyncWrite for ErrorWriter { + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &[u8], + ) -> Poll> { + Poll::Ready(Err(io::Error::new(self.error_kind, "mock write error"))) + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + fn permissive_limiter() -> RateLimiter { + let quota = Quota::per_second(NonZeroU32::new(u32::MAX).unwrap()); + RateLimiter::direct(quota) + } + + fn test_limiter( + cells_per_sec: u32, + burst: u32, + ) -> RateLimiter { + let quota = Quota::per_second(NonZeroU32::new(cells_per_sec).unwrap()) + .allow_burst(NonZeroU32::new(burst).unwrap()); + RateLimiter::direct(quota) + } + + // --- Tests --- + + #[tokio::test] + async fn test_ratelimited_writer_basic() { + let limiter = test_limiter(1, 1); + let data: Vec = vec![b'a'; 1024 * 3]; + let buffer = Vec::with_capacity(data.len()); + let cursor = Cursor::new(buffer); + + let writer = cursor.ratelimit_write(&limiter); + let mut writer = std::pin::pin!(writer); + + let start = Instant::now(); + writer.write_all(&data).await.expect("write_all failed"); + let elapsed = start.elapsed(); + + println!("Basic test elapsed: {:?}", elapsed); + assert!( + elapsed >= Duration::from_millis(1900) && elapsed < Duration::from_millis(2500), + "Expected ~2s elapsed time, got {:?}", + elapsed + ); + + let mut final_cursor = RatelimitedWriter::get_mut(writer); + final_cursor.set_position(0); + let mut written_data = Vec::new(); + final_cursor + .read_to_end(&mut written_data) + .expect("read failed"); + assert_eq!(written_data, data); + } + + #[tokio::test] + async fn test_ratelimited_writer_small_writes() { + let limiter = test_limiter(2, 1); + let data_chunk: Vec = vec![b'b'; 100]; + let num_chunks = 5; + let total_size = data_chunk.len() * num_chunks; + + let buffer = Vec::with_capacity(total_size); + let cursor = Cursor::new(buffer); + + let writer = cursor.ratelimit_write(&limiter); + let mut writer = std::pin::pin!(writer); + + let start = Instant::now(); + for _ in 0..num_chunks { + writer + .write_all(&data_chunk) + .await + .expect("write_all chunk failed"); + } + let elapsed = start.elapsed(); + + println!("Small writes test elapsed: {:?}", elapsed); + assert!( + elapsed >= Duration::from_millis(1500) && elapsed < Duration::from_millis(2100), + "Expected ~1.5s elapsed time, got {:?}", + elapsed + ); + + let mut final_cursor = RatelimitedWriter::get_mut(writer); + final_cursor.set_position(0); + let mut written_data = Vec::new(); + final_cursor + .read_to_end(&mut written_data) + .expect("read failed"); + assert_eq!(written_data.len(), total_size); + assert!(written_data.iter().all(|&x| x == b'b')); + } + + #[tokio::test] + async fn test_ratelimited_writer_burst() { + let limiter = test_limiter(1, 3); + let data: Vec = vec![b'c'; 1024 * 5]; + let buffer = Vec::with_capacity(data.len()); + let cursor = Cursor::new(buffer); + + let writer = cursor.ratelimit_write(&limiter); + let mut writer = std::pin::pin!(writer); + + let start = Instant::now(); + writer.write_all(&data).await.expect("write_all failed"); + let elapsed = start.elapsed(); + + println!("Burst test elapsed: {:?}", elapsed); + assert!( + elapsed >= Duration::from_millis(1900) && elapsed < Duration::from_millis(2500), + "Expected ~2s elapsed time, got {:?}", + elapsed + ); + + let mut final_cursor = RatelimitedWriter::get_mut(writer); + final_cursor.set_position(0); + let mut written_data = Vec::new(); + final_cursor + .read_to_end(&mut written_data) + .expect("read failed"); + assert_eq!(written_data, data); + } + + #[tokio::test] + async fn test_ratelimited_writer_inner_error() { + let limiter = permissive_limiter(); + let inner_writer = ErrorWriter::new(io::ErrorKind::BrokenPipe); + let data = vec![0u8; 100]; + + let writer = inner_writer.ratelimit_write(&limiter); + let mut writer = std::pin::pin!(writer); + + let result = writer.write_all(&data).await; + + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!(e.kind(), io::ErrorKind::BrokenPipe); + assert_eq!(e.to_string(), "mock write error"); + } + } + + #[tokio::test] + async fn test_ratelimited_writer_inner_pending() { + let limiter = permissive_limiter(); + let inner_writer = PendingOnceWriter::new(); + let data = vec![1u8; 512]; + + let writer = inner_writer.ratelimit_write(&limiter); + let mut writer = std::pin::pin!(writer); + + let result = writer.write_all(&data).await; + assert!(result.is_ok()); + + let final_inner_writer = &writer.inner; + assert!( + !final_inner_writer.first_poll, + "Inner writer should have been polled more than once" + ); + assert_eq!(final_inner_writer.total_written, data.len()); + assert_eq!(final_inner_writer.inner, data); + } +} From 6a099a9e62947af13dc1b00b91d100b5137d2177 Mon Sep 17 00:00:00 2001 From: Vladimir Petrzhikovskii Date: Fri, 11 Apr 2025 14:48:38 +0200 Subject: [PATCH 2/2] feat(network): apply rate-limit to outgoing connections --- Cargo.lock | 1 + network/Cargo.toml | 3 ++- network/src/network/config.rs | 5 +++++ network/src/network/peer.rs | 17 ++++++++++++++--- 4 files changed, 22 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e10e0c5503..fc111867bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3864,6 +3864,7 @@ dependencies = [ "everscale-crypto", "exponential-backoff", "futures-util", + "governor", "hex", "indexmap 2.6.0", "metrics", diff --git a/network/Cargo.toml b/network/Cargo.toml index eb88990257..fa22d096ba 100644 --- a/network/Cargo.toml +++ b/network/Cargo.toml @@ -27,6 +27,7 @@ ed25519 = { workspace = true, features = ["alloc", "pkcs8"] } everscale-crypto = { workspace = true } exponential-backoff = { workspace = true } futures-util = { workspace = true, features = ["sink"] } +governor = { workspace = true, features = ["std", "jitter"] } hex = { workspace = true } indexmap = { workspace = true } metrics = { workspace = true } @@ -48,7 +49,7 @@ tokio-util = { workspace = true, features = ["time"] } tracing = { workspace = true } # local deps -tycho-util = { workspace = true } +tycho-util = { workspace = true, features = ["governor"] } [dev-dependencies] clap = { workspace = true, features = ["derive"] } diff --git a/network/src/network/config.rs b/network/src/network/config.rs index 39d7d85f86..671f240c64 100644 --- a/network/src/network/config.rs +++ b/network/src/network/config.rs @@ -73,6 +73,10 @@ pub struct NetworkConfig { /// Default: disabled. pub connection_metrics: Option, + + /// Default: 250 mbit/s + /// Limits all outgoing connections to the same peer to this value. + pub bandwidth_cap: bytesize::ByteSize, } impl Default for NetworkConfig { @@ -93,6 +97,7 @@ impl Default for NetworkConfig { shutdown_idle_timeout: Duration::from_secs(60), enable_0rtt: false, connection_metrics: None, + bandwidth_cap: bytesize::ByteSize::mb(250 / 8), // in megabits per second } } } diff --git a/network/src/network/peer.rs b/network/src/network/peer.rs index 1e03d8064b..90a028235b 100644 --- a/network/src/network/peer.rs +++ b/network/src/network/peer.rs @@ -1,7 +1,10 @@ use std::sync::Arc; use anyhow::Result; +use governor::clock::QuantaClock; +use governor::state::{InMemoryState, NotKeyed}; use tokio_util::codec::{FramedRead, FramedWrite}; +use tycho_util::io::ratelimit::{RatelimitAsyncWriteExt, RatelimitedWriter}; use tycho_util::metrics::{GaugeGuard, HistogramGuard}; use crate::network::config::NetworkConfig; @@ -25,11 +28,18 @@ const METRIC_OUT_MESSAGES: &str = "tycho_net_out_messages"; pub struct Peer { connection: Connection, config: Arc, + rate_limiter: Arc>, } impl Peer { pub(crate) fn new(connection: Connection, config: Arc) -> Self { - Self { connection, config } + Self { + connection, + rate_limiter: Arc::new( + tycho_util::io::ratelimit::rate_limiter(config.bandwidth_cap).unwrap(), + ), + config, + } } pub fn peer_id(&self) -> &PeerId { @@ -42,12 +52,13 @@ impl Peer { let _histogram = HistogramGuard::begin(METRIC_OUT_QUERIES_TIME); let (send_stream, recv_stream) = self.connection.open_bi().await?; + let send_stream = send_stream.ratelimit_write(&self.rate_limiter); + let send_stream = std::pin::pin!(send_stream); let mut send_stream = FramedWrite::new(send_stream, make_codec(&self.config)); let mut recv_stream = FramedRead::new(recv_stream, make_codec(&self.config)); send_request(&mut send_stream, request).await?; - send_stream.get_mut().finish()?; - + RatelimitedWriter::get_mut(send_stream.into_inner()).finish()?; recv_response(&mut recv_stream).await.map_err(Into::into) }