diff --git a/Cargo.lock b/Cargo.lock index 5f15ee20e..298f2b919 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4360,6 +4360,7 @@ dependencies = [ "rayon", "rstest", "runtime-format", + "rustc-hash", "rustic_backend", "rustic_cdc", "rustic_testing", diff --git a/crates/backend/src/opendal.rs b/crates/backend/src/opendal.rs index c830c8cea..7ff11f711 100644 --- a/crates/backend/src/opendal.rs +++ b/crates/backend/src/opendal.rs @@ -34,6 +34,14 @@ use crate::reqwest::reqwest_client; mod constants { /// Default number of retries pub(super) const DEFAULT_RETRY: usize = 5; + + /// B2 `b2_list_file_names` page size to request. + /// + /// `OpenDAL` only sends `maxFileCount` when `ListOptions.limit` is set. If we + /// omit it, B2 defaults to 100 names per page (max 10000). Prune (and + /// check) list every pack under `data/`, so the default turns a few dozen + /// round-trips into thousands and dominates runtime against B2. + pub(super) const B2_LIST_PAGE_SIZE: usize = 10_000; } /// `OpenDALBackend` contains a wrapper around an blocking operator of the `OpenDAL` library. @@ -217,6 +225,30 @@ impl OpenDALBackend { Ok(Self { operator }) } + /// Listing options used for repository (and source) listings. + /// + /// For B2 this raises the per-request page size from the API default of 100 + /// to the documented maximum of 10000. + fn list_options(&self, recursive: bool) -> ListOptions { + ListOptions { + recursive, + limit: self.list_page_size(), + ..Default::default() + } + } + + /// Backend-specific list page size, if we should override the service default. + fn list_page_size(&self) -> Option { + let info = self.operator.info(); + if !info.capability().list_with_limit { + return None; + } + match info.scheme() { + "b2" => Some(constants::B2_LIST_PAGE_SIZE), + _ => None, + } + } + /// Return a path for the given file type and id. /// /// # Arguments @@ -249,10 +281,7 @@ impl OpenDALBackend { /// # Errors /// If listing fails or exclude patterns cannot be compiled pub fn as_source(self, excludes: &Excludes) -> RusticResult { - let list_options = ListOptions { - recursive: true, - ..Default::default() - }; + let list_options = self.list_options(true); // openDAL lister may entries in random order; hence we collect and sort them here. // This also allows to handle listing errors directly let mut entries: Vec<_> = self @@ -348,14 +377,9 @@ impl ReadBackend for OpenDALBackend { } let path = tpe.dirname().to_string() + "/"; - let list_options = ListOptions { - recursive: true, - ..Default::default() - }; - let lister = self .operator - .lister_options(&path, list_options) + .lister_options(&path, self.list_options(true)) .map_err(|err| { RusticError::with_source(ErrorKind::Backend, "Listing failed for `{type}`", err) .attach_context("type", tpe.to_string()) @@ -405,13 +429,9 @@ impl ReadBackend for OpenDALBackend { } let path = tpe.dirname().to_string() + "/"; - let list_options = ListOptions { - recursive: true, - ..Default::default() - }; let lister = self .operator - .lister_options(&path, list_options) + .lister_options(&path, self.list_options(true)) .map_err(|err| { RusticError::with_source(ErrorKind::Backend, "Listing failed for `{type}`", err) .attach_context("type", tpe.to_string()) @@ -632,6 +652,28 @@ mod tests { assert!(Throttle::from_str(input).is_err()); } + #[rstest] + #[case("b2", Some(constants::B2_LIST_PAGE_SIZE))] + #[case("s3_aws", None)] + fn list_page_size_matches_scheme( + #[case] fixture: &str, + #[case] expected: Option, + ) -> Result<()> { + #[derive(Deserialize)] + struct TestCase { + path: String, + options: BTreeMap, + } + + let fixture_path = PathBuf::from(format!("tests/fixtures/opendal/{fixture}.toml")); + let test: TestCase = toml::from_str(&fs::read_to_string(fixture_path)?)?; + let backend = OpenDALBackend::new(test.path, test.options)?; + + assert_eq!(backend.list_page_size(), expected); + assert_eq!(backend.list_options(true).limit, expected); + Ok(()) + } + #[rstest] fn new_opendal_backend( #[files("tests/fixtures/opendal/*.toml")] test_case: PathBuf, diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 691bfb2fb..8571e711c 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -53,6 +53,7 @@ log = { workspace = true } crossbeam-channel = "0.5.15" pariter = "0.6.0" rayon = "1.11.0" +rustc-hash = "2.1.1" # crypto aes256ctr_poly1305aes = { version = "0.2.1", features = ["std"] } # we need std here for error impls @@ -75,7 +76,7 @@ cached = { version = "2.0.2", default-features = false, features = ["proc_macro" dunce = "1.0.5" filetime = "0.2.27" ignore = "0.4.25" -nix = { version = "0.31.1", default-features = false, features = ["user", "fs"] } +nix = { version = "0.31.1", default-features = false, features = ["user", "fs", "resource"] } path-dedot = "4.0.1" walkdir = "2.5.0" diff --git a/crates/core/src/backend.rs b/crates/core/src/backend.rs index b1afb4d46..5a95adb9d 100644 --- a/crates/core/src/backend.rs +++ b/crates/core/src/backend.rs @@ -15,6 +15,7 @@ use std::{io::Read, ops::Deref, path::PathBuf, sync::Arc}; use bytes::{Buf, Bytes, buf::Reader}; use enum_map::Enum; use log::trace; +use rayon::current_num_threads; #[cfg(test)] use mockall::mock; @@ -105,6 +106,22 @@ pub trait ReadBackend: Send + Sync + 'static { /// * If the files could not be listed. fn list_with_size(&self, tpe: FileType) -> RusticResult>; + /// How many parallel readers to use for `stream_list` of these files. + /// + /// Remote backends keep extra workers for B2 RTTs. Cached backends use a + /// few workers when the files are already on disk so we do not thrash HDD. + fn prefetch_workers(&self, _tpe: FileType, _ids: &[Id]) -> usize { + (current_num_threads() + 16).clamp(16, 32) + } + + /// How many threads to spawn to load trees (`TreeStreamer`). + /// + /// Remote backends use `2 × CPUs` (8–32) so prune is not RTT-bound on B2. + /// Cached backends use fewer when pack files are already on disk. + fn tree_loader_count(&self) -> usize { + current_num_threads().saturating_mul(2).clamp(8, 32) + } + /// Lists all files of the given type. /// /// # Arguments @@ -448,6 +465,12 @@ impl ReadBackend for Arc { fn list_with_size(&self, tpe: FileType) -> RusticResult> { self.deref().list_with_size(tpe) } + fn prefetch_workers(&self, tpe: FileType, ids: &[Id]) -> usize { + self.deref().prefetch_workers(tpe, ids) + } + fn tree_loader_count(&self) -> usize { + self.deref().tree_loader_count() + } fn list(&self, tpe: FileType) -> RusticResult> { self.deref().list(tpe) } diff --git a/crates/core/src/backend/cache.rs b/crates/core/src/backend/cache.rs index 7cacd6539..2f56cd417 100644 --- a/crates/core/src/backend/cache.rs +++ b/crates/core/src/backend/cache.rs @@ -1,7 +1,8 @@ use std::{ collections::HashMap, + fmt, fs::{self, File}, - io::{self, Read, Seek, SeekFrom}, + io::{self, Read}, path::{Path, PathBuf}, sync::Arc, }; @@ -19,6 +20,101 @@ use crate::{ repofile::configfile::RepositoryId, }; +mod constants { + /// Pack files kept open for cache `read_partial`. + /// + /// Tree walking reads many blobs from the same packs. Opening the cache + /// file per blob was ~65% of prune CPU in a Time Profiler trace. + pub(super) const OPEN_FILE_CAPACITY: usize = 2048; + /// Descriptors left for sockets, index files, and other I/O. + pub(super) const OPEN_FILE_RESERVE: u64 = 64; +} + +type OpenFileCache = quick_cache::sync::Cache>; + +/// Use up to [`constants::OPEN_FILE_CAPACITY`] cached pack FDs, leaving +/// [`constants::OPEN_FILE_RESERVE`] for other I/O. `rlimit/8` mapped a Darwin +/// 8192 soft limit to 1024 handles. +fn open_file_capacity() -> usize { + #[cfg(unix)] + if let Ok((soft, _)) = + nix::sys::resource::getrlimit(nix::sys::resource::Resource::RLIMIT_NOFILE) + { + return open_file_capacity_from_soft_limit(soft); + } + // Conservative fallback on platforms without a queryable descriptor limit. + 32 +} + +fn open_file_capacity_from_soft_limit(soft: impl Into) -> usize { + usize::try_from(soft.into().saturating_sub(constants::OPEN_FILE_RESERVE)) + .unwrap_or(usize::MAX) + .min(constants::OPEN_FILE_CAPACITY) + .max(1) +} + +struct CachedFile { + file: File, + #[cfg(any(test, not(unix)))] + seek_lock: std::sync::Mutex<()>, +} + +impl CachedFile { + fn new(file: File) -> Self { + Self { + file, + #[cfg(any(test, not(unix)))] + seek_lock: std::sync::Mutex::new(()), + } + } + + fn read_range(&self, offset: u32, length: u32) -> io::Result { + #[cfg(unix)] + { + use std::os::unix::fs::FileExt; + let mut vec = vec![0; length as usize]; + self.file.read_exact_at(&mut vec, u64::from(offset))?; + Ok(vec.into()) + } + #[cfg(not(unix))] + self.read_range_seeking(offset, length) + } + + /// Cloned file handles share a cursor. Serialize the entire seek/read pair + /// on platforms without positional reads; never clone the cached handle. + #[cfg(any(test, not(unix)))] + fn read_range_seeking(&self, offset: u32, length: u32) -> io::Result { + use std::io::{Seek, SeekFrom}; + let _guard = self + .seek_lock + .lock() + .map_err(|_| io::Error::other("cache seek lock poisoned"))?; + let mut file = &self.file; + let mut vec = vec![0; length as usize]; + _ = file.seek(SeekFrom::Start(u64::from(offset)))?; + file.read_exact(&mut vec)?; + Ok(vec.into()) + } +} + +fn is_too_many_open_files(err: &io::Error) -> bool { + #[cfg(unix)] + { + // POSIX EMFILE / ENFILE. Do not treat 4 as a match: that is EINTR. + matches!(err.raw_os_error(), Some(24 | 23)) + } + #[cfg(windows)] + { + // ERROR_TOO_MANY_OPEN_FILES + err.raw_os_error() == Some(4) + } + #[cfg(not(any(unix, windows)))] + { + let _ = err; + false + } +} + /// Backend that caches data. /// /// This backend caches data in a directory. @@ -60,6 +156,34 @@ impl ReadBackend for CachedBackend { self.be.location() } + fn prefetch_workers(&self, tpe: FileType, ids: &[Id]) -> usize { + if ids.is_empty() { + return 1; + } + let hits = ids + .iter() + .filter(|id| self.cache.path(tpe, id).is_file()) + .count(); + // Warm cache: few readers so we do not thrash local disk. Cold cache: + // extra workers for high-latency GETs. + if hits.saturating_mul(4) >= ids.len().saturating_mul(3) { + rayon::current_num_threads().clamp(2, 4) + } else { + (rayon::current_num_threads() + 16).clamp(16, 32) + } + } + + fn tree_loader_count(&self) -> usize { + // Warm pack cache: local preads. 2×CPUs loaders oversubscribe QEMU + // vCPUs (KVM PV spinlocks + musl malloc). Cold cache: keep extra + // loaders for B2 RTTs. + if self.cache.has_cached_packs() { + rayon::current_num_threads().clamp(4, 8) + } else { + self.be.tree_loader_count() + } + } + /// Lists all files with their size of the given type. /// /// # Arguments @@ -159,12 +283,18 @@ impl ReadBackend for CachedBackend { length: u32, ) -> RusticResult { if cacheable || tpe.is_cacheable() { - let guard = self.lock_pool.blocking_lock(*id); - if self.cache.path(tpe, id).exists() { - // early drop the lock guard, so we can read the cache in parallel. - drop(guard); + match self.cache.read_partial(tpe, id, offset, length) { + Ok(Some(data)) => return Ok(data), + Ok(None) => {} + Err(err) => warn!( + "Error in cache backend reading {tpe:?},{id}: {}", + err.display_log() + ), } + // Miss: serialize fills of the same pack so two threads don't both + // download it. Hits above skip this lock and the exists()+open storm. + let _guard = self.lock_pool.blocking_lock(*id); match self.cache.read_partial(tpe, id, offset, length) { Ok(Some(data)) => return Ok(data), Ok(None) => {} @@ -263,10 +393,21 @@ impl WriteBackend for CachedBackend { } /// Backend that caches data in a directory. -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct Cache { /// The path to the cache. path: PathBuf, + /// Recently used cache files kept open for `read_partial`. + open_files: Arc, +} + +impl fmt::Debug for Cache { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Cache") + .field("path", &self.path) + .field("open_files", &self.open_files.len()) + .finish() + } } impl Cache { @@ -329,7 +470,19 @@ impl Cache { .attach_context("id", id.to_string()) })?; - Ok(Self { path }) + Ok(Self { + path, + open_files: Arc::new(OpenFileCache::new(open_file_capacity())), + }) + } + + /// True if at least one pack file is already in this cache. + #[must_use] + pub fn has_cached_packs(&self) -> bool { + WalkDir::new(self.path.join(FileType::Pack.dirname())) + .into_iter() + .filter_map(Result::ok) + .any(|e| e.file_type().is_file()) } /// Returns the path to the location of this [`Cache`]. @@ -501,11 +654,23 @@ impl Cache { ) -> RusticResult> { trace!("cache reading tpe: {tpe:?}, id: {id}, offset: {offset}"); - let path = self.path(tpe, id); + if let Some(file) = self.open_files.get(id) { + match file.read_range(offset, length) { + Ok(data) => { + trace!("cache hit!"); + return Ok(Some(data)); + } + Err(_) => { + // Stale or truncated FD; reopen from disk. + _ = self.open_files.remove(id); + } + } + } - let mut file = match File::open(&path) { - Ok(file) => file, - Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), + let path = self.path(tpe, id); + let file = match self.open_cached(id, &path) { + Ok(Some(file)) => file, + Ok(None) => return Ok(None), Err(err) => { return Err(RusticError::with_source( ErrorKind::InputOutput, @@ -518,23 +683,7 @@ impl Cache { } }; - _ = file - .seek(SeekFrom::Start(u64::from(offset))) - .map_err(|err| { - RusticError::with_source( - ErrorKind::InputOutput, - "Failed to seek to `{offset}` in file `{path}`", - err, - ) - .attach_context("path", path.display().to_string()) - .attach_context("tpe", tpe.to_string()) - .attach_context("id", id.to_string()) - .attach_context("offset", offset.to_string()) - })?; - - let mut vec = vec![0; length as usize]; - - file.read_exact(&mut vec).map_err(|err| { + let data = file.read_range(offset, length).map_err(|err| { RusticError::with_source( ErrorKind::InputOutput, "Failed to read at offset `{offset}` from file at `{path}`", @@ -549,7 +698,33 @@ impl Cache { trace!("cache hit!"); - Ok(Some(vec.into())) + Ok(Some(data)) + } + + fn open_cached(&self, id: &Id, path: &Path) -> io::Result>> { + if let Some(file) = self.open_files.get(id) { + return Ok(Some(file)); + } + + match File::open(path) { + Ok(file) => Ok(Some(self.remember_open(id, file))), + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None), + Err(err) if is_too_many_open_files(&err) => { + self.open_files.clear(); + match File::open(path) { + Ok(file) => Ok(Some(self.remember_open(id, file))), + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None), + Err(err) => Err(err), + } + } + Err(err) => Err(err), + } + } + + fn remember_open(&self, id: &Id, file: File) -> Arc { + let file = Arc::new(CachedFile::new(file)); + self.open_files.insert(*id, file.clone()); + file } /// Writes the given data to the given file. @@ -627,6 +802,8 @@ impl Cache { .attach_context("path", filename.display().to_string()) .ask_report() })?; + // Drop any FD pointing at the previous inode. + _ = self.open_files.remove(id); Ok(()) } @@ -643,6 +820,7 @@ impl Cache { /// * If the file could not be removed. pub fn remove(&self, tpe: FileType, id: &Id) -> RusticResult<()> { trace!("cache writing tpe: {tpe:?}, id: {id}"); + _ = self.open_files.remove(id); let filename = self.path(tpe, id); fs::remove_file(&filename).map_err(|err| { RusticError::with_source( @@ -658,3 +836,166 @@ impl Cache { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + + fn new_cache() -> (tempfile::TempDir, Cache) { + let dir = tempfile::tempdir().unwrap(); + let cache = Cache::new(RepositoryId::default(), Some(dir.path().to_path_buf())).unwrap(); + (dir, cache) + } + + #[test] + fn read_partial_reuses_open_file() { + let (_dir, cache) = new_cache(); + let id = Id::random(); + let payload: Vec = (0..=u8::MAX).cycle().take(4096).collect(); + cache + .write_bytes(FileType::Pack, &id, &payload.clone().into()) + .unwrap(); + + for _ in 0..64 { + let got = cache + .read_partial(FileType::Pack, &id, 100, 50) + .unwrap() + .unwrap(); + assert_eq!(got.as_ref(), &payload[100..150]); + } + assert_eq!(cache.open_files.len(), 1); + } + + #[test] + fn write_bytes_invalidates_open_file() { + let (_dir, cache) = new_cache(); + let id = Id::random(); + cache + .write_bytes(FileType::Pack, &id, &vec![0_u8; 128].into()) + .unwrap(); + assert_eq!( + cache + .read_partial(FileType::Pack, &id, 0, 4) + .unwrap() + .unwrap() + .as_ref(), + &[0, 0, 0, 0] + ); + + cache + .write_bytes(FileType::Pack, &id, &vec![0xff_u8; 128].into()) + .unwrap(); + assert_eq!( + cache + .read_partial(FileType::Pack, &id, 0, 4) + .unwrap() + .unwrap() + .as_ref(), + &[0xff, 0xff, 0xff, 0xff] + ); + } + + #[test] + fn remove_closes_and_hides_file() { + let (_dir, cache) = new_cache(); + let id = Id::random(); + cache + .write_bytes(FileType::Pack, &id, &vec![1_u8; 16].into()) + .unwrap(); + _ = cache.read_partial(FileType::Pack, &id, 0, 4).unwrap(); + cache.remove(FileType::Pack, &id).unwrap(); + assert!( + cache + .read_partial(FileType::Pack, &id, 0, 4) + .unwrap() + .is_none() + ); + assert_eq!(cache.open_files.len(), 0); + } + + #[test] + fn concurrent_partial_reads() { + let (_dir, cache) = new_cache(); + let id = Id::random(); + let payload: Vec = (0..=250_u8).cycle().take(8192).collect(); + cache + .write_bytes(FileType::Pack, &id, &payload.clone().into()) + .unwrap(); + thread::scope(|s| { + for t in 0..16 { + let cache = &cache; + let payload = &payload; + _ = s.spawn(move || { + let offset = u32::try_from(t * 16).unwrap(); + for _ in 0..100 { + let got = cache + .read_partial(FileType::Pack, &id, offset, 16) + .unwrap() + .unwrap(); + let start = usize::try_from(offset).unwrap(); + assert_eq!(got.as_ref(), &payload[start..start + 16]); + // Exercise the non-Unix implementation on every test platform. + let file = cache.open_files.get(&id).unwrap(); + let got = file.read_range_seeking(offset, 16).unwrap(); + assert_eq!(got.as_ref(), &payload[start..start + 16]); + } + }); + } + }); + assert_eq!(cache.open_files.len(), 1); + } + + #[test] + fn open_file_capacity_uses_reserve_not_an_eighth() { + assert_eq!(open_file_capacity_from_soft_limit(8192_u64), 2048); + assert_eq!(open_file_capacity_from_soft_limit(1024_u64), 960); + assert_eq!(open_file_capacity_from_soft_limit(64_u64), 1); + assert_eq!(open_file_capacity_from_soft_limit(0_u64), 1); + } + + #[cfg(unix)] + #[test] + fn cache_respects_file_descriptor_limit() { + use nix::sys::resource::{Resource, getrlimit, setrlimit}; + + const CHILD: &str = "RUSTIC_CACHE_FD_LIMIT_TEST"; + if std::env::var_os(CHILD).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "backend::cache::tests::cache_respects_file_descriptor_limit", + "--nocapture", + "--test-threads=1", + ]) + .env(CHILD, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + + let (_, hard) = getrlimit(Resource::RLIMIT_NOFILE).unwrap(); + setrlimit(Resource::RLIMIT_NOFILE, 64.min(hard), hard).unwrap(); + let (_dir, cache) = new_cache(); + let ids: Vec<_> = (0..100).map(|_| Id::random()).collect(); + for id in &ids { + cache + .write_bytes(FileType::Pack, id, &vec![0_u8; 16].into()) + .unwrap(); + } + for id in &ids { + _ = cache + .read_partial(FileType::Pack, id, 0, 4) + .unwrap() + .unwrap(); + // Other backend/file operations must still have descriptor headroom. + let _other_files: Vec<_> = (0..16).map(|_| File::open("/dev/null").unwrap()).collect(); + } + } +} diff --git a/crates/core/src/backend/decrypt.rs b/crates/core/src/backend/decrypt.rs index c36d74ac5..ed6b06140 100644 --- a/crates/core/src/backend/decrypt.rs +++ b/crates/core/src/backend/decrypt.rs @@ -1,10 +1,103 @@ -use std::{num::NonZeroU32, sync::Arc}; +use std::{cell::RefCell, num::NonZeroU32, sync::Arc}; use bytes::Bytes; use crossbeam_channel::{Receiver, bounded}; use rayon::{prelude::*, spawn}; use zstd::stream::{copy_encode, decode_all, encode_all}; +thread_local! { + static ZSTD: RefCell = const { + RefCell::new(ZstdTls { + decompressor: None, + buf: Vec::new(), + }) + }; +} + +struct ZstdTls { + decompressor: Option>, + buf: Vec, +} + +fn zstd_tls_decompressor( + slot: &mut ZstdTls, +) -> RusticResult<&mut zstd::bulk::Decompressor<'static>> { + if slot.decompressor.is_none() { + slot.decompressor = Some(zstd::bulk::Decompressor::new().map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to create zstd decompressor.", + err, + ) + })?); + } + Ok(slot + .decompressor + .as_mut() + .expect("zstd decompressor is initialized")) +} + +/// Decode zstd with a decompressor kept on this thread. +/// +/// `zstd::decode_all` builds a new `DCtx` per call. Tree walking does that for +/// every blob and showed up as `ZSTD_createDCtx` / `munmap` / page faults. +fn zstd_decompress(data: &[u8], uncompressed_len: usize) -> RusticResult> { + ZSTD.with(|slot| { + let mut slot = slot.borrow_mut(); + zstd_tls_decompressor(&mut slot)? + .decompress(data, uncompressed_len) + .map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to decode zstd compressed data. The data may be corrupted.", + err, + ) + }) + }) +} + +/// Decompress into a thread-local buffer and run `f` on the plaintext. +/// +/// Prune tree walking used to allocate a new uncompressed `Vec` per blob, +/// which showed up as `kernel_init_pages`. The buffer keeps capacity on this +/// thread so later trees reuse the same pages. +fn zstd_decompress_with( + data: &[u8], + uncompressed_len: usize, + f: impl FnOnce(&[u8]) -> RusticResult, +) -> RusticResult { + ZSTD.with(|slot| { + let mut slot = slot.borrow_mut(); + let ZstdTls { decompressor, buf } = &mut *slot; + if decompressor.is_none() { + *decompressor = Some(zstd::bulk::Decompressor::new().map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to create zstd decompressor.", + err, + ) + })?); + } + buf.clear(); + if buf.capacity() < uncompressed_len { + buf.reserve(uncompressed_len); + } + let written = decompressor + .as_mut() + .expect("zstd decompressor is initialized") + .decompress_to_buffer(data, buf) + .map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to decode zstd compressed data. The data may be corrupted.", + err, + ) + })?; + buf.truncate(written); + f(buf) + }) +} + pub use zstd::compression_level_range; use crate::{ @@ -75,13 +168,7 @@ pub trait DecryptReadBackend: ReadBackend + Clone + 'static { ) -> RusticResult { let mut data = self.decrypt(data)?; if let Some(length) = uncompressed_length { - data = decode_all(&*data).map_err(|err| { - RusticError::with_source( - ErrorKind::Internal, - "Failed to decode zstd compressed data. The data may be corrupted.", - err, - ) - })?; + data = zstd_decompress(&data, length.get() as usize)?; if data.len() != length.get() as usize { return Err(RusticError::new( @@ -96,6 +183,34 @@ pub trait DecryptReadBackend: ReadBackend + Clone + 'static { Ok(data.into()) } + /// Decrypt and decompress `data`, then run `f` on the plaintext without + /// allocating a new uncompressed `Vec` on every call. + fn with_decoded_from_partial( + &self, + data: &[u8], + uncompressed_length: Option, + f: impl FnOnce(&[u8]) -> RusticResult, + ) -> RusticResult { + let decrypted = self.decrypt(data)?; + if let Some(length) = uncompressed_length { + let expected = length.get() as usize; + zstd_decompress_with(&decrypted, expected, |plain| { + if plain.len() != expected { + return Err(RusticError::new( + ErrorKind::Internal, + "Length of uncompressed data `{actual_length}` does not match the given length `{expected_length}`.", + ) + .attach_context("expected_length", length.get().to_string()) + .attach_context("actual_length", plain.len().to_string()) + .ask_report()); + } + f(plain) + }) + } else { + f(&decrypted) + } + } + /// Reads the given file with the given offset and length. /// /// # Arguments @@ -190,17 +305,26 @@ pub trait DecryptReadBackend: ReadBackend + Clone + 'static { /// If the files could not be read. fn stream_list(&self, list: Vec, p: &Progress) -> StreamResult { p.set_length(list.len() as u64); - // we use a zero-capacity channel; the loading is typically the bottleneck, not the processing. - let (tx, rx) = bounded(0); + // Index/snapshot files are small; on B2 this is RTT-bound (one GET each). + // Cached backends report fewer workers when the files are already local. + let ids: Vec<_> = list.iter().map(|id| **id).collect(); + let workers = self.prefetch_workers(F::TYPE, &ids); + let (tx, rx) = bounded(workers.saturating_mul(2)); let be = self.clone(); let p = p.clone(); spawn(move || { - _ = list.into_par_iter().try_for_each(|id| { - let file = be.get_file::(&id).map(|file| (id, file)); - p.inc(1); - tx.send(file).ok() // abort as soon as possible if sending fails, i.e. if the receiver is dropped - }); + let work = || { + _ = list.into_par_iter().try_for_each(|id| { + let file = be.get_file::(&id).map(|file| (id, file)); + p.inc(1); + tx.send(file).ok() + }); + }; + match rayon::ThreadPoolBuilder::new().num_threads(workers).build() { + Ok(pool) => pool.install(work), + Err(_) => work(), + } }); Ok(rx) } @@ -643,6 +767,14 @@ impl ReadBackend for DecryptBackend { self.be.list_with_size(tpe) } + fn prefetch_workers(&self, tpe: FileType, ids: &[Id]) -> usize { + self.be.prefetch_workers(tpe, ids) + } + + fn tree_loader_count(&self) -> usize { + self.be.tree_loader_count() + } + fn read_full(&self, tpe: FileType, id: &Id) -> RusticResult { self.be.read_full(tpe, id) } diff --git a/crates/core/src/blob.rs b/crates/core/src/blob.rs index 4ba312f22..8b410c698 100644 --- a/crates/core/src/blob.rs +++ b/crates/core/src/blob.rs @@ -13,8 +13,13 @@ pub(super) mod constants { /// The maximum size of pack-part which is read at once from the backend. /// (needed to limit the memory size used for large backends) pub(crate) const LIMIT_PACK_READ: u32 = 40 * 1024 * 1024; // 40 MiB - /// The maximum size of holes which are still read when repacking - pub(crate) const MAX_HOLESIZE: u32 = 256 * 1024; // 256 kiB + /// Maximum unused gap that is still fetched with the surrounding blobs. + /// + /// 256 KiB was too small for high-latency object stores (B2): every larger + /// hole became another HTTP range GET, and prune/restore issued those + /// sequentially. 4 MiB is about one RTT of extra download on a ~100 Mbps + /// link, which is cheaper than an extra request. + pub(crate) const MAX_HOLESIZE: u32 = 4 * 1024 * 1024; // 4 MiB } /// All [`BlobType`]s which are supported by the repository diff --git a/crates/core/src/blob/tree.rs b/crates/core/src/blob/tree.rs index d2fbe4083..ac1804cb7 100644 --- a/crates/core/src/blob/tree.rs +++ b/crates/core/src/blob/tree.rs @@ -1,21 +1,23 @@ pub mod excludes; pub mod modify; pub mod rewrite; +mod used_blobs; use std::{ borrow::Cow, cmp::Ordering, - collections::{BTreeMap, BTreeSet, BinaryHeap}, + collections::{BTreeMap, BinaryHeap}, ffi::OsStr, mem, path::{Component, Path, PathBuf, Prefix}, str::{self, Utf8Error}, }; -use crossbeam_channel::{Receiver, Sender, bounded, unbounded}; +use crossbeam_channel::{Receiver, Sender, TrySendError, bounded}; use derive_setters::Setters; use ignore::Match; use ignore::overrides::Override; +use rustc_hash::FxHashSet; use serde::{Deserialize, Deserializer}; use serde_derive::Serialize; @@ -52,15 +54,52 @@ pub enum TreeErrorKind { pub(crate) type TreeResult = Result; -pub(super) mod constants { - /// The maximum number of trees that are loaded in parallel - pub(super) const MAX_TREE_LOADER: usize = 4; -} - -pub(crate) type TreeStreamItem = RusticResult<(PathBuf, Tree)>; type NodeStreamItem = RusticResult<(PathBuf, Node)>; impl_blobid!(TreeId, BlobType::Tree); +pub(crate) use used_blobs::UsedBlobsTree; + +/// A tree loaded by [`TreeStreamerOnce`]. +/// +/// Prune uses [`UsedBlobsTree`] so it does not allocate full [`Node`]s. +pub(crate) trait LoadedTree: Send + 'static { + fn load( + be: &BE, + index: &I, + id: TreeId, + ) -> RusticResult + where + Self: Sized; + + fn child_trees(&self, parent: &Path) -> Vec<(PathBuf, TreeId)>; +} + +fn read_tree_bytes( + be: &BE, + index: &I, + id: TreeId, +) -> RusticResult { + index + .get_tree(&id) + .ok_or_else(|| { + RusticError::new( + ErrorKind::Internal, + "Tree ID `{tree_id}` not found in index", + ) + .attach_context("tree_id", id.to_string()) + })? + .read_data(be) +} + +fn tree_json_error(err: serde_json::Error) -> Box { + RusticError::with_source( + ErrorKind::Internal, + "Failed to deserialize tree from JSON.", + err, + ) + .ask_report() +} + #[derive(Default, Serialize, Deserialize, Clone, Debug)] /// A [`Tree`] is a list of [`Node`]s pub struct Tree { @@ -139,27 +178,8 @@ impl Tree { index: &impl ReadGlobalIndex, id: TreeId, ) -> RusticResult { - let data = index - .get_tree(&id) - .ok_or_else(|| { - RusticError::new( - ErrorKind::Internal, - "Tree ID `{tree_id}` not found in index", - ) - .attach_context("tree_id", id.to_string()) - })? - .read_data(be)?; - - let tree = serde_json::from_slice(&data).map_err(|err| { - RusticError::with_source( - ErrorKind::Internal, - "Failed to deserialize tree from JSON.", - err, - ) - .ask_report() - })?; - - Ok(tree) + let data = read_tree_bytes(be, index, id)?; + serde_json::from_slice(&data).map_err(tree_json_error) } /// Creates a new node from a path. @@ -615,19 +635,73 @@ where } } +impl LoadedTree for Tree { + fn load( + be: &BE, + index: &I, + id: TreeId, + ) -> RusticResult { + Self::from_backend(be, index, id) + } + + fn child_trees(&self, parent: &Path) -> Vec<(PathBuf, TreeId)> { + self.nodes + .iter() + .filter_map(|node| { + let id = node.subtree?; + let mut path = parent.to_path_buf(); + path.push(node.name()); + Some((path, id)) + }) + .collect() + } +} + +impl LoadedTree for UsedBlobsTree { + fn load( + be: &BE, + index: &I, + id: TreeId, + ) -> RusticResult { + index + .get_tree(&id) + .ok_or_else(|| { + RusticError::new( + ErrorKind::Internal, + "Tree ID `{tree_id}` not found in index", + ) + .attach_context("tree_id", id.to_string()) + })? + .with_decoded(be, |data| { + used_blobs::parse_used_blobs_tree(data).map_err(tree_json_error) + }) + } + + fn child_trees(&self, _parent: &Path) -> Vec<(PathBuf, TreeId)> { + self.dir_trees + .iter() + .copied() + .map(|id| (PathBuf::new(), id)) + .collect() + } +} + /// [`TreeStreamerOnce`] recursively visits all trees and subtrees, but each tree ID only once /// /// # Type Parameters /// +/// * `T` - The loaded tree type. Defaults to a full [`Tree`]. /// * `P` - The progress indicator #[derive(Debug)] -pub struct TreeStreamerOnce { +pub struct TreeStreamer { /// The visited tree IDs - visited: BTreeSet, + visited: FxHashSet, + /// Depth-first backlog of tree IDs not yet sent to a loader. + backlog: Vec<(PathBuf, TreeId, usize)>, /// The queue to send tree IDs to queue_in: Option>, /// The queue to receive trees from - queue_out: Receiver>, + queue_out: Receiver>, /// The progress indicator p: Progress, /// The number of trees that are not yet finished @@ -636,7 +710,28 @@ pub struct TreeStreamerOnce { finished_ids: usize, } -impl TreeStreamerOnce { +/// Recursively visits all trees and subtrees, but each tree ID only once. +pub type TreeStreamerOnce = TreeStreamer; + +fn ignore_loaded_tree(_: &T) {} + +/// Called from each tree-loader thread after a tree is decoded. +pub(crate) trait OnTreeLoad: Send + 'static { + fn on_load(&mut self, tree: &T); + fn finish(self) + where + Self: Sized, + { + } +} + +impl OnTreeLoad for F { + fn on_load(&mut self, tree: &T) { + self(tree); + } +} + +impl TreeStreamer { /// Creates a new `TreeStreamerOnce`. /// /// # Type Parameters @@ -659,31 +754,54 @@ impl TreeStreamerOnce { ids: Vec, p: Progress, ) -> RusticResult { + Self::new_with_on_load(be, index, ids, p, || ignore_loaded_tree) + } + + /// Like [`Self::new`], but each loader thread gets its own `on_load` from + /// `factory` so prune can fill a thread-local used-id map with no locks. + pub fn new_with_on_load( + be: &BE, + index: &I, + ids: Vec, + p: Progress, + mut factory: F, + ) -> RusticResult + where + BE: DecryptReadBackend, + I: ReadGlobalIndex, + F: FnMut() -> H, + H: OnTreeLoad, + { p.set_length(ids.len() as u64); - let (out_tx, out_rx) = bounded(constants::MAX_TREE_LOADER); - let (in_tx, in_rx) = unbounded(); + let loaders = be.tree_loader_count(); + let (out_tx, out_rx) = bounded(loaders.saturating_mul(4).max(32)); + // Bound the loader input so we do not dump every snapshot root at once. + // Combined with a LIFO backlog this keeps workers on recently discovered + // children (depth-first), which hits the same cached packs. + let (in_tx, in_rx) = bounded(loaders); - for _ in 0..constants::MAX_TREE_LOADER { + for _ in 0..loaders { let be = be.clone(); let index = index.clone(); let in_rx = in_rx.clone(); let out_tx = out_tx.clone(); + let mut on_load = factory(); let _join_handle = std::thread::spawn(move || { for (path, id, count) in in_rx { - if out_tx - .send(Tree::from_backend(&be, &index, id).map(|tree| (path, tree, count))) - .is_err() - { + let loaded = T::load(&be, &index, id).inspect(|tree| on_load.on_load(tree)); + if out_tx.send(loaded.map(|tree| (path, tree, count))).is_err() { break; } } + on_load.finish(); }); } let counter = vec![0; ids.len()]; let mut streamer = Self { - visited: BTreeSet::new(), + visited: FxHashSet::default(), + backlog: Vec::new(), queue_in: Some(in_tx), queue_out: out_rx, p, @@ -692,63 +810,64 @@ impl TreeStreamerOnce { }; for (count, id) in ids.into_iter().enumerate() { - if !streamer - .add_pending(PathBuf::new(), id, count) - .map_err(|err| { - RusticError::with_source( - ErrorKind::Internal, - "Failed to add tree ID `{tree_id}` to unbounded pending queue (`{count}`).", - err, - ) - .attach_context("tree_id", id.to_string()) - .attach_context("count", count.to_string()) - .ask_report() - })? - { + if !streamer.add_pending(PathBuf::new(), id, count) { streamer.p.inc(1); streamer.finished_ids += 1; } } + streamer.fill_queue().map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to send tree IDs to loader queue.", + err, + ) + .ask_report() + })?; Ok(streamer) } - /// Adds a tree ID to the queue. - /// - /// # Arguments - /// - /// * `path` - The path of the tree. - /// * `id` - The ID of the tree. - /// * `count` - The index of the tree. + /// Pushes a tree ID onto the depth-first backlog if it has not been seen. /// /// # Returns /// - /// Whether the tree ID was added to the queue. - /// - /// # Errors - /// - /// * If sending the message fails. - fn add_pending(&mut self, path: PathBuf, id: TreeId, count: usize) -> TreeResult { + /// Whether the tree ID was added. + fn add_pending(&mut self, path: PathBuf, id: TreeId, count: usize) -> bool { if self.visited.insert(id) { - self.queue_in - .as_ref() - .unwrap() - .send((path, id, count)) - .map_err(|err| TreeErrorKind::Channel { - kind: "sending crossbeam message", - source: err.into(), - })?; - self.counter[count] += 1; - Ok(true) + self.backlog.push((path, id, count)); + true } else { - Ok(false) + false } } + + /// Sends backlog items to loaders, last-in first, until the input channel is full. + fn fill_queue(&mut self) -> TreeResult<()> { + let Some(tx) = self.queue_in.as_ref() else { + return Ok(()); + }; + while let Some(job) = self.backlog.pop() { + match tx.try_send(job) { + Ok(()) => {} + Err(TrySendError::Full(job)) => { + self.backlog.push(job); + break; + } + Err(TrySendError::Disconnected(_)) => { + return Err(TreeErrorKind::Channel { + kind: "sending crossbeam message", + source: "loader queue disconnected".into(), + }); + } + } + } + Ok(()) + } } -impl Iterator for TreeStreamerOnce { - type Item = TreeStreamItem; +impl Iterator for TreeStreamer { + type Item = RusticResult<(PathBuf, T)>; fn next(&mut self) -> Option { if self.counter.len() == self.finished_ids { @@ -771,27 +890,19 @@ impl Iterator for TreeStreamerOnce { Ok(Err(err)) => return Some(Err(err)), }; - for node in &tree.nodes { - if let Some(id) = node.subtree { - let mut path = path.clone(); - path.push(node.name()); - match self.add_pending(path.clone(), id, count) { - Ok(_) => {} - Err(err) => { - return Some(Err(err).map_err(|err| { - RusticError::with_source( - ErrorKind::Internal, - "Failed to add tree ID `{tree_id}` to pending queue (`{count}`).", - err, - ) - .attach_context("path", path.display().to_string()) - .attach_context("tree_id", id.to_string()) - .attach_context("count", count.to_string()) - .ask_report() - })); - } - } - } + // Push children last-to-first so the next pop is the first child (DFS). + for (child_path, id) in tree.child_trees(&path).into_iter().rev() { + _ = self.add_pending(child_path, id, count); + } + if let Err(err) = self.fill_queue() { + return Some(Err(err).map_err(|err| { + RusticError::with_source( + ErrorKind::Internal, + "Failed to send tree IDs to loader queue.", + err, + ) + .ask_report() + })); } self.counter[count] -= 1; diff --git a/crates/core/src/blob/tree/used_blobs.rs b/crates/core/src/blob/tree/used_blobs.rs new file mode 100644 index 000000000..de6eefa15 --- /dev/null +++ b/crates/core/src/blob/tree/used_blobs.rs @@ -0,0 +1,248 @@ +//! Deserialize restic trees for prune without materializing full `Node`s. +//! +//! Prune only needs file content blob ids and directory subtree ids. A +//! dedicated serde struct keeps `type` / `content` / `subtree` and ignores +//! names, metadata, and xattrs. + +use serde_derive::Deserialize; + +use crate::blob::{DataId, tree::TreeId}; + +/// Compact tree contents used by prune's used-blob walk. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub(crate) struct UsedBlobsTree { + pub file_blobs: Vec, + pub dir_trees: Vec, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "lowercase")] +enum PruneNodeKind { + File, + Dir, + #[default] + #[serde(other)] + Other, +} + +#[derive(Debug, Deserialize)] +struct PruneNode { + #[serde(rename = "type", default)] + kind: PruneNodeKind, + #[serde(default)] + content: Option>, + #[serde(default)] + subtree: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct PruneTree { + #[serde(default, deserialize_with = "super::deserialize_null_default")] + nodes: Vec, +} + +pub(crate) fn parse_used_blobs_tree(data: &[u8]) -> Result { + let parsed: PruneTree = serde_json::from_slice(data)?; + let mut tree = UsedBlobsTree::default(); + for node in parsed.nodes { + match node.kind { + PruneNodeKind::File => { + if let Some(content) = node.content { + tree.file_blobs.extend(content); + } + } + PruneNodeKind::Dir => { + if let Some(id) = node.subtree { + tree.dir_trees.push(id); + } + } + PruneNodeKind::Other => {} + } + } + Ok(tree) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::node::Node; + use crate::blob::tree::Tree; + + const FILE_ID: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + const TREE_ID: &str = "fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210"; + + #[test] + fn extracts_file_and_dir_ids_and_skips_the_rest() { + let json = format!( + r#"{{ + "nodes": [ + {{ + "name": "foo", + "type": "file", + "mtime": "2020-01-01T00:00:00+00:00", + "mode": 420, + "uid": 1000, + "user": "brad", + "inode": 1, + "size": 3, + "links": 1, + "extended_attributes": [{{"name": "user.foo", "value": "YQ=="}}], + "content": ["{FILE_ID}"] + }}, + {{ + "name": "bar", + "type": "dir", + "subtree": "{TREE_ID}" + }}, + {{ + "name": "link", + "type": "symlink", + "linktarget": "/tmp/x" + }} + ] + }}"# + ); + + let used = parse_used_blobs_tree(json.as_bytes()).unwrap(); + let full: Tree = serde_json::from_slice(json.as_bytes()).unwrap(); + + let full_files: Vec<_> = full + .nodes + .iter() + .filter(|n| matches!(n.node_type, crate::backend::node::NodeType::File)) + .flat_map(|n| n.content.iter().flatten().copied()) + .collect(); + let full_dirs: Vec<_> = full.nodes.iter().filter_map(|n| n.subtree).collect(); + + assert_eq!(used.file_blobs, full_files); + assert_eq!(used.dir_trees, full_dirs); + assert_eq!(used.file_blobs, vec![FILE_ID.parse::().unwrap()]); + assert_eq!(used.dir_trees, vec![TREE_ID.parse::().unwrap()]); + let foo: &Node = &full.nodes[0]; + assert_eq!(foo.name, "foo"); + assert_eq!(foo.meta.extended_attributes.len(), 1); + } + + #[test] + fn null_or_missing_nodes_is_empty() { + assert_eq!( + parse_used_blobs_tree(br#"{"nodes":null}"#).unwrap(), + UsedBlobsTree::default() + ); + assert_eq!( + parse_used_blobs_tree(br#"{}"#).unwrap(), + UsedBlobsTree::default() + ); + assert_eq!( + parse_used_blobs_tree(br#"{"nodes":[]}"#).unwrap(), + UsedBlobsTree::default() + ); + } + + #[test] + fn ignores_unknown_tree_keys() { + let json = format!(r#"{{"extra":1,"nodes":[{{"type":"file","content":["{FILE_ID}"]}}]}}"#); + let used = parse_used_blobs_tree(json.as_bytes()).unwrap(); + assert_eq!(used.file_blobs.len(), 1); + assert!(used.dir_trees.is_empty()); + } + + #[test] + fn hex_ids_accept_uppercase_and_reject_garbage() { + let upper = format!( + r#"{{"nodes":[{{"type":"file","content":["{}"]}}]}}"#, + FILE_ID.to_uppercase() + ); + assert_eq!( + parse_used_blobs_tree(upper.as_bytes()).unwrap().file_blobs, + vec![FILE_ID.parse::().unwrap()] + ); + assert!( + parse_used_blobs_tree(br#"{"nodes":[{"type":"file","content":["zzzz"]}]}"#).is_err() + ); + } + + #[test] + fn skips_escaped_names_and_accepts_content_before_type() { + let json = + format!(r#"{{"nodes":[{{"name":"quo\"te","content":["{FILE_ID}"],"type":"file"}}]}}"#); + let used = parse_used_blobs_tree(json.as_bytes()).unwrap(); + assert_eq!(used.file_blobs, vec![FILE_ID.parse::().unwrap()]); + assert!(used.dir_trees.is_empty()); + } + + #[test] + fn file_content_is_ignored_on_dirs_and_other_types() { + let json = format!( + r#"{{"nodes":[{{"type":"dir","content":["{FILE_ID}"],"subtree":"{TREE_ID}"}},{{"type":"symlink","content":["{FILE_ID}"]}}]}}"# + ); + let used = parse_used_blobs_tree(json.as_bytes()).unwrap(); + assert!(used.file_blobs.is_empty()); + assert_eq!(used.dir_trees, vec![TREE_ID.parse::().unwrap()]); + } + + #[test] + fn escaped_live_references_match_full_tree() { + let json = format!( + r#"{{"nodes":[{{"name":"file","type":"file","content":["{FILE_ID}"]}},{{"name":"dir","type":"dir","subtree":"{TREE_ID}"}}]}}"# + ); + for (plain, escaped) in [ + (r#""nodes""#, r#""n\u006fdes""#), + (r#""type""#, r#""t\u0079pe""#), + (r#""content""#, r#""cont\u0065nt""#), + (r#""subtree""#, r#""subtr\u0065e""#), + (r#""file""#, r#""f\u0069le""#), + (r#""dir""#, r#""d\u0069r""#), + ("012345", r"\u003012345"), + ("fedcba", r"\u0066edcba"), + ] { + let escaped_json = json.replace(plain, escaped); + let full: Tree = serde_json::from_str(&escaped_json).unwrap(); + let used = parse_used_blobs_tree(escaped_json.as_bytes()).unwrap(); + let files: Vec<_> = full + .nodes + .iter() + .flat_map(|n| n.content.iter().flatten().copied()) + .collect(); + let dirs: Vec<_> = full.nodes.iter().filter_map(|n| n.subtree).collect(); + assert_eq!(used.file_blobs, files, "{escaped_json}"); + assert_eq!(used.dir_trees, dirs, "{escaped_json}"); + } + } + + #[test] + fn unknown_or_missing_types_are_skipped() { + assert_eq!( + parse_used_blobs_tree(br#"{"nodes":[{"type":"future_file"}]}"#).unwrap(), + UsedBlobsTree::default() + ); + assert_eq!( + parse_used_blobs_tree(br#"{"nodes":[{"name":"missing type"}]}"#).unwrap(), + UsedBlobsTree::default() + ); + } + + #[test] + fn rejects_invalid_json() { + assert!(parse_used_blobs_tree(br#"{"n\qodes":[]}"#).is_err()); + assert!(parse_used_blobs_tree(br#"{"nodes":[{"type":"f\qile"}]}"#).is_err()); + assert!(parse_used_blobs_tree( + br#"{"nodes":[{"type":"file","content":["\q"]}]}"# + ) + .is_err()); + assert!(parse_used_blobs_tree( + br#"{"nodes":[{"type":"file","type":"symlink"}]}"# + ) + .is_err()); + } + + #[test] + fn skips_long_unescaped_names() { + let name = "n".repeat(80); + let json = format!( + r#"{{"nodes":[{{"name":"{name}","mtime":"2020-01-01T00:00:00+00:00","type":"file","content":["{FILE_ID}"]}}]}}"# + ); + let used = parse_used_blobs_tree(json.as_bytes()).unwrap(); + assert_eq!(used.file_blobs, vec![FILE_ID.parse::().unwrap()]); + } +} diff --git a/crates/core/src/commands/prune.rs b/crates/core/src/commands/prune.rs index 306a5ab42..ad1508cd5 100644 --- a/crates/core/src/commands/prune.rs +++ b/crates/core/src/commands/prune.rs @@ -6,6 +6,7 @@ use std::{ cmp::Ordering, collections::{BTreeMap, BTreeSet}, str::FromStr, + sync::mpsc::{self, Sender}, }; use bytesize::ByteSize; @@ -16,18 +17,18 @@ use itertools::Itertools; use jiff::{Span, Timestamp, Zoned}; use log::{info, warn}; use rayon::prelude::{IntoParallelIterator, ParallelIterator}; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use crate::{ backend::{ FileType, ReadBackend, decrypt::{DecryptReadBackend, DecryptWriteBackend}, - node::NodeType, }, blob::{ BlobId, BlobLocations, BlobType, BlobTypeMap, Initialize, packer::{BlobCopier, CopyPackBlobs, PackSizer}, - tree::TreeStreamerOnce, + tree::{OnTreeLoad, TreeStreamer, UsedBlobsTree}, }, error::{ErrorKind, RusticError, RusticResult}, index::{ @@ -43,6 +44,8 @@ use crate::{ repository::{Open, Repository}, }; +type UsedIdMap = FxHashMap; + pub(super) mod constants { /// Minimum size of an index file to be considered for pruning pub(super) const MIN_INDEX_LEN: usize = 10_000; @@ -584,7 +587,7 @@ pub struct PrunePlan { /// The time the plan was created time: Zoned, /// The ids of the blobs which are used - used_ids: BTreeMap, + used_ids: UsedIdMap, /// The ids of the existing packs existing_packs: BTreeMap, /// The packs which should be repacked @@ -604,7 +607,7 @@ impl PrunePlan { /// * `existing_packs` - The ids of the existing packs /// * `index_files` - The index files fn new( - used_ids: BTreeMap, + used_ids: UsedIdMap, existing_packs: BTreeMap, index_files: Vec<(IndexId, IndexFile)>, ) -> Self { @@ -718,6 +721,14 @@ impl PrunePlan { } p.finish(); + // Pack listing does not need the used-blob set. Start it after the + // index so it does not steal B2 from index GETs, and overlap it with + // the tree walk instead. + let pack_list = { + let be = be.clone(); + std::thread::spawn(move || be.list_with_size(FileType::Pack)) + }; + let (used_ids, total_size) = { let index = GlobalIndex::new_from_index(index_collector.into_index()); let total_size = BlobTypeMap::init(|blob_type| index.total_size(blob_type)); @@ -725,13 +736,20 @@ impl PrunePlan { (used_ids, total_size) }; - // list existing pack files + // list existing pack files (started before the tree walk) let p = repo.progress_spinner("getting packs from repository..."); - let existing_packs: BTreeMap<_, _> = be - .list_with_size(FileType::Pack)? - .into_iter() - .map(|(id, size)| (PackId::from(id), size)) - .collect(); + let existing_packs: BTreeMap<_, _> = match pack_list.join() { + Ok(listed) => listed? + .into_iter() + .map(|(id, size)| (PackId::from(id), size)) + .collect(), + Err(_) => { + return Err(RusticError::new( + ErrorKind::Internal, + "Pack listing thread panicked.", + )); + } + }; p.finish(); let mut pruner = Self::new(used_ids, existing_packs, index_files); @@ -1416,14 +1434,15 @@ pub(crate) fn prune_repository( }) .collect(); - // TODO: repack in parallel - for blobs in blob_chunks { + // Range-GETs for holes in the same pack run in parallel. The + // packer already serializes writes via its channel / lock. + blob_chunks.into_par_iter().try_for_each(|blobs| { if opts.fast_repack { - repacker.copy_fast(blobs, &p)?; + repacker.copy_fast(blobs, &p) } else { - repacker.copy(blobs, &p)?; + repacker.copy(blobs, &p) } - } + })?; Ok(()) })?; _ = tree_repacker.finalize()?; @@ -1491,8 +1510,8 @@ impl PackInfo { /// # Arguments /// /// * `pack` - The `PrunePack` to create the `PackInfo` from - /// * `used_ids` - The `BTreeMap` of used ids - fn from_pack(pack: &PrunePack, used_ids: &mut BTreeMap) -> Self { + /// * `used_ids` - The map of used ids + fn from_pack(pack: &PrunePack, used_ids: &mut UsedIdMap) -> Self { let mut pi = Self { blob_type: pack.blob_type, used_blobs: 0, @@ -1568,6 +1587,27 @@ impl PackInfo { } } +/// Per-loader used-id map. Inserts take no lock; maps are merged after the walk. +struct UsedIdAcc { + map: UsedIdMap, + tx: Sender, +} + +impl OnTreeLoad for UsedIdAcc { + fn on_load(&mut self, tree: &UsedBlobsTree) { + for id in &tree.file_blobs { + _ = self.map.insert(BlobId::from(*id), 0); + } + for id in &tree.dir_trees { + _ = self.map.insert(BlobId::from(*id), 0); + } + } + + fn finish(self) { + _ = self.tx.send(self.map); + } +} + /// Find used blobs in repo and return a map of used ids. /// /// # Arguments @@ -1584,7 +1624,7 @@ fn find_used_blobs( be: &impl DecryptReadBackend, index: &impl ReadGlobalIndex, ignore_snaps: &[SnapshotId], -) -> RusticResult> { +) -> RusticResult { let ignore_snaps: BTreeSet<_> = ignore_snaps.iter().collect(); let p = repo.progress_counter("reading snapshots..."); @@ -1601,31 +1641,27 @@ fn find_used_blobs( .try_collect()?; p.finish(); - let mut ids: BTreeMap<_, _> = snap_trees + let mut ids: UsedIdMap = snap_trees .iter() .map(|id| (BlobId::from(**id), 0)) .collect(); let p = repo.progress_counter("finding used blobs..."); - - let mut tree_streamer = TreeStreamerOnce::new(be, index, snap_trees, p)?; - while let Some(item) = tree_streamer.next().transpose()? { - let (_, tree) = item; - for node in tree.nodes { - match node.node_type { - NodeType::File => { - ids.extend( - node.content - .iter() - .flatten() - .map(|id| (BlobId::from(**id), 0)), - ); - } - NodeType::Dir => { - _ = ids.insert(BlobId::from(*node.subtree.unwrap()), 0); - } - _ => {} // nothing to do + let (maps_tx, maps_rx) = mpsc::channel(); + let mut tree_streamer = + TreeStreamer::::new_with_on_load(be, index, snap_trees, p, { + let maps_tx = maps_tx.clone(); + move || UsedIdAcc { + map: UsedIdMap::default(), + tx: maps_tx.clone(), } - } + })?; + drop(maps_tx); + while let Some(item) = tree_streamer.next().transpose()? { + let _ = item; + } + drop(tree_streamer); + while let Ok(map) = maps_rx.recv() { + ids.extend(map); } Ok(ids) diff --git a/crates/core/src/index.rs b/crates/core/src/index.rs index 95f981c89..ff45ceb67 100644 --- a/crates/core/src/index.rs +++ b/crates/core/src/index.rs @@ -65,6 +65,25 @@ impl IndexEntry { Ok(data) } + /// Decrypt and decompress this blob, then run `f` on the plaintext. + /// + /// Prune tree walking uses this so zstd output can stay in a thread-local + /// buffer instead of a new `Vec`/`Bytes` per tree. + pub fn with_decoded(&self, be: &B, f: F) -> RusticResult + where + B: DecryptReadBackend, + F: FnOnce(&[u8]) -> RusticResult, + { + let cipher = be.read_partial( + FileType::Pack, + &self.pack, + self.blob_type.is_cacheable(), + self.location.offset, + self.location.length, + )?; + be.with_decoded_from_partial(&cipher, self.location.uncompressed_length, f) + } + /// Get the length of the data described by the [`IndexEntry`] #[must_use] pub const fn data_length(&self) -> u32 { diff --git a/crates/core/src/index/binarysorted.rs b/crates/core/src/index/binarysorted.rs index f399da24c..4e9e41a8d 100644 --- a/crates/core/src/index/binarysorted.rs +++ b/crates/core/src/index/binarysorted.rs @@ -1,3 +1,5 @@ +use std::cmp::Ordering; + use rayon::prelude::*; use crate::{ @@ -20,6 +22,77 @@ pub(crate) struct SortedEntry { location: BlobLocation, } +/// Max entries in one collector chunk. +/// +/// Growing a single `Vec` of blob ids doubles it. On a large repo that request +/// is hundreds of MiB while the old buffer is still live, and musl aborts: +/// `memory allocation of N bytes failed`. Chunks cap each allocation. +const ENTRY_CHUNK_LEN: usize = if cfg!(test) { 4 } else { 1 << 20 }; + +/// Append-only vec of bounded chunks. Lookups binary-search every chunk. +#[derive(Debug)] +pub(crate) struct Chunked { + chunks: Vec>, +} + +impl Default for Chunked { + fn default() -> Self { + Self { chunks: Vec::new() } + } +} + +impl Chunked { + fn push(&mut self, item: T) { + if self + .chunks + .last() + .is_none_or(|chunk| chunk.len() >= ENTRY_CHUNK_LEN) + { + self.chunks.push(Vec::with_capacity(ENTRY_CHUNK_LEN)); + } + self.chunks + .last_mut() + .expect("chunk is created above") + .push(item); + } + + fn shrink_last(&mut self) { + if let Some(last) = self.chunks.last_mut() { + last.shrink_to_fit(); + } + } + + fn par_sort_unstable(&mut self) + where + T: Ord + Send, + { + self.chunks.par_iter_mut().for_each(|chunk| { + chunk.sort_unstable(); + }); + } + + fn par_sort_unstable_by(&mut self, compare: F) + where + T: Send, + F: Fn(&T, &T) -> Ordering + Sync, + { + self.chunks.par_iter_mut().for_each(|chunk| { + chunk.sort_unstable_by(&compare); + }); + } + + fn par_sort_unstable_by_key(&mut self, f: F) + where + T: Send, + K: Ord, + F: Fn(&T) -> K + Sync, + { + self.chunks.par_iter_mut().for_each(|chunk| { + chunk.sort_unstable_by_key(&f); + }); + } +} + /// `IndexType` determines which information is stored in the index. #[derive(Debug, Clone, Copy)] pub enum IndexType { @@ -36,8 +109,8 @@ pub enum IndexType { pub(crate) enum EntriesVariants { #[default] None, - Ids(Vec), - FullEntries(Vec), + Ids(Chunked), + FullEntries(Chunked), } #[derive(Default, Debug)] @@ -54,7 +127,8 @@ pub struct IndexCollector(BlobTypeMap); pub struct PackIndexes { c: Index, tpe: BlobType, - idx: BlobTypeMap<(u32, usize)>, + pack_idx: BlobTypeMap, + cursors: BlobTypeMap>, } #[derive(Debug)] @@ -89,11 +163,11 @@ impl IndexCollector { pub fn new(tpe: IndexType) -> Self { let mut collector = Self::default(); - collector.0[BlobType::Tree].entries = EntriesVariants::FullEntries(Vec::new()); + collector.0[BlobType::Tree].entries = EntriesVariants::FullEntries(Chunked::default()); collector.0[BlobType::Data].entries = match tpe { IndexType::OnlyTrees => EntriesVariants::None, - IndexType::DataIds => EntriesVariants::Ids(Vec::new()), - IndexType::Full => EntriesVariants::FullEntries(Vec::new()), + IndexType::DataIds => EntriesVariants::Ids(Chunked::default()), + IndexType::Full => EntriesVariants::FullEntries(Chunked::default()), }; collector @@ -110,8 +184,14 @@ impl IndexCollector { Index(self.0.map(|_, mut tc| { match &mut tc.entries { EntriesVariants::None => {} - EntriesVariants::Ids(ids) => ids.par_sort_unstable(), - EntriesVariants::FullEntries(entries) => entries.par_sort_unstable_by_key(|e| e.id), + EntriesVariants::Ids(ids) => { + ids.shrink_last(); + ids.par_sort_unstable(); + } + EntriesVariants::FullEntries(entries) => { + entries.shrink_last(); + entries.par_sort_unstable_by_key(|e| e.id); + } } let packs = tc.packs.into_iter().map(|(id, _)| id).collect(); @@ -130,7 +210,6 @@ impl Extend for IndexCollector { T: IntoIterator, { for p in iter { - let len = p.blobs.len(); let blob_type = p.blob_type(); let size = p.pack_size(); @@ -140,12 +219,6 @@ impl Extend for IndexCollector { self.0[blob_type].total_size += u64::from(size); - match &mut self.0[blob_type].entries { - EntriesVariants::None => {} - EntriesVariants::Ids(idents) => idents.reserve(len), - EntriesVariants::FullEntries(entries) => entries.reserve(len), - } - for blob in &p.blobs { let be = SortedEntry { id: blob.id, @@ -166,39 +239,45 @@ impl Iterator for PackIndexes { type Item = IndexPack; fn next(&mut self) -> Option { - let (pack_idx, idx) = loop { - let (pack_idx, idx) = &mut self.idx[self.tpe]; - let pack_count = u32::try_from(self.c.0[self.tpe].packs.len()) - .expect("pack count should fit into u32"); - if *pack_idx >= pack_count { - if self.tpe == BlobType::Data { + loop { + let tpe = self.tpe; + let pack_count = + u32::try_from(self.c.0[tpe].packs.len()).expect("pack count should fit into u32"); + if self.pack_idx[tpe] >= pack_count { + if tpe == BlobType::Data { return None; } self.tpe = BlobType::Data; - } else { - break (pack_idx, idx); + continue; } - }; - let mut pack = IndexPack { - id: self.c.0[self.tpe].packs[*pack_idx as usize], - ..Default::default() - }; + let pack_idx = self.pack_idx[tpe]; + let mut pack = IndexPack { + id: self.c.0[tpe].packs[pack_idx as usize], + ..Default::default() + }; - if let EntriesVariants::FullEntries(entries) = &self.c.0[self.tpe].entries { - while *idx < entries.len() && entries[*idx].pack_idx == *pack_idx { - let entry = &entries[*idx]; - pack.blobs.push(IndexBlob { - id: entry.id, - tpe: self.tpe, - location: entry.location, - }); - *idx += 1; + if let EntriesVariants::FullEntries(entries) = &self.c.0[tpe].entries { + let cursors = &mut self.cursors[tpe]; + if cursors.len() != entries.chunks.len() { + cursors.resize(entries.chunks.len(), 0); + } + for (chunk, cursor) in entries.chunks.iter().zip(cursors.iter_mut()) { + while *cursor < chunk.len() && chunk[*cursor].pack_idx == pack_idx { + let entry = &chunk[*cursor]; + pack.blobs.push(IndexBlob { + id: entry.id, + tpe, + location: entry.location, + }); + *cursor += 1; + } + } } - } - *pack_idx += 1; - Some(pack) + self.pack_idx[tpe] += 1; + return Some(pack); + } } } @@ -214,34 +293,32 @@ impl IntoIterator for Index { } } PackIndexes { - c: Self(self.0.map(|_, mut tc| { - if let EntriesVariants::FullEntries(entries) = &mut tc.entries { - entries.par_sort_unstable_by(|e1, e2| e1.pack_idx.cmp(&e2.pack_idx)); - } - - tc - })), + c: self, tpe: BlobType::Tree, - idx: BlobTypeMap::default(), + pack_idx: BlobTypeMap::default(), + cursors: BlobTypeMap::default(), } } } impl ReadIndex for Index { fn get_id(&self, blob_type: BlobType, id: &BlobId) -> Option { - let EntriesVariants::FullEntries(vec) = &self.0[blob_type].entries else { + let EntriesVariants::FullEntries(entries) = &self.0[blob_type].entries else { // get_id() only gives results if index contains full entries return None; }; - vec.binary_search_by_key(id, |e| e.id).ok().map(|index| { - let be = &vec[index]; - IndexEntry::new( - blob_type, - self.0[blob_type].packs[be.pack_idx as usize], - be.location, - ) - }) + for chunk in &entries.chunks { + if let Ok(index) = chunk.binary_search_by_key(id, |e| e.id) { + let be = &chunk[index]; + return Some(IndexEntry::new( + blob_type, + self.0[blob_type].packs[be.pack_idx as usize], + be.location, + )); + } + } + None } fn total_size(&self, blob_type: BlobType) -> u64 { @@ -250,10 +327,14 @@ impl ReadIndex for Index { fn has(&self, blob_type: BlobType, id: &BlobId) -> bool { match &self.0[blob_type].entries { - EntriesVariants::FullEntries(entries) => { - entries.binary_search_by_key(id, |e| e.id).is_ok() - } - EntriesVariants::Ids(ids) => ids.binary_search(id).is_ok(), + EntriesVariants::FullEntries(entries) => entries + .chunks + .iter() + .any(|chunk| chunk.binary_search_by_key(id, |e| e.id).is_ok()), + EntriesVariants::Ids(ids) => ids + .chunks + .iter() + .any(|chunk| chunk.binary_search(id).is_ok()), // has() only gives results if index contains full entries or ids EntriesVariants::None => false, } @@ -439,6 +520,53 @@ mod tests { ); assert!(!index.has(BlobType::Tree, &id)); assert!(index.get_id(BlobType::Tree, &id).is_none()); + + // This id is in the second test-sized chunk (ENTRY_CHUNK_LEN is 4 under cfg(test)). + let id = "ee67585c7c53324e74537ab7aa44f889c0767c1b67e7e336fae6204aef2d4c73".parse()?; + assert!(index.has(BlobType::Data, &id)); + assert_eq!( + index.get_id(BlobType::Data, &id), + Some(IndexEntry { + blob_type: BlobType::Data, + pack: "3b25ec6d16401c31099c259311562160b1b5efbcf70bd69d0463104d3b8148fc".parse()?, + location: BlobLocation { + offset: 7737, + length: 7686, + uncompressed_length: Some(NonZeroU32::new(29928).unwrap()), + } + }), + ); + Ok(()) + } + + #[test] + fn into_iter_groups_blobs_by_pack() -> RusticResult<()> { + let packs: Vec<_> = index(IndexType::Full).into_iter().collect(); + assert_eq!(packs.len(), 3); + assert_eq!( + packs[0].id, + "8431a27d38dd7d192dc37abd43a85d6dc4298de72fc8f583c5d7cdd09fa47274".parse()? + ); + assert_eq!(packs[0].blobs.len(), 2); + assert_eq!( + packs[1].id, + "217f145b63fbc10267f5a686186689ea3389bed0d6a54b50ffc84d71f99eb7fa".parse()? + ); + assert_eq!(packs[1].blobs.len(), 3); + assert_eq!( + packs[2].id, + "3b25ec6d16401c31099c259311562160b1b5efbcf70bd69d0463104d3b8148fc".parse()? + ); + assert_eq!(packs[2].blobs.len(), 4); + Ok(()) + } + + #[test] + fn data_ids_has_across_chunks() -> RusticResult<()> { + let index = index(IndexType::DataIds); + let id = "f2ca1bb6c7e907d06dafe4687e579fce76b37e4e93b7605022da52e6ccc26fd2".parse()?; + assert!(index.has(BlobType::Data, &id)); + assert!(index.get_id(BlobType::Data, &id).is_none()); Ok(()) } }