use std::fs::{self, File};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::ptr::NonNull;
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_buffer::Buffer;
use arrow_ipc::MetadataVersion;
use arrow_ipc::reader::{FileDecoder, read_footer_length};
use arrow_ipc::root_as_footer;
use arrow_ipc::writer::{FileWriter, IpcWriteOptions};
use memmap2::Mmap;
use crate::error::{Error, IoContext, Result};
use crate::signal::Sealed;
const MAGIC: &[u8; 6] = b"ARROW1";
const CRC_KEY: &str = "mira.crc32";
const CRC_LEN_KEY: &str = "mira.crc32.len";
const FORMAT_KEY: &str = "mira.format";
pub const FORMAT_VERSION: u32 = 1;
const LEGACY_VERSION: u32 = 1;
const ZSTD_LEVEL: i32 = 3;
const HEADER_LEN: usize = MAGIC.len().next_multiple_of(ALIGNMENT);
const CONTINUATION: [u8; 4] = [0xff; 4];
fn corrupt(path: &Path, what: String) -> Error {
Error::Io {
path: path.to_path_buf(),
source: io::Error::new(io::ErrorKind::InvalidData, what),
}
}
const ALIGNMENT: usize = 64;
const NANOS_PER_HOUR: i64 = 3_600 * 1_000_000_000;
struct CrcWriter<W> {
inner: W,
hasher: crc32fast::Hasher,
written: u64,
}
impl<W: Write> CrcWriter<W> {
fn new(inner: W) -> Self {
Self {
inner,
hasher: crc32fast::Hasher::new(),
written: 0,
}
}
fn checksum(&self) -> (u64, u32) {
(self.written, self.hasher.clone().finalize())
}
}
impl<W: Write> Write for CrcWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let n = self.inner.write(buf)?;
self.hasher.update(&buf[..n]);
self.written += n as u64;
Ok(n)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
pub fn write_table(path: &Path, batch: &RecordBatch) -> Result<()> {
write_table_with(path, std::slice::from_ref(batch), None)
}
pub fn write_table_zstd(path: &Path, batch: &RecordBatch) -> Result<()> {
write_table_with(
path,
std::slice::from_ref(batch),
Some(arrow_ipc::CompressionType::ZSTD),
)
}
pub fn write_table_lz4(path: &Path, batch: &RecordBatch) -> Result<()> {
write_table_with(
path,
std::slice::from_ref(batch),
Some(arrow_ipc::CompressionType::LZ4_FRAME),
)
}
fn write_table_with(
path: &Path,
batches: &[RecordBatch],
codec: Option<arrow_ipc::CompressionType>,
) -> Result<()> {
let Some(first) = batches.first() else {
return Ok(());
};
let file = File::create(path).ctx(path)?;
let opts = IpcWriteOptions::try_new(ALIGNMENT, false, MetadataVersion::V5)?;
let opts = match codec {
Some(arrow_ipc::CompressionType::ZSTD) => opts
.try_with_compression(codec)?
.try_with_compression_level(Some(ZSTD_LEVEL))?,
Some(c) => opts.try_with_compression(Some(c))?,
None => opts,
};
let mut w = FileWriter::try_new_with_options(
CrcWriter::new(BufWriter::new(file)),
&first.schema(),
opts,
)?;
for batch in batches {
w.write(batch)?;
}
let (len, crc) = w.get_ref().checksum();
w.write_metadata(FORMAT_KEY, FORMAT_VERSION.to_string());
w.write_metadata(CRC_KEY, format!("{crc:08x}"));
w.write_metadata(CRC_LEN_KEY, len.to_string());
w.finish()?;
let mut buf = w.into_inner()?;
buf.flush().ctx(path)?;
let file = buf.inner.into_inner().map_err(|e| Error::Io {
path: path.to_path_buf(),
source: e.into_error(),
})?;
crate::sync_all(&file).ctx(path)?;
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BlockRef {
pub dir: PathBuf,
pub min_ts: i64,
pub max_ts: i64,
pub node: u32,
pub seq: u64,
pub wal_hi: u64,
}
impl BlockRef {
pub fn overlaps(&self, from: i64, to: i64) -> bool {
self.min_ts <= to && self.max_ts >= from
}
}
pub(crate) struct Src<'a> {
pub node: u32,
pub seq: u64,
pub min_ts: i64,
pub max_ts: i64,
pub dir: Option<&'a Path>,
tables: Option<&'a [(&'static str, RecordBatch)]>,
}
impl<'a> Src<'a> {
pub fn disk(b: &'a BlockRef) -> Src<'a> {
Src {
node: b.node,
seq: b.seq,
min_ts: b.min_ts,
max_ts: b.max_ts,
dir: Some(b.dir.as_path()),
tables: None,
}
}
pub fn open(o: &'a crate::signal::Open) -> Src<'a> {
Src {
node: o.node,
seq: o.seq,
min_ts: o.sealed.min_ts,
max_ts: o.sealed.max_ts,
dir: None,
tables: Some(&o.sealed.tables),
}
}
pub fn overlaps(&self, from: i64, to: i64) -> bool {
self.min_ts <= to && self.max_ts >= from
}
pub fn load(&self, name: &str) -> Result<Option<RecordBatch>> {
match (self.tables, self.dir) {
(Some(tables), _) => Ok(tables
.iter()
.find(|(n, _)| *n == name)
.map(|(_, b)| b.clone())),
(None, Some(dir)) => Ok(open_table_opt(&dir.join(format!("{name}.arrow")))?
.and_then(|t| t.batches.first().cloned())),
(None, None) => Ok(None),
}
}
}
pub(crate) fn sources<'a>(
disk: &'a [BlockRef],
open: &'a [std::sync::Arc<crate::signal::Open>],
) -> Vec<Src<'a>> {
disk.iter()
.map(Src::disk)
.chain(
open.iter()
.filter(|o| !disk.iter().any(|b| b.node == o.node && b.seq == o.seq))
.map(|o| Src::open(o)),
)
.collect()
}
pub fn node_id(name: &str) -> u32 {
crate::identity::hash64(name.as_bytes()) as u32
}
fn dir_name(min_ts: i64, max_ts: i64, node: u32, seq: u64, wal_hi: u64) -> String {
format!("{min_ts:020}-{max_ts:020}-{node:08x}-{seq:012}-{wal_hi:020}")
}
fn parse_dir_name(name: &str) -> Option<(i64, i64, u32, u64, u64)> {
let mut parts = name.split('-');
let min = parts.next()?.parse().ok()?;
let max = parts.next()?.parse().ok()?;
let node = u32::from_str_radix(parts.next()?, 16).ok()?;
let seq = parts.next()?.parse().ok()?;
let wal_hi = match parts.next() {
Some(field) => field.parse().ok()?,
None => 0,
};
if parts.next().is_some() {
return None;
}
Some((min, max, node, seq, wal_hi))
}
fn fsync_dir(path: &Path) -> Result<()> {
crate::sync_all(&File::open(path).ctx(path)?).ctx(path)
}
pub fn publish(
root: &Path,
signal: &str,
node: u32,
seq: u64,
wal_hi: u64,
sealed: &Sealed,
) -> Result<BlockRef> {
let (min_ts, max_ts) = (sealed.min_ts, sealed.max_ts);
let staging = root.join(".tmp");
let tmp = staging.join(format!(
"{signal}-{node:08x}-{seq:012}-{min_ts:020}-{max_ts:020}"
));
fs::create_dir_all(&staging).ctx(&staging)?;
fs::create_dir(&tmp).ctx(&tmp)?;
let signal_dir = root.join(signal);
let partition = signal_dir.join(format!("p={}", min_ts.div_euclid(NANOS_PER_HOUR)));
let dir = partition.join(dir_name(min_ts, max_ts, node, seq, wal_hi));
let staged = stage(&tmp, sealed).and_then(|()| {
fs::create_dir_all(&partition).ctx(&partition)?;
fs::rename(&tmp, &dir).ctx(&dir)?;
fsync_dir(&partition)?;
fsync_dir(&signal_dir)
});
if let Err(e) = staged {
unwind_staging(&tmp);
return Err(e);
}
Ok(BlockRef {
dir,
min_ts,
max_ts,
node,
seq,
wal_hi,
})
}
fn unwind_staging(tmp: &Path) {
if let Err(rm) = fs::remove_dir_all(tmp) {
if rm.kind() != io::ErrorKind::NotFound {
tracing::warn!(
path = %tmp.display(),
error = %rm,
"cannot remove the staging directory of a failed publish",
);
}
}
}
fn stage(tmp: &Path, sealed: &Sealed) -> Result<()> {
for (name, batch) in &sealed.tables {
if batch.num_rows() > 0 {
write_table(&tmp.join(format!("{name}.arrow")), batch)?;
}
}
for (name, bytes) in &sealed.sidecars {
let path = tmp.join(name);
let mut f = File::create(&path).ctx(&path)?;
f.write_all(bytes).ctx(&path)?;
crate::sync_all(&f).ctx(&path)?;
}
fsync_dir(tmp)
}
pub fn sweep_staging(root: &Path, signal: &str, node: u32) -> Result<usize> {
let tmp = root.join(".tmp");
let prefix = format!("{signal}-{node:08x}-");
let entries = match fs::read_dir(&tmp) {
Ok(d) => d,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(0),
Err(e) => {
return Err(Error::Io {
path: tmp,
source: e,
});
}
};
let mut removed = 0;
for entry in entries {
let path = entry.ctx(&tmp)?.path();
if path
.file_name()
.and_then(|n| n.to_str())
.is_some_and(|n| n.starts_with(&prefix))
{
fs::remove_dir_all(&path).ctx(&path)?;
removed += 1;
}
}
Ok(removed)
}
pub fn wal_watermarks(root: &Path) -> Result<crate::wal::Watermarks> {
let mut out = [0u64; 3];
for signal in crate::wal::Signal::ALL {
out[signal.index()] = scan(root, signal.as_str())?
.iter()
.map(|b| b.wal_hi)
.max()
.unwrap_or(0);
}
Ok(out)
}
pub fn scan(root: &Path, signal: &str) -> Result<Vec<BlockRef>> {
let base = root.join(signal);
let mut out = Vec::new();
let partitions = match fs::read_dir(&base) {
Ok(d) => d,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(out),
Err(e) => {
return Err(Error::Io {
path: base,
source: e,
});
}
};
for partition in partitions {
let partition = partition.ctx(&base)?.path();
if !partition.is_dir() {
continue;
}
for entry in fs::read_dir(&partition).ctx(&partition)? {
let dir = entry.ctx(&partition)?.path();
let Some((min_ts, max_ts, node, seq, wal_hi)) = dir
.file_name()
.and_then(|n| n.to_str())
.and_then(parse_dir_name)
else {
continue;
};
out.push(BlockRef {
dir,
min_ts,
max_ts,
node,
seq,
wal_hi,
});
}
}
out.sort_by_key(|b| (b.min_ts, b.seq));
Ok(out)
}
pub fn expire(root: &Path, signal: &str, cutoff_ns: i64) -> Result<usize> {
let mut dropped = 0;
for block in scan(root, signal)? {
if block.max_ts < cutoff_ns {
match fs::remove_dir_all(&block.dir) {
Ok(()) => dropped += 1,
Err(e) if e.kind() == io::ErrorKind::NotFound => {}
Err(e) => tracing::warn!(
block = %block.dir.display(),
error = %e,
"cannot expire block; skipping it",
),
}
}
}
Ok(dropped)
}
pub fn check_filesystem(path: &Path) -> Result<()> {
check_fs_type(path, fs_type(path)?)
}
fn check_fs_type(path: &Path, fs: Option<String>) -> Result<()> {
let Some(fs) = fs else {
return Ok(());
};
if fs == "fuse" {
tracing::warn!(
path = %path.display(),
"data directory is on a FUSE filesystem. If it is network-backed \
(gcsfuse, s3fs, rclone), mmap will raise SIGBUS and kill the \
process; if it is local, ignore this."
);
return Ok(());
}
Err(Error::NetworkFilesystem {
path: path.to_path_buf(),
fs,
})
}
pub fn check_writable(path: &Path) -> Result<()> {
let probe = path.join(format!(".mira-write-probe-{}", std::process::id()));
let fail = |source| Error::NotWritable {
path: path.to_path_buf(),
source,
};
fs::write(&probe, []).map_err(fail)?;
fs::remove_file(&probe).map_err(fail)
}
fn statfs(path: &Path) -> Result<libc::statfs> {
let c = std::ffi::CString::new(path.as_os_str().as_encoded_bytes()).map_err(|_| Error::Io {
path: path.to_path_buf(),
source: io::Error::new(io::ErrorKind::InvalidInput, "path contains a NUL byte"),
})?;
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
if unsafe { libc::statfs(c.as_ptr(), &mut buf) } != 0 {
return Err(Error::Io {
path: path.to_path_buf(),
source: io::Error::last_os_error(),
});
}
Ok(buf)
}
pub fn free_fraction(path: &Path) -> Result<f64> {
let buf = statfs(path)?;
if buf.f_blocks == 0 {
return Ok(1.0);
}
Ok(buf.f_bavail as f64 / buf.f_blocks as f64)
}
fn fs_type(path: &Path) -> Result<Option<String>> {
if !path.exists() {
return Ok(None);
}
let buf = statfs(path)?;
#[cfg(target_os = "macos")]
{
let name: Vec<u8> = buf
.f_fstypename
.iter()
.take_while(|&&c| c != 0)
.map(|&c| c as u8)
.collect();
let name = String::from_utf8_lossy(&name).into_owned();
Ok(match name.as_str() {
"nfs" | "smbfs" | "cifs" | "webdav" | "afpfs" | "ftp" => Some(name),
n if n.contains("fuse") => Some("fuse".into()),
_ => None,
})
}
#[cfg(not(target_os = "macos"))]
{
let ty = (buf.f_type as u64) & 0xffff_ffff;
Ok(match ty {
0x6969 => Some("NFS".into()),
0x517b => Some("SMB".into()),
0xff53_4d42 => Some("CIFS".into()),
0xfe53_4d42 => Some("SMB2".into()),
0x0102_1997 => Some("9P".into()),
0x5346_414f => Some("AFS".into()),
0x00c3_6400 => Some("CephFS".into()),
0x0116_1970 => Some("GFS2".into()),
0x7461_636f => Some("OCFS2".into()),
0x0bd0_0bd0 => Some("Lustre".into()),
0x6573_5546 => Some("fuse".into()),
_ => None,
})
}
}
const COLD_MARKER: &str = "cold";
pub const COLD_AFTER_NS: i64 = NANOS_PER_HOUR;
const MAX_COMPACT_PER_SWEEP: usize = 8;
pub fn compact(root: &Path, signal: &str, node: u32, cutoff_ns: i64) -> Result<usize> {
let (mut done, mut tried) = (0, 0);
for block in scan(root, signal)? {
if tried == MAX_COMPACT_PER_SWEEP {
break;
}
if block.max_ts >= cutoff_ns || block.dir.join(COLD_MARKER).exists() {
continue;
}
tried += 1;
match compact_block(&block.dir, node) {
Ok(()) => done += 1,
Err(Error::Io { source, .. }) if source.kind() == io::ErrorKind::NotFound => {}
Err(e) => tracing::warn!(
block = %block.dir.display(),
error = %e,
"cannot compact block; skipping it",
),
}
}
Ok(done)
}
fn compact_block(dir: &Path, node: u32) -> Result<()> {
let mut tables = Vec::new();
for entry in fs::read_dir(dir).ctx(dir)? {
let path = entry.ctx(dir)?.path();
if path.extension().is_some_and(|e| e == "arrow") {
tables.push(path);
}
}
for path in tables {
let table = open_table(&path)?;
let tmp = path.with_extension(format!("{node:08x}.tmp"));
write_table_with(&tmp, &table.batches, Some(arrow_ipc::CompressionType::ZSTD))?;
fs::rename(&tmp, &path).ctx(&path)?;
}
crate::sync_all(&File::create(dir.join(COLD_MARKER)).ctx(dir)?).ctx(dir)?;
fsync_dir(dir)
}
pub struct MappedTable {
pub batches: Vec<RecordBatch>,
mapping: std::ops::Range<usize>,
}
impl MappedTable {
pub fn zero_copy_ratio(&self) -> (usize, usize) {
let (mut inside, mut total) = (0, 0);
for batch in &self.batches {
for col in batch.columns() {
let mut stack = vec![col.to_data()];
while let Some(d) = stack.pop() {
for buf in d.buffers() {
total += 1;
if self.mapping.contains(&(buf.as_ptr() as usize)) {
inside += 1;
}
}
stack.extend(d.child_data().iter().cloned());
}
}
}
(inside, total)
}
}
pub fn open_table_opt(path: &Path) -> Result<Option<MappedTable>> {
match open_table(path) {
Err(Error::Io { source, .. }) if source.kind() == std::io::ErrorKind::NotFound => Ok(None),
other => other.map(Some),
}
}
fn check_format(path: &Path, declared: Option<&str>) -> Result<()> {
let version = match declared {
None => LEGACY_VERSION,
Some(v) => v
.parse::<u32>()
.map_err(|_| corrupt(path, format!("{FORMAT_KEY} is `{v}`, not a version")))?,
};
if version > FORMAT_VERSION {
return Err(Error::Io {
path: path.to_path_buf(),
source: io::Error::new(
io::ErrorKind::Unsupported,
format!(
"block format version {version} was written by a newer Mira; \
this binary reads up to version {FORMAT_VERSION}"
),
),
});
}
Ok(())
}
struct Framed {
block: arrow_ipc::Block,
data: Buffer,
header: arrow_ipc::MessageHeader,
version: MetadataVersion,
}
impl Framed {
fn stride(&self) -> usize {
self.block.metaDataLength() as usize + self.block.bodyLength() as usize
}
}
fn message_at(path: &Path, buffer: &Buffer, offset: usize, body_len: usize) -> Result<Framed> {
let bad = |why: String| corrupt(path, format!("message at offset {offset}: {why}"));
if offset > body_len || body_len - offset < 8 {
return Err(bad(format!(
"outside the {body_len} bytes the checksum covers"
)));
}
let head = &buffer[offset..offset + 8];
if head[..4] != CONTINUATION {
return Err(bad("no message framing here".into()));
}
let declared = i32::from_le_bytes(head[4..].try_into().expect("4 bytes"));
let Some(meta) = declared.checked_add(8).filter(|m| *m >= 8) else {
return Err(bad(format!("{declared} is not a metadata length")));
};
let meta = meta as usize;
if body_len - offset < meta {
return Err(bad(format!(
"{meta} bytes of metadata run past the {body_len} bytes the checksum covers"
)));
}
let message = arrow_ipc::root_as_message(&buffer[offset + 8..offset + meta])
.map_err(|e| bad(format!("not a readable IPC message: {e}")))?;
let body = message.bodyLength();
let room = (body_len - offset - meta) as i64;
if body < 0 || body > room {
return Err(bad(format!(
"a {body}-byte body runs past the {body_len} bytes the checksum covers"
)));
}
Ok(Framed {
block: arrow_ipc::Block::new(0, meta as i32, body),
data: buffer.slice_with_length(offset, meta + body as usize),
header: message.header_type(),
version: message.version(),
})
}
pub fn open_table(path: &Path) -> Result<MappedTable> {
let file = File::open(path).ctx(path)?;
let mmap = unsafe { Mmap::map(&file) }.ctx(path)?;
let _ = mmap.advise(memmap2::Advice::WillNeed);
if mmap.len() < MAGIC.len() + 10 || &mmap[..MAGIC.len()] != MAGIC {
return Err(Error::BadMagic {
path: path.to_path_buf(),
});
}
let len = mmap.len();
let base = mmap.as_ptr() as usize;
let ptr = NonNull::new(mmap.as_ptr().cast_mut()).expect("mmap is never null");
let buffer = unsafe { Buffer::from_custom_allocation(ptr, len, Arc::new(mmap)) };
let trailer = len - 10;
let footer_len = read_footer_length(buffer[trailer..].try_into().expect("10 bytes"))?;
if footer_len > trailer {
return Err(corrupt(
path,
format!("footer says it is {footer_len} bytes, in a {len}-byte file"),
));
}
let footer_start = trailer - footer_len;
let footer = root_as_footer(&buffer[footer_start..trailer])
.map_err(|e| arrow_schema::ArrowError::ParseError(e.to_string()))?;
let find = |key: &str| -> Option<&str> {
footer
.custom_metadata()
.into_iter()
.flatten()
.find(|kv| kv.key() == Some(key))
.and_then(|kv| kv.value())
};
let meta = |key: &'static str| -> Result<&str> {
find(key).ok_or(Error::MissingMetadata {
path: path.to_path_buf(),
key,
})
};
check_format(path, find(FORMAT_KEY))?;
let expected = u32::from_str_radix(meta(CRC_KEY)?, 16).map_err(|_| Error::MissingMetadata {
path: path.to_path_buf(),
key: CRC_KEY,
})?;
let body_len: usize = meta(CRC_LEN_KEY)?
.parse()
.ok()
.filter(|&n: &usize| n <= footer_start)
.ok_or(Error::MissingMetadata {
path: path.to_path_buf(),
key: CRC_LEN_KEY,
})?;
let actual = crc32fast::hash(&buffer[..body_len]);
if actual != expected {
return Err(Error::BadChecksum {
path: path.to_path_buf(),
expected,
actual,
});
}
let first = message_at(path, &buffer, HEADER_LEN, body_len)?;
if first.header != arrow_ipc::MessageHeader::Schema {
return Err(corrupt(
path,
format!(
"the body starts with {:?}, not a schema message",
first.header.variant_name().unwrap_or("an unknown message")
),
));
}
let schema = Arc::new(
arrow_ipc::convert::try_schema_from_ipc_buffer(&first.data).map_err(|e| {
corrupt(
path,
format!("no readable schema at the head of the body: {e}"),
)
})?,
);
let mut decoder = unsafe {
FileDecoder::new(schema, first.version)
.with_require_alignment(true)
.with_skip_validation(true)
};
let mut batches = Vec::new();
let mut offset = HEADER_LEN + first.stride();
while offset < body_len {
let msg = message_at(path, &buffer, offset, body_len)?;
offset += msg.stride();
let fail = |source| Error::Undecodable {
path: path.to_path_buf(),
source,
};
match msg.header {
arrow_ipc::MessageHeader::DictionaryBatch => {
decoder
.read_dictionary(&msg.block, &msg.data)
.map_err(fail)?;
}
arrow_ipc::MessageHeader::RecordBatch => {
if let Some(rb) = decoder
.read_record_batch(&msg.block, &msg.data)
.map_err(fail)?
{
batches.push(rb);
}
}
other => {
return Err(corrupt(
path,
format!(
"{} at offset {offset} is not a dictionary or a record batch",
other.variant_name().unwrap_or("an unknown message")
),
));
}
}
}
Ok(MappedTable {
batches,
mapping: base..base + len,
})
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::{StringArray, UInt32Array};
use arrow_schema::{DataType, Field, Schema};
fn dir(tag: &str) -> PathBuf {
let d = std::env::temp_dir().join(format!("mira-blk-{tag}-{}", std::process::id()));
let _ = fs::remove_dir_all(&d);
fs::create_dir_all(&d).unwrap();
d
}
fn batch() -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::UInt32, false),
Field::new("body", DataType::Utf8, true),
]));
RecordBatch::try_new(
schema,
vec![
Arc::new(UInt32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec![Some("a"), None, Some("ccc")])),
],
)
.unwrap()
}
fn compressible() -> RecordBatch {
let n = 4_096u32;
RecordBatch::try_new(
batch().schema(),
vec![
Arc::new(UInt32Array::from_iter_values(0..n)),
Arc::new(StringArray::from_iter_values(
(0..n).map(|_| "the same line of telemetry, over and over"),
)),
],
)
.unwrap()
}
fn sealed(min_ts: i64, max_ts: i64) -> Sealed {
Sealed {
tables: vec![("logs", batch())],
sidecars: vec![],
min_ts,
max_ts,
num_rows: 3,
}
}
const UNKNOWN_MESSAGE: u8 = 42;
fn footer_start(bytes: &[u8]) -> usize {
let n = bytes.len();
let footer_len = i32::from_le_bytes(bytes[n - 10..n - 6].try_into().unwrap()) as usize;
n - 10 - footer_len
}
fn crc_field(bytes: &[u8]) -> (usize, usize) {
let start = footer_start(bytes);
let footer = root_as_footer(&bytes[start..bytes.len() - 10]).unwrap();
let md = footer.custom_metadata().unwrap();
let find = |key: &str| {
md.iter()
.find(|kv| kv.key() == Some(key))
.and_then(|kv| kv.value())
.unwrap()
};
let hex = find(CRC_KEY);
assert_eq!(hex.len(), 8, "the CRC is written as eight hex digits");
(
hex.as_ptr() as usize - bytes.as_ptr() as usize,
find(CRC_LEN_KEY).parse().unwrap(),
)
}
fn repair_crc(bytes: &mut [u8]) {
let (at, body_len) = crc_field(bytes);
let crc = crc32fast::hash(&bytes[..body_len]);
bytes[at..at + 8].copy_from_slice(format!("{crc:08x}").as_bytes());
}
fn header_tag(bytes: &[u8], offset: usize) -> usize {
let declared =
i32::from_le_bytes(bytes[offset + 4..offset + 8].try_into().unwrap()) as usize;
let meta = offset + 8..offset + 8 + declared;
for i in meta.clone() {
let mut probe = bytes.to_vec();
probe[i] = UNKNOWN_MESSAGE;
let framed = arrow_ipc::root_as_message(&probe[meta.clone()]);
if framed.is_ok_and(|m| m.header_type().0 == UNKNOWN_MESSAGE) {
return i;
}
}
panic!("the message at {offset} has no header tag to corrupt");
}
fn second_message(path: &Path, bytes: &[u8]) -> usize {
let covered = crc_field(bytes).1;
let buffer = Buffer::from_vec(bytes.to_vec());
HEADER_LEN
+ message_at(path, &buffer, HEADER_LEN, covered)
.unwrap()
.stride()
}
fn invalid_data(e: &Error, what: &str) {
let Error::Io { source, .. } = e else {
panic!("{what}: {e}")
};
assert_eq!(source.kind(), io::ErrorKind::InvalidData, "{what}: {e}");
}
#[test]
fn a_corrupt_footer_length_is_an_error_not_a_panic() {
let d = dir("footerlen");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let good = fs::read(&path).unwrap();
for len in [i32::MAX, good.len() as i32, good.len() as i32 - 9] {
let mut bytes = good.clone();
let n = bytes.len();
bytes[n - 10..n - 6].copy_from_slice(&len.to_le_bytes());
fs::write(&path, &bytes).unwrap();
let Err(e) = open_table(&path) else {
panic!("a footer length of {len} in a {n}-byte file read as a block")
};
assert!(
matches!(&e, Error::Io { source, .. }
if source.kind() == io::ErrorKind::InvalidData),
"{e}"
);
}
let _ = fs::remove_dir_all(&d);
}
#[test]
fn no_bit_in_the_unchecked_footer_can_change_what_a_read_returns() {
let d = dir("footerbits");
let path = d.join("logs.arrow");
let want = batch();
write_table(&path, &want).unwrap();
let good = fs::read(&path).unwrap();
let (start, n) = (footer_start(&good), good.len());
assert!(start < n - 10, "no footer to corrupt");
let mut flipped = 0;
for i in start..n {
for bit in 0..8u8 {
let mut bytes = good.clone();
bytes[i] ^= 1 << bit;
fs::write(&path, &bytes).unwrap();
if let Ok(t) = open_table(&path) {
assert_eq!(
t.batches,
vec![want.clone()],
"byte {i} bit {bit} decoded to other data"
);
} else {
flipped += 1;
}
}
}
assert!(flipped > 0, "no corrupt footer was rejected");
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_footer_schema_naming_a_wider_type_is_not_the_one_the_read_uses() {
let d = dir("widen");
let path = d.join("logs.arrow");
let want = batch();
write_table(&path, &want).unwrap();
let good = fs::read(&path).unwrap();
let (start, end) = (footer_start(&good), good.len() - 10);
let body_tag = |bytes: &[u8]| -> Option<u8> {
let fields = root_as_footer(&bytes[start..end])
.ok()?
.schema()?
.fields()?;
let body = (fields.len() == 2).then(|| fields.get(1))?;
(body.name() == Some("body")).then(|| body.type_type().0)
};
assert_eq!(body_tag(&good), Some(arrow_ipc::Type::Utf8.0));
let mut widened = 0;
for i in start..end {
let mut bytes = good.clone();
bytes[i] = arrow_ipc::Type::LargeUtf8.0;
if body_tag(&bytes) != Some(arrow_ipc::Type::LargeUtf8.0) {
continue;
}
widened += 1;
fs::write(&path, &bytes).unwrap();
let t = open_table(&path).expect("a widened footer schema is not read at all");
assert_eq!(t.batches, vec![want.clone()], "byte {i} widened the read");
}
assert!(widened > 0, "the footer schema could not be widened");
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_message_that_does_not_fit_the_checked_body_is_refused() {
let d = dir("extent");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let bytes = fs::read(&path).unwrap();
let covered = footer_start(&bytes) - 8;
let buffer = Buffer::from_vec(bytes);
let schema = message_at(&path, &buffer, HEADER_LEN, covered).unwrap();
assert_eq!(schema.header, arrow_ipc::MessageHeader::Schema);
assert_eq!(schema.block.metaDataLength() % ALIGNMENT as i32, 0);
assert_eq!(schema.block.bodyLength(), 0);
let at = HEADER_LEN + schema.stride();
let rb = message_at(&path, &buffer, at, covered).unwrap();
assert_eq!(rb.header, arrow_ipc::MessageHeader::RecordBatch);
assert_eq!(rb.data.len(), rb.stride());
assert_eq!(
at + rb.stride(),
covered,
"the chain does not end at the CRC"
);
let bad = |offset: usize, body_len: usize| {
let Err(e) = message_at(&path, &buffer, offset, body_len) else {
panic!("offset {offset} framed a message inside {body_len} covered bytes")
};
assert!(
matches!(&e, Error::Io { source, .. }
if source.kind() == io::ErrorKind::InvalidData),
"{e}"
);
};
bad(covered, covered); bad(usize::MAX, covered); bad(HEADER_LEN + 1, covered); bad(HEADER_LEN, HEADER_LEN + 8); bad(at, at + rb.block.metaDataLength() as usize); let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_newer_format_version_is_refused_and_an_older_one_is_not() {
let path = Path::new("logs.arrow");
check_format(path, None).unwrap();
check_format(path, Some("1")).unwrap();
let e = check_format(path, Some("2")).unwrap_err();
assert!(
matches!(&e, Error::Io { source, .. }
if source.kind() == io::ErrorKind::Unsupported),
"{e}"
);
let e = check_format(path, Some("banana")).unwrap_err();
assert!(
matches!(&e, Error::Io { source, .. }
if source.kind() == io::ErrorKind::InvalidData),
"{e}"
);
}
#[test]
fn the_positionally_read_columns_are_pinned_to_the_format_version() {
let names: Vec<&str> = crate::schema::ATTRS
.fields()
.iter()
.map(|f| f.name().as_str())
.collect();
assert_eq!(
names,
[
"parent_id",
"key",
"type",
"str",
"int",
"double",
"bool",
"bytes",
"ser"
],
"the attribute table's column order changed, and three readers take \
those columns by index — so every block already written now decodes \
with the wrong ones. Bump FORMAT_VERSION (currently {FORMAT_VERSION}), \
put the compatibility branch for the old layout next to `check_format`, \
and update this list."
);
}
#[test]
fn a_written_table_carries_its_format_version() {
let d = dir("version");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let bytes = fs::read(&path).unwrap();
let footer = root_as_footer(&bytes[footer_start(&bytes)..bytes.len() - 10]).unwrap();
let stamped = footer
.custom_metadata()
.unwrap()
.iter()
.find(|kv| kv.key() == Some(FORMAT_KEY))
.and_then(|kv| kv.value().map(str::to_string));
assert_eq!(
stamped.as_deref(),
Some(FORMAT_VERSION.to_string().as_str())
);
assert!(open_table(&path).is_ok());
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_failed_publish_leaves_no_staging_directory() {
let root = dir("leak");
fs::write(root.join("logs"), b"not a directory").unwrap();
assert!(publish(&root, "logs", node_id("a"), 0, 0, &sealed(1_000, 2_000)).is_err());
let staged: Vec<_> = fs::read_dir(root.join(".tmp"))
.unwrap()
.map(|e| e.unwrap().path())
.collect();
assert!(staged.is_empty(), "leaked {staged:?}");
let _ = fs::remove_dir_all(&root);
}
#[test]
fn a_block_name_from_before_the_log_parses_with_no_watermark() {
let legacy = format!("{:020}-{:020}-{:08x}-{:012}", 10, 20, 0xabu32, 7u64);
assert_eq!(parse_dir_name(&legacy), Some((10, 20, 0xab, 7, 0)));
assert_eq!(
parse_dir_name(&dir_name(10, 20, 0xab, 7, 99)),
Some((10, 20, 0xab, 7, 99))
);
assert_eq!(parse_dir_name(&format!("{legacy}-1-2")), None);
}
#[test]
fn the_watermark_is_the_highest_per_signal_not_the_newest() {
let root = dir("watermark");
let node = node_id("a");
assert_eq!(wal_watermarks(&root).unwrap(), [0, 0, 0]);
publish(&root, "logs", node, 0, 40, &sealed(5_000, 6_000)).unwrap();
publish(&root, "logs", node, 1, 9, &sealed(1_000, 2_000)).unwrap();
publish(&root, "traces", node, 0, 3, &sealed(1_000, 2_000)).unwrap();
assert_eq!(wal_watermarks(&root).unwrap(), [40, 3, 0]);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn one_unreadable_block_does_not_stop_the_sweep() {
let root = dir("badblock");
let node = node_id("a");
let bad = publish(&root, "logs", node, 0, 0, &sealed(1_000, 2_000))
.unwrap()
.dir;
let good = publish(&root, "logs", node, 1, 0, &sealed(3_000, 4_000))
.unwrap()
.dir;
fs::write(bad.join("logs.arrow"), b"junk").unwrap();
assert_eq!(compact(&root, "logs", node, i64::MAX).unwrap(), 1);
assert!(good.join(COLD_MARKER).exists());
assert!(!bad.join(COLD_MARKER).exists(), "declared cold unread");
assert_eq!(expire(&root, "logs", i64::MAX).unwrap(), 2);
assert!(scan(&root, "logs").unwrap().is_empty());
let _ = fs::remove_dir_all(&root);
}
#[test]
fn one_undeletable_block_does_not_stop_retention() {
use std::os::unix::fs::PermissionsExt;
let root = dir("stuck");
let node = node_id("a");
let stuck = publish(&root, "logs", node, 0, 0, &sealed(1_000, 2_000))
.unwrap()
.dir;
let free = publish(
&root,
"logs",
node,
1,
0,
&sealed(2 * NANOS_PER_HOUR, 2 * NANOS_PER_HOUR + 1),
)
.unwrap()
.dir;
let locked = stuck.parent().unwrap().to_path_buf();
fs::set_permissions(&locked, fs::Permissions::from_mode(0o555)).unwrap();
if fs::write(locked.join("canary"), []).is_err() {
assert_eq!(expire(&root, "logs", i64::MAX).unwrap(), 1);
assert!(stuck.is_dir(), "the locked block is still there");
assert!(!free.exists(), "the block beside it was still dropped");
}
fs::set_permissions(&locked, fs::Permissions::from_mode(0o755)).unwrap();
let _ = fs::remove_dir_all(&root);
}
#[test]
fn free_fraction_is_a_fraction_of_a_real_mount() {
let d = dir("free");
let f = free_fraction(&d).unwrap();
assert!(f > 0.0 && f <= 1.0, "{f} is not a fraction");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
assert!((free_fraction(&path).unwrap() - f).abs() < 0.01);
assert!(free_fraction(&d.join("nope")).is_err());
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_path_with_a_nul_byte_is_refused_rather_than_truncated() {
use std::os::unix::ffi::OsStrExt;
let d = dir("nul");
let mut raw = d.as_os_str().as_encoded_bytes().to_vec();
raw.extend_from_slice(b"\0suffix");
let nul = PathBuf::from(std::ffi::OsStr::from_bytes(&raw));
assert!(free_fraction(&d).is_ok());
let Err(e) = free_fraction(&nul) else {
panic!("a path with a NUL byte was measured")
};
assert!(
matches!(&e, Error::Io { source, .. }
if source.kind() == io::ErrorKind::InvalidInput),
"{e}"
);
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_mount_that_reports_no_blocks_is_empty_not_full() {
let pseudo = ["/proc", "/sys", "/System/Volumes/Data/home"]
.into_iter()
.map(Path::new)
.find(|p| statfs(p).is_ok_and(|b| b.f_blocks == 0));
if let Some(p) = pseudo {
assert_eq!(free_fraction(p).unwrap(), 1.0, "{}", p.display());
}
}
#[test]
fn a_network_filesystem_refuses_to_start_and_fuse_only_warns() {
let path = Path::new("/data");
check_fs_type(path, None).unwrap();
check_fs_type(path, Some("fuse".into())).unwrap();
let e = check_fs_type(path, Some("NFS".into())).unwrap_err();
assert!(
matches!(&e, Error::NetworkFilesystem { fs, .. } if fs == "NFS"),
"{e}"
);
assert!(e.to_string().contains("NFS"), "{e}");
}
#[test]
fn a_staging_directory_that_cannot_be_removed_is_logged_not_returned() {
use std::os::unix::fs::PermissionsExt;
let root = dir("unwind");
let tmp = root.join("logs-0000002a-000000000000-0-0");
fs::create_dir(&tmp).unwrap();
fs::write(tmp.join("logs.arrow"), b"a staged table").unwrap();
fs::set_permissions(&root, fs::Permissions::from_mode(0o555)).unwrap();
if fs::create_dir(root.join("canary")).is_err() {
unwind_staging(&tmp);
assert!(tmp.is_dir(), "the undeletable staging directory went away");
}
fs::set_permissions(&root, fs::Permissions::from_mode(0o755)).unwrap();
unwind_staging(&tmp);
assert!(!tmp.exists(), "the deletable one did not");
unwind_staging(&tmp);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn a_table_with_no_batches_writes_no_file() {
let d = dir("nobatch");
let path = d.join("logs.arrow");
write_table_with(&path, &[], None).unwrap();
assert!(!path.exists(), "an empty table left a file behind");
assert!(open_table_opt(&path).unwrap().is_none());
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_source_with_neither_a_directory_nor_tables_reads_as_empty() {
let src = Src {
node: 0,
seq: 0,
min_ts: 0,
max_ts: 1,
dir: None,
tables: None,
};
assert!(src.load("logs").unwrap().is_none());
assert!(src.overlaps(0, 1));
}
#[test]
fn a_sweep_compacts_at_most_its_budget_and_the_next_one_finishes() {
let root = dir("budget");
let node = node_id("a");
let blocks = MAX_COMPACT_PER_SWEEP + 1;
for seq in 0..blocks as u64 {
publish(&root, "logs", node, seq, 0, &sealed(1_000, 2_000)).unwrap();
}
assert_eq!(
compact(&root, "logs", node, i64::MAX).unwrap(),
MAX_COMPACT_PER_SWEEP
);
let cold = |root: &Path| {
scan(root, "logs")
.unwrap()
.iter()
.filter(|b| b.dir.join(COLD_MARKER).exists())
.count()
};
assert_eq!(cold(&root), MAX_COMPACT_PER_SWEEP);
assert_eq!(compact(&root, "logs", node, i64::MAX).unwrap(), 1);
assert_eq!(cold(&root), blocks);
assert_eq!(compact(&root, "logs", node, i64::MAX).unwrap(), 0);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn a_block_that_vanished_mid_sweep_is_not_a_failure() {
let root = dir("vanished");
let node = node_id("a");
let good = publish(&root, "logs", node, 0, 0, &sealed(1_000, 2_000))
.unwrap()
.dir;
let gone = good
.parent()
.unwrap()
.join(dir_name(1_000, 2_000, node, 99, 0));
std::os::unix::fs::symlink(root.join("no-such-block"), &gone).unwrap();
assert_eq!(scan(&root, "logs").unwrap().len(), 2);
assert!(
matches!(compact_block(&gone, node), Err(Error::Io { source, .. })
if source.kind() == io::ErrorKind::NotFound),
);
assert_eq!(compact(&root, "logs", node, i64::MAX).unwrap(), 1);
assert!(good.join(COLD_MARKER).exists(), "the block beside it");
assert!(!gone.join(COLD_MARKER).exists(), "declared cold unread");
let _ = fs::remove_dir_all(&root);
}
#[test]
fn a_metadata_length_that_is_not_a_length_is_refused() {
let d = dir("metalen");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let good = fs::read(&path).unwrap();
let covered = crc_field(&good).1;
for declared in [-1i32, -8, i32::MIN, i32::MAX, i32::MAX - 7] {
let mut bytes = good.clone();
bytes[HEADER_LEN + 4..HEADER_LEN + 8].copy_from_slice(&declared.to_le_bytes());
let buffer = Buffer::from_vec(bytes);
let Err(e) = message_at(&path, &buffer, HEADER_LEN, covered) else {
panic!("a metadata length of {declared} framed a message")
};
invalid_data(&e, &format!("metadata length {declared}"));
assert!(e.to_string().contains(&declared.to_string()), "{e}");
}
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_body_that_does_not_start_with_a_schema_is_refused() {
let d = dir("noschema");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let mut bytes = fs::read(&path).unwrap();
let tag = header_tag(&bytes, HEADER_LEN);
bytes[tag] = UNKNOWN_MESSAGE;
repair_crc(&mut bytes);
fs::write(&path, &bytes).unwrap();
let Err(e) = open_table(&path) else {
panic!("a body with no schema at the head of it read as a block")
};
invalid_data(&e, "a first message that is not a schema");
assert!(e.to_string().contains("not a schema message"), "{e}");
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_message_after_the_schema_that_is_not_a_batch_is_refused() {
let d = dir("notabatch");
let path = d.join("logs.arrow");
write_table(&path, &batch()).unwrap();
let mut bytes = fs::read(&path).unwrap();
let at = second_message(&path, &bytes);
let tag = header_tag(&bytes, at);
bytes[tag] = UNKNOWN_MESSAGE;
repair_crc(&mut bytes);
fs::write(&path, &bytes).unwrap();
let Err(e) = open_table(&path) else {
panic!("a message that is neither a dictionary nor a batch decoded")
};
invalid_data(&e, "a message that is not a batch");
assert!(
e.to_string()
.contains("is not a dictionary or a record batch"),
"{e}"
);
let _ = fs::remove_dir_all(&d);
}
#[test]
fn a_zstd_frame_that_will_not_decompress_is_an_error() {
let d = dir("zstd");
let path = d.join("logs.arrow");
let want = compressible();
write_table_zstd(&path, &want).unwrap();
assert_eq!(open_table(&path).unwrap().batches, vec![want]);
let mut bytes = fs::read(&path).unwrap();
let covered = crc_field(&bytes).1;
let frame = bytes[..covered]
.windows(4)
.position(|w| w == [0x28, 0xb5, 0x2f, 0xfd])
.expect("no compressed frame in a compressed table");
bytes[frame..frame + 4].copy_from_slice(&[0xff; 4]);
repair_crc(&mut bytes);
fs::write(&path, &bytes).unwrap();
let Err(e) = open_table(&path) else {
panic!("a block with a broken ZSTD frame decoded")
};
assert!(matches!(&e, Error::Undecodable { .. }), "{e}");
assert!(e.to_string().contains("frame"), "{e}");
assert!(!e.to_string().contains("align"), "{e}");
let _ = fs::remove_dir_all(&d);
}
}