annotate rust/hg-core/src/revlog/compression.rs @ 52316:a820a7a1fce0 default tip

setup: require TLS 1.2 support from the Python interpreter (BC)

Before it was optional, and either 1.1 or 1.2 was sufficient.  Now that the
default minimum is 1.2, it needs to be present to work out of the box.

The code here is more convoluted than the corresponding checks in `sslutil.py`,
but I'm leaving it alone because it can all be simplified when py38 is dropped.

//! 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: Send {
    /// 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());
    }
}