use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use tempfile::TempPath;
use crate::hash::Hash;
use crate::store::{
MAX_RAW_OBJECT_SIZE, ObjectSink, ObjectStore, StoreError, StoreResult, sync_parent_dir,
temp_file_in,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SyncPolicy {
PerObject,
#[default]
Batch,
}
pub(crate) trait Syncer: Send + Sync + fmt::Debug {
fn barrier(&self, file: &File, path: &Path) -> io::Result<()>;
fn full(&self, file: &File, path: &Path) -> io::Result<()>;
fn rename(&self, tmp: TempPath, final_path: &Path) -> io::Result<()>;
fn dir_sync(&self, dir: &Path) -> io::Result<()>;
fn dir_barrier(&self, dir: &Path) -> io::Result<()>;
fn device_flush(&self, objects_root: &Path) -> io::Result<()>;
}
#[derive(Debug)]
pub(crate) struct RealSyncer;
impl RealSyncer {
#[cfg(any(target_os = "macos", target_os = "ios"))]
fn file_barrier(file: &File) -> io::Result<()> {
use std::os::unix::io::AsRawFd;
#[allow(unsafe_code)]
let rc = unsafe { libc::fcntl(file.as_raw_fd(), libc::F_BARRIERFSYNC) };
if rc == -1 {
return file.sync_data();
}
Ok(())
}
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
fn file_barrier(file: &File) -> io::Result<()> {
file.sync_data()
}
}
impl Syncer for RealSyncer {
fn barrier(&self, file: &File, _path: &Path) -> io::Result<()> {
Self::file_barrier(file)
}
fn full(&self, file: &File, _path: &Path) -> io::Result<()> {
file.sync_all()
}
fn rename(&self, tmp: TempPath, final_path: &Path) -> io::Result<()> {
tmp.persist(final_path).map_err(|e| e.error)?;
Ok(())
}
fn dir_sync(&self, dir: &Path) -> io::Result<()> {
sync_parent_dir(dir)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
fn dir_barrier(&self, dir: &Path) -> io::Result<()> {
match File::open(dir) {
Ok(d) => d.sync_data().or_else(|_| d.sync_all()),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
}
}
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
fn dir_barrier(&self, dir: &Path) -> io::Result<()> {
sync_parent_dir(dir)
}
#[cfg(unix)]
fn device_flush(&self, objects_root: &Path) -> io::Result<()> {
match File::open(objects_root) {
Ok(d) => d.sync_all(),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
}
}
#[cfg(not(unix))]
#[allow(clippy::unnecessary_wraps)]
fn device_flush(&self, _objects_root: &Path) -> io::Result<()> {
Ok(())
}
}
#[derive(Debug, Default)]
struct BatchState {
staged: HashMap<Hash, TempPath>,
touched_shards: HashSet<PathBuf>,
created_shards: HashSet<PathBuf>,
}
#[derive(Debug)]
pub struct WriteBatch<'s> {
store: &'s ObjectStore,
policy: SyncPolicy,
inner: Mutex<BatchState>,
}
impl<'s> WriteBatch<'s> {
pub(crate) fn new(store: &'s ObjectStore, policy: SyncPolicy) -> Self {
Self {
store,
policy,
inner: Mutex::new(BatchState::default()),
}
}
pub fn write(&self, bytes: &[u8]) -> StoreResult<Hash> {
self.write_parts(&[bytes])
}
pub fn write_parts(&self, parts: &[&[u8]]) -> StoreResult<Hash> {
let mut total: usize = 0;
for p in parts {
total = total
.checked_add(p.len())
.ok_or(StoreError::ObjectTooLarge)?;
}
if total > MAX_RAW_OBJECT_SIZE {
return Err(StoreError::ObjectTooLarge);
}
self.write_prehashed(crate::object::object_id_from_parts(parts), parts)
}
pub(crate) fn write_prehashed(&self, h: Hash, parts: &[&[u8]]) -> StoreResult<Hash> {
let final_path = self.store.path_for(&h);
let shard_dir = final_path
.parent()
.expect("object path always has a 2-hex parent")
.to_path_buf();
let need_mkdir = {
let st = self.inner.lock().expect("batch state mutex poisoned");
if st.staged.contains_key(&h) {
return Ok(h);
}
!st.created_shards.contains(&shard_dir)
};
if final_path.exists() {
self.inner
.lock()
.expect("batch state mutex poisoned")
.touched_shards
.insert(shard_dir);
return Ok(h);
}
if need_mkdir {
fs::create_dir_all(&shard_dir)?;
}
let file_name = final_path
.file_name()
.expect("object path has file name")
.to_string_lossy();
let mut tmp = temp_file_in(&shard_dir, &file_name)?;
for p in parts {
tmp.as_file_mut().write_all(p)?;
}
let syncer = self.store.syncer();
match self.policy {
SyncPolicy::PerObject => {
syncer.full(tmp.as_file(), tmp.path())?;
syncer.rename(tmp.into_temp_path(), &final_path)?;
syncer.dir_sync(&shard_dir)?;
let mut st = self.inner.lock().expect("batch state mutex poisoned");
st.created_shards.insert(shard_dir);
}
SyncPolicy::Batch => {
let mut st = self.inner.lock().expect("batch state mutex poisoned");
st.staged.entry(h).or_insert_with(|| tmp.into_temp_path());
st.touched_shards.insert(shard_dir.clone());
st.created_shards.insert(shard_dir);
}
}
Ok(h)
}
#[must_use]
pub fn contains(&self, h: &Hash) -> bool {
self.inner
.lock()
.expect("batch state mutex poisoned")
.staged
.contains_key(h)
|| self.store.contains(h)
}
pub fn commit(self) -> StoreResult<()> {
let st = self.inner.into_inner().expect("batch state mutex poisoned");
let syncer = self.store.syncer();
match self.policy {
SyncPolicy::PerObject => Ok(()),
SyncPolicy::Batch => {
let staged: Vec<(Hash, TempPath)> = st.staged.into_iter().collect();
parallel_io(staged.len(), |i| {
let f = OpenOptions::new().write(true).open(&staged[i].1)?;
syncer.barrier(&f, &staged[i].1)
})?;
if let Some((_, tmp)) = staged.first() {
let f = OpenOptions::new().write(true).open(tmp)?;
syncer.full(&f, tmp)?;
}
for (h, tmp) in staged {
syncer.rename(tmp, &self.store.path_for(&h))?;
}
let mut shards: Vec<PathBuf> = st.touched_shards.into_iter().collect();
shards.sort();
if !shards.is_empty() {
parallel_io(shards.len(), |i| syncer.dir_barrier(&shards[i]))?;
syncer.device_flush(self.store.objects_root())?;
}
Ok(())
}
}
}
}
const MAX_SYNC_WORKERS: usize = 64;
fn parallel_io(count: usize, op: impl Fn(usize) -> io::Result<()> + Sync) -> io::Result<()> {
if count == 0 {
return Ok(());
}
let workers = MAX_SYNC_WORKERS.min(count);
if workers == 1 {
return (0..count).try_for_each(op);
}
let next = std::sync::atomic::AtomicUsize::new(0);
let mut results: Vec<io::Result<()>> = Vec::new();
std::thread::scope(|scope| {
let handles: Vec<_> = (0..workers)
.map(|_| {
let next = &next;
let op = &op;
scope.spawn(move || -> io::Result<()> {
loop {
let i = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if i >= count {
return Ok(());
}
op(i)?;
}
})
})
.collect();
for h in handles {
results.push(h.join().expect("parallel io worker panicked"));
}
});
results.into_iter().collect()
}
impl ObjectSink for WriteBatch<'_> {
fn put(&self, bytes: &[u8]) -> StoreResult<Hash> {
self.write(bytes)
}
fn put_parts(&self, parts: &[&[u8]]) -> StoreResult<Hash> {
self.write_parts(parts)
}
fn has(&self, h: &Hash) -> bool {
self.contains(h)
}
}
#[cfg(test)]
pub(crate) mod testing {
use super::*;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Ev {
Barrier(PathBuf),
Full(PathBuf),
Rename { tmp: PathBuf, dst: PathBuf },
DirSync(PathBuf),
DirBarrier(PathBuf),
}
#[derive(Debug, Default)]
pub(crate) struct RecordingSyncer {
events: Mutex<Vec<Ev>>,
}
impl RecordingSyncer {
pub(crate) fn events(&self) -> Vec<Ev> {
self.events.lock().unwrap().clone()
}
}
impl Syncer for RecordingSyncer {
fn barrier(&self, _file: &File, path: &Path) -> io::Result<()> {
self.events
.lock()
.unwrap()
.push(Ev::Barrier(path.to_path_buf()));
Ok(())
}
fn full(&self, _file: &File, path: &Path) -> io::Result<()> {
self.events
.lock()
.unwrap()
.push(Ev::Full(path.to_path_buf()));
Ok(())
}
fn rename(&self, tmp: TempPath, final_path: &Path) -> io::Result<()> {
self.events.lock().unwrap().push(Ev::Rename {
tmp: tmp.to_path_buf(),
dst: final_path.to_path_buf(),
});
tmp.persist(final_path).map_err(|e| e.error)?;
Ok(())
}
fn dir_sync(&self, dir: &Path) -> io::Result<()> {
self.events
.lock()
.unwrap()
.push(Ev::DirSync(dir.to_path_buf()));
Ok(())
}
fn dir_barrier(&self, dir: &Path) -> io::Result<()> {
self.events
.lock()
.unwrap()
.push(Ev::DirBarrier(dir.to_path_buf()));
Ok(())
}
fn device_flush(&self, objects_root: &Path) -> io::Result<()> {
self.events
.lock()
.unwrap()
.push(Ev::Full(objects_root.to_path_buf()));
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::testing::{Ev, RecordingSyncer};
use super::*;
use crate::hash;
use proptest::prelude::*;
use std::sync::Arc;
use tempfile::TempDir;
fn fresh_store() -> (TempDir, ObjectStore) {
let dir = TempDir::new().expect("tempdir");
let store =
ObjectStore::init(&crate::layout::RepoLayout::single(dir.path())).expect("init");
(dir, store)
}
fn recording_store() -> (TempDir, ObjectStore, Arc<RecordingSyncer>) {
let (dir, mut store) = fresh_store();
let rec = Arc::new(RecordingSyncer::default());
store.set_syncer(rec.clone());
(dir, store, rec)
}
fn object_file_count(store: &ObjectStore) -> usize {
store.iter_object_hashes().unwrap().len()
}
fn any_file_count(store: &ObjectStore) -> usize {
fn walk(dir: &Path, n: &mut usize) {
if let Ok(rd) = fs::read_dir(dir) {
for e in rd.flatten() {
let p = e.path();
if p.is_dir() {
walk(&p, n);
} else {
*n += 1;
}
}
}
}
let mut n = 0;
walk(store.objects_root(), &mut n);
n
}
#[test]
fn staged_object_is_invisible_until_commit() {
let (dir, store) = fresh_store();
let batch = store.batch();
let h = batch.write(b"staged bytes").unwrap();
let other = ObjectStore::open(&crate::layout::RepoLayout::single(dir.path())).unwrap();
assert!(
!other.contains(&h),
"staged object must be invisible before commit"
);
assert!(other.read(&h).is_err());
batch.commit().unwrap();
assert!(other.contains(&h), "committed object must be visible");
assert_eq!(other.read(&h).unwrap(), b"staged bytes");
}
#[test]
fn batch_contains_sees_staged_and_disk() {
let (_dir, store) = fresh_store();
let on_disk = store.write(b"already stored").unwrap();
let batch = store.batch();
let staged = batch.write(b"only staged").unwrap();
assert!(batch.contains(&on_disk), "must see on-disk objects");
assert!(batch.contains(&staged), "must see its own staged objects");
let phony = hash::hash(b"never written");
assert!(!batch.contains(&phony));
}
#[test]
fn dropped_batch_leaves_no_tmp_files_and_no_objects() {
let (_dir, store) = fresh_store();
{
let batch = store.batch();
batch.write(b"abort me 1").unwrap();
batch.write(b"abort me 2").unwrap();
batch.write(b"abort me 3").unwrap();
}
assert_eq!(object_file_count(&store), 0, "no objects may be visible");
assert_eq!(
any_file_count(&store),
0,
"no temp files may leak from an aborted batch"
);
}
fn fifty_distinct_objects() -> Vec<Vec<u8>> {
(0u32..50)
.map(|i| format!("object #{i}").into_bytes())
.collect()
}
#[test]
fn batch_commit_full_flush_count_is_constant() {
for count in [3usize, 50] {
let (_dir, store, rec) = recording_store();
let batch = store.batch();
for bytes in fifty_distinct_objects().into_iter().take(count) {
batch.write(&bytes).unwrap();
}
batch.write(b"object #0").unwrap();
batch.commit().unwrap();
let fulls = rec
.events()
.iter()
.filter(|e| matches!(e, Ev::Full(_)))
.count();
assert_eq!(
fulls, 2,
"a {count}-object batch must cost exactly two full flushes"
);
}
}
#[test]
fn every_rename_is_preceded_by_its_barrier_and_the_full_flush() {
let (_dir, store, rec) = recording_store();
let batch = store.batch();
for bytes in fifty_distinct_objects() {
batch.write(&bytes).unwrap();
}
batch.commit().unwrap();
let evs = rec.events();
let full_pos = evs
.iter()
.position(|e| matches!(e, Ev::Full(_)))
.expect("one full flush");
let last_barrier = evs
.iter()
.rposition(|e| matches!(e, Ev::Barrier(_)))
.expect("barriers recorded");
assert!(
last_barrier < full_pos,
"all barriers (last at {last_barrier}) must precede the full flush at {full_pos}"
);
for (i, ev) in evs.iter().enumerate() {
if let Ev::Rename { tmp, .. } = ev {
assert!(
i > full_pos,
"rename at {i} must come after the full flush at {full_pos}"
);
let barrier_pos = evs
.iter()
.position(|e| matches!(e, Ev::Barrier(p) if p == tmp))
.unwrap_or_else(|| panic!("no barrier recorded for {}", tmp.display()));
assert!(
barrier_pos < i,
"barrier for {} must precede its rename",
tmp.display()
);
}
}
}
#[test]
fn dir_syncs_come_after_all_renames_and_are_deduped() {
let (_dir, store, rec) = recording_store();
let batch = store.batch();
let mut shards = HashSet::new();
for bytes in fifty_distinct_objects() {
let h = batch.write(&bytes).unwrap();
shards.insert(store.path_for(&h).parent().unwrap().to_path_buf());
}
batch.commit().unwrap();
let evs = rec.events();
let last_rename = evs
.iter()
.rposition(|e| matches!(e, Ev::Rename { .. }))
.expect("renames recorded");
let dir_barriers: Vec<(usize, &PathBuf)> = evs
.iter()
.enumerate()
.filter_map(|(i, e)| match e {
Ev::DirBarrier(p) => Some((i, p)),
_ => None,
})
.collect();
let synced: HashSet<PathBuf> = dir_barriers.iter().map(|(_, p)| (*p).clone()).collect();
assert_eq!(
dir_barriers.len(),
synced.len(),
"each shard dir must be flushed exactly once"
);
assert_eq!(synced, shards, "exactly the touched shards are flushed");
for (i, p) in &dir_barriers {
assert!(
*i > last_rename,
"dir barrier of {} at {i} must come after the last rename at {last_rename}",
p.display()
);
}
let last_full = evs
.iter()
.rposition(|e| matches!(e, Ev::Full(_)))
.expect("device flush recorded");
let last_dir_barrier = dir_barriers.last().expect("dir barriers recorded").0;
assert!(
last_full > last_dir_barrier,
"device flush at {last_full} must follow the last dir barrier at {last_dir_barrier}"
);
}
#[test]
fn dedup_hit_still_dir_syncs_at_commit() {
let (_dir, store, rec) = recording_store();
let h = store.write(b"already present").unwrap();
let shard = store.path_for(&h).parent().unwrap().to_path_buf();
let batch = store.batch();
let h2 = batch.write(b"already present").unwrap();
assert_eq!(h, h2);
let before = rec.events().len();
batch.commit().unwrap();
let evs = rec.events()[before..].to_vec();
assert!(
!evs.iter().any(|e| matches!(e, Ev::Rename { .. })),
"dedup hit must not stage or rename anything"
);
assert!(
evs.contains(&Ev::DirBarrier(shard)),
"commit must still flush the dedup-hit shard dir: its dirent \
may not be durable yet and we are about to reference it"
);
assert!(
matches!(evs.last(), Some(Ev::Full(_))),
"the dir barrier needs a trailing device flush to be durable"
);
}
#[test]
fn sync_policy_per_object_matches_legacy_event_pattern() {
let (_dir, store, rec) = recording_store();
let legacy_h = store.write(b"legacy path").unwrap();
let legacy: Vec<Ev> = rec.events();
let batch = store.batch_with_policy(SyncPolicy::PerObject);
let h = batch.write(b"per object path").unwrap();
assert!(
store.contains(&h),
"PerObject writes must be visible immediately, pre-commit"
);
let per_object: Vec<Ev> = rec.events()[legacy.len()..].to_vec();
let kinds = |evs: &[Ev]| -> Vec<u8> {
evs.iter()
.map(|e| match e {
Ev::Barrier(_) => 0u8,
Ev::Full(_) => 1,
Ev::Rename { .. } => 2,
Ev::DirSync(_) => 3,
Ev::DirBarrier(_) => 4,
})
.collect()
};
assert_eq!(
kinds(&per_object),
kinds(&legacy),
"PerObject batch must reproduce the legacy per-write sync pattern"
);
let before = rec.events().len();
batch.commit().unwrap();
assert_eq!(rec.events().len(), before);
let _ = legacy_h;
}
#[test]
fn idempotent_duplicate_writes_in_one_batch_stage_once() {
let (_dir, store, rec) = recording_store();
let batch = store.batch();
let h1 = batch.write(b"twice staged").unwrap();
let h2 = batch.write(b"twice staged").unwrap();
assert_eq!(h1, h2);
batch.commit().unwrap();
let evs = rec.events();
let barriers = evs.iter().filter(|e| matches!(e, Ev::Barrier(_))).count();
let renames = evs
.iter()
.filter(|e| matches!(e, Ev::Rename { .. }))
.count();
assert_eq!(barriers, 1, "duplicate must not re-stage");
assert_eq!(renames, 1, "duplicate must not re-rename");
}
#[test]
fn batch_write_rejects_oversize() {
let (_dir, store) = fresh_store();
let batch = store.batch();
let oversize = vec![0u8; MAX_RAW_OBJECT_SIZE + 1];
let err = batch.write(&oversize).unwrap_err();
assert!(matches!(err, StoreError::ObjectTooLarge), "got {err:?}");
drop(oversize);
let h = batch.write(&[0u8; 16]).unwrap();
batch.commit().unwrap();
assert!(store.contains(&h));
}
proptest! {
#[test]
fn batch_write_hash_equals_store_write_hash(bytes in proptest::collection::vec(any::<u8>(), 0..4096)) {
let (_dir, store) = fresh_store();
let batch = store.batch();
let h_batch = batch.write(&bytes).unwrap();
batch.commit().unwrap();
let on_disk_via_batch = store.read(&h_batch).unwrap();
let (_dir2, store2) = fresh_store();
let h_store = store2.write(&bytes).unwrap();
prop_assert_eq!(h_batch, h_store, "batch and store writes must agree on the hash");
prop_assert_eq!(on_disk_via_batch, bytes, "on-disk bytes must round-trip");
}
#[test]
fn write_parts_equals_concatenated_write(
parts in proptest::collection::vec(proptest::collection::vec(any::<u8>(), 0..512), 0..8)
) {
let (_dir, store) = fresh_store();
let concatenated: Vec<u8> = parts.iter().flatten().copied().collect();
let batch = store.batch();
let slices: Vec<&[u8]> = parts.iter().map(Vec::as_slice).collect();
let h_parts = batch.write_parts(&slices).unwrap();
let h_whole = batch.write(&concatenated).unwrap();
batch.commit().unwrap();
prop_assert_eq!(h_parts, h_whole, "parts and whole must hash identically");
prop_assert_eq!(store.read(&h_parts).unwrap(), concatenated);
}
}
}