Shard utilities needed for GC pass and server-side xorb rewriting. (#532)

This PR adds a utility that rewrites a shard to include only the
relevant xorb information, dropping unreferenced file information.

In addition, to preserve the global dedup tracking information
associated with the files, this PR also adds a backwards-compatible flag
to the chunk metadata that marks a specific chunk as global dedup
eligible. This allows the global dedup information to be tracked
independently of the file metadata.

Author: hoytak

Cargo.lock

@@ -26,6 +26,7 @@ tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
+more-asserts = { workspace = true }
 
 [target.'cfg(target_family = "wasm")'.dependencies]
 uuid = { workspace = true, features = ["v4", "js"] }

hf_xet/Cargo.lock

@@ -5,8 +5,12 @@ use bytes::Bytes;
 use merklehash::MerkleHash;
 use utils::serialization_utils::*;
 
+use crate::hash_is_global_dedup_eligible;
+
 pub const MDB_DEFAULT_CAS_FLAG: u32 = 0;
 
+pub const MDB_CHUNK_WITH_GLOBAL_DEDUP_FLAG: u32 = 1 << 31;
+
 /// Each CAS consists of a CASChunkSequenceHeader following
 /// a sequence of CASChunkSequenceEntry.
 
@@ -89,7 +93,8 @@ pub struct CASChunkSequenceEntry {
     pub chunk_hash: MerkleHash,
     pub chunk_byte_range_start: u32,
     pub unpacked_segment_bytes: u32,
-    pub _unused: u64,
+    pub flags: u32,
+    pub _unused: u32,
 }
 
 impl CASChunkSequenceEntry {
@@ -106,13 +111,31 @@ impl CASChunkSequenceEntry {
             chunk_hash,
             unpacked_segment_bytes: unpacked_segment_bytes.try_into().unwrap(),
             chunk_byte_range_start: chunk_byte_range_start.try_into().unwrap(),
-            #[cfg(test)]
-            _unused: 216944691646848u64,
-            #[cfg(not(test))]
+            flags: 0,
             _unused: 0,
         }
     }
 
+    /// Mark this chunk as a candidate for population in the global dedup table.
+    pub fn with_global_dedup_flag(self, is_global_dedup_chunk: bool) -> Self {
+        if is_global_dedup_chunk {
+            Self {
+                flags: self.flags | MDB_CHUNK_WITH_GLOBAL_DEDUP_FLAG,
+                ..self
+            }
+        } else {
+            Self {
+                flags: self.flags & !MDB_CHUNK_WITH_GLOBAL_DEDUP_FLAG,
+                ..self
+            }
+        }
+    }
+
+    // Is this chunk elegible for a global dedup query?
+    pub fn is_global_dedup_eligible(&self) -> bool {
+        (self.flags & MDB_CHUNK_WITH_GLOBAL_DEDUP_FLAG) != 0 || hash_is_global_dedup_eligible(&self.chunk_hash)
+    }
+
     pub fn serialize<W: Write>(&self, writer: &mut W) -> Result<usize, std::io::Error> {
         let mut buf = [0u8; size_of::<Self>()];
         {
@@ -122,7 +145,8 @@ impl CASChunkSequenceEntry {
             write_hash(writer, &self.chunk_hash)?;
             write_u32(writer, self.chunk_byte_range_start)?;
             write_u32(writer, self.unpacked_segment_bytes)?;
-            write_u64(writer, self._unused)?;
+            write_u32(writer, self.flags)?;
+            write_u32(writer, self._unused)?;
         }
 
         writer.write_all(&buf[..])?;
@@ -140,7 +164,8 @@ impl CASChunkSequenceEntry {
             chunk_hash: read_hash(reader)?,
             chunk_byte_range_start: read_u32(reader)?,
             unpacked_segment_bytes: read_u32(reader)?,
-            _unused: read_u64(reader)?,
+            flags: read_u32(reader)?,
+            _unused: read_u32(reader)?,
         })
     }
 }
@@ -253,4 +278,21 @@ impl MDBCASInfoView {
         writer.write_all(&self.data[..n_bytes])?;
         Ok(n_bytes)
     }
+
+    #[inline]
+    pub fn serialize_with_chunk_rewrite<W: Write>(
+        &self,
+        writer: &mut W,
+        chunk_rewrite_fn: impl Fn(usize, CASChunkSequenceEntry) -> CASChunkSequenceEntry,
+    ) -> std::io::Result<usize> {
+        let mut n_out_bytes = 0;
+        n_out_bytes += self.header.serialize(writer)?;
+
+        for idx in 0..self.num_entries() {
+            let rewritten_chunk = chunk_rewrite_fn(idx, self.chunk(idx));
+            n_out_bytes += rewritten_chunk.serialize(writer)?;
+        }
+
+        Ok(n_out_bytes)
+    }
 }

mdb_shard/Cargo.toml

@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashMap};
+use std::collections::BTreeMap;
 use std::io::{Read, Seek, SeekFrom, Write, copy};
 use std::mem::size_of;
 use std::ops::Add;
@@ -13,11 +13,11 @@ use tracing::debug;
 use utils::serialization_utils::*;
 
 use crate::cas_structs::*;
-use crate::constants::*;
 use crate::error::{MDBShardError, Result};
 use crate::file_structs::*;
 use crate::interpolation_search::search_on_sorted_u64s;
 use crate::shard_in_memory::MDBInMemoryShard;
+use crate::streaming_shard::MDBMinimalShard;
 use crate::utils::{shard_expiry_time, truncate_hash};
 
 // Same size for FileDataSequenceHeader and FileDataSequenceEntry
@@ -931,44 +931,9 @@ impl MDBShardInfo {
     /// The chunk hashes are either multiple of 'hash_filter_modulues',
     /// or the hash of the first chunk of a file present in the shard.
     pub fn filter_cas_chunks_for_global_dedup<R: Read + Seek>(reader: &mut R) -> Result<Vec<MerkleHash>> {
-        let mut ret = Vec::new();
-
-        // First, go through and get all of the cas chunks.  This allows us to form the lookup for the CAS block
-        // hashes later.
-        let shard = MDBShardInfo::load_from_reader(reader)?;
-
-        let cas_chunks = shard.read_all_cas_blocks_full(reader)?;
-        let mut cas_block_lookup = HashMap::<MerkleHash, usize>::with_capacity(cas_chunks.len());
-
-        for (i, cas_info) in cas_chunks.iter().enumerate() {
-            cas_block_lookup.insert(cas_info.metadata.cas_hash, i);
-            for chunk in cas_info.chunks.iter() {
-                if hash_is_global_dedup_eligible(&chunk.chunk_hash) {
-                    ret.push(chunk.chunk_hash);
-                }
-            }
-        }
-
-        // Now, go through all the files present, collecting the first chunks of the files.
-        // TODO: break this out into a utility if needed.
-        let files = shard.read_all_file_info_sections(reader)?;
+        let shard = MDBMinimalShard::from_reader(reader, true, true)?;
 
-        for fi in files {
-            let Some(entry) = fi.segments.first() else {
-                continue;
-            };
-
-            let Some(cas_block_index) = cas_block_lookup.get(&entry.cas_hash) else {
-                continue;
-            };
-
-            // Scan the cas entries to get the proper index
-            let first_chunk_hash = cas_chunks[*cas_block_index].chunks[entry.chunk_index_start as usize].chunk_hash;
-
-            ret.push(first_chunk_hash);
-        }
-
-        Ok(ret)
+        Ok(shard.global_dedup_eligible_chunks())
     }
 
     /// Export the current shard as an hmac keyed shard, returning the number of bytes written