use std::{ collections::HashMap, future::Future, path::Path, sync::{ atomic::{AtomicU64, Ordering}, Arc, }, }; use anyhow::{anyhow, Error}; use futures::{stream::FuturesUnordered, StreamExt}; use hex::ToHex as _; use humansize::{format_size, BINARY}; use serde::Serialize; use serde_json::json; use sha2::{Digest as _, Sha256}; use tokio::{io::AsyncReadExt as _, join, sync::Mutex}; #[derive(Serialize, Clone)] pub struct FileEntry { filename: String, start: usize, length: usize, permissions: u32, } #[derive(Serialize, Clone)] pub struct ChunkData { files: Vec, checksum: String, } #[derive(Serialize)] pub struct Manifest { version: String, chunks: HashMap, size: u64, } const CHUNK_SIZE: u64 = 1024 * 1024 * 64; const WIGGLE: u64 = 1024 * 1024; use crate::versions::{create_backend_constructor, types::VersionFile}; pub async fn generate_manifest_rusty( dir: &Path, progress_sfn: V, log_sfn: T, ) -> anyhow::Result { let mut backend = create_backend_constructor(dir).ok_or(anyhow!("Could not create backend for path."))?()?; let required_single_file = backend.require_whole_files(); let files = backend.list_files().await?; // Filepath to chunk data let mut chunks: Vec> = Vec::new(); let mut current_chunk: Vec<(VersionFile, u64, u64)> = Vec::new(); log_sfn("organizing files into chunks...".to_string()); for version_file in files { // If we need the whole file, and this file would take up a whole chunk, add it to it's own chunk and move on if required_single_file && version_file.size >= CHUNK_SIZE { let size = version_file.size; chunks.push(vec![(version_file, 0, size)]); continue; } let mut current_size = current_chunk.iter().map(|v| v.2 - v.1).sum::(); // If we need the whole file, add this current file and move on, potentially adding and creating new chunk if need be if required_single_file { let size = version_file.size; current_chunk.push((version_file, 0, size)); current_size += size; if current_size >= CHUNK_SIZE { // Pop current and add, then reset let new_chunk = std::mem::take(&mut current_chunk); chunks.push(new_chunk); } continue; } // Otherwise we calculate how much of the file we need, then use that much let remaining_budget = (CHUNK_SIZE + WIGGLE) - current_size; if version_file.size >= remaining_budget { let remaining_budget = CHUNK_SIZE - current_size; current_chunk.push((version_file.clone(), 0, remaining_budget)); let new_chunk = std::mem::take(&mut current_chunk); chunks.push(new_chunk); let remaining_size = version_file.size - remaining_budget; let mut running_offset = remaining_budget; // Do everything but the last one while running_offset < remaining_size { let chunk_size = CHUNK_SIZE.min(remaining_size); let chunk = vec![(version_file.clone(), running_offset, chunk_size)]; if chunk_size == CHUNK_SIZE { chunks.push(chunk); } else { current_chunk = chunk; } running_offset += chunk_size; } continue; } else { let size = version_file.size; current_chunk.push((version_file, 0, size)); current_size += size; } if current_size >= CHUNK_SIZE { // Pop current and add, then reset let new_chunk = std::mem::take(&mut current_chunk); chunks.push(new_chunk); } } if !current_chunk.is_empty() { chunks.push(current_chunk); } log_sfn(format!( "organized into {} chunks, generating checksums...", chunks.len() )); let manifest: Arc>> = Arc::new(Mutex::new(HashMap::new())); let total_manifest_length = Arc::new(AtomicU64::new(0)); let backend = Arc::new(Mutex::new(backend)); let futures: FuturesUnordered>> = FuturesUnordered::new(); let (send_log, mut recieve_log) = tokio::sync::mpsc::channel(16); let chunks_length = chunks.len(); for (index, chunk) in chunks.into_iter().enumerate() { let send_log = send_log.clone(); let backend = backend.clone(); let total_manifest_length = total_manifest_length.clone(); let manifest = manifest.clone(); futures.push(async move { let mut read_buf = vec![0; 1024 * 1024 * 64]; let uuid = uuid::Uuid::new_v4().to_string(); let mut hasher = Sha256::new(); let mut chunk_data = ChunkData { files: Vec::new(), checksum: String::new(), }; let mut chunk_length = 0; for (file, start, length) in chunk { let mut reader = { let mut backend_lock = backend.lock().await; let reader = backend_lock.reader(&file, start, start + length).await?; reader }; loop { let amount = reader.read(&mut read_buf).await?; if amount == 0 { break; } hasher.update(&read_buf[0..amount]); } chunk_length += length; chunk_data.files.push(FileEntry { filename: file.relative_filename, start: start.try_into().unwrap(), length: length.try_into().unwrap(), permissions: file.permission, }); } send_log .send(format!( "created chunk of size {} from {} files (index {})", format_size(chunk_length, BINARY), chunk_data.files.len(), index )) .await?; total_manifest_length.fetch_add(chunk_length, Ordering::Relaxed); let hash: String = hasher.finalize().encode_hex(); chunk_data.checksum = hash; { let mut manifest_lock = manifest.lock().await; manifest_lock.insert(uuid, chunk_data); }; Ok(()) }); } drop(send_log); join!( async move { let mut current_progress = 0f32; let total_progress = chunks_length as f32; while let Some(message) = recieve_log.recv().await { log_sfn(message); current_progress += 1.0f32; progress_sfn((current_progress / total_progress) * 100.0f32) } }, futures.collect::>>() ); let manifest = manifest.lock().await; let manifest = manifest.clone(); Ok(Manifest { version: "2".to_string(), chunks: manifest, size: total_manifest_length.fetch_add(0, Ordering::Relaxed) }) }