use std::fs::{File, OpenOptions};
use std::io::{BufWriter, Write};
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use anyhow::{Result, anyhow, bail};
use znippy_common::arrow;
use znippy_common::arrow::array::{ArrayRef, UInt64Array};
use znippy_common::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use znippy_common::arrow::ipc::MetadataVersion;
use znippy_common::arrow::ipc::writer::{
DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message,
};
use znippy_common::arrow::record_batch::RecordBatch;
pub type Extent = (u64, u64);
pub trait ArchiveWrite: Send + Sync {
fn append(&self, bytes: &[u8]) -> Result<Extent>;
fn name(&self) -> &'static str;
fn durability(&self) -> &'static str;
}
pub const JOURNAL_ALIGNMENT: u8 = 8;
pub fn journal_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("blob_offset", DataType::UInt64, false),
Field::new("blob_size", DataType::UInt64, false),
]))
}
pub fn journal_options() -> Result<IpcWriteOptions> {
IpcWriteOptions::try_new(JOURNAL_ALIGNMENT as usize, false, MetadataVersion::V5)
.map_err(|e| anyhow!("journal write options: {e}"))
}
pub fn journal_batch(offset: u64, len: u64) -> Result<RecordBatch> {
let off: ArrayRef = Arc::new(UInt64Array::from(vec![offset]));
let size: ArrayRef = Arc::new(UInt64Array::from(vec![len]));
RecordBatch::try_new(journal_schema(), vec![off, size])
.map_err(|e| anyhow!("journal batch: {e}"))
}
pub fn encode_journal_schema() -> Result<Vec<u8>> {
let opts = journal_options()?;
let dg = IpcDataGenerator {};
let mut tracker = DictionaryTracker::new(false);
let encoded =
dg.schema_to_bytes_with_dictionary_tracker(journal_schema().as_ref(), &mut tracker, &opts);
let mut out = Vec::with_capacity(512);
write_message(&mut out, encoded, &opts).map_err(|e| anyhow!("journal schema encode: {e}"))?;
Ok(out)
}
pub fn encode_journal_row(offset: u64, len: u64) -> Result<Vec<u8>> {
let opts = journal_options()?;
let dg = IpcDataGenerator {};
let mut tracker = DictionaryTracker::new(false);
let batch = journal_batch(offset, len)?;
let (dicts, msg) = dg
.encode(&batch, &mut tracker, &opts, &mut Default::default())
.map_err(|e| anyhow!("journal batch encode: {e}"))?;
let mut out = Vec::with_capacity(512);
for d in dicts {
write_message(&mut out, d, &opts).map_err(|e| anyhow!("journal dict encode: {e}"))?;
}
write_message(&mut out, msg, &opts).map_err(|e| anyhow!("journal row encode: {e}"))?;
Ok(out)
}
pub fn read_journal(path: &Path) -> Result<Vec<Extent>> {
let f = File::open(path).map_err(|e| anyhow!("journal open {}: {e}", path.display()))?;
if f.metadata()?.len() == 0 {
return Ok(Vec::new());
}
let reader = match arrow::ipc::reader::StreamReader::try_new(std::io::BufReader::new(f), None) {
Ok(r) => r,
Err(_) => return Ok(Vec::new()),
};
let mut out = Vec::new();
for batch in reader {
let Ok(batch) = batch else { break }; let offs = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("journal column 0 is not u64"))?;
let sizes = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("journal column 1 is not u64"))?;
for i in 0..batch.num_rows() {
out.push((offs.value(i), sizes.value(i)));
}
}
Ok(out)
}
pub const JOURNAL_TOMBSTONE: u64 = u64::MAX;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JournalRow {
Pack(Extent),
Retired(u64),
}
pub fn journal_rows(rows: &[Extent]) -> Vec<JournalRow> {
rows.iter()
.map(|&(offset, len)| {
if len == JOURNAL_TOMBSTONE {
JournalRow::Retired(offset)
} else {
JournalRow::Pack((offset, len))
}
})
.collect()
}
pub fn acked_packs(rows: &[Extent]) -> Vec<Extent> {
journal_rows(rows)
.into_iter()
.filter_map(|r| match r {
JournalRow::Pack(e) => Some(e),
JournalRow::Retired(_) => None,
})
.collect()
}
pub fn retired_offsets(rows: &[Extent]) -> std::collections::HashSet<u64> {
journal_rows(rows)
.into_iter()
.filter_map(|r| match r {
JournalRow::Retired(o) => Some(o),
JournalRow::Pack(_) => None,
})
.collect()
}
pub fn retire_packs(journal: &Path, offsets: &[u64]) -> Result<()> {
if offsets.is_empty() {
return Ok(());
}
let mut buf = Vec::with_capacity(offsets.len() * 256);
for &offset in offsets {
buf.extend_from_slice(&encode_journal_row(offset, JOURNAL_TOMBSTONE)?);
}
let f = OpenOptions::new()
.append(true)
.open(journal)
.map_err(|e| anyhow!("open journal {} to retire a pack: {e}", journal.display()))?;
if f.metadata()?.len() == 0 {
bail!(
"{} is empty — nothing was ever acked here, so there is no pack to retire",
journal.display()
);
}
(&f).write_all(&buf)
.map_err(|e| anyhow!("journal tombstone write: {e}"))?;
f.sync_all()
.map_err(|e| anyhow!("journal tombstone fsync: {e}"))?;
Ok(())
}
fn lock_for_writing(f: &File, path: &Path) -> Result<()> {
match f.try_lock() {
Ok(()) => Ok(()),
Err(std::fs::TryLockError::WouldBlock) => bail!(
"another writer already holds {} — one process at a time appends to a store's \
blob file. Two would share a stale append cursor and overwrite each other's \
acked packs without either one erroring. The lock is an advisory flock and the \
kernel drops it when that process dies, so nothing has to be cleaned up by hand.",
path.display()
),
Err(std::fs::TryLockError::Error(e)) => Err(anyhow!(
"locking {} for writing: {e} — the store is not opened without the writer lock, \
because an unlocked open is the P-014 race",
path.display()
)),
}
}
pub(crate) fn open_blobs(path: &Path) -> Result<(File, u64)> {
let f = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(path)
.map_err(|e| anyhow!("open blobs {}: {e}", path.display()))?;
lock_for_writing(&f, path)?;
let end = f.metadata()?.len();
Ok((f, end))
}
pub struct FastWriter {
blobs: File,
cursor: AtomicU64,
}
impl FastWriter {
pub fn create(archive: &Path) -> Result<Self> {
let (blobs, end) = open_blobs(archive)?;
Ok(Self {
blobs,
cursor: AtomicU64::new(end),
})
}
pub fn blobs(&self) -> &File {
&self.blobs
}
}
impl ArchiveWrite for FastWriter {
fn append(&self, bytes: &[u8]) -> Result<Extent> {
let len = bytes.len() as u64;
let offset = self.cursor.fetch_add(len, Ordering::SeqCst);
self.blobs.write_all_at(bytes, offset)?;
Ok((offset, len))
}
fn name(&self) -> &'static str {
"FastWriter"
}
fn durability(&self) -> &'static str {
"none — page cache only; a machine crash after return loses the bytes"
}
}
struct SafeJournal {
writer: BufWriter<File>,
path: PathBuf,
}
pub struct SafeWriter {
blobs: File,
cursor: AtomicU64,
journal: Mutex<SafeJournal>,
faults: Faults,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Faults {
pub skip_blob_fsync: bool,
pub die_between_fsyncs: bool,
}
impl SafeWriter {
pub fn create(archive: &Path) -> Result<Self> {
Self::create_inner(archive, Faults::default())
}
pub fn create_with_faults(archive: &Path, faults: Faults) -> Result<Self> {
Self::create_inner(archive, faults)
}
fn create_inner(archive: &Path, faults: Faults) -> Result<Self> {
let (blobs, end) = open_blobs(archive)?;
let path = journal_path(archive);
let f = OpenOptions::new()
.read(true)
.append(true)
.create(true)
.open(&path)
.map_err(|e| anyhow!("open journal {}: {e}", path.display()))?;
let already = f.metadata()?.len();
let mut j = SafeJournal {
writer: BufWriter::new(f),
path,
};
if already == 0 {
j.writer
.write_all(&encode_journal_schema()?)
.map_err(|e| anyhow!("journal schema: {e}"))?;
j.writer.flush().map_err(|e| anyhow!("journal flush: {e}"))?;
j.writer
.get_ref()
.sync_all()
.map_err(|e| anyhow!("journal fsync: {e}"))?;
}
Ok(Self {
blobs,
cursor: AtomicU64::new(end),
journal: Mutex::new(j),
faults,
})
}
pub fn journal_path(archive: &Path) -> PathBuf {
journal_path(archive)
}
pub fn blobs(&self) -> &File {
&self.blobs
}
pub fn journal_file(&self) -> PathBuf {
self.journal.lock().expect("journal mutex").path.clone()
}
}
pub(crate) fn journal_path(archive: &Path) -> PathBuf {
let mut s = archive.as_os_str().to_os_string();
s.push(".journal");
PathBuf::from(s)
}
impl ArchiveWrite for SafeWriter {
fn append(&self, bytes: &[u8]) -> Result<Extent> {
let len = bytes.len() as u64;
let offset = self.cursor.fetch_add(len, Ordering::SeqCst);
self.blobs.write_all_at(bytes, offset)?;
if !self.faults.skip_blob_fsync {
self.blobs.sync_all()?;
}
if self.faults.die_between_fsyncs {
bail!("injected crash between the two fsyncs");
}
let mut j = self
.journal
.lock()
.map_err(|_| anyhow!("journal mutex poisoned"))?;
let row = encode_journal_row(offset, len)?;
j.writer
.write_all(&row)
.map_err(|e| anyhow!("journal write: {e}"))?;
j.writer.flush().map_err(|e| anyhow!("journal flush: {e}"))?;
j.writer
.get_ref()
.sync_all()
.map_err(|e| anyhow!("journal fsync: {e}"))?;
Ok((offset, len))
}
fn name(&self) -> &'static str {
"SafeWriter"
}
fn durability(&self) -> &'static str {
"full — blob fsynced, then a journal row fsynced; crash after return keeps both"
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SealReport {
pub archive: PathBuf,
pub packs_sealed: u64,
pub packs_retired: u64,
pub packs_after_copy: u64,
pub blob_bytes: u64,
pub sealed_total_bytes: u64,
}
pub fn seal_generation_zero(
blobs: &Path,
journal: Option<&Path>,
archive: &Path,
reserved: Vec<znippy_common::ReservedSection>,
) -> Result<SealReport> {
use znippy_common::index::{ChunkLoc, data_subindex_schema};
use znippy_common::{ArchiveMetaSink, ArrowIpcSink, GroupKey, base_batch_from_rows};
let staged = staging_sibling(archive);
let sealed = (|| -> Result<SealReport> {
let blob_bytes = std::fs::copy(blobs, &staged).map_err(|e| {
anyhow!(
"copying the blob region {} into {}: {e}",
blobs.display(),
staged.display()
)
})?;
let rows = match journal {
Some(p) if p.exists() => read_journal(p)?,
_ => Vec::new(),
};
let packs = acked_packs(&rows);
let retired = retired_offsets(&rows);
let file = OpenOptions::new()
.read(true)
.write(true)
.open(&staged)
.map_err(|e| anyhow!("reopening {} to seal into: {e}", staged.display()))?;
let mut paths: Vec<String> = Vec::with_capacity(packs.len());
let mut locs: Vec<ChunkLoc> = Vec::with_capacity(packs.len());
let mut packs_retired = 0u64;
let mut packs_after_copy = 0u64;
for (ordinal, &(offset, len)) in packs.iter().enumerate() {
if retired.contains(&offset) {
packs_retired += 1;
continue;
}
if offset.saturating_add(len) > blob_bytes {
packs_after_copy += 1;
continue;
}
paths.push(format!("objects.pack.{ordinal}"));
locs.push(ChunkLoc {
chunk_seq: 0,
fdata_offset: 0,
blob_offset: offset,
blob_size: len,
uncompressed_size: len,
compressed: false,
checksum: blake3_extent(&file, offset, len)?,
});
}
let packs_sealed = paths.len() as u64;
let file = Arc::new(file);
let mut sink = ArrowIpcSink::new(Arc::clone(&file), blob_bytes);
if !reserved.is_empty() {
sink = sink.with_reserved_builder(Box::new(move |_| Ok(reserved)));
}
if !paths.is_empty() {
let batch = base_batch_from_rows(&paths, &locs)?;
sink.push_subindex(
data_subindex_schema().as_ref(),
&[batch],
GroupKey {
pkg_type: 0,
repo: String::new(),
module_name: String::new(),
},
)?;
}
let sealed_total_bytes = Box::new(sink).finish()?;
Ok(SealReport {
archive: archive.to_path_buf(),
packs_sealed,
packs_retired,
packs_after_copy,
blob_bytes,
sealed_total_bytes,
})
})();
let sealed = match sealed {
Ok(s) => s,
Err(e) => {
let _ = std::fs::remove_file(&staged);
return Err(e);
}
};
std::fs::rename(&staged, archive).map_err(|e| {
let _ = std::fs::remove_file(&staged);
anyhow!(
"renaming {} into place as {}: {e}",
staged.display(),
archive.display()
)
})?;
sync_parent_dir(archive);
Ok(sealed)
}
fn blake3_extent(file: &File, offset: u64, len: u64) -> Result<[u8; 32]> {
const WINDOW: usize = 1 << 20;
let mut hasher = znippy_common::blake3::Hasher::new();
let mut buf = vec![0u8; WINDOW.min(len.max(1) as usize)];
let mut done = 0u64;
while done < len {
let want = ((len - done) as usize).min(buf.len());
file.read_exact_at(&mut buf[..want], offset + done)
.map_err(|e| anyhow!("reading the pack at ({offset}, {len}) to checksum it: {e}"))?;
hasher.update(&buf[..want]);
done += want as u64;
}
Ok(*hasher.finalize().as_bytes())
}
fn staging_sibling(archive: &Path) -> PathBuf {
let unique = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let mut p = archive.as_os_str().to_owned();
p.push(format!(".seal-{}-{unique}", std::process::id()));
PathBuf::from(p)
}
fn sync_parent_dir(path: &Path) {
if let Some(parent) = path.parent()
&& let Ok(f) = File::open(parent)
{
let _ = f.sync_all();
}
}