zstd should reuse ctx to avoid huge mmap cost (#941)

This commit is contained in:
Sijie.Sun
2025-06-06 08:59:06 +08:00
committed by GitHub
parent 0314c66635
commit a6773aa549
5 changed files with 196 additions and 109 deletions

223
Cargo.lock generated
View File

@@ -23,6 +23,12 @@ version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
[[package]]
name = "adler2"
version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
[[package]]
name = "aead"
version = "0.5.2"
@@ -196,6 +202,15 @@ dependencies = [
"num-traits",
]
[[package]]
name = "arbitrary"
version = "1.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dde20b3d026af13f561bdd0f15edf01fc734f0dafcedbaf42bba506a9517f223"
dependencies = [
"derive_arbitrary",
]
[[package]]
name = "arboard"
version = "3.4.0"
@@ -229,7 +244,7 @@ dependencies = [
"base64ct",
"blake2",
"cpufeatures",
"password-hash 0.5.0",
"password-hash",
]
[[package]]
@@ -274,8 +289,8 @@ dependencies = [
"memchr",
"pin-project-lite",
"tokio",
"zstd 0.13.2",
"zstd-safe 7.2.1",
"zstd",
"zstd-safe",
]
[[package]]
@@ -632,7 +647,7 @@ dependencies = [
"cc",
"cfg-if",
"libc",
"miniz_oxide",
"miniz_oxide 0.7.4",
"object",
"rustc-demangle",
]
@@ -853,9 +868,9 @@ dependencies = [
[[package]]
name = "bumpalo"
version = "3.16.0"
version = "3.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c"
checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf"
[[package]]
name = "bytecheck"
@@ -924,22 +939,20 @@ dependencies = [
[[package]]
name = "bzip2"
version = "0.4.4"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8"
checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47"
dependencies = [
"bzip2-sys",
"libc",
]
[[package]]
name = "bzip2-sys"
version = "0.1.11+1.0.8"
version = "0.1.13+1.0.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "736a955f3fa7875102d57c82b8cac37ec45224a07fd32d58f9f7a186b6cd4cdc"
checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14"
dependencies = [
"cc",
"libc",
"pkg-config",
]
@@ -1289,9 +1302,9 @@ checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
[[package]]
name = "constant_time_eq"
version = "0.1.5"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6"
[[package]]
name = "conv"
@@ -1668,6 +1681,12 @@ dependencies = [
"thiserror 1.0.63",
]
[[package]]
name = "deflate64"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "da692b8d1080ea3045efaab14434d40468c3d8657e42abddfffca87b428f4c1b"
[[package]]
name = "defmt"
version = "0.3.8"
@@ -1713,14 +1732,25 @@ dependencies = [
[[package]]
name = "deranged"
version = "0.3.11"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4"
checksum = "9c9e6a11ca8224451684bc0d7d5a7adbf8f2fd6887261a1cfc3c0432f9d4068e"
dependencies = [
"powerfmt",
"serde",
]
[[package]]
name = "derive_arbitrary"
version = "1.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30542c1ad912e0e3d22a1935c290e12e8a29d704a420177a31faad4a601a0800"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.87",
]
[[package]]
name = "derive_builder"
version = "0.20.2"
@@ -1916,7 +1946,6 @@ version = "2.3.1"
dependencies = [
"aes-gcm",
"anyhow",
"async-compression",
"async-recursion",
"async-ringbuf",
"async-stream",
@@ -2025,6 +2054,7 @@ dependencies = [
"winreg 0.52.0",
"zerocopy",
"zip",
"zstd",
]
[[package]]
@@ -2396,12 +2426,13 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
[[package]]
name = "flate2"
version = "1.0.31"
version = "1.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f211bbe8e69bbd0cfdea405084f128ae8b4aaa6b0b522fc8f2b009084797920"
checksum = "7ced92e76e966ca2fd84c8f7aa01a4aea65b0eb6648d72f7c8f3e2764a67fece"
dependencies = [
"crc32fast",
"miniz_oxide",
"libz-rs-sys",
"miniz_oxide 0.8.8",
]
[[package]]
@@ -2808,9 +2839,11 @@ 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]]
@@ -3941,10 +3974,11 @@ checksum = "f5d4a7da358eff58addd2877a45865158f0d78c911d43a5784ceb7bbf52833b0"
[[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",
]
@@ -4085,6 +4119,26 @@ dependencies = [
"windows-targets 0.52.6",
]
[[package]]
name = "liblzma"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66352d7a8ac12d4877b6e6ea5a9b7650ee094257dc40889955bea5bc5b08c1d0"
dependencies = [
"liblzma-sys",
]
[[package]]
name = "liblzma-sys"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5839bad90c3cc2e0b8c4ed8296b80e86040240f81d46b9c0e9bc8dd51ddd3af1"
dependencies = [
"cc",
"libc",
"pkg-config",
]
[[package]]
name = "libm"
version = "0.2.8"
@@ -4128,6 +4182,15 @@ version = "0.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "64804cc6a5042d4f05379909ba25b503ec04e2c082151d62122d5dcaa274b961"
[[package]]
name = "libz-rs-sys"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6489ca9bd760fe9642d7644e827b0c9add07df89857b0416ee15c1cc1a3b8c5a"
dependencies = [
"zlib-rs",
]
[[package]]
name = "linux-raw-sys"
version = "0.4.14"
@@ -4159,9 +4222,9 @@ dependencies = [
[[package]]
name = "log"
version = "0.4.22"
version = "0.4.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24"
checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94"
[[package]]
name = "loom"
@@ -4343,6 +4406,15 @@ dependencies = [
"simd-adler32",
]
[[package]]
name = "miniz_oxide"
version = "0.8.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3be647b768db090acb35d5ec5db2b0e1f1de11133ca123b9eacf5137868f892a"
dependencies = [
"adler2",
]
[[package]]
name = "mio"
version = "1.0.2"
@@ -5266,21 +5338,10 @@ checksum = "1a2a4764cc1f8d961d802af27193c6f4f0124bd0e76e8393cf818e18880f0524"
dependencies = [
"argon2",
"getrandom 0.2.15",
"password-hash 0.5.0",
"password-hash",
"rand_core 0.6.4",
]
[[package]]
name = "password-hash"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700"
dependencies = [
"base64ct",
"rand_core 0.6.4",
"subtle",
]
[[package]]
name = "password-hash"
version = "0.5.0"
@@ -5306,14 +5367,12 @@ checksum = "8835116a5c179084a830efb3adc117ab007512b535bc1a21c991d3b32a6b44dd"
[[package]]
name = "pbkdf2"
version = "0.11.0"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917"
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
dependencies = [
"digest",
"hmac",
"password-hash 0.4.2",
"sha2",
]
[[package]]
@@ -5663,7 +5722,7 @@ dependencies = [
"crc32fast",
"fdeflate",
"flate2",
"miniz_oxide",
"miniz_oxide 0.7.4",
]
[[package]]
@@ -8417,9 +8476,9 @@ dependencies = [
[[package]]
name = "time"
version = "0.3.36"
version = "0.3.41"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5dfd88e563464686c916c7e46e623e520ddc6d79fa6641390f2e3fa86e83e885"
checksum = "8a7619e19bc266e0f9c5e6686659d394bc57973859340060a69221e57dbc0c40"
dependencies = [
"deranged",
"itoa 1.0.11",
@@ -8434,15 +8493,15 @@ dependencies = [
[[package]]
name = "time-core"
version = "0.1.2"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3"
checksum = "c9e9a38711f559d9e3ce1cdb06dd7c5b8ea546bc90052da6d06bb76da74bb07c"
[[package]]
name = "time-macros"
version = "0.2.18"
version = "0.2.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f252a68540fde3a3877aeea552b832b40ab9a69e318efd078774a01ddee1ccf"
checksum = "3526739392ec93fd8b359c8e98514cb3e8e021beb4e5f597b00a0221f8ed8a49"
dependencies = [
"num-conv",
"time-core",
@@ -9290,24 +9349,24 @@ checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b"
[[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.87",
@@ -9328,9 +9387,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",
@@ -9338,9 +9397,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",
@@ -9351,9 +9410,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 = "wasm-streams"
@@ -10307,31 +10369,46 @@ dependencies = [
[[package]]
name = "zip"
version = "0.6.6"
version = "4.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
checksum = "153a6fff49d264c4babdcfa6b4d534747f520e56e8f0f384f3b808c4b64cc1fd"
dependencies = [
"aes",
"byteorder",
"arbitrary",
"bzip2",
"constant_time_eq",
"crc32fast",
"crossbeam-utils",
"deflate64",
"flate2",
"getrandom 0.3.2",
"hmac",
"indexmap 2.7.1",
"liblzma",
"memchr",
"pbkdf2",
"sha1",
"time",
"zstd 0.11.2+zstd.1.5.2",
"zeroize",
"zopfli",
"zstd",
]
[[package]]
name = "zstd"
version = "0.11.2+zstd.1.5.2"
name = "zlib-rs"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4"
checksum = "868b928d7949e09af2f6086dfc1e01936064cc7a819253bce650d4e2a2d63ba8"
[[package]]
name = "zopfli"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "edfc5ee405f504cd4984ecc6f14d02d55cfda60fa4b689434ef4102aae150cd7"
dependencies = [
"zstd-safe 5.0.2+zstd.1.5.2",
"bumpalo",
"crc32fast",
"log",
"simd-adler32",
]
[[package]]
@@ -10340,17 +10417,7 @@ version = "0.13.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcf2b778a664581e31e389454a7072dab1647606d44f7feea22cd5abb9c9f3f9"
dependencies = [
"zstd-safe 7.2.1",
]
[[package]]
name = "zstd-safe"
version = "5.0.2+zstd.1.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db"
dependencies = [
"libc",
"zstd-sys",
"zstd-safe",
]
[[package]]

View File

@@ -186,10 +186,7 @@ async-ringbuf = "0.3.1"
service-manager = { git = "https://github.com/chipsenkbeil/service-manager-rs.git", branch = "main" }
async-compression = { version = "0.4.17", default-features = false, features = [
"zstd",
"tokio",
] }
zstd = { version = "0.13" }
kcp-sys = { git = "https://github.com/EasyTier/kcp-sys" }
@@ -261,7 +258,7 @@ prost-reflect-build = { version = "0.14.0" }
[target.'cfg(windows)'.build-dependencies]
reqwest = { version = "0.12.12", features = ["blocking"] }
zip = "0.6.6"
zip = "4.0.0"
# enable thunk-rs when compiling for x86_64 or i686 windows
[target.x86_64-pc-windows-msvc.build-dependencies]

View File

@@ -1,5 +1,10 @@
use async_compression::tokio::write::{ZstdDecoder, ZstdEncoder};
use tokio::io::AsyncWriteExt;
use std::io::{Read, Write};
use dashmap::DashMap;
use std::cell::RefCell;
use zstd::stream::read::Decoder;
use zstd::stream::write::Encoder;
use zstd::zstd_safe::{CCtx, DCtx};
use zerocopy::{AsBytes as _, FromBytes as _};
@@ -29,17 +34,20 @@ impl DefaultCompressor {
data: &[u8],
compress_algo: CompressorAlgo,
) -> Result<Vec<u8>, Error> {
let buf = match compress_algo {
match compress_algo {
CompressorAlgo::ZstdDefault => {
let mut o = ZstdEncoder::new(Vec::new());
o.write_all(data).await?;
o.shutdown().await?;
o.into_inner()
let ret = CTX_MAP.with(|map_cell| {
let map = map_cell.borrow();
let mut ctx_entry = map.entry(compress_algo).or_default();
let writer = Vec::new();
let mut o = Encoder::with_context(writer, ctx_entry.value_mut());
o.write_all(data)?;
o.finish()
});
Ok(ret?)
}
CompressorAlgo::None => data.to_vec(),
};
Ok(buf)
CompressorAlgo::None => Ok(data.to_vec()),
}
}
pub async fn decompress_raw(
@@ -47,17 +55,17 @@ impl DefaultCompressor {
data: &[u8],
compress_algo: CompressorAlgo,
) -> Result<Vec<u8>, Error> {
let buf = match compress_algo {
CompressorAlgo::ZstdDefault => {
let mut o = ZstdDecoder::new(Vec::new());
o.write_all(data).await?;
o.shutdown().await?;
o.into_inner()
}
CompressorAlgo::None => data.to_vec(),
};
Ok(buf)
match compress_algo {
CompressorAlgo::ZstdDefault => DCTX_MAP.with(|map_cell| {
let map = map_cell.borrow();
let mut ctx_entry = map.entry(compress_algo).or_default();
let mut decoder = Decoder::with_context(data, ctx_entry.value_mut());
let mut output = Vec::new();
decoder.read_to_end(&mut output)?;
Ok(output)
}),
CompressorAlgo::None => Ok(data.to_vec()),
}
}
}
@@ -146,6 +154,11 @@ impl Compressor for DefaultCompressor {
}
}
thread_local! {
static CTX_MAP: RefCell<DashMap<CompressorAlgo, CCtx<'static>>> = RefCell::new(DashMap::new());
static DCTX_MAP: RefCell<DashMap<CompressorAlgo, DCtx<'static>>> = RefCell::new(DashMap::new());
}
#[cfg(test)]
pub mod tests {
use super::*;
@@ -158,10 +171,21 @@ pub mod tests {
let compressor = DefaultCompressor {};
println!(
"Uncompressed packet: {:?}, len: {}",
packet,
packet.payload_len()
);
compressor
.compress(&mut packet, CompressorAlgo::ZstdDefault)
.await
.unwrap();
println!(
"Compressed packet: {:?}, len: {}",
packet,
packet.payload_len()
);
assert_eq!(packet.peer_manager_header().unwrap().is_compressed(), true);
compressor.decompress(&mut packet).await.unwrap();

View File

@@ -4,6 +4,7 @@ use std::{
io::Write as _,
sync::{Arc, Mutex},
};
use time::util::refresh_tz;
use tokio::{task::JoinSet, time::timeout};
use tracing::Instrument;
@@ -24,9 +25,7 @@ pub mod stun_codec_ext;
pub fn get_logger_timer<F: time::formatting::Formattable>(
format: F,
) -> tracing_subscriber::fmt::time::OffsetTime<F> {
unsafe {
time::util::local_offset::set_soundness(time::util::local_offset::Soundness::Unsound)
};
refresh_tz();
let local_offset = time::UtcOffset::current_local_offset()
.unwrap_or(time::UtcOffset::from_whole_seconds(0).unwrap());
tracing_subscriber::fmt::time::OffsetTime::new(local_offset, format)

View File

@@ -234,7 +234,7 @@ pub struct AesGcmTail {
}
pub const AES_GCM_ENCRYPTION_RESERVED: usize = std::mem::size_of::<AesGcmTail>();
#[derive(AsBytes, FromZeroes, Clone, Debug, Copy)]
#[derive(AsBytes, FromZeroes, Clone, Debug, Copy, PartialEq, Hash, Eq)]
#[repr(u8)]
pub enum CompressorAlgo {
None = 0,