|
|
//! Helpers around revlog compression
|
|
|
|
|
|
use std::cell::RefCell;
|
|
|
use std::collections::HashSet;
|
|
|
use std::io::Read;
|
|
|
|
|
|
use flate2::bufread::ZlibEncoder;
|
|
|
use flate2::read::ZlibDecoder;
|
|
|
|
|
|
use crate::config::Config;
|
|
|
use crate::errors::HgError;
|
|
|
use crate::exit_codes;
|
|
|
|
|
|
use super::corrupted;
|
|
|
use super::RevlogError;
|
|
|
|
|
|
/// Header byte used to identify ZSTD-compressed data
|
|
|
pub const ZSTD_BYTE: u8 = b'\x28';
|
|
|
/// Header byte used to identify Zlib-compressed data
|
|
|
pub const ZLIB_BYTE: u8 = b'x';
|
|
|
|
|
|
const ZSTD_DEFAULT_LEVEL: u8 = 3;
|
|
|
const ZLIB_DEFAULT_LEVEL: u8 = 6;
|
|
|
/// The length of data below which we don't even try to compress it when using
|
|
|
/// Zstandard.
|
|
|
const MINIMUM_LENGTH_ZSTD: usize = 50;
|
|
|
/// The length of data below which we don't even try to compress it when using
|
|
|
/// Zlib.
|
|
|
const MINIMUM_LENGTH_ZLIB: usize = 44;
|
|
|
|
|
|
/// Defines the available compression engines and their options.
|
|
|
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
|
|
pub enum CompressionConfig {
|
|
|
Zlib {
|
|
|
/// Between 0 and 9 included
|
|
|
level: u8,
|
|
|
},
|
|
|
Zstd {
|
|
|
/// Between 0 and 22 included
|
|
|
level: u8,
|
|
|
/// Never used in practice for now
|
|
|
threads: u8,
|
|
|
},
|
|
|
/// No compression is performed
|
|
|
None,
|
|
|
}
|
|
|
|
|
|
impl CompressionConfig {
|
|
|
pub fn new(
|
|
|
config: &Config,
|
|
|
requirements: &HashSet<String>,
|
|
|
) -> Result<Self, HgError> {
|
|
|
let mut new = Self::default();
|
|
|
|
|
|
let zlib_level = config.get_u32(b"storage", b"revlog.zlib.level")?;
|
|
|
let zstd_level = config.get_u32(b"storage", b"revlog.zstd.level")?;
|
|
|
|
|
|
for requirement in requirements {
|
|
|
if requirement.starts_with("revlog-compression-")
|
|
|
|| requirement.starts_with("exp-compression-")
|
|
|
{
|
|
|
let split = &mut requirement.splitn(3, '-');
|
|
|
split.next();
|
|
|
split.next();
|
|
|
new = match split.next().unwrap() {
|
|
|
"zstd" => CompressionConfig::zstd(zstd_level)?,
|
|
|
e => {
|
|
|
return Err(HgError::UnsupportedFeature(format!(
|
|
|
"Unsupported compression engine '{e}'"
|
|
|
)))
|
|
|
}
|
|
|
};
|
|
|
}
|
|
|
}
|
|
|
if let Some(level) = zlib_level {
|
|
|
if matches!(new, CompressionConfig::Zlib { .. }) {
|
|
|
new.set_level(level as usize)?;
|
|
|
}
|
|
|
}
|
|
|
Ok(new)
|
|
|
}
|
|
|
|
|
|
/// Sets the level of the current compression engine
|
|
|
pub fn set_level(&mut self, new_level: usize) -> Result<(), HgError> {
|
|
|
match self {
|
|
|
CompressionConfig::Zlib { level } => {
|
|
|
if new_level > 9 {
|
|
|
return Err(HgError::abort(
|
|
|
format!(
|
|
|
"invalid compression zlib compression level {}, \
|
|
|
expected between 0 and 9 included",
|
|
|
new_level
|
|
|
),
|
|
|
exit_codes::ABORT,
|
|
|
None,
|
|
|
));
|
|
|
}
|
|
|
*level = new_level as u8;
|
|
|
}
|
|
|
CompressionConfig::Zstd { level, .. } => {
|
|
|
if new_level > 22 {
|
|
|
return Err(HgError::abort(
|
|
|
format!(
|
|
|
"invalid compression zstd compression level {}, \
|
|
|
expected between 0 and 22 included",
|
|
|
new_level
|
|
|
),
|
|
|
exit_codes::ABORT,
|
|
|
None,
|
|
|
));
|
|
|
}
|
|
|
*level = new_level as u8;
|
|
|
}
|
|
|
CompressionConfig::None => {}
|
|
|
}
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
/// Return a ZSTD compression config
|
|
|
pub fn zstd(
|
|
|
zstd_level: Option<u32>,
|
|
|
) -> Result<CompressionConfig, HgError> {
|
|
|
let mut engine = CompressionConfig::Zstd {
|
|
|
level: ZSTD_DEFAULT_LEVEL,
|
|
|
threads: 0,
|
|
|
};
|
|
|
if let Some(level) = zstd_level {
|
|
|
engine.set_level(level as usize)?;
|
|
|
}
|
|
|
Ok(engine)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
impl Default for CompressionConfig {
|
|
|
fn default() -> Self {
|
|
|
Self::Zlib {
|
|
|
level: ZLIB_DEFAULT_LEVEL,
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
/// A high-level trait to define compressors that should be able to compress
|
|
|
/// and decompress arbitrary bytes.
|
|
|
pub trait Compressor {
|
|
|
/// Returns a new [`Vec`] with the compressed data.
|
|
|
/// Should return `Ok(None)` if compression does not apply (e.g. too small)
|
|
|
fn compress(
|
|
|
&mut self,
|
|
|
data: &[u8],
|
|
|
) -> Result<Option<Vec<u8>>, RevlogError>;
|
|
|
/// Returns a new [`Vec`] with the decompressed data.
|
|
|
fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, RevlogError>;
|
|
|
}
|
|
|
|
|
|
/// A compressor that does nothing (useful in tests)
|
|
|
pub struct NoneCompressor;
|
|
|
|
|
|
impl Compressor for NoneCompressor {
|
|
|
fn compress(
|
|
|
&mut self,
|
|
|
_data: &[u8],
|
|
|
) -> Result<Option<Vec<u8>>, RevlogError> {
|
|
|
Ok(None)
|
|
|
}
|
|
|
|
|
|
fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, RevlogError> {
|
|
|
Ok(data.to_owned())
|
|
|
}
|
|
|
}
|
|
|
|
|
|
/// A compressor for Zstandard
|
|
|
pub struct ZstdCompressor {
|
|
|
/// Level of compression to use
|
|
|
level: u8,
|
|
|
/// How many threads are used (not implemented yet)
|
|
|
threads: u8,
|
|
|
/// The underlying zstd compressor
|
|
|
compressor: zstd::bulk::Compressor<'static>,
|
|
|
}
|
|
|
|
|
|
impl ZstdCompressor {
|
|
|
pub fn new(level: u8, threads: u8) -> Self {
|
|
|
Self {
|
|
|
level,
|
|
|
threads,
|
|
|
compressor: zstd::bulk::Compressor::new(level.into())
|
|
|
.expect("invalid zstd arguments"),
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
impl Compressor for ZstdCompressor {
|
|
|
fn compress(
|
|
|
&mut self,
|
|
|
data: &[u8],
|
|
|
) -> Result<Option<Vec<u8>>, RevlogError> {
|
|
|
if self.threads != 0 {
|
|
|
// TODO use a zstd builder + zstd cargo feature to support this
|
|
|
unimplemented!("zstd parallel compression is not implemented");
|
|
|
}
|
|
|
if data.len() < MINIMUM_LENGTH_ZSTD {
|
|
|
return Ok(None);
|
|
|
}
|
|
|
let level = self.level as i32;
|
|
|
if data.len() <= 1000000 {
|
|
|
let compressed = self.compressor.compress(data).map_err(|e| {
|
|
|
corrupted(format!("revlog compress error: {}", e))
|
|
|
})?;
|
|
|
Ok(if compressed.len() < data.len() {
|
|
|
Some(compressed)
|
|
|
} else {
|
|
|
None
|
|
|
})
|
|
|
} else {
|
|
|
Ok(Some(zstd::stream::encode_all(data, level).map_err(
|
|
|
|e| corrupted(format!("revlog compress error: {}", e)),
|
|
|
)?))
|
|
|
}
|
|
|
}
|
|
|
|
|
|
fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, RevlogError> {
|
|
|
zstd::stream::decode_all(data).map_err(|e| {
|
|
|
corrupted(format!("revlog decompress error: {}", e)).into()
|
|
|
})
|
|
|
}
|
|
|
}
|
|
|
|
|
|
/// A compressor for Zlib
|
|
|
pub struct ZlibCompressor {
|
|
|
/// Level of compression to use
|
|
|
level: flate2::Compression,
|
|
|
}
|
|
|
|
|
|
impl ZlibCompressor {
|
|
|
pub fn new(level: u8) -> Self {
|
|
|
Self {
|
|
|
level: flate2::Compression::new(level.into()),
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
impl Compressor for ZlibCompressor {
|
|
|
fn compress(
|
|
|
&mut self,
|
|
|
data: &[u8],
|
|
|
) -> Result<Option<Vec<u8>>, RevlogError> {
|
|
|
assert!(!data.is_empty());
|
|
|
if data.len() < MINIMUM_LENGTH_ZLIB {
|
|
|
return Ok(None);
|
|
|
}
|
|
|
let mut buf = Vec::with_capacity(data.len());
|
|
|
ZlibEncoder::new(data, self.level)
|
|
|
.read_to_end(&mut buf)
|
|
|
.map_err(|e| corrupted(format!("revlog compress error: {}", e)))?;
|
|
|
|
|
|
Ok(if buf.len() < data.len() {
|
|
|
buf.shrink_to_fit();
|
|
|
Some(buf)
|
|
|
} else {
|
|
|
None
|
|
|
})
|
|
|
}
|
|
|
|
|
|
fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, RevlogError> {
|
|
|
let mut decoder = ZlibDecoder::new(data);
|
|
|
// TODO reuse the allocation somehow?
|
|
|
let mut buf = vec![];
|
|
|
decoder.read_to_end(&mut buf).map_err(|e| {
|
|
|
corrupted(format!("revlog decompress error: {}", e))
|
|
|
})?;
|
|
|
Ok(buf)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
thread_local! {
|
|
|
// seems fine to [unwrap] here: this can only fail due to memory allocation
|
|
|
// failing, and it's normal for that to cause panic.
|
|
|
static ZSTD_DECODER : RefCell<zstd::bulk::Decompressor<'static>> =
|
|
|
RefCell::new(zstd::bulk::Decompressor::new().ok().unwrap());
|
|
|
}
|
|
|
|
|
|
/// Util to wrap the reuse of a zstd decoder while controlling its buffer size.
|
|
|
fn zstd_decompress_to_buffer(
|
|
|
bytes: &[u8],
|
|
|
buf: &mut Vec<u8>,
|
|
|
) -> Result<usize, std::io::Error> {
|
|
|
ZSTD_DECODER
|
|
|
.with(|decoder| decoder.borrow_mut().decompress_to_buffer(bytes, buf))
|
|
|
}
|
|
|
|
|
|
/// Specialized revlog decompression to use less memory for deltas while
|
|
|
/// keeping performance acceptable.
|
|
|
pub(super) fn uncompressed_zstd_data(
|
|
|
bytes: &[u8],
|
|
|
is_delta: bool,
|
|
|
uncompressed_len: i32,
|
|
|
) -> Result<Vec<u8>, HgError> {
|
|
|
let cap = uncompressed_len.max(0) as usize;
|
|
|
if is_delta {
|
|
|
// [cap] is usually an over-estimate of the space needed because
|
|
|
// it's the length of delta-decoded data, but we're interested
|
|
|
// in the size of the delta.
|
|
|
// This means we have to [shrink_to_fit] to avoid holding on
|
|
|
// to a large chunk of memory, but it also means we must have a
|
|
|
// fallback branch, for the case when the delta is longer than
|
|
|
// the original data (surprisingly, this does happen in practice)
|
|
|
let mut buf = Vec::with_capacity(cap);
|
|
|
match zstd_decompress_to_buffer(bytes, &mut buf) {
|
|
|
Ok(_) => buf.shrink_to_fit(),
|
|
|
Err(_) => {
|
|
|
buf.clear();
|
|
|
zstd::stream::copy_decode(bytes, &mut buf)
|
|
|
.map_err(|e| corrupted(e.to_string()))?;
|
|
|
}
|
|
|
};
|
|
|
Ok(buf)
|
|
|
} else {
|
|
|
let mut buf = Vec::with_capacity(cap);
|
|
|
let len = zstd_decompress_to_buffer(bytes, &mut buf)
|
|
|
.map_err(|e| corrupted(e.to_string()))?;
|
|
|
if len != uncompressed_len as usize {
|
|
|
Err(corrupted("uncompressed length does not match"))
|
|
|
} else {
|
|
|
Ok(buf)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
#[cfg(test)]
|
|
|
mod tests {
|
|
|
use super::*;
|
|
|
|
|
|
const LARGE_TEXT: &[u8] = b"
|
|
|
Patali Dirapata, Cromda Cromda Ripalo, Pata Pata, Ko Ko Ko
|
|
|
Bokoro Dipoulito, Rondi Rondi Pepino, Pata Pata, Ko Ko Ko
|
|
|
Emana Karassoli, Loucra Loucra Nonponto, Pata Pata, Ko Ko Ko.";
|
|
|
|
|
|
#[test]
|
|
|
fn test_zlib_compressor() {
|
|
|
// Can return `Ok(None)`
|
|
|
let mut compressor = ZlibCompressor::new(1);
|
|
|
assert_eq!(compressor.compress(b"too small").unwrap(), None);
|
|
|
|
|
|
// Compression returns bytes
|
|
|
let compressed_with_1 =
|
|
|
compressor.compress(LARGE_TEXT).unwrap().unwrap();
|
|
|
assert!(compressed_with_1.len() < LARGE_TEXT.len());
|
|
|
// Round trip works
|
|
|
assert_eq!(
|
|
|
compressor.decompress(&compressed_with_1).unwrap(),
|
|
|
LARGE_TEXT
|
|
|
);
|
|
|
|
|
|
// Compression levels mean something
|
|
|
let mut compressor = ZlibCompressor::new(9);
|
|
|
// Compression returns bytes
|
|
|
let compressed = compressor.compress(LARGE_TEXT).unwrap().unwrap();
|
|
|
assert!(compressed.len() < compressed_with_1.len());
|
|
|
}
|
|
|
|
|
|
#[test]
|
|
|
fn test_zstd_compressor() {
|
|
|
// Can return `Ok(None)`
|
|
|
let mut compressor = ZstdCompressor::new(1, 0);
|
|
|
assert_eq!(compressor.compress(b"too small").unwrap(), None);
|
|
|
|
|
|
// Compression returns bytes
|
|
|
let compressed_with_1 =
|
|
|
compressor.compress(LARGE_TEXT).unwrap().unwrap();
|
|
|
assert!(compressed_with_1.len() < LARGE_TEXT.len());
|
|
|
// Round trip works
|
|
|
assert_eq!(
|
|
|
compressor.decompress(&compressed_with_1).unwrap(),
|
|
|
LARGE_TEXT
|
|
|
);
|
|
|
|
|
|
// Compression levels mean something
|
|
|
let mut compressor = ZstdCompressor::new(22, 0);
|
|
|
// Compression returns bytes
|
|
|
let compressed = compressor.compress(LARGE_TEXT).unwrap().unwrap();
|
|
|
assert!(compressed.len() < compressed_with_1.len());
|
|
|
}
|
|
|
}
|
|
|
|
