xet-core-structures 1.6.0

// The shard structure for the in memory querying

use std::collections::BTreeMap;
use std::io::{BufWriter, Read, Write};
use std::mem::size_of;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use tracing::debug;

use super::file_structs::*;
#[cfg(debug_assertions)]
use super::shard_file_handle::{MDBShardFile, new_shard_file_cache};
use super::shard_format::{MDBShardFileHeader, MDBShardInfo};
use super::streaming_shard::{process_shard_file_info_section, process_shard_xorb_info_section};
use super::utils::{shard_file_name, temp_shard_file_name};
use super::xorb_structs::*;
use crate::MerkleHashMap;
use crate::error::Result;
use crate::merklehash::{HashedWrite, MerkleHash};

#[allow(clippy::type_complexity)]
#[derive(Clone, Default, Debug)]
pub struct MDBInMemoryShard {
    pub xorb_content: BTreeMap<MerkleHash, Arc<MDBXorbInfo>>,
    pub file_content: BTreeMap<MerkleHash, MDBFileInfo>,
    pub chunk_hash_lookup: MerkleHashMap<(Arc<MDBXorbInfo>, u64)>,
    current_shard_file_size: u64,
}

impl MDBInMemoryShard {
    pub fn add_xorb_block(&mut self, xorb_block_contents: impl Into<Arc<MDBXorbInfo>>) -> Result<()> {
        let dest_content_v: Arc<MDBXorbInfo> = xorb_block_contents.into();

        self.xorb_content
            .insert(dest_content_v.metadata.xorb_hash, dest_content_v.clone());

        for (i, chunk) in dest_content_v.chunks.iter().enumerate() {
            self.chunk_hash_lookup
                .insert(chunk.chunk_hash, (dest_content_v.clone(), i as u64));
            self.current_shard_file_size += (size_of::<u64>() + 2 * size_of::<u32>()) as u64;
        }
        self.current_shard_file_size += dest_content_v.num_bytes();
        self.current_shard_file_size += (size_of::<u64>() + size_of::<u32>()) as u64;

        Ok(())
    }

    pub fn add_file_reconstruction_info(&mut self, file_info: MDBFileInfo) -> Result<()> {
        self.current_shard_file_size += file_info.num_bytes();
        self.current_shard_file_size += (size_of::<u64>() + size_of::<u32>()) as u64;

        self.file_content.insert(file_info.metadata.file_hash, file_info);

        Ok(())
    }

    pub fn union(&self, other: &Self) -> Result<Self> {
        let mut xorb_content = self.xorb_content.clone();
        other.xorb_content.iter().for_each(|(k, v)| {
            xorb_content.insert(*k, v.clone());
        });

        let mut file_content = self.file_content.clone();
        for (k, v) in &other.file_content {
            if let Some(mut old_v) = file_content.insert(*k, v.clone()) {
                // merge the contents to ensure we have all the information between
                // the two file contents (e.g. verification and metadata_ext)
                old_v.merge_from(v)?;
                file_content.insert(*k, old_v);
            };
        }

        let mut chunk_hash_lookup = self.chunk_hash_lookup.clone();
        other.chunk_hash_lookup.iter().for_each(|(k, v)| {
            chunk_hash_lookup.insert(*k, v.clone());
        });

        let mut s = Self {
            xorb_content,
            file_content,
            current_shard_file_size: 0,
            chunk_hash_lookup,
        };

        s.recalculate_shard_size();
        Ok(s)
    }

    pub fn recalculate_shard_size(&mut self) {
        // Calculate the size
        let mut num_bytes = 0u64;
        for xorb_block_contents in self.xorb_content.values() {
            num_bytes += xorb_block_contents.num_bytes();

            // The xorb lookup table
            num_bytes += (size_of::<u64>() + size_of::<u32>()) as u64;
        }

        for file_info in self.file_content.values() {
            num_bytes += file_info.num_bytes();
            num_bytes += (size_of::<u64>() + size_of::<u32>()) as u64;
        }

        num_bytes += ((size_of::<u64>() + 2 * size_of::<u32>()) * self.chunk_hash_lookup.len()) as u64;

        self.current_shard_file_size = num_bytes;
    }

    pub fn difference(&self, other: &Self) -> Result<Self> {
        let mut s = Self {
            xorb_content: other
                .xorb_content
                .iter()
                .filter(|(k, _)| !self.xorb_content.contains_key(k))
                .map(|(k, v)| (*k, v.clone()))
                .collect(),
            file_content: other
                .file_content
                .iter()
                .filter(|(k, _)| !self.file_content.contains_key(k))
                .map(|(k, v)| (*k, v.clone()))
                .collect(),
            chunk_hash_lookup: other
                .chunk_hash_lookup
                .iter()
                .filter(|(k, _)| !self.chunk_hash_lookup.contains_key(k))
                .map(|(k, v)| (*k, v.clone()))
                .collect(),
            current_shard_file_size: 0,
        };
        s.recalculate_shard_size();
        Ok(s)
    }

    /// Given a file pointer, returns the information needed to reconstruct the file.
    /// Returns the file info if the file hash was found, and None otherwise.
    pub fn get_file_reconstruction_info(&self, file_hash: &MerkleHash) -> Option<MDBFileInfo> {
        if let Some(mdb_file) = self.file_content.get(file_hash) {
            return Some(mdb_file.clone());
        }

        None
    }

    pub fn chunk_hash_dedup_query(&self, query_hashes: &[MerkleHash]) -> Option<(usize, FileDataSequenceEntry)> {
        if query_hashes.is_empty() {
            return None;
        }

        let (chunk_ref, chunk_index_start) = self.chunk_hash_lookup.get(&query_hashes[0])?;

        let chunk_index_start = *chunk_index_start as usize;

        let mut query_idx = 0;

        loop {
            if chunk_index_start + query_idx >= chunk_ref.chunks.len() {
                break;
            }
            if query_idx >= query_hashes.len()
                || chunk_ref.chunks[chunk_index_start + query_idx].chunk_hash != query_hashes[query_idx]
            {
                break;
            }
            query_idx += 1;
        }

        Some((
            query_idx,
            FileDataSequenceEntry::from_xorb_entries(
                &chunk_ref.metadata,
                &chunk_ref.chunks[chunk_index_start..(chunk_index_start + query_idx)],
                chunk_index_start,
                chunk_index_start + query_idx,
            ),
        ))
    }

    pub fn num_xorb_entries(&self) -> usize {
        self.xorb_content.len()
    }

    pub fn num_file_entries(&self) -> usize {
        self.file_content.len()
    }

    pub fn stored_bytes_on_disk(&self) -> u64 {
        self.xorb_content
            .iter()
            .fold(0u64, |acc, (_, xorb)| acc + xorb.metadata.num_bytes_on_disk as u64)
    }

    pub fn materialized_bytes(&self) -> u64 {
        self.file_content.iter().fold(0u64, |acc, (_, file)| {
            acc + file
                .segments
                .iter()
                .fold(0u64, |acc, entry| acc + entry.unpacked_segment_bytes as u64)
        })
    }

    pub fn stored_bytes(&self) -> u64 {
        self.xorb_content
            .iter()
            .fold(0u64, |acc, (_, xorb)| acc + xorb.metadata.num_bytes_in_xorb as u64)
    }

    pub fn is_empty(&self) -> bool {
        self.xorb_content.is_empty() && self.file_content.is_empty()
    }

    /// Returns the number of bytes required
    pub fn shard_file_size(&self) -> u64 {
        self.current_shard_file_size + MDBShardInfo::non_content_byte_size()
    }

    /// Writes the shard out to a file.
    pub fn write_to_temp_shard_file(&self, temp_file_name: &Path, expiration: Option<Duration>) -> Result<MerkleHash> {
        let mut hashed_write; // Need to access after file is closed.

        {
            // Scoped so that file is closed and flushed before name is changed.

            let out_file = std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open(temp_file_name)?;

            hashed_write = HashedWrite::new(out_file);

            let mut buf_write = BufWriter::new(&mut hashed_write);

            // Ask for write access, as we'll flush this at the end
            MDBShardInfo::serialize_from(&mut buf_write, self, expiration)?;

            debug!("Writing out in-memory shard to {temp_file_name:?}.");

            buf_write.flush()?;
        }

        // Get the hash
        hashed_write.flush()?;
        let shard_hash = hashed_write.hash();

        Ok(shard_hash)
    }
    pub fn write_to_directory(&self, directory: &Path, expiration: Option<Duration>) -> Result<PathBuf> {
        // First, create a temporary shard structure in that directory.
        let temp_file_name = directory.join(temp_shard_file_name());

        let shard_hash = self.write_to_temp_shard_file(&temp_file_name, expiration)?;

        let full_file_name = directory.join(shard_file_name(&shard_hash));

        std::fs::rename(&temp_file_name, &full_file_name)?;

        debug!("Wrote out in-memory shard to {full_file_name:?}.");

        #[cfg(debug_assertions)]
        {
            let cache = new_shard_file_cache();
            let shard_file = MDBShardFile::load_from_file(&full_file_name, &cache)?;
            shard_file.verify_shard_integrity();
        }

        Ok(full_file_name)
    }

    /// Serializes the shard to a vector of bytes.
    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
        MDBShardInfo::serialize_from(&mut buf, self, None)?;
        Ok(buf)
    }

    /// Deserializes a shard from a reader, populating xorb / file / chunk lookup
    /// tables in memory. The reader is consumed up to the end of the xorb section
    /// (the footer is not read or validated). Used by the wasm target to ingest
    /// global-dedup shards fetched from CAS without going through disk.
    pub fn from_reader<R: Read>(reader: &mut R) -> Result<Self> {
        let _ = MDBShardFileHeader::deserialize(reader)?;

        let mut shard = Self::default();

        process_shard_file_info_section(reader, |view| shard.add_file_reconstruction_info(MDBFileInfo::from(&view)))?;

        process_shard_xorb_info_section(reader, |view| shard.add_xorb_block(MDBXorbInfo::from(&view)))?;

        Ok(shard)
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;
    use crate::metadata_shard::shard_format::test_routines::{gen_specific_shard, simple_hash};

    #[test]
    fn from_reader_roundtrips_dedup_lookups() {
        let xorb_nodes: &[(u64, &[(u64, u32)])] = &[(101, &[(1, 100), (2, 200), (3, 300)]), (102, &[(4, 50), (5, 75)])];
        let file_nodes: &[(u64, &[(u64, (u32, u32))])] = &[(201, &[(101, (0, 600))]), (202, &[(102, (0, 125))])];

        let original = gen_specific_shard(xorb_nodes, file_nodes, None, None).unwrap();
        let bytes = original.to_bytes().unwrap();
        let reloaded = MDBInMemoryShard::from_reader(&mut Cursor::new(&bytes)).unwrap();

        assert_eq!(reloaded.num_xorb_entries(), original.num_xorb_entries());
        assert_eq!(reloaded.num_file_entries(), original.num_file_entries());
        assert_eq!(reloaded.chunk_hash_lookup.len(), original.chunk_hash_lookup.len());

        for xorb_hash in original.xorb_content.keys() {
            assert!(reloaded.xorb_content.contains_key(xorb_hash), "missing xorb {xorb_hash}");
        }
        for file_hash in original.file_content.keys() {
            assert!(reloaded.file_content.contains_key(file_hash), "missing file {file_hash}");
        }

        // Dedup queries must succeed against the reloaded shard.
        let query = [simple_hash(1), simple_hash(2), simple_hash(3)];
        let (n, fse) = reloaded.chunk_hash_dedup_query(&query).expect("chunk lookup");
        assert_eq!(n, 3);
        assert_eq!(fse.xorb_hash, simple_hash(101));
    }

    // Mirrors the wasm global-dedup path: production CAS returns HMAC-keyed
    // shards, so lookups against a shard ingested via from_reader only match
    // when the query hashes are keyed with the footer's HMAC key.
    #[test]
    fn from_reader_keyed_shard_dedup_requires_hmac() {
        let xorb_nodes: &[(u64, &[(u64, u32)])] = &[(101, &[(1, 100), (2, 200), (3, 300)])];
        let file_nodes: &[(u64, &[(u64, (u32, u32))])] = &[(201, &[(101, (0, 600))])];

        let original = gen_specific_shard(xorb_nodes, file_nodes, None, None).unwrap();
        let bytes = original.to_bytes().unwrap();

        let key = simple_hash(42);
        let mut keyed_bytes = Vec::new();
        MDBShardInfo::export_as_keyed_shard_streaming(
            &mut Cursor::new(&bytes),
            &mut keyed_bytes,
            key,
            std::time::Duration::from_secs(3600),
            true,
            true,
            true,
        )
        .unwrap();

        let mut reader = Cursor::new(&keyed_bytes);
        let info = MDBShardInfo::load_from_reader(&mut reader).unwrap();
        assert_eq!(info.chunk_hmac_key(), Some(key));

        reader.set_position(0);
        let reloaded = MDBInMemoryShard::from_reader(&mut reader).unwrap();

        let raw_query = [simple_hash(1), simple_hash(2), simple_hash(3)];
        assert!(reloaded.chunk_hash_dedup_query(&raw_query).is_none());

        let keyed_query: Vec<_> = raw_query.iter().map(|h| h.hmac(key)).collect();
        let (n, fse) = reloaded.chunk_hash_dedup_query(&keyed_query).expect("keyed chunk lookup");
        assert_eq!(n, 3);
        assert_eq!(fse.xorb_hash, simple_hash(101));
    }
}