use crate::block::{Block, MAX_BLOCK_SIZE};
use crate::chunking::{
ChunkingStrategy, DagLink, DagNode, DeduplicationStats, DEFAULT_CHUNK_SIZE, MAX_LINKS_PER_NODE,
MIN_CHUNK_SIZE,
};
use crate::cid::{Cid, HashAlgorithm};
use crate::error::{Error, Result};
use crate::metrics::global_metrics;
use bytes::Bytes;
use rayon::prelude::*;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[cfg(test)]
use crate::cid::CidBuilder;
#[derive(Debug, Clone)]
pub struct ParallelChunkingConfig {
pub chunk_size: usize,
pub strategy: ChunkingStrategy,
pub max_links_per_node: usize,
pub hash_algorithm: HashAlgorithm,
pub num_threads: Option<usize>,
}
impl Default for ParallelChunkingConfig {
fn default() -> Self {
Self {
chunk_size: DEFAULT_CHUNK_SIZE,
strategy: ChunkingStrategy::FixedSize,
max_links_per_node: MAX_LINKS_PER_NODE,
hash_algorithm: HashAlgorithm::Sha256,
num_threads: None,
}
}
}
impl ParallelChunkingConfig {
pub fn with_chunk_size(chunk_size: usize) -> Result<Self> {
if chunk_size < MIN_CHUNK_SIZE {
return Err(Error::InvalidInput(format!(
"Chunk size {} is below minimum {}",
chunk_size, MIN_CHUNK_SIZE
)));
}
if chunk_size > MAX_BLOCK_SIZE {
return Err(Error::InvalidInput(format!(
"Chunk size {} exceeds maximum {}",
chunk_size, MAX_BLOCK_SIZE
)));
}
Ok(Self {
chunk_size,
..Default::default()
})
}
pub fn with_threads(mut self, num_threads: usize) -> Self {
self.num_threads = Some(num_threads);
self
}
pub fn with_hash_algorithm(mut self, algorithm: HashAlgorithm) -> Self {
self.hash_algorithm = algorithm;
self
}
pub fn with_content_defined(mut self) -> Self {
self.strategy = ChunkingStrategy::ContentDefined;
self
}
}
#[derive(Debug, Clone)]
pub struct ParallelChunkingResult {
pub root_cid: Cid,
pub chunk_count: usize,
pub total_bytes: usize,
pub dedup_stats: DeduplicationStats,
pub duration: Duration,
pub chunk_cids: Vec<Cid>,
pub dag_nodes: Vec<DagNode>,
}
pub struct ParallelChunker {
config: ParallelChunkingConfig,
}
impl ParallelChunker {
pub fn new() -> Self {
Self {
config: ParallelChunkingConfig::default(),
}
}
pub fn with_config(config: ParallelChunkingConfig) -> Self {
Self { config }
}
pub fn chunk_parallel(&self, data: &[u8]) -> Result<ParallelChunkingResult> {
let start = Instant::now();
let metrics = global_metrics();
if data.len() < 1_000_000 {
return self.chunk_sequential(data, start);
}
let chunk_ranges = self.calculate_chunk_ranges(data.len());
let chunk_results: Vec<_> = chunk_ranges
.par_iter()
.map(|(start, end)| {
let chunk_data = &data[*start..*end];
let block = Block::new(Bytes::copy_from_slice(chunk_data))
.map_err(|e| Error::InvalidData(e.to_string()))?;
Ok((*block.cid(), block.data().len()))
})
.collect::<Result<Vec<_>>>()?;
let dag_result = self.build_dag_parallel(&chunk_results)?;
let duration = start.elapsed();
metrics.record_chunking(chunk_results.len(), duration.as_micros() as u64);
Ok(ParallelChunkingResult {
root_cid: dag_result.root_cid,
chunk_count: chunk_results.len(),
total_bytes: data.len(),
dedup_stats: DeduplicationStats {
unique_chunks: chunk_results.len(),
total_chunks: chunk_results.len(),
reused_chunks: 0,
space_savings_percent: 0.0,
total_data_size: data.len() as u64,
deduplicated_size: data.len() as u64,
},
duration,
chunk_cids: chunk_results.iter().map(|(cid, _)| *cid).collect(),
dag_nodes: dag_result.nodes,
})
}
fn calculate_chunk_ranges(&self, data_len: usize) -> Vec<(usize, usize)> {
let chunk_size = self.config.chunk_size;
let mut ranges = Vec::new();
let mut offset = 0;
while offset < data_len {
let end = (offset + chunk_size).min(data_len);
ranges.push((offset, end));
offset = end;
}
ranges
}
fn build_dag_parallel(&self, chunks: &[(Cid, usize)]) -> Result<DagBuildResult> {
if chunks.is_empty() {
return Err(Error::InvalidInput(
"no chunks to build DAG from".to_string(),
));
}
if chunks.len() == 1 {
return Ok(DagBuildResult {
root_cid: chunks[0].0,
nodes: vec![],
});
}
let mut current_level: Vec<Cid> = chunks.iter().map(|(cid, _)| *cid).collect();
let all_nodes = Arc::new(Mutex::new(Vec::new()));
while current_level.len() > 1 {
let max_links = self.config.max_links_per_node;
let groups: Vec<_> = current_level.chunks(max_links).collect();
let parent_results: Vec<_> = groups
.par_iter()
.map(|group| {
let links: Vec<DagLink> = group
.iter()
.enumerate()
.map(|(idx, cid)| DagLink::with_name(*cid, 0, format!("chunk-{}", idx)))
.collect();
let node = DagNode {
links,
total_size: 0, data: None,
};
let ipld = node.to_ipld();
let cbor = ipld
.to_dag_cbor()
.map_err(|e| Error::Serialization(e.to_string()))?;
let block = Block::new(Bytes::from(cbor))
.map_err(|e| Error::InvalidData(e.to_string()))?;
Ok((*block.cid(), node))
})
.collect::<Result<Vec<_>>>()?;
let mut nodes_lock = all_nodes.lock().unwrap_or_else(|e| e.into_inner());
nodes_lock.extend(parent_results.iter().map(|(_, node)| node.clone()));
drop(nodes_lock);
current_level = parent_results.into_iter().map(|(cid, _)| cid).collect();
}
let nodes = Arc::try_unwrap(all_nodes)
.expect("no other Arc references to all_nodes at this point")
.into_inner()
.expect("Mutex is not poisoned");
Ok(DagBuildResult {
root_cid: current_level[0],
nodes,
})
}
fn chunk_sequential(&self, data: &[u8], start: Instant) -> Result<ParallelChunkingResult> {
let chunk_ranges = self.calculate_chunk_ranges(data.len());
let mut chunk_cids = Vec::new();
for (start_offset, end_offset) in chunk_ranges {
let chunk_data = &data[start_offset..end_offset];
let block = Block::new(Bytes::copy_from_slice(chunk_data))?;
chunk_cids.push((*block.cid(), block.data().len()));
}
let dag_result = self.build_dag_parallel(&chunk_cids)?;
Ok(ParallelChunkingResult {
root_cid: dag_result.root_cid,
chunk_count: chunk_cids.len(),
total_bytes: data.len(),
dedup_stats: DeduplicationStats {
unique_chunks: chunk_cids.len(),
total_chunks: chunk_cids.len(),
reused_chunks: 0,
space_savings_percent: 0.0,
total_data_size: data.len() as u64,
deduplicated_size: data.len() as u64,
},
duration: start.elapsed(),
chunk_cids: chunk_cids.iter().map(|(cid, _)| *cid).collect(),
dag_nodes: dag_result.nodes,
})
}
pub fn chunk_files_parallel(&self, files: &[Vec<u8>]) -> Result<Vec<ParallelChunkingResult>> {
files
.par_iter()
.map(|data| self.chunk_parallel(data))
.collect()
}
}
impl Default for ParallelChunker {
fn default() -> Self {
Self::new()
}
}
struct DagBuildResult {
root_cid: Cid,
nodes: Vec<DagNode>,
}
pub struct ParallelDeduplicator {
seen_cids: Arc<Mutex<std::collections::HashSet<Cid>>>,
stats: Arc<Mutex<DeduplicationStats>>,
}
impl ParallelDeduplicator {
pub fn new() -> Self {
Self {
seen_cids: Arc::new(Mutex::new(std::collections::HashSet::new())),
stats: Arc::new(Mutex::new(DeduplicationStats {
unique_chunks: 0,
total_chunks: 0,
reused_chunks: 0,
space_savings_percent: 0.0,
total_data_size: 0,
deduplicated_size: 0,
})),
}
}
pub fn check_unique(&self, cid: &Cid, size: usize) -> bool {
let mut seen = self.seen_cids.lock().unwrap_or_else(|e| e.into_inner());
let mut stats = self.stats.lock().unwrap_or_else(|e| e.into_inner());
stats.total_chunks += 1;
stats.total_data_size += size as u64;
if seen.insert(*cid) {
stats.unique_chunks += 1;
stats.deduplicated_size += size as u64;
true
} else {
stats.reused_chunks += 1;
false
}
}
pub fn stats(&self) -> DeduplicationStats {
let stats = self.stats.lock().unwrap_or_else(|e| e.into_inner());
let mut result = stats.clone();
if result.total_data_size > 0 {
result.space_savings_percent =
(1.0 - (result.deduplicated_size as f64 / result.total_data_size as f64)) * 100.0;
}
result
}
}
impl Default for ParallelDeduplicator {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parallel_chunking_basic() {
let data = vec![0u8; 1_000_000]; let chunker = ParallelChunker::new();
let result = chunker.chunk_parallel(&data).unwrap();
assert!(result.chunk_count > 0);
assert_eq!(result.total_bytes, 1_000_000);
assert!(result.duration.as_micros() > 0);
}
#[test]
fn test_parallel_chunking_small_file() {
let data = vec![0u8; 1024]; let chunker = ParallelChunker::new();
let result = chunker.chunk_parallel(&data).unwrap();
assert_eq!(result.chunk_count, 1);
assert_eq!(result.total_bytes, 1024);
}
#[test]
fn test_parallel_chunking_custom_size() {
let config = ParallelChunkingConfig::with_chunk_size(128 * 1024).unwrap();
let chunker = ParallelChunker::with_config(config);
let data = vec![0u8; 1_000_000];
let result = chunker.chunk_parallel(&data).unwrap();
assert!(result.chunk_count > 0);
}
#[test]
fn test_parallel_chunking_multiple_files() {
let files = vec![vec![0u8; 500_000], vec![1u8; 500_000], vec![2u8; 500_000]];
let chunker = ParallelChunker::new();
let results = chunker.chunk_files_parallel(&files).unwrap();
assert_eq!(results.len(), 3);
for result in results {
assert!(result.chunk_count > 0);
}
}
#[test]
fn test_chunk_ranges() {
let chunker = ParallelChunker::new();
let ranges = chunker.calculate_chunk_ranges(1_000_000);
assert!(!ranges.is_empty());
assert_eq!(ranges[0].0, 0);
for i in 1..ranges.len() {
assert_eq!(ranges[i - 1].1, ranges[i].0);
}
assert_eq!(ranges.last().unwrap().1, 1_000_000);
}
#[test]
fn test_parallel_deduplicator() {
let dedup = ParallelDeduplicator::new();
let cid = CidBuilder::new().build(b"test").unwrap();
assert!(dedup.check_unique(&cid, 100));
assert!(!dedup.check_unique(&cid, 100));
let stats = dedup.stats();
assert_eq!(stats.unique_chunks, 1);
assert_eq!(stats.total_chunks, 2);
assert!(stats.space_savings_percent > 0.0);
}
#[test]
fn test_config_validation() {
assert!(ParallelChunkingConfig::with_chunk_size(100).is_err());
assert!(ParallelChunkingConfig::with_chunk_size(128 * 1024).is_ok());
assert!(ParallelChunkingConfig::with_chunk_size(10_000_000).is_err());
}
#[test]
fn test_config_builder() {
let config = ParallelChunkingConfig::default()
.with_threads(4)
.with_hash_algorithm(HashAlgorithm::Sha3_256)
.with_content_defined();
assert_eq!(config.num_threads, Some(4));
assert_eq!(config.hash_algorithm, HashAlgorithm::Sha3_256);
assert_eq!(config.strategy, ChunkingStrategy::ContentDefined);
}
#[test]
fn test_empty_data() {
let chunker = ParallelChunker::new();
let data: Vec<u8> = vec![];
let result = chunker.chunk_parallel(&data);
assert!(result.is_err());
}
#[test]
fn test_single_chunk() {
let data = vec![42u8; 1024];
let chunker = ParallelChunker::new();
let result = chunker.chunk_parallel(&data).unwrap();
assert_eq!(result.chunk_count, 1);
assert!(!result.chunk_cids.is_empty());
}
#[test]
fn test_dag_building() {
let data = vec![0u8; 5_000_000]; let chunker = ParallelChunker::new();
let result = chunker.chunk_parallel(&data).unwrap();
assert!(result.chunk_count > 1);
assert!(!result.chunk_cids.is_empty());
}
}