Something went wrong. Try again.
Restore a rocks database from object storage
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447//! Push a rocks checkpoint to a backup-engine-compatible backup//!//! The caller makes a [checkpoint](https://github.com/facebook/rocksdb/wiki/Checkpoints)//! (cheap hard links when on same fs as db) and hands us the directory; we//! assemble the `meta/`, `private/`, and `shared_checksum/` layout that//! rocks' `BackupEngine` writes.//!//! a `meta/.<id>.tmp` file marks the in-flight push, `meta/<id>` commits it.//!//! while there are a few attempts to avoid some concurrent-writer issues, eat-//! rocks is not meant to be run with concurrent writers on the same bucket and//! prefix.
use std::collections::HashMap;use std::num::NonZeroUsize;use std::path::{Path, PathBuf};use std::sync::Arc;use std::time::Instant;
use futures::StreamExt;use object_store::buffered::BufWriter;use object_store::limit::LimitStore;use object_store::path::Path as StorePath;use object_store::{ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload};use tokio::io::{AsyncReadExt, AsyncWriteExt};use tracing::{info, warn};
use crate::meta::{BackupFile, BackupMeta, MetaSchema};use crate::restore::crc32c_local_file;use crate::retention::{self, RetentionOptions};use crate::sst;
/// default max concurrent object store operations for a pushpub const DEFAULT_PUSH_CONCURRENCY: usize = 16;
/// multipart upload part sizeconst UPLOAD_PART_SIZE: usize = 8 * 2_usize.pow(20);
/// rocks caps app metadata at 1 MiB raw (`kMaxAppMetaSize`)const MAX_APP_METADATA: usize = 2_usize.pow(20);
/// How `shared_checksum/` files are named.#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]pub enum SharedNaming { /// `<number>_s<db_session_id>_<size>.sst`, rocks' modern default. the /// session id comes from a few KB of SST table properties, so unchanged /// files can be recognized (and their recorded checksums reused) without /// reading them. falls back per-file to [`SharedNaming::LegacyCrc`] when /// no session id is readable, like rocks does. blob files always use legacy /// naming (the blob format carries no session id). #[default] SessionId, /// `<number>_<crc32c>_<size>.sst`, rocks' legacy scheme. naming any file /// requires checksumming it (full read of every shared file on any push). LegacyCrc,}
#[derive(Debug, Clone)]pub struct PushOptions { /// max concurrent object store operations for push /// /// default: [`DEFAULT_PUSH_CONCURRENCY`] pub concurrency: usize, /// sequence number recorded in the meta file, informational. /// /// rocks documents it as "approximate, should not be used by clients", /// never reads it during restore. /// /// default: 0 pub sequence_number: u64, /// unix-seconds timestamp to record in the meta file /// /// default: `None` (current system time) pub timestamp: Option<u64>, /// arbitrary metadata to record in the meta file (max 1 MiB) /// /// default: None pub app_metadata: Option<Vec<u8>>, /// after a successful push, delete old backups, keeping N most recent /// /// default: None (keep everything) pub retain: Option<NonZeroUsize>, /// delete files as soon as they are pushed, then the checkpoint dir itself /// /// if the push fails partway, some files may already be gone and the /// checkpoint will be left incomplete, so a fresh checkpoint is needed to /// retry. /// /// default: false (the checkpoint is left alone) pub consume_checkpoint: bool, /// meta file schema to write /// /// default: [`MetaSchema::V1`] (matches current rocks) pub schema: MetaSchema, /// shared filename scheme /// /// default: [`SharedNaming::SessionId`] pub naming: SharedNaming,}
impl Default for PushOptions { fn default() -> Self { Self { concurrency: DEFAULT_PUSH_CONCURRENCY, sequence_number: 0, timestamp: None, app_metadata: None, retain: None, consume_checkpoint: false, schema: MetaSchema::default(), naming: SharedNaming::default(), } }}
#[derive(Debug)]pub struct PushOutcome { pub backup_id: u64, pub total_files: usize, pub uploaded_files: usize, /// shared files recognized from prior backups pub reused_files: usize, pub uploaded_bytes: u64, pub reused_bytes: u64, /// backups deleted by `retain` pub removed: Vec<u64>, pub elapsed: std::time::Duration,}
/// Backup push possible errors#[derive(Debug, thiserror::Error)]pub enum PushError { #[error("{}: {source}", path.display())] Io { path: PathBuf, #[source] source: std::io::Error, },
#[error( "unexpected file in checkpoint dir: {name:?}: a checkpoint contains only \ CURRENT, MANIFEST-*, OPTIONS-*, *.log, *.sst, *.blob. (a live db dir has \ IDENTITY/LOCK/LOG files)" )] UnexpectedCheckpointFile { name: String },
#[error("checkpoint dirs are flat but found subdirectory: {name:?}")] CheckpointSubdirectory { name: String },
#[error("refusing symlink in checkpoint dir: {name:?}")] CheckpointSymlink { name: String },
#[error("non-utf8 filename in checkpoint dir: {name:?}")] NonUtf8Filename { name: String },
#[error("checkpoint has no CURRENT file: {}", dir.display())] MissingCurrent { dir: PathBuf },
#[error("CURRENT should contain exactly `MANIFEST-<n>\\n`, got {content:?}")] InvalidCurrent { content: String },
#[error("CURRENT points at {current:?} but that MANIFEST is not in the checkpoint")] CurrentManifestMismatch { current: String },
#[error("checkpoint has {count} MANIFEST files, expected exactly 1")] ManifestCount { count: usize },
#[error("app metadata is {size} bytes; rocks caps it at {MAX_APP_METADATA}")] MetadataTooLarge { size: usize },
#[error("failed to list objects under {prefix:?}")] List { prefix: StorePath, #[source] source: object_store::Error, },
#[error("failed to read prior backup meta {backup_id} (needed for checksum reuse)")] MetaFetch { backup_id: u64, #[source] source: Box<crate::Error>, },
#[error("failed to check for existing {key:?}")] Head { key: StorePath, #[source] source: object_store::Error, },
#[error("failed to upload {key:?}")] Upload { key: StorePath, #[source] source: std::io::Error, },
#[error("failed to write {key:?}")] MetaPut { key: StorePath, #[source] source: object_store::Error, },
#[error( "backup id {id} is already claimed ({key:?} exists) — is another push \ running against this prefix?" )] BackupIdClaimed { id: u64, key: StorePath },
#[error("{} changed size during push: expected {expected} bytes, read {actual}", path.display())] SizeChanged { path: PathBuf, expected: u64, actual: u64, },
#[error("{} changed content during push (crc32c mismatch on re-read)", path.display())] ChecksumChanged { path: PathBuf },
#[error("file hashing task for {} failed: {source}", path.display())] BlockingHash { path: PathBuf, #[source] source: tokio::task::JoinError, },
#[error("failed to serialize backup meta")] Serialize(#[from] crate::meta::SerializeError),
#[error( "backup {backup_id} pushed ok, but deleting all but the newest \ {keep} backups afterwards failed" )] Retain { backup_id: u64, keep: usize, #[source] source: Box<retention::RetentionError>, },
#[error( "failed to delete {} after pushing it (consume_checkpoint). the \ checkpoint is now partly consumed: make a fresh one to retry from", path.display() )] ConsumeFile { path: PathBuf, #[source] source: std::io::Error, },}
/// what kind of checkpoint file, and where it goes in the backup#[derive(Debug, Clone, Copy, PartialEq, Eq)]enum FileKind { /// `<digits>.sst` goes in `shared_checksum/` SharedSst, /// `<digits>.blob` goes in `shared_checksum/`, legacy-named SharedBlob, /// `CURRENT`, `MANIFEST-*`, `OPTIONS-*`, `*.log` go in `private/<id>/` Private,}
fn classify(name: &str) -> Option<FileKind> { let digits = crate::restore::is_digit_string; if name == "CURRENT" { return Some(FileKind::Private); } if let Some(rest) = name.strip_prefix("MANIFEST-") && digits(rest) { return Some(FileKind::Private); } if let Some(rest) = name.strip_prefix("OPTIONS-") && digits(rest) { return Some(FileKind::Private); } if let Some(stem) = name.strip_suffix(".log") && digits(stem) { return Some(FileKind::Private); } if let Some(stem) = name.strip_suffix(".sst") && digits(stem) { return Some(FileKind::SharedSst); } if let Some(stem) = name.strip_suffix(".blob") && digits(stem) { return Some(FileKind::SharedBlob); } None}
#[derive(Debug)]struct LocalFile { /// filename within the checkpoint dir name: String, path: PathBuf, size: u64, kind: FileKind,}
/// scan + validate the checkpoint directory (strictly: unknown files error)async fn scan_checkpoint(dir: &Path) -> Result<Vec<LocalFile>, PushError> { let io = |source, path: &Path| PushError::Io { path: path.to_path_buf(), source, };
let mut rd = tokio::fs::read_dir(dir).await.map_err(|e| io(e, dir))?; let mut files = Vec::new(); while let Some(entry) = rd.next_entry().await.map_err(|e| io(e, dir))? { let name = entry .file_name() .to_str() .ok_or_else(|| PushError::NonUtf8Filename { name: entry.file_name().to_string_lossy().into_owned(), })? .to_string(); // symlink_metadata: never follow links out of the checkpoint let meta = tokio::fs::symlink_metadata(entry.path()) .await .map_err(|e| io(e, &entry.path()))?; if meta.is_symlink() { return Err(PushError::CheckpointSymlink { name }); } if meta.is_dir() { return Err(PushError::CheckpointSubdirectory { name }); } let kind = classify(&name) .ok_or_else(|| PushError::UnexpectedCheckpointFile { name: name.clone() })?; files.push(LocalFile { path: entry.path(), size: meta.len(), name, kind, }); }
// sanity: exactly one MANIFEST, and CURRENT names it let manifests: Vec<&str> = files .iter() .filter(|f| f.name.starts_with("MANIFEST-")) .map(|f| f.name.as_str()) .collect(); if manifests.len() != 1 { return Err(PushError::ManifestCount { count: manifests.len(), }); } let current = files .iter() .find(|f| f.name == "CURRENT") .ok_or_else(|| PushError::MissingCurrent { dir: dir.to_path_buf(), })?; let content = tokio::fs::read_to_string(¤t.path) .await .map_err(|e| io(e, ¤t.path))?; let named = content .strip_suffix('\n') .filter(|n| { n.strip_prefix("MANIFEST-") .is_some_and(crate::restore::is_digit_string) }) .ok_or_else(|| PushError::InvalidCurrent { content: content.clone(), })?; if named != manifests[0] { return Err(PushError::CurrentManifestMismatch { current: named.to_string(), }); }
Ok(files)}
/// mangled shared filename: `000009.sst` + session id + size -> `000009_sXX_59.sst`fn session_name(name: &str, sid: &str, size: u64) -> String { let (stem, ext) = name .rsplit_once('.') .expect("classified shared file has extension"); format!("{stem}_s{sid}_{size}.{ext}")}
/// mangled shared filename, legacy scheme: `000009_<crc32c>_<size>.sst`fn legacy_name(name: &str, crc: u32, size: u64) -> String { let (stem, ext) = name .rsplit_once('.') .expect("classified shared file has extension"); format!("{stem}_{crc}_{size}.{ext}")}
/// a shared file's fate for this push#[derive(Debug)]enum SharedAction { /// exists already and referenced by a committed meta: reuse don't re-upload Reuse { crc: u32 }, /// upload: new file, unreferenced orphan at that name, or repair of a /// referenced-but-missing object Upload { known_crc: Option<u32> },}
#[derive(Debug)]struct PlannedFile { local: LocalFile, /// path relative to the backup root, e.g. `shared_checksum/00_s7X_59.sst` backup_path: String, action: SharedAction,}
/// decide a shared file's backup name and whether it needs uploading.////// mirrors rocks' `AddBackupFileWorkItem` (`backup_engine.cc:2540-2718`):/// session-id naming avoids reading the file; a name that's both present/// remotely and referenced by a committed meta is reused; an *unreferenced*/// object at the target name gets overwritten (rocks deletes + re-copies).////// difference from rock: for legacy names a present-but-unreferenced object is/// identical since its crc matches, so we skip the re-upload. rocks re-copies/// in that case.async fn plan_shared( store: Arc<dyn ObjectStore>, store_prefix: StorePath, local: LocalFile, naming: SharedNaming, reused_crcs: Arc<HashMap<String, u32>>,) -> Result<PlannedFile, PushError> { let use_session = naming == SharedNaming::SessionId && local.kind == FileKind::SharedSst; let session_id = if use_session { sst::db_session_id(&local.path).await } else { None };
let (mangled, action) = match session_id { Some(sid) => { let mangled = session_name(&local.name, &sid, local.size); match reused_crcs.get(&mangled) { Some(&crc) => { // referenced by a committed backup; confirm the object is // really there before trusting it (rocks checks existence // too, and re-copies referenced-but-missing files) let key = full_key(&store_prefix, &format!("shared_checksum/{mangled}")); match store.head(&key).await { Ok(_) => (mangled, SharedAction::Reuse { crc }), Err(object_store::Error::NotFound { .. }) => { warn!(key = %key, "shared file referenced by a backup but missing; re-uploading"); (mangled, SharedAction::Upload { known_crc: None }) } Err(source) => return Err(PushError::Head { key, source }), } } None => (mangled, SharedAction::Upload { known_crc: None }), } } None => { // legacy naming needs the crc32c — this is the full-read path let crc = crc32c_local_file(&local.path) .await .map_err(from_hash_error)?; let mangled = legacy_name(&local.name, crc, local.size); let key = full_key(&store_prefix, &format!("shared_checksum/{mangled}")); match store.head(&key).await { // name embeds crc32c and size: existence certifies content Ok(_) => (mangled, SharedAction::Reuse { crc }), Err(object_store::Error::NotFound { .. }) => ( mangled, SharedAction::Upload { known_crc: Some(crc), }, ), Err(source) => return Err(PushError::Head { key, source }), } } };
Ok(PlannedFile { backup_path: format!("shared_checksum/{mangled}"), local, action, })}
/// crc32c_local_file returns crate::Error::Io or ::BlockingHash; translatefn from_hash_error(e: crate::Error) -> PushError { match e { crate::Error::Io { path, source } => PushError::Io { path, source }, crate::Error::BlockingHash { path, source } => PushError::BlockingHash { path, source }, other => unreachable!("crc32c_local_file only returns Io or BlockingHash: {other}"), }}
fn full_key(store_prefix: &StorePath, backup_path: &str) -> StorePath { let mut key = store_prefix.clone(); key.extend(&StorePath::from(backup_path)); key}
/// stream a local file into the store, returning its crc32c.////// multipart above [`UPLOAD_PART_SIZE`]; the multipart upload is aborted on/// any failure so nothing half-finished is left counting against storage.async fn upload_file( store: Arc<dyn ObjectStore>, key: StorePath, local_path: PathBuf, expected_size: u64,) -> Result<u32, PushError> { let mut writer = BufWriter::with_capacity(store, key.clone(), UPLOAD_PART_SIZE).with_max_concurrency(4);
let result = stream_into(&mut writer, &key, &local_path, expected_size).await; match result { Ok(crc) => { writer .shutdown() .await .map_err(|source| PushError::Upload { key: key.clone(), source, })?; Ok(crc) } Err(e) => { if let Err(abort_err) = writer.abort().await { warn!(key = %key, error = %abort_err, "failed to abort multipart upload"); } Err(e) } }}
async fn stream_into( writer: &mut BufWriter, key: &StorePath, local_path: &Path, expected_size: u64,) -> Result<u32, PushError> { let mut f = tokio::fs::File::open(local_path) .await .map_err(|source| PushError::Io { path: local_path.to_path_buf(), source, })?; let mut buf = vec![0u8; 256 * 1024]; let mut crc = 0u32; let mut total = 0u64; loop { let n = f.read(&mut buf).await.map_err(|source| PushError::Io { path: local_path.to_path_buf(), source, })?; if n == 0 { break; } crc = crc32c::crc32c_append(crc, &buf[..n]); total += n as u64; writer .write_all(&buf[..n]) .await .map_err(|source| PushError::Upload { key: key.clone(), source, })?; } // shared names embed the size we statted; checkpoint files are immutable // so any drift means the "checkpoint" is being written to (or isn't one) if total != expected_size { return Err(PushError::SizeChanged { path: local_path.to_path_buf(), expected: expected_size, actual: total, }); } Ok(crc)}
/// PUT a meta/marker object, requiring the key to not exist yetasync fn put_new( store: &Arc<dyn ObjectStore>, id: u64, key: &StorePath, content: String,) -> Result<(), PushError> { let opts = PutOptions::from(PutMode::Create); match store.put_opts(key, PutPayload::from(content), opts).await { Ok(_) => Ok(()), Err(object_store::Error::AlreadyExists { .. }) => Err(PushError::BackupIdClaimed { id, key: key.clone(), }), Err(source) => Err(PushError::MetaPut { key: key.clone(), source, }), }}
/// Push a checkpoint directory to object storage as backup `max(existing)+1`.////// see the module docs and [`PushOptions`]. the returned/// [`PushOutcome::backup_id`] is committed (visible to `list`/`restore`) when/// this returns `Ok`.pub async fn push( store: Arc<dyn ObjectStore>, prefix: &str, checkpoint_dir: impl AsRef<Path>, opts: PushOptions,) -> Result<PushOutcome, PushError> { let checkpoint_dir = checkpoint_dir.as_ref(); let started = Instant::now();
let metadata_hex = match &opts.app_metadata { Some(bytes) if bytes.len() > MAX_APP_METADATA => { return Err(PushError::MetadataTooLarge { size: bytes.len() }); } Some(bytes) => Some(hex::encode(bytes)), None => None, }; let timestamp = opts.timestamp.unwrap_or_else(|| { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .expect("system clock is set before the unix epoch") .as_secs() });
let files = scan_checkpoint(checkpoint_dir).await?; let total_files = files.len();
// keep an unwrapped handle for the retain-purge (it applies its own limit) let raw_store = Arc::clone(&store); let store: Arc<dyn ObjectStore> = Arc::new(LimitStore::new(store, opts.concurrency)); let store_prefix = StorePath::from(prefix);
// -- allocate an id nothing else has touched ---------------------------- let ids = retention::scan_backup_ids(&*store, prefix) .await .map_err(retire_to_push)?; let new_id = ids.highest_seen().map_or(1, |max| max + 1);
// -- reusable checksums from every committed backup --------------------- // (shared_checksum names embedding a session id can't be checksummed from // the name alone; prior metas vouch for them. `backup_engine.cc:2676` // does the same from its in-memory file registry.) let mut reused_crcs: HashMap<String, u32> = HashMap::new(); { let metas = futures::stream::iter(ids.committed.iter().map(|&id| { let store = Arc::clone(&store); let prefix = prefix.to_string(); async move { let meta = crate::fetch_meta(&*store, &prefix, id).await; (id, meta) } })) .buffer_unordered(opts.concurrency) .collect::<Vec<_>>() .await; for (id, result) in metas { let meta = result.map_err(|e| PushError::MetaFetch { backup_id: id, source: Box::new(e), })?; for f in meta.files { if let Some(crc) = f.crc32c && let Some(name) = f.path.strip_prefix("shared_checksum/") { reused_crcs.insert(name.to_string(), crc); } } } } let reused_crcs = Arc::new(reused_crcs);
info!( backup_id = new_id, files = total_files, prior_backups = ids.committed.len(), "planning push" );
// -- name shared files + decide reuse/upload ---------------------------- let (shared, private): (Vec<_>, Vec<_>) = files.into_iter().partition(|f| f.kind != FileKind::Private);
let mut planned: Vec<PlannedFile> = futures::stream::iter(shared.into_iter().map(|f| { plan_shared( Arc::clone(&store), store_prefix.clone(), f, opts.naming, Arc::clone(&reused_crcs), ) })) .buffer_unordered(opts.concurrency) .collect::<Vec<Result<_, _>>>() .await .into_iter() .collect::<Result<_, _>>()?;
for f in private { planned.push(PlannedFile { backup_path: format!("private/{new_id}/{}", f.name), local: f, action: SharedAction::Upload { known_crc: None }, }); }
let reused_files = planned .iter() .filter(|p| matches!(p.action, SharedAction::Reuse { .. })) .count(); let reused_bytes: u64 = planned .iter() .filter(|p| matches!(p.action, SharedAction::Reuse { .. })) .map(|p| p.local.size) .sum(); let upload_bytes: u64 = planned .iter() .filter(|p| matches!(p.action, SharedAction::Upload { .. })) .map(|p| p.local.size) .sum(); info!( backup_id = new_id, reused_files, upload_files = planned.len() - reused_files, upload_mb = format_args!("{:.1}", upload_bytes as f64 / 2_f64.powf(20.)), "upload plan ready" );
// -- claim the id + advertise our file list to gc ------------------------ // full path list, checksums where already known; the final meta (the // commit point) replaces this after uploads complete. planned.sort_by(|a, b| a.backup_path.cmp(&b.backup_path)); let mut meta = BackupMeta { timestamp, sequence_number: opts.sequence_number, metadata: metadata_hex, files: planned .iter() .map(|p| BackupFile { path: p.backup_path.clone(), crc32c: match p.action { SharedAction::Reuse { crc } => Some(crc), SharedAction::Upload { known_crc } => known_crc, }, size: Some(p.local.size), excluded: false, }) .collect(), };
let marker_key = full_key(&store_prefix, &format!("meta/.{new_id}.tmp")); // markers are only ever read by us (rocks' meta scan skips non-canonical // names), so they're always V2: strict v1 can't hold crc-less file lines put_new( &store, new_id, &marker_key, meta.serialize(MetaSchema::V2WithSizes)?, ) .await?;
// reused files are already safely in the bucket, so their local links can // go immediately — no need to wait for the uploads to finish if opts.consume_checkpoint { for p in &planned { if matches!(p.action, SharedAction::Reuse { .. }) { consume(&p.local.path).await?; } } }
// -- upload -------------------------------------------------------------- let uploads = planned.iter().filter_map(|p| { let known_crc = match p.action { SharedAction::Upload { known_crc } => known_crc, SharedAction::Reuse { .. } => return None, }; let store = Arc::clone(&store); let key = full_key(&store_prefix, &p.backup_path); let backup_path = p.backup_path.clone(); let local_path = p.local.path.clone(); let size = p.local.size; let consume_checkpoint = opts.consume_checkpoint; Some(async move { let crc = upload_file(store, key, local_path.clone(), size).await?; if let Some(expected) = known_crc && crc != expected { return Err(PushError::ChecksumChanged { path: local_path }); } // durably uploaded: drop the link now rather than at the end, so // the filesystem can reclaim space while the push continues if consume_checkpoint { consume(&local_path).await?; } Ok::<_, PushError>((backup_path, crc, size)) }) });
let total_uploads = uploads.clone().count(); let mut uploaded_files = 0usize; let mut uploaded_bytes = 0u64; let mut crcs: HashMap<String, u32> = HashMap::new(); let mut stream = futures::stream::iter(uploads).buffer_unordered(opts.concurrency); while let Some(result) = stream.next().await { let (backup_path, crc, size) = result?; crcs.insert(backup_path, crc); uploaded_files += 1; uploaded_bytes += size; if uploaded_files.is_multiple_of(25) || uploaded_files == total_uploads { let secs = started.elapsed().as_secs_f64(); let mb = uploaded_bytes as f64 / 2_f64.powf(20.); info!( uploaded = uploaded_files, total = total_uploads, uploaded_mb = format_args!("{mb:.1}"), rate_mb_s = format_args!("{:.1}", if secs > 0.0 { mb / secs } else { 0.0 }), "progress" ); } } drop(stream);
// -- commit --------------------------------------------------------------- for f in &mut meta.files { if let Some(&crc) = crcs.get(&f.path) { f.crc32c = Some(crc); } assert!( f.crc32c.is_some(), "bug: no checksum for {} at commit time", f.path ); } let meta_key = full_key(&store_prefix, &format!("meta/{new_id}")); put_new(&store, new_id, &meta_key, meta.serialize(opts.schema)?).await?;
if let Err(e) = store.delete(&marker_key).await { // the backup is committed; a stale marker is invisible to restore and // list, and cleanup removes it once it goes stale warn!(key = %marker_key, error = %e, "failed to delete in-progress marker"); }
// every file is uploaded and the backup is committed, so all that's left // of the checkpoint is an empty directory. past the commit point a failure // here costs nothing but a stray dir, so it's a warning, not an error. if opts.consume_checkpoint && let Err(e) = tokio::fs::remove_dir(checkpoint_dir).await { warn!( dir = %checkpoint_dir.display(), error = %e, "pushed backup ok, but could not remove the consumed checkpoint directory" ); }
// -- optional retention --------------------------------------------------- let removed = match opts.retain { Some(keep) => { let outcome = retention::retain_backups( raw_store, prefix, keep, &RetentionOptions { concurrency: opts.concurrency, ..Default::default() }, ) .await .map_err(|source| PushError::Retain { backup_id: new_id, keep: keep.get(), source: Box::new(source), })?; outcome.deleted_backups } None => Vec::new(), };
let elapsed = started.elapsed(); info!( backup_id = new_id, total_files, uploaded_files, reused_files, uploaded_mb = format_args!("{:.1}", uploaded_bytes as f64 / 2_f64.powf(20.)), reused_mb = format_args!("{:.1}", reused_bytes as f64 / 2_f64.powf(20.)), removed = removed.len(), elapsed_secs = format_args!("{:.1}", elapsed.as_secs_f64()), "push complete" );
Ok(PushOutcome { backup_id: new_id, total_files, uploaded_files, reused_files, uploaded_bytes, reused_bytes, removed, elapsed, })}
/// unlink a checkpoint file we're done with (see `consume_checkpoint`)async fn consume(path: &Path) -> Result<(), PushError> { tokio::fs::remove_file(path) .await .map_err(|source| PushError::ConsumeFile { path: path.to_path_buf(), source, })}
/// scan_backup_ids returns RetentionError; surface as the push List errorfn retire_to_push(e: retention::RetentionError) -> PushError { match e { retention::RetentionError::List { prefix, source } => PushError::List { prefix, source }, other => unreachable!("scan_backup_ids only lists: {other}"), }}
#[cfg(test)]mod tests { use super::*; use object_store::memory::InMemory;
/// lay down a plausible fake checkpoint (contents aren't real rocks /// files, so session-id extraction fails and everything takes the legacy /// naming path — which is exactly what we want to exercise here) fn fake_checkpoint(dir: &Path) { std::fs::write(dir.join("CURRENT"), b"MANIFEST-000005\n").unwrap(); std::fs::write(dir.join("MANIFEST-000005"), b"manifest-bytes").unwrap(); std::fs::write(dir.join("OPTIONS-000007"), b"options-bytes").unwrap(); std::fs::write(dir.join("000004.log"), b"wal-bytes").unwrap(); std::fs::write(dir.join("000009.sst"), b"sst-bytes-here").unwrap(); }
async fn list_keys(store: &InMemory) -> Vec<String> { let mut keys: Vec<String> = store .list(None) .map(|r| r.unwrap().location.to_string()) .collect() .await; keys.sort(); keys }
#[tokio::test] async fn push_creates_expected_layout() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
let outcome = push(store.clone(), "", cp.path(), PushOptions::default()) .await .unwrap(); assert_eq!(outcome.backup_id, 1); assert_eq!(outcome.total_files, 5); assert_eq!(outcome.uploaded_files, 5); assert_eq!(outcome.reused_files, 0);
let sst_crc = crc32c::crc32c(b"sst-bytes-here"); let expected_sst = format!("shared_checksum/000009_{sst_crc}_14.sst"); let keys = list_keys(&store).await; assert_eq!( keys, vec![ "meta/1".to_string(), "private/1/000004.log".to_string(), "private/1/CURRENT".to_string(), "private/1/MANIFEST-000005".to_string(), "private/1/OPTIONS-000007".to_string(), expected_sst.clone(), ] );
// meta parses, has every crc, sorted paths. no sizes: the default // schema is V1, exactly what rocks writes. let meta = crate::fetch_meta(&*store, "", 1).await.unwrap(); assert_eq!(meta.files.len(), 5); assert!(meta.files.iter().all(|f| f.crc32c.is_some())); assert!(meta.files.iter().all(|f| f.size.is_none())); assert_eq!(meta.files.last().unwrap().path, expected_sst);
// and our own restore accepts it end-to-end let target = tempfile::tempdir().unwrap(); crate::restore(store, "", target.path(), Default::default()) .await .unwrap(); assert_eq!( std::fs::read(target.path().join("000009.sst")).unwrap(), b"sst-bytes-here" ); }
#[tokio::test] async fn second_push_reuses_shared_files() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push(store.clone(), "", cp.path(), PushOptions::default()) .await .unwrap(); let second = push(store.clone(), "", cp.path(), PushOptions::default()) .await .unwrap();
assert_eq!(second.backup_id, 2); assert_eq!(second.reused_files, 1); // the sst assert_eq!(second.uploaded_files, 4); // private files always upload
// both metas restorable, shared object stored once let keys = list_keys(&store).await; assert_eq!( keys.iter() .filter(|k| k.starts_with("shared_checksum/")) .count(), 1 ); assert_eq!(keys.iter().filter(|k| k.starts_with("meta/")).count(), 2); }
async fn pushed_meta_text(store: &InMemory) -> String { let raw = store .get(&StorePath::from("meta/1")) .await .unwrap() .bytes() .await .unwrap(); String::from_utf8(raw.to_vec()).unwrap() }
#[tokio::test] async fn default_schema_writes_rocks_identical_lines() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push( store.clone(), "", cp.path(), PushOptions { timestamp: Some(1234), sequence_number: 42, ..Default::default() }, ) .await .unwrap();
// v1 by default: no schema header, no size fields -- byte-for-byte // the shape rocks' own StoreToFile produces let text = pushed_meta_text(&store).await; assert!(text.starts_with("1234\n42\n5\n"), "got: {text:?}"); assert!(!text.contains("schema_version")); assert!(!text.contains("size")); assert!(text.ends_with('\n')); }
#[tokio::test] async fn v2_schema_opt_in_adds_header_and_sizes() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push( store.clone(), "", cp.path(), PushOptions { schema: MetaSchema::V2WithSizes, timestamp: Some(1234), sequence_number: 42, ..Default::default() }, ) .await .unwrap();
let text = pushed_meta_text(&store).await; assert!( text.starts_with("schema_version 2.1\n1234\n42\n5\n"), "got: {text:?}" ); let meta = crate::fetch_meta(&*store, "", 1).await.unwrap(); assert!(meta.files.iter().all(|f| f.size.is_some())); }
#[tokio::test] async fn respects_prefix() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push( store.clone(), "some/prefix", cp.path(), PushOptions::default(), ) .await .unwrap(); let keys = list_keys(&store).await; assert!( keys.iter().all(|k| k.starts_with("some/prefix/")), "{keys:?}" ); crate::fetch_meta(&*store, "some/prefix", 1).await.unwrap(); }
#[tokio::test] async fn rejects_live_db_dir() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); // live dbs have IDENTITY (and LOCK, LOG); checkpoints never do std::fs::write(cp.path().join("IDENTITY"), b"uuid").unwrap();
let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions::default(), ) .await; assert!( matches!(result, Err(PushError::UnexpectedCheckpointFile { name }) if name == "IDENTITY") ); }
#[tokio::test] async fn rejects_subdirectory() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); std::fs::create_dir(cp.path().join("nested")).unwrap();
let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions::default(), ) .await; assert!(matches!( result, Err(PushError::CheckpointSubdirectory { .. }) )); }
#[tokio::test] async fn rejects_missing_current() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); std::fs::remove_file(cp.path().join("CURRENT")).unwrap();
let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions::default(), ) .await; assert!(matches!(result, Err(PushError::MissingCurrent { .. }))); }
#[tokio::test] async fn rejects_current_manifest_mismatch() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); std::fs::write(cp.path().join("CURRENT"), b"MANIFEST-000099\n").unwrap();
let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions::default(), ) .await; assert!(matches!( result, Err(PushError::CurrentManifestMismatch { .. }) )); }
#[tokio::test] async fn rejects_malformed_current() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); std::fs::write(cp.path().join("CURRENT"), b"whatever\n").unwrap();
let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions::default(), ) .await; assert!(matches!(result, Err(PushError::InvalidCurrent { .. }))); }
#[tokio::test] async fn claimed_id_conflicts() { // a marker (or meta) appearing between our LIST and our conditional // PUT — i.e. a concurrent pusher racing the same id — must surface as // BackupIdClaimed, not silently overwrite let store = Arc::new(InMemory::new()); store .put( &StorePath::from("meta/.2.tmp"), PutPayload::from_static(b"0\n0\n0\n"), ) .await .unwrap();
let err = put_new( &(store.clone() as Arc<dyn ObjectStore>), 2, &StorePath::from("meta/.2.tmp"), "x".to_string(), ) .await .unwrap_err(); assert!(matches!(err, PushError::BackupIdClaimed { id: 2, .. })); }
#[tokio::test] async fn id_allocation_skips_all_debris() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
// committed 2, marker 5, private debris 7: next id must be 8 store .put( &StorePath::from("meta/2"), PutPayload::from_static(b"0\n0\n0\n"), ) .await .unwrap(); store .put( &StorePath::from("meta/.5.tmp"), PutPayload::from_static(b"0\n0\n0\n"), ) .await .unwrap(); store .put( &StorePath::from("private/7/CURRENT"), PutPayload::from_static(b"junk"), ) .await .unwrap();
let outcome = push(store.clone(), "", cp.path(), PushOptions::default()) .await .unwrap(); assert_eq!(outcome.backup_id, 8); }
#[tokio::test] async fn marker_gone_after_clean_push() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push(store.clone(), "", cp.path(), PushOptions::default()) .await .unwrap(); let keys = list_keys(&store).await; assert!(!keys.iter().any(|k| k.contains(".tmp")), "{keys:?}"); }
#[tokio::test] async fn consume_checkpoint_removes_everything() { let parent = tempfile::tempdir().unwrap(); let cp = parent.path().join("checkpoint"); std::fs::create_dir(&cp).unwrap(); fake_checkpoint(&cp); let store = Arc::new(InMemory::new());
push( store.clone(), "", &cp, PushOptions { consume_checkpoint: true, ..Default::default() }, ) .await .unwrap();
assert!(!cp.exists(), "checkpoint dir should be gone"); // and the backup is intact and restorable let target = tempfile::tempdir().unwrap(); crate::restore(store, "", target.path(), Default::default()) .await .unwrap(); assert_eq!( std::fs::read(target.path().join("000009.sst")).unwrap(), b"sst-bytes-here" ); }
#[tokio::test] async fn consume_checkpoint_drops_reused_files_too() { // second push of the same checkpoint reuses the sst (no upload), but // consuming must still release its local link let store = Arc::new(InMemory::new()); let first = tempfile::tempdir().unwrap(); fake_checkpoint(first.path()); push(store.clone(), "", first.path(), PushOptions::default()) .await .unwrap();
let parent = tempfile::tempdir().unwrap(); let cp = parent.path().join("checkpoint"); std::fs::create_dir(&cp).unwrap(); fake_checkpoint(&cp);
let second = push( store.clone(), "", &cp, PushOptions { consume_checkpoint: true, ..Default::default() }, ) .await .unwrap();
assert_eq!(second.reused_files, 1); assert!(!cp.exists(), "checkpoint dir should be gone"); }
#[tokio::test] async fn checkpoint_untouched_by_default() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new());
push(store, "", cp.path(), PushOptions::default()) .await .unwrap();
for name in [ "CURRENT", "MANIFEST-000005", "OPTIONS-000007", "000004.log", "000009.sst", ] { assert!( cp.path().join(name).exists(), "{name} should still be there" ); } }
#[tokio::test] async fn oversized_metadata_rejected() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let result = push( Arc::new(InMemory::new()), "", cp.path(), PushOptions { app_metadata: Some(vec![0u8; MAX_APP_METADATA + 1]), ..Default::default() }, ) .await; assert!(matches!(result, Err(PushError::MetadataTooLarge { .. }))); }
#[tokio::test] async fn metadata_lands_in_meta_hex_encoded() { let cp = tempfile::tempdir().unwrap(); fake_checkpoint(cp.path()); let store = Arc::new(InMemory::new()); push( store.clone(), "", cp.path(), PushOptions { app_metadata: Some(b"Hello".to_vec()), ..Default::default() }, ) .await .unwrap(); let meta = crate::fetch_meta(&*store, "", 1).await.unwrap(); assert_eq!(meta.metadata.as_deref(), Some("48656c6c6f")); }
#[test] fn classification() { assert_eq!(classify("CURRENT"), Some(FileKind::Private)); assert_eq!(classify("MANIFEST-000012"), Some(FileKind::Private)); assert_eq!(classify("OPTIONS-000010"), Some(FileKind::Private)); assert_eq!(classify("000011.log"), Some(FileKind::Private)); assert_eq!(classify("000009.sst"), Some(FileKind::SharedSst)); assert_eq!(classify("000012.blob"), Some(FileKind::SharedBlob)); // live-db / junk files assert_eq!(classify("IDENTITY"), None); assert_eq!(classify("LOCK"), None); assert_eq!(classify("LOG"), None); assert_eq!(classify("LOG.old.12345"), None); assert_eq!(classify("OPTIONS-000010.dbtmp"), None); assert_eq!(classify("mydata.sst"), None); assert_eq!(classify(".sst"), None); assert_eq!(classify("MANIFEST-"), None); }
#[test] fn name_mangling() { assert_eq!( session_name("000009.sst", "V90YQ7DEHIVN86FM2QZ0", 59021), "000009_sV90YQ7DEHIVN86FM2QZ0_59021.sst" ); assert_eq!( legacy_name("000012.blob", 2894567812, 1044), "000012_2894567812_1044.blob" ); // round-trips through restore's unmangler assert_eq!( crate::restore::db_filename("shared_checksum/000009_sV90YQ7DEHIVN86FM2QZ0_59021.sst") .unwrap(), Path::new("000009.sst") ); }}