gc: a trained zstd dictionary per pack (refs #41) a17fc622ca43

/// Compressed with the pack's trained dictionary (`pack-ID.dict`), which a

/// full repack makes from a sample of the objects it packs: small files

/// share most of their vocabulary, which each one alone cannot use.

const FLAG_DICT: u8 = 2;

/// The size of a pack's dictionary, and the least a repack must pack for

/// one to pay.

const DICT_BYTES: usize = 110 << 10;

const DICT_MIN_OBJECTS: usize = 1000;

/// Objects up to this size are compressed with the dictionary.

const DICT_OBJECT_MAX: usize = 64 << 10;

    /// The pack's dictionary, prepared once, when it has one.

    dict: Option<std::sync::Arc<zstd::dict::DecoderDictionary<'static>>>,

impl std::fmt::Debug for PackIndex {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("PackIndex")

            .field("path", &self.path)

            .field("objects", &self.refs.len())

            .field("dict", &self.dict.is_some())

            .finish()

    }

}

        let dict = match fs::read(pack_path.with_extension("dict")) {

            Ok(d) => Some(std::sync::Arc::new(zstd::dict::DecoderDictionary::copy(&d))),

            Err(e) if e.kind() == io::ErrorKind::NotFound => None,

            Err(e) => return Err(e.into()),

        };

            dict,

        let body = if e.flags & FLAG_DICT != 0 {

            let dict = self.dict.as_ref().ok_or_else(|| {

                StoreError::Corrupt(*r, "packed with a dictionary the pack lacks".into())

            })?;

            zstd::bulk::Decompressor::with_prepared_dictionary(dict)

                .and_then(|mut d| d.decompress(&stored, e.raw_len as usize))

                .map_err(|err| StoreError::Corrupt(*r, format!("zstd: {err}")))?

        } else if e.flags & FLAG_ZSTD != 0 {

                remove_pack_files(p)?;

        // A full repack of enough objects trains a dictionary for them.

        let dict = if level == REPACK_ZSTD_LEVEL && refs.len() >= DICT_MIN_OBJECTS {

            let _t = crate::timing::phase("pack: train a dictionary");

            self.train_dictionary(&refs)

        } else {

            None

        };

        let encoder = dict

            .as_ref()

            .map(|d| zstd::dict::EncoderDictionary::copy(d, level));

                        .map_init(

                            || {

                                encoder.as_ref().and_then(|e| {

                                    zstd::bulk::Compressor::with_prepared_dictionary(e).ok()

                                })

                            },

                            |with_dict, (r, ty, body, raw)| {

                                let packed = if raw {

                                    None

                                } else {

                                    // Small objects with the dictionary;

                                    // a large one has context enough.

                                    let dicted = with_dict

                                        .as_mut()

                                        .filter(|_| body.len() <= DICT_OBJECT_MAX)

                                        .and_then(|c| c.compress(&body).ok())

                                        .map(|c| (FLAG_ZSTD | FLAG_DICT, c));

                                    dicted

                                        .or_else(|| {

                                            zstd::bulk::compress(&body, level)

                                                .ok()

                                                .map(|c| (FLAG_ZSTD, c))

                                        })

                                        .filter(|(_, c)| c.len() < body.len())

                                };

                                (r, ty, body, packed)

                            },

                        )

                    Some(c) => c,

        }

        // The dictionary before the pack that needs it becomes visible.

        if let Some(d) = &dict {

            crate::durable::write_atomic(&self.dir.join(format!("pack-{id}.dict")), d)?;

            remove_pack_files(p)?;

    /// A dictionary for these objects, trained on a sample of them (at

    /// most 10 000, 32 MB); `None` when training fails.

    fn train_dictionary(&self, refs: &[ObjectRef]) -> Option<Vec<u8>> {

        let step = (refs.len() / 10_000).max(1);

        let mut samples: Vec<Vec<u8>> = Vec::new();

        let mut total = 0usize;

        for r in refs.iter().step_by(step) {

            if total > 32 << 20 {

                break;

            }

            let Ok((ty, body)) = self.get(r) else {

                continue;

            };

            if matches!(

                ty,

                ObjectType::Chunk | ObjectType::Blob | ObjectType::Change

            ) && !body.is_empty()

                && body.len() < 1 << 20

            {

                total += body.len();

                samples.push(body);

            }

        }

        zstd::dict::from_samples(&samples, DICT_BYTES).ok()

    }

                        dict: None,

    }

    #[test]

    fn a_full_repack_trains_a_dictionary_and_reads_back() {

        let dir = tempfile::tempdir().unwrap();

        let s = Store::open(dir.path()).unwrap();

        // Many small, alike files: what a dictionary is for.

        let bodies: Vec<Vec<u8>> = (0..1500)

            .map(|i| {

                format!(

                    "// SPDX-License-Identifier: GPL-2.0\n#include <linux/module.h>\n\

                     static int probe_{i}(struct device *dev)\n{{\n\treturn {i};\n}}\n"

                )

                .into_bytes()

            })

            .collect();

        let refs: Vec<ObjectRef> = bodies

            .iter()

            .map(|b| s.put(ObjectType::Chunk, b).unwrap())

            .collect();

        s.repack_all(REPACK_ZSTD_LEVEL).unwrap();

        let dicts = fs::read_dir(dir.path())

            .unwrap()

            .flatten()

            .filter(|e| e.path().extension().is_some_and(|x| x == "dict"))

            .count();

        assert_eq!(dicts, 1, "the pack has a dictionary");

        let with_dict = refs

            .iter()

            .filter(|r| s.find_packed(r).unwrap().1.flags & FLAG_DICT != 0)

            .count();

        assert!(with_dict > 1000, "most use it: {with_dict}");

        let again = Store::open(dir.path()).unwrap();

        for (r, b) in refs.iter().zip(&bodies) {

            assert_eq!(&again.get(r).unwrap().1, b);

        }

        // A repack after that replaces the pack and its dictionary.

        again.repack_all(REPACK_ZSTD_LEVEL).unwrap();

        let files: Vec<String> = fs::read_dir(dir.path())

            .unwrap()

            .flatten()

            .map(|e| e.file_name().to_string_lossy().into_owned())

            .filter(|n| n.starts_with("pack-"))

            .collect();

        assert_eq!(files.len(), 3, "one pack, index and dictionary: {files:?}");

        assert_eq!(&again.get(&refs[7]).unwrap().1, &bodies[7]);

    }

}

/// Remove a pack with its index and dictionary.

fn remove_pack_files(pack: &Path) -> io::Result<()> {

    fs::remove_file(pack)?;

    fs::remove_file(pack.with_extension("idx"))?;

    match fs::remove_file(pack.with_extension("dict")) {

        Err(e) if e.kind() != io::ErrorKind::NotFound => Err(e),

        _ => Ok(()),