use crate::error::Result;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Arc;
pub const PART_SUFFIX: &str = ".rst-part";
pub struct WriteHandle {
file: File,
part_path: PathBuf,
final_path: PathBuf,
size: u64,
created_new: bool,
}
impl WriteHandle {
pub fn open_cloned(final_path: &Path, size: u64, preallocate: bool) -> Result<Self> {
let part_path = part_path_for(final_path);
if !part_path.exists() && final_path.exists() && clone_file(final_path, &part_path) {
if let Ok(file) = OpenOptions::new().read(true).write(true).open(&part_path) {
if file.metadata()?.len() != size {
file.set_len(size)?;
}
return Ok(Self {
file,
part_path,
final_path: final_path.to_path_buf(),
size,
created_new: false,
});
}
}
Self::open(final_path, size, preallocate)
}
pub fn open(final_path: &Path, size: u64, preallocate: bool) -> Result<Self> {
if let Some(parent) = final_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| io::Error::new(e.kind(), format!("{}: {e}", parent.display())))?;
}
let part_path = part_path_for(final_path);
let (file, created_new) = match OpenOptions::new()
.create_new(true)
.read(true)
.write(true)
.open(&part_path)
{
Ok(f) => (f, true),
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
let f = OpenOptions::new()
.read(true)
.write(true)
.truncate(false)
.open(&part_path)
.map_err(|e| {
io::Error::new(e.kind(), format!("{}: {e}", part_path.display()))
})?;
(f, false)
}
Err(e) => {
return Err(
io::Error::new(e.kind(), format!("{}: {e}", part_path.display())).into(),
)
}
};
if preallocate && size > 0 && created_new {
reserve(&file, size)?;
}
if file.metadata()?.len() != size {
file.set_len(size)?;
}
Ok(Self {
file,
part_path,
final_path: final_path.to_path_buf(),
size,
created_new,
})
}
pub fn size(&self) -> u64 {
self.size
}
pub fn part_path(&self) -> &Path {
&self.part_path
}
pub fn created_new(&self) -> bool {
self.created_new
}
pub fn write_zeros_at(&self, offset: u64, len: usize) -> io::Result<()> {
const BLOCK: usize = 256 * 1024;
let zeros = [0u8; BLOCK];
let mut done = 0usize;
while done < len {
let n = BLOCK.min(len - done);
self.write_at(offset + done as u64, &zeros[..n])?;
done += n;
}
Ok(())
}
pub fn write_at(&self, offset: u64, buf: &[u8]) -> io::Result<()> {
let end = offset
.checked_add(buf.len() as u64)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "write offset overflow"))?;
if end > self.size {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"write of {} bytes at {offset} exceeds declared size {}",
buf.len(),
self.size
),
));
}
let mut written = 0usize;
while written < buf.len() {
let n = pwrite(&self.file, &buf[written..], offset + written as u64)?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::WriteZero,
"pwrite made no progress",
));
}
written += n;
}
Ok(())
}
pub fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
pread(&self.file, buf, offset)
}
pub fn matches_at(&self, offset: u64, want: &[u8]) -> io::Result<bool> {
let mut scratch = vec![0u8; want.len()];
let n = pread(&self.file, &mut scratch, offset)?;
Ok(n == want.len() && scratch == want)
}
pub fn sync(&self) -> io::Result<()> {
self.file.sync_all()
}
pub fn commit(&self, mode: u32, mtime: i64, preserve_metadata: bool) -> Result<()> {
self.file.sync_all()?;
if preserve_metadata {
#[cfg(unix)]
if mode != 0 {
use std::os::unix::fs::PermissionsExt;
let _ = self
.file
.set_permissions(std::fs::Permissions::from_mode(mode & 0o7777));
}
}
std::fs::rename(&self.part_path, &self.final_path).map_err(|e| {
io::Error::new(
e.kind(),
format!(
"rename {} -> {}: {e}",
self.part_path.display(),
self.final_path.display()
),
)
})?;
if preserve_metadata && mtime > 0 {
set_mtime(&self.final_path, mtime);
}
Ok(())
}
pub fn discard(&self) {
let _ = std::fs::remove_file(&self.part_path);
}
}
#[allow(clippy::needless_return)]
fn clone_file(src: &Path, dst: &Path) -> bool {
#[cfg(target_os = "macos")]
{
use std::ffi::CString;
let (Ok(s), Ok(d)) = (
CString::new(src.as_os_str().as_encoded_bytes()),
CString::new(dst.as_os_str().as_encoded_bytes()),
) else {
return false;
};
return unsafe { libc::clonefile(s.as_ptr(), d.as_ptr(), 0) } == 0;
}
#[cfg(target_os = "linux")]
{
use std::os::unix::io::AsRawFd;
const FICLONE: libc::c_ulong = 0x4004_9409;
let (Ok(from), Ok(to)) = (
std::fs::File::open(src),
OpenOptions::new().create_new(true).write(true).open(dst),
) else {
return false;
};
let ok = unsafe { libc::ioctl(to.as_raw_fd(), FICLONE, from.as_raw_fd()) } == 0;
if !ok {
drop(to);
let _ = std::fs::remove_file(dst);
}
ok
}
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
{
let _ = (src, dst);
false
}
}
pub fn part_path_for(final_path: &Path) -> PathBuf {
let mut s = final_path.as_os_str().to_os_string();
s.push(PART_SUFFIX);
PathBuf::from(s)
}
#[allow(clippy::needless_return)]
fn reserve(file: &File, size: u64) -> Result<()> {
#[cfg(target_os = "linux")]
{
use std::os::unix::io::AsRawFd;
let rc = unsafe {
libc::fallocate(
file.as_raw_fd(),
libc::FALLOC_FL_KEEP_SIZE,
0,
size as libc::off_t,
)
};
if rc != 0 {
let e = io::Error::last_os_error();
match e.raw_os_error() {
Some(libc::EOPNOTSUPP) | Some(libc::ENOSYS) | Some(libc::EINVAL) => {}
_ => return Err(crate::error::Error::Io(e)),
}
}
return Ok(());
}
#[cfg(target_os = "macos")]
{
use std::os::unix::fs::MetadataExt;
use std::os::unix::io::AsRawFd;
let allocated = file.metadata()?.blocks().saturating_mul(512);
if allocated >= size {
return Ok(());
}
let want = (size - allocated) as libc::off_t;
let mut store = libc::fstore_t {
fst_flags: libc::F_ALLOCATECONTIG,
fst_posmode: libc::F_PEOFPOSMODE,
fst_offset: 0,
fst_length: want,
fst_bytesalloc: 0,
};
unsafe {
if libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &mut store) == -1 {
store.fst_flags = libc::F_ALLOCATEALL;
let _ = libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &mut store);
}
}
return Ok(());
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
{
let _ = (file, size);
Ok(())
}
}
fn set_mtime(path: &Path, mtime: i64) {
#[cfg(unix)]
{
use std::ffi::CString;
if let Ok(c) = CString::new(path.as_os_str().as_encoded_bytes()) {
let times = [
libc::timeval {
tv_sec: mtime as libc::time_t,
tv_usec: 0,
},
libc::timeval {
tv_sec: mtime as libc::time_t,
tv_usec: 0,
},
];
unsafe {
libc::utimes(c.as_ptr(), times.as_ptr());
}
}
}
#[cfg(not(unix))]
let _ = (path, mtime);
}
#[cfg(unix)]
fn pwrite(file: &File, buf: &[u8], offset: u64) -> io::Result<usize> {
use std::os::unix::fs::FileExt;
file.write_at(buf, offset)
}
#[cfg(windows)]
fn pwrite(file: &File, buf: &[u8], offset: u64) -> io::Result<usize> {
use std::os::windows::fs::FileExt;
file.seek_write(buf, offset)
}
#[cfg(unix)]
fn pread(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
use std::os::unix::fs::FileExt;
file.read_at(buf, offset)
}
#[cfg(windows)]
fn pread(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
use std::os::windows::fs::FileExt;
file.seek_read(buf, offset)
}
#[derive(Default)]
struct Slot(Mutex<Option<Arc<WriteHandle>>>);
#[derive(Default)]
pub struct FileWriters {
map: Mutex<HashMap<u32, Arc<Slot>>>,
}
impl FileWriters {
pub fn new() -> Self {
Self::default()
}
pub fn get_or_open<F>(&self, file_id: u32, make: F) -> Result<Arc<WriteHandle>>
where
F: FnOnce() -> Result<WriteHandle>,
{
let slot = {
let mut map = self.map.lock();
map.entry(file_id).or_default().clone()
};
let mut guard = slot.0.lock();
if let Some(h) = guard.as_ref() {
return Ok(h.clone());
}
let handle = Arc::new(make()?);
*guard = Some(handle.clone());
Ok(handle)
}
pub fn take(&self, file_id: u32) -> Option<Arc<WriteHandle>> {
let slot = self.map.lock().remove(&file_id)?;
let mut guard = slot.0.lock();
guard.take()
}
pub fn len(&self) -> usize {
self.map
.lock()
.values()
.filter(|s| s.0.lock().is_some())
.count()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn discard_all(&self) {
let slots: Vec<_> = self.map.lock().drain().map(|(_, s)| s).collect();
for s in slots {
if let Some(h) = s.0.lock().take() {
h.discard();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn writes_out_of_order_then_commits() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("nested/dir/out.bin");
let h = WriteHandle::open(&dest, 3000, true).unwrap();
h.write_at(2000, &[3u8; 1000]).unwrap();
h.write_at(0, &[1u8; 1000]).unwrap();
h.write_at(1000, &[2u8; 1000]).unwrap();
assert!(h.part_path().exists());
assert!(!dest.exists(), "must not appear until committed");
h.commit(0o640, 1_600_000_000, true).unwrap();
assert!(dest.exists());
assert!(!part_path_for(&dest).exists());
let got = std::fs::read(&dest).unwrap();
assert_eq!(got.len(), 3000);
assert_eq!(&got[0..1000], &[1u8; 1000]);
assert_eq!(&got[1000..2000], &[2u8; 1000]);
assert_eq!(&got[2000..3000], &[3u8; 1000]);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let m = std::fs::metadata(&dest).unwrap();
assert_eq!(m.permissions().mode() & 0o777, 0o640);
}
}
#[test]
fn refuses_writes_past_declared_size() {
let tmp = tempfile::tempdir().unwrap();
let h = WriteHandle::open(&tmp.path().join("a.bin"), 100, false).unwrap();
assert!(h.write_at(90, &[0u8; 20]).is_err());
assert!(h.write_at(u64::MAX, &[0u8; 1]).is_err());
h.write_at(90, &[0u8; 10]).unwrap();
}
#[test]
fn concurrent_writers_share_one_handle() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("big.bin");
let h = Arc::new(WriteHandle::open(&dest, 64 * 4096, true).unwrap());
let threads: Vec<_> = (0..8u64)
.map(|t| {
let h = h.clone();
std::thread::spawn(move || {
for i in 0..8u64 {
let idx = t * 8 + i;
h.write_at(idx * 4096, &vec![idx as u8; 4096]).unwrap();
}
})
})
.collect();
for t in threads {
t.join().unwrap();
}
h.commit(0o644, 0, true).unwrap();
let got = std::fs::read(&dest).unwrap();
for idx in 0..64u64 {
let s = (idx * 4096) as usize;
assert!(
got[s..s + 4096].iter().all(|&b| b == idx as u8),
"chunk {idx}"
);
}
}
#[test]
fn reopening_a_partial_preserves_bytes() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("resume.bin");
{
let h = WriteHandle::open(&dest, 8192, true).unwrap();
h.write_at(0, &[9u8; 4096]).unwrap();
h.sync().unwrap();
}
let h = WriteHandle::open(&dest, 8192, true).unwrap();
let mut buf = [0u8; 4096];
h.read_at(0, &mut buf).unwrap();
assert!(buf.iter().all(|&b| b == 9));
}
#[test]
fn registry_opens_once_under_contention() {
let tmp = tempfile::tempdir().unwrap();
let reg = Arc::new(FileWriters::new());
let opens = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let threads: Vec<_> = (0..16)
.map(|_| {
let reg = reg.clone();
let opens = opens.clone();
let path = tmp.path().join("shared.bin");
std::thread::spawn(move || {
reg.get_or_open(1, || {
opens.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
WriteHandle::open(&path, 1024, false)
})
.unwrap()
})
})
.collect();
let handles: Vec<_> = threads.into_iter().map(|t| t.join().unwrap()).collect();
for h in &handles {
assert!(Arc::ptr_eq(h, &handles[0]));
}
assert_eq!(reg.len(), 1);
assert_eq!(
opens.load(std::sync::atomic::Ordering::SeqCst),
1,
"the destination was opened more than once"
);
}
#[cfg(unix)]
fn allocated_bytes(p: &Path) -> u64 {
use std::os::unix::fs::MetadataExt;
std::fs::metadata(p)
.map(|m| m.blocks().saturating_mul(512))
.unwrap_or(0)
}
#[cfg(unix)]
#[test]
fn preallocation_reserves_exactly_the_file_size() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("prealloc.bin");
let size = 64 * 1024 * 1024;
{
let h = WriteHandle::open(&dest, size, true).unwrap();
h.write_at(0, &[1u8; 4096]).unwrap();
h.sync().unwrap();
}
let after_first = allocated_bytes(&part_path_for(&dest));
assert!(
after_first <= size + size / 8,
"a {size} byte file reserved {after_first} bytes on first open"
);
for _ in 0..3 {
let h = WriteHandle::open(&dest, size, true).unwrap();
h.write_at(0, &[1u8; 4096]).unwrap();
h.sync().unwrap();
}
let after_reopens = allocated_bytes(&part_path_for(&dest));
assert!(
after_reopens <= size + size / 8,
"reopening grew the reservation to {after_reopens} for a {size} byte file"
);
}
#[test]
fn zero_length_file_commits() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("empty.bin");
let h = WriteHandle::open(&dest, 0, true).unwrap();
h.commit(0o644, 0, true).unwrap();
assert!(dest.exists());
assert_eq!(std::fs::metadata(&dest).unwrap().len(), 0);
}
}