diff --git a/src/db/repo.rs b/src/db/repo.rs index 72608bb..7373344 100644 --- a/src/db/repo.rs +++ b/src/db/repo.rs @@ -5,6 +5,10 @@ use rusqlite::{Connection, params}; use std::fs; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, +}; use zstd::stream::write::Encoder; fn parse_license_text(metadata: &toml::Value) -> Option { @@ -28,6 +32,23 @@ pub struct RepoManager { pub repo_dir: PathBuf, } +struct IndexedPackage { + name: String, + version: String, + revision: u32, + description: Option, + homepage: Option, + license: Option, + filename: String, + size: u64, + sha256: String, + sha512: String, + provides: Vec, + runtime_dependencies: Vec, + optional_dependencies: Vec, + archive_files: Vec, +} + /// Search hit returned from a cached binary repository database. #[derive(Debug, Clone)] pub struct BinaryRepoSearchHit { @@ -89,13 +110,10 @@ impl RepoManager { self.init_repo_schema(&mut conn)?; - // Find all .depot.pkg.tar.zst files in repo_dir - for entry in fs::read_dir(&self.repo_dir)? { - let entry = entry?; - let path = entry.path(); - if path.to_string_lossy().ends_with(".depot.pkg.tar.zst") { - self.index_package(&mut conn, &path)?; - } + let package_paths = self.collect_repo_package_paths()?; + let indexed_packages = self.collect_indexed_packages_parallel(&package_paths)?; + for indexed in indexed_packages { + self.insert_indexed_package(&mut conn, indexed)?; } conn.close().map_err(|(_, e)| e)?; @@ -109,6 +127,72 @@ impl RepoManager { Ok(compressed_db_path) } + fn collect_repo_package_paths(&self) -> Result> { + let mut package_paths = Vec::new(); + for entry in fs::read_dir(&self.repo_dir) + .with_context(|| format!("Failed to read {}", self.repo_dir.display()))? + { + let entry = entry?; + let path = entry.path(); + if path.is_file() && path.to_string_lossy().ends_with(".depot.pkg.tar.zst") { + package_paths.push(path); + } + } + package_paths.sort(); + Ok(package_paths) + } + + fn collect_indexed_packages_parallel( + &self, + package_paths: &[PathBuf], + ) -> Result> { + if package_paths.is_empty() { + return Ok(Vec::new()); + } + + let worker_count = num_cpus().min(package_paths.len()); + crate::log_info!( + "Using {} thread(s) to index {} package(s)...", + worker_count, + package_paths.len() + ); + + let next_index = AtomicUsize::new(0); + let mut indexed = std::thread::scope(|scope| -> Result> { + let (tx, rx) = mpsc::channel::<(usize, Result)>(); + + for _ in 0..worker_count { + let tx = tx.clone(); + let next_index = &next_index; + scope.spawn(move || { + loop { + let idx = next_index.fetch_add(1, Ordering::Relaxed); + if idx >= package_paths.len() { + break; + } + let result = self.read_indexed_package(&package_paths[idx]); + if tx.send((idx, result)).is_err() { + break; + } + } + }); + } + drop(tx); + + let mut indexed = Vec::with_capacity(package_paths.len()); + for _ in 0..package_paths.len() { + let (idx, result) = rx + .recv() + .context("Failed to receive package indexing result from worker")?; + indexed.push((idx, result?)); + } + Ok(indexed) + })?; + + indexed.sort_by_key(|(idx, _)| *idx); + Ok(indexed.into_iter().map(|(_, pkg)| pkg).collect()) + } + fn init_repo_schema(&self, conn: &mut Connection) -> Result<()> { conn.execute_batch( "CREATE TABLE packages ( @@ -150,10 +234,14 @@ impl RepoManager { Ok(()) } - fn index_package(&self, conn: &mut Connection, pkg_path: &Path) -> Result<()> { + fn read_indexed_package(&self, pkg_path: &Path) -> Result { crate::log_info!("Indexing package {}...", pkg_path.display()); - let filename = pkg_path.file_name().unwrap().to_string_lossy(); + let filename = pkg_path + .file_name() + .and_then(|name| name.to_str()) + .with_context(|| format!("Invalid package filename: {}", pkg_path.display()))? + .to_string(); let size = pkg_path.metadata()?.len(); let (sha256, sha512) = self.calculate_hashes(pkg_path)?; @@ -264,6 +352,42 @@ impl RepoManager { revision = name_parts[2].parse().unwrap_or(1); } + Ok(IndexedPackage { + name, + version, + revision, + description, + homepage, + license, + filename, + size, + sha256, + sha512, + provides, + runtime_dependencies, + optional_dependencies, + archive_files, + }) + } + + fn insert_indexed_package(&self, conn: &mut Connection, indexed: IndexedPackage) -> Result<()> { + let IndexedPackage { + name, + version, + revision, + description, + homepage, + license, + filename, + size, + sha256, + sha512, + provides, + runtime_dependencies, + optional_dependencies, + archive_files, + } = indexed; + // Insert into database conn.execute( "INSERT INTO packages (name, version, revision, description, homepage, license, filename, size, sha256, sha512) @@ -1758,7 +1882,8 @@ optional = [] let mut conn = Connection::open_in_memory().unwrap(); let manager = RepoManager::new(repo_dir.to_path_buf()); manager.init_repo_schema(&mut conn).unwrap(); - manager.index_package(&mut conn, &pkg_path).unwrap(); + let indexed = manager.read_indexed_package(&pkg_path).unwrap(); + manager.insert_indexed_package(&mut conn, indexed).unwrap(); let (name, version, revision, desc, home, lic, sha256, sha512): ( String, @@ -1832,7 +1957,8 @@ license = ["MIT", "Apache-2.0"] let mut conn = Connection::open_in_memory().unwrap(); let manager = RepoManager::new(repo_dir.to_path_buf()); manager.init_repo_schema(&mut conn).unwrap(); - manager.index_package(&mut conn, &pkg_path).unwrap(); + let indexed = manager.read_indexed_package(&pkg_path).unwrap(); + manager.insert_indexed_package(&mut conn, indexed).unwrap(); let lic: Option = conn .query_row("SELECT license FROM packages", [], |r| r.get(0))